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 HuggingFace wrapper 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 huggingface wrapper

├── hf_src

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.12 - PyTorch >= 1.8 - ffmpeg

Install from PyPI: bash pip install marlin-pytorch

Load MARLIN model from online ```python from marlin_pytorch import Marlin

Load MARLIN model from GitHub Release

model = Marlin.fromonline("marlinvitbaseytf") ```

Load MARLIN model from file ```python from marlin_pytorch import Marlin

Load MARLIN model from local file

model = Marlin.fromfile("marlinvitbaseytf", "path/to/marlin.pt")

Load MARLIN model from the ckpt file trained by the scripts in this repo

model = Marlin.fromfile("marlinvitbaseytf", "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 python from marlin_pytorch import Marlin Marlin.clean_cache()

Extract features from cropped video file ```python

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.extractvideo("path/to/video.mp4", keepseq=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 ```python

Extract features from in-the-wild video with various size

features = model.extractvideo("path/to/video.mp4", cropface=True) print(features.shape) # torch.Size([T, 768]) ```

Extract features from video clip tensor ```python

Extract features from clip tensor with size (B, 3, 16, 224, 224)

x = ... # video clip features = model.extractfeatures(x) # torch.Size([B, k, 768]) features = model.extractfeatures(x, keep_seq=False) # torch.Size([B, 768]) ```

Use transformers (HuggingFace) for Feature Extraction

Requirements: - Python - PyTorch - transformers - einops

Currently the huggingface model is only for direct feature extraction without any video pre-processing (e.g. face detection, cropping, strided window, etc).

```python import torch from transformers import AutoModel

model = AutoModel.frompretrained( "ControlNet/marlinvitbaseytf", # or other variants trustremotecode=True ) tensor = torch.rand([1, 3, 16, 224, 224]) # (B, C, T, H, W) output = model(tensor) # torch.Size([1, 1568, 384]) ```

Paper Implementation

Requirements

• Python >= 3.7, < 3.12
• 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: bash 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: bash 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: bash 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

```python from marlinpytorch import Marlin from marlinpytorch.config import registermodelfrom_yaml

registermodelfromyaml("mymarlinmodel", "path/to/config.yaml") model = Marlin.fromfile("mymarlinmodel", "path/to/marlin.ckpt") ```

Evaluation

CelebV-HQ

#### 1. Download the dataset Download dataset from [CelebV-HQ](https://github.com/CelebV-HQ/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. ```bash 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 `` directory in `CelebV-HQ` directory. ```bash 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. ```bash 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. bibtex @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.