Package: marlin-pytorch
MARLIN: Masked Autoencoder for facial video Representation LearnINgThis repo is the official PyTorch implementation for the paperMARLIN: Masked Autoencoder for facial video Representation LearnINg(CVPR 2023).Repository StructureThe 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 (TODO)
├── requirements.txtUsemarlin-pytorchfor Feature ExtractionRequirements:Python >= 3.6, < 3.11PyTorch >= 1.8ffmpegInstall from PyPI:pipinstallmarlin-pytorchLoad MARLIN model from onlinefrommarlin_pytorchimportMarlin# Load MARLIN model from GitHub Releasemodel=Marlin.from_online("marlin_vit_base_ytf")Load MARLIN model from filefrommarlin_pytorchimportMarlin# Load MARLIN model from local filemodel=Marlin.from_file("marlin_vit_base_ytf","path/to/marlin.pt")# Load MARLIN model from the ckpt file trained by the scripts in this repomodel=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, seeMODEL_ZOO.md.When MARLIN model is retrieved from GitHub Release, it will be cached in.marlin. You can remove marlin cache byfrommarlin_pytorchimportMarlinMarlin.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=Truefeatures=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 = 1568Extract features from in-the-wild video file# Extract features from in-the-wild video with various sizefeatures=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 clipfeatures=model.extract_features(x)# torch.Size([B, k, 768])features=model.extract_features(x,keep_seq=False)# torch.Size([B, 768])Paper ImplementationRequirementsPython >= 3.7, < 3.11PyTorch ~= 1.11Torchvision ~= 0.12InstallationFirstly, make sure you have installed PyTorch and Torchvision with or without CUDA.Clone the repo and install the requirements:gitclonehttps://github.com/ControlNet/MARLIN.gitcdMARLIN
pipinstall-rrequirements.txtMARLIN PretrainingDownload theYoutubeFacesdataset (onlyframe_images_DBis required).Download the face parsing model fromface_parsing.farl.lapaand put it inutils/face_sdk/models/face_parsing/face_parsing_1.0.Download the VideoMAE pretrainedcheckpointfor initializing the weights. (ps. They updated their models in thiscommit, 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:pythonpreprocess/ytf_preprocess.py--data_dir/path/to/youtube_faces--max_workers8After 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.csvThen, run the training script:pythontrain.py\--configconfig/pretrain/marlin_vit_base.yaml\--data_dir/path/to/youtube_faces\--n_gpus4\--num_workers8\--batch_size16\--epochs2000\--official_pretrained/path/to/videomae/checkpoint.pthAfter trained, you can load the checkpoint for inference byfrommarlin_pytorchimportMarlinfrommarlin_pytorch.configimportregister_model_from_yamlregister_model_from_yaml("my_marlin_model","path/to/config.yaml")model=Marlin.from_file("my_marlin_model","path/to/marlin.ckpt")ReferencesIf 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},}LicenseThis project is under the CC BY-NC 4.0 license. SeeLICENSEfor details.AcknowledgementsSome code about model is based onMCG-NJU/VideoMAE. The code related to preprocessing
is borrowed fromJDAI-CV/FaceX-Zoo.