Upload 3 files

README.md

@@ -0,0 +1,14 @@
+## UCF
+
+This model has been adapted from [DeepfakeBench](https://github.com/SCLBD/DeepfakeBench).
+
+## 
+
+- **Train UCF model**:
+   - Use `train_ucf.py`, which will download necessary pretrained `xception` backbone weights from HuggingFace (if not present locally) and start a training job with logging outputs in `.logs/`.
+   - Customize the training job by editing `config/ucf.yaml`
+     - `pm2 start train_ucf.py --no-autorestart` to train a generalist detector on datasets from `DATASET_META`
+     - `pm2 start train_ucf.py --no-autorestart -- --faces_only` to train a face expert detector on preprocessed-face only datasets
+
+- **Miner Neurons**:
+   - The `UCF` class in `pretrained_ucf.py` is used by miner neurons to load and perform inference with pretrained UCF model weights.

logger.py

@@ -0,0 +1,36 @@
+import os
+import logging
+
+import torch.distributed as dist
+
+class RankFilter(logging.Filter):
+    def __init__(self, rank):
+        super().__init__()
+        self.rank = rank
+
+    def filter(self, record):
+        return dist.get_rank() == self.rank
+
+def create_logger(log_path):
+    # Create log path
+    if os.path.isdir(os.path.dirname(log_path)):
+        os.makedirs(os.path.dirname(log_path), exist_ok=True)
+
+    # Create logger object
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+    # Create file handler and set the formatter
+    fh = logging.FileHandler(log_path)
+    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
+    fh.setFormatter(formatter)
+
+    # Add the file handler to the logger
+    logger.addHandler(fh)
+
+    # Add a stream handler to print to console
+    sh = logging.StreamHandler()
+    sh.setLevel(logging.INFO)  # Set logging level for stream handler
+    sh.setFormatter(formatter)
+    logger.addHandler(sh)
+
+    return logger

train_detector.py

@@ -0,0 +1,460 @@
+# This script was adapted from the DeepfakeBench training code,
+# originally authored by Zhiyuan Yan (zhiyuanyan@link.cuhk.edu.cn)
+
+# Original: https://github.com/SCLBD/DeepfakeBench/blob/main/training/train.py
+
+# BitMind's modifications include adding a testing phase, changing the 
+# data load/split pipeline to work with subnet 34's image augmentations
+# and datasets from BitMind HuggingFace repositories, quality of life CLI args,
+# logging changes, etc.
+
+import os
+import sys
+import argparse
+from os.path import join
+import random
+import datetime
+import time
+import yaml
+from tqdm import tqdm
+import numpy as np
+from datetime import timedelta
+from copy import deepcopy
+from PIL import Image as pil_image
+from pathlib import Path
+import gc
+
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.backends.cudnn as cudnn
+import torch.utils.data
+import torch.optim as optim
+from torch.utils.data.distributed import DistributedSampler
+import torch.distributed as dist
+from torch.utils.data import DataLoader
+
+from optimizor.SAM import SAM
+from optimizor.LinearLR import LinearDecayLR
+
+from trainer.trainer import Trainer
+from arena.detectors.UCF.detectors import DETECTOR
+from metrics.utils import parse_metric_for_print
+from logger import create_logger, RankFilter
+
+from huggingface_hub import hf_hub_download
+
+# BitMind imports (not from original Deepfake Bench repo)
+from bitmind.dataset_processing.load_split_data import load_datasets, create_real_fake_datasets
+from bitmind.image_transforms import base_transforms, random_aug_transforms
+from bitmind.constants import DATASET_META, FACE_TRAINING_DATASET_META
+from config.constants import (
+    CONFIG_PATH,
+    WEIGHTS_DIR,
+    HF_REPO,
+    BACKBONE_CKPT
+)
+
+parser = argparse.ArgumentParser(description='Process some paths.')
+parser.add_argument('--detector_path', type=str, default=CONFIG_PATH, help='path to detector YAML file')
+parser.add_argument('--faces_only', dest='faces_only', action='store_true', default=False)
+parser.add_argument('--no-save_ckpt', dest='save_ckpt', action='store_false', default=True)
+parser.add_argument('--no-save_feat', dest='save_feat', action='store_false', default=True)
+parser.add_argument("--ddp", action='store_true', default=False)
+parser.add_argument('--local_rank', type=int, default=0)
+parser.add_argument('--workers', type=int, default=os.cpu_count() - 1,
+                    help='number of workers for data loading')
+parser.add_argument('--epochs', type=int, default=None, help='number of training epochs')
+
+args = parser.parse_args()
+torch.cuda.set_device(args.local_rank)
+print(f"torch.cuda.device(0): {torch.cuda.device(0)}")
+print(f"torch.cuda.get_device_name(0): {torch.cuda.get_device_name(0)}")
+
+def ensure_backbone_is_available(logger,
+                                 weights_dir=WEIGHTS_DIR,
+                                 model_filename=BACKBONE_CKPT,
+                                 hugging_face_repo_name=HF_REPO):
+    
+    destination_path = Path(weights_dir) / Path(model_filename)
+    if not destination_path.parent.exists():
+        destination_path.parent.mkdir(parents=True, exist_ok=True)
+        logger.info(f"Created directory {destination_path.parent}.")
+    if not destination_path.exists():
+        model_path = hf_hub_download(hugging_face_repo_name, model_filename)
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model = torch.load(model_path, map_location=device)
+        torch.save(model, destination_path)
+        del model
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        logger.info(f"Downloaded backbone {model_filename} to {destination_path}.")
+    else:
+        logger.info(f"{model_filename} backbone already present at {destination_path}.")
+
+def init_seed(config):
+    if config['manualSeed'] is None:
+        config['manualSeed'] = random.randint(1, 10000)
+    random.seed(config['manualSeed'])
+    if config['cuda']:
+        torch.manual_seed(config['manualSeed'])
+        torch.cuda.manual_seed_all(config['manualSeed'])
+
+def custom_collate_fn(batch):
+    images, labels, source_labels = zip(*batch)
+    
+    images = torch.stack(images, dim=0)  # Stack image tensors into a single tensor
+    labels = torch.LongTensor(labels) 
+    source_labels = torch.LongTensor(source_labels) 
+    
+    data_dict = {
+        'image': images,
+        'label': labels,
+        'label_spe': source_labels,
+        'landmark': None,
+        'mask': None
+    }    
+    return data_dict
+
+def prepare_datasets(config, logger):
+    start_time = log_start_time(logger, "Loading and splitting individual datasets")
+    
+    real_datasets, fake_datasets = load_datasets(dataset_meta=config['dataset_meta'], 
+                                                 expert=config['faces_only'],
+                                                 split_transforms=config['split_transforms'])
+
+    log_finish_time(logger, "Loading and splitting individual datasets", start_time)
+    
+    start_time = log_start_time(logger, "Creating real fake dataset splits")
+    train_dataset, val_dataset, test_dataset = \
+    create_real_fake_datasets(real_datasets,
+                              fake_datasets,
+                              config['split_transforms']['train']['transform'],
+                              config['split_transforms']['validation']['transform'],
+                              config['split_transforms']['test']['transform'],
+                              source_labels=True)
+
+    log_finish_time(logger, "Creating real fake dataset splits", start_time)
+    
+    train_loader = torch.utils.data.DataLoader(train_dataset,
+                                               batch_size=config['train_batchSize'],
+                                               shuffle=True,
+                                               num_workers=config['workers'],
+                                               drop_last=True,
+                                               collate_fn=custom_collate_fn)
+    val_loader = torch.utils.data.DataLoader(val_dataset,
+                                             batch_size=config['train_batchSize'],
+                                             shuffle=True,
+                                             num_workers=config['workers'],
+                                             drop_last=True,
+                                             collate_fn=custom_collate_fn)
+    test_loader = torch.utils.data.DataLoader(test_dataset,
+                                              batch_size=config['train_batchSize'],
+                                              shuffle=True, 
+                                              num_workers=config['workers'],
+                                              drop_last=True,
+                                              collate_fn=custom_collate_fn)
+
+    print(f"Train size: {len(train_loader.dataset)}")
+    print(f"Validation size: {len(val_loader.dataset)}")
+    print(f"Test size: {len(test_loader.dataset)}")
+
+    return train_loader, val_loader, test_loader
+
+def choose_optimizer(model, config):
+    opt_name = config['optimizer']['type']
+    if opt_name == 'sgd':
+        optimizer = optim.SGD(
+            params=model.parameters(),
+            lr=config['optimizer'][opt_name]['lr'],
+            momentum=config['optimizer'][opt_name]['momentum'],
+            weight_decay=config['optimizer'][opt_name]['weight_decay']
+        )
+        return optimizer
+    elif opt_name == 'adam':
+        optimizer = optim.Adam(
+            params=model.parameters(),
+            lr=config['optimizer'][opt_name]['lr'],
+            weight_decay=config['optimizer'][opt_name]['weight_decay'],
+            betas=(config['optimizer'][opt_name]['beta1'], config['optimizer'][opt_name]['beta2']),
+            eps=config['optimizer'][opt_name]['eps'],
+            amsgrad=config['optimizer'][opt_name]['amsgrad'],
+        )
+        return optimizer
+    elif opt_name == 'sam':
+        optimizer = SAM(
+            model.parameters(), 
+            optim.SGD, 
+            lr=config['optimizer'][opt_name]['lr'],
+            momentum=config['optimizer'][opt_name]['momentum'],
+        )
+    else:
+        raise NotImplementedError('Optimizer {} is not implemented'.format(config['optimizer']))
+    return optimizer
+
+
+def choose_scheduler(config, optimizer):
+    if config['lr_scheduler'] is None:
+        return None
+    elif config['lr_scheduler'] == 'step':
+        scheduler = optim.lr_scheduler.StepLR(
+            optimizer,
+            step_size=config['lr_step'],
+            gamma=config['lr_gamma'],
+        )
+        return scheduler
+    elif config['lr_scheduler'] == 'cosine':
+        scheduler = optim.lr_scheduler.CosineAnnealingLR(
+            optimizer,
+            T_max=config['lr_T_max'],
+            eta_min=config['lr_eta_min'],
+        )
+        return scheduler
+    elif config['lr_scheduler'] == 'linear':
+        scheduler = LinearDecayLR(
+            optimizer,
+            config['nEpochs'],
+            int(config['nEpochs']/4),
+        )
+    else:
+        raise NotImplementedError('Scheduler {} is not implemented'.format(config['lr_scheduler']))
+
+def choose_metric(config):
+    metric_scoring = config['metric_scoring']
+    if metric_scoring not in ['eer', 'auc', 'acc', 'ap']:
+        raise NotImplementedError('metric {} is not implemented'.format(metric_scoring))
+    return metric_scoring
+
+def log_start_time(logger, process_name):
+    """Log the start time of a process."""
+    start_time = time.time()
+    logger.info(f"{process_name} Start Time: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time))}")
+    return start_time
+
+def log_finish_time(logger, process_name, start_time):
+    """Log the finish time and elapsed time of a process."""
+    finish_time = time.time()
+    elapsed_time = finish_time - start_time
+
+    # Convert elapsed time into hours, minutes, and seconds
+    hours, rem = divmod(elapsed_time, 3600)
+    minutes, seconds = divmod(rem, 60)
+
+    # Log the finish time and elapsed time
+    logger.info(f"{process_name} Finish Time: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(finish_time))}")
+    logger.info(f"{process_name} Elapsed Time: {int(hours)} hours, {int(minutes)} minutes, {seconds:.2f} seconds")
+
+def save_config(config, outputs_dir):
+    """
+    Saves a config dictionary as both a pickle file and a YAML file, ensuring only basic types are saved.
+    Also, lists like 'mean' and 'std' are saved in flow style (on a single line).
+    
+    Args:
+        config (dict): The configuration dictionary to save.
+        outputs_dir (str): The directory path where the files will be saved.
+    """
+
+    def is_basic_type(value):
+        """
+        Check if a value is a basic data type that can be saved in YAML.
+        Basic types include int, float, str, bool, list, and dict.
+        """
+        return isinstance(value, (int, float, str, bool, list, dict, type(None)))
+
+    def filter_dict(data_dict):
+        """
+        Recursively filter out any keys from the dictionary whose values contain non-basic types (e.g., objects).
+        """
+        if not isinstance(data_dict, dict):
+            return data_dict
+        
+        filtered_dict = {}
+        for key, value in data_dict.items():
+            if isinstance(value, dict):
+                # Recursively filter nested dictionaries
+                nested_dict = filter_dict(value)
+                if nested_dict:  # Only add non-empty dictionaries
+                    filtered_dict[key] = nested_dict
+            elif is_basic_type(value):
+                # Add if the value is a basic type
+                filtered_dict[key] = value
+            else:
+                # Skip the key if the value is not a basic type (e.g., an object)
+                print(f"Skipping key '{key}' because its value is of type {type(value)}")
+        
+        return filtered_dict
+
+    def save_dict_to_yaml(data_dict, file_path):
+        """
+        Saves a dictionary to a YAML file, excluding any keys where the value is an object or contains an object.
+        Additionally, ensures that specific lists (like 'mean' and 'std') are saved in flow style.
+        
+        Args:
+            data_dict (dict): The dictionary to save.
+            file_path (str): The local file path where the YAML file will be saved.
+        """
+        
+        # Custom representer for lists to force flow style (compact lists)
+        class FlowStyleList(list):
+            pass
+        
+        def flow_style_list_representer(dumper, data):
+            return dumper.represent_sequence('tag:yaml.org,2002:seq', data, flow_style=True)
+        
+        yaml.add_representer(FlowStyleList, flow_style_list_representer)
+
+        # Preprocess specific lists to be in flow style
+        if 'mean' in data_dict:
+            data_dict['mean'] = FlowStyleList(data_dict['mean'])
+        if 'std' in data_dict:
+            data_dict['std'] = FlowStyleList(data_dict['std'])
+
+        try:
+            # Filter the dictionary
+            filtered_dict = filter_dict(data_dict)
+            
+            # Save the filtered dictionary as YAML
+            with open(file_path, 'w') as f:
+                yaml.dump(filtered_dict, f, default_flow_style=False)  # Save with default block style except for FlowStyleList
+            print(f"Filtered dictionary successfully saved to {file_path}")
+        except Exception as e:
+            print(f"Error saving dictionary to YAML: {e}")
+
+    # Save as YAML
+    save_dict_to_yaml(config, outputs_dir + '/config.yaml')
+
+def main():
+    torch.cuda.empty_cache()
+    gc.collect()
+    # parse options and load config
+    with open(args.detector_path, 'r') as f:
+        config = yaml.safe_load(f)
+    with open(os.getcwd() + '/config/train_config.yaml', 'r') as f:
+        config2 = yaml.safe_load(f)
+    if 'label_dict' in config:
+        config2['label_dict']=config['label_dict']
+    config.update(config2)
+
+    config['workers'] = args.workers
+    
+    config['local_rank']=args.local_rank
+    if config['dry_run']:
+        config['nEpochs'] = 0
+        config['save_feat']=False
+
+    if args.epochs: config['nEpochs'] = args.epochs
+
+    config['split_transforms'] = {'train': {'name': 'base_transforms',
+                                            'transform': base_transforms},
+                                  'validation': {'name': 'base_transforms',
+                                                 'transform': base_transforms},
+                                  'test': {'name': 'base_transforms',
+                                           'transform': base_transforms}}
+    config['faces_only'] = args.faces_only
+    config['dataset_meta'] = FACE_TRAINING_DATASET_META if config['faces_only'] else DATASET_META
+    dataset_names = [item["path"] for datasets in config['dataset_meta'].values() for item in datasets]
+    config['train_dataset'] = dataset_names
+    config['save_ckpt'] = args.save_ckpt
+    config['save_feat'] = args.save_feat
+
+    config['specific_task_number'] = len(config['dataset_meta']["fake"]) + 1
+    
+    if config['lmdb']:
+        config['dataset_json_folder'] = 'preprocessing/dataset_json_v3'
+        
+    # create logger
+    timenow=datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
+    
+    outputs_dir =  os.path.join(
+                config['log_dir'],
+                config['model_name'] + '_' + timenow
+            )
+    
+    os.makedirs(outputs_dir, exist_ok=True)
+    logger = create_logger(os.path.join(outputs_dir, 'training.log'))
+    config['log_dir'] = outputs_dir
+    logger.info('Save log to {}'.format(outputs_dir))
+    
+    config['ddp']= args.ddp
+
+    # init seed
+    init_seed(config)
+
+    # set cudnn benchmark if needed
+    if config['cudnn']:
+        cudnn.benchmark = True
+    if config['ddp']:
+        # dist.init_process_group(backend='gloo')
+        dist.init_process_group(
+            backend='nccl',
+            timeout=timedelta(minutes=30)
+        )
+        logger.addFilter(RankFilter(0))
+
+    ensure_backbone_is_available(logger=logger,
+                                 model_filename=config['pretrained'].split('/')[-1],
+                                 hugging_face_repo_name='bitmind/' + config['model_name'])
+    
+    # prepare the model (detector)
+    model_class = DETECTOR[config['model_name']]
+    model = model_class(config)
+    
+    # prepare the optimizer
+    optimizer = choose_optimizer(model, config)
+
+    # prepare the scheduler
+    scheduler = choose_scheduler(config, optimizer)
+
+    # prepare the metric
+    metric_scoring = choose_metric(config)
+
+    # prepare the trainer
+    trainer = Trainer(config, model, optimizer, scheduler, logger, metric_scoring)
+
+    # prepare the data loaders
+    train_loader, val_loader, test_loader = prepare_datasets(config, logger)
+
+    # print configuration
+    logger.info("--------------- Configuration ---------------")
+    params_string = "Parameters: \n"
+    for key, value in config.items():
+        params_string += "{}: {}".format(key, value) + "\n"
+    logger.info(params_string)
+    
+    # save training configs
+    save_config(config, outputs_dir)
+    
+    # start training
+    start_time = log_start_time(logger, "Training")
+    for epoch in range(config['start_epoch'], config['nEpochs'] + 1):
+        trainer.model.epoch = epoch
+        best_metric = trainer.train_epoch(
+                    epoch,
+                    train_data_loader=train_loader,
+                    validation_data_loaders={'val':val_loader}
+                )
+        if best_metric is not None:
+            logger.info(f"===> Epoch[{epoch}] end with validation {metric_scoring}: {parse_metric_for_print(best_metric)}!")
+    logger.info("Stop Training on best Validation metric {}".format(parse_metric_for_print(best_metric))) 
+    log_finish_time(logger, "Training", start_time)
+   
+    # test
+    start_time = log_start_time(logger, "Test")
+    trainer.eval(eval_data_loaders={'test':test_loader}, eval_stage="test")
+    log_finish_time(logger, "Test", start_time)
+    
+    # update
+    if 'svdd' in config['model_name']:
+        model.update_R(epoch)
+    if scheduler is not None:
+        scheduler.step()
+
+    # close the tensorboard writers
+    for writer in trainer.writers.values():
+        writer.close()
+
+    torch.cuda.empty_cache()
+    gc.collect()
+
+if __name__ == '__main__':
+    main()