thelukadragar/marlin_vra

MARLIN: Masked Autoencoder for facial video Representation LearnINg

This repo is the official PyTorch implementation for the paper MARLIN: Masked Autoencoder for facial video Representation LearnINg (CVPR 2023).

Repository Structure

The repository contains 2 parts:

• marlin-pytorch: The PyPI package for MARLIN used for inference.
• The implementation for the paper including training and evaluation scripts.

.
├── assets                # Images for README.md
├── LICENSE
├── README.md
├── MODEL_ZOO.md
├── CITATION.cff
├── .gitignore
├── .github

# below is for the PyPI package marlin-pytorch
├── src                   # Source code for marlin-pytorch
├── tests                 # Unittest
├── requirements.lib.txt
├── setup.py
├── init.py
├── version.txt

# below is for the paper implementation
├── configs              # Configs for experiments settings
├── model                # Marlin models
├── preprocess           # Preprocessing scripts
├── dataset              # Dataloaders
├── utils                # Utility functions
├── train.py             # Training script
├── evaluate.py          # Evaluation script
├── requirements.txt

Use marlin-pytorch for Feature Extraction

Requirements:

• Python >= 3.6, < 3.11
• PyTorch >= 1.8
• ffmpeg

Install from PyPI:

pip install marlin-pytorch

Load MARLIN model from online

from marlin_pytorch import Marlin
# Load MARLIN model from GitHub Release
model = Marlin.from_online("marlin_vit_base_ytf")

Load MARLIN model from file

from marlin_pytorch import Marlin
# Load MARLIN model from local file
model = Marlin.from_file("marlin_vit_base_ytf", "path/to/marlin.pt")
# Load MARLIN model from the ckpt file trained by the scripts in this repo
model = Marlin.from_file("marlin_vit_base_ytf", "path/to/marlin.ckpt")

Current model name list:

• marlin_vit_small_ytf: ViT-small encoder trained on YTF dataset. Embedding 384 dim.
• marlin_vit_base_ytf: ViT-base encoder trained on YTF dataset. Embedding 768 dim.
• marlin_vit_large_ytf: ViT-large encoder trained on YTF dataset. Embedding 1024 dim.

For more details, see MODEL_ZOO.md.

When MARLIN model is retrieved from GitHub Release, it will be cached in .marlin. You can remove marlin cache by

from marlin_pytorch import Marlin
Marlin.clean_cache()

Extract features from cropped video file

# Extract features from facial cropped video with size (224x224)
features = model.extract_video("path/to/video.mp4")
print(features.shape)  # torch.Size([T, 768]) where T is the number of windows

# You can keep output of all elements from the sequence by setting keep_seq=True
features = model.extract_video("path/to/video.mp4", keep_seq=True)
print(features.shape)  # torch.Size([T, k, 768]) where k = T/t * H/h * W/w = 8 * 14 * 14 = 1568

Extract features from in-the-wild video file

# Extract features from in-the-wild video with various size
features = model.extract_video("path/to/video.mp4", crop_face=True)
print(features.shape)  # torch.Size([T, 768])

Extract features from video clip tensor

# Extract features from clip tensor with size (B, 3, 16, 224, 224)
x = ...  # video clip
features = model.extract_features(x)  # torch.Size([B, k, 768])
features = model.extract_features(x, keep_seq=False)  # torch.Size([B, 768])

Paper Implementation

Requirements

• Python >= 3.7, < 3.11
• PyTorch ~= 1.11
• Torchvision ~= 0.12

Installation

Firstly, make sure you have installed PyTorch and Torchvision with or without CUDA.

Clone the repo and install the requirements:

git clone https://github.com/ControlNet/MARLIN.git
cd MARLIN
pip install -r requirements.txt

MARLIN Pretraining

Download the YoutubeFaces dataset (only frame_images_DB is required).

Download the face parsing model from face_parsing.farl.lapa and put it in utils/face_sdk/models/face_parsing/face_parsing_1.0.

Download the VideoMAE pretrained checkpoint for initializing the weights. (ps. They updated their models in this commit, but we are using the old models which are not shared anymore by the authors. So we uploaded this model by ourselves.)

Then run scripts to process the dataset:

python preprocess/ytf_preprocess.py --data_dir /path/to/youtube_faces --max_workers 8

After processing, the directory structure should be like this:

├── YoutubeFaces
│   ├── frame_images_DB
│   │   ├── Aaron_Eckhart
│   │   │   ├── 0
│   │   │   │   ├── 0.555.jpg
│   │   │   │   ├── ...
│   │   │   ├── ...
│   │   ├── ...
│   ├── crop_images_DB
│   │   ├── Aaron_Eckhart
│   │   │   ├── 0
│   │   │   │   ├── 0.555.jpg
│   │   │   │   ├── ...
│   │   │   ├── ...
│   │   ├── ...
│   ├── face_parsing_images_DB
│   │   ├── Aaron_Eckhart
│   │   │   ├── 0
│   │   │   │   ├── 0.555.npy
│   │   │   │   ├── ...
│   │   │   ├── ...
│   │   ├── ...
│   ├── train_set.csv
│   ├── val_set.csv

Then, run the training script:

python train.py \
    --config config/pretrain/marlin_vit_base.yaml \
    --data_dir /path/to/youtube_faces \
    --n_gpus 4 \
    --num_workers 8 \
    --batch_size 16 \
    --epochs 2000 \
    --official_pretrained /path/to/videomae/checkpoint.pth

After trained, you can load the checkpoint for inference by

from marlin_pytorch import Marlin
from marlin_pytorch.config import register_model_from_yaml

register_model_from_yaml("my_marlin_model", "path/to/config.yaml")
model = Marlin.from_file("my_marlin_model", "path/to/marlin.ckpt")

Evaluation

CelebV-HQ

1. Download the dataset

Download dataset from CelebV-HQ and the file structure should be like this:

├── CelebV-HQ
│   ├── downloaded
│   │   ├── ***.mp4
│   │   ├── ...
│   ├── celebvhq_info.json
│   ├── ...

2. Preprocess the dataset

Crop the face region from the raw video and split the train val and test sets.

python preprocess/celebvhq_preprocess.py --data_dir /path/to/CelebV-HQ 

3. Extract MARLIN features (Optional, if linear probing)

Extract MARLIN features from the cropped video and saved to <backbone> directory in CelebV-HQ directory.

python preprocess/celebvhq_extract.py --data_dir /path/to/CelebV-HQ --backbone marlin_vit_base_ytf

4. Train and evaluate

Train and evaluate the model adapted from MARLIN to CelebV-HQ.

Please use the configs in config/celebv_hq/*/*.yaml as the config file.

python evaluate.py \
    --config /path/to/config \
    --data_path /path/to/CelebV-HQ 
    --num_workers 4 
    --batch_size 16

License

This project is under the CC BY-NC 4.0 license. See LICENSE for details.

References

If you find this work useful for your research, please consider citing it.

@inproceedings{cai2022marlin,
  title = {MARLIN: Masked Autoencoder for facial video Representation LearnINg},
  author = {Cai, Zhixi and Ghosh, Shreya and Stefanov, Kalin and Dhall, Abhinav and Cai, Jianfei and Rezatofighi, Hamid and Haffari, Reza and Hayat, Munawar},
  booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  year = {2023},
  month = {June},
  pages = {1493-1504},
  doi = {10.1109/CVPR52729.2023.00150},
  publisher = {IEEE},
}

Acknowledgements

Some code about model is based on MCG-NJU/VideoMAE. The code related to preprocessing is borrowed from JDAI-CV/FaceX-Zoo.