|
--- |
|
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). |
