Refactor Directory Structure and Add SPSL Model

README.md

@@ -56,7 +56,7 @@ conda activate arena
 
 With your virtual env activated, you can install `dfd-arena`:
 ```bash
-cd dfd-arena && pip install -e
+cd dfd-arena && pip install -e .
 ```
 
 ## Usage

arena/architectures/SPSL/config/__init__.py

@@ -0,0 +1,7 @@
+import os
+import sys
+current_file_path = os.path.abspath(__file__)
+parent_dir = os.path.dirname(os.path.dirname(current_file_path))
+project_root_dir = os.path.dirname(parent_dir)
+sys.path.append(parent_dir)
+sys.path.append(project_root_dir)

arena/architectures/SPSL/config/constants.py

@@ -0,0 +1,13 @@
+import os
+
+# Path to the directory containing the constants.py file
+CONFIGS_DIR = os.path.dirname(os.path.abspath(__file__))
+
+# The base directory for related files
+BASE_PATH = os.path.abspath(os.path.join(CONFIGS_DIR, ".."))
+# Absolute paths for the required files and directories
+CONFIG_PATH = os.path.join(CONFIGS_DIR, "spsl.yaml")  # Path to the .yaml file
+WEIGHTS_DIR = os.path.join(BASE_PATH, "weights/") # Path to pretrained weights directory
+
+HF_REPO = "bitmind/spsl"
+BACKBONE_CKPT = "xception_best.pth"

arena/architectures/SPSL/config/spsl.yaml

@@ -0,0 +1,88 @@
+# log dir 
+log_dir: /mntcephfs/lab_data/zhiyuanyan/benchmark_results/logs_final/spsl_4frames
+
+# model setting
+pretrained: ../weights/xception_best.pth   # path to a pre-trained model, if using one
+# pretrained: /home/tianshuoge/resnet34-b627a593.pth   # path to a pre-trained model, if using one
+model_name: spsl   # model name
+backbone_name: xception  # backbone name
+
+#backbone setting
+backbone_config:
+  mode: original  # shallow_xception
+  num_classes: 2
+  inc: 4
+  dropout: false
+
+# dataset
+all_dataset: [FaceForensics++, FF-F2F, FF-DF, FF-FS, FF-NT, FaceShifter, DeepFakeDetection, Celeb-DF-v1, Celeb-DF-v2, DFDCP, DFDC, DeeperForensics-1.0, UADFV]
+train_dataset: [FF-FS]
+test_dataset: [FaceForensics++, FF-F2F, FF-DF, FF-FS, FF-NT]
+
+compression: c23  # compression-level for videos
+train_batchSize: 32   # training batch size
+test_batchSize: 32   # test batch size
+workers: 8   # number of data loading workers
+frame_num: {'train': 4, 'test': 32}   # number of frames to use per video in training and testing
+resolution: 256   # resolution of output image to network
+with_mask: false   # whether to include mask information in the input
+with_landmark: false   # whether to include facial landmark information in the input
+save_ckpt: true   # whether to save checkpoint
+save_feat: true   # whether to save features
+
+
+# data augmentation
+use_data_augmentation: true  # Add this flag to enable/disable data augmentation
+data_aug:
+  flip_prob: 0.5
+  rotate_prob: 0.5
+  rotate_limit: [-10, 10]
+  blur_prob: 0.5
+  blur_limit: [3, 7]
+  brightness_prob: 0.5
+  brightness_limit: [-0.1, 0.1]
+  contrast_limit: [-0.1, 0.1]
+  quality_lower: 40
+  quality_upper: 100
+
+# mean and std for normalization
+mean: [0.5, 0.5, 0.5]
+std: [0.5, 0.5, 0.5]
+
+# optimizer config
+optimizer:
+  # choose between 'adam' and 'sgd'
+  type: adam
+  adam:
+    lr: 0.0002  # learning rate
+    beta1: 0.9  # beta1 for Adam optimizer
+    beta2: 0.999 # beta2 for Adam optimizer
+    eps: 0.00000001  # epsilon for Adam optimizer
+    weight_decay: 0.0005  # weight decay for regularization
+    amsgrad: false
+  sgd:
+    lr: 0.0002  # learning rate
+    momentum: 0.9  # momentum for SGD optimizer
+    weight_decay: 0.0005  # weight decay for regularization
+
+# training config
+lr_scheduler: null   # learning rate scheduler
+nEpochs: 10   # number of epochs to train for
+start_epoch: 0   # manual epoch number (useful for restarts)
+save_epoch: 1   # interval epochs for saving models
+rec_iter: 100   # interval iterations for recording
+logdir: ./logs   # folder to output images and logs
+manualSeed: 1024   # manual seed for random number generation
+save_ckpt: false   # whether to save checkpoint
+
+# loss function
+loss_func: cross_entropy   # loss function to use
+losstype: null
+
+# metric
+metric_scoring: auc   # metric for evaluation (auc, acc, eer, ap)
+
+# cuda
+
+cuda: true   # whether to use CUDA acceleration
+cudnn: true   # whether to use CuDNN for convolution operations

arena/architectures/SPSL/config/train_config.yaml

@@ -0,0 +1,43 @@
+mode: train
+lmdb: True
+dry_run: false
+rgb_dir: './datasets/rgb'
+lmdb_dir:  './datasets/lmdb'
+dataset_json_folder: './preprocessing/dataset_json'
+SWA: False
+save_avg: True
+log_dir: ./logs/training/
+# label settings
+label_dict:
+  # DFD
+  DFD_fake: 1
+  DFD_real: 0
+  # FF++ + FaceShifter(FF-real+FF-FH)
+  FF-SH: 1
+  FF-F2F: 1
+  FF-DF: 1
+  FF-FS: 1
+  FF-NT: 1
+  FF-FH: 1
+  FF-real: 0
+  # CelebDF
+  CelebDFv1_real: 0
+  CelebDFv1_fake: 1
+  CelebDFv2_real: 0
+  CelebDFv2_fake: 1
+  # DFDCP
+  DFDCP_Real: 0
+  DFDCP_FakeA: 1
+  DFDCP_FakeB: 1
+  # DFDC
+  DFDC_Fake: 1
+  DFDC_Real: 0
+  # DeeperForensics-1.0
+  DF_fake: 1
+  DF_real: 0
+  # UADFV
+  UADFV_Fake: 1
+  UADFV_Real: 0
+  # Roop
+  roop_Real: 0
+  roop_Fake: 1

arena/architectures/SPSL/detectors/__init__.py

@@ -0,0 +1,11 @@
+import os
+import sys
+current_file_path = os.path.abspath(__file__)
+parent_dir = os.path.dirname(os.path.dirname(current_file_path))
+project_root_dir = os.path.dirname(parent_dir)
+sys.path.append(parent_dir)
+sys.path.append(project_root_dir)
+
+from metrics.registry import DETECTOR
+
+from .spsl_detector import SpslDetector

arena/architectures/SPSL/detectors/base_detector.py

@@ -0,0 +1,71 @@
+# author: Zhiyuan Yan
+# email: zhiyuanyan@link.cuhk.edu.cn
+# date: 2023-0706
+# description: Abstract Class for the Deepfake Detector
+
+import abc
+import torch
+import torch.nn as nn
+from typing import Union
+
+class AbstractDetector(nn.Module, metaclass=abc.ABCMeta):
+    """
+    All deepfake detectors should subclass this class.
+    """
+    def __init__(self, config=None, load_param: Union[bool, str] = False):
+        """
+        config:   (dict)
+            configurations for the model
+        load_param:  (False | True | Path(str))
+            False Do not read; True Read the default path; Path Read the required path
+        """
+        super().__init__()
+
+    @abc.abstractmethod
+    def features(self, data_dict: dict) -> torch.tensor:
+        """
+        Returns the features from the backbone given the input data.
+        """
+        pass
+
+    @abc.abstractmethod
+    def forward(self, data_dict: dict, inference=False) -> dict:
+        """
+        Forward pass through the model, returning the prediction dictionary.
+        """
+        pass
+
+    @abc.abstractmethod
+    def classifier(self, features: torch.tensor) -> torch.tensor:
+        """
+        Classifies the features into classes.
+        """
+        pass
+
+    @abc.abstractmethod
+    def build_backbone(self, config):
+        """
+        Builds the backbone of the model.
+        """
+        pass
+
+    @abc.abstractmethod
+    def build_loss(self, config):
+        """
+        Builds the loss function for the model.
+        """
+        pass
+
+    @abc.abstractmethod
+    def get_losses(self, data_dict: dict, pred_dict: dict) -> dict:
+        """
+        Returns the losses for the model.
+        """
+        pass
+
+    @abc.abstractmethod
+    def get_train_metrics(self, data_dict: dict, pred_dict: dict) -> dict:
+        """
+        Returns the training metrics for the model.
+        """
+        pass