README.md

---
license: cc
base_model:
- MCG-NJU/videomae-base
- MCG-NJU/videomae-large
---
# 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](https://openaccess.thecvf.com/content/CVPR2023/html/Cai_MARLIN_Masked_Autoencoder_for_Facial_Video_Representation_LearnINg_CVPR_2023_paper) (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.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.from_online("marlin_vit_base_ytf")
```

Load MARLIN model from file
```python
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](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.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
```python
# 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
```python
# 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.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](https://www.cs.tau.ac.il/~wolf/ytfaces/) dataset (only `frame_images_DB` is required). 

Download the face parsing model from [face_parsing.farl.lapa](https://github.com/FacePerceiver/facer/releases/download/models-v1/face_parsing.farl.lapa.main_ema_136500_jit191.pt)
and put it in `utils/face_sdk/models/face_parsing/face_parsing_1.0`.

Download the VideoMAE pretrained [checkpoint](https://github.com/ControlNet/MARLIN/releases/misc) 
for initializing the weights. (ps. They updated their models in this 
[commit](https://github.com/MCG-NJU/VideoMAE/commit/2b56a75d166c619f71019e3d1bb1c4aedafe7a90), 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 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

<details>
<summary>CelebV-HQ</summary>

#### 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 `<backbone>` 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
```

</details>


## License

This project is under the CC BY-NC 4.0 license. See [LICENSE](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},
}
```

The arxiv version available: https://arxiv.org/abs/2211.06627

## Acknowledgements

Some code about model is based on [MCG-NJU/VideoMAE](https://github.com/MCG-NJU/VideoMAE). The code related to preprocessing
is borrowed from [JDAI-CV/FaceX-Zoo](https://github.com/JDAI-CV/FaceX-Zoo).