+
+
+
+
+## 🔥 News
+* [2025-01-23]🎉🎉🎉 AIDE is accepted by ICLR 2025.
+* [2024-12-29]🔥🔥🔥 We release the Chamelon dataset.
+* [2024-06-20]🔥🔥🔥 We release the code and checkpoints of AIDE.
+
+
+## 🔍 Chameleon
+
+
+**License**:
+```
+Chameleon is only used for academic research. Commercial use in any form is prohibited.
+```
+
+🌟🌟🌟 If you need the Chameleon dataset, please send an email to **tattoo.ysl@gmail.com**. 🔥🔥🔥
+
+
+
+**Comparison of `Chameleon` with existing benchmarks.**
+
+
+
+We visualize two contemporary AI-generated image benchmarks, namely:
+
+- **(a) AIGCDetect Benchmark**
+- **(b) GenImage Benchmark**
+
+where all images are generated from publicly available generators, such as ProGAN (GAN-based), SD v1.4 (DM-based), and Midjourney (commercial API). These images are generated by unconditional situations or conditioned on simple prompts (e.g., *photo of a plane*) without delicate manual adjustments, thereby inclined to generate obvious artifacts in consistency and semantics (marked with red boxes).
+
+In contrast, our **`Chameleon`** dataset in **(c)** aims to simulate real-world scenarios by collecting diverse images from online websites, where these online images are carefully adjusted by photographers and AI artists.
+
+
+
+## 👀 Method
+
+We conduct a sanity check on **"whether the task of AI-generated image detection has been solved"**. To start with, we present **Chameleon** dataset, consisting AI-generated images that are genuinely challenging for human perception. To quantify the generalization of existing methods, we evaluate 9 off-the-shelf AI-generated image detectors on **Chameleon** dataset. Upon analysis, almost all models classify AI-generated images as real ones. Later, we propose **AIDE**~(**A**I-generated **I**mage **DE**tector with Hybrid Features), which leverages multiple experts to simultaneously extract visual artifacts and noise patterns.
+
+
+
+
+
+
+## Requirements
+
+We test the codes in the following environments, other versions may also be compatible:
+
+- CUDA 11.8
+- Python 3.10
+- Pytorch 2.0.1
+
+
+## Setup
+
+First, clone the repository locally.
+
+```
+https://github.com/shilinyan99/AIDE
+```
+
+Then, install Pytorch 2.0.1 using the conda environment.
+```
+conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 -c pytorch
+```
+
+Lastly, install the necessary packages and pycocotools.
+
+```
+pip install -r requirements.txt
+```
+
+
+## Get Started
+
+### Training
+
+```
+./scripts/train.sh --data_path [/path/to/train_data] --eval_data_path [/path/to/eval_data] --resnet_path [/path/to/pretrained_resnet_path] --convnext_path [/path/to/pretrained_convnext_path] --output_dir [/path/to/output_dir] [other args]
+```
+
+For example, training on ProGAN, run the following command:
+
+```
+./scripts/train.sh --data_path dataset/progan/train --eval_data_path dataset/progan/eval --resnet_path pretrained_ckpts/resnet50.pth --convnext_path pretrained_ckpts/open_clip_pytorch_model.bin --output_dir results/progan_train
+```
+
+### Inference
+
+Inference using the trained model.
+```
+./scripts/eval.sh --data_path [/path/to/train_data] --eval_data_path [/path/to/eval_data] --resume [/path/to/progan_train] --eval True --output_dir [/path/to/output_dir]
+```
+
+For example, evaluating the progan_train model, run the following command:
+
+```
+./scripts/eval.sh --data_path dataset/progan/train --eval_data_path dataset/progan/eval --resume results/progan_train/progan_train.pth --eval True --output_dir results/progan_train
+```
+
+
+
+## Dataset
+
+### Training Set
+We adopt the training set in [CNNSpot](https://github.com/peterwang512/CNNDetection) and [GenImage](https://github.com/Andrew-Zhu/GenImage).
+
+### Test Set
+The whole test set we used in our experiments can be downloaded from [AIGCDetectBenchmark](https://github.com/Ekko-zn/AIGCDetectBenchmark?tab=readme-ov-file) and [GenImage](https://github.com/Andrew-Zhu/GenImage).
+
+
+## Model Zoo
+
+Our training checkpoints can be downloaded from [link](https://drive.google.com/drive/folders/1qx76UFvDpgCxaPLBCmsA2WY-SSzeJrd4?usp=sharing).
+
+## Acknowledgement
+
+This repo is based on [ConvNeXt](https://github.com/facebookresearch/ConvNeXt-V2). We also refer to the repositories [CNNSpot](https://github.com/peterwang512/CNNDetection)、[AIGCDetectBenchmark](https://github.com/Ekko-zn/AIGCDetectBenchmark?tab=readme-ov-file)、[GenImage](https://github.com/Andrew-Zhu/GenImage) and [DNF](https://github.com/YichiCS/DNF). Thanks for their wonderful works.
+
+## Citation
+
+```
+@article{yan2024sanity,
+ title={A Sanity Check for AI-generated Image Detection},
+ author={Yan, Shilin and Li, Ouxiang and Cai, Jiayin and Hao, Yanbin and Jiang, Xiaolong and Hu, Yao and Xie, Weidi},
+ journal={arXiv preprint arXiv:2406.19435},
+ year={2024}
+}
+```
+
+## Contact
+If you have any question about this project, please feel free to contact tattoo.ysl@gmail.com.
diff --git a/clean/image/aide/SOURCE.md b/clean/image/aide/SOURCE.md
new file mode 100644
index 0000000000000000000000000000000000000000..92fa373cbbeadc1b2334d39e49372e616a822213
--- /dev/null
+++ b/clean/image/aide/SOURCE.md
@@ -0,0 +1,17 @@
+# Source: image/aide
+
+| Field | Value |
+|---|---|
+| Upstream | https://github.com/shilinyan99/AIDE |
+| Paper | https://arxiv.org/abs/2406.19435 |
+| Commit SHA | **not recorded** -- the vendoring step did not preserve it |
+| Mirrored on | 2026-09-15 |
+| Upstream license | LICENSE |
+
+This is a **mirror**, stripped to the files needed for inference. The full
+untouched snapshot is at `archive/image__aide.tar.gz`.
+
+This code is the work of its original authors and is **not** covered by the
+DeepSafe project license. If you are an author and want this removed, open an
+issue on https://github.com/deepsafehq/deepsafe-bench and it will be taken down
+within 48 hours, no questions asked.
diff --git a/clean/image/aide/data/__init__.py b/clean/image/aide/data/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/clean/image/aide/data/dct.py b/clean/image/aide/data/dct.py
new file mode 100644
index 0000000000000000000000000000000000000000..0a0fe80983a4c203e163e06eafd8c8bc349bca5c
--- /dev/null
+++ b/clean/image/aide/data/dct.py
@@ -0,0 +1,107 @@
+"""DCT frequency decomposition module for AIDE.
+
+Source: https://github.com/shilinyan99/AIDE (ICLR 2025)
+Selects the most and least frequency-active image patches for
+multi-view input to the AIDE deepfake detector.
+"""
+
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+def DCT_mat(size):
+ m = [[(np.sqrt(1./size) if i == 0 else np.sqrt(2./size)) * np.cos((j + 0.5) * np.pi * i / size) for j in range(size)] for i in range(size)]
+ return m
+
+def generate_filter(start, end, size):
+ return [[0. if i + j > end or i + j < start else 1. for j in range(size)] for i in range(size)]
+
+def norm_sigma(x):
+ return 2. * torch.sigmoid(x) - 1.
+
+class Filter(nn.Module):
+ def __init__(self, size, band_start, band_end, use_learnable=False, norm=False):
+ super(Filter, self).__init__()
+ self.use_learnable = use_learnable
+ self.base = nn.Parameter(torch.tensor(generate_filter(band_start, band_end, size)), requires_grad=False)
+ if self.use_learnable:
+ self.learnable = nn.Parameter(torch.randn(size, size), requires_grad=True)
+ self.learnable.data.normal_(0., 0.1)
+ self.norm = norm
+ if norm:
+ self.ft_num = nn.Parameter(torch.sum(torch.tensor(generate_filter(band_start, band_end, size))), requires_grad=False)
+
+ def forward(self, x):
+ if self.use_learnable:
+ filt = self.base + norm_sigma(self.learnable)
+ else:
+ filt = self.base
+ if self.norm:
+ y = x * filt / self.ft_num
+ else:
+ y = x * filt
+ return y
+
+class DCT_base_Rec_Module(nn.Module):
+ def __init__(self, window_size=32, stride=16, output=256, grade_N=6, level_fliter=[0]):
+ super().__init__()
+ assert output % window_size == 0
+ assert len(level_fliter) > 0
+ self.window_size = window_size
+ self.grade_N = grade_N
+ self.level_N = len(level_fliter)
+ self.N = (output // window_size) * (output // window_size)
+ self._DCT_patch = nn.Parameter(torch.tensor(DCT_mat(window_size)).float(), requires_grad=False)
+ self._DCT_patch_T = nn.Parameter(torch.transpose(torch.tensor(DCT_mat(window_size)).float(), 0, 1), requires_grad=False)
+ self.unfold = nn.Unfold(kernel_size=(window_size, window_size), stride=stride)
+ self.fold0 = nn.Fold(output_size=(window_size, window_size), kernel_size=(window_size, window_size), stride=window_size)
+ level_f = [Filter(window_size, 0, window_size * 2)]
+ self.level_filters = nn.ModuleList([level_f[i] for i in level_fliter])
+ self.grade_filters = nn.ModuleList([Filter(window_size, window_size * 2. / grade_N * i, window_size * 2. / grade_N * (i+1), norm=True) for i in range(grade_N)])
+
+ def forward(self, x):
+ N = self.N
+ grade_N = self.grade_N
+ level_N = self.level_N
+ window_size = self.window_size
+ C, W, H = x.shape
+ x_unfold = self.unfold(x.unsqueeze(0)).squeeze(0)
+ _, L = x_unfold.shape
+ x_unfold = x_unfold.transpose(0, 1).reshape(L, C, window_size, window_size)
+ x_dct = self._DCT_patch @ x_unfold @ self._DCT_patch_T
+ y_list = []
+ for i in range(self.level_N):
+ x_pass = self.level_filters[i](x_dct)
+ y = self._DCT_patch_T @ x_pass @ self._DCT_patch
+ y_list.append(y)
+ level_x_unfold = torch.cat(y_list, dim=1)
+ grade = torch.zeros(L).to(x.device)
+ w, k = 1, 2
+ for _ in range(grade_N):
+ _x = torch.abs(x_dct)
+ _x = torch.log(_x + 1)
+ _x = self.grade_filters[_](_x)
+ _x = torch.sum(_x, dim=[1,2,3])
+ grade += w * _x
+ w *= k
+ _, idx = torch.sort(grade)
+ max_idx = torch.flip(idx, dims=[0])[:N]
+ maxmax_idx = max_idx[0]
+ maxmax_idx1 = max_idx[1] if len(max_idx) > 1 else max_idx[0]
+ min_idx = idx[:N]
+ minmin_idx = idx[0]
+ minmin_idx1 = idx[1] if len(min_idx) > 1 else idx[0]
+ x_minmin = torch.index_select(level_x_unfold, 0, minmin_idx)
+ x_maxmax = torch.index_select(level_x_unfold, 0, maxmax_idx)
+ x_minmin1 = torch.index_select(level_x_unfold, 0, minmin_idx1)
+ x_maxmax1 = torch.index_select(level_x_unfold, 0, maxmax_idx1)
+ x_minmin = x_minmin.reshape(1, level_N*C*window_size*window_size).transpose(0, 1)
+ x_maxmax = x_maxmax.reshape(1, level_N*C*window_size*window_size).transpose(0, 1)
+ x_minmin1 = x_minmin1.reshape(1, level_N*C*window_size*window_size).transpose(0, 1)
+ x_maxmax1 = x_maxmax1.reshape(1, level_N*C*window_size*window_size).transpose(0, 1)
+ x_minmin = self.fold0(x_minmin)
+ x_maxmax = self.fold0(x_maxmax)
+ x_minmin1 = self.fold0(x_minmin1)
+ x_maxmax1 = self.fold0(x_maxmax1)
+ return x_minmin, x_maxmax, x_minmin1, x_maxmax1
diff --git a/clean/image/aide/engine_finetune.py b/clean/image/aide/engine_finetune.py
new file mode 100644
index 0000000000000000000000000000000000000000..e356fc9d83b6622ace43e5a7c68c4badf908818e
--- /dev/null
+++ b/clean/image/aide/engine_finetune.py
@@ -0,0 +1,197 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import math
+from typing import Iterable, Optional
+
+import torch
+import torch.distributed as dist
+from timm.data import Mixup
+from timm.utils import accuracy, ModelEma
+
+import utils
+from utils import adjust_learning_rate
+from scipy.special import softmax
+from sklearn.metrics import (
+ average_precision_score,
+ accuracy_score
+)
+import numpy as np
+
+
+def train_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module,
+ data_loader: Iterable, optimizer: torch.optim.Optimizer,
+ device: torch.device, epoch: int, loss_scaler, max_norm: float = 0,
+ model_ema: Optional[ModelEma] = None, mixup_fn: Optional[Mixup] = None,
+ log_writer=None, args=None):
+ model.train(True)
+ metric_logger = utils.MetricLogger(delimiter=" ")
+ metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+ header = 'Epoch: [{}]'.format(epoch)
+ print_freq = 20
+
+ update_freq = args.update_freq
+ use_amp = args.use_amp
+ optimizer.zero_grad()
+
+ for data_iter_step, (samples, targets) in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
+ # we use a per iteration (instead of per epoch) lr scheduler
+ if data_iter_step % update_freq == 0:
+ adjust_learning_rate(optimizer, data_iter_step / len(data_loader) + epoch, args)
+
+ samples = samples.to(device, non_blocking=True)
+ targets = targets.to(device, non_blocking=True)
+
+ if mixup_fn is not None:
+ samples, targets = mixup_fn(samples, targets)
+
+ if use_amp:
+ with torch.cuda.amp.autocast():
+ output = model(samples)
+ loss = criterion(output, targets)
+ else: # full precision
+ output = model(samples)
+ loss = criterion(output, targets)
+
+ loss_value = loss.item()
+
+ if not math.isfinite(loss_value):
+ print("Loss is {}, stopping training".format(loss_value))
+ assert math.isfinite(loss_value)
+
+ if use_amp:
+ # this attribute is added by timm on one optimizer (adahessian)
+ is_second_order = hasattr(optimizer, 'is_second_order') and optimizer.is_second_order
+ loss /= update_freq
+ grad_norm = loss_scaler(loss, optimizer, clip_grad=max_norm,
+ parameters=model.parameters(), create_graph=is_second_order,
+ update_grad=(data_iter_step + 1) % update_freq == 0)
+ if (data_iter_step + 1) % update_freq == 0:
+ optimizer.zero_grad()
+ if model_ema is not None:
+ model_ema.update(model)
+ else: # full precision
+ loss /= update_freq
+ loss.backward()
+ if (data_iter_step + 1) % update_freq == 0:
+ optimizer.step()
+ optimizer.zero_grad()
+ if model_ema is not None:
+ model_ema.update(model)
+
+ torch.cuda.synchronize()
+
+ if mixup_fn is None:
+ class_acc = (output.max(-1)[-1] == targets).float().mean()
+ else:
+ class_acc = None
+
+ metric_logger.update(loss=loss_value)
+ metric_logger.update(class_acc=class_acc)
+ min_lr = 10.
+ max_lr = 0.
+ for group in optimizer.param_groups:
+ min_lr = min(min_lr, group["lr"])
+ max_lr = max(max_lr, group["lr"])
+
+ metric_logger.update(lr=max_lr)
+ metric_logger.update(min_lr=min_lr)
+ weight_decay_value = None
+ for group in optimizer.param_groups:
+ if group["weight_decay"] > 0:
+ weight_decay_value = group["weight_decay"]
+ metric_logger.update(weight_decay=weight_decay_value)
+ if use_amp:
+ metric_logger.update(grad_norm=grad_norm)
+ if log_writer is not None:
+ log_writer.update(loss=loss_value, head="loss")
+ log_writer.update(class_acc=class_acc, head="loss")
+ log_writer.update(lr=max_lr, head="opt")
+ log_writer.update(min_lr=min_lr, head="opt")
+ log_writer.update(weight_decay=weight_decay_value, head="opt")
+ if use_amp:
+ log_writer.update(grad_norm=grad_norm, head="opt")
+ log_writer.set_step()
+
+ # gather the stats from all processes
+ metric_logger.synchronize_between_processes()
+ print("Averaged stats:", metric_logger)
+ return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+@torch.no_grad()
+def evaluate(data_loader, model, device, use_amp=False):
+ criterion = torch.nn.CrossEntropyLoss()
+
+ metric_logger = utils.MetricLogger(delimiter=" ")
+ header = 'Test:'
+
+ # switch to evaluation mode
+ model.eval()
+
+ for index, batch in enumerate(metric_logger.log_every(data_loader, 10, header)):
+ images = batch[0]
+ target = batch[-1]
+
+ images = images.to(device, non_blocking=True)
+ target = target.to(device, non_blocking=True)
+
+ # compute output
+ if use_amp:
+ with torch.cuda.amp.autocast(dytpe=torch.bfloat16):
+ output = model(images)
+ if isinstance(output, dict):
+ output = output['logits']
+ loss = criterion(output, target)
+ else:
+ output = model(images) #[bs, num_cls]
+ if isinstance(output, dict):
+ output = output['logits']
+
+ loss = criterion(output, target)
+
+ if index == 0:
+ predictions = output
+ labels = target
+ else:
+ predictions = torch.cat((predictions, output), 0)
+ labels = torch.cat((labels, target), 0)
+
+ torch.cuda.synchronize()
+
+ acc1, acc5 = accuracy(output, target, topk=(1, 2))
+
+ batch_size = images.shape[0]
+ metric_logger.update(loss=loss.item())
+ metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
+ metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
+ # gather the stats from all processes
+ metric_logger.synchronize_between_processes()
+ print('* Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f} loss {losses.global_avg:.3f}'
+ .format(top1=metric_logger.acc1, top5=metric_logger.acc5, losses=metric_logger.loss))
+
+
+ output_ddp = [torch.zeros_like(predictions) for _ in range(utils.get_world_size())]
+ dist.all_gather(output_ddp, predictions)
+ labels_ddp = [torch.zeros_like(labels) for _ in range(utils.get_world_size())]
+ dist.all_gather(labels_ddp, labels)
+
+ output_all = torch.concat(output_ddp, dim=0)
+ labels_all = torch.concat(labels_ddp, dim=0)
+
+
+ y_pred = softmax(output_all.detach().cpu().numpy(), axis=1)[:, 1]
+ y_true = labels_all.detach().cpu().numpy()
+ y_true = y_true.astype(int)
+
+
+ acc = accuracy_score(y_true, y_pred > 0.5)
+ ap = average_precision_score(y_true, y_pred)
+
+
+
+ return {k: meter.global_avg for k, meter in metric_logger.meters.items()}, acc, ap
\ No newline at end of file
diff --git a/clean/image/aide/main_finetune.py b/clean/image/aide/main_finetune.py
new file mode 100644
index 0000000000000000000000000000000000000000..e5e5b47c9ab5348809dcb82b17fd8b038d4dc0d0
--- /dev/null
+++ b/clean/image/aide/main_finetune.py
@@ -0,0 +1,449 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import argparse
+import datetime
+import numpy as np
+import time
+import json
+import os
+from pathlib import Path
+
+import torch
+import torch.backends.cudnn as cudnn
+
+from timm.models.layers import trunc_normal_
+from timm.data.mixup import Mixup
+from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
+from timm.utils import ModelEma
+from optim_factory import create_optimizer, LayerDecayValueAssigner
+
+from data.datasets import TrainDataset, TestDataset
+from engine_finetune import train_one_epoch, evaluate
+
+import utils
+from utils import NativeScalerWithGradNormCount as NativeScaler
+from utils import str2bool, remap_checkpoint_keys
+import models.AIDE as AIDE
+import csv
+import warnings
+
+warnings.filterwarnings('ignore')
+
+def get_args_parser():
+ parser = argparse.ArgumentParser('Resnet fine-tuning', add_help=False)
+ parser.add_argument('--batch_size', default=64, type=int,
+ help='Per GPU batch size')
+ parser.add_argument('--epochs', default=100, type=int)
+ parser.add_argument('--update_freq', default=1, type=int,
+ help='gradient accumulation steps')
+
+ # Model parameters
+ parser.add_argument('--model', default='AIDE', type=str, metavar='MODEL',
+ help='Name of model to train')
+ parser.add_argument('--resnet_path', default=None, type=str, metavar='MODEL',
+ help='Path of resnet model')
+ parser.add_argument('--convnext_path', default=None, type=str, metavar='MODEL',
+ help='Path of ConvNeXt of model ')
+
+ # EMA related parameters
+ parser.add_argument('--model_ema', type=str2bool, default=False)
+ parser.add_argument('--model_ema_decay', type=float, default=0.9999, help='')
+ parser.add_argument('--model_ema_force_cpu', type=str2bool, default=False, help='')
+ parser.add_argument('--model_ema_eval', type=str2bool, default=False, help='Using ema to eval during training.')
+
+ # Optimization parameters
+ parser.add_argument('--clip_grad', type=float, default=None, metavar='NORM',
+ help='Clip gradient norm (default: None, no clipping)')
+ parser.add_argument('--weight_decay', type=float, default=0.,
+ help='weight decay (default: 0.05)')
+ parser.add_argument('--lr', type=float, default=None, metavar='LR',
+ help='learning rate (absolute lr)')
+ parser.add_argument('--blr', type=float, default=5e-4, metavar='LR',
+ help='base learning rate: absolute_lr = base_lr * total_batch_size / 256')
+ parser.add_argument('--layer_decay', type=float, default=1.0)
+ parser.add_argument('--min_lr', type=float, default=1e-6, metavar='LR',
+ help='lower lr bound for cyclic schedulers that hit 0 (1e-6)')
+ parser.add_argument('--warmup_epochs', type=int, default=0, metavar='N',
+ help='epochs to warmup LR, if scheduler supports')
+
+ parser.add_argument('--warmup_steps', type=int, default=-1, metavar='N',
+ help='num of steps to warmup LR, will overload warmup_epochs if set > 0')
+ parser.add_argument('--opt', default='adamw', type=str, metavar='OPTIMIZER',
+ help='Optimizer (default: "adamw"')
+ parser.add_argument('--opt_eps', default=1e-8, type=float, metavar='EPSILON',
+ help='Optimizer Epsilon (default: 1e-8)')
+ parser.add_argument('--opt_betas', default=None, type=float, nargs='+', metavar='BETA',
+ help='Optimizer Betas (default: None, use opt default)')
+ parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
+ help='SGD momentum (default: 0.9)')
+ parser.add_argument('--weight_decay_end', type=float, default=None, help="""Final value of the
+ weight decay. We use a cosine schedule for WD and using a larger decay by
+ the end of training improves performance for ViTs.""")
+
+ # Augmentation parameters
+ parser.add_argument('--color_jitter', type=float, default=None, metavar='PCT',
+ help='Color jitter factor (enabled only when not using Auto/RandAug)')
+ parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1', metavar='NAME',
+ help='Use AutoAugment policy. "v0" or "original". " + "(default: rand-m9-mstd0.5-inc1)')
+ parser.add_argument('--smoothing', type=float, default=0.1,
+ help='Label smoothing (default: 0.1)')
+
+ parser.add_argument('--train_interpolation', type=str, default='bicubic',
+ help='Training interpolation (random, bilinear, bicubic default: "bicubic")')
+
+ # * Random Erase params
+ parser.add_argument('--reprob', type=float, default=0.25, metavar='PCT',
+ help='Random erase prob (default: 0.25)')
+ parser.add_argument('--remode', type=str, default='pixel',
+ help='Random erase mode (default: "pixel")')
+ parser.add_argument('--recount', type=int, default=1,
+ help='Random erase count (default: 1)')
+ parser.add_argument('--resplit', type=str2bool, default=False,
+ help='Do not random erase first (clean) augmentation split')
+
+ # * Mixup params
+ parser.add_argument('--mixup', type=float, default=0.,
+ help='mixup alpha, mixup enabled if > 0.')
+ parser.add_argument('--cutmix', type=float, default=0.,
+ help='cutmix alpha, cutmix enabled if > 0.')
+ parser.add_argument('--cutmix_minmax', type=float, nargs='+', default=None,
+ help='cutmix min/max ratio, overrides alpha and enables cutmix if set (default: None)')
+ parser.add_argument('--mixup_prob', type=float, default=1.0,
+ help='Probability of performing mixup or cutmix when either/both is enabled')
+ parser.add_argument('--mixup_switch_prob', type=float, default=0.5,
+ help='Probability of switching to cutmix when both mixup and cutmix enabled')
+ parser.add_argument('--mixup_mode', type=str, default='batch',
+ help='How to apply mixup/cutmix params. Per "batch", "pair", or "elem"')
+
+ # * Finetuning params
+ parser.add_argument('--finetune', default='',
+ help='finetune from checkpoint')
+ parser.add_argument('--head_init_scale', default=0.001, type=float,
+ help='classifier head initial scale, typically adjusted in fine-tuning')
+ parser.add_argument('--model_key', default='model|module', type=str,
+ help='which key to load from saved state dict, usually model or model_ema')
+ parser.add_argument('--model_prefix', default='', type=str)
+
+ # Dataset parameters
+ parser.add_argument('--data_path', default='path/dataset', type=str,
+ help='dataset path')
+ parser.add_argument('--nb_classes', default=2, type=int,
+ help='number of the classification types')
+ parser.add_argument('--output_dir', default='',
+ help='path where to save, empty for no saving')
+ parser.add_argument('--log_dir', default=None,
+ help='path where to tensorboard log')
+ parser.add_argument('--device', default='cuda',
+ help='device to use for training / testing')
+ parser.add_argument('--seed', default=0, type=int)
+ parser.add_argument('--resume', default='',
+ help='resume from checkpoint')
+
+ parser.add_argument('--eval_data_path', default=None, type=str,
+ help='dataset path for evaluation')
+ parser.add_argument('--imagenet_default_mean_and_std', type=str2bool, default=True)
+ parser.add_argument('--data_set', default='IMNET', choices=['CIFAR', 'IMNET', 'image_folder'],
+ type=str, help='ImageNet dataset path')
+ parser.add_argument('--auto_resume', type=str2bool, default=True)
+ parser.add_argument('--save_ckpt', type=str2bool, default=True)
+ parser.add_argument('--save_ckpt_freq', default=1, type=int)
+ parser.add_argument('--save_ckpt_num', default=100, type=int)
+
+ parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
+ help='start epoch')
+ parser.add_argument('--eval', type=str2bool, default=False,
+ help='Perform evaluation only')
+ parser.add_argument('--dist_eval', type=str2bool, default=True,
+ help='Enabling distributed evaluation')
+ parser.add_argument('--disable_eval', type=str2bool, default=False,
+ help='Disabling evaluation during training')
+ parser.add_argument('--num_workers', default=16, type=int)
+ parser.add_argument('--pin_mem', type=str2bool, default=True,
+ help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
+
+ # Evaluation parameters
+ parser.add_argument('--crop_pct', type=float, default=None)
+
+ # distributed training parameters
+ parser.add_argument('--world_size', default=1, type=int,
+ help='number of distributed processes')
+ parser.add_argument('--local_rank', default=-1, type=int)
+ parser.add_argument('--dist_on_itp', type=str2bool, default=False)
+ parser.add_argument('--dist_url', default='env://',
+ help='url used to set up distributed training')
+
+ parser.add_argument('--use_amp', type=str2bool, default=False,
+ help="Use apex AMP (Automatic Mixed Precision) or not")
+ return parser
+
+def main(args):
+ utils.init_distributed_mode(args)
+ print(args)
+ device = torch.device(args.device)
+
+ # fix the seed for reproducibility
+ seed = args.seed + utils.get_rank()
+ torch.manual_seed(seed)
+ np.random.seed(seed)
+ cudnn.benchmark = True
+
+ dataset_train = TrainDataset(is_train=True, args=args)
+
+ if args.disable_eval:
+ args.dist_eval = False
+ dataset_val = None
+ else:
+ dataset_val = TrainDataset(is_train=False, args=args)
+
+ num_tasks = utils.get_world_size()
+ global_rank = utils.get_rank()
+
+ sampler_train = torch.utils.data.DistributedSampler(
+ dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True, seed=args.seed,
+ )
+ print("Sampler_train = %s" % str(sampler_train))
+ if args.dist_eval:
+ if len(dataset_val) % num_tasks != 0:
+ print('Warning: Enabling distributed evaluation with an eval dataset not divisible by process number. '
+ 'This will slightly alter validation results as extra duplicate entries are added to achieve '
+ 'equal num of samples per-process.')
+ sampler_val = torch.utils.data.DistributedSampler(
+ dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=False)
+ else:
+ sampler_val = torch.utils.data.SequentialSampler(dataset_val)
+
+ if global_rank == 0 and args.log_dir is not None:
+ os.makedirs(args.log_dir, exist_ok=True)
+ log_writer = utils.TensorboardLogger(log_dir=args.log_dir)
+ else:
+ log_writer = None
+
+ data_loader_train = torch.utils.data.DataLoader(
+ dataset_train, sampler=sampler_train,
+ batch_size=args.batch_size,
+ num_workers=args.num_workers,
+ pin_memory=args.pin_mem,
+ drop_last=True,
+ )
+ if dataset_val is not None:
+ data_loader_val = torch.utils.data.DataLoader(
+ dataset_val, sampler=sampler_val,
+ batch_size=args.batch_size,
+ num_workers=args.num_workers,
+ pin_memory=args.pin_mem,
+ drop_last=False
+ )
+ else:
+ data_loader_val = None
+
+ mixup_fn = None
+ mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
+ if mixup_active:
+ print("Mixup is activated!")
+ mixup_fn = Mixup(
+ mixup_alpha=args.mixup, cutmix_alpha=args.cutmix, cutmix_minmax=args.cutmix_minmax,
+ prob=args.mixup_prob, switch_prob=args.mixup_switch_prob, mode=args.mixup_mode,
+ label_smoothing=args.smoothing, num_classes=args.nb_classes)
+
+
+ model = AIDE.__dict__[args.model](
+ resnet_path=args.resnet_path,
+ convnext_path=args.convnext_path
+ )
+
+ model.to(device)
+
+ model_ema = None
+ if args.model_ema:
+ # Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
+ model_ema = ModelEma(
+ model,
+ decay=args.model_ema_decay,
+ device='cpu' if args.model_ema_force_cpu else '',
+ resume='')
+ print("Using EMA with decay = %.8f" % args.model_ema_decay)
+
+ model_without_ddp = model
+ n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+ print("Model = %s" % str(model_without_ddp))
+ print('number of params:', n_parameters)
+
+ eff_batch_size = args.batch_size * args.update_freq * utils.get_world_size()
+ num_training_steps_per_epoch = len(dataset_train) // eff_batch_size
+
+ if args.lr is None:
+ args.lr = args.blr * eff_batch_size / 256
+
+ print("base lr: %.2e" % (args.lr * 256 / eff_batch_size))
+ print("actual lr: %.2e" % args.lr)
+
+ print("accumulate grad iterations: %d" % args.update_freq)
+ print("effective batch size: %d" % eff_batch_size)
+
+ assigner = None
+
+ if args.distributed:
+ model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
+ model_without_ddp = model.module
+
+ optimizer = create_optimizer(
+ args, model_without_ddp, skip_list=None,
+ get_num_layer=assigner.get_layer_id if assigner is not None else None,
+ get_layer_scale=assigner.get_scale if assigner is not None else None)
+ loss_scaler = NativeScaler()
+
+ if mixup_fn is not None:
+ # smoothing is handled with mixup label transform
+ criterion = SoftTargetCrossEntropy()
+ elif args.smoothing > 0.:
+ criterion = LabelSmoothingCrossEntropy(smoothing=args.smoothing)
+ else:
+ criterion = torch.nn.CrossEntropyLoss()
+
+ print("criterion = %s" % str(criterion))
+
+ utils.auto_load_model(
+ args=args, model=model, model_without_ddp=model_without_ddp,
+ optimizer=optimizer, loss_scaler=loss_scaler, model_ema=model_ema)
+
+ if args.eval:
+ print(f"Eval only mode")
+
+ vals = os.listdir(args.eval_data_path)
+ if len(vals) == 16:
+ vals = ["progan", "stylegan", "biggan", "cyclegan", "stargan", "gaugan", "stylegan2", "whichfaceisreal", "ADM", "Glide", "Midjourney", "stable_diffusion_v_1_4", "stable_diffusion_v_1_5", "VQDM", "wukong", "DALLE2"]
+ if len(vals) == 8:
+ vals = ["Midjourney", "stable_diffusion_v_1_4", "stable_diffusion_v_1_5", "ADM", "glide", "wukong", "VQDM", "BigGAN"]
+ eval_data_path = args.eval_data_path
+
+ rows = [["{} model testing on...".format(args.resume)],
+ ['testset', 'accuracy', 'avg precision']]
+
+ for v_id, val in enumerate(vals):
+
+ args.eval_data_path = os.path.join(args.eval_data_path, val)
+ dataset_val = TestDataset(is_train=False, args=args)
+ args.eval_data_path = eval_data_path
+
+ if args.dist_eval:
+ if len(dataset_val) % num_tasks != 0:
+ print('Warning: Enabling distributed evaluation with an eval dataset not divisible by process number. '
+ 'This will slightly alter validation results as extra duplicate entries are added to achieve '
+ 'equal num of samples per-process.')
+ sampler_val = torch.utils.data.DistributedSampler(
+ dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=False)
+ else:
+ sampler_val = torch.utils.data.SequentialSampler(dataset_val)
+
+ data_loader_val = torch.utils.data.DataLoader(
+ dataset_val, sampler=sampler_val,
+ batch_size=args.batch_size,
+ num_workers=args.num_workers,
+ pin_memory=args.pin_mem,
+ drop_last=False
+ )
+
+
+ test_stats, acc, ap = evaluate(data_loader_val, model, device)
+ print(f"Accuracy of the network on {len(dataset_val)} test images: {test_stats['acc1']:.5f}%")
+
+ print(f"test dataset is {val} acc: {acc}, ap: {ap}")
+ print("***********************************")
+
+ rows.append([val, acc, ap])
+
+
+ test_dataset_name = args.eval_data_path.split('/')[-2]
+
+ csv_name = os.path.join(args.output_dir, f'{os.path.basename(args.resume)}_{test_dataset_name}.csv')
+ with open(csv_name, 'w') as f:
+ csv_writer = csv.writer(f, delimiter=',')
+ csv_writer.writerows(rows)
+ return
+
+ max_accuracy = 0.0
+ if args.model_ema and args.model_ema_eval:
+ max_accuracy_ema = 0.0
+
+ print("Start training for %d epochs" % args.epochs)
+ start_time = time.time()
+ for epoch in range(args.start_epoch, args.epochs):
+ if args.distributed:
+ data_loader_train.sampler.set_epoch(epoch)
+ if log_writer is not None:
+ log_writer.set_step(epoch * num_training_steps_per_epoch * args.update_freq)
+ train_stats = train_one_epoch(
+ model, criterion, data_loader_train,
+ optimizer, device, epoch, loss_scaler,
+ args.clip_grad, model_ema, mixup_fn,
+ log_writer=log_writer,
+ args=args
+ )
+ if args.output_dir and args.save_ckpt:
+ if (epoch + 1) % args.save_ckpt_freq == 0 or epoch + 1 == args.epochs:
+ utils.save_model(
+ args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
+ loss_scaler=loss_scaler, epoch=epoch, model_ema=model_ema)
+ if data_loader_val is not None:
+ test_stats, acc, ap = evaluate(data_loader_val, model, device, use_amp=args.use_amp)
+ print(f"Accuracy of the model on the {len(dataset_val)} test images: {test_stats['acc1']:.1f}%, ap: {ap}.")
+ if max_accuracy < test_stats["acc1"]:
+ max_accuracy = test_stats["acc1"]
+ if args.output_dir and args.save_ckpt:
+ utils.save_model(
+ args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
+ loss_scaler=loss_scaler, epoch="best", model_ema=model_ema)
+ print(f'Max accuracy: {max_accuracy:.2f}%')
+
+ if log_writer is not None:
+ log_writer.update(test_acc1=test_stats['acc1'], head="perf", step=epoch)
+ log_writer.update(test_acc5=test_stats['acc5'], head="perf", step=epoch)
+ log_writer.update(test_loss=test_stats['loss'], head="perf", step=epoch)
+
+ log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
+ **{f'test_{k}': v for k, v in test_stats.items()},
+ 'epoch': epoch,
+ 'n_parameters': n_parameters}
+
+ # repeat testing routines for EMA, if ema eval is turned on
+ if args.model_ema and args.model_ema_eval:
+ test_stats_ema, acc, ap = evaluate(data_loader_val, model_ema.ema, device, use_amp=args.use_amp)
+ print(f"Accuracy of the model EMA on {len(dataset_val)} test images: {test_stats_ema['acc1']:.1f}%, ap: {ap}")
+ if max_accuracy_ema < test_stats_ema["acc1"]:
+ max_accuracy_ema = test_stats_ema["acc1"]
+ if args.output_dir and args.save_ckpt:
+ utils.save_model(
+ args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
+ loss_scaler=loss_scaler, epoch="best-ema", model_ema=model_ema)
+ print(f'Max EMA accuracy: {max_accuracy_ema:.2f}%')
+ if log_writer is not None:
+ log_writer.update(test_acc1_ema=test_stats_ema['acc1'], head="perf", step=epoch)
+ log_stats.update({**{f'test_{k}_ema': v for k, v in test_stats_ema.items()}})
+ else:
+ log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
+ 'epoch': epoch,
+ 'n_parameters': n_parameters}
+
+ if args.output_dir and utils.is_main_process():
+ if log_writer is not None:
+ log_writer.flush()
+ with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f:
+ f.write(json.dumps(log_stats) + "\n")
+
+ total_time = time.time() - start_time
+ total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+ print('Training time {}'.format(total_time_str))
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser('AIDE traning', parents=[get_args_parser()])
+ args = parser.parse_args()
+ if args.output_dir:
+ Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+ main(args)
diff --git a/clean/image/aide/models/AIDE.py b/clean/image/aide/models/AIDE.py
new file mode 100644
index 0000000000000000000000000000000000000000..053805afad55edf61c9d9de4d2c499418dbb4674
--- /dev/null
+++ b/clean/image/aide/models/AIDE.py
@@ -0,0 +1,298 @@
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+import torch
+import clip
+import open_clip
+from .srm_filter_kernel import all_normalized_hpf_list
+import numpy as np
+
+class HPF(nn.Module):
+ def __init__(self):
+ super(HPF, self).__init__()
+
+ #Load 30 SRM Filters
+ all_hpf_list_5x5 = []
+
+ for hpf_item in all_normalized_hpf_list:
+ if hpf_item.shape[0] == 3:
+ hpf_item = np.pad(hpf_item, pad_width=((1, 1), (1, 1)), mode='constant')
+
+ all_hpf_list_5x5.append(hpf_item)
+
+ hpf_weight = torch.Tensor(all_hpf_list_5x5).view(30, 1, 5, 5).contiguous()
+ hpf_weight = torch.nn.Parameter(hpf_weight.repeat(1, 3, 1, 1), requires_grad=False)
+
+
+ self.hpf = nn.Conv2d(3, 30, kernel_size=5, padding=2, bias=False)
+ self.hpf.weight = hpf_weight
+
+
+ def forward(self, input):
+
+ output = self.hpf(input)
+
+ return output
+
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+ """3x3 convolution with padding"""
+ return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+ padding=1, bias=False)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+ """1x1 convolution"""
+ return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+ expansion = 1
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None):
+ super(BasicBlock, self).__init__()
+ self.conv1 = conv3x3(inplanes, planes, stride)
+ self.bn1 = nn.BatchNorm2d(planes)
+ self.relu = nn.ReLU(inplace=True)
+ self.conv2 = conv3x3(planes, planes)
+ self.bn2 = nn.BatchNorm2d(planes)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x):
+ identity = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+
+ if self.downsample is not None:
+ identity = self.downsample(x)
+
+ out += identity
+ out = self.relu(out)
+
+ return out
+
+
+class Bottleneck(nn.Module):
+ expansion = 4
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None):
+ super(Bottleneck, self).__init__()
+ self.conv1 = conv1x1(inplanes, planes)
+ self.bn1 = nn.BatchNorm2d(planes)
+ self.conv2 = conv3x3(planes, planes, stride)
+ self.bn2 = nn.BatchNorm2d(planes)
+ self.conv3 = conv1x1(planes, planes * self.expansion)
+ self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+ self.relu = nn.ReLU(inplace=True)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x):
+ identity = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+ out = self.relu(out)
+
+ out = self.conv3(out)
+ out = self.bn3(out)
+
+ if self.downsample is not None:
+ identity = self.downsample(x)
+
+ out += identity
+ out = self.relu(out)
+
+ return out
+
+
+class ResNet(nn.Module):
+
+ def __init__(self, block, layers, num_classes=1000, zero_init_residual=True):
+ super(ResNet, self).__init__()
+
+ self.inplanes = 64
+ self.conv1 = nn.Conv2d(30, 64, kernel_size=7, stride=2, padding=3,
+ bias=False)
+ self.bn1 = nn.BatchNorm2d(64)
+ self.relu = nn.ReLU(inplace=True)
+ self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+ self.layer1 = self._make_layer(block, 64, layers[0])
+ self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+ self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+ self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+ self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+ self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+ for m in self.modules():
+ if isinstance(m, nn.Conv2d):
+ nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+ elif isinstance(m, nn.BatchNorm2d):
+ nn.init.constant_(m.weight, 1)
+ nn.init.constant_(m.bias, 0)
+
+ # Zero-initialize the last BN in each residual branch,
+ # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+ # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+ if zero_init_residual:
+ for m in self.modules():
+ if isinstance(m, Bottleneck):
+ nn.init.constant_(m.bn3.weight, 0)
+ elif isinstance(m, BasicBlock):
+ nn.init.constant_(m.bn2.weight, 0)
+
+ def _make_layer(self, block, planes, blocks, stride=1):
+ downsample = None
+ if stride != 1 or self.inplanes != planes * block.expansion:
+ downsample = nn.Sequential(
+ conv1x1(self.inplanes, planes * block.expansion, stride),
+ nn.BatchNorm2d(planes * block.expansion),
+ )
+
+ layers = []
+ layers.append(block(self.inplanes, planes, stride, downsample))
+ self.inplanes = planes * block.expansion
+ for _ in range(1, blocks):
+ layers.append(block(self.inplanes, planes))
+
+ return nn.Sequential(*layers)
+
+ def forward(self, x):
+
+ x = self.conv1(x)
+ x = self.bn1(x)
+ x = self.relu(x)
+ x = self.maxpool(x)
+
+ x = self.layer1(x)
+ x = self.layer2(x)
+ x = self.layer3(x)
+ x = self.layer4(x)
+
+ x = self.avgpool(x)
+ x = x.view(x.size(0), -1)
+
+
+ return x
+
+class Mlp(nn.Module):
+ """ MLP as used in Vision Transformer, MLP-Mixer and related networks
+ """
+
+ def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU):
+ super().__init__()
+ out_features = out_features or in_features
+ hidden_features = hidden_features or in_features
+
+ self.fc1 = nn.Linear(in_features, hidden_features)
+ self.act = act_layer()
+ self.fc2 = nn.Linear(hidden_features, out_features)
+
+ def forward(self, x):
+ x = self.fc1(x)
+ x = self.act(x)
+ x = self.fc2(x)
+ return x
+
+class AIDE_Model(nn.Module):
+
+ def __init__(self, resnet_path, convnext_path):
+ super(AIDE_Model, self).__init__()
+ self.hpf = HPF()
+ self.model_min = ResNet(Bottleneck, [3, 4, 6, 3])
+ self.model_max = ResNet(Bottleneck, [3, 4, 6, 3])
+
+ if resnet_path is not None:
+ pretrained_dict = torch.load(resnet_path, map_location='cpu')
+
+ model_min_dict = self.model_min.state_dict()
+ model_max_dict = self.model_max.state_dict()
+
+ for k in pretrained_dict.keys():
+ if k in model_min_dict and pretrained_dict[k].size() == model_min_dict[k].size():
+ model_min_dict[k] = pretrained_dict[k]
+ model_max_dict[k] = pretrained_dict[k]
+ else:
+ print(f"Skipping layer {k} because of size mismatch")
+
+ self.fc = Mlp(2048 + 256 , 1024, 2)
+
+ print("build model with convnext_xxl")
+ self.openclip_convnext_xxl, _, _ = open_clip.create_model_and_transforms(
+ "convnext_xxlarge", pretrained=convnext_path
+ )
+
+ self.openclip_convnext_xxl = self.openclip_convnext_xxl.visual.trunk
+ self.openclip_convnext_xxl.head.global_pool = nn.Identity()
+ self.openclip_convnext_xxl.head.flatten = nn.Identity()
+
+ self.openclip_convnext_xxl.eval()
+
+ self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+ self.convnext_proj = nn.Sequential(
+ nn.Linear(3072, 256),
+
+ )
+ for param in self.openclip_convnext_xxl.parameters():
+ param.requires_grad = False
+
+
+
+ def forward(self, x):
+
+ b, t, c, h, w = x.shape
+
+ x_minmin = x[:, 0] #[b, c, h, w]
+ x_maxmax = x[:, 1]
+ x_minmin1 = x[:, 2]
+ x_maxmax1 = x[:, 3]
+ tokens = x[:, 4]
+
+ x_minmin = self.hpf(x_minmin)
+ x_maxmax = self.hpf(x_maxmax)
+ x_minmin1 = self.hpf(x_minmin1)
+ x_maxmax1 = self.hpf(x_maxmax1)
+
+ with torch.no_grad():
+
+ clip_mean = torch.Tensor([0.48145466, 0.4578275, 0.40821073])
+ clip_mean = clip_mean.to(tokens, non_blocking=True).view(3, 1, 1)
+ clip_std = torch.Tensor([0.26862954, 0.26130258, 0.27577711])
+ clip_std = clip_std.to(tokens, non_blocking=True).view(3, 1, 1)
+ dinov2_mean = torch.Tensor([0.485, 0.456, 0.406]).to(tokens, non_blocking=True).view(3, 1, 1)
+ dinov2_std = torch.Tensor([0.229, 0.224, 0.225]).to(tokens, non_blocking=True).view(3, 1, 1)
+
+ local_convnext_image_feats = self.openclip_convnext_xxl(
+ tokens * (dinov2_std / clip_std) + (dinov2_mean - clip_mean) / clip_std
+ ) #[b, 3072, 8, 8]
+ assert local_convnext_image_feats.size()[1:] == (3072, 8, 8)
+ local_convnext_image_feats = self.avgpool(local_convnext_image_feats).view(tokens.size(0), -1)
+ x_0 = self.convnext_proj(local_convnext_image_feats)
+
+ x_min = self.model_min(x_minmin)
+ x_max = self.model_max(x_maxmax)
+ x_min1 = self.model_min(x_minmin1)
+ x_max1 = self.model_max(x_maxmax1)
+
+ x_1 = (x_min + x_max + x_min1 + x_max1) / 4
+
+ x = torch.cat([x_0, x_1], dim=1)
+
+ x = self.fc(x)
+
+ return x
+
+def AIDE(resnet_path, convnext_path):
+ model = AIDE_Model(resnet_path, convnext_path)
+ return model
+
diff --git a/clean/image/aide/models/__init__.py b/clean/image/aide/models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/clean/image/aide/models/srm_filter_kernel.py b/clean/image/aide/models/srm_filter_kernel.py
new file mode 100644
index 0000000000000000000000000000000000000000..36eda56b8ad18f4500ba8826a634bfb8d923ab67
--- /dev/null
+++ b/clean/image/aide/models/srm_filter_kernel.py
@@ -0,0 +1,220 @@
+
+import numpy as np
+
+filter_class_1 = [
+ np.array([
+ [1, 0, 0],
+ [0, -1, 0],
+ [0, 0, 0]
+ ], dtype=np.float32),
+ np.array([
+ [0, 1, 0],
+ [0, -1, 0],
+ [0, 0, 0]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, 1],
+ [0, -1, 0],
+ [0, 0, 0]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, 0],
+ [1, -1, 0],
+ [0, 0, 0]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, 0],
+ [0, -1, 1],
+ [0, 0, 0]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, 0],
+ [0, -1, 0],
+ [1, 0, 0]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, 0],
+ [0, -1, 0],
+ [0, 1, 0]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, 0],
+ [0, -1, 0],
+ [0, 0, 1]
+ ], dtype=np.float32)
+]
+
+
+filter_class_2 = [
+ np.array([
+ [1, 0, 0],
+ [0, -2, 0],
+ [0, 0, 1]
+ ], dtype=np.float32),
+ np.array([
+ [0, 1, 0],
+ [0, -2, 0],
+ [0, 1, 0]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, 1],
+ [0, -2, 0],
+ [1, 0, 0]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, 0],
+ [1, -2, 1],
+ [0, 0, 0]
+ ], dtype=np.float32),
+]
+
+
+filter_class_3 = [
+ np.array([
+ [-1, 0, 0, 0, 0],
+ [0, 3, 0, 0, 0],
+ [0, 0, -3, 0, 0],
+ [0, 0, 0, 1, 0],
+ [0, 0, 0, 0, 0]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, -1, 0, 0],
+ [0, 0, 3, 0, 0],
+ [0, 0, -3, 0, 0],
+ [0, 0, 1, 0, 0],
+ [0, 0, 0, 0, 0]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, 0, 0, -1],
+ [0, 0, 0, 3, 0],
+ [0, 0, -3, 0, 0],
+ [0, 1, 0, 0, 0],
+ [0, 0, 0, 0, 0]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0],
+ [0, 1, -3, 3, -1],
+ [0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, 0, 0, 0],
+ [0, 1, 0, 0, 0],
+ [0, 0, -3, 0, 0],
+ [0, 0, 0, 3, 0],
+ [0, 0, 0, 0, -1]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, 0, 0, 0],
+ [0, 0, 1, 0, 0],
+ [0, 0, -3, 0, 0],
+ [0, 0, 3, 0, 0],
+ [0, 0, -1, 0, 0]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, 0, 0, 0],
+ [0, 0, 0, 1, 0],
+ [0, 0, -3, 0, 0],
+ [0, 3, 0, 0, 0],
+ [-1, 0, 0, 0, 0]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0],
+ [-1, 3, -3, 1, 0],
+ [0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0]
+ ], dtype=np.float32)
+]
+
+
+filter_edge_3x3 = [
+ np.array([
+ [-1, 2, -1],
+ [2, -4, 2],
+ [0, 0, 0]
+ ], dtype=np.float32),
+ np.array([
+ [0, 2, -1],
+ [0, -4, 2],
+ [0, 2, -1]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, 0],
+ [2, -4, 2],
+ [-1, 2, -1]
+ ], dtype=np.float32),
+ np.array([
+ [-1, 2, 0],
+ [2, -4, 0],
+ [-1, 2, 0]
+ ], dtype=np.float32),
+]
+
+filter_edge_5x5 = [
+ np.array([
+ [-1, 2, -2, 2, -1],
+ [2, -6, 8, -6, 2],
+ [-2, 8, -12, 8, -2],
+ [0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, -2, 2, -1],
+ [0, 0, 8, -6, 2],
+ [0, 0, -12, 8, -2],
+ [0, 0, 8, -6, 2],
+ [0, 0, -2, 2, -1]
+ ], dtype=np.float32),
+ np.array([
+ [0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0],
+ [-2, 8, -12, 8, -2],
+ [2, -6, 8, -6, 2],
+ [-1, 2, -2, 2, -1]
+ ], dtype=np.float32),
+ np.array([
+ [-1, 2, -2, 0, 0],
+ [2, -6, 8, 0, 0],
+ [-2, 8, -12, 0, 0],
+ [2, -6, 8, 0, 0],
+ [-1, 2, -2, 0, 0]
+ ], dtype=np.float32),
+]
+
+square_3x3 = np.array([
+ [-1, 2, -1],
+ [2, -4, 2],
+ [-1, 2, -1]
+], dtype=np.float32)
+
+square_5x5 = np.array([
+ [-1, 2, -2, 2, -1],
+ [2, -6, 8, -6, 2],
+ [-2, 8, -12, 8, -2],
+ [2, -6, 8, -6, 2],
+ [-1, 2, -2, 2, -1]
+], dtype=np.float32)
+
+
+all_hpf_list = filter_class_1 + filter_class_2 + filter_class_3 + filter_edge_3x3 + filter_edge_5x5 + [square_3x3, square_5x5]
+
+hpf_3x3_list = filter_class_1 + filter_class_2 + filter_edge_3x3 + [square_3x3]
+hpf_5x5_list = filter_class_3 + filter_edge_5x5 + [square_5x5]
+
+normalized_filter_class_2 = [hpf / 2 for hpf in filter_class_2]
+normalized_filter_class_3 = [hpf / 3 for hpf in filter_class_3]
+normalized_filter_edge_3x3 = [hpf / 4 for hpf in filter_edge_3x3]
+normalized_square_3x3 = square_3x3 / 4
+normalized_filter_edge_5x5 = [hpf / 12 for hpf in filter_edge_5x5]
+normalized_square_5x5 = square_5x5 / 12
+
+all_normalized_hpf_list = filter_class_1 + normalized_filter_class_2 + normalized_filter_class_3 + \
+ normalized_filter_edge_3x3 + normalized_filter_edge_5x5 + [normalized_square_3x3, normalized_square_5x5]
+
+normalized_hpf_3x3_list = filter_class_1 + normalized_filter_class_2 + normalized_filter_edge_3x3 + [normalized_square_3x3]
+normalized_hpf_5x5_list = normalized_filter_class_3 + normalized_filter_edge_5x5 + [normalized_square_5x5]
+
+normalized_3x3_list = normalized_filter_edge_3x3 + [normalized_square_3x3]
+normalized_5x5_list = normalized_filter_edge_5x5 + [normalized_square_5x5]
\ No newline at end of file
diff --git a/clean/image/aide/models/utils.py b/clean/image/aide/models/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..b6fc59d6ccf5652f437a491507c44222e042c3e7
--- /dev/null
+++ b/clean/image/aide/models/utils.py
@@ -0,0 +1,116 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import numpy.random as random
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+# from MinkowskiEngine import SparseTensor
+
+# class MinkowskiGRN(nn.Module):
+# """ GRN layer for sparse tensors.
+# """
+# def __init__(self, dim):
+# super().__init__()
+# self.gamma = nn.Parameter(torch.zeros(1, dim))
+# self.beta = nn.Parameter(torch.zeros(1, dim))
+
+# def forward(self, x):
+# cm = x.coordinate_manager
+# in_key = x.coordinate_map_key
+
+# Gx = torch.norm(x.F, p=2, dim=0, keepdim=True)
+# Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
+# return SparseTensor(
+# self.gamma * (x.F * Nx) + self.beta + x.F,
+# coordinate_map_key=in_key,
+# coordinate_manager=cm)
+
+# class MinkowskiDropPath(nn.Module):
+# """ Drop Path for sparse tensors.
+# """
+
+# def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True):
+# super(MinkowskiDropPath, self).__init__()
+# self.drop_prob = drop_prob
+# self.scale_by_keep = scale_by_keep
+
+# def forward(self, x):
+# if self.drop_prob == 0. or not self.training:
+# return x
+# cm = x.coordinate_manager
+# in_key = x.coordinate_map_key
+# keep_prob = 1 - self.drop_prob
+# mask = torch.cat([
+# torch.ones(len(_)) if random.uniform(0, 1) > self.drop_prob
+# else torch.zeros(len(_)) for _ in x.decomposed_coordinates
+# ]).view(-1, 1).to(x.device)
+# if keep_prob > 0.0 and self.scale_by_keep:
+# mask.div_(keep_prob)
+# return SparseTensor(
+# x.F * mask,
+# coordinate_map_key=in_key,
+# coordinate_manager=cm)
+
+# class MinkowskiLayerNorm(nn.Module):
+# """ Channel-wise layer normalization for sparse tensors.
+# """
+
+# def __init__(
+# self,
+# normalized_shape,
+# eps=1e-6,
+# ):
+# super(MinkowskiLayerNorm, self).__init__()
+# self.ln = nn.LayerNorm(normalized_shape, eps=eps)
+# def forward(self, input):
+# output = self.ln(input.F)
+# return SparseTensor(
+# output,
+# coordinate_map_key=input.coordinate_map_key,
+# coordinate_manager=input.coordinate_manager)
+
+class LayerNorm(nn.Module):
+ """ LayerNorm that supports two data formats: channels_last (default) or channels_first.
+ The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
+ shape (batch_size, height, width, channels) while channels_first corresponds to inputs
+ with shape (batch_size, channels, height, width).
+ """
+ def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(normalized_shape))
+ self.bias = nn.Parameter(torch.zeros(normalized_shape))
+ self.eps = eps
+ self.data_format = data_format
+ if self.data_format not in ["channels_last", "channels_first"]:
+ raise NotImplementedError
+ self.normalized_shape = (normalized_shape, )
+
+ def forward(self, x):
+ if self.data_format == "channels_last":
+ return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+ elif self.data_format == "channels_first":
+ u = x.mean(1, keepdim=True)
+ s = (x - u).pow(2).mean(1, keepdim=True)
+ x = (x - u) / torch.sqrt(s + self.eps)
+ x = self.weight[:, None, None] * x + self.bias[:, None, None]
+ return x
+
+class GRN(nn.Module):
+ """ GRN (Global Response Normalization) layer
+ """
+ def __init__(self, dim):
+ super().__init__()
+ self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
+ self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))
+
+ def forward(self, x):
+ Gx = torch.norm(x, p=2, dim=(1,2), keepdim=True)
+ Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
+ return self.gamma * (x * Nx) + self.beta + x
\ No newline at end of file
diff --git a/clean/image/aide/optim_factory.py b/clean/image/aide/optim_factory.py
new file mode 100644
index 0000000000000000000000000000000000000000..27103c2687210f82233d65ba0a46e22bfc1b671d
--- /dev/null
+++ b/clean/image/aide/optim_factory.py
@@ -0,0 +1,222 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import torch
+from torch import optim as optim
+
+from timm.optim.adafactor import Adafactor
+from timm.optim.adahessian import Adahessian
+from timm.optim.adamp import AdamP
+from timm.optim.lookahead import Lookahead
+from timm.optim.nadam import Nadam
+# from timm.optim.novograd import NovoGrad
+# from timm.optim.nvnovograd import NvNovoGrad
+from timm.optim.radam import RAdam
+from timm.optim.rmsprop_tf import RMSpropTF
+from timm.optim.sgdp import SGDP
+
+import json
+
+try:
+ from apex.optimizers import FusedNovoGrad, FusedAdam, FusedLAMB, FusedSGD
+ has_apex = True
+except ImportError:
+ has_apex = False
+
+
+def get_num_layer_for_convnext_single(var_name, depths):
+ """
+ Each layer is assigned distinctive layer ids
+ """
+ if var_name.startswith("downsample_layers"):
+ stage_id = int(var_name.split('.')[1])
+ layer_id = sum(depths[:stage_id]) + 1
+ return layer_id
+
+ elif var_name.startswith("stages"):
+ stage_id = int(var_name.split('.')[1])
+ block_id = int(var_name.split('.')[2])
+ layer_id = sum(depths[:stage_id]) + block_id + 1
+ return layer_id
+
+ else:
+ return sum(depths) + 1
+
+
+def get_num_layer_for_convnext(var_name):
+ """
+ Divide [3, 3, 27, 3] layers into 12 groups; each group is three
+ consecutive blocks, including possible neighboring downsample layers;
+ adapted from https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py
+ """
+ num_max_layer = 12
+ if var_name.startswith("downsample_layers"):
+ stage_id = int(var_name.split('.')[1])
+ if stage_id == 0:
+ layer_id = 0
+ elif stage_id == 1 or stage_id == 2:
+ layer_id = stage_id + 1
+ elif stage_id == 3:
+ layer_id = 12
+ return layer_id
+
+ elif var_name.startswith("stages"):
+ stage_id = int(var_name.split('.')[1])
+ block_id = int(var_name.split('.')[2])
+ if stage_id == 0 or stage_id == 1:
+ layer_id = stage_id + 1
+ elif stage_id == 2:
+ layer_id = 3 + block_id // 3
+ elif stage_id == 3:
+ layer_id = 12
+ return layer_id
+ else:
+ return num_max_layer + 1
+
+class LayerDecayValueAssigner(object):
+ def __init__(self, values, depths=[3,3,27,3], layer_decay_type='single'):
+ self.values = values
+ self.depths = depths
+ self.layer_decay_type = layer_decay_type
+
+ def get_scale(self, layer_id):
+ return self.values[layer_id]
+
+ def get_layer_id(self, var_name):
+ if self.layer_decay_type == 'single':
+ return get_num_layer_for_convnext_single(var_name, self.depths)
+ else:
+ return get_num_layer_for_convnext(var_name)
+
+
+def get_parameter_groups(model, weight_decay=1e-5, skip_list=(), get_num_layer=None, get_layer_scale=None):
+ parameter_group_names = {}
+ parameter_group_vars = {}
+
+ for name, param in model.named_parameters():
+ if not param.requires_grad:
+ continue # frozen weights
+ if len(param.shape) == 1 or name.endswith(".bias") or name in skip_list or \
+ name.endswith(".gamma") or name.endswith(".beta"):
+ group_name = "no_decay"
+ this_weight_decay = 0.
+ else:
+ group_name = "decay"
+ this_weight_decay = weight_decay
+ if get_num_layer is not None:
+ layer_id = get_num_layer(name)
+ group_name = "layer_%d_%s" % (layer_id, group_name)
+ else:
+ layer_id = None
+
+ if group_name not in parameter_group_names:
+ if get_layer_scale is not None:
+ scale = get_layer_scale(layer_id)
+ else:
+ scale = 1.
+
+ parameter_group_names[group_name] = {
+ "weight_decay": this_weight_decay,
+ "params": [],
+ "lr_scale": scale
+ }
+ parameter_group_vars[group_name] = {
+ "weight_decay": this_weight_decay,
+ "params": [],
+ "lr_scale": scale
+ }
+
+ parameter_group_vars[group_name]["params"].append(param)
+ parameter_group_names[group_name]["params"].append(name)
+ print("Param groups = %s" % json.dumps(parameter_group_names, indent=2))
+ return list(parameter_group_vars.values())
+
+
+def create_optimizer(args, model, get_num_layer=None, get_layer_scale=None, filter_bias_and_bn=True, skip_list=None):
+ opt_lower = args.opt.lower()
+ weight_decay = args.weight_decay
+ # if weight_decay and filter_bias_and_bn:
+ if filter_bias_and_bn:
+ skip = {}
+ if skip_list is not None:
+ skip = skip_list
+ elif hasattr(model, 'no_weight_decay'):
+ skip = model.no_weight_decay()
+ parameters = get_parameter_groups(model, weight_decay, skip, get_num_layer, get_layer_scale)
+ weight_decay = 0.
+ else:
+ parameters = model.parameters()
+
+ if 'fused' in opt_lower:
+ assert has_apex and torch.cuda.is_available(), 'APEX and CUDA required for fused optimizers'
+
+ opt_args = dict(lr=args.lr, weight_decay=weight_decay)
+ if hasattr(args, 'opt_eps') and args.opt_eps is not None:
+ opt_args['eps'] = args.opt_eps
+ if hasattr(args, 'opt_betas') and args.opt_betas is not None:
+ opt_args['betas'] = args.opt_betas
+
+ opt_split = opt_lower.split('_')
+ opt_lower = opt_split[-1]
+ if opt_lower == 'sgd' or opt_lower == 'nesterov':
+ opt_args.pop('eps', None)
+ optimizer = optim.SGD(parameters, momentum=args.momentum, nesterov=True, **opt_args)
+ elif opt_lower == 'momentum':
+ opt_args.pop('eps', None)
+ optimizer = optim.SGD(parameters, momentum=args.momentum, nesterov=False, **opt_args)
+ elif opt_lower == 'adam':
+ optimizer = optim.Adam(parameters, **opt_args)
+ elif opt_lower == 'adamw':
+ optimizer = optim.AdamW(parameters, **opt_args)
+ elif opt_lower == 'nadam':
+ optimizer = Nadam(parameters, **opt_args)
+ elif opt_lower == 'radam':
+ optimizer = RAdam(parameters, **opt_args)
+ elif opt_lower == 'adamp':
+ optimizer = AdamP(parameters, wd_ratio=0.01, nesterov=True, **opt_args)
+ elif opt_lower == 'sgdp':
+ optimizer = SGDP(parameters, momentum=args.momentum, nesterov=True, **opt_args)
+ elif opt_lower == 'adadelta':
+ optimizer = optim.Adadelta(parameters, **opt_args)
+ elif opt_lower == 'adafactor':
+ if not args.lr:
+ opt_args['lr'] = None
+ optimizer = Adafactor(parameters, **opt_args)
+ elif opt_lower == 'adahessian':
+ optimizer = Adahessian(parameters, **opt_args)
+ elif opt_lower == 'rmsprop':
+ optimizer = optim.RMSprop(parameters, alpha=0.9, momentum=args.momentum, **opt_args)
+ elif opt_lower == 'rmsproptf':
+ optimizer = RMSpropTF(parameters, alpha=0.9, momentum=args.momentum, **opt_args)
+ elif opt_lower == 'novograd':
+ optimizer = NovoGrad(parameters, **opt_args)
+ elif opt_lower == 'nvnovograd':
+ optimizer = NvNovoGrad(parameters, **opt_args)
+ elif opt_lower == 'fusedsgd':
+ opt_args.pop('eps', None)
+ optimizer = FusedSGD(parameters, momentum=args.momentum, nesterov=True, **opt_args)
+ elif opt_lower == 'fusedmomentum':
+ opt_args.pop('eps', None)
+ optimizer = FusedSGD(parameters, momentum=args.momentum, nesterov=False, **opt_args)
+ elif opt_lower == 'fusedadam':
+ optimizer = FusedAdam(parameters, adam_w_mode=False, **opt_args)
+ elif opt_lower == 'fusedadamw':
+ optimizer = FusedAdam(parameters, adam_w_mode=True, **opt_args)
+ elif opt_lower == 'fusedlamb':
+ optimizer = FusedLAMB(parameters, **opt_args)
+ elif opt_lower == 'fusednovograd':
+ opt_args.setdefault('betas', (0.95, 0.98))
+ optimizer = FusedNovoGrad(parameters, **opt_args)
+ else:
+ assert False and "Invalid optimizer"
+
+ if len(opt_split) > 1:
+ if opt_split[0] == 'lookahead':
+ optimizer = Lookahead(optimizer)
+
+ return optimizer
diff --git a/clean/image/aide/requirements.txt b/clean/image/aide/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..d072b8dd2b1fccc6dfd685d5e7dbaa38598753de
--- /dev/null
+++ b/clean/image/aide/requirements.txt
@@ -0,0 +1,37 @@
+einops==0.6.1
+fairscale==0.4.13
+filelock==3.13.1
+ftfy==6.1.3
+h5py==3.10.0
+imgaug==0.2.6
+keras==2.11.0
+kornia==0.7.2
+kornia_rs==0.1.2
+lmdb==1.4.1
+matplotlib==3.7.4
+matplotlib-inline==0.1.6
+numpy==1.24.3
+omegaconf==2.3.0
+open-clip-torch==2.24.0
+openai-clip==1.0.1
+openpyxl==3.1.2
+pandas==2.0.3
+Pillow==9.5.0
+safetensors==0.4.1
+scikit-image==0.20.0
+scikit-learn==1.3.2
+scipy==1.9.1
+sentencepiece==0.2.0
+streamlit==1.30.0
+tenacity==8.2.3
+tensorboard==2.11.2
+tensorboard-data-server==0.6.1
+tensorboard-plugin-wit==1.8.1
+tensorboardX==2.6.2.2
+timm==0.9.6
+torch==1.11.0
+torch-fidelity==0.3.0
+torchmetrics==0.6.0
+torchsummary==1.5.1
+torchvision==0.12.0
+tqdm==4.66.1
diff --git a/clean/image/aide/scripts/eval.sh b/clean/image/aide/scripts/eval.sh
new file mode 100644
index 0000000000000000000000000000000000000000..6438f26d982fa4548b1299436966a54ec2e39998
--- /dev/null
+++ b/clean/image/aide/scripts/eval.sh
@@ -0,0 +1,21 @@
+
+GPU_NUM=8
+WORLD_SIZE=1
+RANK=0
+MASTER_ADDR=localhost
+MASTER_PORT=29572
+
+DISTRIBUTED_ARGS="
+ --nproc_per_node $GPU_NUM \
+ --nnodes $WORLD_SIZE \
+ --node_rank $RANK \
+ --master_addr $MASTER_ADDR \
+ --master_port $MASTER_PORT
+"
+PY_ARGS=${@:1} # Any other arguments
+python -m torch.distributed.launch $DISTRIBUTED_ARGS main_finetune.py \
+ --model AIDE \
+ --batch_size 32 \
+ --blr 5e-4 \
+ --epochs 5 \
+ ${PY_ARGS}
diff --git a/clean/image/aide/scripts/train.sh b/clean/image/aide/scripts/train.sh
new file mode 100644
index 0000000000000000000000000000000000000000..d781a2a63c06acfcb8123369486404ec3a4aa9ee
--- /dev/null
+++ b/clean/image/aide/scripts/train.sh
@@ -0,0 +1,24 @@
+
+GPU_NUM=8
+WORLD_SIZE=1
+RANK=0
+MASTER_ADDR=localhost
+MASTER_PORT=29512
+
+
+DISTRIBUTED_ARGS="
+ --nproc_per_node $GPU_NUM \
+ --nnodes $WORLD_SIZE \
+ --node_rank $RANK \
+ --master_addr $MASTER_ADDR \
+ --master_port $MASTER_PORT
+"
+
+PY_ARGS=${@:1} # Any other arguments
+
+python -m torch.distributed.launch $DISTRIBUTED_ARGS main_finetune.py \
+ --model AIDE \
+ --batch_size 32 \
+ --blr 1e-4 \
+ --epochs 20 \
+ ${PY_ARGS}
diff --git a/clean/image/aide/utils.py b/clean/image/aide/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..923e5d0242e3032e2fd3e9af2ffa6cf0ffa425da
--- /dev/null
+++ b/clean/image/aide/utils.py
@@ -0,0 +1,578 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import os
+import math
+import time
+from collections import defaultdict, deque
+import datetime
+import numpy as np
+from timm.utils import get_state_dict
+
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+from torch import inf
+
+from tensorboardX import SummaryWriter
+from collections import OrderedDict
+
+def str2bool(v):
+ """
+ Converts string to bool type; enables command line
+ arguments in the format of '--arg1 true --arg2 false'
+ """
+ if isinstance(v, bool):
+ return v
+ if v.lower() in ('yes', 'true', 't', 'y', '1'):
+ return True
+ elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+ return False
+ else:
+ raise argparse.ArgumentTypeError('Boolean value expected.')
+
+class SmoothedValue(object):
+ """Track a series of values and provide access to smoothed values over a
+ window or the global series average.
+ """
+
+ def __init__(self, window_size=20, fmt=None):
+ if fmt is None:
+ fmt = "{median:.4f} ({global_avg:.4f})"
+ self.deque = deque(maxlen=window_size)
+ self.total = 0.0
+ self.count = 0
+ self.fmt = fmt
+
+ def update(self, value, n=1):
+ self.deque.append(value)
+ self.count += n
+ self.total += value * n
+
+ def synchronize_between_processes(self):
+ """
+ Warning: does not synchronize the deque!
+ """
+ if not is_dist_avail_and_initialized():
+ return
+ t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
+ dist.barrier()
+ dist.all_reduce(t)
+ t = t.tolist()
+ self.count = int(t[0])
+ self.total = t[1]
+
+ @property
+ def median(self):
+ d = torch.tensor(list(self.deque))
+ return d.median().item()
+
+ @property
+ def avg(self):
+ d = torch.tensor(list(self.deque), dtype=torch.float32)
+ return d.mean().item()
+
+ @property
+ def global_avg(self):
+ return self.total / self.count
+
+ @property
+ def max(self):
+ return max(self.deque)
+
+ @property
+ def value(self):
+ return self.deque[-1]
+
+ def __str__(self):
+ return self.fmt.format(
+ median=self.median,
+ avg=self.avg,
+ global_avg=self.global_avg,
+ max=self.max,
+ value=self.value)
+
+
+class MetricLogger(object):
+ def __init__(self, delimiter="\t"):
+ self.meters = defaultdict(SmoothedValue)
+ self.delimiter = delimiter
+
+ def update(self, **kwargs):
+ for k, v in kwargs.items():
+ if v is None:
+ continue
+ if isinstance(v, torch.Tensor):
+ v = v.item()
+ assert isinstance(v, (float, int))
+ self.meters[k].update(v)
+
+ def __getattr__(self, attr):
+ if attr in self.meters:
+ return self.meters[attr]
+ if attr in self.__dict__:
+ return self.__dict__[attr]
+ raise AttributeError("'{}' object has no attribute '{}'".format(
+ type(self).__name__, attr))
+
+ def __str__(self):
+ loss_str = []
+ for name, meter in self.meters.items():
+ loss_str.append(
+ "{}: {}".format(name, str(meter))
+ )
+ return self.delimiter.join(loss_str)
+
+ def synchronize_between_processes(self):
+ for meter in self.meters.values():
+ meter.synchronize_between_processes()
+
+ def add_meter(self, name, meter):
+ self.meters[name] = meter
+
+ def log_every(self, iterable, print_freq, header=None):
+ i = 0
+ if not header:
+ header = ''
+ start_time = time.time()
+ end = time.time()
+ iter_time = SmoothedValue(fmt='{avg:.4f}')
+ data_time = SmoothedValue(fmt='{avg:.4f}')
+ space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
+ log_msg = [
+ header,
+ '[{0' + space_fmt + '}/{1}]',
+ 'eta: {eta}',
+ '{meters}',
+ 'time: {time}',
+ 'data: {data}'
+ ]
+ if torch.cuda.is_available():
+ log_msg.append('max mem: {memory:.0f}')
+ log_msg = self.delimiter.join(log_msg)
+ MB = 1024.0 * 1024.0
+ for obj in iterable:
+ data_time.update(time.time() - end)
+ yield obj
+ iter_time.update(time.time() - end)
+ if i % print_freq == 0 or i == len(iterable) - 1:
+ eta_seconds = iter_time.global_avg * (len(iterable) - i)
+ eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+ if torch.cuda.is_available():
+ print(log_msg.format(
+ i, len(iterable), eta=eta_string,
+ meters=str(self),
+ time=str(iter_time), data=str(data_time),
+ memory=torch.cuda.max_memory_allocated() / MB))
+ else:
+ print(log_msg.format(
+ i, len(iterable), eta=eta_string,
+ meters=str(self),
+ time=str(iter_time), data=str(data_time)))
+ i += 1
+ end = time.time()
+ total_time = time.time() - start_time
+ total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+ print('{} Total time: {} ({:.4f} s / it)'.format(
+ header, total_time_str, total_time / len(iterable)))
+
+
+class TensorboardLogger(object):
+ def __init__(self, log_dir):
+ self.writer = SummaryWriter(logdir=log_dir)
+ self.step = 0
+
+ def set_step(self, step=None):
+ if step is not None:
+ self.step = step
+ else:
+ self.step += 1
+
+ def update(self, head='scalar', step=None, **kwargs):
+ for k, v in kwargs.items():
+ if v is None:
+ continue
+ if isinstance(v, torch.Tensor):
+ v = v.item()
+ assert isinstance(v, (float, int))
+ self.writer.add_scalar(head + "/" + k, v, self.step if step is None else step)
+
+ def flush(self):
+ self.writer.flush()
+
+
+class WandbLogger(object):
+ def __init__(self, args):
+ self.args = args
+
+ try:
+ import wandb
+ self._wandb = wandb
+ except ImportError:
+ raise ImportError(
+ "To use the Weights and Biases Logger please install wandb."
+ "Run `pip install wandb` to install it."
+ )
+
+ # Initialize a W&B run
+ if self._wandb.run is None:
+ self._wandb.init(
+ project=args.project,
+ config=args
+ )
+
+ def log_epoch_metrics(self, metrics, commit=True):
+ """
+ Log train/test metrics onto W&B.
+ """
+ # Log number of model parameters as W&B summary
+ self._wandb.summary['n_parameters'] = metrics.get('n_parameters', None)
+ metrics.pop('n_parameters', None)
+
+ # Log current epoch
+ self._wandb.log({'epoch': metrics.get('epoch')}, commit=False)
+ metrics.pop('epoch')
+
+ for k, v in metrics.items():
+ if 'train' in k:
+ self._wandb.log({f'Global Train/{k}': v}, commit=False)
+ elif 'test' in k:
+ self._wandb.log({f'Global Test/{k}': v}, commit=False)
+
+ self._wandb.log({})
+
+ def log_checkpoints(self):
+ output_dir = self.args.output_dir
+ model_artifact = self._wandb.Artifact(
+ self._wandb.run.id + "_model", type="model"
+ )
+
+ model_artifact.add_dir(output_dir)
+ self._wandb.log_artifact(model_artifact, aliases=["latest", "best"])
+
+ def set_steps(self):
+ # Set global training step
+ self._wandb.define_metric('Rank-0 Batch Wise/*', step_metric='Rank-0 Batch Wise/global_train_step')
+ # Set epoch-wise step
+ self._wandb.define_metric('Global Train/*', step_metric='epoch')
+ self._wandb.define_metric('Global Test/*', step_metric='epoch')
+
+
+def setup_for_distributed(is_master):
+ """
+ This function disables printing when not in master process
+ """
+ import builtins as __builtin__
+ builtin_print = __builtin__.print
+
+ def print(*args, **kwargs):
+ force = kwargs.pop('force', False)
+ if is_master or force:
+ builtin_print(*args, **kwargs)
+
+ __builtin__.print = print
+
+
+def is_dist_avail_and_initialized():
+ if not dist.is_available():
+ return False
+ if not dist.is_initialized():
+ return False
+ return True
+
+
+def get_world_size():
+ if not is_dist_avail_and_initialized():
+ return 1
+ return dist.get_world_size()
+
+
+def get_rank():
+ if not is_dist_avail_and_initialized():
+ return 0
+ return dist.get_rank()
+
+
+def is_main_process():
+ return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+ if is_main_process():
+ torch.save(*args, **kwargs)
+
+def init_distributed_mode(args):
+
+ if args.dist_on_itp:
+ args.rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
+ args.world_size = int(os.environ['OMPI_COMM_WORLD_SIZE'])
+ args.gpu = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])
+ args.dist_url = "tcp://%s:%s" % (os.environ['MASTER_ADDR'], os.environ['MASTER_PORT'])
+ os.environ['LOCAL_RANK'] = str(args.gpu)
+ os.environ['RANK'] = str(args.rank)
+ os.environ['WORLD_SIZE'] = str(args.world_size)
+ # ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"]
+ elif 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+ args.rank = int(os.environ["RANK"])
+ args.world_size = int(os.environ['WORLD_SIZE'])
+ args.gpu = int(os.environ['LOCAL_RANK'])
+ elif 'SLURM_PROCID' in os.environ:
+ args.rank = int(os.environ['SLURM_PROCID'])
+ args.gpu = args.rank % torch.cuda.device_count()
+
+ os.environ['RANK'] = str(args.rank)
+ os.environ['LOCAL_RANK'] = str(args.gpu)
+ os.environ['WORLD_SIZE'] = str(args.world_size)
+ else:
+ print('Not using distributed mode')
+ args.distributed = False
+ return
+
+ args.distributed = True
+
+ torch.cuda.set_device(args.gpu)
+ args.dist_backend = 'nccl'
+ print('| distributed init (rank {}): {}, gpu {}'.format(
+ args.rank, args.dist_url, args.gpu), flush=True)
+ torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
+ world_size=args.world_size, rank=args.rank)
+ torch.distributed.barrier()
+ setup_for_distributed(args.rank == 0)
+
+def all_reduce_mean(x):
+ world_size = get_world_size()
+ if world_size > 1:
+ x_reduce = torch.tensor(x).cuda()
+ dist.all_reduce(x_reduce)
+ x_reduce /= world_size
+ return x_reduce.item()
+ else:
+ return x
+
+def load_state_dict(model, state_dict, prefix='', ignore_missing="relative_position_index"):
+ missing_keys = []
+ unexpected_keys = []
+ error_msgs = []
+ # copy state_dict so _load_from_state_dict can modify it
+ metadata = getattr(state_dict, '_metadata', None)
+ state_dict = state_dict.copy()
+ if metadata is not None:
+ state_dict._metadata = metadata
+
+ def load(module, prefix=''):
+ local_metadata = {} if metadata is None else metadata.get(
+ prefix[:-1], {})
+ module._load_from_state_dict(
+ state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
+ for name, child in module._modules.items():
+ if child is not None:
+ load(child, prefix + name + '.')
+
+ load(model, prefix=prefix)
+
+ warn_missing_keys = []
+ ignore_missing_keys = []
+ for key in missing_keys:
+ keep_flag = True
+ for ignore_key in ignore_missing.split('|'):
+ if ignore_key in key:
+ keep_flag = False
+ break
+ if keep_flag:
+ warn_missing_keys.append(key)
+ else:
+ ignore_missing_keys.append(key)
+
+ missing_keys = warn_missing_keys
+
+ if len(missing_keys) > 0:
+ print("Weights of {} not initialized from pretrained model: {}".format(
+ model.__class__.__name__, missing_keys))
+ if len(unexpected_keys) > 0:
+ print("Weights from pretrained model not used in {}: {}".format(
+ model.__class__.__name__, unexpected_keys))
+ if len(ignore_missing_keys) > 0:
+ print("Ignored weights of {} not initialized from pretrained model: {}".format(
+ model.__class__.__name__, ignore_missing_keys))
+ if len(error_msgs) > 0:
+ print('\n'.join(error_msgs))
+
+
+class NativeScalerWithGradNormCount:
+ state_dict_key = "amp_scaler"
+
+ def __init__(self):
+ self._scaler = torch.cuda.amp.GradScaler()
+
+ def __call__(self, loss, optimizer, clip_grad=None, parameters=None, create_graph=False, update_grad=True):
+ self._scaler.scale(loss).backward(create_graph=create_graph)
+ if update_grad:
+ if clip_grad is not None:
+ assert parameters is not None
+ self._scaler.unscale_(optimizer) # unscale the gradients of optimizer's assigned params in-place
+ norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad)
+ else:
+ self._scaler.unscale_(optimizer)
+ norm = get_grad_norm_(parameters)
+ self._scaler.step(optimizer)
+ self._scaler.update()
+ else:
+ norm = None
+ return norm
+
+ def state_dict(self):
+ return self._scaler.state_dict()
+
+ def load_state_dict(self, state_dict):
+ self._scaler.load_state_dict(state_dict)
+
+
+def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
+ if isinstance(parameters, torch.Tensor):
+ parameters = [parameters]
+ parameters = [p for p in parameters if p.grad is not None]
+ norm_type = float(norm_type)
+ if len(parameters) == 0:
+ return torch.tensor(0.)
+ device = parameters[0].grad.device
+ if norm_type == inf:
+ total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
+ else:
+ total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type)
+ return total_norm
+
+def save_model(args, epoch, model, model_without_ddp, optimizer, loss_scaler, model_ema=None):
+ output_dir = Path(args.output_dir)
+ epoch_name = str(epoch)
+ checkpoint_paths = [output_dir / ('checkpoint-%s.pth' % epoch_name)]
+ for checkpoint_path in checkpoint_paths:
+ to_save = {
+ 'model': model_without_ddp.state_dict(),
+ 'optimizer': optimizer.state_dict(),
+ 'epoch': epoch,
+ 'scaler': loss_scaler.state_dict(),
+ 'args': args,
+ }
+
+ if model_ema is not None:
+ to_save['model_ema'] = get_state_dict(model_ema)
+
+ save_on_master(to_save, checkpoint_path)
+
+ if is_main_process() and isinstance(epoch, int):
+ to_del = epoch - args.save_ckpt_num * args.save_ckpt_freq
+ old_ckpt = output_dir / ('checkpoint-%s.pth' % to_del)
+ if os.path.exists(old_ckpt):
+ os.remove(old_ckpt)
+
+def auto_load_model(args, model, model_without_ddp, optimizer, loss_scaler, model_ema=None):
+ output_dir = Path(args.output_dir)
+ if args.auto_resume and len(args.resume) == 0:
+ import glob
+ all_checkpoints = glob.glob(os.path.join(output_dir, 'checkpoint-*.pth'))
+ latest_ckpt = -1
+ for ckpt in all_checkpoints:
+ t = ckpt.split('-')[-1].split('.')[0]
+ if t.isdigit():
+ latest_ckpt = max(int(t), latest_ckpt)
+ if latest_ckpt >= 0:
+ args.resume = os.path.join(output_dir, 'checkpoint-%d.pth' % latest_ckpt)
+ print("Auto resume checkpoint: %s" % args.resume)
+
+ if args.resume:
+ if args.resume.startswith('https'):
+ checkpoint = torch.hub.load_state_dict_from_url(
+ args.resume, map_location='cpu', check_hash=True)
+ else:
+ checkpoint = torch.load(args.resume, map_location='cpu')
+
+ model_without_ddp.load_state_dict(checkpoint['model'])
+ print("Resume checkpoint %s" % args.resume)
+ if 'optimizer' in checkpoint and 'epoch' in checkpoint:
+ optimizer.load_state_dict(checkpoint['optimizer'])
+ if not isinstance(checkpoint['epoch'], str): # does not support resuming with 'best', 'best-ema'
+ args.start_epoch = checkpoint['epoch'] + 1
+ else:
+ assert args.eval, 'Does not support resuming with checkpoint-best'
+ if hasattr(args, 'model_ema') and args.model_ema:
+ if 'model_ema' in checkpoint.keys():
+ model_ema.ema.load_state_dict(checkpoint['model_ema'])
+ else:
+ model_ema.ema.load_state_dict(checkpoint['model'])
+ if 'scaler' in checkpoint:
+ loss_scaler.load_state_dict(checkpoint['scaler'])
+ print("With optim & sched!")
+
+def cosine_scheduler(base_value, final_value, epochs, niter_per_ep, warmup_epochs=0,
+ start_warmup_value=0, warmup_steps=-1):
+ warmup_schedule = np.array([])
+ warmup_iters = warmup_epochs * niter_per_ep
+ if warmup_steps > 0:
+ warmup_iters = warmup_steps
+ print("Set warmup steps = %d" % warmup_iters)
+ if warmup_epochs > 0:
+ warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters)
+
+ iters = np.arange(epochs * niter_per_ep - warmup_iters)
+ schedule = np.array(
+ [final_value + 0.5 * (base_value - final_value) * (1 + math.cos(math.pi * i / (len(iters)))) for i in iters])
+
+ schedule = np.concatenate((warmup_schedule, schedule))
+
+ assert len(schedule) == epochs * niter_per_ep
+ return schedule
+
+def adjust_learning_rate(optimizer, epoch, args):
+ """Decay the learning rate with half-cycle cosine after warmup"""
+ if epoch < args.warmup_epochs:
+ lr = args.lr * epoch / args.warmup_epochs
+ else:
+ lr = args.min_lr + (args.lr - args.min_lr) * 0.5 * \
+ (1. + math.cos(math.pi * (epoch - args.warmup_epochs) / (args.epochs - args.warmup_epochs)))
+ for param_group in optimizer.param_groups:
+ if "lr_scale" in param_group:
+ param_group["lr"] = lr * param_group["lr_scale"]
+ else:
+ param_group["lr"] = lr
+ return lr
+
+def remap_checkpoint_keys(ckpt):
+ new_ckpt = OrderedDict()
+ for k, v in ckpt.items():
+ if k.startswith('encoder'):
+ k = '.'.join(k.split('.')[1:]) # remove encoder in the name
+ if k.endswith('kernel'):
+ k = '.'.join(k.split('.')[:-1]) # remove kernel in the name
+ new_k = k + '.weight'
+ if len(v.shape) == 3: # resahpe standard convolution
+ kv, in_dim, out_dim = v.shape
+ ks = int(math.sqrt(kv))
+ new_ckpt[new_k] = v.permute(2, 1, 0).\
+ reshape(out_dim, in_dim, ks, ks).transpose(3, 2)
+ elif len(v.shape) == 2: # reshape depthwise convolution
+ kv, dim = v.shape
+ ks = int(math.sqrt(kv))
+ new_ckpt[new_k] = v.permute(1, 0).\
+ reshape(dim, 1, ks, ks).transpose(3, 2)
+ continue
+ elif 'ln' in k or 'linear' in k:
+ k = k.split('.')
+ k.pop(-2) # remove ln and linear in the name
+ new_k = '.'.join(k)
+ else:
+ new_k = k
+ new_ckpt[new_k] = v
+
+ # reshape grn affine parameters and biases
+ for k, v in new_ckpt.items():
+ if k.endswith('bias') and len(v.shape) != 1:
+ new_ckpt[k] = v.reshape(-1)
+ # elif 'grn' in k:
+ # new_ckpt[k] = v.unsqueeze(0).unsqueeze(1)
+ return new_ckpt
\ No newline at end of file
diff --git a/clean/image/cospy/.gitignore b/clean/image/cospy/.gitignore
new file mode 100644
index 0000000000000000000000000000000000000000..e590da532e2c928cbc4f3fab8dd5cd48f5051a58
--- /dev/null
+++ b/clean/image/cospy/.gitignore
@@ -0,0 +1,11 @@
+__pycache__/
+ckpt/
+data/test/AIGCDetectionBenchMark/test
+data/test/Co-Spy-Bench/real_image_examples
+data/test/Co-Spy-Bench/synthetic
+data/train/progan/train
+data/train/progan/val
+data/train/sd-v1_4/mscoco2017
+data/train/sd-v1_4/stable-diffusion-v1-4
+pretrained/progan
+pretrained/sd-v1_4
diff --git a/clean/image/cospy/LICENSE b/clean/image/cospy/LICENSE
new file mode 100644
index 0000000000000000000000000000000000000000..4ba4b5e9c15b4c9be378ae45506f49839fd33716
--- /dev/null
+++ b/clean/image/cospy/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Siyuan Cheng
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/clean/image/cospy/README.md b/clean/image/cospy/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..308b0ca684fa2fea7207024dbe1243f9808ce7fa
--- /dev/null
+++ b/clean/image/cospy/README.md
@@ -0,0 +1,337 @@
+
+
+
+
+# CO-SPY: Combining Semantic and Pixel Features to Detect Synthetic Images by AI
+
+
+
+
+
+
+
+Table of Contents
+=================
+- [Table of Contents](#table-of-contents)
+ - [Overview](#overview)
+ - [CO-SPY-Bench](#co-spy-bench)
+ - [Main Code Architecutre](#main-code-architecture)
+ - [Environments](#environments)
+ - [Prerequisites](#prerequisites)
+ - [Download Datasets](#download-datasets)
+ - [Download Pre-trained Weights](#download-pre-trained-weights)
+ - [Experiments](#experiments)
+ - [Evaluation of Pre-trained Detectors](#evaluation-of-pre-trained-detectors)
+ - [Inference on a Single Image](#inference-on-a-single-image)
+ - [Training](#training)
+ - [Citation](#citation)
+ - [Acknowledgement](#acknowledgement)
+
+## Overview
+- This is the official implementation for CVPR 2025 paper "[CO-SPY: Combining Semantic and Pixel Features to Detect Synthetic Images by AI](https://openaccess.thecvf.com/content/CVPR2025/html/Cheng_CO-SPY_Combining_Semantic_and_Pixel_Features_to_Detect_Synthetic_Images_CVPR_2025_paper.html)".
+- [[arXiv](https://arxiv.org/abs/2503.18286)\] | \[[poster](https://www.cs.purdue.edu/homes/cheng535/static/slides/COSPY_poster.pdf)\]
+
+
+
+## CO-SPY-Bench
+We have released our benchmark on [Huggingface](https://huggingface.co/datasets/ruojiruoli/Co-Spy-Bench), designed to offer diverse and comprehensive coverage of the latest generative models:
+- Captions are sourced from five real-world datasets: MSCOCO2017, CC3M, Flickr, TextCaps, and SBU.
+- Synthetic images are generated using 22 different models, covering a wide range of architectures.
+- Diverse generation parameters, such as diffusion steps and guidance scales, are used to enrich variability.
+
+## Main Code Architecture
+ .
+ ├── data # Dataset folder
+ │ ├── in_the_wild # CO-SPY-Bench in-the-wild synthetic samples
+ │ ├── test # Test dataset (CO-SPY-Bench & AIGCDetectionBenchMark)
+ │ └── train # Training dataset (DRCT-2M & CNNDet)
+ ├── dataSets # Various dataset classes
+ ├── detectors # Various detector classes
+ │ ├── progan # Detectors for CNNDet training set
+ │ └── sd-v1_4 # Detectors for DRCT-2M training set
+ ├── pretrained # Pre-trained weights
+ ├── main.py # Main function
+ ├── evaluate.py # Evaluation function
+ ├── main.py # Main function (Entry point)
+ ├── train.py # Training function
+ ├── train.sh # Recommended training script
+ └── utils.py # Utility functions
+
+## Environments
+```bash
+# Create python environment (optional)
+conda env create -f environment.yml
+source activate cospy
+```
+
+## Prerequisites
+Please download the required datasets and pre-trained weights for full evaluation.
+
+### Download Datasets
+Make sure you have `7z` an `unzip` installed. You can install them via conda:
+```bash
+conda install p7zip
+conda install unzip
+```
+
+To download the training and test datasets, run the following commands respectively:
+```bash
+###############################
+# Download the training dataset
+###############################
+# Download CNNDet (ProGAN) training set
+cd data/train/progan
+sh download.sh
+
+# Download DRCT-2M (Stable Diffusion v1.4) training set
+cd data/train/sd-v1_4
+sh download.sh
+
+###############################
+# Download the test dataset
+###############################
+# Download AIGCDetectionBenchMark test set
+cd data/test/AIGCDetectionBenchMark
+python download.py
+
+# Download CO-SPY-Bench test set
+cd data/test/Co-Spy-Bench
+sh download.sh
+```
+
+Finally, the directory structure should look like this:
+```
+.
+├── data
+│ ├── test
+│ │ ├── AIGCDetectionBenchMark
+│ │ │ └── test
+│ │ └── Co-Spy-Bench
+│ │ ├── real_image_examples
+│ │ └── synthetic
+│ └── train
+│ ├── progan
+│ │ ├── train
+│ │ └── val
+│ └── sd-v1_4
+│ ├── mscoco2017
+│ └── stable_diffusion_v1-4
+```
+*Note: Please ensure the use of these datasets complies with the original licenses.*
+
+
+> Please refer to `data/in_the_wild/README.md` for detailed instructions on accessing the **CO-SPY-Bench in-the-wild synthetic samples**.
+>
+> *Note: Some samples are temporarily unavailable (e.g., [instavibe.ai](https://www.instavibe.ai/)). As we cannot confirm whether redistributing these images would violate the original sources’ intellectual property rights, we choose not to release them at this time. This decision and its rationale are discussed in [Issue #6](https://github.com/Megum1/CO-SPY/issues/6). We apologize for any inconvenience this may cause.*
+
+
+### Download Pre-trained Weights
+```bash
+# Download the pre-trained weights
+cd pretrained
+sh download.sh
+# It contains the pre-trained weights on CNNDet (`progan`) and DRCT-2M (`sd-v1_4`) training sets.
+```
+
+
+
+## Experiments
+We provide the source code for training and evaluating the CO-SPY detector.
+
+### Evaluation of Pre-trained Detectors
+To evaluate the pre-trained detector (trained on DRCT-2M (`sd-v1_4`)) on [CO-SPY-Bench](https://huggingface.co/datasets/ruojiruoli/Co-Spy-Bench), run
+```bash
+python main.py --gpu 0 --phase eval --train_dataset sd-v1_4 --pretrain
+```
+When finished, evaluation results will be saved to `ckpt/sd-v1_4/fusion/pretrain_Co-Spy-Bench`.
+
+
+📊 Results (click to expand)
+
+
+Results are highly consistent with the original results reported in Table 2, with slight differences arising from improvements in hyper-parameter settings.
+Original results can be referred to in Issue #7.
+
+Results are highly consistent with the original results reported in Table 8 in Appendix K, with slight differences arising from improvements in hyper-parameter settings.
+
+
+| Model | AP | Accuracy |
+|-------|------|----------|
+| ADM | 89.98 | 79.17 |
+| biggan | 98.00 | 94.30 |
+| cyclegan | 99.47 | 98.98 |
+| DALLE2 | 97.14 | 87.80 |
+| gaugan | 98.24 | 95.05 |
+| Glide | 96.71 | 90.40 |
+| Midjourney | 94.55 | 87.19 |
+| progan | 100.00 | 100.00 |
+| stable_diffusion_v_1_4 | 92.99 | 85.34 |
+| stable_diffusion_v_1_5 | 92.96 | 85.44 |
+| stargan | 99.98 | 99.45 |
+| stylegan | 99.69 | 94.81 |
+| stylegan2 | 99.85 | 94.89 |
+| VQDM | 91.96 | 85.83 |
+| whichfaceisreal | 82.27 | 80.55 |
+| wukong | 89.59 | 80.06 |
+| **Average** | **95.21** | **90.58** |
+
+
+
+Evaluation results contain three files, including:
+- `evaluation.log`: detailed evaluation log.
+- `output.json`: predicted synthetic probabilities for each sample.
+- `result.json`: evaluation metrics (dataset size, AP and accuracy) for each test source.
+
+### Inference on a Single Image
+To run inference on a single image (e.g., using the pre-trained detector trained on DRCT-2M):
+```bash
+python main.py --gpu 0 --phase test --train_dataset sd-v1_4 --pretrain
+# The script will prompt for the image file path:
+# "Please enter the image filepath for scanning: "
+imgs/test.png
+# Output (probability - decision):
+# "CO-SPY Prediction: 0.854 - AI-Generated"
+```
+
+### Training
+We provide two training pipelines: (1) end-to-end training and (2) best practice of training semantic and artifact branches separately, followed by calibrating the combined detector.
+
+```bash
+# Train an end-to-end CO-SPY detector for 10 epochs on DRCT-2M
+python main.py --gpu 0 --phase train --mode end2end --train_dataset sd-v1_4 --epochs 10
+# Train an end-to-end CO-SPY detector for 10 epochs on CNNDet
+python main.py --gpu 1 --phase train --mode end2end --train_dataset progan --epochs 10
+```
+The trained model will be saved to `ckpt//end2end`.
+
+*The end-to-end training may not yield optimal performance due to the conflicting nature of semantic and artifact features. We recommend the following best practice for training.*
+```bash
+# Train an optimal CO-SPY detector using the best practice script on DRCT-2M
+bash train.sh --gpu 2 --dataset sd-v1_4
+# Train an optimal CO-SPY detector using the best practice script on CNNDet
+bash train.sh --gpu 2 --dataset progan
+```
+The trained models will be saved to `ckpt//fusion`.
+
+You can also customize the training parameters for the training of each branch and the calibration step:
+```bash
+# Train the semantic component (on CNNDet training set as an example)
+python main.py \
+ --phase train \
+ --gpu 0 \
+ --mode branch \
+ --branch semantic \
+ --train_dataset progan \
+ --epochs 10
+# Train the artifact component
+python main.py \
+ --phase train \
+ --gpu 1 \
+ --mode branch \
+ --branch artifact \
+ --train_dataset progan \
+ --epochs 20
+# Calibrate the combined CO-SPY detector
+python main.py \
+ --phase train \
+ --gpu 2 \
+ --mode fusion \
+ --train_dataset progan \
+ --epochs 2
+```
+The trained branch models and the calibrated fusion model will be saved to `ckpt/progan/semantic`, `ckpt/progan/artifact`, and `ckpt/progan/fusion`, respectively.
+
+After training, you can evaluate the trained model on the test datasets by just changing the `--phase` to `eval` and specifying the `--train_dataset` and `--mode` accordingly. For example, to evaluate the fusion model (best practice), trained on CNNDet, on AIGCDetectionBenchMark, run:
+```bash
+python main.py --gpu 0 --phase eval --train_dataset progan --mode fusion
+```
+
+## Citation
+Please cite our paper if you find it useful for your research.😀
+
+```bibtex
+@inproceedings{Cheng_2025_CVPR,
+ author = {Cheng, Siyuan and Lyu, Lingjuan and Wang, Zhenting and Zhang, Xiangyu and Sehwag, Vikash},
+ title = {CO-SPY: Combining Semantic and Pixel Features to Detect Synthetic Images by AI},
+ booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+ month = {June},
+ year = {2025},
+ pages = {13455-13465}
+}
+```
+
+## Acknowledgement
+We gratefully acknowledge these outstanding works, which have deeply inspired our project!
+- [CNNDetection](https://github.com/PeterWang512/CNNDetection)
+- [UniversalFakeDetect](https://github.com/WisconsinAIVision/UniversalFakeDetect)
+- [NPR-DeepfakeDetection](https://github.com/chuangchuangtan/NPR-DeepfakeDetection)
+- [DRCT](https://github.com/beibuwandeluori/DRCT)
diff --git a/clean/image/cospy/SOURCE.md b/clean/image/cospy/SOURCE.md
new file mode 100644
index 0000000000000000000000000000000000000000..61c85f57738190ea4db2227158479e8a952c8f18
--- /dev/null
+++ b/clean/image/cospy/SOURCE.md
@@ -0,0 +1,17 @@
+# Source: image/cospy
+
+| Field | Value |
+|---|---|
+| Upstream | https://github.com/Megum1/CO-SPY |
+| Paper | https://arxiv.org/abs/2503.18286 |
+| Commit SHA | **not recorded** -- the vendoring step did not preserve it |
+| Mirrored on | 2026-09-15 |
+| Upstream license | LICENSE |
+
+This is a **mirror**, stripped to the files needed for inference. The full
+untouched snapshot is at `archive/image__cospy.tar.gz`.
+
+This code is the work of its original authors and is **not** covered by the
+DeepSafe project license. If you are an author and want this removed, open an
+issue on https://github.com/deepsafehq/deepsafe-bench and it will be taken down
+within 48 hours, no questions asked.
diff --git a/clean/image/cospy/dataSets/__init__.py b/clean/image/cospy/dataSets/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ecb7a92c483a5b8dd269d8ba545ca55c5217d308
--- /dev/null
+++ b/clean/image/cospy/dataSets/__init__.py
@@ -0,0 +1,67 @@
+from .dataset import TrainDataset, CoSpyBenchTestDataset, AIGCDetectTestDataset
+
+
+# Co-Spy-Bench: List of evaluated real datasets
+CoSpyBench_DATASET_LIST = ["mscoco", "flickr", "cc3m", "textcaps", "sbu"]
+
+# Co-Spy-Bench: List of evaluated generative models
+CoSpyBench_MODEL_LIST = [
+ # CompVis
+ "CompVis@ldm-text2im-large-256",
+ "CompVis@stable-diffusion-v1-4",
+ # runwayml
+ "runwayml@stable-diffusion-v1-5",
+ # segmind
+ "segmind@SSD-1B",
+ "segmind@tiny-sd",
+ "segmind@SegMoE-SD-4x2-v0",
+ "segmind@small-sd",
+ # stabilityai
+ "stabilityai@stable-diffusion-2-1",
+ "stabilityai@stable-diffusion-3-medium-diffusers",
+ "stabilityai@sdxl-turbo",
+ "stabilityai@stable-diffusion-2",
+ "stabilityai@stable-diffusion-xl-base-1.0",
+ # playgroundai
+ "playgroundai@playground-v2.5-1024px-aesthetic",
+ "playgroundai@playground-v2-1024px-aesthetic",
+ "playgroundai@playground-v2-512px-base",
+ "playgroundai@playground-v2-256px-base",
+ # PixArt-alpha
+ "PixArt-alpha@PixArt-XL-2-1024-MS",
+ "PixArt-alpha@PixArt-XL-2-512x512",
+ # latent-consistency
+ "latent-consistency@lcm-lora-sdxl",
+ "latent-consistency@lcm-lora-sdv1-5",
+ # black-forest-labs
+ "black-forest-labs@FLUX.1-schnell",
+ "black-forest-labs@FLUX.1-dev",
+]
+
+# AIGCDetectionBenchMark: List of evaluated real datasets
+AIGCDetectionBenchMark_DATASET_LIST = ["test"]
+
+# AIGCDetectionBenchMark: List of evaluated generative models
+AIGCDetectionBenchMark_MODEL_LIST = [
+ "ADM",
+ "biggan",
+ "cyclegan",
+ "DALLE2",
+ "gaugan",
+ "Glide",
+ "Midjourney",
+ "progan",
+ "stable_diffusion_v_1_4",
+ "stable_diffusion_v_1_5",
+ "stargan",
+ "stylegan",
+ "stylegan2",
+ "VQDM",
+ "whichfaceisreal",
+ "wukong",
+]
+
+
+__all__ = ["TrainDataset", "CoSpyBenchTestDataset", "AIGCDetectTestDataset",
+ "CoSpyBench_DATASET_LIST", "CoSpyBench_MODEL_LIST",
+ "AIGCDetectionBenchMark_DATASET_LIST", "AIGCDetectionBenchMark_MODEL_LIST"]
diff --git a/clean/image/cospy/dataSets/cospy_real.py b/clean/image/cospy/dataSets/cospy_real.py
new file mode 100644
index 0000000000000000000000000000000000000000..e9f5925ba812e04e7d6edfe2dcefc3267e30ec40
--- /dev/null
+++ b/clean/image/cospy/dataSets/cospy_real.py
@@ -0,0 +1,93 @@
+import os
+import json
+import torch
+import numpy as np
+from PIL import Image
+import datasets as ds
+
+
+class MSCOCO2017(torch.utils.data.Dataset):
+ def __init__(self, split='val', transform=None):
+ # Split [train: 118287, val: 5000]
+ image_dir = f"data/train/sd-v1_4/mscoco2017/{split}2017"
+ self.image_list = [os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.endswith(('.jpg'))]
+ self.image_list.sort()
+ caption_filepath = f"data/train/sd-v1_4/stable-diffusion-v1-4/{split}2017_gen.json"
+ with open(caption_filepath, 'r') as f:
+ self.captions = json.load(f)
+
+ # Preprocess the images
+ self.transform = transform
+
+ def __len__(self):
+ return len(self.image_list)
+
+ def __getitem__(self, idx):
+ image_filepath = self.image_list[idx]
+ image = Image.open(image_filepath).convert('RGB')
+ if self.transform:
+ image = self.transform(image)
+ # Get the captions
+ image_id = os.path.basename(image_filepath)
+ caption = self.captions[image_id]
+ return image, caption
+
+
+class Flickr30k(torch.utils.data.Dataset):
+ def __init__(self, split='test', transform=None):
+ # Split [test: 31014]
+ self.dataset = ds.load_dataset("nlphuji/flickr30k")[split]
+
+ # Preprocess the images
+ self.transform = transform
+
+ def __len__(self):
+ return len(self.dataset)
+
+ def __getitem__(self, idx):
+ example = self.dataset[idx]
+ # PIL RGB image
+ image = example['image']
+ if self.transform:
+ image = self.transform(image)
+ # A list of valid captions
+ caption_list = example['caption']
+ # Randomly select a caption
+ caption = np.random.choice(caption_list)
+ return image, caption
+
+
+class OtherReal(torch.utils.data.Dataset):
+ def __init__(self, dataset_name, transform=None):
+ # Load the test samples from other real datasets (CC3M, SBU, TextCaps)
+ # Note: there are only 2,000 examples for each dataset
+ root_dir = "data/test/Co-Spy-Bench/real_image_examples"
+ if dataset_name == "cc3m":
+ self.image_dir = os.path.join(root_dir, "CC3M")
+ elif dataset_name == "sbu":
+ self.image_dir = os.path.join(root_dir, "SBU")
+ elif dataset_name == "textcaps":
+ self.image_dir = os.path.join(root_dir, "TextCaps")
+ else:
+ raise ValueError(f"Dataset {dataset_name} not supported in OtherReal class.")
+
+ self.image_list = [f for f in os.listdir(self.image_dir) if f.endswith(('.jpg'))]
+ self.image_list.sort()
+
+ # Preprocess the images
+ self.transform = transform
+
+ def __len__(self):
+ return len(self.image_list)
+
+ def __getitem__(self, idx):
+ img_name = self.image_list[idx]
+ img_path = os.path.join(self.image_dir, img_name)
+ image = Image.open(img_path).convert('RGB')
+ if self.transform:
+ image = self.transform(image)
+ caption_filepath = img_path.replace('.jpg', '.json')
+ with open(caption_filepath, 'r') as f:
+ caption_data = json.load(f)
+ caption = caption_data['caption']
+ return image, caption
diff --git a/clean/image/cospy/dataSets/dataset.py b/clean/image/cospy/dataSets/dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..050fadc79b8521cd1c67a5295c229cc8b906f50c
--- /dev/null
+++ b/clean/image/cospy/dataSets/dataset.py
@@ -0,0 +1,148 @@
+import os
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+from torch.utils.data import Dataset
+
+from utils import get_list, png_to_jpeg
+from .cospy_real import MSCOCO2017, Flickr30k, OtherReal
+
+
+class TrainDataset(Dataset):
+ def __init__(self, train_dataset, split="train", add_jpeg=False, transform=None):
+ assert split in ["train", "val"]
+ # Root directory of the training datasets
+ root_dir = f"data/train/{train_dataset}"
+
+ # Load the dataset for training
+ if train_dataset == "progan":
+ real_list = get_list(os.path.join(root_dir, split), must_contain='0_real')
+ fake_list = get_list(os.path.join(root_dir, split), must_contain='1_fake')
+ elif train_dataset == "sd-v1_4":
+ real_list = get_list(os.path.join(root_dir, "mscoco2017", f"{split}2017"))
+ fake_list = get_list(os.path.join(root_dir, "stable-diffusion-v1-4", f"{split}2017"))
+
+ # Setting the labels for the dataset
+ self.labels_dict = {}
+ for i in real_list:
+ self.labels_dict[i] = 0
+ for i in fake_list:
+ self.labels_dict[i] = 1
+
+ # Construct the entire dataset
+ self.total_list = real_list + fake_list
+ np.random.shuffle(self.total_list)
+
+ # JPEG compression
+ self.add_jpeg = add_jpeg
+
+ # Transformations
+ self.transform = transform
+
+ def __len__(self):
+ return len(self.total_list)
+
+ def __getitem__(self, idx):
+ img_path = self.total_list[idx]
+ label = self.labels_dict[img_path]
+ image = Image.open(img_path).convert("RGB")
+
+ # Add JPEG compression
+ if self.add_jpeg:
+ image = png_to_jpeg(image, quality=95)
+
+ # Apply the transformation
+ if self.transform is not None:
+ image = self.transform(image)
+ return image, label
+
+
+class CoSpyBenchTestDataset(Dataset):
+ def __init__(self, dataset, model, num_real=2000, add_jpeg=True, transform=None):
+ # Root path of the Co-Spy-Bench dataset
+ root_path = "data/test/Co-Spy-Bench/synthetic"
+
+ # Load fake images
+ fake_dir = os.path.join(root_path, dataset, model)
+ fake_list = [i for i in os.listdir(fake_dir) if i.endswith(".png")]
+ fake_list.sort()
+ self.fake = [os.path.join(fake_dir, i) for i in fake_list]
+
+ # Take the real images from the dataset
+ if dataset == "mscoco":
+ self.real = MSCOCO2017()
+ elif dataset == "flickr":
+ self.real = Flickr30k()
+ else:
+ self.real = OtherReal(dataset)
+
+ # Ensure the number of real and fake images are the same
+ self.num_real = min(num_real, len(self.real), len(self.fake))
+ self.image_idx = list(range(self.num_real * 2))
+ # First half is real, second half is fake
+ self.labels = [0] * self.num_real + [1] * self.num_real
+
+ # JPEG compression
+ self.add_jpeg = add_jpeg
+
+ # Transformations
+ self.transform = transform
+
+ def __len__(self):
+ return len(self.image_idx)
+
+ def __getitem__(self, idx):
+ if idx < self.num_real:
+ image, _ = self.real[idx]
+ else:
+ image = Image.open(self.fake[idx - self.num_real]).convert("RGB")
+
+ # JPEG compression
+ if self.add_jpeg:
+ image = png_to_jpeg(image, quality=95)
+
+ # Transformations
+ if self.transform is not None:
+ image = self.transform(image)
+ label = self.labels[idx]
+
+ return image, label
+
+
+class AIGCDetectTestDataset(Dataset):
+ def __init__(self, dataset, model, transform=None):
+ # Root path of the AIGCDetectionBenchMark dataset
+ root_path = "data/test/AIGCDetectionBenchMark"
+
+ # Load images
+ image_dir = os.path.join(root_path, dataset, model)
+ real_list = get_list(image_dir, must_contain='0_real')
+ fake_list = get_list(image_dir, must_contain='1_fake')
+
+ # Setting the labels for the dataset
+ self.labels_dict = {}
+ for i in real_list:
+ self.labels_dict[i] = 0
+ for i in fake_list:
+ self.labels_dict[i] = 1
+
+ # Construct the entire dataset
+ self.total_list = real_list + fake_list
+ np.random.shuffle(self.total_list)
+
+ # Transformations
+ self.transform = transform
+
+ def __len__(self):
+ return len(self.total_list)
+
+ def __getitem__(self, idx):
+ img_path = self.total_list[idx]
+ label = self.labels_dict[img_path]
+ image = Image.open(img_path).convert("RGB")
+
+ # Transformations
+ if self.transform is not None:
+ image = self.transform(image)
+
+ return image, label
diff --git a/clean/image/cospy/detectors/progan/__init__.py b/clean/image/cospy/detectors/progan/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c9411800ca98dd45775d08a0160c377ab08421ed
--- /dev/null
+++ b/clean/image/cospy/detectors/progan/__init__.py
@@ -0,0 +1,7 @@
+from .base import BaseDetector
+from .artifact import ArtifactDetector
+from .semantic import SemanticDetector
+from .fusion import CoSpyFusionDetector
+from .end2end import End2EndDetector
+
+__all__ = ["BaseDetector", "ArtifactDetector", "SemanticDetector", "CoSpyFusionDetector", "End2EndDetector"]
diff --git a/clean/image/cospy/detectors/progan/artifact.py b/clean/image/cospy/detectors/progan/artifact.py
new file mode 100644
index 0000000000000000000000000000000000000000..e3382b3248280b19d3cb18e7eaf8434846e53b4a
--- /dev/null
+++ b/clean/image/cospy/detectors/progan/artifact.py
@@ -0,0 +1,145 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from .base import BaseDetector
+
+
+class ArtifactDetector(BaseDetector):
+ def __init__(self, dim_artifact=512, num_classes=1):
+ super(ArtifactDetector, self).__init__()
+
+ # Load the artifact encoder based on NPR
+ self.artifact_encoder = ResNet(Bottleneck, [3, 4, 6, 3])
+
+ # Classifier
+ self.fc = torch.nn.Linear(dim_artifact, num_classes)
+
+ # Build transforms
+ self._build_transforms()
+
+ def forward(self, x, return_feat=False):
+ feat = self.artifact_encoder(x)
+ out = self.fc(feat)
+ if return_feat:
+ return feat, out
+ return out
+
+
+# Define the artifact encoder (based on NPR)
+def conv1x1(in_planes, out_planes, stride=1):
+ """1x1 convolution"""
+ return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+ """3x3 convolution with padding"""
+ return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+ padding=1, bias=False)
+
+
+class Bottleneck(nn.Module):
+ expansion = 4
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None):
+ super(Bottleneck, self).__init__()
+ self.conv1 = conv1x1(inplanes, planes)
+ self.bn1 = nn.BatchNorm2d(planes)
+ self.conv2 = conv3x3(planes, planes, stride)
+ self.bn2 = nn.BatchNorm2d(planes)
+ self.conv3 = conv1x1(planes, planes * self.expansion)
+ self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+ self.relu = nn.ReLU(inplace=True)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x):
+ identity = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+ out = self.relu(out)
+
+ out = self.conv3(out)
+ out = self.bn3(out)
+
+ if self.downsample is not None:
+ identity = self.downsample(x)
+
+ out += identity
+ out = self.relu(out)
+
+ return out
+
+
+class ResNet(nn.Module):
+
+ def __init__(self, block, layers, num_classes=1):
+ super(ResNet, self).__init__()
+
+ self.unfoldSize = 2
+ self.unfoldIndex = 0
+ assert self.unfoldSize > 1
+ assert -1 < self.unfoldIndex and self.unfoldIndex < self.unfoldSize*self.unfoldSize
+ self.inplanes = 64
+ self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(64)
+ self.relu = nn.ReLU(inplace=True)
+ self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+ self.layer1 = self._make_layer(block, 64 , layers[0])
+ self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+ self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+ self.fc1 = nn.Linear(512, num_classes)
+
+ for m in self.modules():
+ if isinstance(m, nn.Conv2d):
+ nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+ elif isinstance(m, nn.BatchNorm2d):
+ nn.init.constant_(m.weight, 1)
+ nn.init.constant_(m.bias, 0)
+
+ def _make_layer(self, block, planes, blocks, stride=1):
+ downsample = None
+ if stride != 1 or self.inplanes != planes * block.expansion:
+ downsample = nn.Sequential(
+ conv1x1(self.inplanes, planes * block.expansion, stride),
+ nn.BatchNorm2d(planes * block.expansion),
+ )
+
+ layers = []
+ layers.append(block(self.inplanes, planes, stride, downsample))
+ self.inplanes = planes * block.expansion
+ for _ in range(1, blocks):
+ layers.append(block(self.inplanes, planes))
+
+ return nn.Sequential(*layers)
+
+ def interpolate(self, img, factor):
+ return F.interpolate(
+ F.interpolate(img,
+ scale_factor=factor,
+ mode='nearest',
+ recompute_scale_factor=True),
+ scale_factor=1 / factor,
+ mode='nearest',
+ recompute_scale_factor=True)
+
+ def forward(self, x):
+ artifact = x - self.interpolate(x, 0.5)
+
+ x = self.conv1(artifact * 2.0 / 3.0)
+ x = self.bn1(x)
+ x = self.relu(x)
+ x = self.maxpool(x)
+
+ x = self.layer1(x)
+ x = self.layer2(x)
+
+ x = self.avgpool(x)
+ x = x.view(x.size(0), -1)
+
+ return x
diff --git a/clean/image/cospy/detectors/progan/base.py b/clean/image/cospy/detectors/progan/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..03b6d6cf90375b699cdcd25bbce96f9d0a591233
--- /dev/null
+++ b/clean/image/cospy/detectors/progan/base.py
@@ -0,0 +1,81 @@
+import torch
+from abc import ABC, abstractmethod
+from torchvision import transforms
+from utils import data_augment
+
+
+class BaseDetector(torch.nn.Module, ABC):
+ """Abstract base class for all detectors."""
+
+ def __init__(self):
+ super(BaseDetector, self).__init__()
+
+ # Default normalization (ImageNet)
+ self.mean = [0.485, 0.456, 0.406]
+ self.std = [0.229, 0.224, 0.225]
+
+ # Default resolution
+ self.loadSize = 256
+ self.cropSize = 224
+
+ # Default data augmentation config
+ self.blur_prob = 0.0
+ self.blur_sig = [0.0, 3.0]
+ self.jpg_prob = 0.0
+ self.jpg_method = ['cv2', 'pil']
+ self.jpg_qual = list(range(30, 101))
+
+ @property
+ def aug_config(self):
+ return {
+ "blur_prob": self.blur_prob,
+ "blur_sig": self.blur_sig,
+ "jpg_prob": self.jpg_prob,
+ "jpg_method": self.jpg_method,
+ "jpg_qual": self.jpg_qual,
+ }
+
+ def _build_transforms(self):
+ """Build train and test transforms based on current config."""
+ crop_func = transforms.RandomCrop(self.cropSize)
+ flip_func = transforms.RandomHorizontalFlip()
+ rz_func = transforms.Resize(self.loadSize)
+ aug_func = transforms.Lambda(lambda x: data_augment(x, self.aug_config))
+
+ self.train_transform = transforms.Compose([
+ rz_func,
+ aug_func,
+ crop_func,
+ flip_func,
+ transforms.ToTensor(),
+ transforms.Normalize(mean=self.mean, std=self.std),
+ ])
+
+ self.test_transform = transforms.Compose([
+ rz_func,
+ crop_func,
+ transforms.ToTensor(),
+ transforms.Normalize(mean=self.mean, std=self.std),
+ ])
+
+ @abstractmethod
+ def forward(self, x, return_feat=False):
+ """Forward pass of the detector."""
+ pass
+
+ def predict(self, inputs):
+ """Prediction function."""
+ inputs = inputs.to(next(self.parameters()).device)
+ outputs = self.forward(inputs)
+ prediction = outputs.sigmoid().flatten().tolist()
+ return prediction
+
+ def save_weights(self, weights_path):
+ """Save model weights to a file."""
+ save_params = {k: v.cpu() for k, v in self.state_dict().items()}
+ torch.save(save_params, weights_path)
+
+ def load_weights(self, weights_path):
+ """Load model weights from a file."""
+ weights = torch.load(weights_path, map_location='cpu')
+ self.load_state_dict(weights, strict=False)
diff --git a/clean/image/cospy/detectors/progan/end2end.py b/clean/image/cospy/detectors/progan/end2end.py
new file mode 100644
index 0000000000000000000000000000000000000000..3aacd4ae438778afb36e495969085687a5b40849
--- /dev/null
+++ b/clean/image/cospy/detectors/progan/end2end.py
@@ -0,0 +1,101 @@
+import torch
+import random
+from torchvision import transforms
+from utils import data_augment, weights2cpu
+from .base import BaseDetector
+from .semantic import SemanticDetector
+from .artifact import ArtifactDetector
+
+
+class End2EndDetector(BaseDetector):
+ """CO-SPY End-to-End Detector."""
+
+ def __init__(self, num_classes=1):
+ super(End2EndDetector, self).__init__()
+
+ # Load the semantic detector
+ self.sem = SemanticDetector()
+ self.sem_dim = self.sem.fc.in_features
+
+ # Load the artifact detector
+ self.art = ArtifactDetector()
+ self.art_dim = self.art.fc.in_features
+
+ # Classifier
+ self.fc = torch.nn.Linear(self.sem_dim + self.art_dim, num_classes)
+
+ # Transformations inside the forward function
+ # Including the normalization and resizing (only for the artifact detector)
+ self.sem_transform = transforms.Compose([
+ transforms.Normalize(self.sem.mean, self.sem.std)
+ ])
+ self.art_transform = transforms.Compose([
+ transforms.Normalize(self.art.mean, self.art.std)
+ ])
+
+ # Build transforms (no normalization, done in forward)
+ self._build_transforms()
+
+ def _build_transforms(self):
+ """Build transforms without normalization (normalization done per-branch in forward)."""
+
+ crop_func = transforms.RandomCrop(self.cropSize)
+ flip_func = transforms.RandomHorizontalFlip()
+ rz_func = transforms.Resize(self.loadSize)
+ aug_func = transforms.Lambda(lambda x: data_augment(x, self.aug_config))
+
+ self.train_transform = transforms.Compose([
+ flip_func,
+ aug_func,
+ rz_func,
+ crop_func,
+ transforms.ToTensor(),
+ ])
+
+ self.test_transform = transforms.Compose([
+ rz_func,
+ crop_func,
+ transforms.ToTensor(),
+ ])
+
+ def forward(self, x, return_feat=False, dropout_rate=0.3):
+ x_sem = self.sem_transform(x)
+ x_art = self.art_transform(x)
+
+ # Forward pass
+ sem_feat, sem_coeff = self.sem(x_sem, return_feat=True)
+ art_feat, art_coeff = self.art(x_art, return_feat=True)
+
+ # Dropout during training
+ if self.training:
+ if random.random() < dropout_rate:
+ # Randomly select a feature to drop
+ idx_drop = random.randint(0, 1)
+ if idx_drop == 0:
+ sem_coeff = torch.zeros_like(sem_coeff)
+ else:
+ art_coeff = torch.zeros_like(art_coeff)
+
+ # Concatenate the features
+ feat = torch.cat([sem_coeff * sem_feat, art_coeff * art_feat], dim=1)
+ out = self.fc(feat)
+
+ if return_feat:
+ return feat, out
+ return out
+
+ def save_weights(self, weights_path):
+ save_params = {
+ "sem_fc": weights2cpu(self.sem.fc.state_dict()),
+ "art_fc": weights2cpu(self.art.fc.state_dict()),
+ "art_encoder": weights2cpu(self.art.artifact_encoder.state_dict()),
+ "classifier": weights2cpu(self.fc.state_dict()),
+ }
+ torch.save(save_params, weights_path)
+
+ def load_weights(self, weights_path):
+ weights = torch.load(weights_path, map_location='cpu')
+ self.sem.fc.load_state_dict(weights["sem_fc"])
+ self.art.fc.load_state_dict(weights["art_fc"])
+ self.art.artifact_encoder.load_state_dict(weights["art_encoder"])
+ self.fc.load_state_dict(weights["classifier"])
diff --git a/clean/image/cospy/detectors/progan/fusion.py b/clean/image/cospy/detectors/progan/fusion.py
new file mode 100644
index 0000000000000000000000000000000000000000..c33ed927a310da4b382552af7c2f046ba46b11ab
--- /dev/null
+++ b/clean/image/cospy/detectors/progan/fusion.py
@@ -0,0 +1,86 @@
+import torch
+from torchvision import transforms
+
+from .base import BaseDetector
+from .semantic import SemanticDetector
+from .artifact import ArtifactDetector
+
+
+class CoSpyFusionDetector(BaseDetector):
+ """CO-SPY Fusion Detector (calibrate and fuse semantic and artifact detectors)."""
+
+ def __init__(self, semantic_weights_path, artifact_weights_path, num_classes=1):
+ super(CoSpyFusionDetector, self).__init__()
+
+ # Load the semantic detector
+ self.sem = SemanticDetector()
+ self.sem.load_weights(semantic_weights_path)
+
+ # Load the artifact detector
+ self.art = ArtifactDetector()
+ self.art.load_weights(artifact_weights_path)
+
+ # Freeze the two pre-trained models
+ for param in self.sem.parameters():
+ param.requires_grad = False
+ for param in self.art.parameters():
+ param.requires_grad = False
+
+ # Classifier
+ self.fc = torch.nn.Linear(2, num_classes)
+
+ # Transformations inside the forward function
+ self.sem_transform = transforms.Compose([
+ transforms.Normalize(self.sem.mean, self.sem.std)
+ ])
+ self.art_transform = transforms.Compose([
+ transforms.Normalize(self.art.mean, self.art.std)
+ ])
+
+ # Build transforms (no normalization, done in forward)
+ self._build_fusion_transforms()
+
+ def _build_fusion_transforms(self):
+ """Build transforms without normalization (normalization done per-branch in forward)."""
+ from utils import data_augment
+
+ crop_func = transforms.RandomCrop(self.cropSize)
+ flip_func = transforms.RandomHorizontalFlip()
+ rz_func = transforms.Resize(self.loadSize)
+ aug_func = transforms.Lambda(lambda x: data_augment(x, self.aug_config))
+
+ self.train_transform = transforms.Compose([
+ flip_func,
+ aug_func,
+ rz_func,
+ crop_func,
+ transforms.ToTensor(),
+ ])
+
+ self.test_transform = transforms.Compose([
+ rz_func,
+ crop_func,
+ transforms.ToTensor(),
+ ])
+
+ def forward(self, x, return_feat=False):
+ x_sem = self.sem_transform(x)
+ x_art = self.art_transform(x)
+ pred_sem = self.sem(x_sem)
+ pred_art = self.art(x_art)
+ feat = torch.cat([pred_sem, pred_art], dim=1)
+ out = self.fc(feat)
+ if return_feat:
+ return feat, out
+ return out
+
+ def save_weights(self, weights_path):
+ # Only save the fc layer (sub-detectors are frozen)
+ save_params = {"fc.weight": self.fc.weight.cpu(), "fc.bias": self.fc.bias.cpu()}
+ torch.save(save_params, weights_path)
+
+ def load_weights(self, weights_path):
+ # Load only the fc layer
+ weights = torch.load(weights_path, map_location='cpu')
+ self.fc.weight.data = weights["fc.weight"]
+ self.fc.bias.data = weights["fc.bias"]
diff --git a/clean/image/cospy/detectors/progan/semantic.py b/clean/image/cospy/detectors/progan/semantic.py
new file mode 100644
index 0000000000000000000000000000000000000000..16dc26ec3f2d36e4e4ccdaf142bb678b6c4d368a
--- /dev/null
+++ b/clean/image/cospy/detectors/progan/semantic.py
@@ -0,0 +1,48 @@
+import torch
+from transformers import CLIPModel
+
+from .base import BaseDetector
+
+
+class SemanticDetector(BaseDetector):
+ def __init__(self, dim_clip=768, num_classes=1):
+ super(SemanticDetector, self).__init__()
+
+ # Get the pre-trained CLIP
+ model_name = "openai/clip-vit-large-patch14"
+ self.clip = CLIPModel.from_pretrained(model_name)
+
+ # Freeze the CLIP visual encoder
+ self.clip.requires_grad_(False)
+
+ # Classifier
+ self.fc = torch.nn.Linear(dim_clip, num_classes)
+
+ # CLIP normalization
+ self.mean = [0.48145466, 0.45782750, 0.40821073]
+ self.std = [0.26862954, 0.26130258, 0.27577711]
+
+ # Data augmentation (override defaults)
+ self.blur_prob = 0.5
+ self.jpg_prob = 0.5
+
+ # Build transforms
+ self._build_transforms()
+
+ def forward(self, x, return_feat=False):
+ feat = self.clip.get_image_features(x)
+ out = self.fc(feat)
+ if return_feat:
+ return feat, out
+ return out
+
+ def save_weights(self, weights_path):
+ # Only save the fc layer (CLIP is frozen)
+ save_params = {"fc.weight": self.fc.weight.cpu(), "fc.bias": self.fc.bias.cpu()}
+ torch.save(save_params, weights_path)
+
+ def load_weights(self, weights_path):
+ # Load only the fc layer
+ weights = torch.load(weights_path, map_location='cpu')
+ self.fc.weight.data = weights["fc.weight"]
+ self.fc.bias.data = weights["fc.bias"]
diff --git a/clean/image/cospy/detectors/sd-v1_4/__init__.py b/clean/image/cospy/detectors/sd-v1_4/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c9411800ca98dd45775d08a0160c377ab08421ed
--- /dev/null
+++ b/clean/image/cospy/detectors/sd-v1_4/__init__.py
@@ -0,0 +1,7 @@
+from .base import BaseDetector
+from .artifact import ArtifactDetector
+from .semantic import SemanticDetector
+from .fusion import CoSpyFusionDetector
+from .end2end import End2EndDetector
+
+__all__ = ["BaseDetector", "ArtifactDetector", "SemanticDetector", "CoSpyFusionDetector", "End2EndDetector"]
diff --git a/clean/image/cospy/detectors/sd-v1_4/artifact.py b/clean/image/cospy/detectors/sd-v1_4/artifact.py
new file mode 100644
index 0000000000000000000000000000000000000000..a808c1900f7696afde77a022d1629b3648e64290
--- /dev/null
+++ b/clean/image/cospy/detectors/sd-v1_4/artifact.py
@@ -0,0 +1,200 @@
+import torch
+import torch.nn as nn
+from diffusers import StableDiffusionPipeline
+from torchvision import transforms
+from utils import data_augment
+from .base import BaseDetector
+
+
+class ArtifactDetector(BaseDetector):
+ def __init__(self, dim_artifact=512, num_classes=1):
+ super(ArtifactDetector, self).__init__()
+
+ # Load the pre-trained VAE
+ model_id = "CompVis/stable-diffusion-v1-4"
+ vae = StableDiffusionPipeline.from_pretrained(model_id).vae
+ vae.requires_grad_(False)
+ self.artifact_encoder = VAEReconEncoder(vae)
+
+ # Classifier
+ self.fc = torch.nn.Linear(dim_artifact, num_classes)
+
+ # Build transforms
+ self._build_transforms()
+
+ def _build_transforms(self):
+ # Normalization
+ self.mean = [0.0, 0.0, 0.0]
+ self.std = [1.0, 1.0, 1.0]
+
+ # Resolution
+ self.loadSize = 256
+ self.cropSize = 224
+
+ # Data augmentation
+ self.blur_prob = 0.0
+ self.blur_sig = [0.0, 3.0]
+ self.jpg_prob = 0.5
+ self.jpg_method = ['cv2', 'pil']
+ self.jpg_qual = list(range(70, 96))
+
+ # Define the augmentation configuration
+ self.aug_config = {
+ "blur_prob": self.blur_prob,
+ "blur_sig": self.blur_sig,
+ "jpg_prob": self.jpg_prob,
+ "jpg_method": self.jpg_method,
+ "jpg_qual": self.jpg_qual,
+ }
+
+ # Pre-processing
+ crop_func = transforms.RandomCrop(self.cropSize)
+ flip_func = transforms.RandomHorizontalFlip()
+ rz_func = transforms.Resize(self.loadSize)
+ aug_func = transforms.Lambda(lambda x: data_augment(x, self.aug_config))
+
+ self.train_transform = transforms.Compose([
+ aug_func,
+ rz_func,
+ crop_func,
+ flip_func,
+ transforms.ToTensor(),
+ transforms.Normalize(mean=self.mean, std=self.std),
+ ])
+
+ self.test_transform = transforms.Compose([
+ rz_func,
+ crop_func,
+ transforms.ToTensor(),
+ transforms.Normalize(mean=self.mean, std=self.std),
+ ])
+
+ def forward(self, x, return_feat=False):
+ feat = self.artifact_encoder(x)
+ out = self.fc(feat)
+ if return_feat:
+ return feat, out
+ return out
+
+
+# Helper functions for ResNet
+def conv3x3(in_planes, out_planes, stride=1):
+ """3x3 convolution with padding"""
+ return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+ """1x1 convolution"""
+ return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class Bottleneck(nn.Module):
+ expansion = 4
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None):
+ super(Bottleneck, self).__init__()
+ self.conv1 = conv1x1(inplanes, planes)
+ self.bn1 = nn.BatchNorm2d(planes)
+ self.conv2 = conv3x3(planes, planes, stride)
+ self.bn2 = nn.BatchNorm2d(planes)
+ self.conv3 = conv1x1(planes, planes * self.expansion)
+ self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+ self.relu = nn.ReLU(inplace=True)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x):
+ identity = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+ out = self.relu(out)
+
+ out = self.conv3(out)
+ out = self.bn3(out)
+
+ if self.downsample is not None:
+ identity = self.downsample(x)
+
+ out += identity
+ out = self.relu(out)
+
+ return out
+
+
+class VAEReconEncoder(nn.Module):
+ def __init__(self, vae, block=Bottleneck):
+ super(VAEReconEncoder, self).__init__()
+
+ # Define the ResNet model
+ self.inplanes = 64
+ self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(64)
+ self.relu = nn.ReLU(inplace=True)
+ self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+ # ResNet-50 is [3, 4, 6, 3]
+ self.layer1 = self._make_layer(block, 64, 3)
+ self.layer2 = self._make_layer(block, 128, 4, stride=2)
+
+ self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+
+ # Kaiming initialization
+ for m in self.modules():
+ if isinstance(m, nn.Conv2d):
+ nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+ elif isinstance(m, nn.BatchNorm2d):
+ nn.init.constant_(m.weight, 1)
+ nn.init.constant_(m.bias, 0)
+
+ # Load the VAE model
+ self.vae = vae
+
+ def _make_layer(self, block, planes, blocks, stride=1):
+ downsample = None
+ if stride != 1 or self.inplanes != planes * block.expansion:
+ downsample = nn.Sequential(
+ conv1x1(self.inplanes, planes * block.expansion, stride),
+ nn.BatchNorm2d(planes * block.expansion),
+ )
+
+ layers = []
+ layers.append(block(self.inplanes, planes, stride, downsample))
+ self.inplanes = planes * block.expansion
+ for _ in range(1, blocks):
+ layers.append(block(self.inplanes, planes))
+
+ return nn.Sequential(*layers)
+
+ def reconstruct(self, x):
+ with torch.no_grad():
+ latent = self.vae.encode(x).latent_dist.mean
+ decoded = self.vae.decode(latent).sample
+ return decoded
+
+ def forward(self, x):
+ # Reconstruct
+ x_recon = self.reconstruct(x)
+ # Compute the artifacts
+ x = x - x_recon
+
+ # Scale the artifacts
+ x = x / 7. * 100.
+
+ # Forward pass
+ x = self.conv1(x)
+ x = self.bn1(x)
+ x = self.relu(x)
+ x = self.maxpool(x)
+
+ x = self.layer1(x)
+ x = self.layer2(x)
+
+ x = self.avgpool(x)
+ x = x.view(x.size(0), -1)
+
+ return x
diff --git a/clean/image/cospy/detectors/sd-v1_4/base.py b/clean/image/cospy/detectors/sd-v1_4/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..79313e4e60b262eb616de38336002b266842327b
--- /dev/null
+++ b/clean/image/cospy/detectors/sd-v1_4/base.py
@@ -0,0 +1,36 @@
+import torch
+from abc import ABC, abstractmethod
+
+
+class BaseDetector(torch.nn.Module, ABC):
+ """Abstract base class for all detectors."""
+
+ def __init__(self):
+ super(BaseDetector, self).__init__()
+
+ @abstractmethod
+ def _build_transforms(self):
+ """Build train and test transforms."""
+ pass
+
+ @abstractmethod
+ def forward(self, x, return_feat=False):
+ """Forward pass of the detector."""
+ pass
+
+ def predict(self, inputs):
+ """Prediction function."""
+ inputs = inputs.to(next(self.parameters()).device)
+ outputs = self.forward(inputs)
+ prediction = outputs.sigmoid().flatten().tolist()
+ return prediction
+
+ def save_weights(self, weights_path):
+ """Save model weights to a file."""
+ save_params = {k: v.cpu() for k, v in self.state_dict().items()}
+ torch.save(save_params, weights_path)
+
+ def load_weights(self, weights_path):
+ """Load model weights from a file."""
+ weights = torch.load(weights_path, map_location='cpu')
+ self.load_state_dict(weights, strict=False)
diff --git a/clean/image/cospy/detectors/sd-v1_4/end2end.py b/clean/image/cospy/detectors/sd-v1_4/end2end.py
new file mode 100644
index 0000000000000000000000000000000000000000..c931c0cd4389fa916e80e210dd89e15dea3b14bc
--- /dev/null
+++ b/clean/image/cospy/detectors/sd-v1_4/end2end.py
@@ -0,0 +1,122 @@
+import torch
+import random
+from torchvision import transforms
+from utils import data_augment, weights2cpu
+from .base import BaseDetector
+from .semantic import SemanticDetector
+from .artifact import ArtifactDetector
+
+
+class End2EndDetector(BaseDetector):
+ """CO-SPY End-to-End Detector."""
+
+ def __init__(self, num_classes=1):
+ super(End2EndDetector, self).__init__()
+
+ # Load the semantic detector
+ self.sem = SemanticDetector()
+ self.sem_dim = self.sem.fc.in_features
+
+ # Load the artifact detector
+ self.art = ArtifactDetector()
+ self.art_dim = self.art.fc.in_features
+
+ # Classifier
+ self.fc = torch.nn.Linear(self.sem_dim + self.art_dim, num_classes)
+
+ # Transformations inside the forward function
+ # Including the normalization and resizing (only for the artifact detector)
+ self.sem_transform = transforms.Compose([
+ transforms.Normalize(self.sem.mean, self.sem.std)
+ ])
+ self.art_transform = transforms.Compose([
+ transforms.Resize(self.art.cropSize, antialias=False),
+ transforms.Normalize(self.art.mean, self.art.std)
+ ])
+
+ # Build transforms (no normalization, done in forward)
+ self._build_transforms()
+
+ def _build_transforms(self):
+ """Build transforms without normalization (normalization done per-branch in forward)."""
+
+ # Resolution
+ self.loadSize = 384
+ self.cropSize = 384
+
+ # Data augmentation
+ self.blur_prob = 0.0
+ self.blur_sig = [0.0, 3.0]
+ self.jpg_prob = 0.5
+ self.jpg_method = ['cv2', 'pil']
+ self.jpg_qual = list(range(70, 96))
+
+ # Define the augmentation configuration
+ self.aug_config = {
+ "blur_prob": self.blur_prob,
+ "blur_sig": self.blur_sig,
+ "jpg_prob": self.jpg_prob,
+ "jpg_method": self.jpg_method,
+ "jpg_qual": self.jpg_qual,
+ }
+
+ crop_func = transforms.RandomCrop(self.cropSize)
+ flip_func = transforms.RandomHorizontalFlip()
+ rz_func = transforms.Resize(self.loadSize)
+ aug_func = transforms.Lambda(lambda x: data_augment(x, self.aug_config))
+
+ self.train_transform = transforms.Compose([
+ flip_func,
+ aug_func,
+ rz_func,
+ crop_func,
+ transforms.ToTensor(),
+ ])
+
+ self.test_transform = transforms.Compose([
+ rz_func,
+ crop_func,
+ transforms.ToTensor(),
+ ])
+
+ def forward(self, x, return_feat=False, dropout_rate=0.3):
+ x_sem = self.sem_transform(x)
+ x_art = self.art_transform(x)
+
+ # Forward pass
+ sem_feat, sem_coeff = self.sem(x_sem, return_feat=True)
+ art_feat, art_coeff = self.art(x_art, return_feat=True)
+
+ # Dropout during training
+ if self.training:
+ if random.random() < dropout_rate:
+ # Randomly select a feature to drop
+ idx_drop = random.randint(0, 1)
+ if idx_drop == 0:
+ sem_coeff = torch.zeros_like(sem_coeff)
+ else:
+ art_coeff = torch.zeros_like(art_coeff)
+
+ # Concatenate the features
+ feat = torch.cat([sem_coeff * sem_feat, art_coeff * art_feat], dim=1)
+ out = self.fc(feat)
+
+ if return_feat:
+ return feat, out
+ return out
+
+ def save_weights(self, weights_path):
+ save_params = {
+ "sem_fc": weights2cpu(self.sem.fc.state_dict()),
+ "art_fc": weights2cpu(self.art.fc.state_dict()),
+ "art_encoder": weights2cpu(self.art.artifact_encoder.state_dict()),
+ "classifier": weights2cpu(self.fc.state_dict()),
+ }
+ torch.save(save_params, weights_path)
+
+ def load_weights(self, weights_path):
+ weights = torch.load(weights_path, map_location='cpu')
+ self.sem.fc.load_state_dict(weights["sem_fc"])
+ self.art.fc.load_state_dict(weights["art_fc"])
+ self.art.artifact_encoder.load_state_dict(weights["art_encoder"])
+ self.fc.load_state_dict(weights["classifier"])
diff --git a/clean/image/cospy/detectors/sd-v1_4/fusion.py b/clean/image/cospy/detectors/sd-v1_4/fusion.py
new file mode 100644
index 0000000000000000000000000000000000000000..cd1a84ff383eb6ef47289cf4bb759b21e8e35669
--- /dev/null
+++ b/clean/image/cospy/detectors/sd-v1_4/fusion.py
@@ -0,0 +1,107 @@
+import torch
+from torchvision import transforms
+from utils import data_augment
+from .base import BaseDetector
+from .semantic import SemanticDetector
+from .artifact import ArtifactDetector
+
+
+class CoSpyFusionDetector(BaseDetector):
+ """CO-SPY Fusion Detector (calibrate and fuse semantic and artifact detectors)."""
+
+ def __init__(self, semantic_weights_path, artifact_weights_path, num_classes=1):
+ super(CoSpyFusionDetector, self).__init__()
+
+ # Load the semantic detector
+ self.sem = SemanticDetector()
+ self.sem.load_weights(semantic_weights_path)
+
+ # Load the artifact detector
+ self.art = ArtifactDetector()
+ self.art.load_weights(artifact_weights_path)
+
+ # Freeze the two pre-trained models
+ for param in self.sem.parameters():
+ param.requires_grad = False
+ for param in self.art.parameters():
+ param.requires_grad = False
+
+ # Classifier
+ self.fc = torch.nn.Linear(2, num_classes)
+
+ # Transformations inside the forward function
+ # Including the normalization and resizing (only for the artifact detector)
+ self.sem_transform = transforms.Compose([
+ transforms.Normalize(self.sem.mean, self.sem.std)
+ ])
+ self.art_transform = transforms.Compose([
+ transforms.Resize(self.art.cropSize, antialias=False),
+ transforms.Normalize(self.art.mean, self.art.std)
+ ])
+
+ # Build transforms (no normalization, done in forward)
+ self._build_transforms()
+
+ def _build_transforms(self):
+ """Build transforms without normalization (normalization done per-branch in forward)."""
+
+ # Resolution
+ self.loadSize = 384
+ self.cropSize = 384
+
+ # Data augmentation
+ self.blur_prob = 0.0
+ self.blur_sig = [0.0, 3.0]
+ self.jpg_prob = 0.5
+ self.jpg_method = ['cv2', 'pil']
+ self.jpg_qual = list(range(70, 96))
+
+ # Define the augmentation configuration
+ self.aug_config = {
+ "blur_prob": self.blur_prob,
+ "blur_sig": self.blur_sig,
+ "jpg_prob": self.jpg_prob,
+ "jpg_method": self.jpg_method,
+ "jpg_qual": self.jpg_qual,
+ }
+
+ crop_func = transforms.RandomCrop(self.cropSize)
+ flip_func = transforms.RandomHorizontalFlip()
+ rz_func = transforms.Resize(self.loadSize)
+ aug_func = transforms.Lambda(lambda x: data_augment(x, self.aug_config))
+
+ self.train_transform = transforms.Compose([
+ flip_func,
+ aug_func,
+ rz_func,
+ crop_func,
+ transforms.ToTensor(),
+ ])
+
+ self.test_transform = transforms.Compose([
+ rz_func,
+ crop_func,
+ transforms.ToTensor(),
+ ])
+
+ def forward(self, x, return_feat=False):
+ x_sem = self.sem_transform(x)
+ x_art = self.art_transform(x)
+ pred_sem = self.sem(x_sem)
+ pred_art = self.art(x_art)
+ feat = torch.cat([pred_sem, pred_art], dim=1)
+ out = self.fc(feat)
+ if return_feat:
+ return feat, out
+ return out
+
+ def save_weights(self, weights_path):
+ # Only save the fc layer (sub-detectors are frozen)
+ save_params = {"fc.weight": self.fc.weight.cpu(), "fc.bias": self.fc.bias.cpu()}
+ torch.save(save_params, weights_path)
+
+ def load_weights(self, weights_path):
+ # Load only the fc layer
+ weights = torch.load(weights_path, map_location='cpu')
+ self.fc.weight.data = weights["fc.weight"]
+ self.fc.bias.data = weights["fc.bias"]
diff --git a/clean/image/cospy/detectors/sd-v1_4/semantic.py b/clean/image/cospy/detectors/sd-v1_4/semantic.py
new file mode 100644
index 0000000000000000000000000000000000000000..ff9b08138da36a11512a15fb0362e2459ecb7e45
--- /dev/null
+++ b/clean/image/cospy/detectors/sd-v1_4/semantic.py
@@ -0,0 +1,89 @@
+import torch
+import open_clip
+from torchvision import transforms
+from utils import data_augment
+from .base import BaseDetector
+
+
+class SemanticDetector(BaseDetector):
+ def __init__(self, dim_clip=1152, num_classes=1):
+ super(SemanticDetector, self).__init__()
+
+ # Get the pre-trained CLIP (SigLIP)
+ model_name = "ViT-SO400M-14-SigLIP-384"
+ version = "webli"
+ self.clip, _, _ = open_clip.create_model_and_transforms(model_name, pretrained=version)
+
+ # Freeze the CLIP visual encoder
+ self.clip.requires_grad_(False)
+
+ # Classifier
+ self.fc = torch.nn.Linear(dim_clip, num_classes)
+
+ # Build transforms
+ self._build_transforms()
+
+ def _build_transforms(self):
+ # Normalization
+ self.mean = [0.5, 0.5, 0.5]
+ self.std = [0.5, 0.5, 0.5]
+
+ # Resolution
+ self.loadSize = 384
+ self.cropSize = 384
+
+ # Data augmentation
+ self.blur_prob = 0.5
+ self.blur_sig = [0.0, 3.0]
+ self.jpg_prob = 0.5
+ self.jpg_method = ['cv2', 'pil']
+ self.jpg_qual = list(range(30, 101))
+
+ # Define the augmentation configuration
+ self.aug_config = {
+ "blur_prob": self.blur_prob,
+ "blur_sig": self.blur_sig,
+ "jpg_prob": self.jpg_prob,
+ "jpg_method": self.jpg_method,
+ "jpg_qual": self.jpg_qual,
+ }
+
+ # Pre-processing
+ crop_func = transforms.RandomCrop(self.cropSize)
+ flip_func = transforms.RandomHorizontalFlip()
+ rz_func = transforms.Resize(self.loadSize)
+ aug_func = transforms.Lambda(lambda x: data_augment(x, self.aug_config))
+
+ self.train_transform = transforms.Compose([
+ rz_func,
+ aug_func,
+ crop_func,
+ flip_func,
+ transforms.ToTensor(),
+ transforms.Normalize(mean=self.mean, std=self.std),
+ ])
+
+ self.test_transform = transforms.Compose([
+ rz_func,
+ crop_func,
+ transforms.ToTensor(),
+ transforms.Normalize(mean=self.mean, std=self.std),
+ ])
+
+ def forward(self, x, return_feat=False):
+ feat = self.clip.encode_image(x)
+ out = self.fc(feat)
+ if return_feat:
+ return feat, out
+ return out
+
+ def save_weights(self, weights_path):
+ # Only save the fc layer (CLIP is frozen)
+ save_params = {"fc.weight": self.fc.weight.cpu(), "fc.bias": self.fc.bias.cpu()}
+ torch.save(save_params, weights_path)
+
+ def load_weights(self, weights_path):
+ # Only load the fc layer
+ weights = torch.load(weights_path, map_location='cpu')
+ self.fc.weight.data = weights["fc.weight"]
+ self.fc.bias.data = weights["fc.bias"]
diff --git a/clean/image/cospy/environment.yml b/clean/image/cospy/environment.yml
new file mode 100644
index 0000000000000000000000000000000000000000..b9ab51459e14e99fe1f06d3b20feea84fc46638a
--- /dev/null
+++ b/clean/image/cospy/environment.yml
@@ -0,0 +1,105 @@
+name: cospy
+channels:
+ - defaults
+dependencies:
+ - _libgcc_mutex=0.1=main
+ - _openmp_mutex=5.1=1_gnu
+ - ca-certificates=2025.2.25=h06a4308_0
+ - ld_impl_linux-64=2.40=h12ee557_0
+ - libffi=3.4.4=h6a678d5_1
+ - libgcc-ng=11.2.0=h1234567_1
+ - libgomp=11.2.0=h1234567_1
+ - libstdcxx-ng=11.2.0=h1234567_1
+ - ncurses=6.4=h6a678d5_0
+ - openssl=3.0.16=h5eee18b_0
+ - pip=24.2=py38h06a4308_0
+ - python=3.8.18=h955ad1f_0
+ - readline=8.2=h5eee18b_0
+ - setuptools=75.1.0=py38h06a4308_0
+ - sqlite=3.45.3=h5eee18b_0
+ - tk=8.6.14=h39e8969_0
+ - wheel=0.44.0=py38h06a4308_0
+ - xz=5.6.4=h5eee18b_1
+ - zlib=1.2.13=h5eee18b_1
+ - pip:
+ - accelerate==1.0.1
+ - aiohappyeyeballs==2.4.4
+ - aiohttp==3.10.11
+ - aiosignal==1.3.1
+ - async-timeout==5.0.1
+ - attrs==25.3.0
+ - certifi==2025.1.31
+ - charset-normalizer==3.4.1
+ - contourpy==1.1.1
+ - cycler==0.12.1
+ - datasets==3.1.0
+ - diffusers==0.32.2
+ - dill==0.3.8
+ - filelock==3.16.1
+ - fonttools==4.56.0
+ - frozenlist==1.5.0
+ - fsspec==2024.9.0
+ - ftfy==6.2.3
+ - huggingface-hub==0.29.3
+ - idna==3.10
+ - importlib-metadata==8.5.0
+ - importlib-resources==6.4.5
+ - jinja2==3.1.6
+ - joblib==1.4.2
+ - kiwisolver==1.4.7
+ - loguru==0.7.3
+ - markupsafe==2.1.5
+ - matplotlib==3.7.5
+ - mpmath==1.3.0
+ - multidict==6.1.0
+ - multiprocess==0.70.16
+ - networkx==3.1
+ - numpy==1.24.4
+ - nvidia-cublas-cu12==12.1.3.1
+ - nvidia-cuda-cupti-cu12==12.1.105
+ - nvidia-cuda-nvrtc-cu12==12.1.105
+ - nvidia-cuda-runtime-cu12==12.1.105
+ - nvidia-cudnn-cu12==9.1.0.70
+ - nvidia-cufft-cu12==11.0.2.54
+ - nvidia-curand-cu12==10.3.2.106
+ - nvidia-cusolver-cu12==11.4.5.107
+ - nvidia-cusparse-cu12==12.1.0.106
+ - nvidia-nccl-cu12==2.20.5
+ - nvidia-nvjitlink-cu12==12.8.93
+ - nvidia-nvtx-cu12==12.1.105
+ - open-clip-torch==2.31.0
+ - opencv-python==4.11.0.86
+ - packaging==24.2
+ - pandas==2.0.3
+ - pillow==10.4.0
+ - propcache==0.2.0
+ - psutil==7.0.0
+ - pyarrow==17.0.0
+ - pycocotools==2.0.7
+ - pyparsing==3.1.4
+ - python-dateutil==2.9.0.post0
+ - pytz==2025.1
+ - pyyaml==6.0.2
+ - regex==2024.11.6
+ - requests==2.32.3
+ - safetensors==0.5.3
+ - scikit-learn==1.3.2
+ - scipy==1.10.1
+ - six==1.17.0
+ - sympy==1.13.3
+ - threadpoolctl==3.5.0
+ - timm==1.0.15
+ - tokenizers==0.20.3
+ - torch==2.4.1
+ - torchvision==0.19.1
+ - tqdm==4.67.1
+ - transformers==4.46.3
+ - triton==3.0.0
+ - typing-extensions==4.12.2
+ - tzdata==2025.1
+ - urllib3==2.2.3
+ - wcwidth==0.2.13
+ - xxhash==3.5.0
+ - yarl==1.15.2
+ - zipp==3.20.2
+prefix: /connect4/cheng535-new/anaconda3/envs/cospy
diff --git a/clean/image/cospy/evaluate.py b/clean/image/cospy/evaluate.py
new file mode 100644
index 0000000000000000000000000000000000000000..e809ad3d5e33d688d236c3b9603a544e668307ad
--- /dev/null
+++ b/clean/image/cospy/evaluate.py
@@ -0,0 +1,170 @@
+import os
+import json
+import torch
+import importlib
+import numpy as np
+from tqdm import tqdm
+from PIL import Image
+from loguru import logger
+
+from dataSets import *
+from utils import seed_torch, evaluate
+
+import warnings
+warnings.filterwarnings("ignore")
+
+
+# Load pre-trained models for evaluation
+class Detector:
+ def __init__(self,
+ device: str,
+ mode: str = "fusion",
+ train_dataset: str = "sd-v1_4",
+ pretrain: bool = False,
+ ckpt: str = "ckpt",
+ batch_size: int = 32):
+
+ # Device
+ self.device = device
+ self.mode = mode
+ self.train_dataset = train_dataset
+ self.pretrain = pretrain
+ self.ckpt = ckpt
+ self.batch_size = batch_size
+
+ # Dynamically import the detector module from either "progan" or "sd-v1_4"
+ detector_module = importlib.import_module(f"detectors.{train_dataset}")
+
+ # Get the detector and load weights based on mode
+ if pretrain:
+ # Only provide pre-trained weights for fusion mode
+ # Hardcode to fusion mode
+ self.mode = "fusion"
+ # Load the fusion detector with pre-trained weights
+ semantic_weights_path = f"pretrained/{self.train_dataset}/semantic_weights.pth"
+ artifact_weights_path = f"pretrained/{self.train_dataset}/artifact_weights.pth"
+ fusion_weights_path = f"pretrained/{self.train_dataset}/fusion_weights.pth"
+ if not os.path.exists(semantic_weights_path) or not os.path.exists(artifact_weights_path) or not os.path.exists(fusion_weights_path):
+ raise ValueError("The pre-trained weights are not complete for evaluation")
+ CoSpyFusionDetector = getattr(detector_module, "CoSpyFusionDetector")
+ self.model = CoSpyFusionDetector(
+ semantic_weights_path=semantic_weights_path,
+ artifact_weights_path=artifact_weights_path)
+ self.model.load_weights(fusion_weights_path)
+ else:
+ if self.mode == "fusion":
+ semantic_weights_path = os.path.join(self.ckpt, self.train_dataset, "semantic", "best_model.pth")
+ artifact_weights_path = os.path.join(self.ckpt, self.train_dataset, "artifact", "best_model.pth")
+ fusion_weights_path = os.path.join(self.ckpt, self.train_dataset, "fusion", "best_model.pth")
+ if not os.path.exists(semantic_weights_path) or not os.path.exists(artifact_weights_path) or not os.path.exists(fusion_weights_path):
+ raise ValueError("Semantic, Artifact or Fusion weights path does not exist for fusion mode")
+ CoSpyFusionDetector = getattr(detector_module, "CoSpyFusionDetector")
+ self.model = CoSpyFusionDetector(
+ semantic_weights_path=semantic_weights_path,
+ artifact_weights_path=artifact_weights_path)
+ self.model.load_weights(fusion_weights_path)
+ elif self.mode == "end2end":
+ End2EndDetector = getattr(detector_module, "End2EndDetector")
+ self.model = End2EndDetector()
+ end2end_weights_path = os.path.join(self.ckpt, self.train_dataset, "end2end", "best_model.pth")
+ if not os.path.exists(end2end_weights_path):
+ raise ValueError("End2End weights path does not exist for end2end mode")
+ self.model.load_weights(end2end_weights_path)
+ else:
+ raise ValueError(f"Unknown mode: {self.mode}")
+
+ # Put the model on the device and set to eval
+ self.model.to(self.device)
+ self.model.eval()
+
+ def evaluate_benchmark(self):
+ # Select the appropriate test dataset and evaluation lists
+ if self.train_dataset == "progan":
+ benchmark_name = "AIGCDetectionBenchMark"
+ TestDataset = AIGCDetectTestDataset
+ eval_dataset_list = AIGCDetectionBenchMark_DATASET_LIST
+ eval_model_list = AIGCDetectionBenchMark_MODEL_LIST
+ elif self.train_dataset == "sd-v1_4":
+ benchmark_name = "Co-Spy-Bench"
+ TestDataset = CoSpyBenchTestDataset
+ eval_dataset_list = CoSpyBench_DATASET_LIST
+ eval_model_list = CoSpyBench_MODEL_LIST
+ else:
+ raise ValueError(f"Unknown train dataset: {self.train_dataset}")
+
+ # Set the saving directory
+ if self.pretrain:
+ save_dir = os.path.join(self.ckpt, self.train_dataset, self.mode, f"pretrain_{benchmark_name}")
+ else:
+ save_dir = os.path.join(self.ckpt, self.train_dataset, self.mode, f"eval_{benchmark_name}")
+ if not os.path.exists(save_dir):
+ os.makedirs(save_dir)
+
+ # Setup logger
+ log_path = f"{save_dir}/evaluation.log"
+ if os.path.exists(log_path):
+ os.remove(log_path)
+
+ logger_id = logger.add(
+ log_path,
+ format="{time:MM-DD at HH:mm:ss} | {level} | {module}:{line} | {message}",
+ level="DEBUG",
+ )
+
+ # Save raw model prediction
+ save_output_path = os.path.join(save_dir, "output.json")
+ # Save summarized evaluation result
+ save_result_path = os.path.join(save_dir, "result.json")
+
+ # Begin the evaluation
+ result_all = {}
+ output_all = {}
+ for dataset_name in eval_dataset_list:
+ result_all[dataset_name] = {}
+ output_all[dataset_name] = {}
+ for model_name in eval_model_list:
+ test_dataset = TestDataset(dataset=dataset_name, model=model_name, transform=self.model.test_transform)
+ test_loader = torch.utils.data.DataLoader(test_dataset,
+ batch_size=self.batch_size,
+ shuffle=False,
+ num_workers=4,
+ pin_memory=True)
+
+ # Evaluate the model
+ y_pred, y_true = [], []
+ for (images, labels) in tqdm(test_loader, desc=f"Evaluating {benchmark_name} - {dataset_name} - {model_name}"):
+ y_pred.extend(self.model.predict(images))
+ y_true.extend(labels.tolist())
+
+ ap, accuracy = evaluate(y_pred, y_true)
+ logger.info(f"Evaluate on {benchmark_name} - {dataset_name} - {model_name} | Size {len(y_true)} | AP {ap*100:.2f}% | Accuracy {accuracy*100:.2f}%")
+
+ result_all[dataset_name][model_name] = {"size": len(y_true), "AP": ap, "Accuracy": accuracy}
+ output_all[dataset_name][model_name] = {"y_pred": y_pred, "y_true": y_true}
+
+ # Save the results
+ with open(save_result_path, "w") as f:
+ json.dump(result_all, f, indent=4)
+
+ with open(save_output_path, "w") as f:
+ json.dump(output_all, f, indent=4)
+
+ def scan(self):
+ # Load the image
+ image_filepath = input("Please enter the image filepath for scanning: ")
+ if not os.path.exists(image_filepath):
+ print(f"Image file not found: {image_filepath}")
+ image_filepath = input("Please enter the image filepath for scanning: ")
+
+ image = Image.open(image_filepath).convert("RGB")
+ image = self.model.test_transform(image)
+ image = image.unsqueeze(0)
+ image = image.to(self.device)
+
+ # Make the prediction
+ prediction = self.model.predict(image)[0]
+
+ if prediction > 0.5:
+ print(f"Co-Spy Prediction: {prediction:.3f} - AI-Generated")
+ else:
+ print(f"Co-Spy Prediction: {prediction:.3f} - Real")
diff --git a/clean/image/cospy/main.py b/clean/image/cospy/main.py
new file mode 100644
index 0000000000000000000000000000000000000000..21cbf2390779679d4a7ca467c5b3637b61c5793e
--- /dev/null
+++ b/clean/image/cospy/main.py
@@ -0,0 +1,117 @@
+import torch
+from train import Trainer
+from evaluate import Detector
+from utils import seed_torch
+
+
+def main(args):
+ #########################################
+ # Phase 1: Training
+ #########################################
+ if args.phase == "train":
+ # Initialize Trainer
+ trainer = Trainer(
+ mode=args.mode,
+ device=args.device,
+ branch=args.branch,
+ train_dataset=args.train_dataset,
+ label_smooth=args.label_smooth,
+ ckpt=args.ckpt,
+ epochs=args.epochs,
+ batch_size=args.batch_size
+ )
+ # Start training
+ trainer.train()
+ #########################################
+ # Phase 2: Evaluation
+ #########################################
+ elif args.phase == "eval":
+ # Initialize Detector
+ detector = Detector(
+ device=args.device,
+ mode=args.mode,
+ train_dataset=args.train_dataset,
+ pretrain=args.pretrain,
+ ckpt=args.ckpt,
+ batch_size=args.batch_size
+ )
+ # Start evaluation
+ detector.evaluate_benchmark()
+ ##########################################
+ # Phase 3: Test on a single image
+ ##########################################
+ elif args.phase == "test":
+ # Initialize Detector
+ detector = Detector(
+ device=args.device,
+ mode=args.mode,
+ train_dataset=args.train_dataset,
+ pretrain=args.pretrain,
+ ckpt=args.ckpt,
+ batch_size=args.batch_size
+ )
+ # Test on a single image
+ score = detector.scan()
+ else:
+ raise ValueError(f"Unknown phase: {args.phase}")
+
+
+if __name__ == "__main__":
+ import argparse
+ parser = argparse.ArgumentParser("Co-Spy: Combining Semantic and Pixel Features to Detect Synthetic Images by AI")
+ parser.add_argument("--gpu",
+ type=int,
+ default=0,
+ help="GPU id to use")
+ parser.add_argument("--phase",
+ type=str,
+ default="test",
+ choices=["train", "eval", "test"],
+ help="Select the phase to run Co-Spy: train / eval / test")
+ parser.add_argument("--mode",
+ type=str,
+ default="fusion",
+ choices=["branch", "fusion", "end2end"],
+ help="Select the mode of Co-Spy training")
+ parser.add_argument("--train_dataset",
+ type=str,
+ default="sd-v1_4",
+ help="Training dataset")
+ parser.add_argument("--branch",
+ type=str,
+ default="artifact",
+ choices=["artifact", "semantic"],
+ help="Branch detector (for branch mode)")
+ parser.add_argument("--label_smooth",
+ action="store_true",
+ help="Whether to use label smoothing during training")
+ parser.add_argument("--pretrain",
+ action="store_true",
+ help="Whether to use pre-trained weights for evaluation")
+ parser.add_argument("--ckpt",
+ type=str,
+ default="ckpt",
+ help="Checkpoint directory")
+ parser.add_argument("--epochs",
+ type=int,
+ default=20,
+ help="Number of training epochs")
+ parser.add_argument("--batch_size",
+ type=int,
+ default=32,
+ help="Batch size")
+ parser.add_argument("--seed",
+ type=int,
+ default=1024,
+ help="Random seed")
+
+ args = parser.parse_args()
+
+ # Set random seed
+ seed_torch(args.seed)
+
+ # Set GPU device
+ args.device = f"cuda:{args.gpu}" if torch.cuda.is_available() else "cpu"
+
+ # Run the experiment
+ main(args)
diff --git a/clean/image/cospy/pretrained/download.sh b/clean/image/cospy/pretrained/download.sh
new file mode 100644
index 0000000000000000000000000000000000000000..1661445ce12f1020fd182c8c8f40478c604785b2
--- /dev/null
+++ b/clean/image/cospy/pretrained/download.sh
@@ -0,0 +1,13 @@
+# Download pre-trained weights for CO-SPY
+# Pre-trained weights on CNNDet (ProGAN)
+wget https://huggingface.co/ruojiruoli/Co-Spy-Pretrained-Weights/resolve/main/progan.zip
+unzip progan.zip
+rm progan.zip
+
+# Pre-trained weights on DRCT-2M (Stable Diffusion v1.4)
+wget https://huggingface.co/ruojiruoli/Co-Spy-Pretrained-Weights/resolve/main/sd-v1_4.zip
+unzip sd-v1_4.zip
+rm sd-v1_4.zip
+
+# Clean up unnecessary files
+rm -rf __MACOSX
diff --git a/clean/image/cospy/train.py b/clean/image/cospy/train.py
new file mode 100644
index 0000000000000000000000000000000000000000..80d0f5347c7c12c2d704f18f89849cb8a13e4dfe
--- /dev/null
+++ b/clean/image/cospy/train.py
@@ -0,0 +1,211 @@
+import os
+import time
+import torch
+import importlib
+from loguru import logger
+
+from utils import seed_torch, evaluate
+from dataSets import TrainDataset
+
+import warnings
+warnings.filterwarnings("ignore")
+
+
+class Trainer:
+ def __init__(self,
+ mode: str,
+ device: str,
+ branch: str,
+ train_dataset: str,
+ label_smooth: bool = False,
+ ckpt: str = "ckpt",
+ epochs: int = 20,
+ batch_size: int = 32):
+
+ self.device = device
+ self.mode = mode
+ self.branch = branch
+ self.train_dataset = train_dataset
+ self.label_smooth = label_smooth
+ self.ckpt = ckpt
+ self.epochs = epochs
+ self.batch_size = batch_size
+
+ # Dynamically import the detector module from either "progan" or "sd-v1_4"
+ detector_module = importlib.import_module(f"detectors.{train_dataset}")
+
+ # Get the detector based on mode
+ if self.mode == "branch":
+ ArtifactDetector = getattr(detector_module, "ArtifactDetector")
+ SemanticDetector = getattr(detector_module, "SemanticDetector")
+ if self.branch == "artifact":
+ self.model = ArtifactDetector()
+ elif self.branch == "semantic":
+ self.model = SemanticDetector()
+ else:
+ raise ValueError(f"Unknown detector: {self.branch}")
+ elif self.mode == "fusion":
+ semantic_weights_path = os.path.join(self.ckpt, self.train_dataset, "semantic", "best_model.pth")
+ artifact_weights_path = os.path.join(self.ckpt, self.train_dataset, "artifact", "best_model.pth")
+ if not os.path.exists(semantic_weights_path) or not os.path.exists(artifact_weights_path):
+ raise ValueError("Semantic or Artifact weights path does not exist for fusion mode")
+ CoSpyFusionDetector = getattr(detector_module, "CoSpyFusionDetector")
+ self.model = CoSpyFusionDetector(
+ semantic_weights_path=semantic_weights_path,
+ artifact_weights_path=artifact_weights_path)
+ elif self.mode == "end2end":
+ End2EndDetector = getattr(detector_module, "End2EndDetector")
+ self.model = End2EndDetector()
+ else:
+ raise ValueError(f"Unknown mode: {self.mode}")
+
+ self.model.to(self.device)
+
+ # Initialize the fc layer
+ torch.nn.init.normal_(self.model.fc.weight.data, 0.0, 0.02)
+ if self.mode == "end2end":
+ torch.nn.init.normal_(self.model.sem.fc.weight.data, 0.0, 0.02)
+ torch.nn.init.normal_(self.model.art.fc.weight.data, 0.0, 0.02)
+
+ # Optimizer
+ _beta1 = 0.9
+ _weight_decay = 0.0
+ params = [p for p in self.model.parameters() if p.requires_grad]
+ logger.info(f"Trainable parameters: {len(params)}")
+
+ self._lr = 1e-4 if self.mode != "fusion" else 1e-1
+ self.optimizer = torch.optim.AdamW(params, lr=self._lr, betas=(_beta1, 0.999), weight_decay=_weight_decay)
+
+ # Loss function
+ if self.label_smooth:
+ self.criterion = LabelSmoothingBCEWithLogits(smoothing=0.1)
+ else:
+ self.criterion = torch.nn.BCEWithLogitsLoss()
+
+ # Scheduler
+ self.delr_freq = 10
+
+ def train_step(self, batch_data):
+ inputs, labels = batch_data
+ inputs, labels = inputs.to(self.device), labels.to(self.device)
+
+ self.optimizer.zero_grad()
+ outputs = self.model(inputs)
+ loss = self.criterion(outputs, labels.unsqueeze(1).float())
+ loss.backward()
+ self.optimizer.step()
+
+ eval_loss = loss.item()
+ y_pred = outputs.sigmoid().flatten().tolist()
+ y_true = labels.tolist()
+ return eval_loss, y_pred, y_true
+
+ def scheduler(self, status_dict):
+ epoch = status_dict["epoch"]
+ if epoch % self.delr_freq == 0 and epoch != 0:
+ for param_group in self.optimizer.param_groups:
+ param_group["lr"] *= 0.9
+ self._lr = param_group["lr"]
+
+ def train(self):
+ # Determine data split and transform based on mode
+ if self.mode == "fusion":
+ train_split, val_split = "val", "val"
+ train_transform = self.model.test_transform
+ test_transform = self.model.test_transform
+ else: # branch or end2end mode
+ train_split, val_split = "train", "val"
+ train_transform = self.model.train_transform
+ test_transform = self.model.test_transform
+
+ # Determine save directory
+ if self.mode == "branch":
+ subdir = self.branch
+ else:
+ subdir = self.mode
+ # Set the saving directory
+ model_dir = os.path.join(self.ckpt, self.train_dataset, subdir)
+ if not os.path.exists(model_dir):
+ os.makedirs(model_dir)
+
+ # Setup logger
+ log_path = f"{model_dir}/training.log"
+ if os.path.exists(log_path):
+ os.remove(log_path)
+
+ logger_id = logger.add(
+ log_path,
+ format="{time:MM-DD at HH:mm:ss} | {level} | {module}:{line} | {message}",
+ level="DEBUG",
+ )
+
+ # Add JPEG compression for sd-v1_4 dataset
+ self.add_jpeg = True if self.train_dataset == "sd-v1_4" else False
+
+ # Load the training and validation dataset
+ train_dataset = TrainDataset(train_dataset=self.train_dataset,
+ split=train_split,
+ add_jpeg=self.add_jpeg,
+ transform=train_transform)
+ train_loader = torch.utils.data.DataLoader(train_dataset,
+ batch_size=self.batch_size,
+ shuffle=True,
+ num_workers=4,
+ pin_memory=True)
+ val_dataset = TrainDataset(train_dataset=self.train_dataset,
+ split=val_split,
+ add_jpeg=self.add_jpeg,
+ transform=test_transform)
+ val_loader = torch.utils.data.DataLoader(val_dataset,
+ batch_size=self.batch_size,
+ shuffle=False,
+ num_workers=4,
+ pin_memory=True)
+
+ logger.info(f"Train size {len(train_dataset)} | Val size {len(val_dataset)}")
+
+ # Train the detector
+ best_acc = 0
+ for epoch in range(self.epochs):
+ self.model.train()
+ time_start = time.time()
+ for step_id, batch_data in enumerate(train_loader):
+ eval_loss, y_pred, y_true = self.train_step(batch_data)
+ ap, accuracy = evaluate(y_pred, y_true)
+
+ if (step_id + 1) % 100 == 0:
+ time_end = time.time()
+ logger.info(f"Epoch {epoch} | Batch {step_id + 1}/{len(train_loader)} | Loss {eval_loss:.4f} | AP {ap*100:.2f}% | Accuracy {accuracy*100:.2f}% | Time {time_end-time_start:.2f}s")
+ time_start = time.time()
+
+ # Evaluate the model
+ self.model.eval()
+ y_pred, y_true = [], []
+ for (images, labels) in val_loader:
+ y_pred.extend(self.model.predict(images))
+ y_true.extend(labels.tolist())
+
+ ap, accuracy = evaluate(y_pred, y_true)
+ eval_type = "Test" if self.mode == "branch" else "Total"
+ logger.info(f"Epoch {epoch} | {eval_type} AP {ap*100:.2f}% | {eval_type} Accuracy {accuracy*100:.2f}%")
+
+ # Schedule the training
+ status_dict = {"epoch": epoch, "AP": ap, "Accuracy": accuracy}
+ self.scheduler(status_dict)
+
+ # Save the model
+ if accuracy >= best_acc:
+ best_acc = accuracy
+ self.model.save_weights(f"{model_dir}/best_model.pth")
+ logger.info(f"Best model saved with accuracy {best_acc*100:.2f}%")
+
+ if epoch % 5 == 0:
+ self.model.save_weights(f"{model_dir}/epoch_{epoch}.pth")
+ logger.info(f"Model saved at epoch {epoch}")
+
+ # Save the final model
+ self.model.save_weights(f"{model_dir}/final_model.pth")
+ logger.info("Final model saved")
+
+ # Remove the logger
+ logger.remove(logger_id)
diff --git a/clean/image/cospy/train.sh b/clean/image/cospy/train.sh
new file mode 100644
index 0000000000000000000000000000000000000000..bf886e2fc777f36c94e8338f83fc9b84aa063a5b
--- /dev/null
+++ b/clean/image/cospy/train.sh
@@ -0,0 +1,41 @@
+#!/bin/bash
+
+while [[ $# -gt 0 ]]; do
+ case $1 in
+ --dataset)
+ dataset="$2"
+ shift 2
+ ;;
+ --gpu)
+ gpu="$2"
+ shift 2
+ ;;
+ *)
+ echo "Unknown argument: $1"
+ exit 1
+ ;;
+ esac
+done
+
+python main.py \
+ --phase train \
+ --gpu $gpu \
+ --mode branch \
+ --branch artifact \
+ --train_dataset $dataset \
+ --epochs 20
+
+python main.py \
+ --phase train \
+ --gpu $gpu \
+ --mode branch \
+ --branch semantic \
+ --train_dataset $dataset \
+ --epochs 10
+
+python main.py \
+ --phase train \
+ --gpu $gpu \
+ --mode fusion \
+ --train_dataset $dataset \
+ --epochs 2
diff --git a/clean/image/cospy/utils.py b/clean/image/cospy/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..e72b7b71bafaa00434e50886617bfbe64cc8af47
--- /dev/null
+++ b/clean/image/cospy/utils.py
@@ -0,0 +1,182 @@
+import os
+import cv2
+import torch
+import pickle
+import random
+import numpy as np
+from io import BytesIO
+from PIL import Image, ImageFile
+from torchvision import transforms
+import torchvision.transforms.functional as TF
+from scipy.ndimage.filters import gaussian_filter
+from sklearn.metrics import average_precision_score
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+# Set random seed
+def seed_torch(seed):
+ random.seed(seed)
+ os.environ['PYTHONHASHSEED'] = str(seed)
+ np.random.seed(seed)
+ torch.manual_seed(seed)
+ torch.cuda.manual_seed(seed)
+ torch.cuda.manual_seed_all(seed)
+ torch.backends.cudnn.benchmark = False
+ torch.backends.cudnn.deterministic = True
+
+
+# Load dataset
+def recursively_read(rootdir, must_contain, exts=["png", "PNG", "jpg", "JPG", "jpeg", "JPEG"]):
+ out = []
+ for r, d, f in os.walk(rootdir):
+ for file in f:
+ if (file.split('.')[1] in exts) and (must_contain in os.path.join(r, file)):
+ out.append(os.path.join(r, file))
+ return out
+
+
+def get_list(path, must_contain=''):
+ if ".pickle" in path:
+ with open(path, 'rb') as f:
+ image_list = pickle.load(f)
+ image_list = [item for item in image_list if must_contain in item]
+ else:
+ image_list = recursively_read(path, must_contain)
+ return image_list
+
+
+# Data augmentation techniques
+def data_augment(img, aug_config):
+ img = np.array(img)
+ if img.ndim == 2:
+ img = np.expand_dims(img, axis=2)
+ img = np.repeat(img, 3, axis=2)
+
+ if random.random() < aug_config["blur_prob"]:
+ sig = sample_continuous(aug_config["blur_sig"])
+ gaussian_blur(img, sig)
+
+ if random.random() < aug_config["jpg_prob"]:
+ method = sample_discrete(aug_config["jpg_method"])
+ qual = sample_discrete(aug_config["jpg_qual"])
+ img = jpeg_from_key(img, qual, method)
+
+ return Image.fromarray(img)
+
+
+# Data augmentation techniques
+def tensor_data_augment(images, aug_config):
+ device = images.device
+ images = images.detach().cpu().permute(0, 2, 3, 1).numpy()
+ images = np.uint8(images * 255.)
+ outputs = []
+ for img in images:
+ if random.random() < aug_config["blur_prob"]:
+ sig = sample_continuous(aug_config["blur_sig"])
+ gaussian_blur(img, sig)
+
+ if random.random() < aug_config["jpg_prob"]:
+ method = sample_discrete(aug_config["jpg_method"])
+ qual = sample_discrete(aug_config["jpg_qual"])
+ img = jpeg_from_key(img, qual, method)
+ outputs.append(img)
+ outputs = np.stack(outputs)
+ outputs = torch.from_numpy(outputs).to(device).permute(0, 3, 1, 2).float() / 255.
+ return outputs
+
+
+# Sample continuous or discrete values
+def sample_continuous(s):
+ if len(s) == 1:
+ return s[0]
+ if len(s) == 2:
+ rg = s[1] - s[0]
+ return random.random() * rg + s[0]
+ raise ValueError("Length of iterable s should be 1 or 2.")
+
+
+def sample_discrete(s):
+ if len(s) == 1:
+ return s[0]
+ return random.choice(s)
+
+
+# Gaussian blur
+def gaussian_blur(img, sigma):
+ gaussian_filter(img[:,:,0], output=img[:,:,0], sigma=sigma)
+ gaussian_filter(img[:,:,1], output=img[:,:,1], sigma=sigma)
+ gaussian_filter(img[:,:,2], output=img[:,:,2], sigma=sigma)
+
+
+# JPEG compression
+def cv2_jpg(img, compress_val):
+ img_cv2 = img[:,:,::-1]
+ encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), compress_val]
+ result, encimg = cv2.imencode('.jpg', img_cv2, encode_param)
+ decimg = cv2.imdecode(encimg, 1)
+ return decimg[:,:,::-1]
+
+
+def pil_jpg(img, compress_val):
+ out = BytesIO()
+ img = Image.fromarray(img)
+ img.save(out, format='jpeg', quality=compress_val)
+ img = Image.open(out)
+ # load from memory before ByteIO closes
+ img = np.array(img)
+ out.close()
+ return img
+
+
+def png_to_jpeg(img, quality=95):
+ # Convert the PNG image to JPEG
+ # Input: PIL image
+ # Output: PIL image
+ out = BytesIO()
+ img.save(out, format='jpeg', quality=quality)
+ img = np.array(Image.open(out))
+ # Load from memory before ByteIO closes
+ out.close()
+ img = Image.fromarray(img)
+ return img
+
+
+def jpeg_from_key(img, compress_val, key):
+ jpeg_dict = {'cv2': cv2_jpg, 'pil': pil_jpg}
+ method = jpeg_dict[key]
+ return method(img, compress_val)
+
+
+# Custom resize function
+def custom_resize(img, rz_interp, loadSize):
+ rz_dict = {'bilinear': Image.BILINEAR,
+ 'bicubic': Image.BICUBIC,
+ 'lanczos': Image.LANCZOS,
+ 'nearest': Image.NEAREST}
+ interp = sample_discrete(rz_interp)
+ return TF.resize(img, loadSize, interpolation=rz_dict[interp])
+
+
+def weights2cpu(weights):
+ for key in weights:
+ weights[key] = weights[key].cpu()
+ return weights
+
+
+# Define the label smoothing loss
+class LabelSmoothingBCEWithLogits(torch.nn.Module):
+ def __init__(self, smoothing=0.1):
+ super(LabelSmoothingBCEWithLogits, self).__init__()
+ self.smoothing = smoothing
+
+ def forward(self, pred, target):
+ target = target.float() * (1.0 - self.smoothing) + 0.5 * self.smoothing
+ loss = torch.nn.functional.binary_cross_entropy_with_logits(pred, target, reduction='mean')
+ return loss
+
+
+def evaluate(y_pred, y_true):
+ ap = average_precision_score(y_true, y_pred)
+ accuracy = ((np.array(y_pred) > 0.5) == y_true).mean()
+ return ap, accuracy
diff --git a/clean/image/effort/DeepfakeBench/analysis/logits_decision_boundary.py b/clean/image/effort/DeepfakeBench/analysis/logits_decision_boundary.py
new file mode 100644
index 0000000000000000000000000000000000000000..4c7899d240a99cf9b802696aa0a92082e1f87960
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/analysis/logits_decision_boundary.py
@@ -0,0 +1,404 @@
+"""
+eval pretained model.
+"""
+import os
+import numpy as np
+from os.path import join
+import cv2
+import random
+import datetime
+import time
+import yaml
+import pickle
+from tqdm import tqdm
+from copy import deepcopy
+from PIL import Image as pil_image
+from metrics.utils import get_test_metrics
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.backends.cudnn as cudnn
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+import torch.utils.data
+import torch.optim as optim
+
+from dataset.abstract_dataset import DeepfakeAbstractBaseDataset
+from dataset.ff_blend import FFBlendDataset
+from dataset.fwa_blend import FWABlendDataset
+from dataset.pair_dataset import pairDataset
+
+from trainer.trainer import Trainer
+from detectors import DETECTOR
+from metrics.base_metrics_class import Recorder
+from collections import defaultdict
+
+import argparse
+from logger import create_logger
+
+parser = argparse.ArgumentParser(description='Process some paths.')
+parser.add_argument('--detector_path', type=str,
+ default='/home/zhiyuanyan/DeepfakeBench/training/config/detector/resnet34.yaml',
+ help='path to detector YAML file')
+parser.add_argument("--test_dataset", nargs="+")
+parser.add_argument("--model_name", nargs="+")
+parser.add_argument('--weights_path', type=str,
+ default='/mntcephfs/lab_data/zhiyuanyan/benchmark_results/auc_draw/cnn_aug/resnet34_2023-05-20-16-57-22/test/FaceForensics++/ckpt_epoch_9_best.pth')
+#parser.add_argument("--lmdb", action='store_true', default=False)
+args = parser.parse_args()
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+on_2060 = "2060" in torch.cuda.get_device_name()
+def init_seed(config):
+ if config['manualSeed'] is None:
+ config['manualSeed'] = random.randint(1, 10000)
+ random.seed(config['manualSeed'])
+ torch.manual_seed(config['manualSeed'])
+ if config['cuda']:
+ torch.cuda.manual_seed_all(config['manualSeed'])
+
+
+def prepare_testing_data(config):
+ def get_test_data_loader(config, test_name):
+ # update the config dictionary with the specific testing dataset
+ config = config.copy() # create a copy of config to avoid altering the original one
+ config['test_dataset'] = test_name # specify the current test dataset
+ test_set = DeepfakeAbstractBaseDataset(
+ config=config,
+ mode='test',
+ )
+ test_data_loader = \
+ torch.utils.data.DataLoader(
+ dataset=test_set,
+ batch_size=config['test_batchSize'],
+ shuffle=False,
+ num_workers=int(config['workers']),
+ collate_fn=test_set.collate_fn,
+ drop_last=False
+ )
+ return test_data_loader
+
+ assert len(config['test_dataset'])==1
+ test_data_loader = get_test_data_loader(config, config['test_dataset'][0])
+ return test_data_loader
+
+
+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 plot_logit_distribution(real_logits, fake_logits, save_path):
+ plt.figure(figsize=(8, 8))
+
+ # 绘制散点图
+ plt.scatter(real_logits[:, 0], real_logits[:, 1], color='blue', label='Real', alpha=0.2)
+ plt.scatter(fake_logits[:, 0], fake_logits[:, 1], color='red', label='Fake', alpha=0.2)
+
+ # 添加过原点的决策边界 x = y
+ min_val = min(real_logits[:, 0].min(), real_logits[:, 1].min(), fake_logits[:, 0].min(), fake_logits[:, 1].min())
+ max_val = max(real_logits[:, 0].max(), real_logits[:, 1].max(), fake_logits[:, 0].max(), fake_logits[:, 1].max())
+ x_vals = np.linspace(min_val, max_val, 100)
+ plt.plot(x_vals, x_vals, color='black', linestyle='--', label='Decision Boundary (x = y)')
+
+ # 设置标题和标签
+ plt.title('Logit Outputs for Real and Fake Classes')
+ plt.xlabel('Logit for Real Class')
+ plt.ylabel('Logit for Fake Class')
+ plt.legend()
+ plt.grid(True)
+
+ # 保存图像
+ plt.tight_layout()
+ os.makedirs(save_path, exist_ok=True)
+ plt.savefig(os.path.join(save_path, args.test_dataset[0] + '.png'))
+ plt.show()
+
+
+def calculate_logits_and_confidence(model, data_loader):
+ model.eval()
+
+ real_logits = []
+ fake_logits = []
+ real_confidences = []
+ fake_confidences = []
+
+ for i, data_dict in tqdm(enumerate(data_loader), total=len(data_loader)):
+ # get data
+ data, label, mask, landmark = \
+ data_dict['image'], data_dict['label'], data_dict['mask'], data_dict['landmark']
+ label = torch.where(data_dict['label'] != 0, 1, 0)
+ # move data to GPU
+ data_dict['image'], data_dict['label'] = data.to(device), label.to(device)
+ if mask is not None:
+ data_dict['mask'] = mask.to(device)
+ if landmark is not None:
+ data_dict['landmark'] = landmark.to(device)
+
+ # model forward without considering gradient computation
+ predictions = inference(model, data_dict)
+ output = predictions['cls']
+ logits = output.detach().cpu()
+ confidences = F.softmax(logits, dim=-1)
+
+ for i in range(logits.size(0)):
+ if label[i] == 0:
+ real_logits.append(logits[i])
+ real_confidences.append(confidences[i])
+ else:
+ fake_logits.append(logits[i])
+ fake_confidences.append(confidences[i])
+
+
+ real_logits = torch.stack(real_logits, dim=0)
+ fake_logits = torch.stack(fake_logits, dim=0)
+ real_confidences = torch.stack(real_confidences, dim=0)
+ fake_confidences = torch.stack(fake_confidences, dim=0)
+
+ return real_logits, fake_logits, real_confidences, fake_confidences
+
+
+
+def test_one_dataset(model, data_loader):
+ prediction_lists = []
+ feature_lists = []
+ label_lists = []
+ for i, data_dict in tqdm(enumerate(data_loader), total=len(data_loader)):
+ # get data
+ data, label, mask, landmark = \
+ data_dict['image'], data_dict['label'], data_dict['mask'], data_dict['landmark']
+ label = torch.where(data_dict['label'] != 0, 1, 0)
+ # move data to GPU
+ data_dict['image'], data_dict['label'] = data.to(device), label.to(device)
+ if mask is not None:
+ data_dict['mask'] = mask.to(device)
+ if landmark is not None:
+ data_dict['landmark'] = landmark.to(device)
+
+ # model forward without considering gradient computation
+ predictions = inference(model, data_dict)
+ label_lists += list(data_dict['label'].cpu().detach().numpy())
+ prediction_lists += list(predictions['prob'].cpu().detach().numpy())
+ feature_lists += list(predictions['feat'].cpu().detach().numpy())
+
+ return np.array(prediction_lists), np.array(label_lists),np.array(feature_lists)
+
+def test_epoch(model, test_data_loaders):
+ # set model to eval mode
+ model.eval()
+
+ # define test recorder
+ metrics_all_datasets = {}
+
+ # testing for all test data
+ keys = test_data_loaders.keys()
+ for key in keys:
+ data_dict = test_data_loaders[key].dataset.data_dict
+ # compute loss for each dataset
+ predictions_nps, label_nps,feat_nps = test_one_dataset(model, test_data_loaders[key])
+
+ # compute metric for each dataset
+ metric_one_dataset = get_test_metrics(y_pred=predictions_nps, y_true=label_nps,
+ img_names=data_dict['image'])
+ metrics_all_datasets[key] = metric_one_dataset
+
+ # info for each dataset
+ tqdm.write(f"dataset: {key}")
+ for k, v in metric_one_dataset.items():
+ tqdm.write(f"{k}: {v}")
+
+ return metrics_all_datasets
+
+@torch.no_grad()
+def inference(model, data_dict):
+ predictions = model(data_dict, inference=True)
+ return predictions
+
+
+def main():
+ # parse options and load config
+ with open(args.detector_path, 'r') as f:
+ config = yaml.safe_load(f)
+ with open('./training/config/test_config.yaml', 'r') as f:
+ config2 = yaml.safe_load(f)
+ config.update(config2)
+ if on_2060:
+ config['lmdb_dir'] = r'I:\transform_2_lmdb'
+ config['train_batchSize'] = 10
+ config['workers'] = 0
+ else:
+ config['workers'] = 8
+ config['lmdb_dir'] = r'/mnt/chongqinggeminiceph1fs/geminicephfs/mm-base-vision/jikangcheng/data/LMDBs'
+ weights_path = None
+ # If arguments are provided, they will overwrite the yaml settings
+ if args.test_dataset:
+ config['test_dataset'] = args.test_dataset
+ if args.weights_path:
+ config['weights_path'] = args.weights_path
+ weights_path = args.weights_path
+
+ # init seed
+ init_seed(config)
+
+ # set cudnn benchmark if needed
+ if config['cudnn']:
+ cudnn.benchmark = True
+
+ # prepare the testing data loader
+ test_data_loader = prepare_testing_data(config)
+
+ # prepare the model (detector)
+ model_class = DETECTOR[config['model_name']]
+ model = model_class(config).to(device).eval()
+ if weights_path:
+ try:
+ epoch = int(weights_path.split('/')[-1].split('.')[0].split('_')[2])
+ except:
+ epoch = 0
+ ckpt = torch.load(weights_path, map_location=device)
+ new_weights = {}
+ for key, value in ckpt.items():
+ new_key = key.replace('module.', '')
+ new_weights[new_key] = value
+
+ model.load_state_dict(new_weights, strict=True)
+ print('===> Load checkpoint done!')
+ else:
+ print('Fail to load the pre-trained weights')
+
+ # # start testing
+ # best_metric = test_epoch(model, test_data_loaders)
+ # print('===> Test Done!')
+
+ # Box plot
+ real_logits, fake_logits, real_confidences, fake_confidences = calculate_logits_and_confidence(model, test_data_loader)
+
+ # Plot the logit distribution in 2D space
+ save_path = f"scatter_plot/visualization/{args.model_name[0]}"
+ if not os.path.exists(save_path):
+ os.makedirs(save_path)
+
+ plot_logit_distribution(real_logits, fake_logits, save_path)
+
+
+ save_path = f"box_plot/visualization/{args.model_name[0]}"
+ if not os.path.exists(save_path):
+ os.makedirs(save_path)
+
+ import pandas as pd
+ import matplotlib.pyplot as plt
+ import seaborn as sns
+
+ # Convert tensors to numpy arrays
+ # logits_real = real_logits.cpu().numpy().reshape(-1)
+ # logits_fake = fake_logits.cpu().numpy().reshape(-1)
+
+ # real_confidences = real_confidences.cpu().numpy().reshape(-1)
+ # fake_confidences = fake_confidences.cpu().numpy().reshape(-1)
+
+ # logits_real = torch.cat([real_logits[0], fake_logits[0]]).cpu().numpy().reshape(-1)
+ # logits_fake = torch.cat([real_logits[1], fake_logits[1]]).cpu().numpy().reshape(-1)
+
+
+
+ # Set Seaborn style
+ sns.set_style("whitegrid")
+ plt.rcParams['font.size'] = 14
+ plt.rcParams['axes.labelsize'] = 14
+ plt.rcParams['axes.titlesize'] = 16
+ plt.rcParams['xtick.labelsize'] = 12
+ plt.rcParams['ytick.labelsize'] = 12
+ plt.rcParams['legend.fontsize'] = 14
+
+ # Create subplots
+ fig, axs = plt.subplots(2, 1, figsize=(8, 8))
+
+
+
+ if real_logits.shape[-1] == 2:
+ logits_real = (real_logits[:, 0].cpu().numpy(), fake_logits[:, 0].cpu().numpy())
+ logits_fake = (real_logits[:, 1].cpu().numpy(), fake_logits[:, 1].cpu().numpy())
+
+ real_confidences_list = (real_confidences[:, 0].cpu().numpy(), fake_confidences[:, 0].cpu().numpy())
+ fake_confidences_list = (real_confidences[:, 1].cpu().numpy(), fake_confidences[:, 1].cpu().numpy())
+
+ # Plot histograms and KDE for logits
+ sns.histplot(logits_real[0], bins=30, kde=True, stat='probability', label='real', color=sns.color_palette("Blues")[-1], ax=axs[0])
+ sns.histplot(logits_real[1], bins=30, kde=True, stat='probability', label='fake', color=sns.color_palette("Reds")[-1], ax=axs[0])
+ axs[0].set_ylim(0, 0.15)
+ axs[0].legend()
+
+ # sns.histplot(logits_fake[0], bins=30, kde=True, stat='probability', color=sns.color_palette("Blues")[-1], ax=axs[0, 1])
+ # sns.histplot(logits_fake[1], bins=30, kde=True, stat='probability', color=sns.color_palette("Reds")[-1], ax=axs[0, 1])
+ # axs[0, 1].set_ylim(0, 0.15)
+ # axs[0, 1].legend()
+
+ # Plot histograms and KDE for confidences
+ sns.histplot(real_confidences_list[0], bins=30, kde=True, stat='probability', label='real', color=sns.color_palette("Blues")[-1], ax=axs[1])
+ sns.histplot(real_confidences_list[1], bins=30, kde=True, stat='probability', label='fake', color=sns.color_palette("Reds")[-1], ax=axs[1])
+ axs[1].legend()
+ # sns.histplot(fake_confidences_list[0], bins=30, kde=True, stat='probability', color=sns.color_palette("Blues")[-1], ax=axs[1])
+ # sns.histplot(fake_confidences_list[1], bins=30, kde=True, stat='probability', color=sns.color_palette("Reds")[-1], ax=axs[1])
+
+ # Set titles
+ axs[0].set_title('Logit Distribution of Real and Fake Classes')
+ # axs[0, 1].set_title('Fake Logits')
+ axs[1].set_title('Confidence Distribution of Real and Fake Classes')
+ # axs[1, 1].set_title('Fake Confidences')
+
+ # Show the plot
+ plt.tight_layout()
+ save_ = os.path.join(save_path, args.test_dataset[0] + '.png')
+ plt.savefig(save_)
+
+ # # Create a DataFrame with logits and confidences as columns
+ # df_real = pd.DataFrame({'logits_real': logits_real, 'real_confidences': real_confidences})
+ # df_fake = pd.DataFrame({'logits_fake': logits_fake, 'fake_confidences': fake_confidences})
+
+ # # Write DataFrame to csv file
+ # df_real.to_csv('output_logits_real.csv', index=False)
+ # df_fake.to_csv('output_logits_fake.csv', index=False)
+
+
+
+
+ elif real_logits.shape[-1] == 1:
+
+ logits_real = real_logits.cpu().numpy().flatten()
+ logits_fake = fake_logits.cpu().numpy().flatten()
+ real_confidences_list = real_confidences.cpu().numpy().flatten()
+ fake_confidences_list = fake_confidences.cpu().numpy().flatten()
+
+ # Plot histograms and KDE for logits
+ sns.histplot(logits_fake, bins=30, kde=True, stat='probability', label='real', color=sns.color_palette("Blues")[-1], ax=axs[0])
+ sns.histplot(logits_real, bins=30, kde=True, stat='probability', label='fake', color=sns.color_palette("Reds")[-1], ax=axs[0])
+ axs[0].set_ylim(0, 0.7)
+ axs[0].legend()
+
+ # Plot histograms and KDE for confidences
+ sns.histplot(real_confidences_list, bins=30, kde=True, stat='probability', label='real', color=sns.color_palette("Blues")[-1], ax=axs[1])
+ sns.histplot(fake_confidences_list, bins=30, kde=True, stat='probability', label='fake', color=sns.color_palette("Reds")[-1], ax=axs[1])
+ axs[1].legend()
+
+ # Set titles
+ axs[0].set_title('Logit Distribution of Real and Fake Classes')
+ axs[1].set_title('Confidence Distribution of Real and Fake Classes')
+
+ # Show the plot
+ plt.tight_layout()
+ save_ = os.path.join(save_path, args.test_dataset[0] + '.png')
+ plt.savefig(save_)
+
+
+
+
+ else:
+ raise ValueError('The number of classes is not 1 or 2')
+
+
+
+if __name__ == '__main__':
+ main()
diff --git a/clean/image/effort/DeepfakeBench/analysis/pca_rank.py b/clean/image/effort/DeepfakeBench/analysis/pca_rank.py
new file mode 100644
index 0000000000000000000000000000000000000000..cd7d14193f543c68f6bea2535f7aedc375579dad
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/analysis/pca_rank.py
@@ -0,0 +1,122 @@
+import os
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+import seaborn as sns
+from sklearn.decomposition import PCA
+import pickle
+import torch
+import torch.nn as nn
+
+color_map = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red', 'tab:purple']
+label_dict = {
+ 0: 'FF-Real', 1: 'FF-Fake', 2: 'CDF-Real', 3: 'CDF-Fake', 4: 'FF-Fsh'
+}
+
+
+
+
+import argparse
+parser = argparse.ArgumentParser(description='Process some paths.')
+parser.add_argument("--test_dataset", nargs="+", default=['FF-DF'])
+parser.add_argument("--model_name", nargs="+")
+args = parser.parse_args()
+
+
+
+def plot_explained_variance(explained_variance_ratio, cumulative_explained_variance_ratio, save_path):
+ plt.figure(figsize=(10, 8))
+ plt.plot(explained_variance_ratio[:200], marker='o', alpha=0.4, markersize=12)
+ plt.title('Explained Variance Ratio of Principal Components', fontsize=24, weight='bold', y=1.05)
+ plt.xlabel('Principal Component Index', fontsize=24, weight='bold')
+ plt.ylabel('Explained Variance Ratio', fontsize=24, weight='bold')
+ plt.tick_params(axis='both', labelsize=20) # 或者设置成24，如果希望和轴标签一样大
+ plt.grid(True)
+ plt.savefig(os.path.join(save_path, 'explained_variance_ratio.png'), dpi=300)
+
+ plt.figure(figsize=(10, 8))
+ plt.plot(cumulative_explained_variance_ratio[:200], marker='o', alpha=0.4, markersize=12)
+ plt.title('Cumulative Explained Variance Ratio', fontsize=24, weight='bold', y=1.05)
+ plt.xlabel('Number of Principal Components', fontsize=24, weight='bold')
+ plt.ylabel('Cumulative Explained Variance Ratio', fontsize=24, weight='bold')
+ plt.grid(True)
+ plt.savefig(os.path.join(save_path, 'cumulative_explained_variance_ratio.png'), dpi=300)
+
+def main():
+ base_dir = './zhiyuanyan/DeepfakeBenchv2/features'
+ model_name = args.model_name[0]
+ base_dir = os.path.join(base_dir, model_name)
+ test_dataset = args.test_dataset
+ all_testing_data = []
+ all_testing_label = []
+ pool = nn.AdaptiveAvgPool2d((1, 1))
+
+ for name in os.listdir(base_dir):
+ if name in test_dataset:
+ with open(os.path.join(base_dir, name, 'tsne.pkl'), 'rb') as f:
+ data = pickle.load(f)
+ feat = data['feat']
+ print(f'shape of {name}:', feat.shape)
+ if feat.ndim == 4: # pool
+ print(f'shape of {name} before pooling:', feat.shape)
+ feat = torch.from_numpy(feat)
+ feat = pool(feat).squeeze()
+ feat = feat.numpy()
+ print(f'shape of {name} after pooling:', feat.shape)
+ label = data['label']
+ if name == 'Celeb-DF-v2':
+ label = label + 2
+ elif name == 'FaceShifter':
+ label_mask = (label == 1)
+ feat = feat[label_mask]
+ label = label[label_mask] * 4
+ all_testing_data.append(feat)
+ all_testing_label.append(label)
+ all_testing_data = np.concatenate(all_testing_data, axis=0)
+ all_testing_label = np.concatenate(all_testing_label, axis=0)
+ print('Total number of samples:', len(all_testing_label))
+ print('Label distribution:', np.unique(all_testing_label, return_counts=True))
+
+ # Perform PCA analysis
+ feat = all_testing_data
+ pca = PCA(n_components=min(feat.shape[0], feat.shape[1]))
+ pca.fit(feat)
+ explained_variance_ratio = pca.explained_variance_ratio_
+ cumulative_explained_variance_ratio = np.cumsum(explained_variance_ratio)
+
+ # Plot explained variance ratio and cumulative explained variance ratio
+ save_path = './zhiyuanyan/DeepfakeBenchv2/pca_results/' + model_name
+ os.makedirs(save_path, exist_ok=True)
+ plot_explained_variance(explained_variance_ratio, cumulative_explained_variance_ratio, save_path)
+
+ # Output the number of principal components needed to explain 90% variance
+ num_components_90 = np.argmax(cumulative_explained_variance_ratio >= 0.9) + 1
+ print(f'Number of principal components explaining 90% variance: {num_components_90}')
+ print(f'Feature dimension: {feat.shape[1]}')
+
+ # 计算特征值衰减系数
+ eigenvalues = pca.explained_variance_
+ decay_rate = np.diff(np.log(eigenvalues)) / np.diff(np.arange(len(eigenvalues)))
+ print("平均衰减率:", np.mean(decay_rate[:100]))
+
+ # Optional: Plot scatter plot of the first two principal components
+ feat_transformed = pca.transform(feat)[:, :2]
+ numerical_labels = all_testing_label
+ labels = [label_dict[label] for label in numerical_labels]
+
+ plt.figure(figsize=(10, 8))
+ sns.scatterplot(
+ x=feat_transformed[:, 0],
+ y=feat_transformed[:, 1],
+ hue=labels,
+ palette=color_map[:len(np.unique(numerical_labels))],
+ alpha=0.4
+ )
+ plt.title('Scatter Plot of the First Two Principal Components', fontsize=20)
+ plt.xlabel('Principal Component 1', fontsize=20)
+ plt.ylabel('Principal Component 2', fontsize=20)
+ plt.legend(title='Classes', loc='best')
+ plt.savefig(os.path.join(save_path, 'pca_scatter.png'), dpi=300)
+
+if __name__ == '__main__':
+ main()
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/preprocessing/dataset_json/readme.txt b/clean/image/effort/DeepfakeBench/preprocessing/dataset_json/readme.txt
new file mode 100644
index 0000000000000000000000000000000000000000..44b4905491d315a315daf57d954eaac3be8ae737
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/preprocessing/dataset_json/readme.txt
@@ -0,0 +1,9 @@
+-----
+
+Links of the provided JSON Files：
+
+DF40：https://drive.google.com/drive/folders/1rh-82Rn0pqQ7xzLDKBv9gdjy-sbxYC8-?usp=drive_link
+
+DeepfakeBench: https://drive.google.com/drive/folders/1T4ajtH-7PTAmDYWjn27XNOONXs7KjQWk?usp=drive_link
+
+-----
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/__init__.py b/clean/image/effort/DeepfakeBench/training/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/clean/image/effort/DeepfakeBench/training/config/__init__.py b/clean/image/effort/DeepfakeBench/training/config/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..eec7c2d57b94a1fff5f3af445959378441f9c2cf
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/config/__init__.py
@@ -0,0 +1,12 @@
+# ------------------------------------------------------------------------------
+# Copyright (c) Microsoft
+# Licensed under the MIT License.
+# Written by Ke Sun (sunk@mail.ustc.edu.cn)
+# ------------------------------------------------------------------------------
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from .default import _C as config
+from .default import update_config
+from .models import MODEL_EXTRAS
diff --git a/clean/image/effort/DeepfakeBench/training/config/detector/effort.yaml b/clean/image/effort/DeepfakeBench/training/config/detector/effort.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..780c1411382859b4c175ac7ce1b4deb84d4c1502
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/config/detector/effort.yaml
@@ -0,0 +1,88 @@
+# log dir
+log_dir: ./zhiyuanyan/logs/benchv2/icml25/release
+
+# model setting
+pretrained: 'no need to provide this for the effort model'
+model_name: effort # model name
+backbone_name: vit # backbone name
+
+#backbone setting
+backbone_config:
+ mode: original
+ num_classes: 2
+ inc: 3
+ 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: [FaceForensics++]
+test_dataset: [Celeb-DF-v2, FaceShifter, DeeperForensics-1.0]
+
+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': 8, 'test': 8} # number of frames to use per video in training and testing
+resolution: 224 # 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
+
+
+# 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.48145466, 0.4578275, 0.40821073]
+std: [0.26862954, 0.26130258, 0.27577711]
+
+# 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: true # whether to save checkpoint
+save_feat: true # whether to save features
+
+# loss function
+loss_func: cross_entropy # loss function to use
+losstype: null
+
+# metric
+metric_scoring: auc # metric for evaluation (auc, acc, eer, ap)
+
+# cuda
+ngpu: 1 # number of GPUs to use
+cuda: true # whether to use CUDA acceleration
+cudnn: true # whether to use CuDNN for convolution operations
+
+save_avg: true
diff --git a/clean/image/effort/DeepfakeBench/training/config/test_config.yaml b/clean/image/effort/DeepfakeBench/training/config/test_config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..281dc31d2223d40df80e3bc3e2dd8b448ccf23d4
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/config/test_config.yaml
@@ -0,0 +1,232 @@
+mode: test
+lmdb: False
+dataset_root_rgb: './datasets'
+lmdb_dir: 'I:\transform_2_lmdb'
+dataset_json_folder: '/Youtu_Pangu_Security_Public/youtu-pangu-public/zhiyuanyan/DeepfakeBenchv2/preprocessing/dataset_json'
+
+# 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_Fake: 1
+ roop_Real: 0
+
+ # DF40 all
+ FSAll_Fake: 1
+ FSAll_Real: 0
+ FRAll_Fake: 1
+ FRAll_Real: 0
+ EFSAll_Fake: 1
+ EFSAll_Real: 0
+
+
+ DFR_Real: 0
+ DFR_Fake: 1
+ WDF_Real: 0
+ WDF_Fake: 1
+ FFIW_Real: 0
+ FFIW_Fake: 1
+
+ SRI_fake: 1
+ SRI_real: 0
+
+ # DF40
+ e4s_Fake: 1
+ e4s_Real: 0
+
+ MidJourney_Fake: 1
+ MidJourney_Real: 0
+
+ deepfacelab_Fake: 1
+ deepfacelab_Real: 0
+
+ danet_Fake: 1
+ danet_Real: 0
+
+ fomm_Fake: 1
+ fomm_Real: 0
+
+ e4e_Fake: 1
+ e4e_Real: 0
+
+ hyperreenact_Fake: 1
+ hyperreenact_Real: 0
+
+ MRAA_Fake: 1
+ MRAA_Real: 0
+
+ one_shot_free_Fake: 1
+ one_shot_free_Real: 0
+
+ pirender_Fake: 1
+ pirender_Real: 0
+
+ tpsm_Fake: 1
+ tpsm_Real: 0
+
+ facedancer_Fake: 1
+ facedancer_Real: 0
+
+ facevid2vid_Fake: 1
+ facevid2vid_Real: 0
+
+ mcnet_Fake: 1
+ mcnet_Real: 0
+
+ fsgan_Fake: 1
+ fsgan_Real: 0
+
+ lia_Fake: 1
+ lia_Real: 0
+
+ mraa_Fake: 1
+ mraa_Real: 0
+
+ inswap_Fake: 1
+ inswap_Real: 0
+
+ simswap_Fake: 1
+ simswap_Real: 0
+
+ sadtalker_Fake: 1
+ sadtalker_Real: 0
+
+ wav2lip_Fake: 1
+ wav2lip_Real: 0
+
+ uniface_Fake: 1
+ uniface_Real: 0
+
+ blendface_Fake: 1
+ blendface_Real: 0
+
+ mobileswap_Fake: 1
+ mobileswap_Real: 0
+
+ faceswap_Fake: 1
+ faceswap_Real: 0
+
+ dalle2_face_Fake: 1
+ dalle2_face_Real: 0
+
+ MidJourney_Fake: 1
+ MidJourney_Real: 0
+
+ heygen_Fake: 1
+ heygen_Real: 0
+
+ whichisreal_Fake: 1
+ whichisreal_Real: 0
+
+ StyleGAN2_Fake: 1
+ StyleGAN2_Real: 0
+
+ StyleGAN3_Fake: 1
+ StyleGAN3_Real: 0
+
+ StyleGANXL_Fake: 1
+ StyleGANXL_Real: 0
+
+ ddim_Fake: 1
+ ddim_Real: 0
+
+ DiT_Fake: 1
+ DiT_Real: 0
+
+ pixart_Fake: 1
+ pixart_Real: 0
+
+ SiT_Fake: 1
+ SiT_Real: 0
+
+ sd1.5_Fake: 1
+ sd1.5_Real: 0
+
+ sd2.1_Fake: 1
+ sd2.1_Real: 0
+
+ VQGAN_Fake: 1
+ VQGAN_Real: 0
+
+ stargan_Fake: 1
+ stargan_Real: 0
+
+ starganv2_Fake: 1
+ starganv2_Real: 0
+
+ styleclip_Fake: 1
+ styleclip_Real: 0
+
+ CollabDiff_Fake: 1
+ CollabDiff_Real: 0
+
+ rddm_Fake: 1
+ rddm_Real: 0
+
+ # GenImage
+ adm_Fake: 1
+ adm_Real: 0
+
+ pixart_Fake: 1
+ pixart_Real: 0
+
+ biggan_Fake: 1
+ biggan_Real: 0
+
+ glide_Fake: 1
+ glide_Real: 0
+
+ midjourney_Fake: 1
+ midjourney_Real: 0
+
+ sdv4_Fake: 1
+ sdv4_Real: 0
+
+ sdv5_Fake: 1
+ sdv5_Real: 0
+
+ vqdm_Fake: 1
+ vqdm_Real: 0
+
+ wukong_Fake: 1
+ wukong_Real: 0
+
+ GPA_fake: 1
+ GPA_real: 0
+
+ Chameleon_Fake: 1
+ Chameleon_Real: 0
+
+ GPT4o_Fake: 1
+ GPT4o_Real: 0
+
+ VIP_Real: 0
+ VIP_Fake: 1
diff --git a/clean/image/effort/DeepfakeBench/training/config/train_config.yaml b/clean/image/effort/DeepfakeBench/training/config/train_config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..ae17f48ec54cfadb8a71dbd2fa8f4e6a557880c8
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/config/train_config.yaml
@@ -0,0 +1,225 @@
+mode: train
+lmdb: False
+dry_run: false
+dataset_root_rgb: './datasets'
+lmdb_dir: 'I:\transform_2_lmdb'
+dataset_json_folder: '/Youtu_Pangu_Security_Public/youtu-pangu-public/zhiyuanyan/DeepfakeBenchv2/preprocessing/dataset_json'
+SWA: False
+
+# 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
+
+
+ # DF40
+ FSAll_Fake: 1
+ FSAll_Real: 0
+
+ FRAll_Fake: 1
+ FRAll_Real: 0
+
+ EFSAll_Fake: 1
+ EFSAll_Real: 0
+
+ DF40_train_Fake: 1
+ DF40_train_Real: 0
+
+ e4s_Fake: 1
+ e4s_Real: 0
+
+ danet_Fake: 1
+ danet_Real: 0
+
+ fomm_Fake: 1
+ fomm_Real: 0
+
+ Collaborative_Diffusion_Fake: 1
+ Collaborative_Diffusion_Real: 0
+
+ e4e_Fake: 1
+ e4e_Real: 0
+
+ hyperreenact_Fake: 1
+ hyperreenact_Real: 0
+
+ MRAA_Fake: 1
+ MRAA_Real: 0
+
+ one_shot_free_Fake: 1
+ one_shot_free_Real: 0
+
+ pirender_Fake: 1
+ pirender_Real: 0
+
+ tpsm_Fake: 1
+ tpsm_Real: 0
+
+ facedancer_Fake: 1
+ facedancer_Real: 0
+
+ facevid2vid_Fake: 1
+ facevid2vid_Real: 0
+
+ mcnet_Fake: 1
+ mcnet_Real: 0
+
+ mraa_Fake: 1
+ mraa_Real: 0
+
+ fsgan_Fake: 1
+ fsgan_Real: 0
+
+ lia_Fake: 1
+ lia_Real: 0
+
+ inswap_Fake: 1
+ inswap_Real: 0
+
+ simswap_Fake: 1
+ simswap_Real: 0
+
+ sadtalker_Fake: 1
+ sadtalker_Real: 0
+
+ wav2lip_Fake: 1
+ wav2lip_Real: 0
+
+ uniface_Fake: 1
+ uniface_Real: 0
+
+ blendface_Fake: 1
+ blendface_Real: 0
+
+ mobileswap_Fake: 1
+ mobileswap_Real: 0
+
+ faceswap_Fake: 1
+ faceswap_Real: 0
+
+ dalle2_face_Fake: 1
+ dalle2_face_Real: 0
+
+ MidJourney_Fake: 1
+ MidJourney_Real: 0
+
+ heygen_Fake: 1
+ heygen_Real: 0
+
+ whichisreal_Fake: 1
+ whichisreal_Real: 0
+
+ StyleGAN2_Fake: 1
+ StyleGAN2_Real: 0
+
+ StyleGAN3_Fake: 1
+ StyleGAN3_Real: 0
+
+ StyleGANXL_Fake: 1
+ StyleGANXL_Real: 0
+
+ ddim_Fake: 1
+ ddim_Real: 0
+
+ DiT_Fake: 1
+ DiT_Real: 0
+
+ pixart_Fake: 1
+ pixart_Real: 0
+
+ SiT_Fake: 1
+ SiT_Real: 0
+
+ sd1.5_Fake: 1
+ sd1.5_Real: 0
+
+ sd2.1_Fake: 1
+ sd2.1_Real: 0
+
+ VQGAN_Fake: 1
+ VQGAN_Real: 0
+
+ rddm_Fake: 1
+ rddm_Real: 0
+
+ pixart_Fake: 1
+ pixart_Real: 0
+
+ stargan_Fake: 1
+ stargan_Real: 0
+
+ starganv2_Fake: 1
+ starganv2_Real: 0
+
+ styleclip_Fake: 1
+ styleclip_Real: 0
+
+ deepfacelab_Fake: 1
+ deepfacelab_Real: 0
+
+ CollabDiff_Fake: 1
+ CollabDiff_Real: 0
+
+ e4e_Fake: 1
+ e4e_Real: 0
+
+ # GenImage
+ adm_Fake: 1
+ adm_Real: 0
+
+ biggan_Fake: 2
+ biggan_Real: 0
+
+ glide_Fake: 3
+ glide_Real: 0
+
+ midjourney_Fake: 4
+ midjourney_Real: 0
+
+ sdv4_Fake: 5
+ sdv4_Real: 0
+
+ sdv5_Fake: 6
+ sdv5_Real: 0
+
+ vqdm_Fake: 7
+ vqdm_Real: 0
+
+ wukong_Fake: 8
+ wukong_Real: 0
+
+ GPA_fake: 1
+ GPA_real: 0
+
+ Chameleon_Fake: 1
+ Chameleon_Real: 0
diff --git a/clean/image/effort/DeepfakeBench/training/demo.py b/clean/image/effort/DeepfakeBench/training/demo.py
new file mode 100644
index 0000000000000000000000000000000000000000..d4ad5efe45e3459883d5f2b07d36abf830bf593b
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/demo.py
@@ -0,0 +1,268 @@
+import numpy as np
+import cv2
+import random
+import yaml
+import pickle
+from tqdm import tqdm
+from PIL import Image as pil_image
+import dlib
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.nn.functional as F
+import torch.utils.data
+from torchvision import transforms
+from trainer.trainer import Trainer
+from detectors import DETECTOR
+from collections import defaultdict
+from PIL import Image as pil_image
+from imutils import face_utils
+from skimage import transform as trans
+import torchvision.transforms as T
+import os
+from os.path import join
+from typing import Tuple, List
+from pathlib import Path
+
+"""
+Usage:
+ python infer.py \
+ --detector_config ./training/config/detector/effort.yaml \
+ --weights ../../DeepfakeBenchv2/training/weights/easy_clipl14_cdf.pth \
+ --image ./id9_id6_0009.jpg \
+ --landmark_model ../../DeepfakeBenchv2/preprocessing/dlib_tools/shape_predictor_81_face_landmarks.dat
+"""
+
+import argparse
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+@torch.no_grad()
+def inference(model, data_dict):
+ data, label = data_dict['image'], data_dict['label']
+ # move data to GPU
+ data_dict['image'], data_dict['label'] = data.to(device), label.to(device)
+ predictions = model(data_dict, inference=True)
+ return predictions
+
+
+# preprocess the input image --> cropped face, resize = 256, adding a dimension of batch (output shape: 1x3x256x256)
+def get_keypts(image, face, predictor, face_detector):
+ # detect the facial landmarks for the selected face
+ shape = predictor(image, face)
+
+ # select the key points for the eyes, nose, and mouth
+ leye = np.array([shape.part(37).x, shape.part(37).y]).reshape(-1, 2)
+ reye = np.array([shape.part(44).x, shape.part(44).y]).reshape(-1, 2)
+ nose = np.array([shape.part(30).x, shape.part(30).y]).reshape(-1, 2)
+ lmouth = np.array([shape.part(49).x, shape.part(49).y]).reshape(-1, 2)
+ rmouth = np.array([shape.part(55).x, shape.part(55).y]).reshape(-1, 2)
+
+ pts = np.concatenate([leye, reye, nose, lmouth, rmouth], axis=0)
+
+ return pts
+
+def extract_aligned_face_dlib(face_detector, predictor, image, res=224, mask=None):
+ def img_align_crop(img, landmark=None, outsize=None, scale=1.3, mask=None):
+ """
+ align and crop the face according to the given bbox and landmarks
+ landmark: 5 key points
+ """
+
+ M = None
+ target_size = [112, 112]
+ dst = np.array([
+ [30.2946, 51.6963],
+ [65.5318, 51.5014],
+ [48.0252, 71.7366],
+ [33.5493, 92.3655],
+ [62.7299, 92.2041]], dtype=np.float32)
+
+ if target_size[1] == 112:
+ dst[:, 0] += 8.0
+
+ dst[:, 0] = dst[:, 0] * outsize[0] / target_size[0]
+ dst[:, 1] = dst[:, 1] * outsize[1] / target_size[1]
+
+ target_size = outsize
+
+ margin_rate = scale - 1
+ x_margin = target_size[0] * margin_rate / 2.
+ y_margin = target_size[1] * margin_rate / 2.
+
+ # move
+ dst[:, 0] += x_margin
+ dst[:, 1] += y_margin
+
+ # resize
+ dst[:, 0] *= target_size[0] / (target_size[0] + 2 * x_margin)
+ dst[:, 1] *= target_size[1] / (target_size[1] + 2 * y_margin)
+
+ src = landmark.astype(np.float32)
+
+ # use skimage tranformation
+ tform = trans.SimilarityTransform()
+ tform.estimate(src, dst)
+ M = tform.params[0:2, :]
+
+ # M: use opencv
+ # M = cv2.getAffineTransform(src[[0,1,2],:],dst[[0,1,2],:])
+
+ img = cv2.warpAffine(img, M, (target_size[1], target_size[0]))
+
+ if outsize is not None:
+ img = cv2.resize(img, (outsize[1], outsize[0]))
+
+ if mask is not None:
+ mask = cv2.warpAffine(mask, M, (target_size[1], target_size[0]))
+ mask = cv2.resize(mask, (outsize[1], outsize[0]))
+ return img, mask
+ else:
+ return img
+
+ # Image size
+ height, width = image.shape[:2]
+
+ # Convert to rgb
+ rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+ # Detect with dlib
+ faces = face_detector(rgb, 1)
+ if len(faces):
+ # For now only take the biggest face
+ face = max(faces, key=lambda rect: rect.width() * rect.height())
+
+ # Get the landmarks/parts for the face in box d only with the five key points
+ landmarks = get_keypts(rgb, face, predictor, face_detector)
+
+ # Align and crop the face
+ cropped_face = img_align_crop(rgb, landmarks, outsize=(res, res), mask=mask)
+ cropped_face = cv2.cvtColor(cropped_face, cv2.COLOR_RGB2BGR)
+
+ # Extract the all landmarks from the aligned face
+ face_align = face_detector(cropped_face, 1)
+ landmark = predictor(cropped_face, face_align[0])
+ landmark = face_utils.shape_to_np(landmark)
+
+ return cropped_face, landmark,face
+
+ else:
+ return None, None
+
+
+def load_detector(detector_cfg: str, weights: str):
+ with open(detector_cfg, "r") as f:
+ cfg = yaml.safe_load(f)
+
+ model_cls = DETECTOR[cfg["model_name"]]
+ model = model_cls(cfg).to(device)
+
+ ckpt = torch.load(weights, map_location=device)
+ state = ckpt.get("state_dict", ckpt)
+ state = {k.replace("module.", ""): v for k, v in state.items()}
+ model.load_state_dict(state, strict=False) # FIXME ⚠
+ model.eval()
+ print("[✓] Detector loaded.")
+ return model
+
+
+def preprocess_face(img_bgr: np.ndarray):
+ """BGR → tensor"""
+ img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+ img_rgb = cv2.resize(img_rgb, (224, 224), interpolation=cv2.INTER_LINEAR)
+ transform = T.Compose([
+ T.ToTensor(),
+ T.Normalize([0.48145466, 0.4578275, 0.40821073],
+ [0.26862954, 0.26130258, 0.27577711]),
+ ])
+ return transform(pil_image.fromarray(img_rgb)).unsqueeze(0) # 1×3×H×W
+
+
+@torch.inference_mode()
+def infer_single_image(
+ img_bgr: np.ndarray,
+ face_detector,
+ landmark_predictor,
+ model,
+) -> Tuple[int, float]:
+ """Return (cls_out, prob)"""
+ if face_detector is None or landmark_predictor is None:
+ face_aligned = img_bgr
+ else:
+ face_aligned, _, _ = extract_aligned_face_dlib(
+ face_detector, landmark_predictor, img_bgr, res=224
+ )
+
+ face_tensor = preprocess_face(face_aligned).to(device)
+ data = {"image": face_tensor, "label": torch.tensor([0]).to(device)}
+ preds = inference(model, data)
+ cls_out = preds["cls"].squeeze().cpu().numpy() # 0/1
+ prob = preds["prob"].squeeze().cpu().numpy() # prob
+ return cls_out, prob
+
+
+IMG_EXTS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
+def collect_image_paths(path_str: str) -> List[Path]:
+ p = Path(path_str)
+ if not p.exists():
+ raise FileNotFoundError(f"[Error] Path does not exist: {path_str}")
+
+ if p.is_file():
+ if p.suffix.lower() not in IMG_EXTS:
+ raise ValueError(f"[Error] Invalid image format: {p.name}")
+ return [p]
+
+ img_list = [fp for fp in p.iterdir() if fp.is_file() and fp.suffix.lower() in IMG_EXTS]
+ if not img_list:
+ raise RuntimeError(f"[Error] No valid image files found in directory: {path_str}")
+
+ return sorted(img_list)
+
+
+def parse_args():
+ p = argparse.ArgumentParser(
+ description="Deepfake image inference (single image version)"
+ )
+ p.add_argument("--detector_config", default='training/config/detector/effort.yaml',
+ help="YAML 配置文件路径")
+ p.add_argument("--weights", required=True,
+ help="Detector 预训练权重")
+ p.add_argument("--image", required=True,
+ help="tested image")
+ p.add_argument("--landmark_model", default=False,
+ help="dlib 81 landmarks dat 文件 / 如果不需要裁剪人脸就是False")
+ return p.parse_args()
+
+
+def main():
+ args = parse_args()
+
+ model = load_detector(args.detector_config, args.weights)
+ if args.landmark_model:
+ face_det = dlib.get_frontal_face_detector()
+ shape_predictor = dlib.shape_predictor(args.landmark_model)
+ else:
+ face_det, shape_predictor = None, None
+
+ img_paths = collect_image_paths(args.image)
+ multiple = len(img_paths) > 1
+ if multiple:
+ print(f"Collected {len(img_paths)} images in total，let's infer them...\n")
+
+ # ---------- infer ----------
+ for idx, img_path in enumerate(img_paths, 1):
+ img = cv2.imread(str(img_path))
+ if img is None:
+ print(f"[Warning] loading wrong，skip: {img_path}", file=sys.stderr)
+ continue
+
+ cls, prob = infer_single_image(img, face_det, shape_predictor, model)
+ print(
+ f"[{idx}/{len(img_paths)}] {img_path.name:>30} | Pred Label: {cls} "
+ f"(0=Real, 1=Fake) | Fake Prob: {prob:.4f}"
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/clean/image/effort/DeepfakeBench/training/detectors/__init__.py b/clean/image/effort/DeepfakeBench/training/detectors/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..95452704d19f9f74d793b3df6031b8c82d30c1a3
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/detectors/__init__.py
@@ -0,0 +1,10 @@
+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 utils.registry import DETECTOR
+from .effort_detector import EffortDetector
diff --git a/clean/image/effort/DeepfakeBench/training/detectors/base_detector.py b/clean/image/effort/DeepfakeBench/training/detectors/base_detector.py
new file mode 100644
index 0000000000000000000000000000000000000000..935d102215d37f735419ce32d5b1a8b50b817531
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/detectors/base_detector.py
@@ -0,0 +1,78 @@
+# 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
+
+ @abc.abstractmethod
+ def get_test_metrics(self):
+ """
+ Returns the testing metrics for the model.
+ """
+ pass
diff --git a/clean/image/effort/DeepfakeBench/training/detectors/effort_detector.py b/clean/image/effort/DeepfakeBench/training/detectors/effort_detector.py
new file mode 100644
index 0000000000000000000000000000000000000000..c9784ba965addaf69b094d4c8786267bdd218e3c
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/detectors/effort_detector.py
@@ -0,0 +1,347 @@
+import os
+import math
+import datetime
+import logging
+import numpy as np
+from sklearn import metrics
+from typing import Union
+from collections import defaultdict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.nn import DataParallel
+from torch.utils.tensorboard import SummaryWriter
+
+from metrics.base_metrics_class import calculate_metrics_for_train
+
+from .base_detector import AbstractDetector
+from detectors import DETECTOR
+from networks import BACKBONE
+from loss import LOSSFUNC
+
+import loralib as lora
+from transformers import AutoProcessor, CLIPModel, ViTModel, ViTConfig
+
+logger = logging.getLogger(__name__)
+
+
+@DETECTOR.register_module(module_name='effort')
+class EffortDetector(nn.Module):
+ def __init__(self, config=None):
+ super(EffortDetector, self).__init__()
+ self.config = config
+ self.backbone = self.build_backbone(config)
+ self.head = nn.Linear(1024, 2)
+ self.loss_func = nn.CrossEntropyLoss()
+ self.prob, self.label = [], []
+ self.correct, self.total = 0, 0
+
+ def build_backbone(self, config):
+ # ⚠⚠⚠ Download CLIP model using the below link
+ # https://drive.google.com/drive/folders/1fm3Jd8lFMiSP1qgdmsxfqlJZGpr_bXsx?usp=drive_link
+
+ # mean: [0.48145466, 0.4578275, 0.40821073]
+ # std: [0.26862954, 0.26130258, 0.27577711]
+
+ # ViT-L/14 224*224
+ clip_model = CLIPModel.from_pretrained("../models--openai--clip-vit-large-patch14") # the path of this folder in your disk (download from the above link)
+
+ # Apply SVD to self_attn layers only
+ # ViT-L/14 224*224: 1024-1
+ clip_model.vision_model = apply_svd_residual_to_self_attn(clip_model.vision_model, r=1024-1)
+
+ #for name, param in clip_model.vision_model.named_parameters():
+ # print('{}: {}'.format(name, param.requires_grad))
+ #num_param = sum(p.numel() for p in clip_model.vision_model.parameters() if p.requires_grad)
+ #num_total_param = sum(p.numel() for p in clip_model.vision_model.parameters())
+ #print('Number of total parameters: {}, tunable parameters: {}'.format(num_total_param, num_param))
+
+ return clip_model.vision_model
+
+ def features(self, data_dict: dict) -> torch.tensor:
+ feat = self.backbone(data_dict['image'])['pooler_output']
+ return feat
+
+ def classifier(self, features: torch.tensor) -> torch.tensor:
+ return self.head(features)
+
+ #def get_losses(self, data_dict: dict, pred_dict: dict) -> dict:
+ # label = data_dict['label']
+ # pred = pred_dict['cls']
+ # loss = self.loss_func(pred, label)
+ #
+ # if self.training:
+ # # Regularization term
+ # lambda_reg = 1.0
+ # reg_term = 0.0
+ # num_reg = 0
+ # for module in self.backbone.modules():
+ # if isinstance(module, SVDResidualLinear):
+ # reg_term += module.compute_orthogonal_loss()
+ # reg_term += module.compute_keepsv_loss()
+ # num_reg += 1
+ #
+ # loss += lambda_reg * reg_term / num_reg
+ #
+ # loss_dict = {'overall': loss}
+ # return loss_dict
+
+ def compute_weight_loss(self):
+ weight_sum_dict = {}
+ num_weight_dict = {}
+ for name, module in self.backbone.named_modules():
+ if isinstance(module, SVDResidualLinear):
+ weight_curr = module.compute_current_weight()
+ if str(weight_curr.size()) not in weight_sum_dict.keys():
+ weight_sum_dict[str(weight_curr.size())] = weight_curr
+ num_weight_dict[str(weight_curr.size())] = 1
+ else:
+ weight_sum_dict[str(weight_curr.size())] += weight_curr
+ num_weight_dict[str(weight_curr.size())] += 1
+
+ loss2 = 0.0
+ for k in weight_sum_dict.keys():
+ _, S_sum, _ = torch.linalg.svd(weight_sum_dict[k], full_matrices=False)
+ loss2 += -torch.mean(S_sum)
+ loss2 /= len(weight_sum_dict.keys())
+ return loss2
+
+ def get_losses(self, data_dict: dict, pred_dict: dict) -> dict:
+ label = data_dict['label'] # Tensor of shape [batch_size]
+ pred = pred_dict['cls'] # Tensor of shape [batch_size, num_classes]
+
+ # Compute overall loss using all samples
+ loss = self.loss_func(pred, label)
+
+ # Create masks for real and fake classes
+ mask_real = label == 0 # Boolean tensor
+ mask_fake = label == 1 # Boolean tensor
+
+ # Compute loss for real class
+ if mask_real.sum() > 0:
+ pred_real = pred[mask_real]
+ label_real = label[mask_real]
+ loss_real = self.loss_func(pred_real, label_real)
+ else:
+ # No real samples in batch
+ loss_real = torch.tensor(0.0, device=pred.device)
+
+ # Compute loss for fake class
+ if mask_fake.sum() > 0:
+ pred_fake = pred[mask_fake]
+ label_fake = label[mask_fake]
+ loss_fake = self.loss_func(pred_fake, label_fake)
+ else:
+ # No fake samples in batch
+ loss_fake = torch.tensor(0.0, device=pred.device)
+
+
+ # loss2 = self.compute_weight_loss()
+ # overall_loss = loss + loss2
+
+ # Return a dictionary with all losses
+ loss_dict = {
+ 'overall': loss,
+ 'real_loss': loss_real,
+ 'fake_loss': loss_fake,
+ # 'erank_loss': loss2
+ }
+ return loss_dict
+
+ def get_train_metrics(self, data_dict: dict, pred_dict: dict) -> dict:
+ label = data_dict['label']
+ pred = pred_dict['cls']
+ # compute metrics for batch data
+ auc, eer, acc, ap = calculate_metrics_for_train(label.detach(), pred.detach())
+ metric_batch_dict = {'acc': acc, 'auc': auc, 'eer': eer, 'ap': ap}
+ return metric_batch_dict
+
+ def forward(self, data_dict: dict, inference=False) -> dict:
+ # get the features by backbone
+ features = self.features(data_dict)
+ # get the prediction by classifier
+ pred = self.classifier(features)
+ # get the probability of the pred
+ prob = torch.softmax(pred, dim=1)[:, 1]
+ # build the prediction dict for each output
+ pred_dict = {'cls': pred, 'prob': prob, 'feat': features}
+
+ return pred_dict
+
+
+# Custom module to represent the residual using SVD components
+class SVDResidualLinear(nn.Module):
+ def __init__(self, in_features, out_features, r, bias=True, init_weight=None):
+ super(SVDResidualLinear, self).__init__()
+ self.in_features = in_features
+ self.out_features = out_features
+ self.r = r # Number of top singular values to exclude
+
+ # Original weights (fixed)
+ self.weight_main = nn.Parameter(torch.Tensor(out_features, in_features), requires_grad=False)
+ if init_weight is not None:
+ self.weight_main.data.copy_(init_weight)
+ else:
+ nn.init.kaiming_uniform_(self.weight_main, a=math.sqrt(5))
+
+ # Bias
+ if bias:
+ self.bias = nn.Parameter(torch.Tensor(out_features))
+ nn.init.zeros_(self.bias)
+ else:
+ self.register_parameter('bias', None)
+
+ def compute_current_weight(self):
+ if self.S_residual is not None:
+ return self.weight_main + self.U_residual @ torch.diag(self.S_residual) @ self.V_residual
+ else:
+ return self.weight_main
+
+ def forward(self, x):
+ if hasattr(self, 'U_residual') and hasattr(self, 'V_residual') and self.S_residual is not None:
+ # Reconstruct the residual weight
+ residual_weight = self.U_residual @ torch.diag(self.S_residual) @ self.V_residual
+ # Total weight is the fixed main weight plus the residual
+ weight = self.weight_main + residual_weight
+ else:
+ # If residual components are not set, use only the main weight
+ weight = self.weight_main
+
+ return F.linear(x, weight, self.bias)
+
+ def compute_orthogonal_loss(self):
+ if self.S_residual is not None:
+ # According to the properties of orthogonal matrices: A^TA = I
+ UUT = torch.cat((self.U_r, self.U_residual), dim=1) @ torch.cat((self.U_r, self.U_residual), dim=1).t()
+ VVT = torch.cat((self.V_r, self.V_residual), dim=0) @ torch.cat((self.V_r, self.V_residual), dim=0).t()
+ # print(self.U_r.size(), self.U_residual.size()) # torch.Size([1024, 1023]) torch.Size([1024, 1])
+ # print(self.V_r.size(), self.V_residual.size()) # torch.Size([1023, 1024]) torch.Size([1, 1024])
+ # UUT = self.U_residual @ self.U_residual.t()
+ # VVT = self.V_residual @ self.V_residual.t()
+
+ # Construct an identity matrix
+ UUT_identity = torch.eye(UUT.size(0), device=UUT.device)
+ VVT_identity = torch.eye(VVT.size(0), device=VVT.device)
+
+ # Using frobenius norm to compute loss
+ loss = 0.5 * torch.norm(UUT - UUT_identity, p='fro') + 0.5 * torch.norm(VVT - VVT_identity, p='fro')
+ else:
+ loss = 0.0
+
+ return loss
+
+ def compute_keepsv_loss(self):
+ if (self.S_residual is not None) and (self.weight_original_fnorm is not None):
+ # Total current weight is the fixed main weight plus the residual
+ weight_current = self.weight_main + self.U_residual @ torch.diag(self.S_residual) @ self.V_residual
+ # Frobenius norm of current weight
+ weight_current_fnorm = torch.norm(weight_current, p='fro')
+
+ loss = torch.abs(weight_current_fnorm ** 2 - self.weight_original_fnorm ** 2)
+ # loss = torch.abs(weight_current_fnorm ** 2 + 0.01 * self.weight_main_fnorm ** 2 - 1.01 * self.weight_original_fnorm ** 2)
+ else:
+ loss = 0.0
+
+ return loss
+
+ def compute_fn_loss(self):
+ if (self.S_residual is not None):
+ weight_current = self.weight_main + self.U_residual @ torch.diag(self.S_residual) @ self.V_residual
+ weight_current_fnorm = torch.norm(weight_current, p='fro')
+
+ loss = weight_current_fnorm ** 2
+ else:
+ loss = 0.0
+
+ return loss
+
+# Function to replace nn.Linear modules within self_attn modules with SVDResidualLinear
+def apply_svd_residual_to_self_attn(model, r):
+ for name, module in model.named_children():
+ if 'self_attn' in name:
+ # Replace nn.Linear layers in this module
+ for sub_name, sub_module in module.named_modules():
+ if isinstance(sub_module, nn.Linear):
+ # Get parent module within self_attn
+ parent_module = module
+ sub_module_names = sub_name.split('.')
+ for module_name in sub_module_names[:-1]:
+ parent_module = getattr(parent_module, module_name)
+ # Replace the nn.Linear layer with SVDResidualLinear
+ setattr(parent_module, sub_module_names[-1], replace_with_svd_residual(sub_module, r))
+ else:
+ # Recursively apply to child modules
+ apply_svd_residual_to_self_attn(module, r)
+ # After replacing, set requires_grad for residual components
+ for param_name, param in model.named_parameters():
+ if any(x in param_name for x in ['S_residual', 'U_residual', 'V_residual']):
+ param.requires_grad = True
+ else:
+ param.requires_grad = False
+ return model
+
+
+# Function to replace a module with SVDResidualLinear
+def replace_with_svd_residual(module, r):
+ if isinstance(module, nn.Linear):
+ in_features = module.in_features
+ out_features = module.out_features
+ bias = module.bias is not None
+
+ # Create SVDResidualLinear module
+ new_module = SVDResidualLinear(in_features, out_features, r, bias=bias, init_weight=module.weight.data.clone())
+
+ if bias and module.bias is not None:
+ new_module.bias.data.copy_(module.bias.data)
+
+ new_module.weight_original_fnorm = torch.norm(module.weight.data, p='fro')
+
+ # Perform SVD on the original weight
+ _svd_device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+ U, S, Vh = torch.linalg.svd(module.weight.data.to(_svd_device), full_matrices=False)
+ U, S, Vh = U.cpu(), S.cpu(), Vh.cpu()
+
+ # Determine r based on the rank of the weight matrix
+ r = min(r, len(S)) # Ensure r does not exceed the number of singular values
+
+ # Keep top r singular components (main weight)
+ U_r = U[:, :r] # Shape: (out_features, r)
+ S_r = S[:r] # Shape: (r,)
+ Vh_r = Vh[:r, :] # Shape: (r, in_features)
+
+ # Reconstruct the main weight (fixed)
+ weight_main = U_r @ torch.diag(S_r) @ Vh_r
+
+ # Calculate the frobenius norm of main weight
+ new_module.weight_main_fnorm = torch.norm(weight_main.data, p='fro')
+
+ # Set the main weight
+ new_module.weight_main.data.copy_(weight_main)
+
+ # Residual components (trainable)
+ U_residual = U[:, r:] # Shape: (out_features, n - r)
+ S_residual = S[r:] # Shape: (n - r,)
+ Vh_residual = Vh[r:, :] # Shape: (n - r, in_features)
+
+ if len(S_residual) > 0:
+ new_module.S_residual = nn.Parameter(S_residual.clone())
+ new_module.U_residual = nn.Parameter(U_residual.clone())
+ new_module.V_residual = nn.Parameter(Vh_residual.clone())
+
+ new_module.S_r = nn.Parameter(S_r.clone(), requires_grad=False)
+ new_module.U_r = nn.Parameter(U_r.clone(), requires_grad=False)
+ new_module.V_r = nn.Parameter(Vh_r.clone(), requires_grad=False)
+ else:
+ new_module.S_residual = None
+ new_module.U_residual = None
+ new_module.V_residual = None
+
+ new_module.S_r = None
+ new_module.U_r = None
+ new_module.V_r = None
+
+ return new_module
+ else:
+ return module
diff --git a/clean/image/effort/DeepfakeBench/training/logger.py b/clean/image/effort/DeepfakeBench/training/logger.py
new file mode 100644
index 0000000000000000000000000000000000000000..831c8679c128d91740d1ca3fb9fdae56dfe1b1b4
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/logger.py
@@ -0,0 +1,27 @@
+import os
+import logging
+
+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
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/loss/__init__.py b/clean/image/effort/DeepfakeBench/training/loss/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d25bdf89a79b7a322c9ad4a1bdfa489d39ee4310
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/loss/__init__.py
@@ -0,0 +1,12 @@
+from utils.registry import LOSSFUNC
+from .cross_entropy_loss import CrossEntropyLoss
+from .consistency_loss import ConsistencyCos
+from .capsule_loss import CapsuleLoss
+from .bce_loss import BCELoss
+from .am_softmax import AMSoftmaxLoss
+from .am_softmax import AMSoftmax_OHEM
+from .contrastive_regularization import ContrastiveLoss
+from .l1_loss import L1Loss
+from .id_loss import IDLoss
+from .vgg_loss import VGGLoss
+from .js_loss import JS_Loss
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/loss/abstract_loss_func.py b/clean/image/effort/DeepfakeBench/training/loss/abstract_loss_func.py
new file mode 100644
index 0000000000000000000000000000000000000000..45d3324ed53be4310867b326e9eaabd265634138
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/loss/abstract_loss_func.py
@@ -0,0 +1,17 @@
+import torch.nn as nn
+
+class AbstractLossClass(nn.Module):
+ """Abstract class for loss functions."""
+ def __init__(self):
+ super(AbstractLossClass, self).__init__()
+
+ def forward(self, pred, label):
+ """
+ Args:
+ pred: prediction of the model
+ label: ground truth label
+
+ Return:
+ loss: loss value
+ """
+ raise NotImplementedError('Each subclass should implement the forward method.')
diff --git a/clean/image/effort/DeepfakeBench/training/loss/am_softmax.py b/clean/image/effort/DeepfakeBench/training/loss/am_softmax.py
new file mode 100644
index 0000000000000000000000000000000000000000..8a43a4287f7c036392c69c30e1a518a7710b41e1
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/loss/am_softmax.py
@@ -0,0 +1,145 @@
+"""
+ Copyright (c) 2018 Intel Corporation
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+ http://www.apache.org/licenses/LICENSE-2.0
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import Parameter
+import torch as th
+
+from loss.abstract_loss_func import AbstractLossClass
+from utils.registry import LOSSFUNC
+
+
+#------------ AMSoftmax Loss ----------------------
+
+def focal_loss(input_values, gamma):
+ """Computes the focal loss"""
+ p = torch.exp(-input_values)
+ loss = (1 - p) ** gamma * input_values
+ return loss.mean()
+
+
+@LOSSFUNC.register_module(module_name="am_softmax")
+class AMSoftmaxLoss(AbstractLossClass):
+ """Computes the AM-Softmax loss with cos or arc margin"""
+ margin_types = ['cos', 'arc']
+
+ def __init__(self, margin_type='cos', gamma=0., m=0.5, s=30, t=1.):
+ super().__init__()
+ assert margin_type in AMSoftmaxLoss.margin_types
+ self.margin_type = margin_type
+ assert gamma >= 0
+ self.gamma = gamma
+ assert m > 0
+ self.m = m
+ assert s > 0
+ self.s = s
+ self.cos_m = math.cos(m)
+ self.sin_m = math.sin(m)
+ self.th = math.cos(math.pi - m)
+ assert t >= 1
+ self.t = t
+
+ def forward(self, cos_theta, target):
+ if self.margin_type == 'cos':
+ phi_theta = cos_theta - self.m
+ else:
+ sine = torch.sqrt(1.0 - torch.pow(cos_theta, 2))
+ phi_theta = cos_theta * self.cos_m - sine * self.sin_m #cos(theta+m)
+ phi_theta = torch.where(cos_theta > self.th, phi_theta, cos_theta - self.sin_m * self.m)
+
+ index = torch.zeros_like(cos_theta, dtype=torch.uint8)
+ index.scatter_(1, target.data.view(-1, 1), 1)
+ output = torch.where(index, phi_theta, cos_theta)
+
+ if self.gamma == 0 and self.t == 1.:
+ return F.cross_entropy(self.s*output, target)
+
+ if self.t > 1:
+ h_theta = self.t - 1 + self.t*cos_theta
+ support_vecs_mask = (1 - index) * \
+ torch.lt(torch.masked_select(phi_theta, index).view(-1, 1).repeat(1, h_theta.shape[1]) - cos_theta, 0)
+ output = torch.where(support_vecs_mask, h_theta, output)
+ return F.cross_entropy(self.s*output, target)
+
+ return focal_loss(F.cross_entropy(self.s*output, target, reduction='none'), self.gamma)
+
+
+@LOSSFUNC.register_module(module_name="am_softmax_ohem")
+class AMSoftmax_OHEM(AbstractLossClass):
+ """Computes the AM-Softmax loss with cos or arc margin"""
+ margin_types = ['cos', 'arc']
+
+ def __init__(self, margin_type='cos', gamma=0., m=0.5, s=30, t=1., ratio=1.):
+ super(self).__init__()
+ assert margin_type in AMSoftmaxLoss.margin_types
+ self.margin_type = margin_type
+ assert gamma >= 0
+ self.gamma = gamma
+ assert m > 0
+ self.m = m
+ assert s > 0
+ self.s = s
+ self.cos_m = math.cos(m)
+ self.sin_m = math.sin(m)
+ self.th = math.cos(math.pi - m)
+ assert t >= 1
+ self.t = t
+ self.ratio = ratio
+
+
+ # ------- online hard example mining --------------------
+ def get_subidx(self,x,y,ratio):
+ num_inst = x.size(0)
+ num_hns = int(ratio * num_inst)
+ x_ = x.clone()
+ inst_losses = th.autograd.Variable(th.zeros(num_inst)).cuda()
+
+ for idx, label in enumerate(y.data):
+ inst_losses[idx] = -x_.data[idx, label]
+
+ _, idxs = inst_losses.topk(num_hns)
+ return idxs
+
+
+ def forward(self, cos_theta, target):
+ if self.margin_type == 'cos':
+ phi_theta = cos_theta - self.m
+ else:
+ sine = torch.sqrt(1.0 - torch.pow(cos_theta, 2))
+ phi_theta = cos_theta * self.cos_m - sine * self.sin_m #cos(theta+m)
+ phi_theta = torch.where(cos_theta > self.th, phi_theta, cos_theta - self.sin_m * self.m)
+
+ index = torch.zeros_like(cos_theta, dtype=torch.uint8)
+ index.scatter_(1, target.data.view(-1, 1), 1)
+ output = torch.where(index, phi_theta, cos_theta)
+
+ out = F.log_softmax(output,dim=1)
+ idxs = self.get_subidx(out,target,self.ratio) # select hard examples
+
+ output2 = output.index_select(0, idxs)
+ target2 = target.index_select(0, idxs)
+
+ if self.gamma == 0 and self.t == 1.:
+ return F.cross_entropy(self.s*output2, target2)
+
+ if self.t > 1:
+ h_theta = self.t - 1 + self.t*cos_theta
+ support_vecs_mask = (1 - index) * \
+ torch.lt(torch.masked_select(phi_theta, index).view(-1, 1).repeat(1, h_theta.shape[1]) - cos_theta, 0)
+ output2 = torch.where(support_vecs_mask, h_theta, output2)
+ return F.cross_entropy(self.s*output2, target2)
+
+ return focal_loss(F.cross_entropy(self.s*output2, target2, reduction='none'), self.gamma)
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/loss/bce_loss.py b/clean/image/effort/DeepfakeBench/training/loss/bce_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..4d71b5e6e4b2b9d6bc25fdd9262d9f66d475e619
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/loss/bce_loss.py
@@ -0,0 +1,26 @@
+import torch.nn as nn
+from loss.abstract_loss_func import AbstractLossClass
+from utils.registry import LOSSFUNC
+
+
+@LOSSFUNC.register_module(module_name="bce")
+class BCELoss(AbstractLossClass):
+ def __init__(self):
+ super().__init__()
+ self.loss_fn = nn.BCELoss()
+
+ def forward(self, inputs, targets):
+ """
+ Computes the bce loss.
+
+ Args:
+ inputs: A PyTorch tensor of size (batch_size, num_classes) containing the predicted scores.
+ targets: A PyTorch tensor of size (batch_size) containing the ground-truth class indices.
+
+ Returns:
+ A scalar tensor representing the bce loss.
+ """
+ # Compute the bce loss
+ loss = self.loss_fn(inputs, targets.float())
+
+ return loss
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/loss/capsule_loss.py b/clean/image/effort/DeepfakeBench/training/loss/capsule_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..7e2adad9ced2d1eade1ad80bf7b1cd84cf60ea9b
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/loss/capsule_loss.py
@@ -0,0 +1,28 @@
+import torch.nn as nn
+from loss.abstract_loss_func import AbstractLossClass
+from utils.registry import LOSSFUNC
+
+
+@LOSSFUNC.register_module(module_name="capsule_loss")
+class CapsuleLoss(AbstractLossClass):
+ def __init__(self):
+ super().__init__()
+ self.cross_entropy_loss = nn.CrossEntropyLoss()
+
+ def forward(self, inputs, targets):
+ """
+ Computes the capsule loss.
+
+ Args:
+ inputs: A PyTorch tensor of size (batch_size, num_classes) containing the predicted scores.
+ targets: A PyTorch tensor of size (batch_size) containing the ground-truth class indices.
+
+ Returns:
+ A scalar tensor representing the capsule loss.
+ """
+ # Compute the capsule loss
+ loss_t = self.cross_entropy_loss(inputs[:,0,:], targets)
+
+ for i in range(inputs.size(1) - 1):
+ loss_t = loss_t + self.cross_entropy_loss(inputs[:,i+1,:], targets)
+ return loss_t
diff --git a/clean/image/effort/DeepfakeBench/training/loss/classNseg_loss.py b/clean/image/effort/DeepfakeBench/training/loss/classNseg_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..4f81f2b577805e5d45f88ccf2ce21c7cd73d3802
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/loss/classNseg_loss.py
@@ -0,0 +1,64 @@
+import torch
+from torch import nn
+from loss.abstract_loss_func import AbstractLossClass
+from utils.registry import LOSSFUNC
+
+
+@LOSSFUNC.register_module(module_name="classNseg_loss")
+class ClassNsegLoss(AbstractLossClass):
+ def __init__(self):
+ super().__init__()
+ self.gamma = 1.0 # weight decay. default=5.0
+ self.act_loss_fn = ActivationLoss()
+ self.rect_loss_fn = ReconstructionLoss()
+ self.seg_loss_fn = SegmentationLoss()
+
+ def forward(self, inputs, targets):
+
+ zero_abs = torch.abs(latent[:,0]).view(latent.shape[0], -1)
+ zero = zero_abs.mean(dim=1)
+
+ one_abs = torch.abs(latent[:,1]).view(latent.shape[0], -1)
+ one = one_abs.mean(dim=1)
+
+ loss_act = self.act_loss_fn(zero, one, labels_data)
+ loss_act_data = loss_act.item()
+
+ loss_seg = self.seg_loss_fn(seg, mask)
+ loss_seg = loss_seg * self.gamma
+ loss_seg_data = loss_seg.item()
+
+ loss_rect = self.rect_loss_fn(rect, rgb)
+ loss_rect = loss_rect * self.gamma
+ loss_rect_data = loss_rect.item()
+ loss_total = loss_act + loss_seg + loss_rect
+ return loss_total, loss_act_data, loss_seg_data, loss_rect_data
+
+
+class ActivationLoss(nn.Module):
+ def __init__(self):
+ super(ActivationLoss, self).__init__()
+
+ def forward(self, zero, one, labels):
+
+ loss_act = torch.abs(one - labels.data) + torch.abs(zero - (1.0 - labels.data))
+ return 1 / labels.shape[0] * loss_act.sum()
+
+class ReconstructionLoss(nn.Module):
+ def __init__(self):
+ super(ReconstructionLoss, self).__init__()
+ self.loss = nn.MSELoss()
+
+ def forward(self, reconstruction, groundtruth):
+
+ return self.loss(reconstruction, groundtruth.data)
+
+class SegmentationLoss(nn.Module):
+ def __init__(self):
+ super(SegmentationLoss, self).__init__()
+ self.loss = nn.CrossEntropyLoss()
+
+ def forward(self, segment, groundtruth):
+
+ return self.loss(segment.view(segment.shape[0], segment.shape[1], segment.shape[2] * segment.shape[3]),
+ groundtruth.data.view(groundtruth.shape[0], groundtruth.shape[1] * groundtruth.shape[2]))
diff --git a/clean/image/effort/DeepfakeBench/training/loss/consistency_loss.py b/clean/image/effort/DeepfakeBench/training/loss/consistency_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..50e666aad8bbdd659c136d5506d1dbd2a7675685
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/loss/consistency_loss.py
@@ -0,0 +1,54 @@
+import torch.nn as nn
+import torch
+from loss.abstract_loss_func import AbstractLossClass
+from utils.registry import LOSSFUNC
+
+
+@LOSSFUNC.register_module(module_name="consistency_loss")
+class ConsistencyCos(nn.Module):
+ def __init__(self):
+ super(ConsistencyCos, self).__init__()
+ # # CrossEntropy Loss
+ # weight=torch.Tensor([4.0, 1.0])
+ # if torch.cuda.is_available():
+ # weight = weight.cuda()
+ # self.loss_fn = nn.CrossEntropyLoss(weight)
+ self.loss_fn = nn.CrossEntropyLoss()
+ self.mse_fn = nn.MSELoss()
+
+ def forward(self, feat, inputs, targets):
+ feat = nn.functional.normalize(feat, dim=1)
+ feat_0 = feat[:int(feat.size(0)/2),:]
+ feat_1 = feat[int(feat.size(0)/2): 2*int(feat.size(0)/2),:]
+
+ cos = torch.einsum('nc,nc->n', [feat_0, feat_1]).unsqueeze(-1)
+ labels = torch.ones((cos.shape[0],1), dtype=torch.float, requires_grad=False)
+ if torch.cuda.is_available():
+ labels = labels.cuda()
+ self.consistency_rate = 1.0
+ loss = self.consistency_rate * self.mse_fn(cos, labels) + self.loss_fn(inputs, targets)
+ return loss
+
+#
+##FIXME to be implemented
+class ConsistencyL2(nn.Module):
+ def __init__(self):
+ super(ConsistencyL2, self).__init__()
+ self.mse_fn = nn.MSELoss()
+
+ def forward(self, feat):
+ feat_0 = feat[:int(feat.size(0)/2),:]
+ feat_1 = feat[int(feat.size(0)/2):,:]
+ loss = self.mse_fn(feat_0, feat_1)
+ return loss
+
+class ConsistencyL1(nn.Module):
+ def __init__(self):
+ super(ConsistencyL1, self).__init__()
+ self.L1_fn = nn.L1Loss()
+
+ def forward(self, feat):
+ feat_0 = feat[:int(feat.size(0)/2),:]
+ feat_1 = feat[int(feat.size(0)/2):,:]
+ loss = self.L1_fn(feat_0, feat_1)
+ return loss
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/loss/contrastive_regularization.py b/clean/image/effort/DeepfakeBench/training/loss/contrastive_regularization.py
new file mode 100644
index 0000000000000000000000000000000000000000..e067dac81fc9befb6b8e434df8a70cadf3b9807c
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/loss/contrastive_regularization.py
@@ -0,0 +1,78 @@
+import random
+from collections import defaultdict
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from loss.abstract_loss_func import AbstractLossClass
+from utils.registry import LOSSFUNC
+
+
+def swap_spe_features(type_list, value_list):
+ type_list = type_list.cpu().numpy().tolist()
+ # get index
+ index_list = list(range(len(type_list)))
+
+ # init a dict, where its key is the type and value is the index
+ spe_dict = defaultdict(list)
+
+ # do for-loop to get spe dict
+ for i, one_type in enumerate(type_list):
+ spe_dict[one_type].append(index_list[i])
+
+ # shuffle the value list of each key
+ for keys in spe_dict.keys():
+ random.shuffle(spe_dict[keys])
+
+ # generate a new index list for the value list
+ new_index_list = []
+ for one_type in type_list:
+ value = spe_dict[one_type].pop()
+ new_index_list.append(value)
+
+ # swap the value_list by new_index_list
+ value_list_new = value_list[new_index_list]
+
+ return value_list_new
+
+
+@LOSSFUNC.register_module(module_name="contrastive_regularization")
+class ContrastiveLoss(AbstractLossClass):
+ def __init__(self, margin=1.0):
+ super().__init__()
+ self.margin = margin
+
+ def contrastive_loss(self, anchor, positive, negative):
+ dist_pos = F.pairwise_distance(anchor, positive)
+ dist_neg = F.pairwise_distance(anchor, negative)
+ # Compute loss as the distance between anchor and negative minus the distance between anchor and positive
+ loss = torch.mean(torch.clamp(dist_pos - dist_neg + self.margin, min=0.0))
+ return loss
+
+ def forward(self, common, specific, spe_label):
+ # prepare
+ bs = common.shape[0]
+ real_common, fake_common = common.chunk(2)
+ ### common real
+ idx_list = list(range(0, bs//2))
+ random.shuffle(idx_list)
+ real_common_anchor = common[idx_list]
+ ### common fake
+ idx_list = list(range(bs//2, bs))
+ random.shuffle(idx_list)
+ fake_common_anchor = common[idx_list]
+ ### specific
+ specific_anchor = swap_spe_features(spe_label, specific)
+ real_specific_anchor, fake_specific_anchor = specific_anchor.chunk(2)
+ real_specific, fake_specific = specific.chunk(2)
+
+ # Compute the contrastive loss of common between real and fake
+ loss_realcommon = self.contrastive_loss(real_common, real_common_anchor, fake_common_anchor)
+ loss_fakecommon = self.contrastive_loss(fake_common, fake_common_anchor, real_common_anchor)
+
+ # Comupte the constrastive loss of specific between real and fake
+ loss_realspecific = self.contrastive_loss(real_specific, real_specific_anchor, fake_specific_anchor)
+ loss_fakespecific = self.contrastive_loss(fake_specific, fake_specific_anchor, real_specific_anchor)
+
+ # Compute the final loss as the sum of all contrastive losses
+ loss = loss_realcommon + loss_fakecommon + loss_fakespecific + loss_realspecific
+ return loss
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/loss/cross_entropy_loss.py b/clean/image/effort/DeepfakeBench/training/loss/cross_entropy_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..bb1c11c168a7f1ac61bec91f1ed68dd77524e795
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/loss/cross_entropy_loss.py
@@ -0,0 +1,26 @@
+import torch.nn as nn
+from loss.abstract_loss_func import AbstractLossClass
+from utils.registry import LOSSFUNC
+
+
+@LOSSFUNC.register_module(module_name="cross_entropy")
+class CrossEntropyLoss(AbstractLossClass):
+ def __init__(self):
+ super().__init__()
+ self.loss_fn = nn.CrossEntropyLoss()
+
+ def forward(self, inputs, targets):
+ """
+ Computes the cross-entropy loss.
+
+ Args:
+ inputs: A PyTorch tensor of size (batch_size, num_classes) containing the predicted scores.
+ targets: A PyTorch tensor of size (batch_size) containing the ground-truth class indices.
+
+ Returns:
+ A scalar tensor representing the cross-entropy loss.
+ """
+ # Compute the cross-entropy loss
+ loss = self.loss_fn(inputs, targets)
+
+ return loss
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/loss/det_loss.py b/clean/image/effort/DeepfakeBench/training/loss/det_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..a6acfa30d8c48871c467c3fb97e2b0d88ba874b6
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/loss/det_loss.py
@@ -0,0 +1,165 @@
+# import torch.nn as nn
+# from loss.abstract_loss_func import AbstractLossClass
+# from utils.registry import LOSSFUNC
+
+
+# @LOSSFUNC.register_module(module_name="det_loss")
+# class DETLoss(AbstractLossClass):
+# def __init__(self):
+# super().__init__()
+# self.loss_fn = nn.BCELoss()
+
+# def forward(self, inputs, targets):
+# """
+# Computes the bce loss.
+
+# Args:
+# inputs: A PyTorch tensor of size (batch_size, num_classes) containing the predicted scores.
+# targets: A PyTorch tensor of size (batch_size) containing the ground-truth class indices.
+
+# Returns:
+# A scalar tensor representing the bce loss.
+# """
+# # Compute the bce loss
+# loss = self.loss_fn(inputs, targets.float())
+
+# return loss
+
+
+# # loss init
+# det_criterion = MultiBoxLoss(
+# cfg['det_loss']['num_classes'],
+# cfg['det_loss']['overlap_thresh'],
+# cfg['det_loss']['prior_for_matching'],
+# cfg['det_loss']['bkg_label'],
+# cfg['det_loss']['neg_mining'],
+# cfg['det_loss']['neg_pos'],
+# cfg['det_loss']['neg_overlap'],
+# cfg['det_loss']['encode_target'],
+# cfg['det_loss']['use_gpu']
+# )
+# criterion = nn.CrossEntropyLoss()
+
+# loss_end_cls = criterion(outputs, labels)
+# loss_l, loss_c = det_criterion(
+# (locations, confidence),
+# confidence_labels, location_labels
+# )
+# acc = sum(outputs.max(-1).indices == labels).item() / labels.shape[0]
+# det_loss = 0.1 * (loss_l + loss_c)
+# loss = det_loss + loss_end_cls
+# loss.backward()
+
+
+# class MultiBoxLoss(nn.Module):
+# """SSD Weighted Loss Function
+# Compute Targets:
+# 1) Produce Confidence Target Indices by matching ground truth boxes
+# with (default) 'priorboxes' that have jaccard index > threshold parameter
+# (default threshold: 0.5).
+# 2) Produce localization target by 'encoding' variance into offsets of ground
+# truth boxes and their matched 'priorboxes'.
+# 3) Hard negative mining to filter the excessive number of negative examples
+# that comes with using a large number of default bounding boxes.
+# (default negative:positive ratio 3:1)
+# Objective Loss:
+# L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N
+# Where, Lconf is the CrossEntropy Loss and Lloc is the SmoothL1 Loss
+# weighted by α which is set to 1 by cross val.
+# Args:
+# c: class confidences,
+# l: predicted boxes,
+# g: ground truth boxes
+# N: number of matched default boxes
+# See: https://arxiv.org/pdf/1512.02325.pdf for more details.
+# """
+
+# def __init__(self, num_classes, overlap_thresh, prior_for_matching,
+# bkg_label, neg_mining, neg_pos, neg_overlap, encode_target,
+# use_gpu=True):
+# super(MultiBoxLoss, self).__init__()
+# self.use_gpu = use_gpu
+# self.num_classes = num_classes
+# self.threshold = overlap_thresh
+# self.background_label = bkg_label
+# self.encode_target = encode_target
+# self.use_prior_for_matching = prior_for_matching
+# self.do_neg_mining = neg_mining
+# self.negpos_ratio = neg_pos
+# self.neg_overlap = neg_overlap
+# self.variance = [0.1, 0.2] # cfg['variance']
+
+# # def forward(self, predictions, targets):
+# def forward(self, predictions, conf_t, loc_t):
+# """Multibox Loss
+# Args:
+# predictions (tuple): A tuple containing loc preds, conf preds,
+# and prior boxes from SSD net.
+# conf shape: torch.size(batch_size,num_priors,num_classes)
+# loc shape: torch.size(batch_size,num_priors,4)
+# priors shape: torch.size(num_priors,4)
+
+# targets (tensor): Ground truth boxes and labels for a batch,
+# shape: [batch_size,num_objs,5] (last idx is the label).
+# """
+# '''
+# priors = priors[:loc_data.size(1), :]
+# num_priors = (priors.size(0))
+# num_classes = self.num_classes
+
+# # match priors (default boxes) and ground truth boxes
+# loc_t = torch.Tensor(num, num_priors, 4)
+# conf_t = torch.LongTensor(num, num_priors)
+# for idx in range(num):
+# truths = targets[idx][:, :-1].data
+# labels = targets[idx][:, -1].data
+# defaults = priors.data
+# match(self.threshold, truths, defaults, self.variance, labels,
+# loc_t, conf_t, idx)
+# '''
+# loc_data, conf_data = predictions
+# num = loc_data.size(0)
+# if self.use_gpu:
+# loc_t = loc_t.cuda()
+# conf_t = conf_t.cuda()
+# # wrap targets
+# loc_t = Variable(loc_t, requires_grad=False)
+# conf_t = Variable(conf_t, requires_grad=False)
+
+# pos = conf_t > 0
+# num_pos = pos.sum(dim=1, keepdim=True)
+
+# # Localization Loss (Smooth L1)
+# # Shape: [batch,num_priors,4]
+# pos_idx = pos.unsqueeze(pos.dim()).expand_as(loc_data)
+# loc_p = loc_data[pos_idx].view(-1, 4)
+# loc_t = loc_t[pos_idx].view(-1, 4)
+# loss_l = F.smooth_l1_loss(loc_p, loc_t, size_average=False)
+
+# # Compute max conf across batch for hard negative mining
+# batch_conf = conf_data.view(-1, self.num_classes)
+# loss_c = log_sum_exp(batch_conf) - batch_conf.gather(1, conf_t.view(-1, 1))
+
+# # Hard Negative Mining
+# # loss_c[pos] = 0 # filter out pos boxes for now
+# loss_c[pos.view(-1, 1)] = 0
+# loss_c = loss_c.view(num, -1)
+# _, loss_idx = loss_c.sort(1, descending=True)
+# _, idx_rank = loss_idx.sort(1)
+# num_pos = pos.long().sum(1, keepdim=True)
+# num_neg = torch.clamp(self.negpos_ratio*num_pos, max=pos.size(1)-1)
+# neg = idx_rank < num_neg.expand_as(idx_rank)
+
+# # Confidence Loss Including Positive and Negative Examples
+# pos_idx = pos.unsqueeze(2).expand_as(conf_data)
+# neg_idx = neg.unsqueeze(2).expand_as(conf_data)
+# conf_p = conf_data[(pos_idx+neg_idx).gt(0)].view(-1, self.num_classes)
+# targets_weighted = conf_t[(pos+neg).gt(0)]
+# loss_c = F.cross_entropy(conf_p, targets_weighted, size_average=False)
+
+# # Sum of losses: L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N
+
+# N = num_pos.data.sum() if num_pos.data.sum() else 1
+# loss_l /= N
+# loss_c /= N
+# return loss_l, loss_c
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/loss/id_loss.py b/clean/image/effort/DeepfakeBench/training/loss/id_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..1c10579957a6a34213766d3809051c99da4a1bd1
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/loss/id_loss.py
@@ -0,0 +1,16 @@
+import torch
+import torch.nn as nn
+from loss.abstract_loss_func import AbstractLossClass
+from utils.registry import LOSSFUNC
+
+@LOSSFUNC.register_module(module_name="id_loss")
+class IDLoss(AbstractLossClass):
+ def __init__(self, margin=0.5):
+ super().__init__()
+ self.cosine_similarity = nn.CosineSimilarity(dim=1, eps=1e-6)
+ self.margin = margin
+
+ def forward(self, x1, x2):
+ cosine_similarity = self.cosine_similarity(x1, x2)
+ theta = torch.acos(cosine_similarity)
+ return 1 - torch.cos(theta + self.margin)
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/loss/js_loss.py b/clean/image/effort/DeepfakeBench/training/loss/js_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..e70a3d7485bba8fcd51503a9fdc94327ce3cea78
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/loss/js_loss.py
@@ -0,0 +1,32 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from loss.abstract_loss_func import AbstractLossClass
+from utils.registry import LOSSFUNC
+
+
+@LOSSFUNC.register_module(module_name="jsloss")
+class JS_Loss(AbstractLossClass):
+ def __init__(self):
+ super().__init__()
+
+ def forward(self, inputs, targets):
+ """
+ Computes the Jensen-Shannon divergence loss.
+ """
+ # Compute the probability distributions
+ inputs_prob = F.softmax(inputs, dim=1)
+ targets_prob = F.softmax(targets, dim=1)
+
+ # Compute the average probability distribution
+ avg_prob = (inputs_prob + targets_prob) / 2
+
+ # Compute the KL divergence component for each distribution
+ kl_div_loss = nn.KLDivLoss(reduction='batchmean')
+ kl_inputs = kl_div_loss(inputs_prob.log(), avg_prob)
+ kl_targets = kl_div_loss(targets_prob.log(), avg_prob)
+
+ # Compute the Jensen-Shannon divergence
+ loss = 0.5 * (kl_inputs + kl_targets)
+
+ return loss
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/loss/l1_loss.py b/clean/image/effort/DeepfakeBench/training/loss/l1_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..4db12aae3cec995471bfc24b69ba563d49fa6e58
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/loss/l1_loss.py
@@ -0,0 +1,19 @@
+import torch.nn as nn
+from loss.abstract_loss_func import AbstractLossClass
+from utils.registry import LOSSFUNC
+
+
+@LOSSFUNC.register_module(module_name="l1loss")
+class L1Loss(AbstractLossClass):
+ def __init__(self):
+ super().__init__()
+ self.loss_fn = nn.L1Loss()
+
+ def forward(self, inputs, targets):
+ """
+ Computes the l1 loss.
+ """
+ # Compute the l1 loss
+ loss = self.loss_fn(inputs, targets)
+
+ return loss
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/loss/vgg_loss.py b/clean/image/effort/DeepfakeBench/training/loss/vgg_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..62353ea538b962a49a07d5c3a07cf63106a76906
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/loss/vgg_loss.py
@@ -0,0 +1,152 @@
+"""A VGG-based perceptual loss function for PyTorch."""
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torchvision import models, transforms
+import torch
+import torch.nn as nn
+from loss.abstract_loss_func import AbstractLossClass
+from utils.registry import LOSSFUNC
+
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+class Lambda(nn.Module):
+ """Wraps a callable in an :class:`nn.Module` without registering it."""
+
+ def __init__(self, func):
+ super().__init__()
+ object.__setattr__(self, 'forward', func)
+
+ def extra_repr(self):
+ return getattr(self.forward, '__name__', type(self.forward).__name__) + '()'
+
+
+class WeightedLoss(nn.ModuleList):
+ """A weighted combination of multiple loss functions."""
+
+ def __init__(self, losses, weights, verbose=False):
+ super().__init__()
+ for loss in losses:
+ self.append(loss if isinstance(loss, nn.Module) else Lambda(loss))
+ self.weights = weights
+ self.verbose = verbose
+
+ def _print_losses(self, losses):
+ for i, loss in enumerate(losses):
+ print(f'({i}) {type(self[i]).__name__}: {loss.item()}')
+
+ def forward(self, *args, **kwargs):
+ losses = []
+ for loss, weight in zip(self, self.weights):
+ losses.append(loss(*args, **kwargs) * weight)
+ if self.verbose:
+ self._print_losses(losses)
+ return sum(losses)
+
+
+class TVLoss(nn.Module):
+ """Total variation loss (Lp penalty on image gradient magnitude).
+ The input must be 4D. If a target (second parameter) is passed in, it is
+ ignored.
+ ``p=1`` yields the vectorial total variation norm. It is a generalization
+ of the originally proposed (isotropic) 2D total variation norm (see
+ (see https://en.wikipedia.org/wiki/Total_variation_denoising) for color
+ images. On images with a single channel it is equal to the 2D TV norm.
+ ``p=2`` yields a variant that is often used for smoothing out noise in
+ reconstructions of images from neural network feature maps (see Mahendran
+ and Vevaldi, "Understanding Deep Image Representations by Inverting
+ Them", https://arxiv.org/abs/1412.0035)
+ :attr:`reduction` can be set to ``'mean'``, ``'sum'``, or ``'none'``
+ similarly to the loss functions in :mod:`torch.nn`. The default is
+ ``'mean'``.
+ """
+
+ def __init__(self, p, reduction='mean', eps=1e-8):
+ super().__init__()
+ if p not in {1, 2}:
+ raise ValueError('p must be 1 or 2')
+ if reduction not in {'mean', 'sum', 'none'}:
+ raise ValueError("reduction must be 'mean', 'sum', or 'none'")
+ self.p = p
+ self.reduction = reduction
+ self.eps = eps
+
+ def forward(self, input, target=None):
+ input = F.pad(input, (0, 1, 0, 1), 'replicate')
+ x_diff = input[..., :-1, :-1] - input[..., :-1, 1:]
+ y_diff = input[..., :-1, :-1] - input[..., 1:, :-1]
+ diff = x_diff**2 + y_diff**2
+ if self.p == 1:
+ diff = (diff + self.eps).mean(dim=1, keepdims=True).sqrt()
+ if self.reduction == 'mean':
+ return diff.mean()
+ if self.reduction == 'sum':
+ return diff.sum()
+ return diff
+
+
+@LOSSFUNC.register_module(module_name="vgg_loss")
+class VGGLoss(AbstractLossClass):
+ """Computes the VGG perceptual loss between two batches of images.
+ The input and target must be 4D tensors with three channels
+ ``(B, 3, H, W)`` and must have equivalent shapes. Pixel values should be
+ normalized to the range 0–1.
+ The VGG perceptual loss is the mean squared difference between the features
+ computed for the input and target at layer :attr:`layer` (default 8, or
+ ``relu2_2``) of the pretrained model specified by :attr:`model` (either
+ ``'vgg16'`` (default) or ``'vgg19'``).
+ If :attr:`shift` is nonzero, a random shift of at most :attr:`shift`
+ pixels in both height and width will be applied to all images in the input
+ and target. The shift will only be applied when the loss function is in
+ training mode, and will not be applied if a precomputed feature map is
+ supplied as the target.
+ :attr:`reduction` can be set to ``'mean'``, ``'sum'``, or ``'none'``
+ similarly to the loss functions in :mod:`torch.nn`. The default is
+ ``'mean'``.
+ :meth:`get_features()` may be used to precompute the features for the
+ target, to speed up the case where inputs are compared against the same
+ target over and over. To use the precomputed features, pass them in as
+ :attr:`target` and set :attr:`target_is_features` to :code:`True`.
+ Instances of :class:`VGGLoss` must be manually converted to the same
+ device and dtype as their inputs.
+ """
+
+ models = {'vgg16': models.vgg16, 'vgg19': models.vgg19}
+
+ def __init__(self, model='vgg16', layer=8, shift=0, reduction='mean'):
+ super().__init__()
+ self.instancenorm = nn.InstanceNorm2d(512, affine=False)
+ self.shift = shift
+ self.reduction = reduction
+ self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
+ std=[0.229, 0.224, 0.225])
+ self.model = self.models[model](pretrained=True).features[:layer+1]
+ self.model.eval()
+ self.model.requires_grad_(False)
+ self.model.to(device)
+
+ def get_features(self, input):
+ return self.model(self.normalize(input))
+
+ def train(self, mode=True):
+ self.training = mode
+
+ def forward(self, input, target, target_is_features=False):
+ if target_is_features:
+ input_feats = self.get_features(input)
+ target_feats = target
+ else:
+ sep = input.shape[0]
+ batch = torch.cat([input, target])
+ if self.shift and self.training:
+ padded = F.pad(batch, [self.shift] * 4, mode='replicate')
+ batch = transforms.RandomCrop(batch.shape[2:])(padded)
+ feats = self.get_features(batch)
+ input_feats, target_feats = feats[:sep], feats[sep:]
+ # input_feats, target_feats = \
+ # self.instancenorm(input_feats), \
+ # self.instancenorm(target_feats)
+ return F.mse_loss(input_feats, target_feats, reduction=self.reduction)
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/metrics/__init__.py b/clean/image/effort/DeepfakeBench/training/metrics/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..676145d777810e4a51bdaf59fdec4f5358aae349
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/metrics/__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)
diff --git a/clean/image/effort/DeepfakeBench/training/metrics/base_metrics_class.py b/clean/image/effort/DeepfakeBench/training/metrics/base_metrics_class.py
new file mode 100644
index 0000000000000000000000000000000000000000..4e3f33b3707e0adc53e5c68848bc4c51132e53ca
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/metrics/base_metrics_class.py
@@ -0,0 +1,205 @@
+import numpy as np
+from sklearn import metrics
+from collections import defaultdict
+import torch
+import torch.nn as nn
+
+
+def get_accracy(output, label):
+ _, prediction = torch.max(output, 1) # argmax
+ correct = (prediction == label).sum().item()
+ accuracy = correct / prediction.size(0)
+ return accuracy
+
+
+def get_prediction(output, label):
+ prob = nn.functional.softmax(output, dim=1)[:, 1]
+ prob = prob.view(prob.size(0), 1)
+ label = label.view(label.size(0), 1)
+ #print(prob.size(), label.size())
+ datas = torch.cat((prob, label.float()), dim=1)
+ return datas
+
+
+def calculate_metrics_for_train(label, output):
+ if output.size(1) == 2:
+ prob = torch.softmax(output, dim=1)[:, 1]
+ else:
+ prob = output
+
+ # Accuracy
+ _, prediction = torch.max(output, 1)
+ correct = (prediction == label).sum().item()
+ accuracy = correct / prediction.size(0)
+
+ # Average Precision
+ y_true = label.cpu().detach().numpy()
+ y_pred = prob.cpu().detach().numpy()
+ ap = metrics.average_precision_score(y_true, y_pred)
+
+ # AUC and EER
+ try:
+ fpr, tpr, thresholds = metrics.roc_curve(label.squeeze().cpu().numpy(),
+ prob.squeeze().cpu().numpy(),
+ pos_label=1)
+ except:
+ # for the case when we only have one sample
+ return None, None, accuracy, ap
+
+ if np.isnan(fpr[0]) or np.isnan(tpr[0]):
+ # for the case when all the samples within a batch is fake/real
+ auc, eer = None, None
+ else:
+ auc = metrics.auc(fpr, tpr)
+ fnr = 1 - tpr
+ eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+
+ return auc, eer, accuracy, ap
+
+
+# ------------ compute average metrics of batches---------------------
+class Metrics_batch():
+ def __init__(self):
+ self.tprs = []
+ self.mean_fpr = np.linspace(0, 1, 100)
+ self.aucs = []
+ self.eers = []
+ self.aps = []
+
+ self.correct = 0
+ self.total = 0
+ self.losses = []
+
+ def update(self, label, output):
+ acc = self._update_acc(label, output)
+ if output.size(1) == 2:
+ prob = torch.softmax(output, dim=1)[:, 1]
+ else:
+ prob = output
+ #label = 1-label
+ #prob = torch.softmax(output, dim=1)[:, 1]
+ auc, eer = self._update_auc(label, prob)
+ ap = self._update_ap(label, prob)
+
+ return acc, auc, eer, ap
+
+ def _update_auc(self, lab, prob):
+ fpr, tpr, thresholds = metrics.roc_curve(lab.squeeze().cpu().numpy(),
+ prob.squeeze().cpu().numpy(),
+ pos_label=1)
+ if np.isnan(fpr[0]) or np.isnan(tpr[0]):
+ return -1, -1
+
+ auc = metrics.auc(fpr, tpr)
+ interp_tpr = np.interp(self.mean_fpr, fpr, tpr)
+ interp_tpr[0] = 0.0
+ self.tprs.append(interp_tpr)
+ self.aucs.append(auc)
+
+ # return auc
+
+ # EER
+ fnr = 1 - tpr
+ eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+ self.eers.append(eer)
+
+ return auc, eer
+
+ def _update_acc(self, lab, output):
+ _, prediction = torch.max(output, 1) # argmax
+ correct = (prediction == lab).sum().item()
+ accuracy = correct / prediction.size(0)
+ # self.accs.append(accuracy)
+ self.correct = self.correct+correct
+ self.total = self.total+lab.size(0)
+ return accuracy
+
+ def _update_ap(self, label, prob):
+ y_true = label.cpu().detach().numpy()
+ y_pred = prob.cpu().detach().numpy()
+ ap = metrics.average_precision_score(y_true,y_pred)
+ self.aps.append(ap)
+
+ return np.mean(ap)
+
+ def get_mean_metrics(self):
+ mean_acc, std_acc = self.correct/self.total, 0
+ mean_auc, std_auc = self._mean_auc()
+ mean_err, std_err = np.mean(self.eers), np.std(self.eers)
+ mean_ap, std_ap = np.mean(self.aps), np.std(self.aps)
+
+ return {'acc':mean_acc, 'auc':mean_auc, 'eer':mean_err, 'ap':mean_ap}
+
+ def _mean_auc(self):
+ mean_tpr = np.mean(self.tprs, axis=0)
+ mean_tpr[-1] = 1.0
+ mean_auc = metrics.auc(self.mean_fpr, mean_tpr)
+ std_auc = np.std(self.aucs)
+ return mean_auc, std_auc
+
+ def clear(self):
+ self.tprs.clear()
+ self.aucs.clear()
+ # self.accs.clear()
+ self.correct=0
+ self.total=0
+ self.eers.clear()
+ self.aps.clear()
+ self.losses.clear()
+
+
+# ------------ compute average metrics of all data ---------------------
+class Metrics_all():
+ def __init__(self):
+ self.probs = []
+ self.labels = []
+ self.correct = 0
+ self.total = 0
+
+ def store(self, label, output):
+ prob = torch.softmax(output, dim=1)[:, 1]
+ _, prediction = torch.max(output, 1) # argmax
+ correct = (prediction == label).sum().item()
+ self.correct += correct
+ self.total += label.size(0)
+ self.labels.append(label.squeeze().cpu().numpy())
+ self.probs.append(prob.squeeze().cpu().numpy())
+
+ def get_metrics(self):
+ y_pred = np.concatenate(self.probs)
+ y_true = np.concatenate(self.labels)
+ # auc
+ fpr, tpr, thresholds = metrics.roc_curve(y_true,y_pred,pos_label=1)
+ auc = metrics.auc(fpr, tpr)
+ # eer
+ fnr = 1 - tpr
+ eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+ # ap
+ ap = metrics.average_precision_score(y_true,y_pred)
+ # acc
+ acc = self.correct / self.total
+ return {'acc':acc, 'auc':auc, 'eer':eer, 'ap':ap}
+
+ def clear(self):
+ self.probs.clear()
+ self.labels.clear()
+ self.correct = 0
+ self.total = 0
+
+
+# only used to record a series of scalar value
+class Recorder:
+ def __init__(self):
+ self.sum = 0
+ self.num = 0
+ def update(self, item, num=1):
+ if item is not None:
+ self.sum += item * num
+ self.num += num
+ def average(self):
+ if self.num == 0:
+ return None
+ return self.sum/self.num
+ def clear(self):
+ self.sum = 0
+ self.num = 0
diff --git a/clean/image/effort/DeepfakeBench/training/metrics/registry.py b/clean/image/effort/DeepfakeBench/training/metrics/registry.py
new file mode 100644
index 0000000000000000000000000000000000000000..86e256c18d0ad522de79149a154f676fd0bdb414
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/metrics/registry.py
@@ -0,0 +1,20 @@
+class Registry(object):
+ def __init__(self):
+ self.data = {}
+
+ def register_module(self, module_name=None):
+ def _register(cls):
+ name = module_name
+ if module_name is None:
+ name = cls.__name__
+ self.data[name] = cls
+ return cls
+ return _register
+
+ def __getitem__(self, key):
+ return self.data[key]
+
+BACKBONE = Registry()
+DETECTOR = Registry()
+TRAINER = Registry()
+LOSSFUNC = Registry()
diff --git a/clean/image/effort/DeepfakeBench/training/metrics/utils.py b/clean/image/effort/DeepfakeBench/training/metrics/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..55fe8665d6883fdd6b043bbc8366d65b19441a47
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/metrics/utils.py
@@ -0,0 +1,176 @@
+from sklearn import metrics
+import numpy as np
+
+
+def parse_metric_for_print(metric_dict):
+ if metric_dict is None:
+ return "\n"
+ str = "\n"
+ str += "================================ Each dataset best metric ================================ \n"
+ for key, value in metric_dict.items():
+ if key != 'avg':
+ str= str+ f"| {key}: "
+ for k,v in value.items():
+ str = str + f" {k}={v} "
+ str= str+ "| \n"
+ else:
+ str += "============================================================================================= \n"
+ str += "================================== Average best metric ====================================== \n"
+ avg_dict = value
+ for avg_key, avg_value in avg_dict.items():
+ if avg_key == 'dataset_dict':
+ for key,value in avg_value.items():
+ str = str + f"| {key}: {value} | \n"
+ else:
+ str = str + f"| avg {avg_key}: {avg_value} | \n"
+ str += "============================================================================================="
+ return str
+
+
+# def get_test_metrics(y_pred, y_true, img_names):
+# def get_video_metrics(image, pred, label):
+# result_dict = {}
+# new_label = []
+# new_pred = []
+# # print(image[0])
+# # print(pred.shape)
+# # print(label.shape)
+# for item in np.transpose(np.stack((image, pred, label)), (1, 0)):
+# # 分割字符串，获取'a'和'b'的值
+# s = item[0]
+# if '\\' in s:
+# parts = s.split('\\')
+# else:
+# parts = s.split('/')
+# a = parts[-2]
+# b = parts[-1]
+
+# # 如果'a'的值还没有在字典中，添加一个新的键值对
+# if a not in result_dict:
+# result_dict[a] = []
+
+# # 将'b'的值添加到'a'的列表中
+# result_dict[a].append(item)
+# image_arr = list(result_dict.values())
+# # 将字典的值转换为一个列表，得到二维数组
+
+# for video in image_arr:
+# pred_sum = 0
+# label_sum = 0
+# leng = 0
+# for frame in video:
+# pred_sum += float(frame[1])
+# label_sum += int(frame[2])
+# leng += 1
+# new_pred.append(pred_sum / leng)
+# new_label.append(int(label_sum / leng))
+# fpr, tpr, thresholds = metrics.roc_curve(new_label, new_pred)
+# v_auc = metrics.auc(fpr, tpr)
+# fnr = 1 - tpr
+# v_eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+# return v_auc, v_eer
+
+
+# y_pred = y_pred.squeeze()
+# # auc
+# fpr, tpr, thresholds = metrics.roc_curve(y_true, y_pred, pos_label=1)
+# auc = metrics.auc(fpr, tpr)
+# # eer
+# fnr = 1 - tpr
+# eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+# # ap
+# ap = metrics.average_precision_score(y_true, y_pred)
+# # acc
+# prediction_class = (y_pred > 0.5).astype(int)
+# correct = (prediction_class == np.clip(y_true, a_min=0, a_max=1)).sum().item()
+# acc = correct / len(prediction_class)
+# if type(img_names[0]) is not list:
+# # calculate video-level auc for the frame-level methods.
+# try:
+# v_auc, _ = get_video_metrics(img_names, y_pred, y_true)
+# except Exception as e:
+# print(e)
+# v_auc=auc
+# else:
+# # video-level methods
+# v_auc=auc
+
+# return {'acc': acc, 'auc': auc, 'eer': eer, 'ap': ap, 'pred': y_pred, 'video_auc': v_auc, 'label': y_true}
+
+
+
+
+
+def get_test_metrics(y_pred, y_true, img_names):
+ def get_video_metrics(image, pred, label):
+ result_dict = {}
+ new_label = []
+ new_pred = []
+ for item in np.transpose(np.stack((image, pred, label)), (1, 0)):
+ s = item[0]
+ if '\\' in s:
+ parts = s.split('\\')
+ else:
+ parts = s.split('/')
+ a = parts[-2]
+ b = parts[-1]
+
+ if a not in result_dict:
+ result_dict[a] = []
+
+ result_dict[a].append(item)
+ image_arr = list(result_dict.values())
+
+ for video in image_arr:
+ pred_sum = 0
+ label_sum = 0
+ leng = 0
+ for frame in video:
+ pred_sum += float(frame[1])
+ label_sum += int(frame[2])
+ leng += 1
+ new_pred.append(pred_sum / leng)
+ new_label.append(int(label_sum / leng))
+
+ fpr, tpr, thresholds = metrics.roc_curve(new_label, new_pred)
+ v_auc = metrics.auc(fpr, tpr)
+ fnr = 1 - tpr
+ v_eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+
+ # Calculate video-level acc
+ prediction_class = (np.array(new_pred) > 0.5).astype(int)
+ correct = (prediction_class == np.array(new_label)).sum().item()
+ v_acc = correct / len(prediction_class)
+
+ return v_auc, v_eer, v_acc
+
+
+ y_pred = y_pred.squeeze()
+ # auc
+ fpr, tpr, thresholds = metrics.roc_curve(y_true, y_pred, pos_label=1)
+ auc = metrics.auc(fpr, tpr)
+ # eer
+ fnr = 1 - tpr
+ eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+ # ap
+ ap = metrics.average_precision_score(y_true, y_pred)
+ # acc
+ prediction_class = (y_pred > 0.5).astype(int)
+ correct = (prediction_class == np.clip(y_true, a_min=0, a_max=1)).sum().item()
+ acc = correct / len(prediction_class)
+ if type(img_names[0]) is not list:
+ # calculate video-level auc for the frame-level methods.
+ try:
+ v_auc, v_eer, v_acc = get_video_metrics(img_names, y_pred, y_true)
+ return {'acc': acc, 'auc': auc, 'eer': eer, 'ap': ap, 'pred': y_pred, 'video_auc': v_auc, 'video_eer': v_eer, 'video_acc': v_acc, 'label': y_true}
+ except Exception as e:
+ print(e)
+ v_auc=auc
+ return {'acc': acc, 'auc': auc, 'eer': eer, 'ap': ap, 'pred': y_pred, 'label': y_true}
+ else:
+ # video-level methods
+ v_auc=auc
+ v_eer=eer
+ v_acc=acc
+ return {'acc': acc, 'auc': auc, 'eer': eer, 'ap': ap, 'pred': y_pred, 'video_auc': v_auc, 'video_eer': v_eer, 'video_acc': v_acc, 'label': y_true}
+
diff --git a/clean/image/effort/DeepfakeBench/training/networks/__init__.py b/clean/image/effort/DeepfakeBench/training/networks/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..9daf1a387c03c3ab553d68dde34fa6fd4b495627
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/networks/__init__.py
@@ -0,0 +1,5 @@
+from utils.registry import BACKBONE
+from .xception import Xception
+from .mesonet import Meso4, MesoInception4
+from .resnet34 import ResNet34
+from .efficientnetb4 import EfficientNetB4
diff --git a/clean/image/effort/DeepfakeBench/training/networks/adaface.py b/clean/image/effort/DeepfakeBench/training/networks/adaface.py
new file mode 100644
index 0000000000000000000000000000000000000000..21730fdac4bed733ba2cde876884f5b46d8ea230
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/networks/adaface.py
@@ -0,0 +1,414 @@
+from collections import namedtuple
+import torch
+import torch.nn as nn
+from torch.nn import Dropout
+from torch.nn import MaxPool2d
+from torch.nn import Sequential
+from torch.nn import Conv2d, Linear
+from torch.nn import BatchNorm1d, BatchNorm2d
+from torch.nn import ReLU, Sigmoid
+from torch.nn import Module
+from torch.nn import PReLU
+import os
+
+def build_model(model_name='ir_50'):
+ if model_name == 'ir_101':
+ return IR_101(input_size=(112,112))
+ elif model_name == 'ir_50':
+ return IR_50(input_size=(112,112))
+ elif model_name == 'ir_se_50':
+ return IR_SE_50(input_size=(112,112))
+ elif model_name == 'ir_34':
+ return IR_34(input_size=(112,112))
+ elif model_name == 'ir_18':
+ return IR_18(input_size=(112,112))
+ else:
+ raise ValueError('not a correct model name', model_name)
+
+def initialize_weights(modules):
+ """ Weight initilize, conv2d and linear is initialized with kaiming_normal
+ """
+ for m in modules:
+ if isinstance(m, nn.Conv2d):
+ nn.init.kaiming_normal_(m.weight,
+ mode='fan_out',
+ nonlinearity='relu')
+ if m.bias is not None:
+ m.bias.data.zero_()
+ elif isinstance(m, nn.BatchNorm2d):
+ m.weight.data.fill_(1)
+ m.bias.data.zero_()
+ elif isinstance(m, nn.Linear):
+ nn.init.kaiming_normal_(m.weight,
+ mode='fan_out',
+ nonlinearity='relu')
+ if m.bias is not None:
+ m.bias.data.zero_()
+
+
+class Flatten(Module):
+ """ Flat tensor
+ """
+ def forward(self, input):
+ return input.view(input.size(0), -1)
+
+
+class LinearBlock(Module):
+ """ Convolution block without no-linear activation layer
+ """
+ def __init__(self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1):
+ super(LinearBlock, self).__init__()
+ self.conv = Conv2d(in_c, out_c, kernel, stride, padding, groups=groups, bias=False)
+ self.bn = BatchNorm2d(out_c)
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.bn(x)
+ return x
+
+
+class GNAP(Module):
+ """ Global Norm-Aware Pooling block
+ """
+ def __init__(self, in_c):
+ super(GNAP, self).__init__()
+ self.bn1 = BatchNorm2d(in_c, affine=False)
+ self.pool = nn.AdaptiveAvgPool2d((1, 1))
+ self.bn2 = BatchNorm1d(in_c, affine=False)
+
+ def forward(self, x):
+ x = self.bn1(x)
+ x_norm = torch.norm(x, 2, 1, True)
+ x_norm_mean = torch.mean(x_norm)
+ weight = x_norm_mean / x_norm
+ x = x * weight
+ x = self.pool(x)
+ x = x.view(x.shape[0], -1)
+ feature = self.bn2(x)
+ return feature
+
+
+class GDC(Module):
+ """ Global Depthwise Convolution block
+ """
+ def __init__(self, in_c, embedding_size):
+ super(GDC, self).__init__()
+ self.conv_6_dw = LinearBlock(in_c, in_c,
+ groups=in_c,
+ kernel=(7, 7),
+ stride=(1, 1),
+ padding=(0, 0))
+ self.conv_6_flatten = Flatten()
+ self.linear = Linear(in_c, embedding_size, bias=False)
+ self.bn = BatchNorm1d(embedding_size, affine=False)
+
+ def forward(self, x):
+ x = self.conv_6_dw(x)
+ x = self.conv_6_flatten(x)
+ x = self.linear(x)
+ x = self.bn(x)
+ return x
+
+
+class SEModule(Module):
+ """ SE block
+ """
+ def __init__(self, channels, reduction):
+ super(SEModule, self).__init__()
+ self.avg_pool = nn.AdaptiveAvgPool2d(1)
+ self.fc1 = Conv2d(channels, channels // reduction,
+ kernel_size=1, padding=0, bias=False)
+
+ nn.init.xavier_uniform_(self.fc1.weight.data)
+
+ self.relu = ReLU(inplace=True)
+ self.fc2 = Conv2d(channels // reduction, channels,
+ kernel_size=1, padding=0, bias=False)
+
+ self.sigmoid = Sigmoid()
+
+ def forward(self, x):
+ module_input = x
+ x = self.avg_pool(x)
+ x = self.fc1(x)
+ x = self.relu(x)
+ x = self.fc2(x)
+ x = self.sigmoid(x)
+
+ return module_input * x
+
+
+
+class BasicBlockIR(Module):
+ """ BasicBlock for IRNet
+ """
+ def __init__(self, in_channel, depth, stride):
+ super(BasicBlockIR, self).__init__()
+ if in_channel == depth:
+ self.shortcut_layer = MaxPool2d(1, stride)
+ else:
+ self.shortcut_layer = Sequential(
+ Conv2d(in_channel, depth, (1, 1), stride, bias=False),
+ BatchNorm2d(depth))
+ self.res_layer = Sequential(
+ BatchNorm2d(in_channel),
+ Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False),
+ BatchNorm2d(depth),
+ PReLU(depth),
+ Conv2d(depth, depth, (3, 3), stride, 1, bias=False),
+ BatchNorm2d(depth))
+
+ def forward(self, x):
+ shortcut = self.shortcut_layer(x)
+ res = self.res_layer(x)
+
+ return res + shortcut
+
+
+class BottleneckIR(Module):
+ """ BasicBlock with bottleneck for IRNet
+ """
+ def __init__(self, in_channel, depth, stride):
+ super(BottleneckIR, self).__init__()
+ reduction_channel = depth // 4
+ if in_channel == depth:
+ self.shortcut_layer = MaxPool2d(1, stride)
+ else:
+ self.shortcut_layer = Sequential(
+ Conv2d(in_channel, depth, (1, 1), stride, bias=False),
+ BatchNorm2d(depth))
+ self.res_layer = Sequential(
+ BatchNorm2d(in_channel),
+ Conv2d(in_channel, reduction_channel, (1, 1), (1, 1), 0, bias=False),
+ BatchNorm2d(reduction_channel),
+ PReLU(reduction_channel),
+ Conv2d(reduction_channel, reduction_channel, (3, 3), (1, 1), 1, bias=False),
+ BatchNorm2d(reduction_channel),
+ PReLU(reduction_channel),
+ Conv2d(reduction_channel, depth, (1, 1), stride, 0, bias=False),
+ BatchNorm2d(depth))
+
+ def forward(self, x):
+ shortcut = self.shortcut_layer(x)
+ res = self.res_layer(x)
+
+ return res + shortcut
+
+
+class BasicBlockIRSE(BasicBlockIR):
+ def __init__(self, in_channel, depth, stride):
+ super(BasicBlockIRSE, self).__init__(in_channel, depth, stride)
+ self.res_layer.add_module("se_block", SEModule(depth, 16))
+
+
+class BottleneckIRSE(BottleneckIR):
+ def __init__(self, in_channel, depth, stride):
+ super(BottleneckIRSE, self).__init__(in_channel, depth, stride)
+ self.res_layer.add_module("se_block", SEModule(depth, 16))
+
+
+class Bottleneck(namedtuple('Block', ['in_channel', 'depth', 'stride'])):
+ '''A named tuple describing a ResNet block.'''
+
+
+def get_block(in_channel, depth, num_units, stride=2):
+
+ return [Bottleneck(in_channel, depth, stride)] +\
+ [Bottleneck(depth, depth, 1) for i in range(num_units - 1)]
+
+
+def get_blocks(num_layers):
+ if num_layers == 18:
+ blocks = [
+ get_block(in_channel=64, depth=64, num_units=2),
+ get_block(in_channel=64, depth=128, num_units=2),
+ get_block(in_channel=128, depth=256, num_units=2),
+ get_block(in_channel=256, depth=512, num_units=2)
+ ]
+ elif num_layers == 34:
+ blocks = [
+ get_block(in_channel=64, depth=64, num_units=3),
+ get_block(in_channel=64, depth=128, num_units=4),
+ get_block(in_channel=128, depth=256, num_units=6),
+ get_block(in_channel=256, depth=512, num_units=3)
+ ]
+ elif num_layers == 50:
+ blocks = [
+ get_block(in_channel=64, depth=64, num_units=3),
+ get_block(in_channel=64, depth=128, num_units=4),
+ get_block(in_channel=128, depth=256, num_units=14),
+ get_block(in_channel=256, depth=512, num_units=3)
+ ]
+ elif num_layers == 100:
+ blocks = [
+ get_block(in_channel=64, depth=64, num_units=3),
+ get_block(in_channel=64, depth=128, num_units=13),
+ get_block(in_channel=128, depth=256, num_units=30),
+ get_block(in_channel=256, depth=512, num_units=3)
+ ]
+ elif num_layers == 152:
+ blocks = [
+ get_block(in_channel=64, depth=256, num_units=3),
+ get_block(in_channel=256, depth=512, num_units=8),
+ get_block(in_channel=512, depth=1024, num_units=36),
+ get_block(in_channel=1024, depth=2048, num_units=3)
+ ]
+ elif num_layers == 200:
+ blocks = [
+ get_block(in_channel=64, depth=256, num_units=3),
+ get_block(in_channel=256, depth=512, num_units=24),
+ get_block(in_channel=512, depth=1024, num_units=36),
+ get_block(in_channel=1024, depth=2048, num_units=3)
+ ]
+
+ return blocks
+
+
+class Backbone(Module):
+ def __init__(self, input_size, num_layers, mode='ir'):
+ """ Args:
+ input_size: input_size of backbone
+ num_layers: num_layers of backbone
+ mode: support ir or irse
+ """
+ super(Backbone, self).__init__()
+ assert input_size[0] in [112, 224], \
+ "input_size should be [112, 112] or [224, 224]"
+ assert num_layers in [18, 34, 50, 100, 152, 200], \
+ "num_layers should be 18, 34, 50, 100 or 152"
+ assert mode in ['ir', 'ir_se'], \
+ "mode should be ir or ir_se"
+ self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False),
+ BatchNorm2d(64), PReLU(64))
+ blocks = get_blocks(num_layers)
+ if num_layers <= 100:
+ if mode == 'ir':
+ unit_module = BasicBlockIR
+ elif mode == 'ir_se':
+ unit_module = BasicBlockIRSE
+ output_channel = 512
+ else:
+ if mode == 'ir':
+ unit_module = BottleneckIR
+ elif mode == 'ir_se':
+ unit_module = BottleneckIRSE
+ output_channel = 2048
+
+ if input_size[0] == 112:
+ self.output_layer = Sequential(BatchNorm2d(output_channel),
+ Dropout(0.4), Flatten(),
+ Linear(output_channel * 7 * 7, 512),
+ BatchNorm1d(512, affine=False))
+ else:
+ self.output_layer = Sequential(
+ BatchNorm2d(output_channel), Dropout(0.4), Flatten(),
+ Linear(output_channel * 14 * 14, 512),
+ BatchNorm1d(512, affine=False))
+
+ modules = []
+ for block in blocks:
+ for bottleneck in block:
+ modules.append(
+ unit_module(bottleneck.in_channel, bottleneck.depth,
+ bottleneck.stride))
+ self.body = Sequential(*modules)
+
+ initialize_weights(self.modules())
+
+
+ def forward(self, x):
+
+ # current code only supports one extra image
+ # it comes with a extra dimension for number of extra image. We will just squeeze it out for now
+ x = self.input_layer(x)
+
+ for idx, module in enumerate(self.body):
+ x = module(x)
+
+ x = self.output_layer(x)
+ norm = torch.norm(x, 2, 1, True)
+ output = torch.div(x, norm)
+
+ return output, norm
+
+
+
+def IR_18(input_size):
+ """ Constructs a ir-18 model.
+ """
+ model = Backbone(input_size, 18, 'ir')
+
+ return model
+
+
+def IR_34(input_size):
+ """ Constructs a ir-34 model.
+ """
+ model = Backbone(input_size, 34, 'ir')
+
+ return model
+
+
+def IR_50(input_size):
+ """ Constructs a ir-50 model.
+ """
+ model = Backbone(input_size, 50, 'ir')
+
+ return model
+
+
+def IR_101(input_size):
+ """ Constructs a ir-101 model.
+ """
+ model = Backbone(input_size, 100, 'ir')
+
+ return model
+
+
+def IR_152(input_size):
+ """ Constructs a ir-152 model.
+ """
+ model = Backbone(input_size, 152, 'ir')
+
+ return model
+
+
+def IR_200(input_size):
+ """ Constructs a ir-200 model.
+ """
+ model = Backbone(input_size, 200, 'ir')
+
+ return model
+
+
+def IR_SE_50(input_size):
+ """ Constructs a ir_se-50 model.
+ """
+ model = Backbone(input_size, 50, 'ir_se')
+
+ return model
+
+
+def IR_SE_101(input_size):
+ """ Constructs a ir_se-101 model.
+ """
+ model = Backbone(input_size, 100, 'ir_se')
+
+ return model
+
+
+def IR_SE_152(input_size):
+ """ Constructs a ir_se-152 model.
+ """
+ model = Backbone(input_size, 152, 'ir_se')
+
+ return model
+
+
+def IR_SE_200(input_size):
+ """ Constructs a ir_se-200 model.
+ """
+ model = Backbone(input_size, 200, 'ir_se')
+
+ return model
+
diff --git a/clean/image/effort/DeepfakeBench/training/networks/base_backbone.py b/clean/image/effort/DeepfakeBench/training/networks/base_backbone.py
new file mode 100644
index 0000000000000000000000000000000000000000..8cbb14439c4c8da93d12edd0f0af442d8f460698
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/networks/base_backbone.py
@@ -0,0 +1,32 @@
+import abc
+import torch
+from typing import Union
+
+class AbstractBackbone(abc.ABC):
+ """
+ All backbones for detectors should subclass this class.
+ """
+ def __init__(self, config, 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
+ """
+ pass
+
+ @abc.abstractmethod
+ def features(self, data_dict: dict) -> torch.tensor:
+ """
+ """
+
+ @abc.abstractmethod
+ def classifier(self, features: torch.tensor) -> torch.tensor:
+ """
+ """
+
+ def init_weights(self, pretrained_path: Union[bool, str]):
+ """
+ This method can be optionally implemented by subclasses.
+ """
+ pass
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/networks/cls_hrnet.py b/clean/image/effort/DeepfakeBench/training/networks/cls_hrnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..2c04a0e05817e7a3d91dd2af895aae7c5dc67403
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/networks/cls_hrnet.py
@@ -0,0 +1,569 @@
+'''
+# author: Zhiyuan Yan
+# email: zhiyuanyan@link.cuhk.edu.cn
+# date: 2023-0706
+
+# ------------------------------------------------------------------------------
+# Copyright (c) Microsoft
+# Licensed under the MIT License.
+# Written by Bin Xiao (Bin.Xiao@microsoft.com)
+# Modified by Ke Sun (sunk@mail.ustc.edu.cn)
+# ------------------------------------------------------------------------------
+
+The code is mainly modified from the below link:
+https://github.com/HRNet/HRNet-Image-Classification/tree/master
+'''
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import logging
+import functools
+
+import numpy as np
+from typing import Union
+
+import torch
+import torch.nn as nn
+import torch._utils
+import torch.nn.functional as F
+
+BN_MOMENTUM = 0.1
+logger = logging.getLogger(__name__)
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+ """3x3 convolution with padding"""
+ return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+ padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+ expansion = 1
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None):
+ super(BasicBlock, self).__init__()
+ self.conv1 = conv3x3(inplanes, planes, stride)
+ self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
+ self.relu = nn.ReLU(inplace=True)
+ self.conv2 = conv3x3(planes, planes)
+ self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x):
+ residual = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+
+ if self.downsample is not None:
+ residual = self.downsample(x)
+
+ out += residual
+ out = self.relu(out)
+
+ return out
+
+
+class Bottleneck(nn.Module):
+ expansion = 4
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None):
+ super(Bottleneck, self).__init__()
+ self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
+ self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+ padding=1, bias=False)
+ self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
+ self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
+ bias=False)
+ self.bn3 = nn.BatchNorm2d(planes * self.expansion,
+ momentum=BN_MOMENTUM)
+ self.relu = nn.ReLU(inplace=True)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x):
+ residual = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+ out = self.relu(out)
+
+ out = self.conv3(out)
+ out = self.bn3(out)
+
+ if self.downsample is not None:
+ residual = self.downsample(x)
+
+ out += residual
+ out = self.relu(out)
+
+ return out
+
+
+class HighResolutionModule(nn.Module):
+ def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
+ num_channels, fuse_method, multi_scale_output=True):
+ super(HighResolutionModule, self).__init__()
+ self._check_branches(
+ num_branches, blocks, num_blocks, num_inchannels, num_channels)
+
+ self.num_inchannels = num_inchannels
+ self.fuse_method = fuse_method
+ self.num_branches = num_branches
+
+ self.multi_scale_output = multi_scale_output
+
+ self.branches = self._make_branches(
+ num_branches, blocks, num_blocks, num_channels)
+ self.fuse_layers = self._make_fuse_layers()
+ self.relu = nn.ReLU(False)
+
+ def _check_branches(self, num_branches, blocks, num_blocks,
+ num_inchannels, num_channels):
+ if num_branches != len(num_blocks):
+ error_msg = 'NUM_BRANCHES({}) <> NUM_BLOCKS({})'.format(
+ num_branches, len(num_blocks))
+ logger.error(error_msg)
+ raise ValueError(error_msg)
+
+ if num_branches != len(num_channels):
+ error_msg = 'NUM_BRANCHES({}) <> NUM_CHANNELS({})'.format(
+ num_branches, len(num_channels))
+ logger.error(error_msg)
+ raise ValueError(error_msg)
+
+ if num_branches != len(num_inchannels):
+ error_msg = 'NUM_BRANCHES({}) <> NUM_INCHANNELS({})'.format(
+ num_branches, len(num_inchannels))
+ logger.error(error_msg)
+ raise ValueError(error_msg)
+
+ def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
+ stride=1):
+ downsample = None
+ if stride != 1 or \
+ self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
+ downsample = nn.Sequential(
+ nn.Conv2d(self.num_inchannels[branch_index],
+ num_channels[branch_index] * block.expansion,
+ kernel_size=1, stride=stride, bias=False),
+ nn.BatchNorm2d(num_channels[branch_index] * block.expansion,
+ momentum=BN_MOMENTUM),
+ )
+
+ layers = []
+ layers.append(block(self.num_inchannels[branch_index],
+ num_channels[branch_index], stride, downsample))
+ self.num_inchannels[branch_index] = \
+ num_channels[branch_index] * block.expansion
+ for i in range(1, num_blocks[branch_index]):
+ layers.append(block(self.num_inchannels[branch_index],
+ num_channels[branch_index]))
+
+ return nn.Sequential(*layers)
+
+ def _make_branches(self, num_branches, block, num_blocks, num_channels):
+ branches = []
+
+ for i in range(num_branches):
+ branches.append(
+ self._make_one_branch(i, block, num_blocks, num_channels))
+
+ return nn.ModuleList(branches)
+
+ def _make_fuse_layers(self):
+ if self.num_branches == 1:
+ return None
+
+ num_branches = self.num_branches
+ num_inchannels = self.num_inchannels
+ fuse_layers = []
+ for i in range(num_branches if self.multi_scale_output else 1):
+ fuse_layer = []
+ for j in range(num_branches):
+ if j > i:
+ fuse_layer.append(nn.Sequential(
+ nn.Conv2d(num_inchannels[j],
+ num_inchannels[i],
+ 1,
+ 1,
+ 0,
+ bias=False),
+ nn.BatchNorm2d(num_inchannels[i],
+ momentum=BN_MOMENTUM),
+ nn.Upsample(scale_factor=2**(j-i), mode='nearest')))
+ elif j == i:
+ fuse_layer.append(None)
+ else:
+ conv3x3s = []
+ for k in range(i-j):
+ if k == i - j - 1:
+ num_outchannels_conv3x3 = num_inchannels[i]
+ conv3x3s.append(nn.Sequential(
+ nn.Conv2d(num_inchannels[j],
+ num_outchannels_conv3x3,
+ 3, 2, 1, bias=False),
+ nn.BatchNorm2d(num_outchannels_conv3x3,
+ momentum=BN_MOMENTUM)))
+ else:
+ num_outchannels_conv3x3 = num_inchannels[j]
+ conv3x3s.append(nn.Sequential(
+ nn.Conv2d(num_inchannels[j],
+ num_outchannels_conv3x3,
+ 3, 2, 1, bias=False),
+ nn.BatchNorm2d(num_outchannels_conv3x3,
+ momentum=BN_MOMENTUM),
+ nn.ReLU(False)))
+ fuse_layer.append(nn.Sequential(*conv3x3s))
+ fuse_layers.append(nn.ModuleList(fuse_layer))
+
+ return nn.ModuleList(fuse_layers)
+
+ def get_num_inchannels(self):
+ return self.num_inchannels
+
+ def forward(self, x):
+ if self.num_branches == 1:
+ return [self.branches[0](x[0])]
+
+ for i in range(self.num_branches):
+ x[i] = self.branches[i](x[i])
+
+ x_fuse = []
+ for i in range(len(self.fuse_layers)):
+ y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
+ for j in range(1, self.num_branches):
+ if i == j:
+ y = y + x[j]
+ else:
+ y = y + self.fuse_layers[i][j](x[j])
+ x_fuse.append(self.relu(y))
+
+ return x_fuse
+
+
+blocks_dict = {
+ 'BASIC': BasicBlock,
+ 'BOTTLENECK': Bottleneck
+}
+
+
+class HighResolutionNet(nn.Module):
+
+ def __init__(self, cfg):
+ super(HighResolutionNet, self).__init__()
+
+ self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1,
+ bias=False)
+ self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
+ self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1,
+ bias=False)
+ self.bn2 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
+ self.relu = nn.ReLU(inplace=True)
+
+ self.stage1_cfg = cfg['MODEL']['EXTRA']['STAGE1']
+ num_channels = self.stage1_cfg['NUM_CHANNELS'][0]
+ block = blocks_dict[self.stage1_cfg['BLOCK']]
+ num_blocks = self.stage1_cfg['NUM_BLOCKS'][0]
+ self.layer1 = self._make_layer(block, 64, num_channels, num_blocks)
+ stage1_out_channel = block.expansion*num_channels
+
+ self.stage2_cfg = cfg['MODEL']['EXTRA']['STAGE2']
+ num_channels = self.stage2_cfg['NUM_CHANNELS']
+ block = blocks_dict[self.stage2_cfg['BLOCK']]
+ num_channels = [
+ num_channels[i] * block.expansion for i in range(len(num_channels))]
+ self.transition1 = self._make_transition_layer(
+ [stage1_out_channel], num_channels)
+ self.stage2, pre_stage_channels = self._make_stage(
+ self.stage2_cfg, num_channels)
+
+ self.stage3_cfg = cfg['MODEL']['EXTRA']['STAGE3']
+ num_channels = self.stage3_cfg['NUM_CHANNELS']
+ block = blocks_dict[self.stage3_cfg['BLOCK']]
+ num_channels = [
+ num_channels[i] * block.expansion for i in range(len(num_channels))]
+ self.transition2 = self._make_transition_layer(
+ pre_stage_channels, num_channels)
+ self.stage3, pre_stage_channels = self._make_stage(
+ self.stage3_cfg, num_channels)
+
+ self.stage4_cfg = cfg['MODEL']['EXTRA']['STAGE4']
+ num_channels = self.stage4_cfg['NUM_CHANNELS']
+ block = blocks_dict[self.stage4_cfg['BLOCK']]
+ num_channels = [
+ num_channels[i] * block.expansion for i in range(len(num_channels))]
+ self.transition3 = self._make_transition_layer(
+ pre_stage_channels, num_channels)
+ self.stage4, pre_stage_channels = self._make_stage(
+ self.stage4_cfg, num_channels, multi_scale_output=True)
+
+ # Classification Head
+ self.incre_modules, self.downsamp_modules, \
+ self.final_layer = self._make_head(pre_stage_channels)
+
+ self.fc = nn.Linear(2048, 1000)
+
+
+ def _make_head(self, pre_stage_channels):
+ head_block = Bottleneck
+ head_channels = [32, 64, 128, 256]
+
+ # Increasing the #channels on each resolution
+ # from C, 2C, 4C, 8C to 128, 256, 512, 1024
+ incre_modules = []
+ for i, channels in enumerate(pre_stage_channels):
+ incre_module = self._make_layer(head_block,
+ channels,
+ head_channels[i],
+ 1,
+ stride=1)
+ incre_modules.append(incre_module)
+ incre_modules = nn.ModuleList(incre_modules)
+
+ # downsampling modules
+ downsamp_modules = []
+ for i in range(len(pre_stage_channels)-1):
+ in_channels = head_channels[i] * head_block.expansion
+ out_channels = head_channels[i+1] * head_block.expansion
+
+ downsamp_module = nn.Sequential(
+ nn.Conv2d(in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ stride=2,
+ padding=1),
+ nn.BatchNorm2d(out_channels, momentum=BN_MOMENTUM),
+ nn.ReLU(inplace=True)
+ )
+
+ downsamp_modules.append(downsamp_module)
+ downsamp_modules = nn.ModuleList(downsamp_modules)
+
+ final_layer = nn.Sequential(
+ nn.Conv2d(
+ in_channels=head_channels[3] * head_block.expansion,
+ out_channels=2048,
+ kernel_size=1,
+ stride=1,
+ padding=0
+ ),
+ nn.BatchNorm2d(2048, momentum=BN_MOMENTUM),
+ nn.ReLU(inplace=True)
+ )
+
+ return incre_modules, downsamp_modules, final_layer
+
+ def _make_transition_layer(
+ self, num_channels_pre_layer, num_channels_cur_layer):
+ num_branches_cur = len(num_channels_cur_layer)
+ num_branches_pre = len(num_channels_pre_layer)
+
+ transition_layers = []
+ for i in range(num_branches_cur):
+ if i < num_branches_pre:
+ if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
+ transition_layers.append(nn.Sequential(
+ nn.Conv2d(num_channels_pre_layer[i],
+ num_channels_cur_layer[i],
+ 3,
+ 1,
+ 1,
+ bias=False),
+ nn.BatchNorm2d(
+ num_channels_cur_layer[i], momentum=BN_MOMENTUM),
+ nn.ReLU(inplace=True)))
+ else:
+ transition_layers.append(None)
+ else:
+ conv3x3s = []
+ for j in range(i+1-num_branches_pre):
+ inchannels = num_channels_pre_layer[-1]
+ outchannels = num_channels_cur_layer[i] \
+ if j == i-num_branches_pre else inchannels
+ conv3x3s.append(nn.Sequential(
+ nn.Conv2d(
+ inchannels, outchannels, 3, 2, 1, bias=False),
+ nn.BatchNorm2d(outchannels, momentum=BN_MOMENTUM),
+ nn.ReLU(inplace=True)))
+ transition_layers.append(nn.Sequential(*conv3x3s))
+
+ return nn.ModuleList(transition_layers)
+
+ def _make_layer(self, block, inplanes, planes, blocks, stride=1):
+ downsample = None
+ if stride != 1 or inplanes != planes * block.expansion:
+ downsample = nn.Sequential(
+ nn.Conv2d(inplanes, planes * block.expansion,
+ kernel_size=1, stride=stride, bias=False),
+ nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),
+ )
+
+ layers = []
+ layers.append(block(inplanes, planes, stride, downsample))
+ inplanes = planes * block.expansion
+ for i in range(1, blocks):
+ layers.append(block(inplanes, planes))
+
+ return nn.Sequential(*layers)
+
+ def _make_stage(self, layer_config, num_inchannels,
+ multi_scale_output=True):
+ num_modules = layer_config['NUM_MODULES']
+ num_branches = layer_config['NUM_BRANCHES']
+ num_blocks = layer_config['NUM_BLOCKS']
+ num_channels = layer_config['NUM_CHANNELS']
+ block = blocks_dict[layer_config['BLOCK']]
+ fuse_method = layer_config['FUSE_METHOD']
+
+ modules = []
+ for i in range(num_modules):
+ # multi_scale_output is only used last module
+ if not multi_scale_output and i == num_modules - 1:
+ reset_multi_scale_output = False
+ else:
+ reset_multi_scale_output = True
+
+ modules.append(
+ HighResolutionModule(num_branches,
+ block,
+ num_blocks,
+ num_inchannels,
+ num_channels,
+ fuse_method,
+ reset_multi_scale_output)
+ )
+ num_inchannels = modules[-1].get_num_inchannels()
+
+ return nn.Sequential(*modules), num_inchannels
+
+ def forward(self, x):
+ x = self.conv1(x)
+ x = self.bn1(x)
+ x = self.relu(x)
+ x = self.conv2(x)
+ x = self.bn2(x)
+ x = self.relu(x)
+ x = self.layer1(x)
+
+ x_list = []
+ for i in range(self.stage2_cfg['NUM_BRANCHES']):
+ if self.transition1[i] is not None:
+ x_list.append(self.transition1[i](x))
+ else:
+ x_list.append(x)
+ y_list = self.stage2(x_list)
+
+ x_list = []
+ for i in range(self.stage3_cfg['NUM_BRANCHES']):
+ if self.transition2[i] is not None:
+ x_list.append(self.transition2[i](y_list[-1]))
+ else:
+ x_list.append(y_list[i])
+ y_list = self.stage3(x_list)
+
+ x_list = []
+ for i in range(self.stage4_cfg['NUM_BRANCHES']):
+ if self.transition3[i] is not None:
+ x_list.append(self.transition3[i](y_list[-1]))
+ else:
+ x_list.append(y_list[i])
+ y_list = self.stage4(x_list)
+
+ # Classification Head
+ y = self.incre_modules[0](y_list[0])
+ for i in range(len(self.downsamp_modules)):
+ y = self.incre_modules[i+1](y_list[i+1]) + \
+ self.downsamp_modules[i](y)
+
+ y = self.final_layer(y)
+
+ if torch._C._get_tracing_state():
+ y = y.flatten(start_dim=2).mean(dim=2)
+ else:
+ y = F.avg_pool2d(y, kernel_size=y.size()
+ [2:]).view(y.size(0), -1)
+
+ y = self.fc(y)
+
+ return y
+
+ def features(self, x):
+ x = self.conv1(x)
+ x = self.bn1(x)
+ x = self.relu(x)
+ x = self.conv2(x)
+ x = self.bn2(x)
+ x = self.relu(x)
+ x = self.layer1(x)
+
+ x_list = []
+ for i in range(self.stage2_cfg['NUM_BRANCHES']):
+ if self.transition1[i] is not None:
+ x_list.append(self.transition1[i](x))
+ else:
+ x_list.append(x)
+ y_list = self.stage2(x_list)
+
+ x_list = []
+ for i in range(self.stage3_cfg['NUM_BRANCHES']):
+ if self.transition2[i] is not None:
+ x_list.append(self.transition2[i](y_list[-1]))
+ else:
+ x_list.append(y_list[i])
+ y_list = self.stage3(x_list)
+
+ x_list = []
+ for i in range(self.stage4_cfg['NUM_BRANCHES']):
+ if self.transition3[i] is not None:
+ x_list.append(self.transition3[i](y_list[-1]))
+ else:
+ x_list.append(y_list[i])
+ y_list = self.stage4(x_list)
+
+ # Upsampling
+ x0, x1, x2, x3 = y_list
+ x0_h, x0_w = x0.size(2), x0.size(3)
+ x1 = F.upsample(x1, size=(x0_h, x0_w), mode='bilinear')
+ x2 = F.upsample(x2, size=(x0_h, x0_w), mode='bilinear')
+ x3 = F.upsample(x3, size=(x0_h, x0_w), mode='bilinear')
+
+ x_out = torch.cat([x0, x1, x2, x3], 1)
+
+ #print(x_out.size())
+
+ return x_out
+
+ def classifier(self, x):
+ # Classification Head
+ y = self.incre_modules[0](x[0])
+ for i in range(len(self.downsamp_modules)):
+ y = self.incre_modules[i+1](x[i+1]) + \
+ self.downsamp_modules[i](y)
+
+ y = self.final_layer(y)
+
+ if torch._C._get_tracing_state():
+ y = y.flatten(start_dim=2).mean(dim=2)
+ else:
+ y = F.avg_pool2d(y, kernel_size=y.size()
+ [2:]).view(y.size(0), -1)
+
+ y = self.fc(y)
+
+def get_cls_net(config, **kwargs):
+ model = HighResolutionNet(config, **kwargs)
+ return model
diff --git a/clean/image/effort/DeepfakeBench/training/networks/efficientnetb4.py b/clean/image/effort/DeepfakeBench/training/networks/efficientnetb4.py
new file mode 100644
index 0000000000000000000000000000000000000000..e39fa07896903006aa5782d5e26109ef3704cc2e
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/networks/efficientnetb4.py
@@ -0,0 +1,79 @@
+'''
+# author: Zhiyuan Yan
+# email: zhiyuanyan@link.cuhk.edu.cn
+# date: 2023-0706
+
+The code is for EfficientNetB4 backbone.
+'''
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Union
+from efficientnet_pytorch import EfficientNet
+from utils.registry import BACKBONE
+
+
+@BACKBONE.register_module(module_name="efficientnetb4")
+class EfficientNetB4(nn.Module):
+ def __init__(self, efficientnetb4_config):
+ super(EfficientNetB4, self).__init__()
+ """ Constructor
+ Args:
+ efficientnetb4_config: configuration file with the dict format
+ """
+ self.num_classes = efficientnetb4_config["num_classes"]
+ inc = efficientnetb4_config["inc"]
+ self.dropout = efficientnetb4_config["dropout"]
+ self.mode = efficientnetb4_config["mode"]
+
+ # Load the EfficientNet-B4 model without pre-trained weights
+ self.efficientnet = EfficientNet.from_pretrained('efficientnet-b4') # FIXME: load the pretrained weights from online
+ # self.efficientnet = EfficientNet.from_name('efficientnet-b4')
+
+ # Modify the first convolutional layer to accept input tensors with 'inc' channels
+ self.efficientnet._conv_stem = nn.Conv2d(inc, 48, kernel_size=3, stride=2, bias=False)
+
+ # Remove the last layer (the classifier) from the EfficientNet-B4 model
+ self.efficientnet._fc = nn.Identity()
+
+ if self.dropout:
+ # Add dropout layer if specified
+ self.dropout_layer = nn.Dropout(p=self.dropout)
+
+ # Initialize the last_layer layer
+ self.last_layer = nn.Linear(1792, self.num_classes)
+
+ if self.mode == 'adjust_channel':
+ self.adjust_channel = nn.Sequential(
+ nn.Conv2d(1792, 512, 1, 1),
+ nn.BatchNorm2d(512),
+ nn.ReLU(inplace=True),
+ )
+
+ def features(self, x):
+ # Extract features from the EfficientNet-B4 model
+ x = self.efficientnet.extract_features(x)
+ if self.mode == 'adjust_channel':
+ x = self.adjust_channel(x)
+ return x
+
+ def classifier(self, x):
+ x = F.adaptive_avg_pool2d(x, (1, 1))
+ x = x.view(x.size(0), -1)
+
+ # Apply dropout if specified
+ if self.dropout:
+ x = self.dropout_layer(x)
+
+ # Apply last_layer layer
+ x = self.last_layer(x)
+ return x
+
+ def forward(self, x):
+ # Extract features and apply classifier layer
+ x = self.features(x)
+ x = F.adaptive_avg_pool2d(x, (1, 1))
+ x = x.view(x.size(0), -1)
+ x = self.classifier(x)
+ return x
diff --git a/clean/image/effort/DeepfakeBench/training/networks/iresnet.py b/clean/image/effort/DeepfakeBench/training/networks/iresnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..12a454cb9acc227b8968806cf3c47c4145e6304b
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/networks/iresnet.py
@@ -0,0 +1,191 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+__all__ = ['iresnet18', 'iresnet34', 'iresnet50', 'iresnet100', 'iresnet200']
+
+def set_requires_grad(model, val):
+ for p in model.parameters():
+ p.requires_grad = val
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+ """3x3 convolution with padding"""
+ return nn.Conv2d(in_planes,
+ out_planes,
+ kernel_size=3,
+ stride=stride,
+ padding=dilation,
+ groups=groups,
+ bias=False,
+ dilation=dilation)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+ """1x1 convolution"""
+ return nn.Conv2d(in_planes,
+ out_planes,
+ kernel_size=1,
+ stride=stride,
+ bias=False)
+
+
+class IBasicBlock(nn.Module):
+ expansion = 1
+ def __init__(self, inplanes, planes, stride=1, downsample=None,
+ groups=1, base_width=64, dilation=1):
+ super(IBasicBlock, self).__init__()
+ if groups != 1 or base_width != 64:
+ raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+ if dilation > 1:
+ raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+ self.bn1 = nn.BatchNorm2d(inplanes, eps=1e-05,)
+ self.conv1 = conv3x3(inplanes, planes)
+ self.bn2 = nn.BatchNorm2d(planes, eps=1e-05,)
+ self.prelu = nn.PReLU(planes)
+ self.conv2 = conv3x3(planes, planes, stride)
+ self.bn3 = nn.BatchNorm2d(planes, eps=1e-05,)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x):
+ identity = x
+ out = self.bn1(x)
+ out = self.conv1(out)
+ out = self.bn2(out)
+ out = self.prelu(out)
+ out = self.conv2(out)
+ out = self.bn3(out)
+ if self.downsample is not None:
+ identity = self.downsample(x)
+ out += identity
+ return out
+
+
+class IResNet(nn.Module):
+ def __init__(self,
+ block, layers, dropout=0, num_features=512, zero_init_residual=False,
+ groups=1, width_per_group=64, replace_stride_with_dilation=None, fp16=False, fc_scale=7*7):
+ super(IResNet, self).__init__()
+ self.fp16 = fp16
+ self.inplanes = 64
+ self.dilation = 1
+ if replace_stride_with_dilation is None:
+ replace_stride_with_dilation = [False, False, False]
+ if len(replace_stride_with_dilation) != 3:
+ raise ValueError("replace_stride_with_dilation should be None "
+ "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+ self.groups = groups
+ self.base_width = width_per_group
+ self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(self.inplanes, eps=1e-05)
+ self.prelu = nn.PReLU(self.inplanes)
+ self.layer1 = self._make_layer(block, 64, layers[0], stride=2)
+ self.layer2 = self._make_layer(block,
+ 128,
+ layers[1],
+ stride=2,
+ dilate=replace_stride_with_dilation[0])
+
+ self.layer3 = self._make_layer(block,
+ 256,
+ layers[2],
+ stride=2,
+ dilate=replace_stride_with_dilation[1])
+ set_requires_grad(self.layer1, False)
+ set_requires_grad(self.layer2, False)
+ set_requires_grad(self.layer3, False)
+ self.layer4 = self._make_layer(block,
+ 512,
+ layers[3],
+ stride=2,
+ dilate=replace_stride_with_dilation[2])
+ self.bn2 = nn.BatchNorm2d(512 * block.expansion, eps=1e-05,)
+ self.dropout = nn.Dropout(p=dropout, inplace=True)
+ self.fc = nn.Linear(512 * block.expansion * fc_scale, num_features)
+ self.features = nn.BatchNorm1d(num_features, eps=1e-05)
+ nn.init.constant_(self.features.weight, 1.0)
+ self.features.weight.requires_grad = False
+
+ for m in self.modules():
+ if isinstance(m, nn.Conv2d):
+ nn.init.normal_(m.weight, 0, 0.1)
+ elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+ nn.init.constant_(m.weight, 1)
+ nn.init.constant_(m.bias, 0)
+
+ if zero_init_residual:
+ for m in self.modules():
+ if isinstance(m, IBasicBlock):
+ nn.init.constant_(m.bn2.weight, 0)
+
+ def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+ downsample = None
+ previous_dilation = self.dilation
+ if dilate:
+ self.dilation *= stride
+ stride = 1
+ if stride != 1 or self.inplanes != planes * block.expansion:
+ downsample = nn.Sequential(
+ conv1x1(self.inplanes, planes * block.expansion, stride),
+ nn.BatchNorm2d(planes * block.expansion, eps=1e-05, ),
+ )
+ layers = []
+ layers.append(
+ block(self.inplanes, planes, stride, downsample, self.groups,
+ self.base_width, previous_dilation))
+ self.inplanes = planes * block.expansion
+ for _ in range(1, blocks):
+ layers.append(
+ block(self.inplanes,
+ planes,
+ groups=self.groups,
+ base_width=self.base_width,
+ dilation=self.dilation))
+
+ return nn.Sequential(*layers)
+
+ def forward(self, x):
+ with torch.cuda.amp.autocast(self.fp16):
+ x = self.conv1(x)
+ x = self.bn1(x)
+ x = self.prelu(x)
+ x = self.layer1(x)
+ x = self.layer2(x)
+ x = self.layer3(x)
+ x = self.layer4(x)
+ x = self.bn2(x)
+ x = self.dropout(x)
+ x = F.avg_pool2d(x, kernel_size=2, stride=2, padding=0)
+ return x
+
+
+def _iresnet(arch, block, layers, pretrained, progress, **kwargs):
+ model = IResNet(block, layers, **kwargs)
+ if pretrained:
+ raise ValueError()
+ return model
+
+
+def iresnet18(pretrained=False, progress=True, **kwargs):
+ return _iresnet('iresnet18', IBasicBlock, [2, 2, 2, 2], pretrained,
+ progress, **kwargs)
+
+
+def iresnet34(pretrained=False, progress=True, **kwargs):
+ return _iresnet('iresnet34', IBasicBlock, [3, 4, 6, 3], pretrained,
+ progress, **kwargs)
+
+
+def iresnet50(pretrained=False, progress=True, **kwargs):
+ return _iresnet('iresnet50', IBasicBlock, [3, 4, 14, 3], pretrained,
+ progress, **kwargs)
+
+
+def iresnet100(pretrained=False, progress=True, **kwargs):
+ return _iresnet('iresnet100', IBasicBlock, [3, 13, 30, 3], pretrained,
+ progress, **kwargs)
+
+
+def iresnet200(pretrained=False, progress=True, **kwargs):
+ return _iresnet('iresnet200', IBasicBlock, [6, 26, 60, 6], pretrained,
+ progress, **kwargs)
diff --git a/clean/image/effort/DeepfakeBench/training/networks/mesonet.py b/clean/image/effort/DeepfakeBench/training/networks/mesonet.py
new file mode 100644
index 0000000000000000000000000000000000000000..bb4df43794f301ea6802ca16e1265ab45d183ab8
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/networks/mesonet.py
@@ -0,0 +1,189 @@
+'''
+# author: Zhiyuan Yan
+# email: zhiyuanyan@link.cuhk.edu.cn
+# date: 2023-0706
+
+The code is mainly modified from the below link:
+https://github.com/HongguLiu/MesoNet-Pytorch
+'''
+
+import os
+import argparse
+import logging
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import torch.utils.model_zoo as model_zoo
+from torch.nn import init
+from typing import Union
+from utils.registry import BACKBONE
+
+logger = logging.getLogger(__name__)
+
+@BACKBONE.register_module(module_name="meso4")
+class Meso4(nn.Module):
+ def __init__(self, meso4_config):
+ super(Meso4, self).__init__()
+ self.num_classes = meso4_config["num_classes"]
+ inc = meso4_config["inc"]
+ self.conv1 = nn.Conv2d(inc, 8, 3, padding=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(8)
+ self.relu = nn.ReLU(inplace=True)
+ self.leakyrelu = nn.LeakyReLU(0.1)
+
+ self.conv2 = nn.Conv2d(8, 8, 5, padding=2, bias=False)
+ self.bn2 = nn.BatchNorm2d(16)
+ self.conv3 = nn.Conv2d(8, 16, 5, padding=2, bias=False)
+ self.conv4 = nn.Conv2d(16, 16, 5, padding=2, bias=False)
+ self.maxpooling1 = nn.MaxPool2d(kernel_size=(2, 2))
+ self.maxpooling2 = nn.MaxPool2d(kernel_size=(4, 4))
+ #flatten: x = x.view(x.size(0), -1)
+ self.dropout = nn.Dropout2d(0.5)
+ self.fc1 = nn.Linear(16*8*8, 16)
+ self.fc2 = nn.Linear(16, self.num_classes)
+
+
+ def features(self, input):
+ x = self.conv1(input) #(8, 256, 256)
+ x = self.relu(x)
+ x = self.bn1(x)
+ x = self.maxpooling1(x) #(8, 128, 128)
+
+ x = self.conv2(x) #(8, 128, 128)
+ x = self.relu(x)
+ x = self.bn1(x)
+ x = self.maxpooling1(x) #(8, 64, 64)
+
+ x = self.conv3(x) #(16, 64, 64)
+ x = self.relu(x)
+ x = self.bn2(x)
+ x = self.maxpooling1(x) #(16, 32, 32)
+
+ x = self.conv4(x) #(16, 32, 32)
+ x = self.relu(x)
+ x = self.bn2(x)
+ x = self.maxpooling2(x) #(16, 8, 8)
+ x = x.view(x.size(0), -1) #(Batch, 16*8*8)
+
+ return x
+
+ def classifier(self, feature):
+ out = self.dropout(feature)
+ out = self.fc1(out) #(Batch, 16)
+ out = self.leakyrelu(out)
+ out = self.dropout(out)
+ out = self.fc2(out)
+ return out
+
+ def forward(self, input):
+ x = self.features(input)
+ out = self.classifier(x)
+ return out, x
+
+
+@BACKBONE.register_module(module_name="meso4Inception")
+class MesoInception4(nn.Module):
+ def __init__(self, mesoInception4_config):
+ super(MesoInception4, self).__init__()
+ self.num_classes = mesoInception4_config["num_classes"]
+ inc = mesoInception4_config["inc"]
+ #InceptionLayer1
+ self.Incption1_conv1 = nn.Conv2d(3, 1, 1, padding=0, bias=False)
+ self.Incption1_conv2_1 = nn.Conv2d(3, 4, 1, padding=0, bias=False)
+ self.Incption1_conv2_2 = nn.Conv2d(4, 4, 3, padding=1, bias=False)
+ self.Incption1_conv3_1 = nn.Conv2d(3, 4, 1, padding=0, bias=False)
+ self.Incption1_conv3_2 = nn.Conv2d(4, 4, 3, padding=2, dilation=2, bias=False)
+ self.Incption1_conv4_1 = nn.Conv2d(3, 2, 1, padding=0, bias=False)
+ self.Incption1_conv4_2 = nn.Conv2d(2, 2, 3, padding=3, dilation=3, bias=False)
+ self.Incption1_bn = nn.BatchNorm2d(11)
+
+
+ #InceptionLayer2
+ self.Incption2_conv1 = nn.Conv2d(11, 2, 1, padding=0, bias=False)
+ self.Incption2_conv2_1 = nn.Conv2d(11, 4, 1, padding=0, bias=False)
+ self.Incption2_conv2_2 = nn.Conv2d(4, 4, 3, padding=1, bias=False)
+ self.Incption2_conv3_1 = nn.Conv2d(11, 4, 1, padding=0, bias=False)
+ self.Incption2_conv3_2 = nn.Conv2d(4, 4, 3, padding=2, dilation=2, bias=False)
+ self.Incption2_conv4_1 = nn.Conv2d(11, 2, 1, padding=0, bias=False)
+ self.Incption2_conv4_2 = nn.Conv2d(2, 2, 3, padding=3, dilation=3, bias=False)
+ self.Incption2_bn = nn.BatchNorm2d(12)
+
+ #Normal Layer
+ self.conv1 = nn.Conv2d(12, 16, 5, padding=2, bias=False)
+ self.relu = nn.ReLU(inplace=True)
+ self.leakyrelu = nn.LeakyReLU(0.1)
+ self.bn1 = nn.BatchNorm2d(16)
+ self.maxpooling1 = nn.MaxPool2d(kernel_size=(2, 2))
+
+ self.conv2 = nn.Conv2d(16, 16, 5, padding=2, bias=False)
+ self.maxpooling2 = nn.MaxPool2d(kernel_size=(4, 4))
+
+ self.dropout = nn.Dropout2d(0.5)
+ self.fc1 = nn.Linear(16*8*8, 16)
+ self.fc2 = nn.Linear(16, self.num_classes)
+
+
+ #InceptionLayer
+ def InceptionLayer1(self, input):
+ x1 = self.Incption1_conv1(input)
+ x2 = self.Incption1_conv2_1(input)
+ x2 = self.Incption1_conv2_2(x2)
+ x3 = self.Incption1_conv3_1(input)
+ x3 = self.Incption1_conv3_2(x3)
+ x4 = self.Incption1_conv4_1(input)
+ x4 = self.Incption1_conv4_2(x4)
+ y = torch.cat((x1, x2, x3, x4), 1)
+ y = self.Incption1_bn(y)
+ y = self.maxpooling1(y)
+
+ return y
+
+ def InceptionLayer2(self, input):
+ x1 = self.Incption2_conv1(input)
+ x2 = self.Incption2_conv2_1(input)
+ x2 = self.Incption2_conv2_2(x2)
+ x3 = self.Incption2_conv3_1(input)
+ x3 = self.Incption2_conv3_2(x3)
+ x4 = self.Incption2_conv4_1(input)
+ x4 = self.Incption2_conv4_2(x4)
+ y = torch.cat((x1, x2, x3, x4), 1)
+ y = self.Incption2_bn(y)
+ y = self.maxpooling1(y)
+
+ return y
+
+
+ def features(self, input):
+ x = self.InceptionLayer1(input) #(Batch, 11, 128, 128)
+ x = self.InceptionLayer2(x) #(Batch, 12, 64, 64)
+
+ x = self.conv1(x) #(Batch, 16, 64 ,64)
+ x = self.relu(x)
+ x = self.bn1(x)
+ x = self.maxpooling1(x) #(Batch, 16, 32, 32)
+
+ x = self.conv2(x) #(Batch, 16, 32, 32)
+ x = self.relu(x)
+ x = self.bn1(x)
+ x = self.maxpooling2(x) #(Batch, 16, 8, 8)
+
+ x = x.view(x.size(0), -1) #(Batch, 16*8*8)
+
+ return x
+
+ def classifier(self, feature):
+
+ out = self.dropout(feature)
+ out = self.fc1(out) #(Batch, 16)
+ out = self.leakyrelu(out)
+ out = self.dropout(out)
+ out = self.fc2(out)
+ return out
+
+ def forward(self, input):
+ x = self.features(input)
+ out = self.classifier(x)
+ return out, x
diff --git a/clean/image/effort/DeepfakeBench/training/networks/resnet.py b/clean/image/effort/DeepfakeBench/training/networks/resnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..64d649ab427b06444a9b4b17c53869640ebbe2b6
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/networks/resnet.py
@@ -0,0 +1,501 @@
+# -*- coding: utf-8 -*-
+"""
+Created on 18-5-21 下午5:26
+
+@author: ronghuaiyang
+"""
+import torch
+import torch.nn as nn
+import math
+import torch.utils.model_zoo as model_zoo
+import torch.nn.utils.weight_norm as weight_norm
+import torch.nn.functional as F
+
+
+# __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
+# 'resnet152']
+
+
+model_urls = {
+ 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+ 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+ 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+ 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+ 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+ """3x3 convolution with padding"""
+ return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+ padding=1, bias=False)
+
+
+class AdaIN(nn.Module):
+ def __init__(self, eps=1e-5):
+ super().__init__()
+ self.eps = eps
+ # self.l1 = nn.Linear(num_classes, in_channel*4, bias=True) #bias is good :)
+
+ def c_norm(self, x, bs, ch, eps=1e-7):
+ # assert isinstance(x, torch.cuda.FloatTensor)
+ x_var = x.var(dim=-1) + eps
+ x_std = x_var.sqrt().view(bs, ch, 1, 1)
+ x_mean = x.mean(dim=-1).view(bs, ch, 1, 1)
+ return x_std, x_mean
+
+ def forward(self, x, y):
+ assert x.size(0)==y.size(0)
+ size = x.size()
+ bs, ch = size[:2]
+ x_ = x.view(bs, ch, -1)
+ y_ = y.reshape(bs, ch, -1)
+ x_std, x_mean = self.c_norm(x_, bs, ch, eps=self.eps)
+ y_std, y_mean = self.c_norm(y_, bs, ch, eps=self.eps)
+ out = ((x - x_mean.expand(size)) / x_std.expand(size)) \
+ * y_std.expand(size) + y_mean.expand(size)
+ return out
+
+
+class BasicBlock(nn.Module):
+ expansion = 1
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None):
+ super(BasicBlock, self).__init__()
+ self.conv1 = conv3x3(inplanes, planes, stride)
+ self.bn1 = nn.BatchNorm2d(planes)
+ self.relu = nn.ReLU(inplace=True)
+ self.conv2 = conv3x3(planes, planes)
+ self.bn2 = nn.BatchNorm2d(planes)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x):
+ residual = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+
+ if self.downsample is not None:
+ residual = self.downsample(x)
+
+ out += residual
+ out = self.relu(out)
+
+ return out
+
+
+class BasicBlock_adain(nn.Module):
+ expansion = 1
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None):
+ super(BasicBlock_adain, self).__init__()
+ self.conv1 = conv3x3(inplanes, planes, stride)
+ self.adain1 = AdaIN()
+ self.relu = nn.ReLU(inplace=True)
+ self.conv2 = conv3x3(planes, planes)
+ self.adain2 = AdaIN()
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, feat): # x is content, c is style
+ x, c = feat
+ residual = x
+
+ x = self.conv1(x)
+ out = self.adain1(x, c)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.adain2(out, c)
+
+ if self.downsample is not None:
+ residual = self.downsample(residual)
+
+ out += residual
+ out = self.relu(out)
+
+ return (out, c)
+
+
+class IRBlock(nn.Module):
+ expansion = 1
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None, use_se=True):
+ super(IRBlock, self).__init__()
+ self.bn0 = nn.BatchNorm2d(inplanes)
+ self.conv1 = conv3x3(inplanes, inplanes)
+ self.bn1 = nn.BatchNorm2d(inplanes)
+ self.prelu = nn.PReLU()
+ self.conv2 = conv3x3(inplanes, planes, stride)
+ self.bn2 = nn.BatchNorm2d(planes)
+ self.downsample = downsample
+ self.stride = stride
+ self.use_se = use_se
+ if self.use_se:
+ self.se = SEBlock(planes)
+
+ def forward(self, x):
+ residual = x
+ out = self.bn0(x)
+ out = self.conv1(out)
+ out = self.bn1(out)
+ out = self.prelu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+ if self.use_se:
+ out = self.se(out)
+
+ if self.downsample is not None:
+ residual = self.downsample(x)
+
+ out += residual
+ out = self.prelu(out)
+
+ return out
+
+
+class IRBlock_3conv(nn.Module):
+ expansion = 1
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None, use_se=True):
+ super(IRBlock_3conv, self).__init__()
+ self.bn0 = nn.BatchNorm2d(inplanes)
+ self.conv1 = conv3x3(inplanes, inplanes)
+ self.bn1 = nn.BatchNorm2d(inplanes)
+ self.prelu1 = nn.PReLU()
+ self.conv2 = conv3x3(inplanes, planes, stride)
+ self.bn2 = nn.BatchNorm2d(planes)
+ self.prelu2 = nn.PReLU()
+ self.conv3 = conv3x3(planes, planes)
+ self.bn3 = nn.BatchNorm2d(planes)
+ self.downsample = downsample
+ self.stride = stride
+ self.use_se = use_se
+ if self.use_se:
+ self.se = SEBlock(planes)
+ self.prelu = nn.PReLU()
+
+ def forward(self, x):
+ residual = x
+ out = self.bn0(x)
+ out = self.conv1(out)
+ out = self.bn1(out)
+ out = self.prelu1(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+ out = self.prelu2(out)
+
+ out = self.conv3(out)
+ out = self.bn3(out)
+ if self.use_se:
+ out = self.se(out)
+
+ if self.downsample is not None:
+ residual = self.downsample(x)
+
+ out += residual
+ out = self.prelu(out)
+
+ return out
+
+
+class Bottleneck(nn.Module):
+ expansion = 4
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None):
+ super(Bottleneck, self).__init__()
+ self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(planes)
+ self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+ padding=1, bias=False)
+ self.bn2 = nn.BatchNorm2d(planes)
+ self.conv3 = nn.Conv2d(
+ planes, planes * self.expansion, kernel_size=1, bias=False)
+ self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+ self.relu = nn.ReLU(inplace=True)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x):
+ residual = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+ out = self.relu(out)
+
+ out = self.conv3(out)
+ out = self.bn3(out)
+
+ if self.downsample is not None:
+ residual = self.downsample(x)
+
+ out += residual
+ out = self.relu(out)
+
+ return out
+
+
+class SEBlock(nn.Module):
+ def __init__(self, channel, reduction=16):
+ super(SEBlock, self).__init__()
+ self.avg_pool = nn.AdaptiveAvgPool2d(1)
+ self.fc = nn.Sequential(
+ nn.Linear(channel, channel // reduction),
+ nn.PReLU(),
+ nn.Linear(channel // reduction, channel),
+ nn.Sigmoid()
+ )
+
+ def forward(self, x):
+ b, c, _, _ = x.size()
+ y = self.avg_pool(x).view(b, c)
+ y = self.fc(y).view(b, c, 1, 1)
+ return x * y
+
+
+class ResNetFace(nn.Module):
+ def __init__(self, block, layers, use_se=True, inc=3):
+ self.inplanes = 64
+ self.use_se = use_se
+ super(ResNetFace, self).__init__()
+ self.conv1 = nn.Conv2d(inc, 64, kernel_size=3, padding=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(64)
+ self.prelu = nn.PReLU()
+ self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
+ self.layer1 = self._make_layer(block, 64, layers[0])
+ self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+ self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+ self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+ self.bn4 = nn.BatchNorm2d(512)
+ #self.dropout = nn.Dropout()
+ self.fc5 = nn.Linear(512 * 8 * 8, 512)
+ #self.bn5 = nn.BatchNorm1d(512)
+
+ for m in self.modules():
+ if isinstance(m, nn.Conv2d):
+ nn.init.xavier_normal_(m.weight)
+ elif isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm1d):
+ nn.init.constant_(m.weight, 1)
+ nn.init.constant_(m.bias, 0)
+ elif isinstance(m, nn.Linear):
+ nn.init.xavier_normal_(m.weight)
+ nn.init.constant_(m.bias, 0)
+
+ def _make_layer(self, block, planes, blocks, stride=1):
+ downsample = None
+ if stride != 1 or self.inplanes != planes * block.expansion:
+ downsample = nn.Sequential(
+ nn.Conv2d(self.inplanes, planes * block.expansion,
+ kernel_size=1, stride=stride, bias=False),
+ nn.BatchNorm2d(planes * block.expansion),
+ )
+ layers = []
+ layers.append(block(self.inplanes, planes, stride,
+ downsample, use_se=self.use_se))
+ self.inplanes = planes
+ for i in range(1, blocks):
+ layers.append(block(self.inplanes, planes, use_se=self.use_se))
+
+ return nn.Sequential(*layers)
+
+ def features(self, x):
+ x = self.conv1(x)
+ x = self.bn1(x)
+ x = self.prelu(x)
+ x = self.maxpool(x)
+
+ x = self.layer1(x)
+ x = self.layer2(x)
+ x = self.layer3(x)
+ x = self.layer4(x)
+ x = self.bn4(x)
+
+ return x
+
+ def classifier(self, x):
+ x = x.view(x.size(0), -1)
+ x = self.fc5(x)
+
+ return x
+
+ def forward(self, x):
+ x = self.conv1(x)
+ x = self.bn1(x)
+ x = self.prelu(x)
+ x = self.maxpool(x)
+
+ x = self.layer1(x)
+ x = self.layer2(x)
+ x = self.layer3(x)
+ x = self.layer4(x)
+ x = self.bn4(x)
+ #x = self.dropout(x)
+ x = x.view(x.size(0), -1)
+ x = self.fc5(x)
+ #x = self.bn5(x)
+
+ return x
+
+
+class ResNet(nn.Module):
+
+ def __init__(self, block, layers, basedim=32, inc=1):
+ self.inplanes = basedim
+ super(ResNet, self).__init__()
+ # self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+ # bias=False)
+ self.conv1 = nn.Conv2d(inc, self.inplanes, kernel_size=3, stride=1, padding=1,
+ bias=False)
+ self.bn1 = nn.BatchNorm2d(self.inplanes)
+ self.relu = nn.ReLU(inplace=True)
+ # self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+ self.layer1 = self._make_layer(block, basedim, layers[0], stride=2)
+ self.layer2 = self._make_layer(block, 2*basedim, layers[1], stride=2)
+ self.layer3 = self._make_layer(block, 4*basedim, layers[2], stride=2)
+ self.layer4 = self._make_layer(block, 8*basedim, layers[3], stride=2)
+ # self.avgpool = nn.AvgPool2d(8, stride=1)
+ # self.fc = nn.Linear(512 * block.expansion, num_classes)
+ self.fc5 = nn.Linear(512 * 8 * 8, 512)
+
+ for m in self.modules():
+ if isinstance(m, nn.Conv2d):
+ nn.init.kaiming_normal_(
+ m.weight, mode='fan_out', nonlinearity='relu')
+ elif isinstance(m, nn.BatchNorm2d):
+ nn.init.constant_(m.weight, 1)
+ nn.init.constant_(m.bias, 0)
+
+ def _make_layer(self, block, planes, blocks, stride=1):
+ downsample = None
+ if stride != 1 or self.inplanes != planes * block.expansion:
+ downsample = nn.Sequential(
+ nn.Conv2d(self.inplanes, planes * block.expansion,
+ kernel_size=1, stride=stride, bias=False),
+ nn.BatchNorm2d(planes * block.expansion),
+ )
+
+ layers = []
+ layers.append(block(self.inplanes, planes, stride, downsample))
+ self.inplanes = planes * block.expansion
+ for i in range(1, blocks):
+ layers.append(block(self.inplanes, planes))
+
+ return nn.Sequential(*layers)
+
+ def features(self, x):
+ x = self.conv1(x)
+ x = self.bn1(x)
+ x = self.relu(x)
+ # x = self.maxpool(x)
+
+ x = self.layer1(x)
+ x = self.layer2(x)
+ x = self.layer3(x)
+ x = self.layer4(x)
+
+ return x
+
+ def classifier(self, x):
+ x = x.view(x.size(0), -1)
+ x = self.fc5(x)
+
+ return x
+
+ def forward(self, x):
+ x = self.conv1(x)
+ x = self.bn1(x)
+ x = self.relu(x)
+ # x = self.maxpool(x)
+
+ x = self.layer1(x)
+ x = self.layer2(x)
+ x = self.layer3(x)
+ x = self.layer4(x)
+ # x = nn.AvgPool2d(kernel_size=x.size()[2:])(x)
+ # x = self.avgpool(x)
+ x = x.view(x.size(0), -1)
+ x = self.fc5(x)
+
+ return x
+
+
+def resnet18(pretrained=False, **kwargs):
+ """Constructs a ResNet-18 model.
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ """
+ model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
+ if pretrained:
+ model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
+ return model
+
+
+def resnet34(pretrained=False, **kwargs):
+ """Constructs a ResNet-34 model.
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ """
+ model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
+ if pretrained:
+ model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
+ return model
+
+
+def resnet50(pretrained=False, **kwargs):
+ """Constructs a ResNet-50 model.
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ """
+ model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
+ if pretrained:
+ model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
+ return model
+
+
+def resnet101(pretrained=False, **kwargs):
+ """Constructs a ResNet-101 model.
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ """
+ model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
+ if pretrained:
+ model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
+ return model
+
+
+def resnet152(pretrained=False, **kwargs):
+ """Constructs a ResNet-152 model.
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ """
+ model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
+ if pretrained:
+ model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
+ return model
+
+
+def resnet_face18(use_se=True, **kwargs):
+ model = ResNetFace(IRBlock, [2, 2, 2, 2], use_se=use_se, **kwargs)
+ return model
+
+
+def resnet_face62(use_se=True, **kwargs):
+ model = ResNetFace(IRBlock_3conv, [3, 4, 10, 3], use_se=use_se, **kwargs)
+ return model
+
+if __name__ == "__main__":
+ net = HR_resnet()
+ dummy = torch.rand(10,3,256,256)
+ x = net(dummy)
+ print('output:', x.size())
diff --git a/clean/image/effort/DeepfakeBench/training/networks/resnet34.py b/clean/image/effort/DeepfakeBench/training/networks/resnet34.py
new file mode 100644
index 0000000000000000000000000000000000000000..e6fc2d2c42e1e6ae26f3b92ab0f419e0db183392
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/networks/resnet34.py
@@ -0,0 +1,60 @@
+'''
+# author: Zhiyuan Yan
+# email: zhiyuanyan@link.cuhk.edu.cn
+# date: 2023-0706
+
+The code is for ResNet34 backbone.
+'''
+
+import os
+import logging
+from typing import Union
+import torch
+import torchvision
+import torch.nn as nn
+import torch.nn.functional as F
+from utils.registry import BACKBONE
+
+logger = logging.getLogger(__name__)
+
+@BACKBONE.register_module(module_name="resnet34")
+class ResNet34(nn.Module):
+ def __init__(self, resnet_config):
+ super(ResNet34, self).__init__()
+ """ Constructor
+ Args:
+ resnet_config: configuration file with the dict format
+ """
+ self.num_classes = resnet_config["num_classes"]
+ inc = resnet_config["inc"]
+ self.mode = resnet_config["mode"]
+
+ # Define layers of the backbone
+ resnet = torchvision.models.resnet34(pretrained=True) # FIXME: download the pretrained weights from online
+ # resnet.conv1 = nn.Conv2d(inc, 64, kernel_size=7, stride=2, padding=3, bias=False)
+ self.resnet = torch.nn.Sequential(*list(resnet.children())[:-2])
+ self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+ self.fc = nn.Linear(512, self.num_classes)
+
+ if self.mode == 'adjust_channel':
+ self.adjust_channel = nn.Sequential(
+ nn.Conv2d(512, 512, 1, 1),
+ nn.BatchNorm2d(512),
+ nn.ReLU(inplace=True),
+ )
+
+
+ def features(self, inp):
+ x = self.resnet(inp)
+ return x
+
+ def classifier(self, features):
+ x = self.avgpool(features)
+ x = x.view(x.size(0), -1)
+ x = self.fc(x)
+ return x
+
+ def forward(self, inp):
+ x = self.features(inp)
+ out = self.classifier(x)
+ return out
diff --git a/clean/image/effort/DeepfakeBench/training/networks/vgg.py b/clean/image/effort/DeepfakeBench/training/networks/vgg.py
new file mode 100644
index 0000000000000000000000000000000000000000..da74f9c6a1ca1a93470cefeb68781ddee05d3460
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/networks/vgg.py
@@ -0,0 +1,143 @@
+"""A VGG-based perceptual loss function for PyTorch."""
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torchvision import models, transforms
+
+
+class Lambda(nn.Module):
+ """Wraps a callable in an :class:`nn.Module` without registering it."""
+
+ def __init__(self, func):
+ super().__init__()
+ object.__setattr__(self, 'forward', func)
+
+ def extra_repr(self):
+ return getattr(self.forward, '__name__', type(self.forward).__name__) + '()'
+
+
+class WeightedLoss(nn.ModuleList):
+ """A weighted combination of multiple loss functions."""
+
+ def __init__(self, losses, weights, verbose=False):
+ super().__init__()
+ for loss in losses:
+ self.append(loss if isinstance(loss, nn.Module) else Lambda(loss))
+ self.weights = weights
+ self.verbose = verbose
+
+ def _print_losses(self, losses):
+ for i, loss in enumerate(losses):
+ print(f'({i}) {type(self[i]).__name__}: {loss.item()}')
+
+ def forward(self, *args, **kwargs):
+ losses = []
+ for loss, weight in zip(self, self.weights):
+ losses.append(loss(*args, **kwargs) * weight)
+ if self.verbose:
+ self._print_losses(losses)
+ return sum(losses)
+
+
+class TVLoss(nn.Module):
+ """Total variation loss (Lp penalty on image gradient magnitude).
+ The input must be 4D. If a target (second parameter) is passed in, it is
+ ignored.
+ ``p=1`` yields the vectorial total variation norm. It is a generalization
+ of the originally proposed (isotropic) 2D total variation norm (see
+ (see https://en.wikipedia.org/wiki/Total_variation_denoising) for color
+ images. On images with a single channel it is equal to the 2D TV norm.
+ ``p=2`` yields a variant that is often used for smoothing out noise in
+ reconstructions of images from neural network feature maps (see Mahendran
+ and Vevaldi, "Understanding Deep Image Representations by Inverting
+ Them", https://arxiv.org/abs/1412.0035)
+ :attr:`reduction` can be set to ``'mean'``, ``'sum'``, or ``'none'``
+ similarly to the loss functions in :mod:`torch.nn`. The default is
+ ``'mean'``.
+ """
+
+ def __init__(self, p, reduction='mean', eps=1e-8):
+ super().__init__()
+ if p not in {1, 2}:
+ raise ValueError('p must be 1 or 2')
+ if reduction not in {'mean', 'sum', 'none'}:
+ raise ValueError("reduction must be 'mean', 'sum', or 'none'")
+ self.p = p
+ self.reduction = reduction
+ self.eps = eps
+
+ def forward(self, input, target=None):
+ input = F.pad(input, (0, 1, 0, 1), 'replicate')
+ x_diff = input[..., :-1, :-1] - input[..., :-1, 1:]
+ y_diff = input[..., :-1, :-1] - input[..., 1:, :-1]
+ diff = x_diff**2 + y_diff**2
+ if self.p == 1:
+ diff = (diff + self.eps).mean(dim=1, keepdims=True).sqrt()
+ if self.reduction == 'mean':
+ return diff.mean()
+ if self.reduction == 'sum':
+ return diff.sum()
+ return diff
+
+
+class VGGLoss(nn.Module):
+ """Computes the VGG perceptual loss between two batches of images.
+ The input and target must be 4D tensors with three channels
+ ``(B, 3, H, W)`` and must have equivalent shapes. Pixel values should be
+ normalized to the range 0–1.
+ The VGG perceptual loss is the mean squared difference between the features
+ computed for the input and target at layer :attr:`layer` (default 8, or
+ ``relu2_2``) of the pretrained model specified by :attr:`model` (either
+ ``'vgg16'`` (default) or ``'vgg19'``).
+ If :attr:`shift` is nonzero, a random shift of at most :attr:`shift`
+ pixels in both height and width will be applied to all images in the input
+ and target. The shift will only be applied when the loss function is in
+ training mode, and will not be applied if a precomputed feature map is
+ supplied as the target.
+ :attr:`reduction` can be set to ``'mean'``, ``'sum'``, or ``'none'``
+ similarly to the loss functions in :mod:`torch.nn`. The default is
+ ``'mean'``.
+ :meth:`get_features()` may be used to precompute the features for the
+ target, to speed up the case where inputs are compared against the same
+ target over and over. To use the precomputed features, pass them in as
+ :attr:`target` and set :attr:`target_is_features` to :code:`True`.
+ Instances of :class:`VGGLoss` must be manually converted to the same
+ device and dtype as their inputs.
+ """
+
+ models = {'vgg16': models.vgg16, 'vgg19': models.vgg19}
+
+ def __init__(self, model='vgg16', layer=8, shift=0, reduction='mean'):
+ super().__init__()
+ self.instancenorm = nn.InstanceNorm2d(512, affine=False)
+ self.shift = shift
+ self.reduction = reduction
+ self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
+ std=[0.229, 0.224, 0.225])
+ self.model = self.models[model](pretrained=True).features[:layer+1]
+ self.model.eval()
+ self.model.requires_grad_(False)
+
+ def get_features(self, input):
+ return self.model(self.normalize(input))
+
+ def train(self, mode=True):
+ self.training = mode
+
+ def forward(self, input, target, target_is_features=False):
+ if target_is_features:
+ input_feats = self.get_features(input)
+ target_feats = target
+ else:
+ sep = input.shape[0]
+ batch = torch.cat([input, target])
+ if self.shift and self.training:
+ padded = F.pad(batch, [self.shift] * 4, mode='replicate')
+ batch = transforms.RandomCrop(batch.shape[2:])(padded)
+ feats = self.get_features(batch)
+ input_feats, target_feats = feats[:sep], feats[sep:]
+ # input_feats, target_feats = \
+ # self.instancenorm(input_feats), \
+ # self.instancenorm(target_feats)
+ return F.mse_loss(input_feats, target_feats, reduction=self.reduction)
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/networks/xception.py b/clean/image/effort/DeepfakeBench/training/networks/xception.py
new file mode 100644
index 0000000000000000000000000000000000000000..8c61d9d37d959d0d99100f139ab64d8b5d62140a
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/networks/xception.py
@@ -0,0 +1,272 @@
+'''
+# author: Zhiyuan Yan
+# email: zhiyuanyan@link.cuhk.edu.cn
+# date: 2023-0706
+
+The code is mainly modified from GitHub link below:
+https://github.com/ondyari/FaceForensics/blob/master/classification/network/xception.py
+'''
+
+import os
+import argparse
+import logging
+
+import math
+import torch
+# import pretrainedmodels
+import torch.nn as nn
+import torch.nn.functional as F
+
+import torch.utils.model_zoo as model_zoo
+from torch.nn import init
+from typing import Union
+from utils.registry import BACKBONE
+
+logger = logging.getLogger(__name__)
+
+
+
+class SeparableConv2d(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False):
+ super(SeparableConv2d, self).__init__()
+
+ self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size,
+ stride, padding, dilation, groups=in_channels, bias=bias)
+ self.pointwise = nn.Conv2d(
+ in_channels, out_channels, 1, 1, 0, 1, 1, bias=bias)
+
+ def forward(self, x):
+ x = self.conv1(x)
+ x = self.pointwise(x)
+ return x
+
+
+class Block(nn.Module):
+ def __init__(self, in_filters, out_filters, reps, strides=1, start_with_relu=True, grow_first=True):
+ super(Block, self).__init__()
+
+ if out_filters != in_filters or strides != 1:
+ self.skip = nn.Conv2d(in_filters, out_filters,
+ 1, stride=strides, bias=False)
+ self.skipbn = nn.BatchNorm2d(out_filters)
+ else:
+ self.skip = None
+
+ self.relu = nn.ReLU(inplace=True)
+ rep = []
+
+ filters = in_filters
+ if grow_first: # whether the number of filters grows first
+ rep.append(self.relu)
+ rep.append(SeparableConv2d(in_filters, out_filters,
+ 3, stride=1, padding=1, bias=False))
+ rep.append(nn.BatchNorm2d(out_filters))
+ filters = out_filters
+
+ for i in range(reps-1):
+ rep.append(self.relu)
+ rep.append(SeparableConv2d(filters, filters,
+ 3, stride=1, padding=1, bias=False))
+ rep.append(nn.BatchNorm2d(filters))
+
+ if not grow_first:
+ rep.append(self.relu)
+ rep.append(SeparableConv2d(in_filters, out_filters,
+ 3, stride=1, padding=1, bias=False))
+ rep.append(nn.BatchNorm2d(out_filters))
+
+ if not start_with_relu:
+ rep = rep[1:]
+ else:
+ rep[0] = nn.ReLU(inplace=False)
+
+ if strides != 1:
+ rep.append(nn.MaxPool2d(3, strides, 1))
+ self.rep = nn.Sequential(*rep)
+
+ def forward(self, inp):
+ x = self.rep(inp)
+
+ if self.skip is not None:
+ skip = self.skip(inp)
+ skip = self.skipbn(skip)
+ else:
+ skip = inp
+
+ x += skip
+ return x
+
+def add_gaussian_noise(ins, mean=0, stddev=0.2):
+ noise = ins.data.new(ins.size()).normal_(mean, stddev)
+ return ins + noise
+
+
+@BACKBONE.register_module(module_name="xception")
+class Xception(nn.Module):
+ """
+ Xception optimized for the ImageNet dataset, as specified in
+ https://arxiv.org/pdf/1610.02357.pdf
+ """
+
+ def __init__(self, xception_config):
+ """ Constructor
+ Args:
+ xception_config: configuration file with the dict format
+ """
+ super(Xception, self).__init__()
+ self.num_classes = xception_config["num_classes"]
+ self.mode = xception_config["mode"]
+ inc = xception_config["inc"]
+ dropout = xception_config["dropout"]
+
+ # Entry flow
+ self.conv1 = nn.Conv2d(inc, 32, 3, 2, 0, bias=False)
+
+ self.bn1 = nn.BatchNorm2d(32)
+ self.relu = nn.ReLU(inplace=True)
+
+ self.conv2 = nn.Conv2d(32, 64, 3, bias=False)
+ self.bn2 = nn.BatchNorm2d(64)
+ # do relu here
+
+ self.block1 = Block(
+ 64, 128, 2, 2, start_with_relu=False, grow_first=True)
+ self.block2 = Block(
+ 128, 256, 2, 2, start_with_relu=True, grow_first=True)
+ self.block3 = Block(
+ 256, 728, 2, 2, start_with_relu=True, grow_first=True)
+
+ # middle flow
+ self.block4 = Block(
+ 728, 728, 3, 1, start_with_relu=True, grow_first=True)
+ self.block5 = Block(
+ 728, 728, 3, 1, start_with_relu=True, grow_first=True)
+ self.block6 = Block(
+ 728, 728, 3, 1, start_with_relu=True, grow_first=True)
+ self.block7 = Block(
+ 728, 728, 3, 1, start_with_relu=True, grow_first=True)
+
+ self.block8 = Block(
+ 728, 728, 3, 1, start_with_relu=True, grow_first=True)
+ self.block9 = Block(
+ 728, 728, 3, 1, start_with_relu=True, grow_first=True)
+ self.block10 = Block(
+ 728, 728, 3, 1, start_with_relu=True, grow_first=True)
+ self.block11 = Block(
+ 728, 728, 3, 1, start_with_relu=True, grow_first=True)
+
+ # Exit flow
+ self.block12 = Block(
+ 728, 1024, 2, 2, start_with_relu=True, grow_first=False)
+
+ self.conv3 = SeparableConv2d(1024, 1536, 3, 1, 1)
+ self.bn3 = nn.BatchNorm2d(1536)
+
+ # do relu here
+ self.conv4 = SeparableConv2d(1536, 2048, 3, 1, 1)
+ self.bn4 = nn.BatchNorm2d(2048)
+
+ self.last_linear = nn.Linear(2048, self.num_classes)
+ if dropout:
+ self.last_linear = nn.Sequential(
+ nn.Dropout(p=dropout),
+ nn.Linear(2048, self.num_classes)
+ )
+
+ self.adjust_channel = nn.Sequential(
+ nn.Conv2d(2048, 512, 1, 1),
+ nn.BatchNorm2d(512),
+ nn.ReLU(inplace=True),
+ )
+
+ def fea_part1_0(self, x):
+ x = self.conv1(x)
+ x = self.bn1(x)
+ x = self.relu(x)
+
+ return x
+
+ def fea_part1_1(self, x):
+
+ x = self.conv2(x)
+ x = self.bn2(x)
+ x = self.relu(x)
+
+ return x
+
+ def fea_part1(self, x):
+ x = self.conv1(x)
+ x = self.bn1(x)
+ x = self.relu(x)
+
+ x = self.conv2(x)
+ x = self.bn2(x)
+ x = self.relu(x)
+
+ return x
+
+ def fea_part2(self, x):
+ x = self.block1(x)
+ x = self.block2(x)
+ x = self.block3(x)
+
+ return x
+
+ def fea_part3(self, x):
+ if self.mode == "shallow_xception":
+ return x
+ else:
+ x = self.block4(x)
+ x = self.block5(x)
+ x = self.block6(x)
+ x = self.block7(x)
+ return x
+
+ def fea_part4(self, x):
+ if self.mode == "shallow_xception":
+ x = self.block12(x)
+ else:
+ x = self.block8(x)
+ x = self.block9(x)
+ x = self.block10(x)
+ x = self.block11(x)
+ x = self.block12(x)
+ return x
+
+ def fea_part5(self, x):
+ x = self.conv3(x)
+ x = self.bn3(x)
+ x = self.relu(x)
+
+ x = self.conv4(x)
+ x = self.bn4(x)
+
+ return x
+
+ def features(self, input):
+ x = self.fea_part1(input)
+
+ x = self.fea_part2(x)
+ x = self.fea_part3(x)
+ x = self.fea_part4(x)
+
+ x = self.fea_part5(x)
+
+ if self.mode == 'adjust_channel':
+ x = self.adjust_channel(x)
+
+ return x
+
+ def classifier(self, features):
+ x = self.relu(features)
+
+ if len(x.shape) == 4:
+ x = F.adaptive_avg_pool2d(x, (1, 1))
+ x = x.view(x.size(0), -1)
+ out = self.last_linear(x)
+ return out
+
+ def forward(self, input):
+ x = self.features(input)
+ out = self.classifier(x)
+ return out, x
diff --git a/clean/image/effort/DeepfakeBench/training/networks/xception_ffd.py b/clean/image/effort/DeepfakeBench/training/networks/xception_ffd.py
new file mode 100644
index 0000000000000000000000000000000000000000..5f23ddae50da43390081167dd199d31a1118c889
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/networks/xception_ffd.py
@@ -0,0 +1,267 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import os
+import sys
+
+class SeparableConv2d(nn.Module):
+ def __init__(self, c_in, c_out, ks, stride=1, padding=0, dilation=1, bias=False):
+ super(SeparableConv2d, self).__init__()
+ self.c = nn.Conv2d(c_in, c_in, ks, stride, padding, dilation, groups=c_in, bias=bias)
+ self.pointwise = nn.Conv2d(c_in, c_out, 1, 1, 0, 1, 1, bias=bias)
+
+ def forward(self, x):
+ x = self.c(x)
+ x = self.pointwise(x)
+ return x
+
+class Block(nn.Module):
+ def __init__(self, c_in, c_out, reps, stride=1, start_with_relu=True, grow_first=True):
+ super(Block, self).__init__()
+
+ self.skip = None
+ self.skip_bn = None
+ if c_out != c_in or stride!= 1:
+ self.skip = nn.Conv2d(c_in, c_out, 1, stride=stride, bias=False)
+ self.skip_bn = nn.BatchNorm2d(c_out)
+
+ self.relu = nn.ReLU(inplace=True)
+
+ rep = []
+ c = c_in
+ if grow_first:
+ rep.append(self.relu)
+ rep.append(SeparableConv2d(c_in, c_out, 3, stride=1, padding=1, bias=False))
+ rep.append(nn.BatchNorm2d(c_out))
+ c = c_out
+
+ for i in range(reps - 1):
+ rep.append(self.relu)
+ rep.append(SeparableConv2d(c, c, 3, stride=1, padding=1, bias=False))
+ rep.append(nn.BatchNorm2d(c))
+
+ if not grow_first:
+ rep.append(self.relu)
+ rep.append(SeparableConv2d(c_in, c_out, 3, stride=1, padding=1, bias=False))
+ rep.append(nn.BatchNorm2d(c_out))
+
+ if not start_with_relu:
+ rep = rep[1:]
+ else:
+ rep[0] = nn.ReLU(inplace=False)
+
+ if stride != 1:
+ rep.append(nn.MaxPool2d(3, stride, 1))
+ self.rep = nn.Sequential(*rep)
+
+ def forward(self, inp):
+ x = self.rep(inp)
+
+ if self.skip is not None:
+ y = self.skip(inp)
+ y = self.skip_bn(y)
+ else:
+ y = inp
+
+ x += y
+ return x
+
+class RegressionMap(nn.Module):
+ def __init__(self, c_in):
+ super(RegressionMap, self).__init__()
+ self.c = SeparableConv2d(c_in, 1, 3, stride=1, padding=1, bias=False)
+ self.s = nn.Sigmoid()
+
+ def forward(self, x):
+ mask = self.c(x)
+ mask = self.s(mask)
+ return mask, None
+
+class TemplateMap(nn.Module):
+ def __init__(self, c_in, templates):
+ super(TemplateMap, self).__init__()
+ self.c = Block(c_in, 364, 2, 2, start_with_relu=True, grow_first=False)
+ self.l = nn.Linear(364, 10)
+ self.relu = nn.ReLU(inplace=True)
+
+ self.templates = templates
+
+ def forward(self, x):
+ v = self.c(x)
+ v = self.relu(v)
+ v = F.adaptive_avg_pool2d(v, (1,1))
+ v = v.view(v.size(0), -1)
+ v = self.l(v)
+ mask = torch.mm(v, self.templates.reshape(10,361))
+ mask = mask.reshape(x.shape[0], 1, 19, 19)
+
+ return mask, v
+
+class PCATemplateMap(nn.Module):
+ def __init__(self, templates):
+ super(PCATemplateMap, self).__init__()
+ self.templates = templates
+
+ def forward(self, x):
+ fe = x.view(x.shape[0], x.shape[1], x.shape[2]*x.shape[3])
+ fe = torch.transpose(fe, 1, 2)
+ mu = torch.mean(fe, 2, keepdim=True)
+ fea_diff = fe - mu
+
+ cov_fea = torch.bmm(fea_diff, torch.transpose(fea_diff, 1, 2))
+ B = self.templates.reshape(1, 10, 361).repeat(x.shape[0], 1, 1)
+ D = torch.bmm(torch.bmm(B, cov_fea), torch.transpose(B, 1, 2))
+ eigen_value, eigen_vector = D.symeig(eigenvectors=True)
+ index = torch.tensor([9]).cuda()
+ eigen = torch.index_select(eigen_vector, 2, index)
+
+ v = eigen.squeeze(-1)
+ mask = torch.mm(v, self.templates.reshape(10, 361))
+ mask = mask.reshape(x.shape[0], 1, 19, 19)
+ return mask, v
+
+class Xception(nn.Module):
+ """
+ Xception optimized for the ImageNet dataset, as specified in
+ https://arxiv.org/pdf/1610.02357.pdf
+ """
+ def __init__(self, maptype, templates, num_classes=1000):
+ super(Xception, self).__init__()
+ self.num_classes = num_classes
+
+ self.conv1 = nn.Conv2d(3, 32, 3,2, 0, bias=False)
+ self.bn1 = nn.BatchNorm2d(32)
+ self.relu = nn.ReLU(inplace=True)
+
+ self.conv2 = nn.Conv2d(32,64,3,bias=False)
+ self.bn2 = nn.BatchNorm2d(64)
+
+ self.block1=Block(64,128,2,2,start_with_relu=False,grow_first=True)
+ self.block2=Block(128,256,2,2,start_with_relu=True,grow_first=True)
+ self.block3=Block(256,728,2,2,start_with_relu=True,grow_first=True)
+ self.block4=Block(728,728,3,1,start_with_relu=True,grow_first=True)
+ self.block5=Block(728,728,3,1,start_with_relu=True,grow_first=True)
+ self.block6=Block(728,728,3,1,start_with_relu=True,grow_first=True)
+ self.block7=Block(728,728,3,1,start_with_relu=True,grow_first=True)
+ self.block8=Block(728,728,3,1,start_with_relu=True,grow_first=True)
+ self.block9=Block(728,728,3,1,start_with_relu=True,grow_first=True)
+ self.block10=Block(728,728,3,1,start_with_relu=True,grow_first=True)
+ self.block11=Block(728,728,3,1,start_with_relu=True,grow_first=True)
+ self.block12=Block(728,1024,2,2,start_with_relu=True,grow_first=False)
+
+ self.conv3 = SeparableConv2d(1024,1536,3,1,1)
+ self.bn3 = nn.BatchNorm2d(1536)
+
+ self.conv4 = SeparableConv2d(1536,2048,3,1,1)
+ self.bn4 = nn.BatchNorm2d(2048)
+
+ self.last_linear = nn.Linear(2048, num_classes)
+
+ if maptype == 'none':
+ self.map = [1, None]
+ elif maptype == 'reg':
+ self.map = RegressionMap(728)
+ elif maptype == 'tmp':
+ self.map = TemplateMap(728, templates)
+ elif maptype == 'pca_tmp':
+ self.map = PCATemplateMap(728)
+ else:
+ print('Unknown map type: `{0}`'.format(maptype))
+ sys.exit()
+
+ def features(self, input):
+ x = self.conv1(input)
+ x = self.bn1(x)
+ x = self.relu(x)
+
+ x = self.conv2(x)
+ x = self.bn2(x)
+ x = self.relu(x)
+
+ x = self.block1(x)
+ x = self.block2(x)
+ x = self.block3(x)
+ x = self.block4(x)
+ x = self.block5(x)
+ x = self.block6(x)
+ x = self.block7(x)
+ mask, vec = self.map(x)
+ x = x * mask
+ x = self.block8(x)
+ x = self.block9(x)
+ x = self.block10(x)
+ x = self.block11(x)
+ x = self.block12(x)
+ x = self.conv3(x)
+ x = self.bn3(x)
+ x = self.relu(x)
+
+ x = self.conv4(x)
+ x = self.bn4(x)
+ return x, mask, vec
+
+ def logits(self, features):
+ x = self.relu(features)
+ x = F.adaptive_avg_pool2d(x, (1, 1))
+ x = x.view(x.size(0), -1)
+ x = self.last_linear(x)
+ return x
+
+ def forward(self, input):
+ x, mask, vec = self.features(input)
+ x = self.logits(x)
+ return x, mask, vec
+
+def init_weights(m):
+ classname = m.__class__.__name__
+ if classname.find('SeparableConv2d') != -1:
+ m.c.weight.data.normal_(0.0, 0.01)
+ if m.c.bias is not None:
+ m.c.bias.data.fill_(0)
+ m.pointwise.weight.data.normal_(0.0, 0.01)
+ if m.pointwise.bias is not None:
+ m.pointwise.bias.data.fill_(0)
+ elif classname.find('Conv') != -1 or classname.find('Linear') != -1:
+ m.weight.data.normal_(0.0, 0.01)
+ if m.bias is not None:
+ m.bias.data.fill_(0)
+ elif classname.find('BatchNorm') != -1:
+ m.weight.data.normal_(1.0, 0.01)
+ m.bias.data.fill_(0)
+ elif classname.find('LSTM') != -1:
+ for i in m._parameters:
+ if i.__class__.__name__.find('weight') != -1:
+ i.data.normal_(0.0, 0.01)
+ elif i.__class__.__name__.find('bias') != -1:
+ i.bias.data.fill_(0)
+
+class Model:
+ def __init__(self, maptype='None', templates=None, num_classes=2, load_pretrain=True):
+ model = Xception(maptype, templates, num_classes=num_classes)
+ if load_pretrain:
+ state_dict = torch.load('./xception-b5690688.pth')
+ for name, weights in state_dict:
+ if 'pointwise' in name:
+ state_dict[name] = weights.unsqueeze(-1).unsqueeze(-1)
+ del state_dict['fc.weight']
+ del state_dict['fc.bias']
+ model.load_state_dict(state_dict, False)
+ else:
+ model.apply(init_weights)
+ self.model = model
+
+ def save(self, epoch, optim, model_dir):
+ state = {'net': self.model.state_dict(), 'optim': optim.state_dict()}
+ torch.save(state, '{0}/{1:06d}.tar'.format(model_dir, epoch))
+ print('Saved model `{0}`'.format(epoch))
+
+ def load(self, epoch, model_dir):
+ filename = '{0}{1:06d}.tar'.format(model_dir, epoch)
+ print('Loading model from {0}'.format(filename))
+ if os.path.exists(filename):
+ state = torch.load(filename)
+ self.model.load_state_dict(state['net'])
+ else:
+ print('Failed to load model from {0}'.format(filename))
+
diff --git a/clean/image/effort/DeepfakeBench/training/optimizor/LinearLR.py b/clean/image/effort/DeepfakeBench/training/optimizor/LinearLR.py
new file mode 100644
index 0000000000000000000000000000000000000000..80bc70dbae46bb9f76aa65afe6f4a1b95dd25619
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/optimizor/LinearLR.py
@@ -0,0 +1,20 @@
+import torch
+from torch.optim import SGD
+from torch.optim.lr_scheduler import _LRScheduler
+
+class LinearDecayLR(_LRScheduler):
+ def __init__(self, optimizer, n_epoch, start_decay, last_epoch=-1):
+ self.start_decay=start_decay
+ self.n_epoch=n_epoch
+ super(LinearDecayLR, self).__init__(optimizer, last_epoch)
+
+ def get_lr(self):
+ last_epoch = self.last_epoch
+ n_epoch=self.n_epoch
+ b_lr=self.base_lrs[0]
+ start_decay=self.start_decay
+ if last_epoch>start_decay:
+ lr=b_lr-b_lr/(n_epoch-start_decay)*(last_epoch-start_decay)
+ else:
+ lr=b_lr
+ return [lr]
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/optimizor/SAM.py b/clean/image/effort/DeepfakeBench/training/optimizor/SAM.py
new file mode 100644
index 0000000000000000000000000000000000000000..7b8d1dc52726ffea22553ce96a6e0d37a902fbff
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/optimizor/SAM.py
@@ -0,0 +1,77 @@
+# borrowed from
+
+import torch
+
+import torch
+import torch.nn as nn
+
+def disable_running_stats(model):
+ def _disable(module):
+ if isinstance(module, nn.BatchNorm2d):
+ module.backup_momentum = module.momentum
+ module.momentum = 0
+
+ model.apply(_disable)
+
+def enable_running_stats(model):
+ def _enable(module):
+ if isinstance(module, nn.BatchNorm2d) and hasattr(module, "backup_momentum"):
+ module.momentum = module.backup_momentum
+
+ model.apply(_enable)
+
+class SAM(torch.optim.Optimizer):
+ def __init__(self, params, base_optimizer, rho=0.05, **kwargs):
+ assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}"
+
+ defaults = dict(rho=rho, **kwargs)
+ super(SAM, self).__init__(params, defaults)
+
+ self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
+ self.param_groups = self.base_optimizer.param_groups
+
+ @torch.no_grad()
+ def first_step(self, zero_grad=False):
+ grad_norm = self._grad_norm()
+ for group in self.param_groups:
+ scale = group["rho"] / (grad_norm + 1e-12)
+
+ for p in group["params"]:
+ if p.grad is None: continue
+ e_w = p.grad * scale.to(p)
+ p.add_(e_w) # climb to the local maximum "w + e(w)"
+ self.state[p]["e_w"] = e_w
+
+ if zero_grad: self.zero_grad()
+
+ @torch.no_grad()
+ def second_step(self, zero_grad=False):
+ for group in self.param_groups:
+ for p in group["params"]:
+ if p.grad is None: continue
+ p.sub_(self.state[p]["e_w"]) # get back to "w" from "w + e(w)"
+
+ self.base_optimizer.step() # do the actual "sharpness-aware" update
+
+ if zero_grad: self.zero_grad()
+
+ @torch.no_grad()
+ def step(self, closure=None):
+ assert closure is not None, "Sharpness Aware Minimization requires closure, but it was not provided"
+ closure = torch.enable_grad()(closure) # the closure should do a full forward-backward pass
+
+ self.first_step(zero_grad=True)
+ closure()
+ self.second_step()
+
+ def _grad_norm(self):
+ shared_device = self.param_groups[0]["params"][0].device # put everything on the same device, in case of model parallelism
+ norm = torch.norm(
+ torch.stack([
+ p.grad.norm(p=2).to(shared_device)
+ for group in self.param_groups for p in group["params"]
+ if p.grad is not None
+ ]),
+ p=2
+ )
+ return norm
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/test.py b/clean/image/effort/DeepfakeBench/training/test.py
new file mode 100644
index 0000000000000000000000000000000000000000..13a22ed9533b18454609b41ee81dbbca9a2a848e
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/test.py
@@ -0,0 +1,253 @@
+"""
+eval pretained model.
+"""
+import os
+import numpy as np
+from os.path import join
+import cv2
+import random
+import datetime
+import time
+import yaml
+import pickle
+from tqdm import tqdm
+from copy import deepcopy
+from PIL import Image as pil_image
+from metrics.utils import get_test_metrics
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.backends.cudnn as cudnn
+import torch.nn.functional as F
+import torch.utils.data
+import torch.optim as optim
+
+from dataset.abstract_dataset import DeepfakeAbstractBaseDataset
+
+from trainer.trainer import Trainer
+from detectors import DETECTOR
+from metrics.base_metrics_class import Recorder
+from collections import defaultdict
+
+import argparse
+from logger import create_logger
+
+parser = argparse.ArgumentParser(description='Process some paths.')
+parser.add_argument('--detector_path', type=str,
+ default='./training/config/detector/effort.yaml',
+ help='path to detector YAML file')
+parser.add_argument("--test_dataset", nargs="+")
+parser.add_argument('--weights_path', type=str,
+ default='./weights/effort_ckpt.pth')
+#parser.add_argument("--lmdb", action='store_true', default=False)
+args = parser.parse_args()
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+on_2060 = "2060" in torch.cuda.get_device_name()
+def init_seed(config):
+ if config['manualSeed'] is None:
+ config['manualSeed'] = random.randint(1, 10000)
+ random.seed(config['manualSeed'])
+ torch.manual_seed(config['manualSeed'])
+ if config['cuda']:
+ torch.cuda.manual_seed_all(config['manualSeed'])
+
+
+def prepare_testing_data(config):
+ def get_test_data_loader(config, test_name):
+ # update the config dictionary with the specific testing dataset
+ config = config.copy() # create a copy of config to avoid altering the original one
+ config['test_dataset'] = test_name # specify the current test dataset
+ test_set = DeepfakeAbstractBaseDataset(
+ config=config,
+ mode='test',
+ )
+ test_data_loader = \
+ torch.utils.data.DataLoader(
+ dataset=test_set,
+ batch_size=config['test_batchSize'],
+ shuffle=False,
+ num_workers=int(config['workers']),
+ collate_fn=test_set.collate_fn,
+ drop_last=False
+ )
+ return test_data_loader
+
+ test_data_loaders = {}
+ for one_test_name in config['test_dataset']:
+ test_data_loaders[one_test_name] = get_test_data_loader(config, one_test_name)
+ return test_data_loaders
+
+
+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 test_one_dataset(model, data_loader):
+ prediction_lists = []
+ feature_lists = []
+ label_lists = []
+ for i, data_dict in tqdm(enumerate(data_loader), total=len(data_loader)):
+ # get data
+ data, label, mask, landmark = \
+ data_dict['image'], data_dict['label'], data_dict['mask'], data_dict['landmark']
+ label = torch.where(data_dict['label'] != 0, 1, 0)
+ # move data to GPU
+ data_dict['image'], data_dict['label'] = data.to(device), label.to(device)
+ if mask is not None:
+ data_dict['mask'] = mask.to(device)
+ if landmark is not None:
+ data_dict['landmark'] = landmark.to(device)
+
+ # model forward without considering gradient computation
+ predictions = inference(model, data_dict)
+ label_lists += list(data_dict['label'].cpu().detach().numpy())
+ prediction_lists += list(predictions['prob'].cpu().detach().numpy())
+ feature_lists += list(predictions['feat'].cpu().detach().numpy())
+
+ return np.array(prediction_lists), np.array(label_lists),np.array(feature_lists)
+
+def test_epoch(model, test_data_loaders):
+ # set model to eval mode
+ model.eval()
+
+ # define test recorder
+ metrics_all_datasets = {}
+
+ # testing for all test data
+ keys = test_data_loaders.keys()
+ for key in keys:
+ data_dict = test_data_loaders[key].dataset.data_dict
+ # compute loss for each dataset
+ predictions_nps, label_nps,feat_nps = test_one_dataset(model, test_data_loaders[key])
+
+ # compute metric for each dataset
+ metric_one_dataset = get_test_metrics(y_pred=predictions_nps, y_true=label_nps,
+ img_names=data_dict['image'])
+ metrics_all_datasets[key] = metric_one_dataset
+
+ # info for each dataset
+ tqdm.write(f"dataset: {key}")
+ for k, v in metric_one_dataset.items():
+ tqdm.write(f"{k}: {v}")
+
+ return metrics_all_datasets
+
+@torch.no_grad()
+def inference(model, data_dict):
+ predictions = model(data_dict, inference=True)
+ return predictions
+
+
+
+
+def main():
+ # parse options and load config
+ with open(args.detector_path, 'r') as f:
+ config = yaml.safe_load(f)
+ with open('./training/config/test_config.yaml', 'r') as f:
+ config2 = yaml.safe_load(f)
+ config.update(config2)
+ if on_2060:
+ config['lmdb_dir'] = r'I:\transform_2_lmdb'
+ config['train_batchSize'] = 10
+ config['workers'] = 0
+ else:
+ config['workers'] = 8
+ config['lmdb_dir'] = r'/mnt/chongqinggeminiceph1fs/geminicephfs/mm-base-vision/jikangcheng/data/LMDBs'
+ weights_path = None
+ # If arguments are provided, they will overwrite the yaml settings
+ if args.test_dataset:
+ config['test_dataset'] = args.test_dataset
+ if args.weights_path:
+ config['weights_path'] = args.weights_path
+ weights_path = args.weights_path
+
+ # init seed
+ init_seed(config)
+
+ # set cudnn benchmark if needed
+ if config['cudnn']:
+ cudnn.benchmark = True
+
+ # prepare the testing data loader
+ test_data_loaders = prepare_testing_data(config)
+
+ # prepare the model (detector)
+ model_class = DETECTOR[config['model_name']]
+ model = model_class(config).to(device)
+ total_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+ print(f"Total number of trainable parameters in the model: {total_trainable_params}")
+ epoch = 0
+ if weights_path:
+ try:
+ epoch = int(weights_path.split('/')[-1].split('.')[0].split('_')[2])
+ except:
+ epoch = 0
+ ckpt = torch.load(weights_path, map_location=device)
+ if 'state_dict' in ckpt:
+ ckpt = ckpt['state_dict']
+
+ # # 加载模型的状态字典
+ # model_dict = model.state_dict()
+ # new_ckpt={}
+ # for key in ckpt.keys():
+ # # 替换键
+ # new_key = key.replace('common_encoder_f','student_encoder')
+ # # 将旧的值复制到新的键下
+ # new_ckpt[new_key] = ckpt[key]
+ # # 获取ckpt和model的key集合
+ # ckpt_keys = set(new_ckpt.keys())
+ # model_keys = set(model_dict.keys())
+ #
+ # # 找出共同的key
+ # common_keys = ckpt_keys & model_keys
+ # print("Common keys:")
+ # for key in common_keys:
+ # print(key)
+ #
+ # # 找出只在ckpt中的key
+ # ckpt_unique_keys = ckpt_keys - model_keys
+ # print("\nKeys only in ckpt:")
+ # for key in ckpt_unique_keys:
+ # print(key)
+ #
+ # # 找出只在model中的key
+ # model_unique_keys = model_keys - ckpt_keys
+ # print("\nKeys only in model:")
+ # for key in model_unique_keys:
+ # print(key)
+
+ # 创建一个新的字典，删除module前缀
+ new_weights = {}
+ for key, value in ckpt.items():
+ new_key = key.replace('module.', '') # 删除module前缀
+ #new_key = 'backbone.' + new_key # 删除module前缀
+ # if 'base_model.' in new_key:
+ # new_key = new_key.replace('base_model.', 'backbone.')
+ # if 'classifier.' in new_key:
+ # new_key = new_key.replace('classifier.', 'head.')
+ new_weights[new_key] = value
+
+
+ model.load_state_dict(new_weights, strict=True)
+ print('===> Load checkpoint done!')
+ else:
+ print('Fail to load the pre-trained weights')
+
+ # clip_rank_results = analyze_clip_effective_rank(model.backbone)
+
+ # for layer_name, rank in clip_rank_results.items():
+ # print(f"Layer: {layer_name}, Effective Rank: {rank}")
+
+
+ # start testing
+ best_metric = test_epoch(model, test_data_loaders)
+ print('===> Test Done!')
+
+if __name__ == '__main__':
+ main()
diff --git a/clean/image/effort/DeepfakeBench/training/train.py b/clean/image/effort/DeepfakeBench/training/train.py
new file mode 100644
index 0000000000000000000000000000000000000000..607a613e1d0e6f7ae2fed743cb116cfb654722e0
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/train.py
@@ -0,0 +1,277 @@
+# author: Zhiyuan Yan
+# email: zhiyuanyan@link.cuhk.edu.cn
+# date: 2023-03-30
+# description: training code.
+
+import os
+import argparse
+from os.path import join
+import cv2
+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
+
+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 optimizor.SAM import SAM
+from optimizor.LinearLR import LinearDecayLR
+
+from trainer.trainer import Trainer
+from detectors import DETECTOR
+from dataset import *
+from metrics.utils import parse_metric_for_print
+from logger import create_logger, RankFilter
+
+
+parser = argparse.ArgumentParser(description='Process some paths.')
+parser.add_argument('--detector_path', type=str,
+ default='/data/home/zhiyuanyan/DeepfakeBenchv2/training/config/detector/effort.yaml',
+ help='path to detector YAML file')
+parser.add_argument("--train_dataset", nargs="+")
+parser.add_argument("--test_dataset", nargs="+")
+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)
+args = parser.parse_args()
+torch.cuda.set_device(args.local_rank)
+
+
+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 prepare_training_data(config):
+ # Only use the blending dataset class in training
+ train_set = DeepfakeAbstractBaseDataset(
+ config=config,
+ mode='train',
+ )
+ if config['ddp']:
+ sampler = DistributedSampler(train_set)
+ train_data_loader = \
+ torch.utils.data.DataLoader(
+ dataset=train_set,
+ batch_size=config['train_batchSize'],
+ num_workers=int(config['workers']),
+ collate_fn=train_set.collate_fn,
+ sampler=sampler
+ )
+ else:
+ train_data_loader = \
+ torch.utils.data.DataLoader(
+ dataset=train_set,
+ batch_size=config['train_batchSize'],
+ shuffle=True,
+ num_workers=int(config['workers']),
+ collate_fn=train_set.collate_fn,
+ )
+ return train_data_loader
+
+
+def prepare_testing_data(config):
+ def get_test_data_loader(config, test_name):
+ # update the config dictionary with the specific testing dataset
+ config = config.copy() # create a copy of config to avoid altering the original one
+ config['test_dataset'] = test_name # specify the current test dataset
+
+ test_set = DeepfakeAbstractBaseDataset(
+ config=config,
+ mode='test',
+ )
+
+ test_data_loader = \
+ torch.utils.data.DataLoader(
+ dataset=test_set,
+ batch_size=config['test_batchSize'],
+ shuffle=False,
+ num_workers=int(config['workers']),
+ collate_fn=test_set.collate_fn,
+ drop_last = (test_name=='DeepFakeDetection'),
+ )
+
+ return test_data_loader
+
+ test_data_loaders = {}
+ for one_test_name in config['test_dataset']:
+ test_data_loaders[one_test_name] = get_test_data_loader(config, one_test_name)
+ return test_data_loaders
+
+
+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 main():
+ # parse options and load config
+ with open(args.detector_path, 'r') as f:
+ config = yaml.safe_load(f)
+ with open('./training/config/train_config.yaml', 'r') as f:
+ config2 = yaml.safe_load(f)
+ config.update(config2)
+ config['local_rank']=args.local_rank
+ if config['dry_run']:
+ config['nEpochs'] = 0
+ config['save_feat']=False
+ # If arguments are provided, they will overwrite the yaml settings
+ if args.train_dataset:
+ config['train_dataset'] = args.train_dataset
+ if args.test_dataset:
+ config['test_dataset'] = args.test_dataset
+ config['save_ckpt'] = args.save_ckpt
+ config['save_feat'] = args.save_feat
+ if config['lmdb']:
+ config['dataset_json_folder'] = 'preprocessing/dataset_json_v3'
+ # create logger
+ logger_path = config['log_dir']
+ os.makedirs(logger_path, exist_ok=True)
+ logger = create_logger(os.path.join(logger_path, 'training.log'))
+ logger.info('Save log to {}'.format(logger_path))
+ config['ddp']= args.ddp
+ # 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)
+
+ # 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))
+ # prepare the training data loader
+ train_data_loader = prepare_training_data(config)
+
+ # prepare the testing data loader
+ test_data_loaders = prepare_testing_data(config)
+
+ # 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)
+
+ # start training
+ for epoch in range(config['start_epoch'], config['nEpochs'] + 1):
+ trainer.model.epoch = epoch
+ best_metric = trainer.train_epoch(
+ epoch=epoch,
+ train_data_loader=train_data_loader,
+ test_data_loaders=test_data_loaders,
+ )
+ if best_metric is not None:
+ logger.info(f"===> Epoch[{epoch}] end with testing {metric_scoring}: {parse_metric_for_print(best_metric)}!")
+ logger.info("Stop Training on best Testing metric {}".format(parse_metric_for_print(best_metric)))
+ # 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()
+
+
+
+if __name__ == '__main__':
+ main()
diff --git a/clean/image/effort/DeepfakeBench/training/trainer/__init__.py b/clean/image/effort/DeepfakeBench/training/trainer/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c96d5a5f7af796495ee04bbc29924f239a89ded8
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/trainer/__init__.py
@@ -0,0 +1,2 @@
+from trainer.trainer import Trainer
+from utils.registry import TRAINER
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/trainer/base_trainer.py b/clean/image/effort/DeepfakeBench/training/trainer/base_trainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..e1402752b994f82702349e5c7c8d98b1fc7d1519
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/trainer/base_trainer.py
@@ -0,0 +1,50 @@
+import datetime
+from copy import deepcopy
+from abc import ABC, abstractmethod
+
+
+class BaseTrainer(ABC):
+ """
+ """
+
+ def __init__(
+ self,
+ config,
+ model,
+ optimizer,
+ scheduler,
+ writer,
+ ):
+ # check if all the necessary components are implemented
+ if config is None or model is None or optimizer is None or scheduler is None or writer is None:
+ raise NotImplementedError("config, model, optimizier, scheduler, and tensorboard writer must be implemented")
+
+ self.config = config
+ self.model = model
+ self.optimizer = optimizer
+ self.scheduler = scheduler
+ self.writer = writer
+
+ @abstractmethod
+ def speed_up(self):
+ pass
+
+ @abstractmethod
+ def setTrain(self):
+ pass
+
+ @abstractmethod
+ def setEval(self):
+ pass
+
+ @abstractmethod
+ def load_ckpt(self, model_path):
+ pass
+
+ @abstractmethod
+ def save_ckpt(self, dataset, epoch, iters, best=False):
+ pass
+
+ @abstractmethod
+ def inference(self, data_dict):
+ pass
diff --git a/clean/image/effort/DeepfakeBench/training/trainer/trainer.py b/clean/image/effort/DeepfakeBench/training/trainer/trainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..07063c231056887cdaa84ed62a1892199e521401
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/trainer/trainer.py
@@ -0,0 +1,467 @@
+# author: Zhiyuan Yan
+# email: zhiyuanyan@link.cuhk.edu.cn
+# date: 2023-03-30
+# description: trainer
+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)
+
+import pickle
+import datetime
+import logging
+import numpy as np
+from copy import deepcopy
+from collections import defaultdict
+from tqdm import tqdm
+import time
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.nn import DataParallel
+from torch.utils.tensorboard import SummaryWriter
+from metrics.base_metrics_class import Recorder
+from torch.optim.swa_utils import AveragedModel, SWALR
+from torch import distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+from sklearn import metrics
+from metrics.utils import get_test_metrics
+
+FFpp_pool=['FaceForensics++','FF-DF','FF-F2F','FF-FS','FF-NT']#
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class Trainer(object):
+ def __init__(
+ self,
+ config,
+ model,
+ optimizer,
+ scheduler,
+ logger,
+ metric_scoring='auc',
+ time_now = datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S'),
+ swa_model=None
+ ):
+ # check if all the necessary components are implemented
+ if config is None or model is None or optimizer is None or logger is None:
+ raise ValueError("config, model, optimizier, logger, and tensorboard writer must be implemented")
+
+ self.config = config
+ self.model = model
+ self.optimizer = optimizer
+ self.scheduler = scheduler
+ self.swa_model = swa_model
+ self.writers = {} # dict to maintain different tensorboard writers for each dataset and metric
+ self.logger = logger
+ self.metric_scoring = metric_scoring
+ # maintain the best metric of all epochs
+ self.best_metrics_all_time = defaultdict(
+ lambda: defaultdict(lambda: float('-inf')
+ if self.metric_scoring != 'eer' else float('inf'))
+ )
+ self.speed_up() # move model to GPU
+
+ # get current time
+ self.timenow = time_now
+ # create directory path
+ if 'task_target' not in config:
+ self.log_dir = os.path.join(
+ self.config['log_dir'],
+ self.config['model_name'] + '_' + self.timenow
+ )
+ else:
+ task_str = f"_{config['task_target']}" if config['task_target'] is not None else ""
+ self.log_dir = os.path.join(
+ self.config['log_dir'],
+ self.config['model_name'] + task_str + '_' + self.timenow
+ )
+ os.makedirs(self.log_dir, exist_ok=True)
+
+ def get_writer(self, phase, dataset_key, metric_key):
+ writer_key = f"{phase}-{dataset_key}-{metric_key}"
+ if writer_key not in self.writers:
+ # update directory path
+ writer_path = os.path.join(
+ self.log_dir,
+ phase,
+ dataset_key,
+ metric_key,
+ "metric_board"
+ )
+ os.makedirs(writer_path, exist_ok=True)
+ # update writers dictionary
+ self.writers[writer_key] = SummaryWriter(writer_path)
+ return self.writers[writer_key]
+
+
+ def speed_up(self):
+ self.model.to(device)
+ self.model.device = device
+ if self.config['ddp'] == True:
+ num_gpus = torch.cuda.device_count()
+ print(f'avai gpus: {num_gpus}')
+ # local_rank=[i for i in range(0,num_gpus)]
+ self.model = DDP(self.model, device_ids=[self.config['local_rank']],find_unused_parameters=True, output_device=self.config['local_rank'])
+ #self.optimizer = nn.DataParallel(self.optimizer, device_ids=[int(os.environ['LOCAL_RANK'])])
+
+ def setTrain(self):
+ self.model.train()
+ self.train = True
+
+ def setEval(self):
+ self.model.eval()
+ self.train = False
+
+ def load_ckpt(self, model_path):
+ if os.path.isfile(model_path):
+ saved = torch.load(model_path, map_location='cpu')
+ suffix = model_path.split('.')[-1]
+ if suffix == 'p':
+ self.model.load_state_dict(saved.state_dict())
+ else:
+ self.model.load_state_dict(saved)
+ self.logger.info('Model found in {}'.format(model_path))
+ else:
+ raise NotImplementedError(
+ "=> no model found at '{}'".format(model_path))
+
+ def save_ckpt(self, phase, dataset_key,ckpt_info=None):
+ save_dir = os.path.join(self.log_dir, phase, dataset_key)
+ os.makedirs(save_dir, exist_ok=True)
+ ckpt_name = f"ckpt_best.pth"
+ save_path = os.path.join(save_dir, ckpt_name)
+ if self.config['ddp'] == True:
+ torch.save(self.model.state_dict(), save_path)
+ else:
+ if 'svdd' in self.config['model_name']:
+ torch.save({'R': self.model.R,
+ 'c': self.model.c,
+ 'state_dict': self.model.state_dict(),}, save_path)
+ else:
+ torch.save(self.model.state_dict(), save_path)
+ self.logger.info(f"Checkpoint saved to {save_path}, current ckpt is {ckpt_info}")
+
+ def save_swa_ckpt(self):
+ save_dir = self.log_dir
+ os.makedirs(save_dir, exist_ok=True)
+ ckpt_name = f"swa.pth"
+ save_path = os.path.join(save_dir, ckpt_name)
+ torch.save(self.swa_model.state_dict(), save_path)
+ self.logger.info(f"SWA Checkpoint saved to {save_path}")
+
+
+ def save_feat(self, phase, fea, dataset_key):
+ save_dir = os.path.join(self.log_dir, phase, dataset_key)
+ os.makedirs(save_dir, exist_ok=True)
+ features = fea
+ feat_name = f"feat_best.npy"
+ save_path = os.path.join(save_dir, feat_name)
+ np.save(save_path, features)
+ self.logger.info(f"Feature saved to {save_path}")
+
+ def save_data_dict(self, phase, data_dict, dataset_key):
+ save_dir = os.path.join(self.log_dir, phase, dataset_key)
+ os.makedirs(save_dir, exist_ok=True)
+ file_path = os.path.join(save_dir, f'data_dict_{phase}.pickle')
+ with open(file_path, 'wb') as file:
+ pickle.dump(data_dict, file)
+ self.logger.info(f"data_dict saved to {file_path}")
+
+ def save_metrics(self, phase, metric_one_dataset, dataset_key):
+ save_dir = os.path.join(self.log_dir, phase, dataset_key)
+ os.makedirs(save_dir, exist_ok=True)
+ file_path = os.path.join(save_dir, 'metric_dict_best.pickle')
+ with open(file_path, 'wb') as file:
+ pickle.dump(metric_one_dataset, file)
+ self.logger.info(f"Metrics saved to {file_path}")
+
+ def train_step(self,data_dict):
+ if self.config['optimizer']['type']=='sam':
+ for i in range(2):
+ predictions = self.model(data_dict)
+ losses = self.model.get_losses(data_dict, predictions)
+ if i == 0:
+ pred_first = predictions
+ losses_first = losses
+ self.optimizer.zero_grad()
+ losses['overall'].backward()
+ if i == 0:
+ self.optimizer.first_step(zero_grad=True)
+ else:
+ self.optimizer.second_step(zero_grad=True)
+ return losses_first, pred_first
+ else:
+
+ predictions = self.model(data_dict)
+ if type(self.model) is DDP:
+ losses = self.model.module.get_losses(data_dict, predictions)
+ else:
+ losses = self.model.get_losses(data_dict, predictions)
+ self.optimizer.zero_grad()
+ losses['overall'].backward()
+ #self.model.module.set_mask_grad()
+ self.optimizer.step()
+
+
+ return losses,predictions
+
+
+ def train_epoch(
+ self,
+ epoch,
+ train_data_loader,
+ test_data_loaders=None,
+ ):
+
+ self.logger.info("===> Epoch[{}] start!".format(epoch))
+ if epoch>=1:
+ times_per_epoch = 2
+ else:
+ times_per_epoch = 2
+
+
+ #times_per_epoch=4
+
+ test_step = len(train_data_loader) // times_per_epoch # test 10 times per epoch
+ step_cnt = epoch * len(train_data_loader)
+
+ # save the training data_dict
+ data_dict = train_data_loader.dataset.data_dict
+ self.save_data_dict('train', data_dict, ','.join(self.config['train_dataset']))
+ # define training recorder
+ train_recorder_loss = defaultdict(Recorder)
+ train_recorder_metric = defaultdict(Recorder)
+
+ for iteration, data_dict in tqdm(enumerate(train_data_loader),total=len(train_data_loader)):
+ self.setTrain()
+ # more elegant and more scalable way of moving data to GPU
+ for key in data_dict.keys():
+ if data_dict[key]!=None and key!='name':
+ data_dict[key]=data_dict[key].cuda()
+
+ losses,predictions=self.train_step(data_dict)
+
+ # update learning rate
+
+ if 'SWA' in self.config and self.config['SWA'] and epoch>self.config['swa_start']:
+ self.swa_model.update_parameters(self.model)
+
+ # compute training metric for each batch data
+ if type(self.model) is DDP:
+ batch_metrics = self.model.module.get_train_metrics(data_dict, predictions)
+ else:
+ batch_metrics = self.model.get_train_metrics(data_dict, predictions)
+
+ # store data by recorder
+ ## store metric
+ for name, value in batch_metrics.items():
+ train_recorder_metric[name].update(value)
+ ## store loss
+ for name, value in losses.items():
+ train_recorder_loss[name].update(value)
+
+ # run tensorboard to visualize the training process
+ if iteration % 300 == 0 and self.config['local_rank']==0:
+ if self.config['SWA'] and (epoch>self.config['swa_start'] or self.config['dry_run']):
+ self.scheduler.step()
+ # info for loss
+ loss_str = f"Iter: {step_cnt} "
+ for k, v in train_recorder_loss.items():
+ v_avg = v.average()
+ if v_avg == None:
+ loss_str += f"training-loss, {k}: not calculated"
+ continue
+ loss_str += f"training-loss, {k}: {v_avg} "
+ # tensorboard-1. loss
+ writer = self.get_writer('train', ','.join(self.config['train_dataset']), k)
+ writer.add_scalar(f'train_loss/{k}', v_avg, global_step=step_cnt)
+ self.logger.info(loss_str)
+ # info for metric
+ metric_str = f"Iter: {step_cnt} "
+ for k, v in train_recorder_metric.items():
+ v_avg = v.average()
+ if v_avg == None:
+ metric_str += f"training-metric, {k}: not calculated "
+ continue
+ metric_str += f"training-metric, {k}: {v_avg} "
+ # tensorboard-2. metric
+ writer = self.get_writer('train', ','.join(self.config['train_dataset']), k)
+ writer.add_scalar(f'train_metric/{k}', v_avg, global_step=step_cnt)
+ self.logger.info(metric_str)
+
+
+
+ # clear recorder.
+ # Note we only consider the current 300 samples for computing batch-level loss/metric
+ for name, recorder in train_recorder_loss.items(): # clear loss recorder
+ recorder.clear()
+ for name, recorder in train_recorder_metric.items(): # clear metric recorder
+ recorder.clear()
+
+ # run test
+ #if True:
+ if (step_cnt+1) % test_step == 0:
+ if test_data_loaders is not None and (not self.config['ddp'] ):
+ self.logger.info("===> Test start!")
+ test_best_metric = self.test_epoch(
+ epoch,
+ iteration,
+ test_data_loaders,
+ step_cnt,
+ )
+ elif test_data_loaders is not None and (self.config['ddp'] and dist.get_rank() == 0):
+ self.logger.info("===> Test start!")
+ test_best_metric = self.test_epoch(
+ epoch,
+ iteration,
+ test_data_loaders,
+ step_cnt,
+ )
+ else:
+ test_best_metric = None
+
+ # total_end_time = time.time()
+ # total_elapsed_time = total_end_time - total_start_time
+ # print("总花费的时间: {:.2f} 秒".format(total_elapsed_time))
+ step_cnt += 1
+ return test_best_metric
+
+ def get_respect_acc(self,prob,label):
+ pred = np.where(prob > 0.5, 1, 0)
+ judge = (pred == label)
+ zero_num = len(label) - np.count_nonzero(label)
+ acc_fake = np.count_nonzero(judge[zero_num:]) / len(judge[zero_num:])
+ acc_real = np.count_nonzero(judge[:zero_num]) / len(judge[:zero_num])
+ return acc_real,acc_fake
+
+ def test_one_dataset(self, data_loader):
+ # define test recorder
+ test_recorder_loss = defaultdict(Recorder)
+ prediction_lists = []
+ feature_lists=[]
+ label_lists = []
+ for i, data_dict in tqdm(enumerate(data_loader),total=len(data_loader)):
+ # get data
+ if 'label_spe' in data_dict:
+ data_dict.pop('label_spe') # remove the specific label
+ data_dict['label'] = torch.where(data_dict['label']!=0, 1, 0) # fix the label to 0 and 1 only
+ # move data to GPU elegantly
+ for key in data_dict.keys():
+ if data_dict[key]!=None:
+ data_dict[key]=data_dict[key].cuda()
+ # model forward without considering gradient computation
+ predictions = self.inference(data_dict)
+ label_lists += list(data_dict['label'].cpu().detach().numpy())
+ prediction_lists += list(predictions['prob'].cpu().detach().numpy())
+ feature_lists += list(predictions['feat'].cpu().detach().numpy())
+ if type(self.model) is not AveragedModel:
+ # compute all losses for each batch data
+ with torch.no_grad():
+ if type(self.model) is DDP:
+ losses = self.model.module.get_losses(data_dict, predictions)
+ else:
+ losses = self.model.get_losses(data_dict, predictions)
+
+ # store data by recorder
+ for name, value in losses.items():
+ test_recorder_loss[name].update(value)
+
+ return test_recorder_loss, np.array(prediction_lists), np.array(label_lists),np.array(feature_lists)
+
+ def save_best(self,epoch,iteration,step,losses_one_dataset_recorder,key,metric_one_dataset):
+ best_metric = self.best_metrics_all_time[key].get(self.metric_scoring,
+ float('-inf') if self.metric_scoring != 'eer' else float(
+ 'inf'))
+ # Check if the current score is an improvement
+ improved = (metric_one_dataset[self.metric_scoring] > best_metric) if self.metric_scoring != 'eer' else (
+ metric_one_dataset[self.metric_scoring] < best_metric)
+ if improved:
+ # Update the best metric
+ self.best_metrics_all_time[key][self.metric_scoring] = metric_one_dataset[self.metric_scoring]
+ if key == 'avg':
+ self.best_metrics_all_time[key]['dataset_dict'] = metric_one_dataset['dataset_dict']
+ # Save checkpoint, feature, and metrics if specified in config
+ if self.config['save_ckpt'] and key not in FFpp_pool:
+ self.save_ckpt('test', key, f"{epoch}+{iteration}")
+ self.save_metrics('test', metric_one_dataset, key)
+ if losses_one_dataset_recorder is not None:
+ # info for each dataset
+ loss_str = f"dataset: {key} step: {step} "
+ for k, v in losses_one_dataset_recorder.items():
+ writer = self.get_writer('test', key, k)
+ v_avg = v.average()
+ if v_avg == None:
+ print(f'{k} is not calculated')
+ continue
+ # tensorboard-1. loss
+ writer.add_scalar(f'test_losses/{k}', v_avg, global_step=step)
+ loss_str += f"testing-loss, {k}: {v_avg} "
+ self.logger.info(loss_str)
+ # tqdm.write(loss_str)
+ metric_str = f"dataset: {key} step: {step} "
+ for k, v in metric_one_dataset.items():
+ if k == 'pred' or k == 'label' or k=='dataset_dict':
+ continue
+ metric_str += f"testing-metric, {k}: {v} "
+ # tensorboard-2. metric
+ writer = self.get_writer('test', key, k)
+ writer.add_scalar(f'test_metrics/{k}', v, global_step=step)
+ if 'pred' in metric_one_dataset:
+ acc_real, acc_fake = self.get_respect_acc(metric_one_dataset['pred'], metric_one_dataset['label'])
+ metric_str += f'testing-metric, acc_real:{acc_real}; acc_fake:{acc_fake}'
+ writer.add_scalar(f'test_metrics/acc_real', acc_real, global_step=step)
+ writer.add_scalar(f'test_metrics/acc_fake', acc_fake, global_step=step)
+ self.logger.info(metric_str)
+
+ def test_epoch(self, epoch, iteration, test_data_loaders, step):
+ # set model to eval mode
+ self.setEval()
+
+ # define test recorder
+ losses_all_datasets = {}
+ metrics_all_datasets = {}
+ best_metrics_per_dataset = defaultdict(dict) # best metric for each dataset, for each metric
+ avg_metric = {'acc': 0, 'auc': 0, 'eer': 0, 'ap': 0,'video_auc': 0,'dataset_dict':{}}
+ # testing for all test data
+ keys = test_data_loaders.keys()
+ for key in keys:
+ # save the testing data_dict
+ data_dict = test_data_loaders[key].dataset.data_dict
+ self.save_data_dict('test', data_dict, key)
+
+ # compute loss for each dataset
+ losses_one_dataset_recorder, predictions_nps, label_nps, feature_nps = self.test_one_dataset(test_data_loaders[key])
+ # print(f'stack len:{predictions_nps.shape};{label_nps.shape};{len(data_dict["image"])}')
+ losses_all_datasets[key] = losses_one_dataset_recorder
+ metric_one_dataset=get_test_metrics(y_pred=predictions_nps,y_true=label_nps,img_names=data_dict['image'])
+ for metric_name, value in metric_one_dataset.items():
+ if metric_name in avg_metric:
+ avg_metric[metric_name]+=value
+ avg_metric['dataset_dict'][key] = metric_one_dataset[self.metric_scoring]
+ if type(self.model) is AveragedModel:
+ metric_str = f"Iter Final for SWA: "
+ for k, v in metric_one_dataset.items():
+ metric_str += f"testing-metric, {k}: {v} "
+ self.logger.info(metric_str)
+ continue
+ self.save_best(epoch,iteration,step,losses_one_dataset_recorder,key,metric_one_dataset)
+
+ if len(keys)>0 and self.config.get('save_avg',False):
+ # calculate avg value
+ for key in avg_metric:
+ if key != 'dataset_dict':
+ avg_metric[key] /= len(keys)
+ self.save_best(epoch, iteration, step, None, 'avg', avg_metric)
+
+ self.logger.info('===> Test Done!')
+ return self.best_metrics_all_time # return all types of mean metrics for determining the best ckpt
+
+ @torch.no_grad()
+ def inference(self, data_dict):
+ predictions = self.model(data_dict, inference=True)
+ return predictions
diff --git a/clean/image/effort/DeepfakeBench/training/utils/__init__.py b/clean/image/effort/DeepfakeBench/training/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/clean/image/effort/DeepfakeBench/training/utils/metrics.py b/clean/image/effort/DeepfakeBench/training/utils/metrics.py
new file mode 100644
index 0000000000000000000000000000000000000000..b9c69c71b6be0e26822c313fb52b0893808be1ee
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/utils/metrics.py
@@ -0,0 +1,170 @@
+import numpy as np
+from sklearn import metrics
+import torch
+import torch.nn as nn
+
+def get_accracy(output, label):
+ _, prediction = torch.max(output, 1) # argmax
+ correct = (prediction == label).sum().item()
+ accuracy = correct / prediction.size(0)
+
+ return accuracy
+
+
+def get_prediction(output, label):
+ prob = nn.functional.softmax(output, dim=1)[:, 1]
+ prob = prob.view(prob.size(0), 1)
+ label = label.view(label.size(0), 1)
+ #print(prob.size(), label.size())
+ datas = torch.cat((prob, label.float()), dim=1)
+ return datas
+
+
+def calculate_metrics(label, output):
+ if output.size(1) == 2:
+ prob = torch.softmax(output, dim=1)[:, 1]
+ else:
+ prob = output
+
+ # Accuracy
+ _, prediction = torch.max(output, 1)
+ correct = (prediction == label).sum().item()
+ accuracy = correct / prediction.size(0)
+
+ # AUC and EER
+ fpr, tpr, thresholds = metrics.roc_curve(label.squeeze().cpu().numpy(),
+ prob.squeeze().cpu().numpy(),
+ pos_label=1)
+ if np.isnan(fpr[0]) or np.isnan(tpr[0]):
+ auc, eer = -1, -1
+ else:
+ auc = metrics.auc(fpr, tpr)
+ fnr = 1 - tpr
+ eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+
+ # Average Precision
+ y_true = label.cpu().detach().numpy()
+ y_pred = prob.cpu().detach().numpy()
+ ap = metrics.average_precision_score(y_true, y_pred)
+
+ return auc, eer, accuracy, ap
+
+
+# ------------ compute average metrics of batches---------------------
+class Metrics_batches_mean():
+ def __init__(self):
+ self.tprs = []
+ self.mean_fpr = np.linspace(0, 1, 100)
+ self.aucs = []
+ self.eers = []
+ self.aps = []
+
+ self.correct = 0
+ self.total = 0
+
+ def update(self, label, output):
+ acc = self._update_acc(label, output)
+ if output.size(1) == 2:
+ prob = torch.softmax(output, dim=1)[:, 1]
+ else:
+ prob = output
+ #label = 1-label
+ #prob = torch.softmax(output, dim=1)[:, 1]
+ auc, eer = self._update_auc(label, prob)
+ ap = self._update_ap(label, prob)
+
+ return acc, auc, eer, ap
+
+ def _update_auc(self, lab, prob):
+ fpr, tpr, thresholds = metrics.roc_curve(lab.squeeze().cpu().numpy(),
+ prob.squeeze().cpu().numpy(),
+ pos_label=1)
+ if np.isnan(fpr[0]) or np.isnan(tpr[0]):
+ return -1, -1
+
+ auc = metrics.auc(fpr, tpr)
+ interp_tpr = np.interp(self.mean_fpr, fpr, tpr)
+ interp_tpr[0] = 0.0
+ self.tprs.append(interp_tpr)
+ self.aucs.append(auc)
+
+ # return auc
+
+ # EER
+ fnr = 1 - tpr
+ eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+ self.eers.append(eer)
+
+ return auc, eer
+
+ def _update_acc(self, lab, output):
+ _, prediction = torch.max(output, 1) # argmax
+ correct = (prediction == lab).sum().item()
+ accuracy = correct / prediction.size(0)
+ # self.accs.append(accuracy)
+ self.correct = self.correct+correct
+ self.total = self.total+lab.size(0)
+ return accuracy
+
+ def _update_ap(self, label, prob):
+ y_true = label.cpu().detach().numpy()
+ y_pred = prob.cpu().detach().numpy()
+ ap = metrics.average_precision_score(y_true,y_pred)
+ self.aps.append(ap)
+
+ return np.mean(ap)
+
+
+ def get_mean_metrics(self):
+ mean_acc, std_acc = self.correct/self.total, 0
+ mean_auc, std_auc = self._mean_auc()
+ mean_err, std_err = np.mean(self.eers), np.std(self.eers)
+ mean_ap, std_ap = np.mean(self.aps), np.std(self.aps)
+
+ return mean_acc, std_acc, mean_auc, std_auc, mean_err, std_err, mean_ap, std_ap
+
+ def _mean_auc(self):
+ mean_tpr = np.mean(self.tprs, axis=0)
+ mean_tpr[-1] = 1.0
+ mean_auc = metrics.auc(self.mean_fpr, mean_tpr)
+ std_auc = np.std(self.aucs)
+ return mean_auc, std_auc
+
+ def clear(self):
+ self.tprs.clear()
+ self.aucs.clear()
+ # self.accs.clear()
+ self.correct=0
+ self.total=0
+ self.eers.clear()
+ self.aps.clear()
+
+
+class Metrics_all():
+ def __init__(self):
+ self.probs = []
+ self.labels = []
+
+ def store(self, label, output):
+ prob = torch.softmax(output, dim=1)[:, 1]
+ self.labels.append(label.squeeze().cpu().numpy())
+ self.probs.append(prob.squeeze().cpu().numpy())
+
+
+ def get_metrics(self):
+ y_pred = np.concatenate(self.probs)
+ y_true = np.concatenate(self.labels)
+ fpr, tpr, thresholds = metrics.roc_curve(y_true,y_pred,pos_label=1)
+ auc = metrics.auc(fpr, tpr)
+
+ # EER
+ fnr = 1 - tpr
+ eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+
+ ap = metrics.average_precision_score(y_true,y_pred)
+
+ return 0, 0, auc, 0, eer, 0, ap, 0
+
+ def clear(self):
+ self.probs.clear()
+ self.labels.clear()
\ No newline at end of file
diff --git a/clean/image/effort/DeepfakeBench/training/utils/registry.py b/clean/image/effort/DeepfakeBench/training/utils/registry.py
new file mode 100644
index 0000000000000000000000000000000000000000..86e256c18d0ad522de79149a154f676fd0bdb414
--- /dev/null
+++ b/clean/image/effort/DeepfakeBench/training/utils/registry.py
@@ -0,0 +1,20 @@
+class Registry(object):
+ def __init__(self):
+ self.data = {}
+
+ def register_module(self, module_name=None):
+ def _register(cls):
+ name = module_name
+ if module_name is None:
+ name = cls.__name__
+ self.data[name] = cls
+ return cls
+ return _register
+
+ def __getitem__(self, key):
+ return self.data[key]
+
+BACKBONE = Registry()
+DETECTOR = Registry()
+TRAINER = Registry()
+LOSSFUNC = Registry()
diff --git a/clean/image/effort/README.md b/clean/image/effort/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..1b6920c73acefc6c7992e8e59f92924ce2f481d3
--- /dev/null
+++ b/clean/image/effort/README.md
@@ -0,0 +1,182 @@
+# Effort: Orthogonal Subspace Decomposition for Generalizable AI-Generated Image Detection ([Paper](https://arxiv.org/abs/2411.15633); [Checkpoints](https://drive.google.com/drive/folders/19kQwGDjF18uk78EnnypxxOLaG4Aa4v1h?usp=sharing))
+
+[](https://creativecommons.org/licenses/by-nc/4.0/)   
+
+> 🎉🎉🎉 **Our paper has been accepted by ICML 2025 Oral 🏆!**
+
+Welcome to our work **Effort**, for detecting AI-generated images (AIGIs).
+
+In this work, we propose: (1) a **very very easy and effective method** for generalization AIGI detection😀; and (2) a **novel analysis tool** for quantifying the "degree of model's overfitting"😊.
+
+
+The figure below provides a brief introduction to our method: our method can be **plug-and-play inserted** into *any* vit-based large models such as CLIP.
+
+
+
+
+
+
+
+
+If you want to know a **very high-level code implementation of our method**, see below.
+
+```
+# 🟩 Perform SVD on the original weight
+U, S, Vh = torch.linalg.svd(module.weight.data, full_matrices=False)
+
+# 🟨 Keep top r singular components (main weight)
+U_r = U[:, :r] # 🔵 Shape: (out_features, r)
+S_r = S[:r] # 🔵 Shape: (r,)
+Vh_r = Vh[:r, :] # 🔵 Shape: (r, in_features)
+
+# 🟪 Reconstruct the main weight (fixed)
+weight_main = U_r @ torch.diag(S_r) @ Vh_r
+
+# 🟥 Residual components (trainable)
+U_residual = U[:, r:] # 🔵 Shape: (out_features, n - r)
+S_residual = S[r:] # 🔵 Shape: (n - r,)
+Vh_residual = Vh[r:, :] # 🔵 Shape: (n - r, in_features)
+```
+
+If you want to see **more method-specific implementation details**, please see the file [effort_implementation.py](https://github.com/YZY-stack/Effort-AIGI-Detection/blob/main/DeepfakeBench/training/detectors/effort_detector.py).
+
+---
+
+
+The following two tables display the **part results** of our method on **both the (face) deepfake detection benchmark and the (natural) AIGI detection benchmark**. Please refer to our paper for more results.
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+---
+
+
+## ⏳ Quick Start (if you just want to do the *inference*)
+[Back to top]
+
+
+### 1. Installation
+Please run the following script to install the required libraries:
+
+```
+sh install.sh
+```
+
+### 2. Download checkpoints
+If you are a deepfake player, more interested in face deepfake detection:
+- The checkpoint of "CLIP-L14 + our Effort" **training on FaceForensics++** are released at [Google Drive](https://drive.google.com/file/d/1m4fyJecABU-Yk3bJ4b1WhUwQa0xCkMLI/view?usp=drive_link).
+
+If you are interested in detecting general AI-generated images, we provide two checkpoints that are trained on GenImage and Chameleon datasets, respectively:
+- The checkpoint of "CLIP-L14 + our Effort" **training on GenImage (sdv1.4)** are released at [Google Drive](https://drive.google.com/file/d/1UXf1hC9FC1yV93uKwXSkdtepsgpIAU9d/view?usp=sharing).
+- The checkpoint of "CLIP-L14 + our Effort" **training on Chameleon (sdv1.4)** are released at [Google Drive](https://drive.google.com/file/d/1GlJ1y4xmTdqV0FfIcyBwNNU6cQird9DR/view?usp=sharing).
+
+
+### 3. Run demo
+You can then infer **one image *or* one folder with several images** using the pretrained weights.
+
+Specifically, run the following line:
+
+```
+cd DeepfakeBench/
+
+python3 training/demo.py --detector_config training/config/detector/effort.yaml --weights ./training/weights/{NAME_OF_THE_CKPT}.pth --image {IMAGE_PATH or IMAGE_FOLDER}
+```
+
+After running the above line, you can obtain the prediction results (fake probabilities) for each image.
+
+
+Note, you are processing a **face image**, please add the ``--landmark_model ./preprocessing/shape_predictor_81_face_landmarks.dat`` to **extract the facial region** for inference, as our model (trained on face deepfakes) used this face extractor for processing faces.
+
+
+---
+
+
+## 💻 Reproduction and Benchmarking Evaluation
+
+[Back to top]
+
+Since I am the creator and developer of [DeepfakeBench](https://github.com/SCLBD/DeepfakeBench), the **data, codebase, and benchmarking protocols are mainly used from it**. If you are a researcher in (face) deepfake detection, I highly recommend you try DeepfakeBench.
+
+If you want to **try other codebases, such as UnivFD**, we provide a folder `UniversalFakeDetect_Benchmark/` to reproduce and implement our method using its codebase. Using this codebase, you can then **reproduce the results of Table 2 of our manuscript**.
+
+Below, we provide the **detailed procedure to use DeepfakeBench to reproduce the results** of our paper, such as Table 1.
+
+
+### 1. Download datasets
+
+If you want to reproduce the results of each deepfake dataset, you can download the processed datasets (have already finished preprocessing such as frame extraction and face cropping) from [DeepfakeBench](https://github.com/SCLBD/DeepfakeBench). For evaluating more diverse fake methods (such as SimSwap, BlendFace, DeepFaceLab, etc), you are recommended to use the just-released [DF40 dataset](https://github.com/YZY-stack/DF40) (with 40 distinct forgery methods implemented).
+
+
+
+### 2. Preprocessing (**optional**)
+
+If you only want to use the processed data we provided, you can skip this step.
+
+Otherwise, you need to use the following codes for doing **data preprocessing strictly following DeepfakeBench**.
+
+
+### 3. Rearrangement (**optional**)
+
+> "Rearrangment" here means that we need to **create a *JSON file* for each dataset for collecting all frames within different folders**. Please refer to **DeepfakeBench** and **DF40** for the provided JSON files for each dataset.
+
+After running the above line, you will obtain the JSON files for each dataset in the `./preprocessing/dataset_json` folder. The rearranged structure organizes the data in a hierarchical manner, grouping videos based on their labels and data splits (*i.e.,* train, test, validation). Each video is represented as a dictionary entry containing relevant metadata, including file paths, labels, compression levels (if applicable), *etc*.
+
+
+
+### 4. Training
+
+First, you can run the following lines to train the model:
+- For multiple GPUs:
+```
+python3 -m torch.distributed.launch --nproc_per_node=4 training/train.py \
+--detector_path ./training/config/detector/effort.yaml \
+--train_dataset FaceForensics++ \
+--test_dataset Celeb-DF-v2 \
+--ddp
+```
+- For a single GPU:
+```
+python3 training/train.py \
+--detector_path ./training/config/detector/effort.yaml \
+--train_dataset FaceForensics++ \
+--test_dataset Celeb-DF-v2 \
+```
+
+### 5. Testing
+
+Once you finish training, you can test the model on several deepfake datasets such as DF40.
+
+```
+python3 training/test.py \
+--detector_path ./training/config/detector/effort.yaml \
+--test_dataset simswap_ff blendface_ff uniface_ff fomm_ff deepfacelab \
+--weights_path ./training/weights/{CKPT}.pth
+```
+Then, you can obtain similar evaluation results reported in our manuscript.
+
+
+---
+
+## 📕 Citation
+If you find our work helpful to your research, please consider citing our paper as follows:
+```
+@article{yan2024effort,
+ title={Effort: Efficient Orthogonal Modeling for Generalizable AI-Generated Image Detection},
+ author={Yan, Zhiyuan and Wang, Jiangming and Wang, Zhendong and Jin, Peng and Zhang, Ke-Yue and Chen, Shen and Yao, Taiping and Ding, Shouhong and Wu, Baoyuan and Yuan, Li},
+ journal={arXiv preprint arXiv:2411.15633},
+ year={2024}
+}
+```
+
+
diff --git a/clean/image/effort/SOURCE.md b/clean/image/effort/SOURCE.md
new file mode 100644
index 0000000000000000000000000000000000000000..63cebec5b0661dad8fa530ff89f58439a9cd0ad2
--- /dev/null
+++ b/clean/image/effort/SOURCE.md
@@ -0,0 +1,17 @@
+# Source: image/effort
+
+| Field | Value |
+|---|---|
+| Upstream | https://github.com/YZY-stack/Effort-AIGI-Detection |
+| Paper | https://arxiv.org/abs/2411.15633 |
+| Commit SHA | **not recorded** -- the vendoring step did not preserve it |
+| Mirrored on | 2026-09-15 |
+| Upstream license | no license file present upstream (all rights reserved) |
+
+This is a **mirror**, stripped to the files needed for inference. The full
+untouched snapshot is at `archive/image__effort.tar.gz`.
+
+This code is the work of its original authors and is **not** covered by the
+DeepSafe project license. If you are an author and want this removed, open an
+issue on https://github.com/deepsafehq/deepsafe-bench and it will be taken down
+within 48 hours, no questions asked.
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/dataset_paths.py b/clean/image/effort/UniversalFakeDetect_Benchmark/dataset_paths.py
new file mode 100644
index 0000000000000000000000000000000000000000..6e5f7ce81c71a410f4b7254760fa713e871b2d38
--- /dev/null
+++ b/clean/image/effort/UniversalFakeDetect_Benchmark/dataset_paths.py
@@ -0,0 +1,136 @@
+DATASET_PATHS = [
+
+ dict(
+ real_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/progan',
+ fake_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/progan',
+ data_mode='wang2020',
+ key='progan'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/cyclegan',
+ fake_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/cyclegan',
+ data_mode='wang2020',
+ key='cyclegan'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/biggan/', # Imagenet
+ fake_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/biggan/',
+ data_mode='wang2020',
+ key='biggan'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/stylegan',
+ fake_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/stylegan',
+ data_mode='wang2020',
+ key='stylegan'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/gaugan', # It is COCO
+ fake_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/gaugan',
+ data_mode='wang2020',
+ key='gaugan'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/stargan',
+ fake_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/stargan',
+ data_mode='wang2020',
+ key='stargan'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/deepfake',
+ fake_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/deepfake',
+ data_mode='wang2020',
+ key='deepfake'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/seeingdark',
+ fake_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/seeingdark',
+ data_mode='wang2020',
+ key='sitd'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/san',
+ fake_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/san',
+ data_mode='wang2020',
+ key='san'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/crn', # Images from some video games
+ fake_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/crn',
+ data_mode='wang2020',
+ key='crn'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/imle', # Images from some video games
+ fake_path='/Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/test/imle',
+ data_mode='wang2020',
+ key='imle'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security/jeremiewang/dataset/open/diffusion_datasets/imagenet',
+ fake_path='/Youtu_Pangu_Security/jeremiewang/dataset/open/diffusion_datasets/guided',
+ data_mode='wang2020',
+ key='guided'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security/jeremiewang/dataset/open/diffusion_datasets/laion',
+ fake_path='/Youtu_Pangu_Security/jeremiewang/dataset/open/diffusion_datasets/ldm_200',
+ data_mode='wang2020',
+ key='ldm_200'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security/jeremiewang/dataset/open/diffusion_datasets/laion',
+ fake_path='/Youtu_Pangu_Security/jeremiewang/dataset/open/diffusion_datasets/ldm_200_cfg',
+ data_mode='wang2020',
+ key='ldm_200_cfg'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security/jeremiewang/dataset/open/diffusion_datasets/laion',
+ fake_path='/Youtu_Pangu_Security/jeremiewang/dataset/open/diffusion_datasets/ldm_100',
+ data_mode='wang2020',
+ key='ldm_100'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security/jeremiewang/dataset/open/diffusion_datasets/laion',
+ fake_path='/Youtu_Pangu_Security/jeremiewang/dataset/open/diffusion_datasets/glide_100_27',
+ data_mode='wang2020',
+ key='glide_100_27'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security/jeremiewang/dataset/open/diffusion_datasets/laion',
+ fake_path='/Youtu_Pangu_Security/jeremiewang/dataset/open/diffusion_datasets/glide_50_27',
+ data_mode='wang2020',
+ key='glide_50_27'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security/jeremiewang/dataset/open/diffusion_datasets/laion',
+ fake_path='/Youtu_Pangu_Security/jeremiewang/dataset/open/diffusion_datasets/glide_100_10',
+ data_mode='wang2020',
+ key='glide_100_10'
+ ),
+
+ dict(
+ real_path='/Youtu_Pangu_Security/jeremiewang/dataset/open/diffusion_datasets/laion',
+ fake_path='/Youtu_Pangu_Security/jeremiewang/dataset/open/diffusion_datasets/dalle',
+ data_mode='wang2020',
+ key='dalle'
+ ),
+
+]
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/earlystop.py b/clean/image/effort/UniversalFakeDetect_Benchmark/earlystop.py
new file mode 100644
index 0000000000000000000000000000000000000000..01296e3419acddaf0e92a69477fddb3be9440cf5
--- /dev/null
+++ b/clean/image/effort/UniversalFakeDetect_Benchmark/earlystop.py
@@ -0,0 +1,44 @@
+import numpy as np
+import torch
+
+
+class EarlyStopping:
+ """Early stops the training if validation loss doesn't improve after a given patience."""
+ def __init__(self, patience=1, verbose=False, delta=0):
+ """
+ Args:
+ patience (int): How long to wait after last time validation loss improved.
+ Default: 7
+ verbose (bool): If True, prints a message for each validation loss improvement.
+ Default: False
+ delta (float): Minimum change in the monitored quantity to qualify as an improvement.
+ Default: 0
+ """
+ self.patience = patience
+ self.verbose = verbose
+ self.counter = 0
+ self.best_score = None
+ self.early_stop = False
+ self.score_max = -np.Inf
+ self.delta = delta
+
+ def __call__(self, score, model):
+ if self.best_score is None:
+ self.best_score = score
+ self.save_checkpoint(score, model)
+ elif score < self.best_score - self.delta:
+ self.counter += 1
+ print(f'EarlyStopping counter: {self.counter} out of {self.patience}')
+ if self.counter >= self.patience:
+ self.early_stop = True
+ else:
+ self.best_score = score
+ self.save_checkpoint(score, model)
+ self.counter = 0
+
+ def save_checkpoint(self, score, model):
+ '''Saves model when validation loss decrease.'''
+ if self.verbose:
+ print(f'Validation accuracy increased ({self.score_max:.6f} --> {score:.6f}). Saving model ...')
+ model.save_networks('best')
+ self.score_max = score
\ No newline at end of file
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/models/__init__.py b/clean/image/effort/UniversalFakeDetect_Benchmark/models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..4417c51d2232e5bf7febbdf1b07582f7a4d1524e
--- /dev/null
+++ b/clean/image/effort/UniversalFakeDetect_Benchmark/models/__init__.py
@@ -0,0 +1,18 @@
+from .clip_models import ClipModel
+
+
+VALID_NAMES = {
+ 'CLIP:ViT-B/16_svd':'/Youtu_Pangu_Security_Public/youtu-pangu-public/jeremiewang/pretrained_model/huggingface/openai/clip-vit-base-patch16/',
+ 'CLIP:ViT-B/32_svd':'/Youtu_Pangu_Security_Public/youtu-pangu-public/jeremiewang/pretrained_model/huggingface/openai/clip-vit-base-patch32/',
+ 'CLIP:ViT-L/14_svd':'/Youtu_Pangu_Security_Public/youtu-pangu-public/zhiyuanyan/huggingface/hub/models--openai--clip-vit-large-patch14/snapshots/32bd64288804d66eefd0ccbe215aa642df71cc41/',
+ 'SigLIP:ViT-L/16_256_svd':'/Youtu_Pangu_Security_Public/youtu-pangu-public/jeremiewang/pretrained_model/huggingface/google/siglip-large-patch16-256/',
+ 'BEiTv2:ViT-L/16_svd':'/Youtu_Pangu_Security_Public/youtu-pangu-public/jeremiewang/pretrained_model/BEiT-v2/beitv2_large_patch16_224_pt1k_ft21k.pth',
+}
+
+
+def get_model(name, opt):
+ assert name in VALID_NAMES.keys()
+ if name.startswith("CLIP:"):
+ return ClipModel(VALID_NAMES[name], opt)
+ else:
+ assert False
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/models/base_model.py b/clean/image/effort/UniversalFakeDetect_Benchmark/models/base_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..684bdd31004eb9d5664da1aba08dc3ba3b7c4d80
--- /dev/null
+++ b/clean/image/effort/UniversalFakeDetect_Benchmark/models/base_model.py
@@ -0,0 +1,58 @@
+import os
+import torch
+import torch.nn as nn
+from torch.nn import init
+from torch.optim import lr_scheduler
+
+
+class BaseModel(nn.Module):
+ def __init__(self, opt):
+ super(BaseModel, self).__init__()
+ self.opt = opt
+ self.total_steps = 0
+ self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
+ self.device = torch.device('cuda:{}'.format(opt.gpu_ids[0])) if opt.gpu_ids else torch.device('cpu')
+
+ def save_networks(self, save_filename):
+ save_path = os.path.join(self.save_dir, save_filename)
+
+ # serialize model and optimizer to dict
+ state_dict = {
+ 'model': self.model.state_dict(),
+ 'optimizer' : self.optimizer.state_dict(),
+ 'total_steps' : self.total_steps,
+ }
+
+ torch.save(state_dict, save_path)
+
+
+ def eval(self):
+ self.model.eval()
+
+ def test(self):
+ with torch.no_grad():
+ self.forward()
+
+
+def init_weights(net, init_type='normal', gain=0.02):
+ def init_func(m):
+ classname = m.__class__.__name__
+ if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
+ if init_type == 'normal':
+ init.normal_(m.weight.data, 0.0, gain)
+ elif init_type == 'xavier':
+ init.xavier_normal_(m.weight.data, gain=gain)
+ elif init_type == 'kaiming':
+ init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+ elif init_type == 'orthogonal':
+ init.orthogonal_(m.weight.data, gain=gain)
+ else:
+ raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
+ if hasattr(m, 'bias') and m.bias is not None:
+ init.constant_(m.bias.data, 0.0)
+ elif classname.find('BatchNorm2d') != -1:
+ init.normal_(m.weight.data, 1.0, gain)
+ init.constant_(m.bias.data, 0.0)
+
+ print('initialize network with %s' % init_type)
+ net.apply(init_func)
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/models/clip_models.py b/clean/image/effort/UniversalFakeDetect_Benchmark/models/clip_models.py
new file mode 100644
index 0000000000000000000000000000000000000000..1b787f0ebb2ec915e2960ce462e35d0356d91ccf
--- /dev/null
+++ b/clean/image/effort/UniversalFakeDetect_Benchmark/models/clip_models.py
@@ -0,0 +1,167 @@
+import math
+from PIL import Image
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoProcessor, CLIPModel, ViTModel, ViTConfig
+import loralib as lora
+
+
+class ClipModel(nn.Module):
+ def __init__(self, name, opt, num_classes=1):
+ super(ClipModel, self).__init__()
+ self.use_svd = opt.use_svd
+
+ if self.use_svd:
+ self.model = CLIPModel.from_pretrained(name)
+ self.model.vision_model = apply_svd_residual_to_self_attn(self.model.vision_model, r=1024-1)
+
+ for name, param in self.model.vision_model.named_parameters():
+ print('{}: {}'.format(name, param.requires_grad))
+ num_param = sum(p.numel() for p in self.model.vision_model.parameters() if p.requires_grad)
+ num_total_param = sum(p.numel() for p in self.model.vision_model.parameters())
+ print('Number of total parameters: {}, tunable parameters: {}'.format(num_total_param, num_param))
+
+ self.fc = nn.Linear( 1024, num_classes )
+ else:
+ self.model = CLIPModel.from_pretrained(name)
+
+ for name, param in self.model.vision_model.named_parameters():
+ print('{}: {}'.format(name, param.requires_grad))
+ num_param = sum(p.numel() for p in self.model.vision_model.parameters() if p.requires_grad)
+ num_total_param = sum(p.numel() for p in self.model.vision_model.parameters())
+ print('Number of total parameters: {}, tunable parameters: {}'.format(num_total_param, num_param))
+
+ self.fc = nn.Linear( 1024, num_classes )
+
+ def forward(self, x, return_feature=False):
+ features = self.model.vision_model(x)['pooler_output']
+
+ if return_feature:
+ return features
+ return self.fc(features)
+
+
+# Custom module to represent the residual using SVD components
+class SVDResidualLinear(nn.Module):
+ def __init__(self, in_features, out_features, r, bias=True, init_weight=None):
+ super(SVDResidualLinear, self).__init__()
+ self.in_features = in_features
+ self.out_features = out_features
+ self.r = r # Number of top singular values to exclude
+
+ # Main weight (fixed)
+ self.weight_main = nn.Parameter(torch.Tensor(out_features, in_features), requires_grad=False)
+ if init_weight is not None:
+ self.weight_main.data.copy_(init_weight)
+ else:
+ nn.init.kaiming_uniform_(self.weight_main, a=math.sqrt(5))
+
+ # Bias
+ if bias:
+ self.bias = nn.Parameter(torch.Tensor(out_features))
+ nn.init.zeros_(self.bias)
+ else:
+ self.register_parameter('bias', None)
+
+ # SVD components
+ self.S_r = None
+ self.U_r = None
+ self.V_r = None
+ self.S_residual = None
+ self.U_residual = None
+ self.V_residual = None
+
+ def forward(self, x):
+ if self.S_residual is not None:
+ # Reconstruct the residual weight
+ residual_weight = self.U_residual @ torch.diag(self.S_residual) @ self.V_residual
+ # Total weight is the fixed main weight plus the residual
+ weight = self.weight_main + residual_weight
+ else:
+ # If residual components are not set, use only the main weight
+ weight = self.weight_main
+
+ return F.linear(x, weight, self.bias)
+
+
+# Function to replace nn.Linear modules within self_attn modules with SVDResidualLinear
+def apply_svd_residual_to_self_attn(model, r):
+ for name, module in model.named_children():
+ # if ('self_attn' in name) or ('mlp' in name):
+ if ('self_attn' in name):
+ # Replace nn.Linear layers in this module
+ for sub_name, sub_module in module.named_modules():
+ if isinstance(sub_module, nn.Linear):
+ # Get parent module within self_attn
+ parent_module = module
+ sub_module_names = sub_name.split('.')
+ for module_name in sub_module_names[:-1]:
+ parent_module = getattr(parent_module, module_name)
+ # Replace the nn.Linear layer with SVDResidualLinear
+ setattr(parent_module, sub_module_names[-1], replace_with_svd_residual(sub_module, r))
+ else:
+ # Recursively apply to child modules
+ apply_svd_residual_to_self_attn(module, r)
+ # After replacing, set requires_grad for residual components
+ for param_name, param in model.named_parameters():
+ if any(x in param_name for x in ['S_residual', 'U_residual', 'V_residual']):
+ param.requires_grad = True
+ else:
+ param.requires_grad = False
+ return model
+
+
+# Function to replace a module with SVDResidualLinear
+def replace_with_svd_residual(module, r):
+ if isinstance(module, nn.Linear):
+ in_features = module.in_features
+ out_features = module.out_features
+ bias = module.bias is not None
+
+ # Create SVDResidualLinear module
+ new_module = SVDResidualLinear(in_features, out_features, r, bias=bias, init_weight=module.weight.data.clone())
+
+ if bias and module.bias is not None:
+ new_module.bias.data.copy_(module.bias.data)
+
+ # Calculate the frobenius norm of original weight
+ new_module.weight_original_fnorm = torch.norm(module.weight.data, p='fro')
+
+ # Perform SVD on the original weight
+ U, S, Vh = torch.linalg.svd(module.weight.data, full_matrices=False)
+
+ # Determine r based on the rank of the weight matrix
+ r = min(r, len(S)) # Ensure r does not exceed the number of singular values
+
+ # Keep top r singular components (main weight)
+ U_r = U[:, :r] # Shape: (out_features, r)
+ S_r = S[:r] # Shape: (r,)
+ Vh_r = Vh[:r, :] # Shape: (r, in_features)
+
+ # Reconstruct the main weight (fixed)
+ weight_main = U_r @ torch.diag(S_r) @ Vh_r
+
+ # Calculate the frobenius norm of main weight
+ new_module.weight_main_fnorm = torch.norm(weight_main.data, p='fro')
+
+ # Set the main weight
+ new_module.weight_main.data.copy_(weight_main)
+
+ # Residual components (trainable)
+ U_residual = U[:, r:] # Shape: (out_features, n - r)
+ S_residual = S[r:] # Shape: (n - r,)
+ Vh_residual = Vh[r:, :] # Shape: (n - r, in_features)
+
+ if len(S_residual) > 0:
+ new_module.S_residual = nn.Parameter(S_residual.clone())
+ new_module.U_residual = nn.Parameter(U_residual.clone())
+ new_module.V_residual = nn.Parameter(Vh_residual.clone())
+ else:
+ new_module.S_residual = None
+ new_module.U_residual = None
+ new_module.V_residual = None
+
+ return new_module
+ else:
+ return module
\ No newline at end of file
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/models/trainer.py b/clean/image/effort/UniversalFakeDetect_Benchmark/models/trainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..fd9d5e7cd0c6a1aa1adcafb0c3550ec4ce002404
--- /dev/null
+++ b/clean/image/effort/UniversalFakeDetect_Benchmark/models/trainer.py
@@ -0,0 +1,71 @@
+import functools
+import torch
+import torch.nn as nn
+from .base_model import BaseModel, init_weights
+import sys
+from models import get_model
+
+
+class Trainer(BaseModel):
+ def name(self):
+ return 'Trainer'
+
+ def __init__(self, opt):
+ super(Trainer, self).__init__(opt)
+ self.opt = opt
+ self.model = get_model(opt.arch, opt)
+ self.lr = opt.lr
+ torch.nn.init.normal_(self.model.fc.weight.data, 0.0, opt.init_gain)
+
+ if opt.fix_backbone:
+ params = []
+ for name, p in self.model.named_parameters():
+ if name=="fc.weight" or name=="fc.bias":
+ params.append(p)
+ else:
+ p.requires_grad = False
+ else:
+ print("Your backbone is not fixed. Are you sure you want to proceed? If this is a mistake, enable the --fix_backbone command during training and rerun")
+ # import time
+ # time.sleep(3)
+ params = self.model.parameters()
+
+ if opt.optim == 'adam':
+ self.optimizer = torch.optim.AdamW(params, lr=opt.lr, betas=(opt.beta1, 0.999), weight_decay=opt.weight_decay)
+ elif opt.optim == 'sgd':
+ self.optimizer = torch.optim.SGD(params, lr=opt.lr, momentum=0.0, weight_decay=opt.weight_decay)
+ else:
+ raise ValueError("optim should be [adam, sgd]")
+
+ self.loss_fn = nn.BCEWithLogitsLoss()
+
+ self.model.to(opt.gpu_ids[0])
+
+ def adjust_learning_rate(self, min_lr=1e-6):
+ for param_group in self.optimizer.param_groups:
+ param_group['lr'] *= 0.8
+ self.lr = param_group['lr']
+ if param_group['lr'] < min_lr:
+ return False
+ return True
+
+ def set_input(self, input):
+ self.input = input[0].to(self.device)
+ self.label = input[1].to(self.device).float()
+
+ def forward(self):
+ self.output = self.model(self.input)
+ self.output = self.output.view(-1).unsqueeze(1)
+
+ def get_loss(self):
+ return self.loss_fn(self.output.squeeze(1), self.label)
+
+ def optimize_parameters(self):
+ self.forward()
+ self.loss = self.loss_fn(self.output.squeeze(1), self.label)
+ self.optimizer.zero_grad()
+ self.loss.backward()
+ self.optimizer.step()
+
+
+
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/options/__init__.py b/clean/image/effort/UniversalFakeDetect_Benchmark/options/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/options/base_options.py b/clean/image/effort/UniversalFakeDetect_Benchmark/options/base_options.py
new file mode 100644
index 0000000000000000000000000000000000000000..e669d7256f54d7f75e2a11c9e892ca3547f2be29
--- /dev/null
+++ b/clean/image/effort/UniversalFakeDetect_Benchmark/options/base_options.py
@@ -0,0 +1,118 @@
+import argparse
+import os
+import util
+import torch
+
+
+class BaseOptions():
+ def __init__(self):
+ self.initialized = False
+
+ def initialize(self, parser):
+ parser.add_argument('--mode', default='binary')
+ parser.add_argument('--arch', type=str, default='res50', help='see my_models/__init__.py')
+ parser.add_argument('--fix_backbone', action='store_true')
+ parser.add_argument('--use_svd', action='store_true')
+
+ # data augmentation
+ parser.add_argument('--rz_interp', default='bilinear')
+ parser.add_argument('--blur_prob', type=float, default=0.5)
+ parser.add_argument('--blur_sig', default='0.0,3.0')
+ parser.add_argument('--jpg_prob', type=float, default=0.5)
+ parser.add_argument('--jpg_method', default='cv2,pil')
+ parser.add_argument('--jpg_qual', default='30,100')
+
+
+ parser.add_argument('--real_list_path', default=None, help='only used if data_mode==ours: path for the list of real images, which should contain train.pickle and val.pickle')
+ parser.add_argument('--fake_list_path', default=None, help='only used if data_mode==ours: path for the list of fake images, which should contain train.pickle and val.pickle')
+ parser.add_argument('--wang2020_data_path', default=None, help='only used if data_mode==wang2020 it should contain train and test folders')
+ parser.add_argument('--data_mode', default='ours', help='wang2020 or ours')
+ parser.add_argument('--data_label', default='train', help='label to decide whether train or validation dataset')
+ parser.add_argument('--weight_decay', type=float, default=0.0, help='loss weight for l2 reg')
+
+ parser.add_argument('--class_bal', action='store_true') # what is this ?
+ parser.add_argument('--batch_size', type=int, default=256, help='input batch size')
+ parser.add_argument('--loadSize', type=int, default=256, help='scale images to this size')
+ parser.add_argument('--cropSize', type=int, default=224, help='then crop to this size')
+ parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU')
+ parser.add_argument('--name', type=str, default='experiment_name', help='name of the experiment. It decides where to store samples and models')
+ parser.add_argument('--num_threads', default=4, type=int, help='# threads for loading data')
+ parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
+ parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')
+ parser.add_argument('--resize_or_crop', type=str, default='scale_and_crop', help='scaling and cropping of images at load time [resize_and_crop|crop|scale_width|scale_width_and_crop|none]')
+ parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data augmentation')
+ parser.add_argument('--init_type', type=str, default='normal', help='network initialization [normal|xavier|kaiming|orthogonal]')
+ parser.add_argument('--init_gain', type=float, default=0.02, help='scaling factor for normal, xavier and orthogonal.')
+ parser.add_argument('--suffix', default='', type=str, help='customized suffix: opt.name = opt.name + suffix: e.g., {model}_{netG}_size{loadSize}')
+ self.initialized = True
+ return parser
+
+ def gather_options(self):
+ # initialize parser with basic options
+ if not self.initialized:
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+ parser = self.initialize(parser)
+
+ # get the basic options
+ opt, _ = parser.parse_known_args()
+ self.parser = parser
+
+ return parser.parse_args()
+
+ def print_options(self, opt):
+ message = ''
+ message += '----------------- Options ---------------\n'
+ for k, v in sorted(vars(opt).items()):
+ comment = ''
+ default = self.parser.get_default(k)
+ if v != default:
+ comment = '\t[default: %s]' % str(default)
+ message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
+ message += '----------------- End -------------------'
+ print(message)
+
+ # save to the disk
+ expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
+ util.mkdirs(expr_dir)
+ file_name = os.path.join(expr_dir, 'opt.txt')
+ with open(file_name, 'wt') as opt_file:
+ opt_file.write(message)
+ opt_file.write('\n')
+
+ def parse(self, print_options=True):
+
+ opt = self.gather_options()
+ opt.isTrain = self.isTrain # train or test
+
+ # process opt.suffix
+ if opt.suffix:
+ suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else ''
+ opt.name = opt.name + suffix
+
+ if print_options:
+ self.print_options(opt)
+
+ # set gpu ids
+ str_ids = opt.gpu_ids.split(',')
+ opt.gpu_ids = []
+ for str_id in str_ids:
+ id = int(str_id)
+ if id >= 0:
+ opt.gpu_ids.append(id)
+ if len(opt.gpu_ids) > 0:
+ torch.cuda.set_device(opt.gpu_ids[0])
+
+ # additional
+ #opt.classes = opt.classes.split(',')
+ opt.rz_interp = opt.rz_interp.split(',')
+ opt.blur_sig = [float(s) for s in opt.blur_sig.split(',')]
+ opt.jpg_method = opt.jpg_method.split(',')
+ opt.jpg_qual = [int(s) for s in opt.jpg_qual.split(',')]
+ if len(opt.jpg_qual) == 2:
+ opt.jpg_qual = list(range(opt.jpg_qual[0], opt.jpg_qual[1] + 1))
+ elif len(opt.jpg_qual) > 2:
+ raise ValueError("Shouldn't have more than 2 values for --jpg_qual.")
+
+ self.opt = opt
+ return self.opt
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/options/test_options.py b/clean/image/effort/UniversalFakeDetect_Benchmark/options/test_options.py
new file mode 100644
index 0000000000000000000000000000000000000000..f824c7aae81d325bf81e16c5a5e2c1931ccd5b34
--- /dev/null
+++ b/clean/image/effort/UniversalFakeDetect_Benchmark/options/test_options.py
@@ -0,0 +1,13 @@
+from .base_options import BaseOptions
+
+
+class TestOptions(BaseOptions):
+ def initialize(self, parser):
+ parser = BaseOptions.initialize(self, parser)
+ parser.add_argument('--model_path')
+ parser.add_argument('--no_resize', action='store_true')
+ parser.add_argument('--no_crop', action='store_true')
+ parser.add_argument('--eval', action='store_true', help='use eval mode during test time.')
+
+ self.isTrain = False
+ return parser
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/options/train_options.py b/clean/image/effort/UniversalFakeDetect_Benchmark/options/train_options.py
new file mode 100644
index 0000000000000000000000000000000000000000..032cd8129f09fb6ed487c74d8d9eb72eb7798cc8
--- /dev/null
+++ b/clean/image/effort/UniversalFakeDetect_Benchmark/options/train_options.py
@@ -0,0 +1,22 @@
+from .base_options import BaseOptions
+
+
+class TrainOptions(BaseOptions):
+ def initialize(self, parser):
+ parser = BaseOptions.initialize(self, parser)
+ parser.add_argument('--earlystop_epoch', type=int, default=5)
+ parser.add_argument('--data_aug', action='store_true', help='if specified, perform additional data augmentation (photometric, blurring, jpegging)')
+ parser.add_argument('--optim', type=str, default='adam', help='optim to use [sgd, adam]')
+ parser.add_argument('--new_optim', action='store_true', help='new optimizer instead of loading the optim state')
+ parser.add_argument('--loss_freq', type=int, default=500, help='frequency of showing loss on tensorboard')
+ parser.add_argument('--save_epoch_freq', type=int, default=1, help='frequency of saving checkpoints at the end of epochs')
+ parser.add_argument('--epoch_count', type=int, default=1, help='the starting epoch count, we save the model by , +, ...')
+ parser.add_argument('--last_epoch', type=int, default=-1, help='starting epoch count for scheduler intialization')
+ parser.add_argument('--train_split', type=str, default='train', help='train, val, test, etc')
+ parser.add_argument('--val_split', type=str, default='val', help='train, val, test, etc')
+ parser.add_argument('--niter', type=int, default=100, help='total epoches')
+ parser.add_argument('--beta1', type=float, default=0.9, help='momentum term of adam')
+ parser.add_argument('--lr', type=float, default=0.0001, help='initial learning rate for adam')
+
+ self.isTrain = True
+ return parser
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/results/clip_vitl14_svd/iters_18000/acc0.txt b/clean/image/effort/UniversalFakeDetect_Benchmark/results/clip_vitl14_svd/iters_18000/acc0.txt
new file mode 100644
index 0000000000000000000000000000000000000000..20d1b1750ede60d5629fe80e2f9113fbe7a4615c
--- /dev/null
+++ b/clean/image/effort/UniversalFakeDetect_Benchmark/results/clip_vitl14_svd/iters_18000/acc0.txt
@@ -0,0 +1,41 @@
+-----------------------------------------
+thres: 0.5
+progan: 100.0 100.0 100.0
+thres: 0.5
+cyclegan: 99.7 100.0 99.85
+thres: 0.5
+biggan: 99.5 99.7 99.6
+thres: 0.5
+stylegan: 100.0 90.1 95.05
+thres: 0.5
+gaugan: 99.2 100.0 99.6
+thres: 0.5
+stargan: 100.0 100.0 100.0
+thres: 0.5
+deepfake: 100.0 75.2 87.6
+thres: 0.5
+sitd: 91.0 94.0 92.5
+thres: 0.5
+san: 100.0 62.0 81.0
+thres: 0.5
+crn: 97.8 100.0 98.9
+thres: 0.5
+imle: 97.8 100.0 98.9
+thres: 0.5
+guided: 99.7 38.6 69.15
+thres: 0.5
+ldm_200: 100.0 98.6 99.3
+thres: 0.5
+ldm_200_cfg: 100.0 93.6 96.8
+thres: 0.5
+ldm_100: 100.0 98.9 99.45
+thres: 0.5
+glide_100_27: 100.0 94.9 97.45
+thres: 0.5
+glide_50_27: 100.0 95.6 97.8
+thres: 0.5
+glide_100_10: 100.0 94.3 97.15
+thres: 0.5
+dalle: 100.0 96.1 98.05
+avg: 99.19 91.14 95.17
+-----------------------------------------
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/results/clip_vitl14_svd/iters_18000/acc1.txt b/clean/image/effort/UniversalFakeDetect_Benchmark/results/clip_vitl14_svd/iters_18000/acc1.txt
new file mode 100644
index 0000000000000000000000000000000000000000..b7fe24ac7b9f259fc9bfc41c866db2b88c34319e
--- /dev/null
+++ b/clean/image/effort/UniversalFakeDetect_Benchmark/results/clip_vitl14_svd/iters_18000/acc1.txt
@@ -0,0 +1,41 @@
+-----------------------------------------
+thres: 0.5002928899193648
+progan: 100.0 100.0 100.0
+thres: 0.999847412109375
+cyclegan: 99.9 99.9 99.9
+thres: 0.9992189407348633
+biggan: 100.0 99.4 99.7
+thres: 4.657595854951069e-05
+stylegan: 100.0 97.1 98.55
+thres: 0.998432993888855
+gaugan: 100.0 100.0 100.0
+thres: 0.5262151546776295
+stargan: 100.0 100.0 100.0
+thres: 0.002418691059574485
+deepfake: 98.2 93.7 95.95
+thres: 0.892815351486206
+sitd: 97.0 92.0 94.5
+thres: 0.0004463841032702476
+san: 96.0 82.0 89.0
+thres: 0.9999018907546997
+crn: 99.8 99.8 99.8
+thres: 0.9999430179595947
+imle: 99.8 99.7 99.75
+thres: 4.349513255874626e-05
+guided: 87.0 89.2 88.1
+thres: 0.011746786534786224
+ldm_200: 100.0 99.2 99.6
+thres: 0.001133577199652791
+ldm_200_cfg: 99.7 98.5 99.1
+thres: 0.0033982524182647467
+ldm_100: 99.9 99.8 99.85
+thres: 0.0011845811968669295
+glide_100_27: 99.7 98.4 99.05
+thres: 0.001386432908475399
+glide_50_27: 99.7 99.2 99.45
+thres: 0.0010550597216933966
+glide_100_10: 99.6 99.2 99.4
+thres: 0.0013092183507978916
+dalle: 99.7 98.8 99.25
+avg: 98.74 98.74 97.94
+-----------------------------------------
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/results/clip_vitl14_svd/iters_18000/ap.txt b/clean/image/effort/UniversalFakeDetect_Benchmark/results/clip_vitl14_svd/iters_18000/ap.txt
new file mode 100644
index 0000000000000000000000000000000000000000..4a936910fd1b6d3b4adc30b0d983727ecdb97f31
--- /dev/null
+++ b/clean/image/effort/UniversalFakeDetect_Benchmark/results/clip_vitl14_svd/iters_18000/ap.txt
@@ -0,0 +1,22 @@
+-----------------------------------------
+progan: 100.0
+cyclegan: 100.0
+biggan: 99.99
+stylegan: 99.77
+gaugan: 100.0
+stargan: 100.0
+deepfake: 98.95
+sitd: 97.53
+san: 96.22
+crn: 100.0
+imle: 100.0
+guided: 95.39
+ldm_200: 99.99
+ldm_200_cfg: 99.89
+ldm_100: 100.0
+glide_100_27: 99.87
+glide_50_27: 99.92
+glide_100_10: 99.98
+dalle: 99.96
+avg: 99.34
+-----------------------------------------
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/script/test.sh b/clean/image/effort/UniversalFakeDetect_Benchmark/script/test.sh
new file mode 100644
index 0000000000000000000000000000000000000000..6c651ebdf9489c1d06f9982d0704d97897210c5f
--- /dev/null
+++ b/clean/image/effort/UniversalFakeDetect_Benchmark/script/test.sh
@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+
+PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \
+CUDA_VISIBLE_DEVICES=0 python validate.py \
+ --arch CLIP:ViT-L/14_svd \
+ --ckpt checkpoints/clip_vitl14_svd/model_iters_18000.pth \
+ --result_folder results/clip_vitl14_svd/iters_18000 \
+ --use_svd
\ No newline at end of file
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/script/train_svd.sh b/clean/image/effort/UniversalFakeDetect_Benchmark/script/train_svd.sh
new file mode 100644
index 0000000000000000000000000000000000000000..7d7f0449c9c17821a74ba9048a50be77fa5ec21d
--- /dev/null
+++ b/clean/image/effort/UniversalFakeDetect_Benchmark/script/train_svd.sh
@@ -0,0 +1,40 @@
+#!/usr/bin/env bash
+
+# clip svd
+PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \
+CUDA_VISIBLE_DEVICES=1 python train.py \
+ --name clip_vitl14_svd \
+ --wang2020_data_path /Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/ \
+ --data_mode wang2020 \
+ --arch CLIP:ViT-L/14_svd \
+ --batch_size 48 \
+ --loadSize 256 \
+ --cropSize 224 \
+ --lr 0.0002 \
+ --use_svd
+
+# siglip svd
+# PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \
+# CUDA_VISIBLE_DEVICES=1 python train.py \
+# --name siglip_vitl14_svd \
+# --wang2020_data_path /Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/ \
+# --data_mode wang2020 \
+# --arch SigLIP:ViT-L/16_256_svd \
+# --batch_size 48 \
+# --loadSize 256 \
+# --cropSize 256 \
+# --lr 0.0002 \
+# --use_svd
+
+# beitv2 svd
+# PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \
+# CUDA_VISIBLE_DEVICES=1 python train.py \
+# --name beitv2_vitl14_svd \
+# --wang2020_data_path /Youtu_Pangu_Security_Public/etoilefu/CNNDetection/dataset/ \
+# --data_mode wang2020 \
+# --arch BEiTv2:ViT-L/16_svd \
+# --batch_size 48 \
+# --loadSize 256 \
+# --cropSize 224 \
+# --lr 0.0002 \
+# --use_svd
\ No newline at end of file
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/train.py b/clean/image/effort/UniversalFakeDetect_Benchmark/train.py
new file mode 100644
index 0000000000000000000000000000000000000000..735b8662aae6f5e4dd851215fa1c2086278a6f26
--- /dev/null
+++ b/clean/image/effort/UniversalFakeDetect_Benchmark/train.py
@@ -0,0 +1,93 @@
+import os
+import time
+import random
+from tensorboardX import SummaryWriter
+
+from validate import validate, find_best_threshold, RealFakeDataset
+from data import create_dataloader
+from earlystop import EarlyStopping
+from models.trainer import Trainer
+from options.train_options import TrainOptions
+from dataset_paths import DATASET_PATHS
+import torch
+import numpy as np
+
+
+SEED = 0
+def set_seed():
+ torch.manual_seed(SEED)
+ torch.cuda.manual_seed(SEED)
+ torch.cuda.manual_seed_all(SEED)
+ np.random.seed(SEED)
+ random.seed(SEED)
+ torch.backends.cudnn.deterministic = True
+ torch.backends.cudnn.benchmark = False
+
+
+def get_val_opt():
+ val_opt = TrainOptions().parse(print_options=False)
+ val_opt.isTrain = False
+ val_opt.no_resize = False
+ val_opt.no_crop = False
+ val_opt.serial_batches = True
+ val_opt.data_label = 'val'
+
+ return val_opt
+
+
+if __name__ == '__main__':
+ opt = TrainOptions().parse()
+ val_opt = get_val_opt()
+
+ set_seed()
+
+ model = Trainer(opt)
+
+ data_loader = create_dataloader(opt)
+ val_loader = create_dataloader(val_opt)
+
+ train_writer = SummaryWriter(os.path.join(opt.checkpoints_dir, opt.name, "train"))
+ val_writer = SummaryWriter(os.path.join(opt.checkpoints_dir, opt.name, "val"))
+
+ early_stopping = EarlyStopping(patience=opt.earlystop_epoch, delta=-0.001, verbose=True)
+ start_time = time.time()
+ print ("Length of data loader: %d" %(len(data_loader)))
+ with open( os.path.join(opt.checkpoints_dir, opt.name,'log.txt'), 'a') as f:
+ f.write("Length of data loader: %d \n" %(len(data_loader)) )
+ for epoch in range(opt.niter):
+ model.save_networks( 'model_epoch_init.pth' )
+
+ for i, data in enumerate(data_loader):
+ model.total_steps += 1
+
+ model.set_input(data)
+ model.optimize_parameters()
+
+ if model.total_steps % opt.loss_freq == 0:
+ print("Train loss: {} at step: {}".format(model.loss, model.total_steps))
+ train_writer.add_scalar('loss', model.loss, model.total_steps)
+ print("Iter time: ", ((time.time()-start_time)/model.total_steps) )
+ with open( os.path.join(opt.checkpoints_dir, opt.name,'log.txt'), 'a') as f:
+ f.write(f"Iter time: {(time.time()-start_time)/model.total_steps}, Lr: {model.lr}, Train loss: {model.loss} at step: {model.total_steps}\n")
+
+ if model.total_steps in [50,100,500,550,600,650,700,800,900,1000,1200,1500,2000,3000,5000,8000,10000,12000,18000,20000,23000,25000]: # save models at these iters
+ model.train()
+ model.save_networks('model_iters_%s.pth' % model.total_steps)
+
+ # if model.total_steps % 500 == 0:
+ # model.adjust_learning_rate()
+
+ if epoch % opt.save_epoch_freq == 0:
+ print('saving the model at the end of epoch %d' % (epoch))
+ model.train()
+ model.save_networks( 'model_epoch_%s.pth' % epoch )
+
+ # Validation
+ model.eval()
+ ap, r_acc, f_acc, acc = validate(model.model, val_loader)
+ val_writer.add_scalar('accuracy', acc, model.total_steps)
+ val_writer.add_scalar('ap', ap, model.total_steps)
+ print("(Val @ epoch {}) acc: {}; ap: {}".format(epoch, acc, ap))
+
+ model.train()
+
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/util.py b/clean/image/effort/UniversalFakeDetect_Benchmark/util.py
new file mode 100644
index 0000000000000000000000000000000000000000..53f23f80e7f9a37c9be67cbd81c7c67f8594706a
--- /dev/null
+++ b/clean/image/effort/UniversalFakeDetect_Benchmark/util.py
@@ -0,0 +1,21 @@
+import os
+import torch
+
+
+def mkdirs(paths):
+ if isinstance(paths, list) and not isinstance(paths, str):
+ for path in paths:
+ mkdir(path)
+ else:
+ mkdir(paths)
+
+
+def mkdir(path):
+ if not os.path.exists(path):
+ os.makedirs(path)
+
+
+def unnormalize(tens, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
+ # assume tensor of shape NxCxHxW
+ return tens * torch.Tensor(std)[None, :, None, None] + torch.Tensor(
+ mean)[None, :, None, None]
\ No newline at end of file
diff --git a/clean/image/effort/UniversalFakeDetect_Benchmark/validate.py b/clean/image/effort/UniversalFakeDetect_Benchmark/validate.py
new file mode 100644
index 0000000000000000000000000000000000000000..53da394816e8658f2da11cee77a557ed0a821948
--- /dev/null
+++ b/clean/image/effort/UniversalFakeDetect_Benchmark/validate.py
@@ -0,0 +1,359 @@
+import argparse
+from ast import arg
+import os
+import math
+import csv
+import torch
+import torchvision.transforms as transforms
+import torch.utils.data
+import numpy as np
+from sklearn.metrics import average_precision_score, precision_recall_curve, accuracy_score
+from torch.utils.data import Dataset
+import sys
+from models import get_model
+from PIL import Image
+import pickle
+from tqdm import tqdm
+from io import BytesIO
+from copy import deepcopy
+from dataset_paths import DATASET_PATHS
+import random
+import shutil
+from scipy.ndimage.filters import gaussian_filter
+
+
+SEED = 0
+def set_seed():
+ torch.manual_seed(SEED)
+ torch.cuda.manual_seed(SEED)
+ np.random.seed(SEED)
+ random.seed(SEED)
+
+
+MEAN = {
+ "imagenet":[0.485, 0.456, 0.406],
+ "clip":[0.48145466, 0.4578275, 0.40821073],
+ "beitv2": [0.485, 0.456, 0.406],
+ "siglip": [0.5, 0.5, 0.5],
+}
+
+STD = {
+ "imagenet":[0.229, 0.224, 0.225],
+ "clip":[0.26862954, 0.26130258, 0.27577711],
+ "beitv2": [0.229, 0.224, 0.225],
+ "siglip": [0.5, 0.5, 0.5],
+}
+
+
+def translate_duplicate(img, cropSize):
+ if min(img.size) < cropSize:
+ width, height = img.size
+
+ new_width = width * math.ceil(cropSize/width)
+ new_height = height * math.ceil(cropSize/height)
+
+ new_img = Image.new('RGB', (new_width, new_height))
+ for i in range(0, new_width, width):
+ for j in range(0, new_height, height):
+ new_img.paste(img, (i, j))
+ return new_img
+ else:
+ return img
+
+
+def find_best_threshold(y_true, y_pred):
+ "We assume first half is real 0, and the second half is fake 1"
+
+ N = y_true.shape[0]
+
+ if y_pred[0:N//2].max() <= y_pred[N//2:N].min(): # perfectly separable case
+ return (y_pred[0:N//2].max() + y_pred[N//2:N].min()) / 2
+
+ best_acc = 0
+ best_thres = 0
+ for thres in y_pred:
+ temp = deepcopy(y_pred)
+ temp[temp>=thres] = 1
+ temp[temp= best_acc:
+ best_thres = thres
+ best_acc = acc
+
+ return best_thres
+
+
+def png2jpg(img, quality):
+ out = BytesIO()
+ img.save(out, format='jpeg', quality=quality) # ranging from 0-95, 75 is default
+ img = Image.open(out)
+ # load from memory before ByteIO closes
+ img = np.array(img)
+ out.close()
+
+ return Image.fromarray(img)
+
+
+def gaussian_blur(img, sigma):
+ img = np.array(img)
+
+ gaussian_filter(img[:,:,0], output=img[:,:,0], sigma=sigma)
+ gaussian_filter(img[:,:,1], output=img[:,:,1], sigma=sigma)
+ gaussian_filter(img[:,:,2], output=img[:,:,2], sigma=sigma)
+
+ return Image.fromarray(img)
+
+
+def calculate_acc(y_true, y_pred, thres):
+ r_acc = accuracy_score(y_true[y_true==0], y_pred[y_true==0] > thres)
+ f_acc = accuracy_score(y_true[y_true==1], y_pred[y_true==1] > thres)
+ acc = accuracy_score(y_true, y_pred > thres)
+
+ return r_acc, f_acc, acc
+
+
+def validate(model, loader, find_thres=False):
+ with torch.no_grad():
+ y_true, y_pred = [], []
+ print ("Length of dataset: %d" %(len(loader)))
+ for img, label in loader:
+ in_tens = img.cuda()
+
+ y_pred.extend(model(in_tens).sigmoid().flatten().tolist())
+ # y_pred.extend(model(in_tens).softmax(dim=1)[:, 1].flatten().tolist())
+ y_true.extend(label.flatten().tolist())
+
+ y_true, y_pred = np.array(y_true), np.array(y_pred)
+
+ # Get AP
+ ap = average_precision_score(y_true, y_pred)
+
+ # Acc based on 0.5
+ r_acc0, f_acc0, acc0 = calculate_acc(y_true, y_pred, 0.5)
+ if not find_thres:
+ return ap, r_acc0, f_acc0, acc0
+
+ # Acc based on the best thres
+ best_thres = find_best_threshold(y_true, y_pred)
+ r_acc1, f_acc1, acc1 = calculate_acc(y_true, y_pred, best_thres)
+
+ return ap, r_acc0, f_acc0, acc0, r_acc1, f_acc1, acc1, best_thres
+
+
+
+
+def recursively_read(rootdir, must_contain, classes=[], exts=["png", "jpg", "JPEG", "jpeg"]):
+ out = []
+ for r, d, f in os.walk(rootdir):
+ for file in f:
+ if (file.split('.')[1] in exts) and (must_contain in os.path.join(r, file)):
+ if len(classes) == 0:
+ out.append(os.path.join(r, file))
+ elif os.path.join(r, file).split('/')[-3] in classes:
+ out.append(os.path.join(r, file))
+
+ return out
+
+
+def get_list(path, must_contain='', classes=[]):
+ if ".pickle" in path:
+ with open(path, 'rb') as f:
+ image_list = pickle.load(f)
+ image_list = [ item for item in image_list if must_contain in item ]
+ else:
+ image_list = recursively_read(path, must_contain, classes)
+
+ return image_list
+
+
+class RealFakeDataset(Dataset):
+ def __init__(self, real_path,
+ fake_path,
+ data_mode,
+ max_sample,
+ arch,
+ jpeg_quality=None,
+ gaussian_sigma=None):
+
+ assert data_mode in ["wang2020", "ours"]
+ self.jpeg_quality = jpeg_quality
+ self.gaussian_sigma = gaussian_sigma
+
+ # = = = = = = data path = = = = = = = = = #
+ if type(real_path) == str and type(fake_path) == str:
+ real_list, fake_list = self.read_path(real_path, fake_path, data_mode, max_sample)
+ else:
+ real_list = []
+ fake_list = []
+ for real_p, fake_p in zip(real_path, fake_path):
+ real_l, fake_l = self.read_path(real_p, fake_p, data_mode, max_sample)
+ real_list += real_l
+ fake_list += fake_l
+
+ self.total_list = real_list + fake_list
+
+
+ # = = = = = = label = = = = = = = = = #
+ self.labels_dict = {}
+ for i in real_list:
+ self.labels_dict[i] = 0
+ for i in fake_list:
+ self.labels_dict[i] = 1
+
+ if arch.lower().startswith("imagenet"):
+ stat_from = "imagenet"
+ elif arch.lower().startswith("clip"):
+ stat_from = "clip"
+ elif arch.lower().startswith("siglip"):
+ stat_from = "siglip"
+ elif arch.lower().startswith("beitv2"):
+ stat_from = "beitv2"
+
+ self.transform = transforms.Compose([
+ # transforms.Resize((256, 256)),
+ transforms.Lambda(lambda img: translate_duplicate(img, 256)),
+ transforms.CenterCrop(224) if stat_from != "siglip" else transforms.CenterCrop(256),
+ transforms.ToTensor(),
+ transforms.Normalize( mean=MEAN[stat_from], std=STD[stat_from] ),
+ ])
+
+ def read_path(self, real_path, fake_path, data_mode, max_sample):
+ if data_mode == 'wang2020':
+ real_list = get_list(real_path, must_contain='0_real')
+ fake_list = get_list(fake_path, must_contain='1_fake')
+ else:
+ real_list = get_list(real_path)
+ fake_list = get_list(fake_path)
+
+ if max_sample is not None:
+ if (max_sample > len(real_list)) or (max_sample > len(fake_list)):
+ max_sample = 100
+ print("not enough images, max_sample falling to 100")
+ random.shuffle(real_list)
+ random.shuffle(fake_list)
+ real_list = real_list[0:max_sample]
+ fake_list = fake_list[0:max_sample]
+
+ assert len(real_list) == len(fake_list)
+
+ return real_list, fake_list
+
+ def __len__(self):
+ return len(self.total_list)
+
+ def __getitem__(self, idx):
+ img_path = self.total_list[idx]
+
+ label = self.labels_dict[img_path]
+ img = Image.open(img_path).convert("RGB")
+
+ if self.gaussian_sigma is not None:
+ img = gaussian_blur(img, self.gaussian_sigma)
+ if self.jpeg_quality is not None:
+ img = png2jpg(img, self.jpeg_quality)
+
+ img = self.transform(img)
+
+ return img, label
+
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+ parser.add_argument('--real_path', type=str, default=None, help='dir name or a pickle')
+ parser.add_argument('--fake_path', type=str, default=None, help='dir name or a pickle')
+ parser.add_argument('--data_mode', type=str, default=None, help='wang2020 or ours')
+ parser.add_argument('--max_sample', type=int, default=1000, help='only check this number of images for both fake/real')
+
+ parser.add_argument('--arch', type=str, default='res50')
+ parser.add_argument('--ckpt', type=str, default='./pretrained_weights/fc_weights.pth')
+
+ parser.add_argument('--result_folder', type=str, default='result', help='')
+ parser.add_argument('--batch_size', type=int, default=128)
+
+ parser.add_argument('--jpeg_quality', type=int, default=None, help="100, 90, 80, ... 30. Used to test robustness of our model. Not apply if None")
+ parser.add_argument('--gaussian_sigma', type=int, default=None, help="0,1,2,3,4. Used to test robustness of our model. Not apply if None")
+
+ parser.add_argument('--use_svd', action='store_true')
+ parser.add_argument('--use_lora', action='store_true')
+
+ opt = parser.parse_args()
+
+ os.makedirs(opt.result_folder, exist_ok=True)
+
+ model = get_model(opt.arch, opt)
+ if opt.use_svd:
+ state_dict = torch.load(opt.ckpt, map_location='cpu')['model']
+ model.load_state_dict(state_dict)
+ else:
+ state_dict = torch.load(opt.ckpt, map_location='cpu')['model']
+ model.fc.load_state_dict(state_dict)
+
+ print ("Model loaded..")
+ model.eval()
+ model.cuda()
+
+ if (opt.real_path == None) or (opt.fake_path == None) or (opt.data_mode == None):
+ dataset_paths = DATASET_PATHS
+ else:
+ dataset_paths = [ dict(real_path=opt.real_path, fake_path=opt.fake_path, data_mode=opt.data_mode) ]
+
+ with open( os.path.join(opt.result_folder,'ap.txt'), 'a') as f:
+ f.write('-----------------------------------------'+'\n' )
+ with open( os.path.join(opt.result_folder,'acc0.txt'), 'a') as f:
+ f.write('-----------------------------------------'+'\n' )
+ with open( os.path.join(opt.result_folder,'acc1.txt'), 'a') as f:
+ f.write('-----------------------------------------'+'\n' )
+
+ ap_avg, r_acc0_avg, f_acc0_avg, acc0_avg, r_acc1_avg, f_acc1_avg, acc1_avg = 0, 0, 0, 0, 0, 0, 0
+ for dataset_path in (dataset_paths):
+ set_seed()
+
+ dataset = RealFakeDataset( dataset_path['real_path'],
+ dataset_path['fake_path'],
+ dataset_path['data_mode'],
+ opt.max_sample,
+ opt.arch,
+ jpeg_quality=opt.jpeg_quality,
+ gaussian_sigma=opt.gaussian_sigma,
+ )
+
+ loader = torch.utils.data.DataLoader(dataset, batch_size=opt.batch_size, shuffle=False, num_workers=4)
+ ap, r_acc0, f_acc0, acc0, r_acc1, f_acc1, acc1, best_thres = validate(model, loader, find_thres=True)
+ with open( os.path.join(opt.result_folder,'acc0.txt'), 'a') as f:
+ f.write(f"thres: 0.5\n")
+ with open( os.path.join(opt.result_folder,'acc1.txt'), 'a') as f:
+ f.write(f"thres: {best_thres}\n")
+
+ ap_avg += ap
+ r_acc0_avg += r_acc0
+ f_acc0_avg += f_acc0
+ acc0_avg += acc0
+ r_acc1_avg += r_acc1
+ f_acc1_avg += r_acc1
+ acc1_avg += acc1
+
+ with open( os.path.join(opt.result_folder,'ap.txt'), 'a') as f:
+ f.write(dataset_path['key']+': ' + str(round(ap*100, 2))+'\n' )
+ with open( os.path.join(opt.result_folder,'acc0.txt'), 'a') as f:
+ f.write(dataset_path['key']+': ' + str(round(r_acc0*100, 2))+' '+str(round(f_acc0*100, 2))+' '+str(round(acc0*100, 2))+'\n' )
+ with open( os.path.join(opt.result_folder,'acc1.txt'), 'a') as f:
+ f.write(dataset_path['key']+': ' + str(round(r_acc1*100, 2))+' '+str(round(f_acc1*100, 2))+' '+str(round(acc1*100, 2))+'\n' )
+
+ ap_avg /= len(dataset_paths)
+ r_acc0_avg /= len(dataset_paths)
+ f_acc0_avg /= len(dataset_paths)
+ acc0_avg /= len(dataset_paths)
+ r_acc1_avg /= len(dataset_paths)
+ f_acc1_avg /= len(dataset_paths)
+ acc1_avg /= len(dataset_paths)
+
+ with open( os.path.join(opt.result_folder,'ap.txt'), 'a') as f:
+ f.write('avg: ' + str(round(ap_avg*100, 2))+'\n' )
+ f.write('-----------------------------------------'+'\n' )
+ with open( os.path.join(opt.result_folder,'acc0.txt'), 'a') as f:
+ f.write('avg: ' + str(round(r_acc0_avg*100, 2))+' '+str(round(f_acc0_avg*100, 2))+' '+str(round(acc0_avg*100, 2))+'\n' )
+ f.write('-----------------------------------------'+'\n' )
+ with open( os.path.join(opt.result_folder,'acc1.txt'), 'a') as f:
+ f.write('avg: ' + str(round(r_acc1_avg*100, 2))+' '+str(round(f_acc1_avg*100, 2))+' '+str(round(acc1_avg*100, 2))+'\n' )
+ f.write('-----------------------------------------'+'\n' )
\ No newline at end of file
diff --git a/clean/image/effort/install.sh b/clean/image/effort/install.sh
new file mode 100644
index 0000000000000000000000000000000000000000..b0ba2738e2a4191cbd0d40eb972cf41aef63b360
--- /dev/null
+++ b/clean/image/effort/install.sh
@@ -0,0 +1,33 @@
+#!/bin/bash
+
+pip install numpy==1.21.5
+pip install pandas==1.4.2
+pip install Pillow==9.0.1
+pip install dlib==19.24.0
+pip install imageio==2.9.0
+pip install imgaug==0.4.0
+pip install tqdm==4.61.0
+pip install scipy==1.7.3
+pip install seaborn==0.11.2
+pip install pyyaml==6.0
+pip install imutils==0.5.4
+pip install opencv-python==4.6.0.66
+pip install scikit-image==0.19.2
+pip install scikit-learn==1.0.2
+pip install albumentations==1.1.0
+pip install torch==1.12.0+cu113 torchvision==0.13.0+cu113 torchaudio==0.12.0 --extra-index-url https://download.pytorch.org/whl/cu113
+pip install efficientnet-pytorch==0.7.1
+pip install timm==0.6.12
+pip install segmentation-models-pytorch==0.3.2
+pip install torchtoolbox==0.1.8.2
+pip install tensorboard==2.10.1
+pip install setuptools==59.5.0
+pip install loralib
+pip install einops
+pip install transformers
+pip install filterpy
+pip install simplejson
+pip install kornia
+pip install fvcore
+pip install imgaug==0.4.0
+pip install git+https://github.com/openai/CLIP.git
diff --git a/clean/image/fsd/.gitignore b/clean/image/fsd/.gitignore
new file mode 100644
index 0000000000000000000000000000000000000000..7579d281285624f503dae661196ea4e3c2e654fb
--- /dev/null
+++ b/clean/image/fsd/.gitignore
@@ -0,0 +1,8 @@
+__pycache__/
+*.pyc
+*.egg-info/
+.venv/
+dist/
+build/
+internal/
+weights/
diff --git a/clean/image/fsd/CITATION.cff b/clean/image/fsd/CITATION.cff
new file mode 100644
index 0000000000000000000000000000000000000000..83eb011baef5a1120bd9832a6ebf84708df78bb6
--- /dev/null
+++ b/clean/image/fsd/CITATION.cff
@@ -0,0 +1,44 @@
+cff-version: 1.2.0
+message: "If you use this software, please cite it as below."
+title: "Forensic Self-Descriptions Are All You Need for Zero-Shot Detection, Open-Set Source Attribution, and Clustering of AI-generated Images"
+type: software
+authors:
+ - given-names: Tai D.
+ family-names: Nguyen
+ email: tdn47@drexel.edu
+ affiliation: Drexel University
+ - given-names: Aref
+ family-names: Azizpour
+ affiliation: Drexel University
+ - given-names: Matthew C.
+ family-names: Stamm
+ affiliation: Drexel University
+repository-code: "https://github.com/ductai199x/Forensic-Self-Descriptions-CVPR25"
+url: "https://arxiv.org/abs/2503.21003"
+license: CC-BY-NC-SA-4.0
+keywords:
+ - deepfake detection
+ - ai-generated image detection
+ - image forensics
+ - zero-shot detection
+ - computer vision
+ - pytorch
+ - cvpr 2025
+ - stable diffusion
+ - midjourney
+ - generative ai
+preferred-citation:
+ type: conference-paper
+ title: "Forensic Self-Descriptions Are All You Need for Zero-Shot Detection, Open-Set Source Attribution, and Clustering of AI-generated Images"
+ authors:
+ - given-names: Tai D.
+ family-names: Nguyen
+ - given-names: Aref
+ family-names: Azizpour
+ - given-names: Matthew C.
+ family-names: Stamm
+ collection-title: "Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)"
+ year: 2025
+ month: 6
+ start: 3040
+ end: 3050
diff --git a/clean/image/fsd/LICENSE b/clean/image/fsd/LICENSE
new file mode 100644
index 0000000000000000000000000000000000000000..682189be62f508be84a0bac99fe8f552a33d4929
--- /dev/null
+++ b/clean/image/fsd/LICENSE
@@ -0,0 +1,22 @@
+Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International
+
+Copyright (c) 2025 Tai D. Nguyen, Multimedia Information Security Lab, Drexel University
+
+This work is licensed under the Creative Commons
+Attribution-NonCommercial-ShareAlike 4.0 International License.
+
+You are free to:
+ - Share: copy and redistribute the material in any medium or format
+ - Adapt: remix, transform, and build upon the material
+
+Under the following terms:
+ - Attribution: You must give appropriate credit, provide a link to the
+ license, and indicate if changes were made.
+ - NonCommercial: You may not use the material for commercial purposes.
+ - ShareAlike: If you remix, transform, or build upon the material, you
+ must distribute your contributions under the same license.
+
+No additional restrictions: You may not apply legal terms or technological
+measures that legally restrict others from doing anything the license permits.
+
+Full license text: https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
diff --git a/clean/image/fsd/README.md b/clean/image/fsd/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..7a5d3c8227d57f1f39a5631595fd2eafdbe3731c
--- /dev/null
+++ b/clean/image/fsd/README.md
@@ -0,0 +1,228 @@
+# Forensic Self-Descriptions (FSD): Zero-Shot AI-Generated Image Detection
+
+[](https://cvpr.thecvf.com/)
+[](https://arxiv.org/abs/2503.21003)
+[](https://creativecommons.org/licenses/by-nc-sa/4.0/)
+[](https://arxiv.org/abs/2503.21003)
+
+**Zero-shot AI-generated image detection -- trained only on real images, generalizes to any unseen generator.**
+
+Official PyTorch implementation of **"Forensic Self-Descriptions Are All You Need for Zero-Shot Detection, Open-Set Source Attribution, and Clustering of AI-generated Images"** (CVPR 2025) by Tai D. Nguyen, Aref Azizpour, and Matthew C. Stamm.
+
+> **TL;DR:** FSD is a deepfake / AI-generated image detector that achieves **96.0% average AUC** across 24 generators (Stable Diffusion, Midjourney, DALL-E, StyleGAN, etc.) while being trained **exclusively on real photographs** — no synthetic training data required.
+
+
+
+
+
+
+
+
+
+## Updates
+
+- **2026-03**: Added source attribution -- identify which AI generator created an image (14 sources supported).
+- **2026-03**: Weights now auto-download from GitHub releases on first use.
+- **2026-02**: Released inference code and pre-trained model weights for AI-generated image detection.
+
+## Roadmap
+
+- [x] Code for open-set source attribution
+- [ ] Code for unsupervised clustering
+
+## Overview
+
+FSD is a self-supervised forensic method that detects AI-generated images without needing to train on any specific generator. It works by:
+
+1. **Forensic Residual Extraction (FRE)**: Constrained prediction-error filters extract pixel-level forensic residuals
+2. **Multi-scale FSD computation**: Residuals are analyzed across scales to produce a compact 960-dimensional forensic descriptor
+3. **GMM scoring**: A Gaussian Mixture Model scores each descriptor, yielding a z-score where more negative values indicate AI-generated content
+
+## Results
+
+FSD achieves **state-of-the-art** zero-shot synthetic image detection while being **completely zero-shot** -- it is trained **only on real images** and has never seen any synthetic image during training. Unlike most competing methods which require synthetic training data from specific generators, FSD generalizes to any generator out of the box.
+
+**Zero-shot detection performance** (average AUC across 24 generators including ProGAN, StyleGAN 1-3, GigaGAN, GLIDE, Stable Diffusion 1.5-3.0, DALLE, Midjourney, Firefly, etc.):
+
+| Method | Training Data | COCO17 | IN-1k | IN-22k | MIDB | Average |
+|--------|:---:|:---:|:---:|:---:|:---:|:---:|
+| CNNDet | Real + Synthetic | 0.756 | 0.714 | 0.733 | 0.683 | 0.722 |
+| PatchFor | Real + Synthetic | 0.833 | 0.823 | 0.845 | 0.790 | 0.823 |
+| UFD | Real + Synthetic | 0.903 | 0.862 | 0.815 | 0.612 | 0.798 |
+| LGrad | Real + Synthetic | 0.819 | 0.770 | 0.866 | 0.824 | 0.820 |
+| DE-FAKE | Real + Synthetic | 0.765 | 0.749 | 0.617 | 0.791 | 0.731 |
+| Aeroblade | Training-Free | 0.728 | 0.741 | 0.582 | 0.646 | 0.674 |
+| ZED | Real Only | 0.751 | 0.676 | 0.716 | 0.747 | 0.723 |
+| NPR | Real + Synthetic | 0.945 | 0.900 | 0.900 | 0.957 | 0.926 |
+| **Ours (FSD)** | **Real Only** | **0.968** | **0.962** | **0.941** | **0.971** | **0.960** |
+
+See [the paper](https://arxiv.org/abs/2503.21003) for full results on source attribution and clustering.
+
+## Installation
+
+This project uses [uv](https://docs.astral.sh/uv/) for dependency management.
+
+```bash
+git clone https://github.com/ductai199x/Forensic-Self-Descriptions-CVPR25.git
+cd Forensic-Self-Descriptions-CVPR25
+
+# Install dependencies and create virtual environment
+uv sync
+
+# Activate the virtual environment
+source .venv/bin/activate
+```
+
+## Quick Start
+
+### Python API
+
+```python
+from fsd import FSDDetector
+
+# Detection only
+detector = FSDDetector.load()
+result = detector.score("photo.jpg")
+
+print(result.z_score) # e.g., -3.5 (negative = likely fake)
+print(result.is_fake) # True/False based on threshold
+```
+
+Score multiple images:
+```python
+results = detector.score_batch(["img1.jpg", "img2.png", "img3.webp"])
+for path, result in zip(paths, results):
+ print(f"{path}: z={result.z_score:.4f} {'FAKE' if result.is_fake else 'REAL'}")
+```
+
+#### Source Attribution
+
+Identify which AI generator created an image:
+
+```python
+# Load with attribution support
+detector = FSDDetector.load(attribution=True)
+result = detector.attribute("suspicious_image.jpg")
+
+print(result.source) # e.g., "Stable Diffusion XL"
+print(result.confidence) # e.g., 0.95
+print(result.is_fake) # True
+print(result.scores) # per-source log-likelihoods
+```
+
+Supported sources: DALL-E 3, Stable Diffusion 1.5/3/XL, Midjourney v6, Adobe Firefly, StyleGAN2/3, ProGAN, GigaGAN, Grok, GPT-Image 1/1.5, and more.
+
+### Command Line
+
+```bash
+# Single image
+fsd-score photo.jpg
+
+# Multiple images
+fsd-score img1.jpg img2.png img3.webp
+
+# Directory of images
+fsd-score --dir path/to/images/
+
+# With source attribution
+fsd-score photo.jpg --attribute
+
+# Custom threshold (default: -2.0, more negative = stricter)
+fsd-score photo.jpg --threshold -3.0
+
+# Use GPU
+fsd-score photo.jpg --device cuda
+
+# CSV output
+fsd-score --dir images/ --csv > results.csv
+```
+
+### Gradio Demo
+
+An interactive web demo for testing images in your browser:
+
+```bash
+# Launch (auto-detects GPU)
+uv run demo.py
+
+# Create a public shareable link
+uv run demo.py --share
+
+# Force CPU-only
+uv run demo.py --device cpu
+```
+
+### Multi-GPU / Ray Serve
+
+For scoring large batches across multiple GPUs, start the Ray Serve service first, then score images against it:
+
+```bash
+# Start the scoring service (auto-detects GPUs)
+fsd-score-ray serve
+
+# In another terminal, score images against the running service
+fsd-score-ray score photo.jpg
+fsd-score-ray score --dir path/to/images/ --csv > results.csv
+```
+
+Configure the service:
+```bash
+# Custom port and GPU allocation
+fsd-score-ray serve --port 9000 --num-gpus 4 --gpu-per-replica 0.5
+
+# Score against non-default port
+fsd-score-ray score --url http://localhost:9000 --dir images/
+```
+
+You can also query the service directly via HTTP:
+```bash
+curl -X POST http://localhost:8000 \
+ -H "Content-Type: application/json" \
+ -d '{"path": "/absolute/path/to/image.jpg"}'
+```
+
+## Interpreting Results
+
+The detector outputs a **z-score** for each image:
+- **z > -2**: Likely real
+- **z < -2**: Likely AI-generated (default threshold)
+- **z < -3**: High confidence AI-generated
+
+> **Note:** Detection is significantly more reliable than attribution. Detection is zero-shot (trained only on real images) and generalizes to any generator with 96% average AUC. Attribution, on the other hand, can only identify sources it has been trained on and may misclassify images from unknown generators. Always trust the detection result over the attribution result.
+
+## Pre-trained Weights
+
+Weights are automatically downloaded from [GitHub releases](https://github.com/ductai199x/Forensic-Self-Descriptions-CVPR25/releases) on first use and cached to `~/.cache/fsd/`. No manual download needed.
+
+**Detection weights:**
+- `fre.pt` -- Forensic Residual Extractor (constrained convolution, ~10 KB)
+- `gmm.pt` -- Gaussian Mixture Model (K=5, tied covariance, ~15 MB)
+- `fsd_transforms.pt` -- Detection feature transforms (~40 MB)
+- `config.json` -- Model configuration and scoring parameters
+
+**Attribution weights** (downloaded when `attribution=True`):
+- `attribution_transforms.pt` -- Attribution feature transform (~26 MB)
+- `source_gmms.pt` -- Per-source GMMs for 14 generators (~207 MB)
+
+## Citation
+
+If you find this work useful, please cite:
+
+```bibtex
+@InProceedings{Nguyen_2025_CVPR,
+ author = {Nguyen, Tai D. and Azizpour, Aref and Stamm, Matthew C.},
+ title = {Forensic Self-Descriptions Are All You Need for Zero-Shot Detection, Open-Set Source Attribution, and Clustering of AI-generated Images},
+ booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+ month = {June},
+ year = {2025},
+ pages = {3040-3050}
+}
+```
+
+## Acknowledgments
+
+This work was conducted at the [Multimedia Information Security Lab (MISL)](https://misl.ece.drexel.edu/) at Drexel University under the supervision of Dr. Matthew C. Stamm.
+
+## License
+
+This project is licensed under [CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/) -- research use only, no commercial use, share-alike.
diff --git a/clean/image/fsd/SOURCE.md b/clean/image/fsd/SOURCE.md
new file mode 100644
index 0000000000000000000000000000000000000000..36e5264e455fbe23604ebae1caac41a19b2e281f
--- /dev/null
+++ b/clean/image/fsd/SOURCE.md
@@ -0,0 +1,17 @@
+# Source: image/fsd
+
+| Field | Value |
+|---|---|
+| Upstream | https://github.com/ductai199x/Forensic-Self-Descriptions-CVPR25 |
+| Paper | https://arxiv.org/abs/2503.21003 |
+| Commit SHA | **not recorded** -- the vendoring step did not preserve it |
+| Mirrored on | 2026-09-15 |
+| Upstream license | LICENSE |
+
+This is a **mirror**, stripped to the files needed for inference. The full
+untouched snapshot is at `archive/image__fsd.tar.gz`.
+
+This code is the work of its original authors and is **not** covered by the
+DeepSafe project license. If you are an author and want this removed, open an
+issue on https://github.com/deepsafehq/deepsafe-bench and it will be taken down
+within 48 hours, no questions asked.
diff --git a/clean/image/fsd/demo.py b/clean/image/fsd/demo.py
new file mode 100644
index 0000000000000000000000000000000000000000..7508c562683acfd287b1978a017c799a0948ce8c
--- /dev/null
+++ b/clean/image/fsd/demo.py
@@ -0,0 +1,469 @@
+"""Gradio demo for FSD: Detecting AI-Generated Images via Forensic Self-Descriptions.
+
+Usage:
+ uv run demo.py
+ uv run demo.py --share
+ uv run demo.py --device cpu
+"""
+
+import argparse
+
+import gradio as gr
+from PIL import Image
+# Register HEIF/HEIC support before any image loading
+try:
+ from pillow_heif import register_heif_opener
+
+ register_heif_opener()
+except ImportError:
+ pass
+
+from fsd import FSDDetector, DetectionResult, AttributionResult
+
+# ---------------------------------------------------------------------------
+# Palette (dark theme, user-provided from coolors.co)
+# ---------------------------------------------------------------------------
+JET_BLACK = "#2d3142" # dark bg
+BEIGE = "#e9edde" # primary text on dark
+BANANA_CREAM = "#e7e247" # warning / uncertain accent
+GLAUCOUS = "#5c80bc" # links, buttons, secondary accent
+PEARL_AQUA = "#69d1c5" # positive / real accent
+
+CARD_BG = "#363b50" # slightly lighter than jet for cards
+MUTED = "#9a9eb0" # subdued text
+
+
+# ---------------------------------------------------------------------------
+# Result rendering
+# ---------------------------------------------------------------------------
+
+def _verdict(z: float, threshold: float):
+ if z >= -1.0:
+ return ("Real", "Forensic signature is consistent with real photographs.", "verdict-real")
+ if z >= threshold:
+ return ("Likely Real", "Leans toward real, but not a definitive match.", "verdict-likely-real")
+ if z >= threshold - 1.0:
+ return ("Likely AI", "Shows signs of AI generation in its forensic signature.", "verdict-likely-ai")
+ return ("AI-Generated", "Forensic signature strongly deviates from real photographs.", "verdict-ai")
+
+
+def _prob_fake(z: float, threshold: float = -2.0, k: float = 2.0) -> float:
+ """Sigmoid centered at the decision threshold so z=threshold -> exactly 50%.
+
+ k=2.0 gives: z=0 -> 2%, z=-1 -> 12%, z=-2 -> 50%, z=-3 -> 88%, z=-4 -> 98%.
+ """
+ import math
+ return 1.0 / (1.0 + math.exp(-k * (threshold - z)))
+
+
+def _build_attribution_html(attr_result: AttributionResult) -> str:
+ """Build horizontal bar chart HTML for source attribution scores."""
+ # Scores are already calibrated probabilities (z-score-normalized softmax)
+ ranked = sorted(attr_result.scores.items(), key=lambda x: x[1], reverse=True)
+
+ bars_html = ""
+ for name, prob in ranked:
+ pct = prob * 100
+ is_best = (name == attr_result.source)
+ bar_cls = "attr-bar-best" if is_best else ""
+ label_cls = "attr-label-best" if is_best else ""
+ bars_html += f"""
+
Forensic Residual Extraction — Learned prediction-error filters
+ capture subtle pixel-level traces that differ between real cameras and AI generators.
+
Self-Description Computation — Multi-scale patch analysis produces
+ a compact 960-dimensional forensic fingerprint of the image.
+
Statistical Scoring — A Gaussian Mixture Model, trained exclusively
+ on real photographs, measures how well the fingerprint matches natural image statistics.
+
Z-Score & Decision — The score is normalized into a z-score
+ (standard deviations from the real-image mean). More negative = less like a real photo.
+
Source Attribution — If an image is flagged as AI-generated,
+ per-source statistical models identify which generator most likely produced it.
+
+
+ Interpreting the results
+
+
Z-score above −1
+
Forensic signature matches real photographs — very likely real.
+
Z-score −1 to −2
+
Still within the real range — likely a genuine photograph.
+
Z-score −2 to −3
+
Crosses the detection threshold — likely AI-generated.
+
Z-score below −3
+
Far beyond the threshold — very likely AI-generated.
+
+
+
+ This detector is trained only on real photographs and has never seen AI-generated images.
+ It generalizes to new generators zero-shot, but accuracy may vary with heavy JPEG
+ compression, screenshots, or other post-processing.
+
+Generative models have enabled the creation of highly realistic facial-synthetic images, raising significant concerns due to their potential for misuse. Despite rapid advancements in the field of deepfake detection, developing efficient approaches to leverage foundation models for improved generalizability to unseen forgery samples remains challenging. To address this challenge, we propose a novel side-network-based decoder that extracts spatial and temporal cues using the CLIP image encoder for generalized video-based Deepfake detection. Additionally, we introduce Facial Component Guidance (FCG) to enhance spatial learning generalizability by encouraging the model to focus on key facial regions. By leveraging the generic features of a vision-language foundation model, our approach demonstrates promising generalizability on challenging Deepfake datasets while also exhibiting superiority in training data efficiency, parameter efficiency, and model robustness.
+
+
+## 📝TODOs
+ - [x] Training + Evaluation Code
+ - [x] Model Weights
+ - [x] Inference Code
+ - [ ] HeyGen Evaluation Dataset
+
+
+## 🙌News
+ - June 08: We have released the [model checkpoint](https://drive.google.com/file/d/1ydD5rnaaF0i2zLE7NidLtAhjonHoVQOk/view?usp=sharing) and provided inference code for single videos! Checkout [this section](#inference---demo-video) for further details!
+
+## 🚀Installation
+```shell
+# conda environment
+conda env create -f environment.yml
+```
+
+## 📂Dataset Structure
+The structure of the **pre-processed datasets** for our project, the video files (*.avi) have been processed to only retain the aligned face. We use soft-links **(ln -s)** to manage and link the folders containing pre-processed videos on different drives.
+```shell
+datasets
+├── cdf
+│ ├── FAKE
+│ │ └── videos
+│ │ └── *.avi
+│ ├── REAL
+│ │ └── videos
+│ │ └── *.avi
+│ └── csv_files
+│ ├── test_fake.csv
+│ └── test_real.csv
+├── dfdc
+│ ├── csv_files
+│ │ └── test.csv
+│ └── videos
+├── dfo
+│ ├── FAKE
+│ │ └── videos
+│ │ └── *.avi
+│ ├── REAL
+│ │ └── videos
+│ │ └── *.avi
+│ └── csv_files
+│ ├── test_fake.csv
+│ └── test_real.csv
+├── ffpp
+│ ├── DF
+│ │ ├── c23
+│ │ │ └── videos
+│ │ │ └── *.avi
+│ │ ├── c40
+│ │ │ └── videos
+│ │ │ └── *.avi
+│ │ └── raw
+│ │ └── videos
+│ │ │ └── *.avi
+│ ├── F2F ...
+│ ├── FS ...
+│ ├── FSh ...
+│ ├── NT ...
+│ ├── real ...
+│ └── csv_files
+│ ├── test.json
+│ ├── train.json
+│ └── val.json
+|
+└── robustness
+ ├── BW
+ │ ├── 1
+ │ │ ├── DF
+ │ │ │ └── c23
+ │ │ │ └── videos
+ │ │ │ └── *.avi
+ │ │ ├── F2F ...
+ │ │ ├── FS ...
+ │ │ ├── FSh ...
+ │ │ ├── NT ...
+ │ │ ├── real ...
+ │ │ │
+ │ │ └── csv_files
+ │ │ ├── test.json
+ │ │ ├── train.json
+ │ │ └── val.json
+ │ │
+ │ │
+ │ │
+ . .
+ . .
+ . .
+```
+
+
+## 🔧Dataset Pre-processing
+### Generic Pre-processing
+This phase performs the required pre-processing for our method, this includes *facial alignment (using the mean face from LRW)* and *facial cropping*.
+```bash
+# First, fetch all the landmarks & bboxes of the video frames.
+python -m src.preprocess.fetch_landmark_bbox \
+--root-dir="/storage/FaceForensicC23" \ # The root folder of the dataset
+--video-dir="videos" \ # The root folder of the videos
+--fdata-dir="frame_data" \ # The folder to save the extracted frame data
+--glob-exp="*/*" \ # The glob expression to search through the root video folder
+--split-num=1 \ # Split the dataset into several parts for parallel process.
+--part-num=1 \ # The part of dataset to process for parallel process.
+--batch=1 \ # The batch size for the 2D-FAN face data extraction. (suggestion: 1)
+--max-res=800 # The maximum resolution for either side of the image
+
+# Then, crop all the faces from the original videos.
+python -m src.preprocess.crop_main_face \
+--root-dir="/storage/FaceForensicC23/" \ # The root folder of the dataset
+--video-dir="videos" \ # The root folder of the videos
+--fdata-dir="frame_data" \ # The folder to fetch the frame data for landmarks and bboxes
+--glob-exp="*/*" \ # The glob expression to search through the root video folder
+--crop-dir="cropped" \ # The folder to save the cropped videos
+--crop-width=150 \ # The width for the cropped videos
+--crop-height=150 \ # The height for the cropped videos
+--mean-face="./misc/20words_mean_face.npy" # The mean face for face aligned cropping.
+--replace \ # Control whether to replace existing cropped videos
+--workers=1 # Number of works to perform parallel process (default: cpu / 2 )
+```
+
+
+### Robustness Pre-processing
+This phase requires the pre-processed facial landmarks to perform facial cropping, please refer to the **Generic Pre-processing** for further detail.
+```bash
+# First, we add perturbation to all the videos.
+python -m src.preprocess.phase1_apply_all_to_videos \
+--dts-root="/storage/FaceForensicC23" \ # The root folder of the dataset
+--vid-dir="videos" \ # The root folder of the videos
+--rob-dir="robustness" \ # The folder to save the perturbed videos
+--glob-exp="*/*.mp4" \ # The glob expression to search through the root video folder
+--split=1 \ # Split the dataset into several parts for parallel process.
+--part=1 \ # The part of dataset to process for parallel process.
+--workers=1 # Number of works to perform parallel process (default: cpu / 2 )
+
+# Then, crop all the faces from the perturbed videos.
+python -m src.preprocess.phase2_face_crop_all_videos \
+(setup/run/clean) # the three phase operations
+--root-dir="/storage/FaceForensicC23/" \ # The root folder of the dataset
+--rob-dir="videos" \ # The root folder of the robustness videos
+--fd-dir="frame_data" \ # The folder to fetch the frame data for landmarks and bboxes
+--glob-exp="*/*/*/*.mp4" \ # The glob expression to search through the root video folder
+--crop-dir="cropped_robust" \ # The folder to save the cropped videos
+--mean-face="./misc/20words_mean_face.npy" \ # The mean face for face aligned cropping.
+--workers=1 # Number of works to perform parallel process (default: cpu / 2 )
+```
+
+## 🤖Training & Evaluation
+### Training - Preset Settings
+In `./scripts`, scripts are provided to start the training process for the settings mentioned in our paper.
+These settings are configured to run on a cluster with `V100*4`.
+```bash
+bash ./scripts/model/ffg_l14.sh # begin training process
+```
+### Training - Custom Settings
+Our project is built on `pytorch-lightning (2.2.0)`, please refer the [official manual](https://lightning.ai/docs/pytorch/2.2.0/common/trainer.html#trainer-class-api) and adjust the following files for advance configurations:
+```bash
+./configs/base.yaml # major training settings (e.g. epochs, optimizer, batch size, mixed-precision ...)
+./configs/data.yaml # settings for the training & validation dataset
+./configs/inference.yaml # settings for the evaluation dataset (extension of data.yaml)
+./configs/logger.yaml # settings for the WandB logger
+./configs/clip/L14/ffg.yaml # settings for the main model
+./configs/test.yaml # settings for debugging (offline logging, small batch size, short epochs ...)
+```
+The following command starts the training process with the provided settings:
+```bash
+# For debugging, add '--config configs/test.yaml' after the '--config configs/clip/L14/ffg.yaml'.
+python main.py \
+--config configs/base.yaml \
+--config configs/clip/L14/ffg.yaml
+
+# Fine-grained control is supported with the pytorch-lightning-cli.
+python main.py \
+--config configs/base.yaml \
+--config configs/clip/L14/ffg.yaml \
+--optimizer.lr=1e-5 \
+--trainer.max_epochs=10 \
+--data.init_args.train_datamodules.init_args.batch_size=5
+```
+### Evaluation - Standard
+To perform evaluation on datasets, run the following command:
+```bash
+python inference.py \
+"logs/fcg_l14/setting.yaml" \ # model settings
+"./configs/inference.yaml" \ # evaluation dataset settings
+"logs/fcg_l14/checkpoint.ckpt" \ # model checkpoint
+"--devices=4" # number of devices to compute in parallel
+```
+### Evaluation - Robustness
+We provide tools in `./scripts/tools/` to simplify the robustness evaluation task: `create-robust-config.sh` creates an evaluation config for each perturbation types and `inference-robust.sh` runs through all the datasets with the specified model.
+
+## 😎Inference - Demo Video
+To run inference on a single video with an indicator, please download our model checkpoint and execute the following commands:
+```bash
+# Pre-Processing: fetch facial landmark and bounding box
+python -m src.preprocess.fetch_landmark_bbox \
+--root-dir="./resources" \
+--video-dir="videos" \
+--fdata-dir="frame_data" \
+--glob-exp="*"
+# Pre-Processing: crop out the facial regions
+python -m src.preprocess.crop_main_face \
+--root-dir="./resources" \
+--video-dir="videos" \
+--fdata-dir="frame_data" \
+--crop-dir="cropped" \
+--glob-exp="*"
+# Main Process
+python -m demo \
+"checkpoint/setting.yaml" \ # the model setting of the checkpoint
+"checkpoint/weights.ckpt" \ # the model weights of the checkpoint
+"resources/videos/000_003.mp4" \ # the video to process
+--out_path="test.avi" \ # the output path of the processed video
+--threshold=0.5 \ # the threshold for the real/fake indicator
+--batch_size=30 # the input batch size of the model (~10G VRAM when batch_size=30 )
+```
+The following is a sample frame from the processed video:
+
+ We introduce a novel method for deepfake video detection that utilizes pixel‑wise temporal frequency spectra. Unlike previous approaches that stack 2D frame‑wise spatial frequency spectra, we extract pixel‑wise temporal frequency by performing a 1D Fourier transform on the time axis per pixel, effectively identifying temporal artifacts. We also propose an Attention Proposal Module (APM) to extract regions of interest for detecting these artifacts. Our method demonstrates outstanding generalizability and robustness in various challenging deepfake video detection scenarios.
+
+
+
+
+
+
Method & Architecture
+
+
+
+
+
+ Frequency Extraction: 1D FFT per‑pixel along time captures subtle temporal artifacts.
+
+
+
+
+
+
+
+ Architecture: Frequency Feature Extractor + APM → Joint Transformer for robust detection.
+
+
+
+
+
+
Pixel‑wise Temporal Frequency: 1D FFT on the time axis per pixel to capture subtle artifacts.
+
Attention Proposal Module (APM): Weakly‑supervised focus on temporal‑artifact regions.
+
Joint Transformer: Fuses global/part‑based frequency signals and spatio‑temporal context.
+
+
+
+
+
+
Experiments & Results
+
+
+
+
+
+ APM highlights regions (eyes, mouth) where temporal incoherence is likely.
+
+
+
+
+
+
+
+ Consistent SOTA video‑level AUC across datasets demonstrates strong generalization.
+
+
+
+
+
+
SOTA performance on FF++, CDF, DFDC, KoDF, etc.
+
Strong cross‑deepfakes and cross‑synthesis generalization.
Low‑frequency components align with raw signals, but high‑frequency components diverge after compression.
+
+
Future Work
+
We will investigate temporal‑frequency regularization to mitigate compression‑induced degradation.
+
+
+
+
+
+ Average temporal‑frequency under H.264/JPEG/WebP compression.
+
+
+
+
+
+
+
+
+
📚 Citation
+
+
+
@InProceedings{Kim_2025_ICCV,
+ author = {Kim, Taehoon and Choi, Jongwook and Jeong, Yonghyun and Noh, Haeun and Yoo, Jaejun and Baek, Seungryul and Choi, Jongwon},
+ title = {Beyond Spatial Frequency: Pixel-wise Temporal Frequency-based Deepfake Video Detection},
+ booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
+ month = {October},
+ year = {2025},
+ pages = {11198-11207}
+}
This work was partly supported by the IITP grant funded by the Korea government (MSIT): No. RS-2025-02263841: Development of a Real-time Multimodal Framework for Comprehensive Deepfake Detection Incorporating Common Sense Error AnalysisRS-2021-II211341: Artificial Intelligence Graduate School Program (Chung-Ang University) No. RS-2020-II201336: AIGS program (UNIST)
+
+
+
+
+
+
+
+
+
+
diff --git a/clean/video/pwtf_dvd/inference/config_ftcn.py b/clean/video/pwtf_dvd/inference/config_ftcn.py
new file mode 100644
index 0000000000000000000000000000000000000000..0469cec71147fafe83d496e20b08acf3cdf94853
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/config_ftcn.py
@@ -0,0 +1,138 @@
+#!/usr/bin/python
+# Borrow from tensorpack,credits goes to yuxin wu
+
+# the loaded sequnce is
+# default config in this file
+# -> provided setting file (you can not add new config after this)
+# -> manully overrided config
+# -> computed config in finalize config (you can change config after this)
+
+import os
+import pprint
+import yaml
+
+__all__ = ["config", "finalize_configs"]
+
+
+class AttrDict:
+ _freezed = False
+ """ Avoid accidental creation of new hierarchies. """
+
+ def __getattr__(self, name):
+ if self._freezed:
+ raise AttributeError(name)
+ ret = AttrDict()
+ setattr(self, name, ret)
+ return ret
+
+ def __setattr__(self, name, value):
+ if self._freezed and name not in self.__dict__:
+ raise AttributeError("Cannot create new attribute!")
+ super().__setattr__(name, value)
+
+ def __str__(self):
+ return pprint.pformat(self.to_dict(), indent=1)
+
+ __repr__ = __str__
+
+ def to_dict(self):
+ """Convert to a nested dict. """
+ return {
+ k: v.to_dict() if isinstance(v, AttrDict) else v
+ for k, v in self.__dict__.items()
+ if not k.startswith("_")
+ }
+
+ def update_args(self, args):
+ """Update from command line args. """
+ for cfg in args:
+ keys, v = cfg.split("=", maxsplit=1)
+ keylist = keys.split(".")
+ dic = self
+ # print(keylist)
+ if len(keylist) == 1:
+ assert keylist[0] in dir(dic), "Unknown config key: {}".format(
+ keylist[0]
+ )
+ for i, k in enumerate(keylist[:-1]):
+ assert k in dir(dic), "Unknown config key: {}".format(k)
+ dic = getattr(dic, k)
+ key = keylist[-1]
+ assert key in dir(dic), "Unknown config key: {}".format(key)
+ oldv = getattr(dic, key)
+ if not isinstance(oldv, str):
+ v = eval(v)
+ setattr(dic, key, v)
+
+ def update_with_yaml(self, rel_path):
+ base_path = os.path.dirname(os.path.abspath(__file__))
+ setting_path = os.path.normpath(os.path.join(base_path, "setting", rel_path))
+ setting_name = os.path.basename(setting_path).split(".")[0]
+
+ with open(setting_path, "r") as f:
+ overrided_setting = yaml.load(f,Loader=yaml.FullLoader)
+ # if 'setting_name' not in overrided_setting:
+ # raise RuntimeError('you must provide a setting name for non root_setting: {}'.format(rel_path))
+ self.update_with_dict(overrided_setting)
+ setattr(self, "setting_name", setting_name)
+
+ def init_with_yaml(self):
+ base_path = os.path.dirname(os.path.abspath(__file__))
+ setting_path = os.path.normpath(os.path.join(base_path, "root_setting.yaml"))
+ with open(setting_path, "r") as f:
+ overrided_setting = yaml.load(f,Loader=yaml.FullLoader)
+ self.update_with_dict(overrided_setting)
+
+ def update_with_text(self,text):
+ overrided_setting = yaml.load(text, Loader=yaml.FullLoader)
+ self.update_with_dict(overrided_setting)
+
+ def update_with_dict(self, dicts):
+ for k, v in dicts.items():
+ if isinstance(v, dict):
+ getattr(self, k).update_with_dict(v)
+ else:
+ setattr(self, k, v)
+
+ def freeze(self):
+ self._freezed = True
+ for v in self.__dict__.values():
+ if isinstance(v, AttrDict):
+ v.freeze()
+
+ # avoid silent bugs
+ def __eq__(self, _):
+ raise NotImplementedError()
+
+ def __ne__(self, _):
+ raise NotImplementedError()
+
+
+config = AttrDict()
+_C = config # short alias to avoid coding
+
+
+# you can directly write setting here as _C.model_dir='.\checkpoint' or in root_setting.yaml
+
+
+#
+
+
+def finalize_configs(input_cfg=_C, freeze=True, verbose=True):
+
+ # _C.base_path = os.path.dirname(os.path.abspath(__file__))
+ input_cfg.base_path = os.path.dirname(__file__)
+
+ # for running in remote server
+ # for k, v in input_cfg.path.__dict__.items():
+ # v = os.path.normpath(os.path.join(input_cfg.base_path, v))
+ # setattr(input_cfg.path, k, v)
+ if freeze:
+ input_cfg.freeze()
+ # if verbose:
+ # logger.info("Config: ------------------------------------------\n" + str(_C))
+
+
+if __name__ == "__main__":
+ print("?")
+ print(os.path.dirname(__file__))
diff --git a/clean/video/pwtf_dvd/inference/model/attention_network.py b/clean/video/pwtf_dvd/inference/model/attention_network.py
new file mode 100644
index 0000000000000000000000000000000000000000..13280c3acbbc7f134274d914d94315d3e12d88d4
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/model/attention_network.py
@@ -0,0 +1,193 @@
+import torch.nn as nn
+import math
+import torch
+import torch.utils.model_zoo as model_zoo
+import torch.nn.functional as F
+from torch.autograd import Variable
+from model.resnet_backbone import resnet50, oneByoneConvNet
+import numpy as np
+
+
+class CONV(nn.Module):
+ def __init__(self, partials_num=5, ori_dim=64):
+ super(CONV, self).__init__()
+ self.partials_num = partials_num
+ self.ori_dim = ori_dim
+ self.conv0 = nn.Conv2d(ori_dim, ori_dim, kernel_size=1)
+ # Create conv layers dynamically
+ self.convs = nn.ModuleList()
+ for i in range(partials_num):
+ self.convs.append(nn.Conv2d(ori_dim, ori_dim, kernel_size=1))
+
+ def forward(self, x_p, x):
+ batch_size = x.size(0)
+ h, w = x.size(2), x.size(3)
+
+ x_sum = self.conv0(x)
+
+ if self.partials_num > 0 and x_p is not None:
+ x_p = F.interpolate(x_p, size=(h, w), mode='bilinear', align_corners=True)
+ x_p = x_p.view(batch_size, self.partials_num, self.ori_dim, h, w)
+ x_p_p = x_p.permute(1, 0, 2, 3, 4)
+
+ for i in range(min(self.partials_num, len(self.convs))):
+ x_sum += self.convs[i](x_p_p[i])
+
+ return x_sum
+class AttentionCropFunction(torch.autograd.Function):
+ @staticmethod
+ def forward(self, images, locs):
+ h = lambda x: 1. / (1. + torch.exp(-10. * x))
+ in_size = images.size()[2]
+ unit = torch.stack([torch.arange(0, in_size)] * in_size).float()
+ x = torch.stack([unit.t()] * 16)
+ y = torch.stack([unit] * 16)
+ if isinstance(images, torch.cuda.FloatTensor):
+ x, y = x.cuda(), y.cuda()
+ in_size = images.size()[2]
+ ret = []
+ for i in range(images.size(0)):
+ tx, ty, tl = locs[i][0]*224, locs[i][1]*224, 44
+ tx = tx if tx > tl else tl
+ tx = tx if tx < in_size-tl else in_size-tl
+ ty = ty if ty > tl else tl
+ ty = ty if ty < in_size-tl else in_size-tl
+
+ w_off = int(tx-tl) if (tx-tl) > 0 else 0
+ h_off = int(ty-tl) if (ty-tl) > 0 else 0
+ w_end = int(tx+tl) if (tx+tl) < in_size else in_size
+ h_end = int(ty+tl) if (ty+tl) < in_size else in_size
+
+ mk = (h(x-w_off) - h(x-w_end)) * (h(y-h_off) - h(y-h_end))
+ xatt = images[i] * mk
+ xatt_cropped = xatt[:, w_off:w_end, h_off:h_end]
+ before_upsample = Variable(xatt_cropped.unsqueeze(0))
+ xamp = F.interpolate(before_upsample, size=(88,88), mode='bilinear', align_corners = True)
+ ret.append(xamp.data.squeeze())
+ ret_tensor = torch.stack(ret)
+ self.save_for_backward(images, ret_tensor)
+ return ret_tensor
+
+ @staticmethod
+ def backward(self, grad_output):
+ pass
+class AttentionCropLayer(nn.Module):
+ """
+ Crop function sholud be implemented with the nn.Function.
+ Detailed description is in 'Attention localization and amplification' part.
+ Forward function will not changed. backward function will not opearate with autograd, but munually implemented function
+ """
+ def forward(self, images, locs):
+ return AttentionCropFunction.apply(images, locs)
+class APN(nn.Module):
+ def __init__(self,depth=18, partials_num = 5):
+ super(APN, self).__init__()
+ self.localization = LocalizationCNN(depth)
+ self.partials_num = partials_num
+ self.region_proposal = nn.Sequential(
+ nn.Linear(2048 * 7 * 7, 1024),
+ nn.Tanh(),
+ nn.Linear(1024, 2*partials_num),
+ nn.Sigmoid(),
+ )
+ def forward(self, x, x_ft):
+ x_batch_size = x.size(0)
+ x_l = self.localization(x,x_ft)
+ x_l = x_l.view(-1, 2048*7*7)
+ xs = self.region_proposal(x_l)
+ xs_t = xs.view(x_batch_size * self.partials_num, 2)
+ return xs_t
+
+class LocalizationCNN(nn.Module):
+ def __init__(self,depth = 18):
+ super(LocalizationCNN, self).__init__()
+ self.depth = depth
+ self.conv1 = nn.Sequential( nn.Conv2d(16, 64, kernel_size=7, stride=2, padding=3, bias=False),
+ nn.BatchNorm2d(64),
+ nn.ReLU(inplace=True),
+ nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
+ )
+ self.conv2 = nn.Sequential( nn.Conv2d(64, 256, kernel_size=3, stride=1, padding=1, bias=False),
+ nn.BatchNorm2d(256),
+ nn.ReLU(inplace=True)
+ )
+ self.conv3 = nn.Sequential( nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1, bias=False),
+ nn.BatchNorm2d(512),
+ nn.ReLU(inplace=True),
+ nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
+ self.conv4 = nn.Sequential( nn.Conv2d(512, 1024, kernel_size=3, stride=1, padding=1, bias=False),
+ nn.BatchNorm2d(1024),
+ nn.ReLU(inplace=True),nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
+ self.conv5 = nn.Sequential( nn.Conv2d(1024, 2048, kernel_size=3, stride=1, padding=1, bias=False),
+ nn.BatchNorm2d(2048),
+ nn.ReLU(inplace=True),nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
+ def forward(self, x,x_ft):
+ x = self.conv1(x) + x_ft[0].detach().float()
+ x = self.conv2(x) + x_ft[1].detach().float()
+ x = self.conv3(x) + x_ft[2].detach().float()
+ x = self.conv4(x) + x_ft[3].detach().float()
+ x = self.conv5(x) + x_ft[4].detach().float()
+ return x
+h = lambda x: 1. / (1. + torch.exp(-100. * x))
+class APNResNet(nn.Module):
+ def __init__(self, partials_num = 5, depth= 18):
+ super(APNResNet, self).__init__()
+ self.partials_num = partials_num
+
+ self.whole_resnet = resnet50(pretrained=True)
+ self.oconv0 = oneByoneConvNet(64,32,64)
+ self.oconv1 = oneByoneConvNet(256,128,256)
+ self.oconv2 = oneByoneConvNet(512,256,512)
+ self.oconv3 = oneByoneConvNet(1024,512,1024)
+ self.oconv4 = oneByoneConvNet(2048,1024,2048)
+
+ self.attn_1 = CONV(partials_num=partials_num, ori_dim=64)
+ self.attn_2 = CONV(partials_num=partials_num, ori_dim=256)
+ self.attn_3 = CONV(partials_num=partials_num, ori_dim=512)
+ self.attn_4 = CONV(partials_num=partials_num, ori_dim=1024)
+ self.fn = nn.Linear(2048, 1024)
+ self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+
+ self.APN = APN(depth, partials_num)
+ self.crop_resize = AttentionCropLayer()
+
+ def forward(self, x):
+ batch_size = x.size(0)
+ out1_w, out2_w, out3_w, out4_w, out5_w, out5_p_w = self.whole_resnet(x)
+
+ out5_p_w = out5_p_w.view(batch_size, -1)
+
+
+ xs_t= self.APN(x.clone(), (out1_w, out2_w, out3_w,out4_w,out5_w))
+ images = x
+ images = images.repeat_interleave(self.partials_num, dim=0)
+ scaled_x = self.crop_resize(images, xs_t)
+ scaled_x_b = scaled_x.view(-1, 16, 88, 88)
+ out1_p_b, out2_p_b, out3_p_b, out4_p_b, out5_p_b, out5_p_p_b = self.whole_resnet(scaled_x_b)
+ whole_box_num = 5
+ out1 = self.attention(out1_w, out1_p_b, 1, whole_box_num)
+ out2 = self.attention(out2_w, out2_p_b, 2, whole_box_num)
+ out3 = self.attention(out3_w, out3_p_b, 3, whole_box_num)
+ out4 = self.attention(out4_w, out4_p_b, 4, whole_box_num)
+ out5 = self.avgpool(self.oconv4(out5_p_b)).view(batch_size, self.partials_num, -1)
+
+ out1 = self.oconv0(out1)
+ out2 = self.oconv1(out2)
+ out3 = self.oconv2(out3)
+ out4 = self.oconv3(out4)
+ out6 = self.fn(out5_p_w)
+
+
+
+ return out1, out2, out3, out4, out5, out6, (xs_t, scaled_x)
+
+ def attention(self, x, x_p, stage, box_num):
+ if stage == 1:
+ out = self.attn_1(x_p, x)
+ if stage == 2:
+ out = self.attn_2(x_p, x)
+ if stage == 3:
+ out = self.attn_3(x_p, x)
+ if stage == 4:
+ out = self.attn_4(x_p, x)
+ return out
diff --git a/clean/video/pwtf_dvd/inference/model/framework.py b/clean/video/pwtf_dvd/inference/model/framework.py
new file mode 100644
index 0000000000000000000000000000000000000000..9839957b2d74c38a43074539388861f453c4ab36
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/model/framework.py
@@ -0,0 +1,38 @@
+
+import torch
+import torch.nn as nn
+from .video_encoder import I3D8x8
+from .attention_network import APNResNet
+from .transformers import TransformerHead, SpatialTransformerE
+
+
+# Model for deployment
+# Note: modified model structure for easier deployment combined with video encoder FTCN official repo code
+def get_model():
+ part_num = 5
+ model= I3D8x8()
+ model_ft = APNResNet(partials_num=part_num,depth =50)
+ params = dict(spatial_size=14, time_size=16, in_channels=1024,num_parts=part_num)
+ TTE = TransformerHead(**params)
+ STE = SpatialTransformerE(**params)
+ MLP = torch.nn.Linear(2048,1)
+ return Framework(model, model_ft, TTE, STE, MLP)
+
+
+class Framework(nn.Module):
+ def __init__(self, model, model_ft, TTE, STE, MLP):
+ super(Framework, self).__init__()
+ self.model = model
+ self.model_ft = model_ft
+ self.TTE = TTE
+ self.STE = STE
+ self.MLP = MLP
+
+ def forward(self, video_sample, ft_sample):
+ out1, out2, out3, out4, out5,out6, (xs, scaled_x) = self.model_ft(ft_sample.float())
+ ft_feats = [out1, out2, out3, out4, None]
+ x,_ = self.model(video_sample,ft_feats)
+ ft_s , _= self.STE(x,out5,xs)
+ ft_t, _ = self.TTE(x,out6)
+ outputs = self.MLP(torch.concat((ft_t, ft_s), dim = 1))
+ return outputs
diff --git a/clean/video/pwtf_dvd/inference/model/resnet_backbone.py b/clean/video/pwtf_dvd/inference/model/resnet_backbone.py
new file mode 100644
index 0000000000000000000000000000000000000000..5cf21e4bf7d02d292fcdd60a0929a5bd0f5772d8
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/model/resnet_backbone.py
@@ -0,0 +1,235 @@
+import torch.nn as nn
+import math
+import torch
+import torch.utils.model_zoo as model_zoo
+import torch.nn.functional as F
+__all__ = ['ResNet', 'resnet18', 'resnet50', ]
+
+
+normalization = nn.BatchNorm2d
+
+model_urls = {
+ 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+ 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+ 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+ 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+ 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+}
+def conv3x3(in_planes, out_planes, stride=1):
+ """3x3 convolution with padding"""
+ return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+ padding=1, bias=False)
+class BasicBlock(nn.Module):
+ expansion = 1
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None):
+ super(BasicBlock, self).__init__()
+ self.conv1 = conv3x3(inplanes, planes, stride)
+ self.bn1 = normalization(planes)
+ self.relu = nn.ReLU(inplace=False)
+ self.conv2 = conv3x3(planes, planes)
+ self.bn2 = normalization(planes)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x):
+ residual = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+
+ if self.downsample is not None:
+ residual = self.downsample(x)
+
+ out = out + residual
+ out = self.relu(out)
+
+ return out
+
+ def forward_masked(self, x, mask=None):
+ residual = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+
+ if self.downsample is not None:
+ residual = self.downsample(x)
+
+
+ out = out + residual
+ if mask is not None:
+ out = out * mask[None,:,None,None]# + self.bn2.bias[None,:,None,None] * (1 - mask[None,:,None,None])
+
+ out = self.relu(out)
+
+ return out
+
+class Bottleneck(nn.Module):
+ expansion = 4
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None):
+ super(Bottleneck, self).__init__()
+ self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+ self.bn1 = normalization(planes)
+ self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+ padding=1, bias=False)
+ self.bn2 = normalization(planes)
+ self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+ self.bn3 = normalization(planes * 4)
+ self.relu = nn.ReLU(inplace=False)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x):
+ residual = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+ out = self.relu(out)
+
+ out = self.conv3(out)
+ out = self.bn3(out)
+
+ if self.downsample is not None:
+ residual = self.downsample(x)
+
+ out += residual
+ out = self.relu(out)
+
+ return out
+
+
+class AbstractResNet(nn.Module):
+
+ def __init__(self, block, layers, num_classes=1000, maxiter=-1):
+ super(AbstractResNet, self).__init__()
+ self.inplanes = 64
+ self.conv1 = nn.Conv2d(16, 64, kernel_size=7, stride=2, padding=3,
+ bias=False)
+ self.bn1 = normalization(64)
+ self.relu = nn.ReLU(inplace=False)
+ self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+ self.layer1 = self._make_layer(block, 64, layers[0])
+ self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+ self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+ self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+ self.avgpool = nn.AdaptiveAvgPool2d(1)
+
+
+ def _make_layer(self, block, planes, blocks, stride=1):
+ downsample = None
+ if stride != 1 or self.inplanes != planes * block.expansion:
+ downsample = nn.Sequential(
+ nn.Conv2d(self.inplanes, planes * block.expansion,
+ kernel_size=1, stride=stride, bias=False),
+ normalization(planes * block.expansion),
+ )
+
+ layers = []
+ layers.append(block(self.inplanes, planes, stride, downsample))
+ self.inplanes = planes * block.expansion
+ for i in range(1, blocks):
+ layers.append(block(self.inplanes, planes))
+
+ return nn.Sequential(*layers)
+
+ def features(self,x):
+ x__ = self.maxpool(self.relu(self.bn1(self.conv1(x))))
+ x__1 = self.layer1(x__)
+ x__2 = self.layer2(x__1)
+ x__3 = self.layer3(x__2)
+ x__4 = self.layer4(x__3)
+ return x__, x__1, x__2, x__3, x__4
+
+
+
+ def load_state_dict(self, state_dict, strict=True):
+ missing_keys = []
+ unexpected_keys = []
+ error_msgs = []
+
+ # copy state_dict so _load_from_state_dict can modify it
+ metadata = getattr(state_dict, '_metadata', None)
+ state_dict = state_dict.copy()
+ if metadata is not None:
+ state_dict._metadata = metadata
+
+ def load(module, prefix=''):
+ local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
+ module._load_from_state_dict(
+ state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+ for name, child in module._modules.items():
+ if child is not None:
+ load(child, prefix + name + '.')
+
+ load(self)
+
+ if strict:
+ error_msg = ''
+ if len(unexpected_keys) > 0:
+ error_msgs.insert(
+ 0, 'Unexpected key(s) in state_dict: {}. '.format(
+ ', '.join('"{}"'.format(k) for k in unexpected_keys)))
+ if len(missing_keys) > 0:
+ error_msgs.insert(
+ 0, 'Missing key(s) in state_dict: {}. '.format(
+ ', '.join('"{}"'.format(k) for k in missing_keys)))
+
+ if len(error_msgs) > 0:
+ pass
+
+
+
+# Removed training-related classes for inference-only code
+class oneByoneConvNet(nn.Module):
+ def __init__(self, in_channels,hidden_channels,out_channels):
+ super(oneByoneConvNet, self).__init__()
+ self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=hidden_channels, kernel_size=1)
+
+ self.conv2 = nn.Conv2d(in_channels=hidden_channels, out_channels=out_channels, kernel_size=1)
+
+ def forward(self, x):
+ x = self.conv1(x)
+ x = F.relu(x)
+ x = self.conv2(x)
+ return x
+class ResNet(AbstractResNet):
+
+ def __init__(self, block, layers, num_classes=2, max_iter=-1, normalized=True, use_bias=False, simsiam=False):
+ super(ResNet, self).__init__(block, layers, num_classes)
+ def forward(self, x):
+ x__, x__1, x__2, x__3, x__4 = self.features(x)
+ x__4_p = self.avgpool(x__4).view(x__.size(0), -1).squeeze()
+ return x__, x__1, x__2, x__3,x__4, x__4_p
+
+def resnet50(pretrained=False, **kwargs):
+ model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
+ if pretrained:
+ weight = model_zoo.load_url(model_urls['resnet50'])
+ conv1_weight = weight['conv1.weight']
+ conv1_weight = conv1_weight.repeat(1, 5, 1, 1)
+ conv1_weight = torch.concat([conv1_weight, conv1_weight[:,:1,:,:]], dim=1)
+ model.load_state_dict(weight)
+ model.conv1.weight = nn.Parameter(conv1_weight)
+ return model
+
+class Normalize(nn.Module):
+ def __init__(self):
+ super(Normalize, self).__init__()
+
+ def forward(self, x):
+ return F.normalize(x, dim=-1)
+
+# Removed get_model function as it's not needed for inference
\ No newline at end of file
diff --git a/clean/video/pwtf_dvd/inference/model/transformers.py b/clean/video/pwtf_dvd/inference/model/transformers.py
new file mode 100644
index 0000000000000000000000000000000000000000..1cd59044dcc707b980e5cfc4bb78b54c435c794a
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/model/transformers.py
@@ -0,0 +1,474 @@
+
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+from einops import rearrange, repeat
+# Removed Rearrange import - not used in inference
+# Removed unused imports for inference
+from math import sqrt
+from config_ftcn import config as my_cfg
+
+# Initialize config
+my_cfg.init_with_yaml()
+my_cfg.update_with_yaml("ftcn_tt.yaml")
+my_cfg.freeze()
+
+# =============================================================================
+# Basic Transformer Components
+# =============================================================================
+
+class Residual(nn.Module):
+ def __init__(self, fn):
+ super().__init__()
+ self.fn = fn
+ def forward(self, x, **kwargs):
+ return self.fn(x, **kwargs) + x
+
+class PreNorm(nn.Module):
+ def __init__(self, dim, fn):
+ super().__init__()
+ self.norm = nn.LayerNorm(dim)
+ self.fn = fn
+ def forward(self, x, **kwargs):
+ return self.fn(self.norm(x), **kwargs)
+
+class FeedForward(nn.Module):
+ def __init__(self, dim, hidden_dim, dropout = 0.):
+ super().__init__()
+ self.net = nn.Sequential(
+ nn.Linear(dim, hidden_dim),
+ nn.GELU(),
+ nn.Dropout(dropout),
+ nn.Linear(hidden_dim, dim),
+ nn.Dropout(dropout)
+ )
+ def forward(self, x):
+ return self.net(x)
+
+class Attention(nn.Module):
+ def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
+ super().__init__()
+ inner_dim = dim_head * heads
+ project_out = not (heads == 1 and dim_head == dim)
+
+ self.heads = heads
+ self.scale = dim_head ** -0.5
+
+ self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+
+ self.to_out = nn.Sequential(
+ nn.Linear(inner_dim, dim),
+ nn.Dropout(dropout)
+ ) if project_out else nn.Identity()
+
+ def forward(self, x, mask = None):
+ b, n, _, h = *x.shape, self.heads
+ qkv = self.to_qkv(x).chunk(3, dim = -1)
+ q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)
+
+ dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
+ mask_value = -torch.finfo(dots.dtype).max
+
+ if mask is not None:
+ mask = F.pad(mask.flatten(1), (1, 0), value = True)
+ assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions'
+ mask = rearrange(mask, 'b i -> b () i ()') * rearrange(mask, 'b j -> b () () j')
+ dots.masked_fill_(~mask, mask_value)
+ del mask
+
+ attn = dots.softmax(dim=-1)
+
+ out = einsum('b h i j, b h j d -> b h i d', attn, v)
+ out = rearrange(out, 'b h n d -> b n (h d)')
+ out = self.to_out(out)
+ return out
+
+class Transformer(nn.Module):
+ def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
+ super().__init__()
+ self.layers = nn.ModuleList([])
+ for _ in range(depth):
+ self.layers.append(nn.ModuleList([
+ Residual(PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout))),
+ Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout)))
+ ]))
+ def forward(self, x, mask = None):
+ for attn, ff in self.layers:
+ x = attn(x, mask = mask)
+ x = ff(x)
+ return x
+
+# =============================================================================
+# Utility Functions
+# =============================================================================
+
+def valid_idx(idx, h):
+ i = idx // h
+ j = idx % h
+ if j == 0 or i == h - 1 or j == h - 1:
+ return False
+ else:
+ return True
+
+
+# =============================================================================
+# Patch Pooling Classes (Simplified for Inference)
+# =============================================================================
+
+class CenterPatchPool(nn.Module):
+ """Simplified patch pooling for inference - always use center patch"""
+ def __init__(self):
+ super().__init__()
+
+ def forward(self, x):
+ # batch,channel,16,7x7
+ b, c, t, h, w = x.shape
+ x = x.reshape(b, c, t, h * w)
+ idx = h * w // 2 # Always use center patch for inference
+ x = x[..., idx]
+ return x
+
+class CenterAvgPool(nn.Module):
+ """Simplified average pooling for inference - use all valid patches"""
+ def __init__(self):
+ super().__init__()
+
+ def forward(self, x):
+ # batch,channel,16,7x7
+ b, c, t, h, w = x.shape
+ x = x.reshape(b, c, t, h * w)
+ candidates = list(range(h * w))
+ candidates = [idx for idx in candidates if valid_idx(idx, h)]
+ x = x[..., candidates].mean(-1)
+ return x
+
+class CenterSelect(nn.Module):
+ """Simplified selection for inference - use all valid patches"""
+ def __init__(self):
+ super().__init__()
+
+ def forward(self, x):
+ # batch,7x7
+ size = x.shape[1]
+ h = int(sqrt(size))
+ candidates = list(range(size))
+ candidates = [idx for idx in candidates if valid_idx(idx, h)]
+ x = x[:, candidates]
+ return x
+
+# =============================================================================
+# Vision Transformer Classes
+# =============================================================================
+
+# Removed ViT and VideoiT classes - not used in inference
+
+# =============================================================================
+# Specialized Transformer Classes
+# =============================================================================
+
+# SpatialTransformer class - used by SpatialTransformerE
+class SpatialTransformer(nn.Module):
+ def __init__(self,num_patches, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', dim_head = 64, dropout = 0., emb_dropout = 0.):
+ super().__init__()
+ self.dim =dim
+ self.num_patches=num_patches
+ self.pos_embedding = nn.Parameter(posemb_sincos_2d(14,14,dim))
+ self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+ self.dropout = nn.Dropout(emb_dropout)
+ self.fre_pos = nn.Parameter(torch.randn(1, 1, dim))
+ self.transformer = Transformer(1024, depth, heads, dim_head, mlp_dim, dropout)
+ self.pool = pool
+ self.to_latent = nn.Identity()
+ self.mlp_head = nn.Sequential(
+ nn.LayerNorm(1024), nn.Linear(1024, 1024)
+ )
+ self.freq_embedding = nn.Sequential(nn.LayerNorm(dim),nn.Linear(dim,dim))
+ self.cls_token_pos = nn.Parameter(torch.randn(1, 1, dim))
+ self.mlp_head_ste = nn.Linear(dim, 1)
+
+ def posemb_loc(self, locs,batch_size):
+ if locs is not None:
+ b,n,_ = locs.shape
+
+ pos224 =self.pos_embedding.view(14,14,-1)
+ pos_locs = torch.zeros(batch_size,n, self.dim)
+ if locs is not None:
+ # locs[locs>=1] = 0.999999
+ loc_x = torch.floor(locs[:,:,0]*6)
+ loc_y = torch.floor(locs[:,:,1]*6)
+
+ loc_x = loc_x.long()
+ loc_y = loc_y.long()
+ x_a = locs[:,:,0]*6 - loc_x
+ y_a = locs[:,:,0]*6 - loc_y
+
+ x1_a = torch.sqrt( torch.pow(1-x_a,2) + torch.pow(1-y_a,2))
+ x2_a = torch.sqrt( torch.pow(1-x_a,2) + torch.pow(y_a,2))
+ x3_a = torch.sqrt( torch.pow(x_a,2) + torch.pow(1-y_a,2))
+ x4_a = torch.sqrt( torch.pow(x_a,2) + torch.pow(y_a,2))
+ pos_locs[:, :, :] = pos224[loc_x,loc_y,:]*x1_a.unsqueeze(-1)\
+ +pos224[loc_x+1,loc_y,:]*x2_a.unsqueeze(-1)\
+ +pos224[loc_x,loc_y+1,:]*x3_a.unsqueeze(-1)\
+ +pos224[loc_x+1,loc_y+1,:]*x4_a.unsqueeze(-1)
+
+ return pos_locs
+
+ def forward(self, x,ft_feature,locs):
+ b, n, _ = x.shape #batch,num_patches,channels #
+ cls_tokens_pos = repeat(self.cls_token_pos, '() n d -> b n d', b = b)
+ cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b)
+ cls_tokens = cls_tokens+cls_tokens_pos
+ x += self.pos_embedding
+ if locs is not None:
+ posemb_loc = self.posemb_loc(locs,b).clone().detach()
+ if torch.cuda.is_available() and x.is_cuda:
+ posemb_loc = posemb_loc.cuda()
+ posemb_loc += self.fre_pos
+
+ ft_feature = ft_feature+posemb_loc
+ ft_feature = self.freq_embedding(ft_feature)
+ x = torch.cat((cls_tokens, x,ft_feature), dim=1)
+ else :
+ x = torch.cat((cls_tokens, x), dim=1)
+ x = self.dropout(x)
+ x = self.transformer(x, mask=None)
+ x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
+
+ x = self.to_latent(x)
+ x = self.mlp_head(x)
+ return x, self.mlp_head_ste(x.clone())
+import torch
+from torch import nn
+import torch.nn.functional as F
+from einops import repeat
+
+def posemb_sincos_2d(h, w, dim, temperature: int = 10000, dtype=torch.float32):
+ assert (dim % 4) == 0, "dim must be multiple of 4"
+ y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
+ omega = torch.arange(dim // 4) / (dim // 4 - 1)
+ omega = 1.0 / (temperature ** omega)
+ y = y.flatten()[:, None] * omega[None, :]
+ x = x.flatten()[:, None] * omega[None, :]
+ pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim=1)
+ return pe.type(dtype)
+
+class SpatialTransformer(nn.Module):
+ def __init__(self, num_patches, num_classes, dim, depth, heads, mlp_dim,
+ pool='cls', dim_head=64, dropout=0., emb_dropout=0.):
+ super().__init__()
+ self.dim = dim
+ self.num_patches = num_patches
+ self.register_buffer('pos_embedding', posemb_sincos_2d(14, 14, dim))
+
+ self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+ self.cls_token_pos = nn.Parameter(torch.randn(1, 1, dim))
+ self.fre_pos = nn.Parameter(torch.randn(1, 1, dim))
+ self.dropout = nn.Dropout(emb_dropout)
+
+ self.spatial_proj = nn.Linear(dim, dim)
+ self.freq_embedding = nn.Sequential(nn.LayerNorm(dim), nn.Linear(dim, dim))
+ self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
+
+ self.pool = pool
+ self.to_latent = nn.Identity()
+
+ # output ES ∈ R^1024
+ self.mlp_head = nn.Sequential(nn.LayerNorm(dim), nn.Linear(dim, 1024))
+ self.mlp_head_ste = nn.Linear(dim, 1)
+
+ @torch.no_grad()
+ def posemb_loc(self, locs, batch_size):
+ if locs is None:
+ return None
+ # locs: (B, P, 2) in [0,1]
+ B, P, _ = locs.shape
+ device, dtype = locs.device, locs.dtype
+
+ pos_grid = self.pos_embedding.view(14, 14, self.dim).to(device) # (14,14,dim) - already [196, dim]
+
+ S = 13.0
+ x_f = (locs[:, :, 0] * S).clamp(0, 13 - 1e-6)
+ y_f = (locs[:, :, 1] * S).clamp(0, 13 - 1e-6)
+
+ ix = torch.floor(x_f).long().clamp(0, 12)
+ iy = torch.floor(y_f).long().clamp(0, 12)
+ ix1 = (ix + 1).clamp(max=13)
+ iy1 = (iy + 1).clamp(max=13)
+
+ x_a = (x_f - ix).to(dtype) # frac
+ y_a = (y_f - iy).to(dtype)
+
+ w11 = (1 - x_a) * (1 - y_a)
+ w12 = x_a * (1 - y_a)
+ w21 = (1 - x_a) * y_a
+ w22 = x_a * y_a
+
+ p11 = pos_grid[ix, iy, :] # (B,P,dim) via advanced indexing
+ p12 = pos_grid[ix1, iy, :]
+ p21 = pos_grid[ix, iy1, :]
+ p22 = pos_grid[ix1, iy1, :]
+
+ pos_locs = (p11 * w11.unsqueeze(-1) +
+ p12 * w12.unsqueeze(-1) +
+ p21 * w21.unsqueeze(-1) +
+ p22 * w22.unsqueeze(-1)) # (B,P,dim)
+ return pos_locs
+
+ def forward(self, x, ft_feature, locs):
+ # x: (B, N=14*14, dim), ft_feature: (B, P, dim), locs: (B, P, 2)
+ B, N, _ = x.shape
+
+ # cls token (+ pos)
+ cls_tokens = repeat(self.cls_token, '() n d -> b n d', b=B)
+ cls_tokens_pos = repeat(self.cls_token_pos, '() n d -> b n d', b=B)
+ cls_tokens = cls_tokens + cls_tokens_pos
+
+ # Wsp z_sp + pos_sp
+ x = self.spatial_proj(x)
+ x = x + self.pos_embedding.unsqueeze(0).to(x.device)
+
+ # part tokens: Wfreq Z_p + (interp(pos_sp) + pos_freq)
+ if locs is not None and ft_feature is not None:
+ pospart = self.posemb_loc(locs, B) # (B,P,dim)
+ if pospart is not None:
+ pospart = pospart + self.fre_pos # + posfreq
+ ft_feature = self.freq_embedding(ft_feature) # Wfreq
+ if pospart is not None:
+ ft_feature = ft_feature + pospart
+ tokens = torch.cat((cls_tokens, x, ft_feature), dim=1)
+ else:
+ tokens = torch.cat((cls_tokens, x), dim=1)
+
+ tokens = self.dropout(tokens)
+ tokens = self.transformer(tokens, mask=None)
+ out = tokens.mean(dim=1) if self.pool == 'mean' else tokens[:, 0]
+
+ out = self.to_latent(out)
+ es = self.mlp_head(out) # ES ∈ R^1024
+ ste_score = self.mlp_head_ste(out) # optional head
+
+ return es, ste_score
+
+class TimeTransformer(nn.Module):
+ def __init__(self,num_patches, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', dim_head = 64, dropout = 0., emb_dropout = 0.):
+ super().__init__()
+ assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
+
+ self.num_patches=num_patches
+ self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
+ self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+ self.dropout = nn.Dropout(emb_dropout)
+
+ self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
+ self.freq_embedding = nn.Sequential(nn.LayerNorm(dim),nn.Linear(dim,dim))
+ self.pool = pool
+ self.to_latent = nn.Identity()
+ self.LN = nn.LayerNorm(dim)
+
+ def forward(self, x,ft_t):
+ b, n, _ = x.shape #batch,num_patches,channels #
+ cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b)
+
+ ft_t = self.freq_embedding(ft_t)
+ x +=ft_t.unsqueeze(1)
+ x = torch.cat((cls_tokens, x), dim=1)
+ x += self.pos_embedding[:, :(n + 1)]
+
+ # x = torch.cat((x,ft_t), dim=1)
+ x = self.dropout(x)
+ x = self.transformer(x, mask=None)
+ x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
+ x = self.to_latent(x)
+ x = self.LN(x)
+
+ return x, None
+
+# =============================================================================
+# FTCN-specific Transformer Classes
+# =============================================================================
+
+class SpatialTransformerE(nn.Module):
+ def __init__(self, spatial_size=14, time_size = 16, in_channels = 1024, num_parts=5):
+ super().__init__()
+ self.num_parts = num_parts
+ self.in_channels = in_channels
+ default_params= dict(
+ dim=self.in_channels, depth=1, heads=16, mlp_dim=2048, dropout=0.1, emb_dropout=0.1,
+ num_patches = spatial_size ** 2, num_classes = 1
+ )
+ self.num_patches = spatial_size ** 2
+ self.freq_embedding = nn.Linear(2048,1024)
+ self.freq_embedding.weight.data.normal_(mean=0.0, std=0.02)
+ self.freq_embedding.bias.data.zero_()
+ self.pool = nn.AvgPool3d((time_size, 1, 1))
+ self.spatial_T = SpatialTransformer( **default_params )
+
+ def forward(self, x, ft, locs):
+ batch_size = x.shape[0]
+ x = self.pool(x)
+ if self.num_parts > 0:
+ ft = self.freq_embedding(ft.reshape(batch_size*self.num_parts,2048))
+ ft = ft.view(batch_size,self.num_parts,1024)
+ locs = locs.reshape(-1, self.num_parts, 2)
+ x = x.view(batch_size,self.num_patches,1024)
+ x = self.spatial_T(x,ft,locs)
+ return x
+
+class TransformerHead(nn.Module):
+ def __init__(self, spatial_size=7, time_size=16, in_channels=1024, num_parts=5):
+ super().__init__()
+ if my_cfg.model.inco.no_time_pool:
+ time_size = time_size * 2
+ patch_type = my_cfg.model.transformer.patch_type # time
+ if patch_type == "time":
+ self.pool = nn.AvgPool3d((1, spatial_size, spatial_size))
+ self.num_patches = time_size
+ elif patch_type == "spatial":
+ self.pool = nn.AvgPool3d((time_size, 1, 1))
+ self.num_patches = spatial_size ** 2
+ elif patch_type == "random":
+ self.pool = CenterPatchPool()
+ self.num_patches = time_size
+ elif patch_type == "random_avg":
+ self.pool = CenterAvgPool()
+ self.num_patches = time_size
+ elif patch_type == "all":
+ self.pool = nn.Identity()
+ self.num_patches = time_size * spatial_size * spatial_size
+ else:
+ raise NotImplementedError(patch_type)
+
+ self.dim = my_cfg.model.transformer.dim # False
+ if self.dim == -1:
+ self.dim = in_channels # 2048
+ my_cfg.model.transformer.dim = self.dim
+
+ self.in_channels = in_channels
+
+ if self.dim != self.in_channels:
+ self.fc = nn.Linear(self.in_channels, self.dim)
+
+ default_params = dict(
+ dim=self.dim, depth=6, heads=16, mlp_dim=2048, dropout=0.1, emb_dropout=0.1,
+ )
+ params = my_cfg.model.transformer.to_dict()
+ for key in default_params:
+ if key in params:
+ default_params[key] = params[key]
+
+ self.time_T = TimeTransformer(
+ num_patches=self.num_patches, num_classes=1, **default_params
+ )
+
+ self.sigmoid = nn.Sigmoid()
+
+ def forward(self, x,ft_t):
+ x = self.pool(x)
+ x = x.reshape(-1, self.in_channels, self.num_patches)
+ x = x.permute(0, 2, 1)
+ if self.dim != self.in_channels:
+ x = self.fc(x.reshape(-1, self.in_channels))
+ x = x.reshape(-1, self.num_patches, self.dim)
+
+ x = self.time_T(x,ft_t)
+ return x
\ No newline at end of file
diff --git a/clean/video/pwtf_dvd/inference/model/video_encoder.py b/clean/video/pwtf_dvd/inference/model/video_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..e512f03f456c6c0351ff82173adaa1a4e1f83fa8
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/model/video_encoder.py
@@ -0,0 +1,227 @@
+import gc
+import math
+
+# Simplified config for inference
+config_text = """
+DATA:
+ NUM_FRAMES: 8
+ SAMPLING_RATE: 8
+ TEST_CROP_SIZE: 256
+ INPUT_CHANNEL_NUM: [3]
+RESNET:
+ ZERO_INIT_FINAL_BN: True
+ WIDTH_PER_GROUP: 64
+ NUM_GROUPS: 1
+ DEPTH: 50
+ TRANS_FUNC: bottleneck_transform
+ STRIDE_1X1: False
+ NUM_BLOCK_TEMP_KERNEL: [[3], [4], [6], [3]]
+NONLOCAL:
+ LOCATION: [[[]], [[]], [[]], [[]]]
+ GROUP: [[1], [1], [1], [1]]
+ INSTANTIATION: softmax
+BN:
+ USE_PRECISE_STATS: True
+ NUM_BATCHES_PRECISE: 200
+MODEL:
+ NUM_CLASSES: 1
+ ARCH: i3d
+ MODEL_NAME: ResNet
+ DROPOUT_RATE: 0.1
+ HEAD_ACT: sigmoid
+TEST:
+ ENABLE: True
+ DATASET: kinetics
+ BATCH_SIZE: 64
+DATA_LOADER:
+ NUM_WORKERS: 8
+ PIN_MEMORY: True
+NUM_GPUS: 8
+NUM_SHARDS: 1
+RNG_SEED: 0
+OUTPUT_DIR: .
+"""
+
+# from .f3net import FAD_Head
+from slowfast.models.video_model_builder import ResNet as ResNetOri
+from slowfast.config.defaults import get_cfg
+import torch
+from torch import nn
+from config_ftcn import config as my_cfg
+from inspect import signature
+# Removed TimeTransformer import - not used directly in this file
+# Removed random import - not used in inference
+
+my_cfg.init_with_yaml()
+my_cfg.update_with_yaml("ftcn_tt.yaml")
+my_cfg.freeze()
+
+
+
+class CenterPatchPool(nn.Module):
+ """Simplified patch pooling for inference - always use center patch"""
+ def __init__(self):
+ super().__init__()
+
+ def forward(self, x):
+ # batch,channel,16,7x7
+ b, c, t, h, w = x.shape
+ x = x.reshape(b, c, t, h * w)
+ idx = h * w // 2 # Always use center patch for inference
+ x = x[..., idx]
+ return x
+
+
+def valid_idx(idx, h):
+ i = idx // h
+ j = idx % h
+ if j == 0 or i == h - 1 or j == h - 1:
+ return False
+ else:
+ return True
+
+
+class CenterAvgPool(nn.Module):
+ """Simplified average pooling for inference - use all valid patches"""
+ def __init__(self):
+ super().__init__()
+
+ def forward(self, x):
+ # batch,channel,16,7x7
+ b, c, t, h, w = x.shape
+ x = x.reshape(b, c, t, h * w)
+ candidates = list(range(h * w))
+ candidates = [idx for idx in candidates if valid_idx(idx, h)]
+ x = x[..., candidates].mean(-1)
+ return x
+
+
+
+# Removed duplicate TransformerHead class - using the one from temporal_transformer.py
+
+
+parameters = [parameter for parameter in signature(nn.Conv3d).parameters]
+
+spatial_count = my_cfg.model.inco.spatial_count
+keep_stride_count = my_cfg.model.inco.keep_stride_count
+
+
+def temporal_only_conv(module, name, removed, stride_removed=0):
+ """
+ Recursively put desired batch norm in nn.module module.
+
+ set module = net to start code.
+ """
+ # go through all attributes of module nn.module (e.g. network or layer) and put batch norms if present
+ for attr_str in dir(module):
+ sub_module = getattr(module, attr_str)
+ if type(sub_module) == nn.Conv3d:
+ target_spatial_size = 1
+ predefine_padding = {1: 0, 3: 1, 5: 2, 7: 3}
+ kernel_size = list(sub_module.kernel_size)
+ assert kernel_size[1] == kernel_size[2]
+ stride = sub_module.stride
+ extra = None
+ if stride[1] == stride[2] == 2:
+ stride_removed += 1
+ if stride_removed > keep_stride_count:
+ stride = [1, 1, 1]
+ extra = nn.MaxPool3d((1, 2, 2))
+
+ if kernel_size[1] == 1 and extra is None:
+ continue
+ padding = list(sub_module.padding)
+
+ kernel_size[1] = kernel_size[2] = target_spatial_size
+ padding[1] = padding[2] = predefine_padding[target_spatial_size]
+ if 'device' in parameters:
+ parameters.remove('device')
+ if 'dtype' in parameters:
+ parameters.remove('dtype')
+ param_dict = {key: getattr(sub_module, key) for key in parameters}
+
+ param_dict.update(kernel_size=kernel_size, padding=padding, stride=stride)
+
+ conv = nn.Conv3d(**param_dict)
+
+ new_module = conv
+
+ removed += 1
+ if removed > spatial_count:
+ setattr(module, attr_str, new_module)
+ if extra is not None:
+ if attr_str == "conv":
+ bn_str = "bn"
+ else:
+ bn_str = f"{attr_str}_bn"
+ bn_module = getattr(module, bn_str)
+ assert isinstance(bn_module, nn.BatchNorm3d)
+ new_bn_module = nn.Sequential(bn_module, extra)
+ setattr(module, bn_str, new_bn_module)
+ else:
+ print("keep spatial")
+ elif type(sub_module) == nn.Dropout:
+ new_module = nn.Dropout(p=0.5)
+ setattr(module, attr_str, new_module)
+ if my_cfg.model.inco.no_time_pool:
+ if type(sub_module) == nn.MaxPool3d:
+ kernel_size = list(sub_module.kernel_size)
+ if kernel_size[0] == 2:
+ kernel_size[0] = 1
+ setattr(module, attr_str, nn.MaxPool3d(kernel_size))
+ elif type(sub_module) == nn.AvgPool3d:
+ kernel_size = list(sub_module.kernel_size)
+ kernel_size[0] = 2 * kernel_size[0]
+ setattr(module, attr_str, nn.AvgPool3d(kernel_size))
+
+ # iterate through immediate child modules. Note, the recursion is done by our code no need to use named_modules()
+ old_name = name
+ for name, immediate_child_module in module.named_children():
+ removed, stride_removed = temporal_only_conv(
+ immediate_child_module, old_name + "." + name, removed, stride_removed
+ )
+ return removed, stride_removed
+
+
+class I3D8x8(nn.Module):
+ def __init__(self) -> None:
+ super(I3D8x8, self).__init__()
+ cfg = get_cfg()
+ cfg.merge_from_str(config_text)
+ cfg.NUM_GPUS = 1
+ cfg.TEST.BATCH_SIZE = 1
+ cfg.TRAIN.BATCH_SIZE = 1
+
+ cfg.DATA.NUM_FRAMES = my_cfg.clip_size
+ SOLVER = my_cfg.model.inco.SOLVER
+ if SOLVER is not None:
+ for key, val in SOLVER.to_dict().items():
+ old_val = getattr(cfg.SOLVER, key)
+ val = type(old_val)(val)
+ setattr(cfg.SOLVER, key, val)
+
+ if my_cfg.model.inco.i3d_routine:
+ self.cfg = cfg
+ self.resnet = ResNetOri(cfg)
+ temporal_only_conv(self.resnet, "model", 0)
+
+ stop_point = my_cfg.model.transformer.stop_point
+
+ for i in [5, 4, 3]:
+ if stop_point <= i:
+ setattr(self.resnet, f"s{i}", nn.Identity())
+ if stop_point == 3:
+ setattr(self.resnet, f"pathway0_pool", nn.Identity())
+ gc.collect()
+ torch.cuda.empty_cache()
+
+ def forward(
+ self,
+ images, ft_features=None,
+ freeze_backbone=False
+ ):
+ assert not freeze_backbone
+
+ inputs = [images]
+ pred = self.resnet(inputs, ft_features)
+ return pred
diff --git a/clean/video/pwtf_dvd/inference/root_setting.yaml b/clean/video/pwtf_dvd/inference/root_setting.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..b66140f63b9d87e5a977503d7eaf65fdc5510f42
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/root_setting.yaml
@@ -0,0 +1,192 @@
+# this is the root setting of all setting and will be loaded in first place
+# the loading sequence is root_setting -> specific setting
+# -> manully overided setting -> computed setting in finalize_config()
+# which means access config before finalize can be dangerous
+# always remember to add a space after :
+
+setting_name: base
+test_dataset_type: null
+clip_size: 8
+var_clip_size: -1
+
+model:
+ teacher: null
+ pretrained: null
+ fc_weight: null
+ fc_bias: null
+ patch:
+ name: null
+ fc_only: false
+ feat_type: image
+ tester: null
+ tester_path: null
+ tester_data_mode: null
+ decodec:
+ reg_loss_type: l1
+ feat_loss_func: l1
+ dis_model: basic
+ gan_type: vanilla
+ inco:
+ kernel_size: 1
+ pool_ref_type: sum
+ temp_kernel_size: 3
+ spatial_sizes: []
+ spatial_count: 0
+ no_time_pool: false
+ keep_stride_count: 0
+ i3d_routine: true
+ SOLVER:
+ BASE_LR: 0.1
+ LR_POLICY: cosine
+ MAX_EPOCH: 196
+ MOMENTUM: 0.9
+ WEIGHT_DECAY: 1e-4
+ WARMUP_EPOCHS: 34.0
+ WARMUP_START_LR: 0.01
+ OPTIMIZING_METHOD: sgd
+ transformer:
+ patch_type: null #time,spatial,all
+ dim: -1
+ stop_point: 6
+ random_select: true
+ k: 8
+ sigmoid_before: false
+ denoise:
+ layers: []
+
+test_on_train: false
+
+test: false
+debug: false
+pre_load_data: false
+# for better compatiable with philly and potential running enviroment
+# all path under path should be rel_path w.r.t the config.py
+# and the abspath will be compute when finalize
+
+# python -m torch.distributed.launch --nproc_per_node=2 main.py --setting get_all_datas_id_emb_old.yaml main.py
+# --config trainer.default.log_step=1 trainer.default.sample_step=20
+
+trash_face: false
+
+strategies: []
+branches: ["continuous_same","continuous_diff","discontinuous"]
+epoch: -1
+
+reg_weight: 10
+class_weight: 1
+final_weight: 1
+feat_weight: 0
+gan_weight: 0
+
+noise_only: false
+error_only: false
+force_ds_scale: 1.0
+
+path:
+ model_dir: ../checkpoint
+ pretrain_dir: ../pretrain
+ log_dir: ../checkpoint
+ data_dir: ../data
+ precomputed_dir: ../precomputed
+ test_data_dir: null
+ extra_data_dir: ../extra_data
+ lmdb_dir: null
+
+patch_input: false
+input_noise: false
+mask_direct: false
+max_to_keep: 50
+base_count: -1
+enable_lmdb_cache: false
+
+aug_in_train: true
+aug_in_test: false
+vis_in_train: true
+
+data_mode: "image"
+data_source: "zip"
+
+aug:
+ flip_prob: 0
+ reverse_prob: 0
+ gray_prob: 0
+ size_aug_prob: 0
+ quantify_prob: 0
+ quantify_steps: [16,32]
+ min_size: 64
+ max_size: 256
+ jpeg_aug_prob: 0
+ dxy_gauss_prob: 0
+ dxy_gauss_scale: 0
+ min_quality: 60
+ max_quality: 100
+ gaussain: false
+ need_img_degrade: true
+ need_mask_distortion: true
+ need_color_match: true
+ feather_range: [0.2,0.2]
+ adaptive_clip_size: false
+ clip_sizes: null
+ cutout: 0
+ earse: 0
+ types: null
+ earse_type: null
+ time_earse_prob: 0
+ time_earse_type: null
+ jitter_prob: 0
+ test_types: null
+ raw_prob: 0
+ bi_types: null
+ bi_count: 1
+ face_sources: null
+ inplace_types: null
+ no_poisson_prob: 0
+ blend:
+ pseudo_prob: 0.5
+ from_real_prob: 1
+ skip_prob: 0
+ shuffle_prob: 0
+ multi_prob: 0
+ multi_config:
+ from_real_prob: 1
+ self_prob: 1
+ self:
+ one: 0
+ continuity: 1
+ other:
+ one: 0
+ all: 0
+ continuity: 1
+ celeb_prob: 0
+ fft:
+ type: null
+ fad:
+ type: null
+ flag: high
+ tta:
+ type: null
+ param: null
+
+
+batch_size: 64
+test_batch_size: 64
+imsize: 256
+next_frame_rate: 0.5
+trainer:
+ default:
+ apex_option: O0
+ n_worker: 12
+ optim: adam
+ model_save_step: 1000
+ log_step: 25
+ sample_step: 1000
+ init_lr: 3e-4
+ validation_step: 5000
+ test_sample_step: 20
+ one_test_step: 500
+ test_freq: 1
+ total_step: 100000
+ lr_step: 20000
+ freeze_backbone_step: 0
+ total_epoch: 200
+fetch_method: prefetch
\ No newline at end of file
diff --git a/clean/video/pwtf_dvd/inference/setting/ftcn_tt.yaml b/clean/video/pwtf_dvd/inference/setting/ftcn_tt.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..a6bb2fbe8b9c92f7d9260ea205ac714c5b3c6b4a
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/setting/ftcn_tt.yaml
@@ -0,0 +1,114 @@
+# this is the root setting of all setting and will be loaded in first place
+# the loading sequence is root_setting -> specific setting
+# -> manully overided setting -> computed setting in finalize_config()
+# which means access config before finalize can be dangerous
+# always remember to add a space after ":"
+
+setting_name: base
+data_mode: image
+# for better compatiable with philly and potential running enviroment
+# all path under path should be rel_path w.r.t the config.py
+# and the abspath will be compute when finalize
+
+# python -m torch.distributed.launch --nproc_per_node=2 main.py --setting naive_raw.yaml
+# --config trainer.default.log_step=1 trainer.default.sample_step=20
+strategies: ["scale_mean","scale_0","scale_1","scale_2","scale_3"]
+mask_direct: true
+
+clip_size: 32
+
+reg_weight: 1
+class_weight: 1
+final_weight: 1
+model:
+ inco:
+ spatial_count: 0
+ SOLVER:
+ BASE_LR: 0.1
+ LR_POLICY: cosine
+ MAX_EPOCH: 100
+ MOMENTUM: 0.9
+ WEIGHT_DECAY: 1e-4
+ WARMUP_EPOCHS: 10
+ WARMUP_START_LR: 0.01
+ OPTIMIZING_METHOD: sgd
+ transformer:
+ patch_type: time
+ stop_point: 5
+ depth: 1
+
+path:
+ model_dir: ../checkpoint
+ pretrain_dir: ../pretrain
+ log_dir: ../checkpoint
+ data_dir: host:lmdb_dir
+ precomputed_dir: null
+
+trainer_type: YL3DIncoPolicyS
+dataset_type: YL_3D_INCO_BASE_ZIP_PNG_S
+classifier_type: i3d_temporal_var_fix_dropout_tt_cfg
+
+imsize: 224
+base_count: 12
+aug_in_train: true
+test_on_train: true
+
+next_frame_rate: 0.0
+aug:
+ min_size: 64
+ max_size: 317
+ min_quality: 60
+ max_quality: 100
+ need_img_degrade: false
+ need_mask_distortion: true
+ need_color_match: true
+ max_step: 4
+ compression: false
+ cutout: 0
+ earse: 1
+ aug_prob: 0
+ types: ["C23_NOISE"]
+ earse_type: ["strong_black"]
+
+dataset:
+ real_train:
+ original_c23: 1
+ fake_train:
+ NeuralTextures_c23: 1
+ Face2Face_c23: 1
+ FaceSwap_c23: 1
+ Deepfakes_c23: 1
+ aug_online:
+ empty: 1
+ tests:
+ NeuralTextures_c23: ["NeuralTextures_c23"]
+ Face2Face_c23: ["Face2Face_c23"]
+ FaceSwap_c23: ["FaceSwap_c23"]
+ Deepfakes_c23: ["Deepfakes_c23"]
+
+max_to_keep: 100
+
+data_source: lmdb
+
+trainer:
+ default:
+ apex_option: O0
+ batch_size: 16
+ test_batch_size: 8
+ model_save_step: 10000
+ log_step: 200
+ sample_step: 1000
+ init_lr: 3e-4
+ total_epoch: 1000
+ one_test_step: 200
+ detach_step: 5000
+ validation_step: 10000
+ freeze_backbone_step: 0
+ total_step: 200000
+ lr_step: 100000
+
+
+
+#classifier:
+# default:
+# pretrained: false
\ No newline at end of file
diff --git a/clean/video/pwtf_dvd/inference/slowfast/__init__.py b/clean/video/pwtf_dvd/inference/slowfast/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..b65a8f9dc0cd5092e2dc879f6bb4aba8fafbb092
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/__init__.py
@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+from slowfast.utils.env import setup_environment
+
+setup_environment()
diff --git a/clean/video/pwtf_dvd/inference/slowfast/config/__init__.py b/clean/video/pwtf_dvd/inference/slowfast/config/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..8dbe96a785072a24a9bcc4841a1934024f2b06a1
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/config/__init__.py
@@ -0,0 +1,2 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
diff --git a/clean/video/pwtf_dvd/inference/slowfast/config/custom_config.py b/clean/video/pwtf_dvd/inference/slowfast/config/custom_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..8131da2951d8cb629f664b39da4675d0ae5adee5
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/config/custom_config.py
@@ -0,0 +1,9 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Add custom configs and default values"""
+
+
+def add_custom_config(_C):
+ # Add your own customized configs.
+ pass
diff --git a/clean/video/pwtf_dvd/inference/slowfast/config/defaults.py b/clean/video/pwtf_dvd/inference/slowfast/config/defaults.py
new file mode 100644
index 0000000000000000000000000000000000000000..73f0f562c3f95d92382d364c8a8ccbe5b6c27ca3
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/config/defaults.py
@@ -0,0 +1,816 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Configs."""
+import yaml
+from fvcore.common.config import CfgNode as CfgNodeOri
+
+from . import custom_config
+def load_yaml_with_base(text: str, allow_unsafe: bool = False):
+ """
+ Just like `yaml.load(open(filename))`, but inherit attributes from its
+ `_BASE_`.
+ Args:
+ text (str): the file name of the current config. Will be used to
+ find the base config file.
+ allow_unsafe (bool): whether to allow loading the config file with
+ `yaml.unsafe_load`.
+ Returns:
+ (dict): the loaded yaml
+ """
+ cfg = yaml.load(text, Loader=yaml.FullLoader)
+ return cfg
+class CfgNode(CfgNodeOri):
+ def merge_from_str(self, text, allow_unsafe=False):
+ loaded_cfg = load_yaml_with_base(text, allow_unsafe=allow_unsafe)
+ loaded_cfg = type(self)(loaded_cfg)
+ self.merge_from_other_cfg(loaded_cfg)
+
+# -----------------------------------------------------------------------------
+# Config definition
+# -----------------------------------------------------------------------------
+_C = CfgNode()
+
+
+# ---------------------------------------------------------------------------- #
+# Batch norm options
+# ---------------------------------------------------------------------------- #
+_C.BN = CfgNode()
+
+# Precise BN stats.
+_C.BN.USE_PRECISE_STATS = False
+
+# Number of samples use to compute precise bn.
+_C.BN.NUM_BATCHES_PRECISE = 200
+
+# Weight decay value that applies on BN.
+_C.BN.WEIGHT_DECAY = 0.0
+
+# Norm type, options include `batchnorm`, `sub_batchnorm`, `sync_batchnorm`
+_C.BN.NORM_TYPE = "batchnorm"
+
+# Parameter for SubBatchNorm, where it splits the batch dimension into
+# NUM_SPLITS splits, and run BN on each of them separately independently.
+_C.BN.NUM_SPLITS = 1
+
+# Parameter for NaiveSyncBatchNorm3d, where the stats across `NUM_SYNC_DEVICES`
+# devices will be synchronized.
+_C.BN.NUM_SYNC_DEVICES = 1
+
+
+# ---------------------------------------------------------------------------- #
+# Training options.
+# ---------------------------------------------------------------------------- #
+_C.TRAIN = CfgNode()
+
+# If True Train the model, else skip training.
+_C.TRAIN.ENABLE = True
+
+# Dataset.
+_C.TRAIN.DATASET = "kinetics"
+
+# Total mini-batch size.
+_C.TRAIN.BATCH_SIZE = 64
+
+_C.TRAIN.SPLIT = "train_subset2.pth"
+# Evaluate model on test data every eval period epochs.
+_C.TRAIN.EVAL_PERIOD = 1
+
+# Save model checkpoint every checkpoint period epochs.
+_C.TRAIN.CHECKPOINT_PERIOD = 1
+
+# Save model checkpoint every checkpoint period iters.
+_C.TRAIN.CHECKPOINT_PERIOD_BY_ITER = 500
+
+
+# Resume training from the latest checkpoint in the output directory.
+_C.TRAIN.AUTO_RESUME = True
+
+# Path to the checkpoint to load the initial weight.
+_C.TRAIN.CHECKPOINT_FILE_PATH = ""
+
+# Checkpoint types include `caffe2` or `pytorch`.
+_C.TRAIN.CHECKPOINT_TYPE = "pytorch"
+
+# If True, perform inflation when loading checkpoint.
+_C.TRAIN.CHECKPOINT_INFLATE = False
+
+
+# ---------------------------------------------------------------------------- #
+# Testing options
+# ---------------------------------------------------------------------------- #
+_C.TEST = CfgNode()
+
+# If True test the model, else skip the testing.
+_C.TEST.ENABLE = True
+
+# Dataset for testing.
+_C.TEST.DATASET = "kinetics"
+
+_C.TEST.SPLIT = "test_subset2.pth"
+# Total mini-batch size
+_C.TEST.BATCH_SIZE = 8
+
+# Path to the checkpoint to load the initial weight.
+_C.TEST.CHECKPOINT_FILE_PATH = ""
+
+# Number of clips to sample from a video uniformly for aggregating the
+# prediction results.
+_C.TEST.NUM_ENSEMBLE_VIEWS = 10
+
+# Number of crops to sample from a frame spatially for aggregating the
+# prediction results.
+_C.TEST.NUM_SPATIAL_CROPS = 3
+
+# Checkpoint types include `caffe2` or `pytorch`.
+_C.TEST.CHECKPOINT_TYPE = "pytorch"
+# Path to saving prediction results file.
+_C.TEST.SAVE_RESULTS_PATH = ""
+# -----------------------------------------------------------------------------
+# ResNet options
+# -----------------------------------------------------------------------------
+_C.RESNET = CfgNode()
+
+# Transformation function.
+_C.RESNET.TRANS_FUNC = "bottleneck_transform"
+
+# Number of groups. 1 for ResNet, and larger than 1 for ResNeXt).
+_C.RESNET.NUM_GROUPS = 1
+
+# Width of each group (64 -> ResNet; 4 -> ResNeXt).
+_C.RESNET.WIDTH_PER_GROUP = 64
+
+# Apply relu in a inplace manner. ########## FIxed by TH for training FT #################
+_C.RESNET.INPLACE_RELU = False # Default True for C2 and torch.
+
+# Apply stride to 1x1 conv.
+_C.RESNET.STRIDE_1X1 = False
+
+# If true, initialize the gamma of the final BN of each block to zero.
+_C.RESNET.ZERO_INIT_FINAL_BN = False
+
+# Number of weight layers.
+_C.RESNET.DEPTH = 50
+
+
+# label of branchs
+_C.RESNET.LABELS = ["continus","discontinus"]
+
+# If the current block has more than NUM_BLOCK_TEMP_KERNEL blocks, use temporal
+# kernel of 1 for the rest of the blocks.
+_C.RESNET.NUM_BLOCK_TEMP_KERNEL = [[3], [4], [6], [3]]
+
+# Size of stride on different res stages.
+_C.RESNET.SPATIAL_STRIDES = [[1], [2], [2], [2]]
+
+# Size of dilation on different res stages.
+_C.RESNET.SPATIAL_DILATIONS = [[1], [1], [1], [1]]
+
+
+# -----------------------------------------------------------------------------
+# Nonlocal options
+# -----------------------------------------------------------------------------
+_C.NONLOCAL = CfgNode()
+
+# Index of each stage and block to add nonlocal layers.
+_C.NONLOCAL.LOCATION = [[[]], [[]], [[]], [[]]]
+
+# Number of group for nonlocal for each stage.
+_C.NONLOCAL.GROUP = [[1], [1], [1], [1]]
+
+# Instatiation to use for non-local layer.
+_C.NONLOCAL.INSTANTIATION = "dot_product"
+
+
+# Size of pooling layers used in Non-Local.
+_C.NONLOCAL.POOL = [
+ # Res2
+ [[1, 2, 2], [1, 2, 2]],
+ # Res3
+ [[1, 2, 2], [1, 2, 2]],
+ # Res4
+ [[1, 2, 2], [1, 2, 2]],
+ # Res5
+ [[1, 2, 2], [1, 2, 2]],
+]
+
+# -----------------------------------------------------------------------------
+# Model options
+# -----------------------------------------------------------------------------
+_C.MODEL = CfgNode()
+
+# Model architecture.
+_C.MODEL.ARCH = "slowfast"
+
+# Model name
+_C.MODEL.MODEL_NAME = "SlowFast"
+
+# The number of classes to predict for the model.
+_C.MODEL.NUM_CLASSES = 400
+
+# Loss function.
+_C.MODEL.LOSS_FUNC = "cross_entropy"
+
+_C.MODEL.MASK_WEIGHT = 100
+
+_C.MODEL.CLASS_WEIGHT = 1
+
+# Model architectures that has one single pathway.
+_C.MODEL.SINGLE_PATHWAY_ARCH = ["c2d", "i3d", "slow"]
+
+# Model architectures that has multiple pathways.
+_C.MODEL.MULTI_PATHWAY_ARCH = ["slowfast"]
+
+# Dropout rate before final projection in the backbone.
+_C.MODEL.DROPOUT_RATE = 0.5
+
+# The std to initialize the fc layer(s).
+_C.MODEL.FC_INIT_STD = 0.01
+
+# Activation layer for the output head.
+_C.MODEL.HEAD_ACT = "softmax"
+
+
+# -----------------------------------------------------------------------------
+# SlowFast options
+# -----------------------------------------------------------------------------
+_C.SLOWFAST = CfgNode()
+
+# Corresponds to the inverse of the channel reduction ratio, $\beta$ between
+# the Slow and Fast pathways.
+_C.SLOWFAST.BETA_INV = 8
+
+# Corresponds to the frame rate reduction ratio, $\alpha$ between the Slow and
+# Fast pathways.
+_C.SLOWFAST.ALPHA = 8
+
+# Ratio of channel dimensions between the Slow and Fast pathways.
+_C.SLOWFAST.FUSION_CONV_CHANNEL_RATIO = 2
+
+# Kernel dimension used for fusing information from Fast pathway to Slow
+# pathway.
+_C.SLOWFAST.FUSION_KERNEL_SZ = 5
+
+
+# -----------------------------------------------------------------------------
+# Data options
+# -----------------------------------------------------------------------------
+_C.DATA = CfgNode()
+
+# The path to the data directory.
+_C.DATA.PATH_TO_DATA_DIR = ""
+
+_C.DATA.DATASET = "faceforensics"
+
+_C.DATA.MODE = ""
+
+_C.DATA.ADAPTIVE = False
+
+_C.DATA.SCALE = 1.0
+# The separator used between path and label.
+_C.DATA.PATH_LABEL_SEPARATOR = " "
+
+# Video path prefix if any.
+_C.DATA.PATH_PREFIX = ""
+
+# The spatial crop size of the input clip.
+_C.DATA.CROP_SIZE = 224
+
+# The number of frames of the input clip.
+_C.DATA.NUM_FRAMES = 8
+
+_C.DATA.NUM_FRAMES_RANGE = [1,2,3,4,5,6,7,8]
+
+# The video sampling rate of the input clip.
+_C.DATA.SAMPLING_RATE = 8
+
+# The mean value of the video raw pixels across the R G B channels.
+_C.DATA.MEAN = [0.45, 0.45, 0.45]
+# List of input frame channel dimensions.
+
+_C.DATA.INPUT_CHANNEL_NUM = [12, 12]
+
+# The std value of the video raw pixels across the R G B channels.
+_C.DATA.STD = [0.225, 0.225, 0.225]
+
+# The spatial augmentation jitter scales for training.
+_C.DATA.TRAIN_JITTER_SCALES = [256, 320]
+
+# The spatial crop size for training.
+_C.DATA.TRAIN_CROP_SIZE = 224
+
+# The spatial crop size for testing.
+_C.DATA.TEST_CROP_SIZE = 256
+
+# Input videos may has different fps, convert it to the target video fps before
+# frame sampling.
+_C.DATA.TARGET_FPS = 30
+
+# Decoding backend, options include `pyav` or `torchvision`
+_C.DATA.DECODING_BACKEND = "pyav"
+
+# if True, sample uniformly in [1 / max_scale, 1 / min_scale] and take a
+# reciprocal to get the scale. If False, take a uniform sample from
+# [min_scale, max_scale].
+_C.DATA.INV_UNIFORM_SAMPLE = False
+
+# If True, perform random horizontal flip on the video frames during training.
+_C.DATA.RANDOM_FLIP = True
+
+# If True, calculdate the map as metric.
+_C.DATA.MULTI_LABEL = False
+
+# Method to perform the ensemble, options include "sum" and "max".
+_C.DATA.ENSEMBLE_METHOD = "sum"
+
+# If True, revert the default input channel (RBG <-> BGR).
+_C.DATA.REVERSE_INPUT_CHANNEL = False
+
+
+# ---------------------------------------------------------------------------- #
+# Optimizer options
+# ---------------------------------------------------------------------------- #
+_C.SOLVER = CfgNode()
+
+# Base learning rate.
+_C.SOLVER.BASE_LR = 0.1
+
+# Learning rate policy (see utils/lr_policy.py for options and examples).
+_C.SOLVER.LR_POLICY = "cosine"
+
+# Exponential decay factor.
+_C.SOLVER.GAMMA = 0.1
+
+# Step size for 'exp' and 'cos' policies (in epochs).
+_C.SOLVER.STEP_SIZE = 1
+
+# Steps for 'steps_' policies (in epochs).
+_C.SOLVER.STEPS = []
+
+# Learning rates for 'steps_' policies.
+_C.SOLVER.LRS = []
+
+# Maximal number of epochs.
+_C.SOLVER.MAX_EPOCH = 300
+
+# Momentum.
+_C.SOLVER.MOMENTUM = 0.9
+
+# Momentum dampening.
+_C.SOLVER.DAMPENING = 0.0
+
+# Nesterov momentum.
+_C.SOLVER.NESTEROV = True
+
+# L2 regularization.
+_C.SOLVER.WEIGHT_DECAY = 1e-4
+
+# Start the warm up from SOLVER.BASE_LR * SOLVER.WARMUP_FACTOR.
+_C.SOLVER.WARMUP_FACTOR = 0.1
+
+# Gradually warm up the SOLVER.BASE_LR over this number of epochs.
+_C.SOLVER.WARMUP_EPOCHS = 0.0
+
+# The start learning rate of the warm up.
+_C.SOLVER.WARMUP_START_LR = 0.01
+
+# Optimization method.
+_C.SOLVER.OPTIMIZING_METHOD = "sgd"
+
+_C.SOLVER.LR_STEP = 50000
+
+_C.SOLVER.TOTAL_STEP = 200000
+
+_C.SOLVER.FREEZE_STEP = 10000
+
+
+# ---------------------------------------------------------------------------- #
+# Misc options
+# ---------------------------------------------------------------------------- #
+
+# Number of GPUs to use (applies to both training and testing).
+_C.NUM_GPUS = 1
+
+# Number of machine to use for the job.
+_C.NUM_SHARDS = 1
+
+# The index of the current machine.
+_C.SHARD_ID = 0
+
+# Output basedir.
+_C.OUTPUT_DIR = "./tmp"
+
+# train module
+_C.TRAIN_MODULE= "train_unet_by_iter"
+
+# Note that non-determinism may still be present due to non-deterministic
+# operator implementations in GPU operator libraries.
+_C.RNG_SEED = 1
+
+# Log period in iters.
+_C.LOG_PERIOD = 10
+
+# If True, log the model info.
+_C.LOG_MODEL_INFO = True
+
+# Distributed backend.
+_C.DIST_BACKEND = "nccl"
+
+# ---------------------------------------------------------------------------- #
+# Benchmark options
+# ---------------------------------------------------------------------------- #
+_C.BENCHMARK = CfgNode()
+
+# Number of epochs for data loading benchmark.
+_C.BENCHMARK.NUM_EPOCHS = 5
+
+# Log period in iters for data loading benchmark.
+_C.BENCHMARK.LOG_PERIOD = 100
+
+# If True, shuffle dataloader for epoch during benchmark.
+_C.BENCHMARK.SHUFFLE = True
+
+
+# ---------------------------------------------------------------------------- #
+# Common train/test data loader options
+# ---------------------------------------------------------------------------- #
+_C.DATA_LOADER = CfgNode()
+
+# Number of data loader workers per training process.
+_C.DATA_LOADER.NUM_WORKERS = 8
+
+# Load data to pinned host memory.
+_C.DATA_LOADER.PIN_MEMORY = True
+
+# Enable multi thread decoding.
+_C.DATA_LOADER.ENABLE_MULTI_THREAD_DECODE = False
+
+
+# ---------------------------------------------------------------------------- #
+# Detection options.
+# ---------------------------------------------------------------------------- #
+_C.DETECTION = CfgNode()
+
+# Whether enable video detection.
+_C.DETECTION.ENABLE = False
+
+# Aligned version of RoI. More details can be found at slowfast/models/head_helper.py
+_C.DETECTION.ALIGNED = True
+
+# Spatial scale factor.
+_C.DETECTION.SPATIAL_SCALE_FACTOR = 16
+
+# RoI tranformation resolution.
+_C.DETECTION.ROI_XFORM_RESOLUTION = 7
+
+
+# -----------------------------------------------------------------------------
+# AVA Dataset options
+# -----------------------------------------------------------------------------
+_C.AVA = CfgNode()
+
+# Directory path of frames.
+_C.AVA.FRAME_DIR = "/mnt/fair-flash3-east/ava_trainval_frames.img/"
+
+# Directory path for files of frame lists.
+_C.AVA.FRAME_LIST_DIR = (
+ "/mnt/vol/gfsai-flash3-east/ai-group/users/haoqifan/ava/frame_list/"
+)
+
+# Directory path for annotation files.
+_C.AVA.ANNOTATION_DIR = (
+ "/mnt/vol/gfsai-flash3-east/ai-group/users/haoqifan/ava/frame_list/"
+)
+
+# Filenames of training samples list files.
+_C.AVA.TRAIN_LISTS = ["train.csv"]
+
+# Filenames of test samples list files.
+_C.AVA.TEST_LISTS = ["val.csv"]
+
+# Filenames of box list files for training. Note that we assume files which
+# contains predicted boxes will have a suffix "predicted_boxes" in the
+# filename.
+_C.AVA.TRAIN_GT_BOX_LISTS = ["ava_train_v2.2.csv"]
+_C.AVA.TRAIN_PREDICT_BOX_LISTS = []
+
+# Filenames of box list files for test.
+_C.AVA.TEST_PREDICT_BOX_LISTS = ["ava_val_predicted_boxes.csv"]
+
+# This option controls the score threshold for the predicted boxes to use.
+_C.AVA.DETECTION_SCORE_THRESH = 0.9
+
+# If use BGR as the format of input frames.
+_C.AVA.BGR = False
+
+# Training augmentation parameters
+# Whether to use color augmentation method.
+_C.AVA.TRAIN_USE_COLOR_AUGMENTATION = False
+
+# Whether to only use PCA jitter augmentation when using color augmentation
+# method (otherwise combine with color jitter method).
+_C.AVA.TRAIN_PCA_JITTER_ONLY = True
+
+# Eigenvalues for PCA jittering. Note PCA is RGB based.
+_C.AVA.TRAIN_PCA_EIGVAL = [0.225, 0.224, 0.229]
+
+# Eigenvectors for PCA jittering.
+_C.AVA.TRAIN_PCA_EIGVEC = [
+ [-0.5675, 0.7192, 0.4009],
+ [-0.5808, -0.0045, -0.8140],
+ [-0.5836, -0.6948, 0.4203],
+]
+
+# Whether to do horizontal flipping during test.
+_C.AVA.TEST_FORCE_FLIP = False
+
+# Whether to use full test set for validation split.
+_C.AVA.FULL_TEST_ON_VAL = False
+
+# The name of the file to the ava label map.
+_C.AVA.LABEL_MAP_FILE = "ava_action_list_v2.2_for_activitynet_2019.pbtxt"
+
+# The name of the file to the ava exclusion.
+_C.AVA.EXCLUSION_FILE = "ava_val_excluded_timestamps_v2.2.csv"
+
+# The name of the file to the ava groundtruth.
+_C.AVA.GROUNDTRUTH_FILE = "ava_val_v2.2.csv"
+
+# Backend to process image, includes `pytorch` and `cv2`.
+_C.AVA.IMG_PROC_BACKEND = "cv2"
+
+# ---------------------------------------------------------------------------- #
+# Multigrid training options
+# See https://arxiv.org/abs/1912.00998 for details about multigrid training.
+# ---------------------------------------------------------------------------- #
+_C.MULTIGRID = CfgNode()
+
+# Multigrid training allows us to train for more epochs with fewer iterations.
+# This hyperparameter specifies how many times more epochs to train.
+# The default setting in paper trains for 1.5x more epochs than baseline.
+_C.MULTIGRID.EPOCH_FACTOR = 1.5
+
+# Enable short cycles.
+_C.MULTIGRID.SHORT_CYCLE = False
+# Short cycle additional spatial dimensions relative to the default crop size.
+_C.MULTIGRID.SHORT_CYCLE_FACTORS = [0.5, 0.5 ** 0.5]
+
+_C.MULTIGRID.LONG_CYCLE = False
+# (Temporal, Spatial) dimensions relative to the default shape.
+_C.MULTIGRID.LONG_CYCLE_FACTORS = [
+ (0.25, 0.5 ** 0.5),
+ (0.5, 0.5 ** 0.5),
+ (0.5, 1),
+ (1, 1),
+]
+
+# While a standard BN computes stats across all examples in a GPU,
+# for multigrid training we fix the number of clips to compute BN stats on.
+# See https://arxiv.org/abs/1912.00998 for details.
+_C.MULTIGRID.BN_BASE_SIZE = 8
+
+# Multigrid training epochs are not proportional to actual training time or
+# computations, so _C.TRAIN.EVAL_PERIOD leads to too frequent or rare
+# evaluation. We use a multigrid-specific rule to determine when to evaluate:
+# This hyperparameter defines how many times to evaluate a model per long
+# cycle shape.
+_C.MULTIGRID.EVAL_FREQ = 3
+
+# No need to specify; Set automatically and used as global variables.
+_C.MULTIGRID.LONG_CYCLE_SAMPLING_RATE = 0
+_C.MULTIGRID.DEFAULT_B = 0
+_C.MULTIGRID.DEFAULT_T = 0
+_C.MULTIGRID.DEFAULT_S = 0
+
+# -----------------------------------------------------------------------------
+# Tensorboard Visualization Options
+# -----------------------------------------------------------------------------
+_C.TENSORBOARD = CfgNode()
+
+# Log to summary writer, this will automatically.
+# log loss, lr and metrics during train/eval.
+_C.TENSORBOARD.ENABLE = False
+# Provide path to prediction results for visualization.
+# This is a pickle file of [prediction_tensor, label_tensor]
+_C.TENSORBOARD.PREDICTIONS_PATH = ""
+# Path to directory for tensorboard logs.
+# Default to to cfg.OUTPUT_DIR/runs-{cfg.TRAIN.DATASET}.
+_C.TENSORBOARD.LOG_DIR = ""
+# Path to a json file providing class_name - id mapping
+# in the format {"class_name1": id1, "class_name2": id2, ...}.
+# This file must be provided to enable plotting confusion matrix
+# by a subset or parent categories.
+_C.TENSORBOARD.CLASS_NAMES_PATH = ""
+
+# Path to a json file for categories -> classes mapping
+# in the format {"parent_class": ["child_class1", "child_class2",...], ...}.
+_C.TENSORBOARD.CATEGORIES_PATH = ""
+
+# Config for confusion matrices visualization.
+_C.TENSORBOARD.CONFUSION_MATRIX = CfgNode()
+# Visualize confusion matrix.
+_C.TENSORBOARD.CONFUSION_MATRIX.ENABLE = False
+# Figure size of the confusion matrices plotted.
+_C.TENSORBOARD.CONFUSION_MATRIX.FIGSIZE = [8, 8]
+# Path to a subset of categories to visualize.
+# File contains class names separated by newline characters.
+_C.TENSORBOARD.CONFUSION_MATRIX.SUBSET_PATH = ""
+
+# Config for histogram visualization.
+_C.TENSORBOARD.HISTOGRAM = CfgNode()
+# Visualize histograms.
+_C.TENSORBOARD.HISTOGRAM.ENABLE = False
+# Path to a subset of classes to plot histograms.
+# Class names must be separated by newline characters.
+_C.TENSORBOARD.HISTOGRAM.SUBSET_PATH = ""
+# Visualize top-k most predicted classes on histograms for each
+# chosen true label.
+_C.TENSORBOARD.HISTOGRAM.TOPK = 10
+# Figure size of the histograms plotted.
+_C.TENSORBOARD.HISTOGRAM.FIGSIZE = [8, 8]
+
+# Config for layers' weights and activations visualization.
+# _C.TENSORBOARD.ENABLE must be True.
+_C.TENSORBOARD.MODEL_VIS = CfgNode()
+
+# If False, skip model visualization.
+_C.TENSORBOARD.MODEL_VIS.ENABLE = False
+
+# If False, skip visualizing model weights.
+_C.TENSORBOARD.MODEL_VIS.MODEL_WEIGHTS = False
+
+# If False, skip visualizing model activations.
+_C.TENSORBOARD.MODEL_VIS.ACTIVATIONS = False
+
+# If False, skip visualizing input videos.
+_C.TENSORBOARD.MODEL_VIS.INPUT_VIDEO = False
+
+
+# List of strings containing data about layer names and their indexing to
+# visualize weights and activations for. The indexing is meant for
+# choosing a subset of activations outputed by a layer for visualization.
+# If indexing is not specified, visualize all activations outputed by the layer.
+# For each string, layer name and indexing is separated by whitespaces.
+# e.g.: [layer1 1,2;1,2, layer2, layer3 150,151;3,4]; this means for each array `arr`
+# along the batch dimension in `layer1`, we take arr[[1, 2], [1, 2]]
+_C.TENSORBOARD.MODEL_VIS.LAYER_LIST = []
+# Top-k predictions to plot on videos
+_C.TENSORBOARD.MODEL_VIS.TOPK_PREDS = 1
+# Colormap to for text boxes and bounding boxes colors
+_C.TENSORBOARD.MODEL_VIS.COLORMAP = "Pastel2"
+# Config for visualization video inputs with Grad-CAM.
+# _C.TENSORBOARD.ENABLE must be True.
+_C.TENSORBOARD.MODEL_VIS.GRAD_CAM = CfgNode()
+# Whether to run visualization using Grad-CAM technique.
+_C.TENSORBOARD.MODEL_VIS.GRAD_CAM.ENABLE = True
+# CNN layers to use for Grad-CAM. The number of layers must be equal to
+# number of pathway(s).
+_C.TENSORBOARD.MODEL_VIS.GRAD_CAM.LAYER_LIST = []
+# If True, visualize Grad-CAM using true labels for each instances.
+# If False, use the highest predicted class.
+_C.TENSORBOARD.MODEL_VIS.GRAD_CAM.USE_TRUE_LABEL = False
+# Colormap to for text boxes and bounding boxes colors
+_C.TENSORBOARD.MODEL_VIS.GRAD_CAM.COLORMAP = "viridis"
+
+# Config for visualization for wrong prediction visualization.
+# _C.TENSORBOARD.ENABLE must be True.
+_C.TENSORBOARD.WRONG_PRED_VIS = CfgNode()
+_C.TENSORBOARD.WRONG_PRED_VIS.ENABLE = False
+# Folder tag to origanize model eval videos under.
+_C.TENSORBOARD.WRONG_PRED_VIS.TAG = "Incorrectly classified videos."
+# Subset of labels to visualize. Only wrong predictions with true labels
+# within this subset is visualized.
+_C.TENSORBOARD.WRONG_PRED_VIS.SUBSET_PATH = ""
+
+
+
+###############
+_C.JITTER = CfgNode()
+
+_C.JITTER.ENABLE = False
+
+_C.JITTER.CONTINUS_METHODS=["blend_diff_person","blend_downsampled","blend_same_person"]
+_C.JITTER.DISCONTINUS_METHODS=["light", "rotate", "skip"]
+
+_C.JITTER.STRONG_INNER_CLIP_MASK_JITTER= False
+
+# ---------------------------------------------------------------------------- #
+# Demo options
+# ---------------------------------------------------------------------------- #
+_C.DEMO = CfgNode()
+
+# Run model in DEMO mode.
+_C.DEMO.ENABLE = False
+
+# Path to a json file providing class_name - id mapping
+# in the format {"class_name1": id1, "class_name2": id2, ...}.
+_C.DEMO.LABEL_FILE_PATH = ""
+
+# Specify a camera device as input. This will be prioritized
+# over input video if set.
+# If -1, use input video instead.
+_C.DEMO.WEBCAM = -1
+
+# Path to input video for demo.
+_C.DEMO.INPUT_VIDEO = ""
+# Custom width for reading input video data.
+_C.DEMO.DISPLAY_WIDTH = 0
+# Custom height for reading input video data.
+_C.DEMO.DISPLAY_HEIGHT = 0
+# Path to Detectron2 object detection model configuration,
+# only used for detection tasks.
+_C.DEMO.DETECTRON2_CFG = "COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"
+# Path to Detectron2 object detection model pre-trained weights.
+_C.DEMO.DETECTRON2_WEIGHTS = "detectron2://COCO-Detection/faster_rcnn_R_50_FPN_3x/137849458/model_final_280758.pkl"
+# Threshold for choosing predicted bounding boxes by Detectron2.
+_C.DEMO.DETECTRON2_THRESH = 0.9
+# Number of overlapping frames between 2 consecutive clips.
+# Increase this number for more frequent action predictions.
+# The number of overlapping frames cannot be larger than
+# half of the sequence length `cfg.DATA.NUM_FRAMES * cfg.DATA.SAMPLING_RATE`
+_C.DEMO.BUFFER_SIZE = 0
+# If specified, the visualized outputs will be written this a video file of
+# this path. Otherwise, the visualized outputs will be displayed in a window.
+_C.DEMO.OUTPUT_FILE = ""
+# Frames per second rate for writing to output video file.
+# If not set (-1), use fps rate from input file.
+_C.DEMO.OUTPUT_FPS = -1
+# Input format from demo video reader ("RGB" or "BGR").
+_C.DEMO.INPUT_FORMAT = "BGR"
+# Draw visualization frames in [keyframe_idx - CLIP_VIS_SIZE, keyframe_idx + CLIP_VIS_SIZE] inclusively.
+_C.DEMO.CLIP_VIS_SIZE = 10
+# Number of processes to run video visualizer.
+_C.DEMO.NUM_VIS_INSTANCES = 2
+
+# Path to pre-computed predicted boxes
+_C.DEMO.PREDS_BOXES = ""
+# Whether to run in with multi-threaded video reader.
+_C.DEMO.THREAD_ENABLE = False
+# Take one clip for every `DEMO.NUM_CLIPS_SKIP` + 1 for prediction and visualization.
+# This is used for fast demo speed by reducing the prediction/visualiztion frequency.
+# If -1, take the most recent read clip for visualization. This mode is only supported
+# if `DEMO.THREAD_ENABLE` is set to True.
+_C.DEMO.NUM_CLIPS_SKIP = 0
+# Path to ground-truth boxes and labels (optional)
+_C.DEMO.GT_BOXES = ""
+# The starting second of the video w.r.t bounding boxes file.
+_C.DEMO.STARTING_SECOND = 900
+# Frames per second of the input video/folder of images.
+_C.DEMO.FPS = 30
+# Visualize with top-k predictions or predictions above certain threshold(s).
+# Option: {"thres", "top-k"}
+_C.DEMO.VIS_MODE = "thres"
+# Threshold for common class names.
+_C.DEMO.COMMON_CLASS_THRES = 0.7
+# Theshold for uncommon class names. This will not be
+# used if `_C.DEMO.COMMON_CLASS_NAMES` is empty.
+_C.DEMO.UNCOMMON_CLASS_THRES = 0.3
+# This is chosen based on distribution of examples in
+# each classes in AVA dataset.
+_C.DEMO.COMMON_CLASS_NAMES = [
+ "watch (a person)",
+ "talk to (e.g., self, a person, a group)",
+ "listen to (a person)",
+ "touch (an object)",
+ "carry/hold (an object)",
+ "walk",
+ "sit",
+ "lie/sleep",
+ "bend/bow (at the waist)",
+]
+# Slow-motion rate for the visualization. The visualized portions of the
+# video will be played `_C.DEMO.SLOWMO` times slower than usual speed.
+_C.DEMO.SLOWMO = 1
+
+# Add custom config with default values.
+custom_config.add_custom_config(_C)
+
+
+def _assert_and_infer_cfg(cfg):
+ # BN assertions.
+ if cfg.BN.USE_PRECISE_STATS:
+ assert cfg.BN.NUM_BATCHES_PRECISE >= 0
+ # TRAIN assertions.
+ assert cfg.TRAIN.CHECKPOINT_TYPE in ["pytorch", "caffe2"]
+ assert cfg.TRAIN.BATCH_SIZE % cfg.NUM_GPUS == 0
+
+ # TEST assertions.
+ assert cfg.TEST.CHECKPOINT_TYPE in ["pytorch", "caffe2"]
+ assert cfg.TEST.BATCH_SIZE % cfg.NUM_GPUS == 0
+ assert cfg.TEST.NUM_SPATIAL_CROPS == 3
+
+ # RESNET assertions.
+ assert cfg.RESNET.NUM_GROUPS > 0
+ assert cfg.RESNET.WIDTH_PER_GROUP > 0
+ assert cfg.RESNET.WIDTH_PER_GROUP % cfg.RESNET.NUM_GROUPS == 0
+
+ # General assertions.
+ assert cfg.SHARD_ID < cfg.NUM_SHARDS
+ return cfg
+
+
+def get_cfg():
+ """
+ Get a copy of the default config.
+ """
+ return _assert_and_infer_cfg(_C.clone())
diff --git a/clean/video/pwtf_dvd/inference/slowfast/models/__init__.py b/clean/video/pwtf_dvd/inference/slowfast/models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f82b3dc1e3af0dbabf4a3a6153e48b977eb1059e
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/models/__init__.py
@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+from .build import MODEL_REGISTRY, build_model # noqa
+from .custom_video_model_builder import * # noqa
+from .video_model_builder import ResNet, SlowFast # noqa
\ No newline at end of file
diff --git a/clean/video/pwtf_dvd/inference/slowfast/models/batchnorm_helper.py b/clean/video/pwtf_dvd/inference/slowfast/models/batchnorm_helper.py
new file mode 100644
index 0000000000000000000000000000000000000000..4e52d50497d9c0a58e5ace0a2fde94b5418ef563
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/models/batchnorm_helper.py
@@ -0,0 +1,218 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""BatchNorm (BN) utility functions and custom batch-size BN implementations"""
+
+from functools import partial
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from torch.autograd.function import Function
+
+import slowfast.utils.distributed as du
+
+
+def get_norm(cfg):
+ """
+ Args:
+ cfg (CfgNode): model building configs, details are in the comments of
+ the config file.
+ Returns:
+ nn.Module: the normalization layer.
+ """
+ if cfg.BN.NORM_TYPE == "batchnorm":
+ return nn.BatchNorm3d
+ elif cfg.BN.NORM_TYPE == "sub_batchnorm":
+ return partial(SubBatchNorm3d, num_splits=cfg.BN.NUM_SPLITS)
+ elif cfg.BN.NORM_TYPE == "sync_batchnorm":
+ return partial(
+ NaiveSyncBatchNorm3d, num_sync_devices=cfg.BN.NUM_SYNC_DEVICES
+ )
+ else:
+ raise NotImplementedError(
+ "Norm type {} is not supported".format(cfg.BN.NORM_TYPE)
+ )
+
+
+class SubBatchNorm3d(nn.Module):
+ """
+ The standard BN layer computes stats across all examples in a GPU. In some
+ cases it is desirable to compute stats across only a subset of examples
+ (e.g., in multigrid training https://arxiv.org/abs/1912.00998).
+ SubBatchNorm3d splits the batch dimension into N splits, and run BN on
+ each of them separately (so that the stats are computed on each subset of
+ examples (1/N of batch) independently. During evaluation, it aggregates
+ the stats from all splits into one BN.
+ """
+
+ def __init__(self, num_splits, **args):
+ """
+ Args:
+ num_splits (int): number of splits.
+ args (list): other arguments.
+ """
+ super(SubBatchNorm3d, self).__init__()
+ self.num_splits = num_splits
+ num_features = args["num_features"]
+ # Keep only one set of weight and bias.
+ if args.get("affine", True):
+ self.affine = True
+ args["affine"] = False
+ self.weight = torch.nn.Parameter(torch.ones(num_features))
+ self.bias = torch.nn.Parameter(torch.zeros(num_features))
+ else:
+ self.affine = False
+ self.bn = nn.BatchNorm3d(**args)
+ args["num_features"] = num_features * num_splits
+ self.split_bn = nn.BatchNorm3d(**args)
+
+ def _get_aggregated_mean_std(self, means, stds, n):
+ """
+ Calculate the aggregated mean and stds.
+ Args:
+ means (tensor): mean values.
+ stds (tensor): standard deviations.
+ n (int): number of sets of means and stds.
+ """
+ mean = means.view(n, -1).sum(0) / n
+ std = (
+ stds.view(n, -1).sum(0) / n
+ + ((means.view(n, -1) - mean) ** 2).view(n, -1).sum(0) / n
+ )
+ return mean.detach(), std.detach()
+
+ def aggregate_stats(self):
+ """
+ Synchronize running_mean, and running_var. Call this before eval.
+ """
+ if self.split_bn.track_running_stats:
+ (
+ self.bn.running_mean.data,
+ self.bn.running_var.data,
+ ) = self._get_aggregated_mean_std(
+ self.split_bn.running_mean,
+ self.split_bn.running_var,
+ self.num_splits,
+ )
+
+ def forward(self, x):
+ if self.training:
+ n, c, t, h, w = x.shape
+ x = x.view(n // self.num_splits, c * self.num_splits, t, h, w)
+ x = self.split_bn(x)
+ x = x.view(n, c, t, h, w)
+ else:
+ x = self.bn(x)
+ if self.affine:
+ x = x * self.weight.view((-1, 1, 1, 1))
+ x = x + self.bias.view((-1, 1, 1, 1))
+ return x
+
+
+class GroupGather(Function):
+ """
+ GroupGather performs all gather on each of the local process/ GPU groups.
+ """
+
+ @staticmethod
+ def forward(ctx, input, num_sync_devices, num_groups):
+ """
+ Perform forwarding, gathering the stats across different process/ GPU
+ group.
+ """
+ ctx.num_sync_devices = num_sync_devices
+ ctx.num_groups = num_groups
+
+ input_list = [
+ torch.zeros_like(input) for k in range(du.get_local_size())
+ ]
+ dist.all_gather(
+ input_list, input, async_op=False, group=du._LOCAL_PROCESS_GROUP
+ )
+
+ inputs = torch.stack(input_list, dim=0)
+ if num_groups > 1:
+ rank = du.get_local_rank()
+ group_idx = rank // num_sync_devices
+ inputs = inputs[
+ group_idx
+ * num_sync_devices : (group_idx + 1)
+ * num_sync_devices
+ ]
+ inputs = torch.sum(inputs, dim=0)
+ return inputs
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ """
+ Perform backwarding, gathering the gradients across different process/ GPU
+ group.
+ """
+ grad_output_list = [
+ torch.zeros_like(grad_output) for k in range(du.get_local_size())
+ ]
+ dist.all_gather(
+ grad_output_list,
+ grad_output,
+ async_op=False,
+ group=du._LOCAL_PROCESS_GROUP,
+ )
+
+ grads = torch.stack(grad_output_list, dim=0)
+ if ctx.num_groups > 1:
+ rank = du.get_local_rank()
+ group_idx = rank // ctx.num_sync_devices
+ grads = grads[
+ group_idx
+ * ctx.num_sync_devices : (group_idx + 1)
+ * ctx.num_sync_devices
+ ]
+ grads = torch.sum(grads, dim=0)
+ return grads, None, None
+
+
+class NaiveSyncBatchNorm3d(nn.BatchNorm3d):
+ def __init__(self, num_sync_devices, **args):
+ """
+ Naive version of Synchronized 3D BatchNorm.
+ Args:
+ num_sync_devices (int): number of device to sync.
+ args (list): other arguments.
+ """
+ self.num_sync_devices = num_sync_devices
+ if self.num_sync_devices > 0:
+ assert du.get_local_size() % self.num_sync_devices == 0, (
+ du.get_local_size(),
+ self.num_sync_devices,
+ )
+ self.num_groups = du.get_local_size() // self.num_sync_devices
+ else:
+ self.num_sync_devices = du.get_local_size()
+ self.num_groups = 1
+ super(NaiveSyncBatchNorm3d, self).__init__(**args)
+
+ def forward(self, input):
+ if du.get_local_size() == 1 or not self.training:
+ return super().forward(input)
+
+ assert input.shape[0] > 0, "SyncBatchNorm does not support empty inputs"
+ C = input.shape[1]
+ mean = torch.mean(input, dim=[0, 2, 3, 4])
+ meansqr = torch.mean(input * input, dim=[0, 2, 3, 4])
+
+ vec = torch.cat([mean, meansqr], dim=0)
+ vec = GroupGather.apply(vec, self.num_sync_devices, self.num_groups) * (
+ 1.0 / self.num_sync_devices
+ )
+
+ mean, meansqr = torch.split(vec, C)
+ var = meansqr - mean * mean
+ self.running_mean += self.momentum * (mean.detach() - self.running_mean)
+ self.running_var += self.momentum * (var.detach() - self.running_var)
+
+ invstd = torch.rsqrt(var + self.eps)
+ scale = self.weight * invstd
+ bias = self.bias - mean * scale
+ scale = scale.reshape(1, -1, 1, 1, 1)
+ bias = bias.reshape(1, -1, 1, 1, 1)
+ return input * scale + bias
diff --git a/clean/video/pwtf_dvd/inference/slowfast/models/build.py b/clean/video/pwtf_dvd/inference/slowfast/models/build.py
new file mode 100644
index 0000000000000000000000000000000000000000..8dd9cca224b3e77bb8c3cd899489358737d4cd35
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/models/build.py
@@ -0,0 +1,53 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Model construction functions."""
+
+import torch
+from fvcore.common.registry import Registry
+
+MODEL_REGISTRY = Registry("MODEL")
+MODEL_REGISTRY.__doc__ = """
+Registry for video model.
+
+The registered object will be called with `obj(cfg)`.
+The call should return a `torch.nn.Module` object.
+"""
+
+
+def build_model(cfg, gpu_id=None):
+ """
+ Builds the video model.
+ Args:
+ cfg (configs): configs that contains the hyper-parameters to build the
+ backbone. Details can be seen in slowfast/config/defaults.py.
+ gpu_id (Optional[int]): specify the gpu index to build model.
+ """
+ if torch.cuda.is_available():
+ assert (
+ cfg.NUM_GPUS <= torch.cuda.device_count()
+ ), "Cannot use more GPU devices than available"
+ else:
+ assert (
+ cfg.NUM_GPUS == 0
+ ), "Cuda is not available. Please set `NUM_GPUS: 0 for running on CPUs."
+
+ # Construct the model
+ name = cfg.MODEL.MODEL_NAME
+ model = MODEL_REGISTRY.get(name)(cfg)
+
+ if cfg.NUM_GPUS:
+ if gpu_id is None:
+ # Determine the GPU used by the current process
+ cur_device = torch.cuda.current_device()
+ else:
+ cur_device = gpu_id
+ # Transfer the model to the current GPU device
+ model = model.cuda(device=cur_device)
+ # Use multi-process data parallel model in the multi-gpu setting
+ if cfg.NUM_GPUS > 1:
+ # Make model replica operate on the current device
+ model = torch.nn.parallel.DistributedDataParallel(
+ module=model, device_ids=[cur_device], output_device=cur_device,find_unused_parameters=True
+ )
+ return model
diff --git a/clean/video/pwtf_dvd/inference/slowfast/models/custom_video_model_builder.py b/clean/video/pwtf_dvd/inference/slowfast/models/custom_video_model_builder.py
new file mode 100644
index 0000000000000000000000000000000000000000..f261f67b95616b8582b10998a290611ee108b2a9
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/models/custom_video_model_builder.py
@@ -0,0 +1,5 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+
+"""A More Flexible Video models."""
diff --git a/clean/video/pwtf_dvd/inference/slowfast/models/head_helper.py b/clean/video/pwtf_dvd/inference/slowfast/models/head_helper.py
new file mode 100644
index 0000000000000000000000000000000000000000..df04b010430b6000005676d52174243383873d05
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/models/head_helper.py
@@ -0,0 +1,95 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""ResNe(X)t Head helper."""
+
+import torch
+import torch.nn as nn
+
+class ResNetBasicHead(nn.Module):
+ """
+ ResNe(X)t 3D head.
+ This layer performs a fully-connected projection during training, when the
+ input size is 1x1x1. It performs a convolutional projection during testing
+ when the input size is larger than 1x1x1. If the inputs are from multiple
+ different pathways, the inputs will be concatenated after pooling.
+ """
+
+ def __init__(
+ self,
+ dim_in,
+ num_classes,
+ pool_size,
+ dropout_rate=0.0,
+ act_func="softmax",
+ ):
+ """
+ The `__init__` method of any subclass should also contain these
+ arguments.
+ ResNetBasicHead takes p pathways as input where p in [1, infty].
+
+ Args:
+ dim_in (list): the list of channel dimensions of the p inputs to the
+ ResNetHead.
+ num_classes (int): the channel dimensions of the p outputs to the
+ ResNetHead.
+ pool_size (list): the list of kernel sizes of p spatial temporal
+ poolings, temporal pool kernel size, spatial pool kernel size,
+ spatial pool kernel size in order.
+ dropout_rate (float): dropout rate. If equal to 0.0, perform no
+ dropout.
+ act_func (string): activation function to use. 'softmax': applies
+ softmax on the output. 'sigmoid': applies sigmoid on the output.
+ """
+ super(ResNetBasicHead, self).__init__()
+ assert (
+ len({len(pool_size), len(dim_in)}) == 1
+ ), "pathway dimensions are not consistent."
+ self.num_pathways = len(pool_size)
+
+ for pathway in range(self.num_pathways):
+ if pool_size[pathway] is None:
+ avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
+ else:
+ avg_pool = nn.AvgPool3d(pool_size[pathway], stride=1)
+ self.add_module("pathway{}_avgpool".format(pathway), avg_pool)
+
+ if dropout_rate > 0.0:
+ self.dropout = nn.Dropout(dropout_rate)
+ # Perform FC in a fully convolutional manner. The FC layer will be
+ # initialized with a different std comparing to convolutional layers.
+ self.projection = nn.Linear(sum(dim_in), num_classes, bias=True)
+
+ # Softmax for evaluation and testing.
+ if act_func == "softmax":
+ self.act = nn.Softmax(dim=4)
+ elif act_func == "sigmoid":
+ self.act = nn.Sigmoid()
+ else:
+ raise NotImplementedError(
+ "{} is not supported as an activation"
+ "function.".format(act_func)
+ )
+
+ def forward(self, inputs):
+ assert (
+ len(inputs) == self.num_pathways
+ ), "Input tensor does not contain {} pathway".format(self.num_pathways)
+ pool_out = []
+ for pathway in range(self.num_pathways):
+ m = getattr(self, "pathway{}_avgpool".format(pathway))
+ pool_out.append(m(inputs[pathway]))
+ x = torch.cat(pool_out, 1)
+ # (N, C, T, H, W) -> (N, T, H, W, C).
+ x = x.permute((0, 2, 3, 4, 1))
+ # Perform dropout.
+ if hasattr(self, "dropout"):
+ x = self.dropout(x)
+ x = self.projection(x)
+
+ # Performs fully convlutional inference.
+ # if not self.training:
+ # x = x.mean([1, 2, 3])
+ x = self.act(x)
+ x = x.view(x.shape[0], -1)
+ return x
diff --git a/clean/video/pwtf_dvd/inference/slowfast/models/losses.py b/clean/video/pwtf_dvd/inference/slowfast/models/losses.py
new file mode 100644
index 0000000000000000000000000000000000000000..7dda4eb19b2cf76275ba1778dc5f8730058a6c31
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/models/losses.py
@@ -0,0 +1,23 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Loss functions."""
+
+import torch.nn as nn
+
+_LOSSES = {
+ "cross_entropy": nn.CrossEntropyLoss,
+ "bce": nn.BCELoss,
+ "bce_logit": nn.BCEWithLogitsLoss,
+}
+
+
+def get_loss_func(loss_name):
+ """
+ Retrieve the loss given the loss name.
+ Args (int):
+ loss_name: the name of the loss to use.
+ """
+ if loss_name not in _LOSSES.keys():
+ raise NotImplementedError("Loss {} is not supported".format(loss_name))
+ return _LOSSES[loss_name]
diff --git a/clean/video/pwtf_dvd/inference/slowfast/models/nonlocal_helper.py b/clean/video/pwtf_dvd/inference/slowfast/models/nonlocal_helper.py
new file mode 100644
index 0000000000000000000000000000000000000000..6e68d05817256a66d0b6ecf0f96292446cf41270
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/models/nonlocal_helper.py
@@ -0,0 +1,148 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Non-local helper"""
+
+import torch
+import torch.nn as nn
+
+
+class Nonlocal(nn.Module):
+ """
+ Builds Non-local Neural Networks as a generic family of building
+ blocks for capturing long-range dependencies. Non-local Network
+ computes the response at a position as a weighted sum of the
+ features at all positions. This building block can be plugged into
+ many computer vision architectures.
+ More details in the paper: https://arxiv.org/pdf/1711.07971.pdf
+ """
+
+ def __init__(
+ self,
+ dim,
+ dim_inner,
+ pool_size=None,
+ instantiation="softmax",
+ zero_init_final_conv=False,
+ zero_init_final_norm=True,
+ norm_eps=1e-5,
+ norm_momentum=0.1,
+ norm_module=nn.BatchNorm3d,
+ ):
+ """
+ Args:
+ dim (int): number of dimension for the input.
+ dim_inner (int): number of dimension inside of the Non-local block.
+ pool_size (list): the kernel size of spatial temporal pooling,
+ temporal pool kernel size, spatial pool kernel size, spatial
+ pool kernel size in order. By default pool_size is None,
+ then there would be no pooling used.
+ instantiation (string): supports two different instantiation method:
+ "dot_product": normalizing correlation matrix with L2.
+ "softmax": normalizing correlation matrix with Softmax.
+ zero_init_final_conv (bool): If true, zero initializing the final
+ convolution of the Non-local block.
+ zero_init_final_norm (bool):
+ If true, zero initializing the final batch norm of the Non-local
+ block.
+ norm_module (nn.Module): nn.Module for the normalization layer. The
+ default is nn.BatchNorm3d.
+ """
+ super(Nonlocal, self).__init__()
+ self.dim = dim
+ self.dim_inner = dim_inner
+ self.pool_size = pool_size
+ self.instantiation = instantiation
+ self.use_pool = (
+ False
+ if pool_size is None
+ else any((size > 1 for size in pool_size))
+ )
+ self.norm_eps = norm_eps
+ self.norm_momentum = norm_momentum
+ self._construct_nonlocal(
+ zero_init_final_conv, zero_init_final_norm, norm_module
+ )
+
+ def _construct_nonlocal(
+ self, zero_init_final_conv, zero_init_final_norm, norm_module
+ ):
+ # Three convolution heads: theta, phi, and g.
+ self.conv_theta = nn.Conv3d(
+ self.dim, self.dim_inner, kernel_size=1, stride=1, padding=0
+ )
+ self.conv_phi = nn.Conv3d(
+ self.dim, self.dim_inner, kernel_size=1, stride=1, padding=0
+ )
+ self.conv_g = nn.Conv3d(
+ self.dim, self.dim_inner, kernel_size=1, stride=1, padding=0
+ )
+
+ # Final convolution output.
+ self.conv_out = nn.Conv3d(
+ self.dim_inner, self.dim, kernel_size=1, stride=1, padding=0
+ )
+ # Zero initializing the final convolution output.
+ self.conv_out.zero_init = zero_init_final_conv
+
+ # TODO: change the name to `norm`
+ self.bn = norm_module(
+ num_features=self.dim,
+ eps=self.norm_eps,
+ momentum=self.norm_momentum,
+ )
+ # Zero initializing the final bn.
+ self.bn.transform_final_bn = zero_init_final_norm
+
+ # Optional to add the spatial-temporal pooling.
+ if self.use_pool:
+ self.pool = nn.MaxPool3d(
+ kernel_size=self.pool_size,
+ stride=self.pool_size,
+ padding=[0, 0, 0],
+ )
+
+ def forward(self, x):
+ x_identity = x
+ N, C, T, H, W = x.size()
+
+ theta = self.conv_theta(x)
+
+ # Perform temporal-spatial pooling to reduce the computation.
+ if self.use_pool:
+ x = self.pool(x)
+
+ phi = self.conv_phi(x)
+ g = self.conv_g(x)
+
+ theta = theta.view(N, self.dim_inner, -1)
+ phi = phi.view(N, self.dim_inner, -1)
+ g = g.view(N, self.dim_inner, -1)
+
+ # (N, C, TxHxW) * (N, C, TxHxW) => (N, TxHxW, TxHxW).
+ theta_phi = torch.einsum("nct,ncp->ntp", (theta, phi))
+ # For original Non-local paper, there are two main ways to normalize
+ # the affinity tensor:
+ # 1) Softmax normalization (norm on exp).
+ # 2) dot_product normalization.
+ if self.instantiation == "softmax":
+ # Normalizing the affinity tensor theta_phi before softmax.
+ theta_phi = theta_phi * (self.dim_inner ** -0.5)
+ theta_phi = nn.functional.softmax(theta_phi, dim=2)
+ elif self.instantiation == "dot_product":
+ spatial_temporal_dim = theta_phi.shape[2]
+ theta_phi = theta_phi / spatial_temporal_dim
+ else:
+ raise NotImplementedError(
+ "Unknown norm type {}".format(self.instantiation)
+ )
+
+ # (N, TxHxW, TxHxW) * (N, C, TxHxW) => (N, C, TxHxW).
+ theta_phi_g = torch.einsum("ntg,ncg->nct", (theta_phi, g))
+
+ # (N, C, TxHxW) => (N, C, T, H, W).
+ theta_phi_g = theta_phi_g.view(N, self.dim_inner, T, H, W)
+
+ p = self.conv_out(theta_phi_g)
+ p = self.bn(p)
+ return x_identity + p
diff --git a/clean/video/pwtf_dvd/inference/slowfast/models/optimizer.py b/clean/video/pwtf_dvd/inference/slowfast/models/optimizer.py
new file mode 100644
index 0000000000000000000000000000000000000000..130f2cebf994741bc45a6519f07c5f0740c106ac
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/models/optimizer.py
@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Optimizer."""
+
+import torch
+
+import slowfast.utils.lr_policy as lr_policy
+
+
+def construct_optimizer(model, cfg):
+ """
+ Construct a stochastic gradient descent or ADAM optimizer with momentum.
+ Details can be found in:
+ Herbert Robbins, and Sutton Monro. "A stochastic approximation method."
+ and
+ Diederik P.Kingma, and Jimmy Ba.
+ "Adam: A Method for Stochastic Optimization."
+
+ Args:
+ model (model): model to perform stochastic gradient descent
+ optimization or ADAM optimization.
+ cfg (config): configs of hyper-parameters of SGD or ADAM, includes base
+ learning rate, momentum, weight_decay, dampening, and etc.
+ """
+ # Batchnorm parameters.
+ bn_params = []
+ # Non-batchnorm parameters.
+ non_bn_parameters = []
+ for name, p in model.named_parameters():
+ if "bn" in name:
+ bn_params.append(p)
+ else:
+ non_bn_parameters.append(p)
+ # Apply different weight decay to Batchnorm and non-batchnorm parameters.
+ # In Caffe2 classification codebase the weight decay for batchnorm is 0.0.
+ # Having a different weight decay on batchnorm might cause a performance
+ # drop.
+ optim_params = [
+ {"params": bn_params, "weight_decay": cfg.BN.WEIGHT_DECAY},
+ {"params": non_bn_parameters, "weight_decay": cfg.SOLVER.WEIGHT_DECAY},
+ ]
+ # Check all parameters will be passed into optimizer.
+ assert len(list(model.parameters())) == len(non_bn_parameters) + len(
+ bn_params
+ ), "parameter size does not match: {} + {} != {}".format(
+ len(non_bn_parameters), len(bn_params), len(list(model.parameters()))
+ )
+
+ if cfg.SOLVER.OPTIMIZING_METHOD == "sgd":
+ return torch.optim.SGD(
+ optim_params,
+ lr=cfg.SOLVER.BASE_LR,
+ momentum=cfg.SOLVER.MOMENTUM,
+ weight_decay=cfg.SOLVER.WEIGHT_DECAY,
+ dampening=cfg.SOLVER.DAMPENING,
+ nesterov=cfg.SOLVER.NESTEROV,
+ )
+ elif cfg.SOLVER.OPTIMIZING_METHOD == "adam":
+ return torch.optim.Adam(
+ optim_params,
+ lr=cfg.SOLVER.BASE_LR,
+ betas=(0.9, 0.999),
+ weight_decay=cfg.SOLVER.WEIGHT_DECAY,
+ )
+ else:
+ raise NotImplementedError(
+ "Does not support {} optimizer".format(cfg.SOLVER.OPTIMIZING_METHOD)
+ )
+
+
+def get_epoch_lr(cur_epoch, cfg):
+ """
+ Retrieves the lr for the given epoch (as specified by the lr policy).
+ Args:
+ cfg (config): configs of hyper-parameters of ADAM, includes base
+ learning rate, betas, and weight decays.
+ cur_epoch (float): the number of epoch of the current training stage.
+ """
+ return lr_policy.get_lr_at_epoch(cfg, cur_epoch)
+
+def get_iter_lr(cur_iter, cfg):
+ """
+ Retrieves the lr for the given iter (as specified by the lr policy).
+ Args:
+ cfg (config): configs of hyper-parameters of ADAM, includes base
+ learning rate, betas, and weight decays.
+ cur_epoch (float): the number of epoch of the current training stage.
+ """
+ lr=lr_policy.get_lr_at_iter(cfg, cur_iter)
+
+ return lr
+
+
+def set_lr(optimizer, new_lr):
+ """
+ Sets the optimizer lr to the specified value.
+ Args:
+ optimizer (optim): the optimizer using to optimize the current network.
+ new_lr (float): the new learning rate to set.
+ """
+ for param_group in optimizer.param_groups:
+ param_group["lr"] = new_lr
diff --git a/clean/video/pwtf_dvd/inference/slowfast/models/resnet_helper.py b/clean/video/pwtf_dvd/inference/slowfast/models/resnet_helper.py
new file mode 100644
index 0000000000000000000000000000000000000000..3e10df460ce94d2c816ab546188412989a4ed628
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/models/resnet_helper.py
@@ -0,0 +1,649 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Video models."""
+
+import torch.nn as nn
+
+from slowfast.models.nonlocal_helper import Nonlocal
+
+
+def get_trans_func(name):
+ """
+ Retrieves the transformation module by name.
+ """
+ trans_funcs = {
+ "bottleneck_transform": BottleneckTransform,
+ "basic_transform": BasicTransform,
+ "temporal_transform":TemporalTransform
+ }
+ assert (
+ name in trans_funcs.keys()
+ ), "Transformation function '{}' not supported".format(name)
+ return trans_funcs[name]
+
+
+class BasicTransform(nn.Module):
+ """
+ Basic transformation: Tx3x3, 1x3x3, where T is the size of temporal kernel.
+ """
+
+ def __init__(
+ self,
+ dim_in,
+ dim_out,
+ temp_kernel_size,
+ stride,
+ dim_inner=None,
+ num_groups=1,
+ stride_1x1=None,
+ inplace_relu=True,
+ eps=1e-5,
+ bn_mmt=0.1,
+ norm_module=nn.BatchNorm3d,
+ ):
+ """
+ Args:
+ dim_in (int): the channel dimensions of the input.
+ dim_out (int): the channel dimension of the output.
+ temp_kernel_size (int): the temporal kernel sizes of the first
+ convolution in the basic block.
+ stride (int): the stride of the bottleneck.
+ dim_inner (None): the inner dimension would not be used in
+ BasicTransform.
+ num_groups (int): number of groups for the convolution. Number of
+ group is always 1 for BasicTransform.
+ stride_1x1 (None): stride_1x1 will not be used in BasicTransform.
+ inplace_relu (bool): if True, calculate the relu on the original
+ input without allocating new memory.
+ eps (float): epsilon for batch norm.
+ bn_mmt (float): momentum for batch norm. Noted that BN momentum in
+ PyTorch = 1 - BN momentum in Caffe2.
+ norm_module (nn.Module): nn.Module for the normalization layer. The
+ default is nn.BatchNorm3d.
+ """
+ super(BasicTransform, self).__init__()
+ self.temp_kernel_size = temp_kernel_size
+ self._inplace_relu = inplace_relu
+ self._eps = eps
+ self._bn_mmt = bn_mmt
+ self._construct(dim_in, dim_out, stride, norm_module)
+
+ def _construct(self, dim_in, dim_out, stride, norm_module):
+ # Tx3x3, BN, ReLU.
+ self.a = nn.Conv3d(
+ dim_in,
+ dim_out,
+ kernel_size=[self.temp_kernel_size, 3, 3],
+ stride=[1, stride, stride],
+ padding=[int(self.temp_kernel_size // 2), 1, 1],
+ bias=False,
+ )
+ self.a_bn = norm_module(
+ num_features=dim_out, eps=self._eps, momentum=self._bn_mmt
+ )
+ self.a_relu = nn.ReLU(inplace=self._inplace_relu)
+ # 1x3x3, BN.
+ self.b = nn.Conv3d(
+ dim_out,
+ dim_out,
+ kernel_size=[1, 3, 3],
+ stride=[1, 1, 1],
+ padding=[0, 1, 1],
+ bias=False,
+ )
+ self.b_bn = norm_module(
+ num_features=dim_out, eps=self._eps, momentum=self._bn_mmt
+ )
+
+ self.b_bn.transform_final_bn = True
+
+ def forward(self, x):
+ x = self.a(x)
+ x = self.a_bn(x)
+ x = self.a_relu(x)
+
+ x = self.b(x)
+ x = self.b_bn(x)
+ return x
+
+class TemporalTransform(nn.Module):
+ """
+ Basic transformation: Tx3x3, 1x3x3, where T is the size of temporal kernel.
+ """
+
+ def __init__(
+ self,
+ dim_in,
+ dim_out,
+ temp_kernel_size,
+ stride,
+ dim_inner=None,
+ num_groups=1,
+ stride_1x1=None,
+ inplace_relu=True,
+ eps=1e-5,
+ bn_mmt=0.1,
+ norm_module=nn.BatchNorm3d,
+ dilation=1
+ ):
+ """
+ Args:
+ dim_in (int): the channel dimensions of the input.
+ dim_out (int): the channel dimension of the output.
+ temp_kernel_size (int): the temporal kernel sizes of the first
+ convolution in the basic block.
+ stride (int): the stride of the bottleneck.
+ dim_inner (None): the inner dimension would not be used in
+ BasicTransform.
+ num_groups (int): number of groups for the convolution. Number of
+ group is always 1 for BasicTransform.
+ stride_1x1 (None): stride_1x1 will not be used in BasicTransform.
+ inplace_relu (bool): if True, calculate the relu on the original
+ input without allocating new memory.
+ eps (float): epsilon for batch norm.
+ bn_mmt (float): momentum for batch norm. Noted that BN momentum in
+ PyTorch = 1 - BN momentum in Caffe2.
+ norm_module (nn.Module): nn.Module for the normalization layer. The
+ default is nn.BatchNorm3d.
+ """
+ super(TemporalTransform, self).__init__()
+ self.temp_kernel_size = temp_kernel_size
+ self._inplace_relu = inplace_relu
+ self._eps = eps
+ self._bn_mmt = bn_mmt
+ self._construct(dim_in, dim_out, stride, norm_module)
+
+ def _construct(self, dim_in, dim_out, stride, norm_module):
+ # Tx3x3, BN, ReLU.
+ self.a = nn.Conv3d(
+ dim_in,
+ dim_out,
+ kernel_size=[self.temp_kernel_size, 3, 3],
+ stride=[1, stride, stride],
+ padding=[int(self.temp_kernel_size // 2), 1, 1],
+ bias=False,
+ )
+ self.a_bn = norm_module(
+ num_features=dim_out, eps=self._eps, momentum=self._bn_mmt
+ )
+ self.a_relu = nn.ReLU(inplace=self._inplace_relu)
+ # 1x3x3, BN.
+ self.b = nn.Conv3d(
+ dim_out,
+ dim_out,
+ kernel_size=[1, 3, 3],
+ stride=[1, 1, 1],
+ padding=[0, 1, 1],
+ bias=False,
+ )
+ self.b_bn = norm_module(
+ num_features=dim_out, eps=self._eps, momentum=self._bn_mmt
+ )
+
+ self.b_bn.transform_final_bn = True
+
+ def forward(self, x):
+ x = self.a(x)
+ x = self.a_bn(x)
+ x = self.a_relu(x)
+
+ x = self.b(x)
+ x = self.b_bn(x)
+ return x
+
+
+class BottleneckTransform(nn.Module):
+ """
+ Bottleneck transformation: Tx1x1, 1x3x3, 1x1x1, where T is the size of
+ temporal kernel.
+ """
+
+ def __init__(
+ self,
+ dim_in,
+ dim_out,
+ temp_kernel_size,
+ stride,
+ dim_inner,
+ num_groups,
+ stride_1x1=False,
+ inplace_relu=True,
+ eps=1e-5,
+ bn_mmt=0.1,
+ dilation=1,
+ norm_module=nn.BatchNorm3d,
+ ):
+ """
+ Args:
+ dim_in (int): the channel dimensions of the input.
+ dim_out (int): the channel dimension of the output.
+ temp_kernel_size (int): the temporal kernel sizes of the first
+ convolution in the bottleneck.
+ stride (int): the stride of the bottleneck.
+ dim_inner (int): the inner dimension of the block.
+ num_groups (int): number of groups for the convolution. num_groups=1
+ is for standard ResNet like networks, and num_groups>1 is for
+ ResNeXt like networks.
+ stride_1x1 (bool): if True, apply stride to 1x1 conv, otherwise
+ apply stride to the 3x3 conv.
+ inplace_relu (bool): if True, calculate the relu on the original
+ input without allocating new memory.
+ eps (float): epsilon for batch norm.
+ bn_mmt (float): momentum for batch norm. Noted that BN momentum in
+ PyTorch = 1 - BN momentum in Caffe2.
+ dilation (int): size of dilation.
+ norm_module (nn.Module): nn.Module for the normalization layer. The
+ default is nn.BatchNorm3d.
+ """
+ super(BottleneckTransform, self).__init__()
+ self.temp_kernel_size = temp_kernel_size
+ self._inplace_relu = inplace_relu
+ self._eps = eps
+ self._bn_mmt = bn_mmt
+ self._stride_1x1 = stride_1x1
+ self._construct(
+ dim_in,
+ dim_out,
+ stride,
+ dim_inner,
+ num_groups,
+ dilation,
+ norm_module,
+ )
+
+ def _construct(
+ self,
+ dim_in,
+ dim_out,
+ stride,
+ dim_inner,
+ num_groups,
+ dilation,
+ norm_module,
+ ):
+ (str1x1, str3x3) = (stride, 1) if self._stride_1x1 else (1, stride)
+
+ # Tx1x1, BN, ReLU.
+ self.a = nn.Conv3d(
+ dim_in,
+ dim_inner,
+ kernel_size=[self.temp_kernel_size, 1, 1],
+ stride=[1, str1x1, str1x1],
+ padding=[int(self.temp_kernel_size // 2), 0, 0],
+ bias=False,
+ )
+ self.a_bn = norm_module(
+ num_features=dim_inner, eps=self._eps, momentum=self._bn_mmt
+ )
+ self.a_relu = nn.ReLU(inplace=self._inplace_relu)
+
+ # 1x3x3, BN, ReLU.
+ self.b = nn.Conv3d(
+ dim_inner,
+ dim_inner,
+ [1, 3, 3],
+ stride=[1, str3x3, str3x3],
+ padding=[0, dilation, dilation],
+ groups=num_groups,
+ bias=False,
+ dilation=[1, dilation, dilation],
+ )
+ self.b_bn = norm_module(
+ num_features=dim_inner, eps=self._eps, momentum=self._bn_mmt
+ )
+ self.b_relu = nn.ReLU(inplace=self._inplace_relu)
+
+ # 1x1x1, BN.
+ self.c = nn.Conv3d(
+ dim_inner,
+ dim_out,
+ kernel_size=[1, 1, 1],
+ stride=[1, 1, 1],
+ padding=[0, 0, 0],
+ bias=False,
+ )
+ self.c_bn = norm_module(
+ num_features=dim_out, eps=self._eps, momentum=self._bn_mmt
+ )
+ self.c_bn.transform_final_bn = True
+
+ def forward(self, x):
+ # Explicitly forward every layer.
+ # Branch2a.
+ x = self.a(x)
+ x = self.a_bn(x)
+ x = self.a_relu(x)
+
+ # Branch2b.
+ x = self.b(x)
+ x = self.b_bn(x)
+ x = self.b_relu(x)
+
+ # Branch2c
+ x = self.c(x)
+ x = self.c_bn(x)
+ return x
+
+
+class ResBlock(nn.Module):
+ """
+ Residual block.
+ """
+
+ def __init__(
+ self,
+ dim_in,
+ dim_out,
+ temp_kernel_size,
+ stride,
+ trans_func,
+ dim_inner,
+ num_groups=1,
+ stride_1x1=False,
+ inplace_relu=True, # default :: False
+ # ################### Fixed by TH for training FT###############################
+ eps=1e-5,
+ bn_mmt=0.1,
+ dilation=1,
+ norm_module=nn.BatchNorm3d,
+ ):
+ """
+ ResBlock class constructs redisual blocks. More details can be found in:
+ Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.
+ "Deep residual learning for image recognition."
+ https://arxiv.org/abs/1512.03385
+ Args:
+ dim_in (int): the channel dimensions of the input.
+ dim_out (int): the channel dimension of the output.
+ temp_kernel_size (int): the temporal kernel sizes of the middle
+ convolution in the bottleneck.
+ stride (int): the stride of the bottleneck.
+ trans_func (string): transform function to be used to construct the
+ bottleneck.
+ dim_inner (int): the inner dimension of the block.
+ num_groups (int): number of groups for the convolution. num_groups=1
+ is for standard ResNet like networks, and num_groups>1 is for
+ ResNeXt like networks.
+ stride_1x1 (bool): if True, apply stride to 1x1 conv, otherwise
+ apply stride to the 3x3 conv.
+ inplace_relu (bool): calculate the relu on the original input
+ without allocating new memory.
+ eps (float): epsilon for batch norm.
+ bn_mmt (float): momentum for batch norm. Noted that BN momentum in
+ PyTorch = 1 - BN momentum in Caffe2.
+ dilation (int): size of dilation.
+ norm_module (nn.Module): nn.Module for the normalization layer. The
+ default is nn.BatchNorm3d.
+ """
+ super(ResBlock, self).__init__()
+ self._inplace_relu = inplace_relu
+ self._eps = eps
+ self._bn_mmt = bn_mmt
+ self._construct(
+ dim_in,
+ dim_out,
+ temp_kernel_size,
+ stride,
+ trans_func,
+ dim_inner,
+ num_groups,
+ stride_1x1,
+ inplace_relu,
+ dilation,
+ norm_module,
+ )
+
+ def _construct(
+ self,
+ dim_in,
+ dim_out,
+ temp_kernel_size,
+ stride,
+ trans_func,
+ dim_inner,
+ num_groups,
+ stride_1x1,
+ inplace_relu,
+ dilation,
+ norm_module,
+ ):
+ # Use skip connection with projection if dim or res change.
+ if (dim_in != dim_out) or (stride != 1):
+ self.branch1 = nn.Conv3d(
+ dim_in,
+ dim_out,
+ kernel_size=1,
+ stride=[1, stride, stride],
+ padding=0,
+ bias=False,
+ dilation=1,
+ )
+ self.branch1_bn = norm_module(
+ num_features=dim_out, eps=self._eps, momentum=self._bn_mmt
+ )
+ self.branch2 = trans_func(
+ dim_in,
+ dim_out,
+ temp_kernel_size,
+ stride,
+ dim_inner,
+ num_groups,
+ stride_1x1=stride_1x1,
+ inplace_relu=inplace_relu,
+ dilation=dilation,
+ norm_module=norm_module,
+ )
+ self.relu = nn.ReLU(self._inplace_relu)
+
+ def forward(self, x):
+ if hasattr(self, "branch1"):
+ x = self.branch1_bn(self.branch1(x)) + self.branch2(x)
+ else:
+ x = x + self.branch2(x)
+
+ x = self.relu(x)
+ return x
+
+
+class ResStage(nn.Module):
+ """
+ Stage of 3D ResNet. It expects to have one or more tensors as input for
+ single pathway (C2D, I3D, Slow), and multi-pathway (SlowFast) cases.
+ More details can be found here:
+
+ Christoph Feichtenhofer, Haoqi Fan, Jitendra Malik, and Kaiming He.
+ "SlowFast networks for video recognition."
+ https://arxiv.org/pdf/1812.03982.pdf
+ """
+
+ def __init__(
+ self,
+ dim_in,
+ dim_out,
+ stride,
+ temp_kernel_sizes,
+ num_blocks,
+ dim_inner,
+ num_groups,
+ num_block_temp_kernel,
+ nonlocal_inds,
+ nonlocal_group,
+ nonlocal_pool,
+ dilation,
+ instantiation="softmax",
+ trans_func_name="bottleneck_transform",
+ stride_1x1=False,
+ inplace_relu=True, ## Fixed by TH
+ norm_module=nn.BatchNorm3d,
+ ):
+ """
+ The `__init__` method of any subclass should also contain these arguments.
+ ResStage builds p streams, where p can be greater or equal to one.
+ Args:
+ dim_in (list): list of p the channel dimensions of the input.
+ Different channel dimensions control the input dimension of
+ different pathways.
+ dim_out (list): list of p the channel dimensions of the output.
+ Different channel dimensions control the input dimension of
+ different pathways.
+ temp_kernel_sizes (list): list of the p temporal kernel sizes of the
+ convolution in the bottleneck. Different temp_kernel_sizes
+ control different pathway.
+ stride (list): list of the p strides of the bottleneck. Different
+ stride control different pathway.
+ num_blocks (list): list of p numbers of blocks for each of the
+ pathway.
+ dim_inner (list): list of the p inner channel dimensions of the
+ input. Different channel dimensions control the input dimension
+ of different pathways.
+ num_groups (list): list of number of p groups for the convolution.
+ num_groups=1 is for standard ResNet like networks, and
+ num_groups>1 is for ResNeXt like networks.
+ num_block_temp_kernel (list): extent the temp_kernel_sizes to
+ num_block_temp_kernel blocks, then fill temporal kernel size
+ of 1 for the rest of the layers.
+ nonlocal_inds (list): If the tuple is empty, no nonlocal layer will
+ be added. If the tuple is not empty, add nonlocal layers after
+ the index-th block.
+ dilation (list): size of dilation for each pathway.
+ nonlocal_group (list): list of number of p nonlocal groups. Each
+ number controls how to fold temporal dimension to batch
+ dimension before applying nonlocal transformation.
+ https://github.com/facebookresearch/video-nonlocal-net.
+ instantiation (string): different instantiation for nonlocal layer.
+ Supports two different instantiation method:
+ "dot_product": normalizing correlation matrix with L2.
+ "softmax": normalizing correlation matrix with Softmax.
+ trans_func_name (string): name of the the transformation function apply
+ on the network.
+ norm_module (nn.Module): nn.Module for the normalization layer. The
+ default is nn.BatchNorm3d.
+ """
+ super(ResStage, self).__init__()
+ assert all(
+ (
+ num_block_temp_kernel[i] <= num_blocks[i]
+ for i in range(len(temp_kernel_sizes))
+ )
+ )
+ self.num_blocks = num_blocks
+ self.nonlocal_group = nonlocal_group
+ self.temp_kernel_sizes = [
+ (temp_kernel_sizes[i] * num_blocks[i])[: num_block_temp_kernel[i]]
+ + [1] * (num_blocks[i] - num_block_temp_kernel[i])
+ for i in range(len(temp_kernel_sizes))
+ ]
+ assert (
+ len(
+ {
+ len(dim_in),
+ len(dim_out),
+ len(temp_kernel_sizes),
+ len(stride),
+ len(num_blocks),
+ len(dim_inner),
+ len(num_groups),
+ len(num_block_temp_kernel),
+ len(nonlocal_inds),
+ len(nonlocal_group),
+ }
+ )
+ == 1
+ )
+ self.num_pathways = len(self.num_blocks)
+ self._construct(
+ dim_in,
+ dim_out,
+ stride,
+ dim_inner,
+ num_groups,
+ trans_func_name,
+ stride_1x1,
+ inplace_relu,
+ nonlocal_inds,
+ nonlocal_pool,
+ instantiation,
+ dilation,
+ norm_module,
+ )
+
+ def _construct(
+ self,
+ dim_in,
+ dim_out,
+ stride,
+ dim_inner,
+ num_groups,
+ trans_func_name,
+ stride_1x1,
+ inplace_relu,
+ nonlocal_inds,
+ nonlocal_pool,
+ instantiation,
+ dilation,
+ norm_module,
+ ):
+ for pathway in range(self.num_pathways):
+ for i in range(self.num_blocks[pathway]):
+ # Retrieve the transformation function.
+ trans_func = get_trans_func(trans_func_name)
+ # Construct the block.
+ res_block = ResBlock(
+ dim_in[pathway] if i == 0 else dim_out[pathway],
+ dim_out[pathway],
+ self.temp_kernel_sizes[pathway][i],
+ stride[pathway] if i == 0 else 1,
+ trans_func,
+ dim_inner[pathway],
+ num_groups[pathway],
+ stride_1x1=stride_1x1,
+ inplace_relu=inplace_relu,
+ dilation=dilation[pathway],
+ norm_module=norm_module,
+ )
+ self.add_module("pathway{}_res{}".format(pathway, i), res_block)
+ if i in nonlocal_inds[pathway]:
+ nln = Nonlocal(
+ dim_out[pathway],
+ dim_out[pathway] // 2,
+ nonlocal_pool[pathway],
+ instantiation=instantiation,
+ norm_module=norm_module,
+ )
+ self.add_module(
+ "pathway{}_nonlocal{}".format(pathway, i), nln
+ )
+
+ def forward(self, inputs):
+ output = []
+ for pathway in range(self.num_pathways):
+ x = inputs[pathway]
+ for i in range(self.num_blocks[pathway]):
+ m = getattr(self, "pathway{}_res{}".format(pathway, i))
+ x = m(x)
+ if hasattr(self, "pathway{}_nonlocal{}".format(pathway, i)):
+ nln = getattr(
+ self, "pathway{}_nonlocal{}".format(pathway, i)
+ )
+ b, c, t, h, w = x.shape
+ if self.nonlocal_group[pathway] > 1:
+ # Fold temporal dimension into batch dimension.
+ x = x.permute(0, 2, 1, 3, 4)
+ x = x.reshape(
+ b * self.nonlocal_group[pathway],
+ t // self.nonlocal_group[pathway],
+ c,
+ h,
+ w,
+ )
+ x = x.permute(0, 2, 1, 3, 4)
+ x = nln(x)
+ if self.nonlocal_group[pathway] > 1:
+ # Fold back to temporal dimension.
+ x = x.permute(0, 2, 1, 3, 4)
+ x = x.reshape(b, t, c, h, w)
+ x = x.permute(0, 2, 1, 3, 4)
+ output.append(x)
+
+ return output
diff --git a/clean/video/pwtf_dvd/inference/slowfast/models/stem_helper.py b/clean/video/pwtf_dvd/inference/slowfast/models/stem_helper.py
new file mode 100644
index 0000000000000000000000000000000000000000..481977b15a13edf54bfdb17fd3627b6657d56262
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/models/stem_helper.py
@@ -0,0 +1,178 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""ResNe(X)t 3D stem helper."""
+
+import torch.nn as nn
+
+
+class VideoModelStem(nn.Module):
+ """
+ Video 3D stem module. Provides stem operations of Conv, BN, ReLU, MaxPool
+ on input data tensor for one or multiple pathways.
+ """
+
+ def __init__(
+ self,
+ dim_in,
+ dim_out,
+ kernel,
+ stride,
+ padding,
+ inplace_relu=True,
+ eps=1e-5,
+ bn_mmt=0.1,
+ norm_module=nn.BatchNorm3d,
+ ):
+ """
+ The `__init__` method of any subclass should also contain these
+ arguments. List size of 1 for single pathway models (C2D, I3D, Slow
+ and etc), list size of 2 for two pathway models (SlowFast).
+
+ Args:
+ dim_in (list): the list of channel dimensions of the inputs.
+ dim_out (list): the output dimension of the convolution in the stem
+ layer.
+ kernel (list): the kernels' size of the convolutions in the stem
+ layers. Temporal kernel size, height kernel size, width kernel
+ size in order.
+ stride (list): the stride sizes of the convolutions in the stem
+ layer. Temporal kernel stride, height kernel size, width kernel
+ size in order.
+ padding (list): the paddings' sizes of the convolutions in the stem
+ layer. Temporal padding size, height padding size, width padding
+ size in order.
+ inplace_relu (bool): calculate the relu on the original input
+ without allocating new memory.
+ eps (float): epsilon for batch norm.
+ bn_mmt (float): momentum for batch norm. Noted that BN momentum in
+ PyTorch = 1 - BN momentum in Caffe2.
+ norm_module (nn.Module): nn.Module for the normalization layer. The
+ default is nn.BatchNorm3d.
+ """
+ super(VideoModelStem, self).__init__()
+
+ assert (
+ len(
+ {
+ len(dim_in),
+ len(dim_out),
+ len(kernel),
+ len(stride),
+ len(padding),
+ }
+ )
+ == 1
+ ), "Input pathway dimensions are not consistent."
+ self.num_pathways = len(dim_in)
+ self.kernel = kernel
+ self.stride = stride
+ self.padding = padding
+ self.inplace_relu = inplace_relu
+ self.eps = eps
+ self.bn_mmt = bn_mmt
+ # Construct the stem layer.
+ self._construct_stem(dim_in, dim_out, norm_module)
+
+ def _construct_stem(self, dim_in, dim_out, norm_module):
+ for pathway in range(len(dim_in)):
+ stem = ResNetBasicStem(
+ dim_in[pathway],
+ dim_out[pathway],
+ self.kernel[pathway],
+ self.stride[pathway],
+ self.padding[pathway],
+ self.inplace_relu,
+ self.eps,
+ self.bn_mmt,
+ norm_module,
+ )
+ self.add_module("pathway{}_stem".format(pathway), stem)
+
+ def forward(self, x):
+ assert (
+ len(x) == self.num_pathways
+ ), "Input tensor does not contain {} pathway".format(self.num_pathways)
+ for pathway in range(len(x)):
+ m = getattr(self, "pathway{}_stem".format(pathway))
+ x[pathway] = m(x[pathway])
+ return x
+
+
+class ResNetBasicStem(nn.Module):
+ """
+ ResNe(X)t 3D stem module.
+ Performs spatiotemporal Convolution, BN, and Relu following by a
+ spatiotemporal pooling.
+ """
+
+ def __init__(
+ self,
+ dim_in,
+ dim_out,
+ kernel,
+ stride,
+ padding,
+ inplace_relu=True,
+ eps=1e-5,
+ bn_mmt=0.1,
+ norm_module=nn.BatchNorm3d,
+ ):
+ """
+ The `__init__` method of any subclass should also contain these arguments.
+
+ Args:
+ dim_in (int): the channel dimension of the input. Normally 3 is used
+ for rgb input, and 2 or 3 is used for optical flow input.
+ dim_out (int): the output dimension of the convolution in the stem
+ layer.
+ kernel (list): the kernel size of the convolution in the stem layer.
+ temporal kernel size, height kernel size, width kernel size in
+ order.
+ stride (list): the stride size of the convolution in the stem layer.
+ temporal kernel stride, height kernel size, width kernel size in
+ order.
+ padding (int): the padding size of the convolution in the stem
+ layer, temporal padding size, height padding size, width
+ padding size in order.
+ inplace_relu (bool): calculate the relu on the original input
+ without allocating new memory.
+ eps (float): epsilon for batch norm.
+ bn_mmt (float): momentum for batch norm. Noted that BN momentum in
+ PyTorch = 1 - BN momentum in Caffe2.
+ norm_module (nn.Module): nn.Module for the normalization layer. The
+ default is nn.BatchNorm3d.
+ """
+ super(ResNetBasicStem, self).__init__()
+ self.kernel = kernel
+ self.stride = stride
+ self.padding = padding
+ self.inplace_relu = inplace_relu
+ self.eps = eps
+ self.bn_mmt = bn_mmt
+ # Construct the stem layer.
+ self._construct_stem(dim_in, dim_out, norm_module)
+
+ def _construct_stem(self, dim_in, dim_out, norm_module):
+ self.conv = nn.Conv3d(
+ dim_in,
+ dim_out,
+ self.kernel,
+ stride=self.stride,
+ padding=self.padding,
+ bias=False,
+ )
+ self.bn = norm_module(
+ num_features=dim_out, eps=self.eps, momentum=self.bn_mmt
+ )
+ self.relu = nn.ReLU(self.inplace_relu)
+ self.pool_layer = nn.MaxPool3d(
+ kernel_size=[1, 3, 3], stride=[1, 2, 2], padding=[0, 1, 1]
+ )
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.bn(x)
+ x = self.relu(x)
+ x = self.pool_layer(x)
+ return x
diff --git a/clean/video/pwtf_dvd/inference/slowfast/models/unet_helper.py b/clean/video/pwtf_dvd/inference/slowfast/models/unet_helper.py
new file mode 100644
index 0000000000000000000000000000000000000000..36b7202cd1936a433b193017f6c363e5dd317b4f
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/models/unet_helper.py
@@ -0,0 +1,157 @@
+InPlaceABN = None
+from torch import nn
+import torch.nn.functional as F
+
+
+class Conv3dReLU(nn.Sequential):
+ def __init__(
+ self,
+ in_channels,
+ out_channels,
+ kernel_size,
+ padding=0,
+ stride=1,
+ use_batchnorm=True,
+ ):
+
+ if use_batchnorm == "inplace" and InPlaceABN is None:
+ raise RuntimeError(
+ "In order to use `use_batchnorm='inplace'` inplace_abn package must be installed. "
+ + "To install see: https://github.com/mapillary/inplace_abn"
+ )
+
+ conv = nn.Conv3d(
+ in_channels,
+ out_channels,
+ kernel_size,
+ stride=stride,
+ padding=padding,
+ bias=not (use_batchnorm),
+ )
+ relu = nn.ReLU(inplace=True)
+
+ if use_batchnorm == "inplace":
+ bn = InPlaceABN(out_channels, activation="leaky_relu", activation_param=0.0)
+ relu = nn.Identity()
+
+ elif use_batchnorm and use_batchnorm != "inplace":
+ bn = nn.BatchNorm3d(out_channels)
+
+ else:
+ bn = nn.Identity()
+
+ super(Conv3dReLU, self).__init__(conv, bn, relu)
+
+
+class DecoderBlock(nn.Module):
+ def __init__(
+ self, in_channels, skip_channels, out_channels, use_batchnorm=True,
+ ):
+ super().__init__()
+ self.conv1 = Conv3dReLU(
+ in_channels + skip_channels,
+ out_channels,
+ kernel_size=3,
+ padding=1,
+ use_batchnorm=use_batchnorm,
+ )
+
+ self.conv2 = Conv3dReLU(
+ out_channels,
+ out_channels,
+ kernel_size=3,
+ padding=1,
+ use_batchnorm=use_batchnorm,
+ )
+
+ def forward(self, x):
+ x = self.conv1(x)
+ x = self.conv2(x)
+ return x
+
+
+class LightDecoderBlock(nn.Module):
+ def __init__(
+ self, in_channels, skip_channels, out_channels, use_batchnorm=True,
+ ):
+ super().__init__()
+ self.conv1 = Conv3dReLU(
+ in_channels + skip_channels,
+ out_channels,
+ kernel_size=3,
+ padding=1,
+ use_batchnorm=use_batchnorm,
+ )
+
+ def forward(self, x):
+ x = self.conv1(x)
+ return x
+
+
+def freeze_net(model: nn.Module, freeze_prefixs):
+ flag = False
+ for name, param in model.named_parameters():
+ items = name.split(".")
+ if items[0] == "module":
+ prefix = items[1]
+ else:
+ prefix = items[0]
+ if prefix in freeze_prefixs:
+ if param.requires_grad is True:
+ param.requires_grad = False
+ flag = True
+ # print("freeze",name)
+
+ assert flag
+
+
+def unfreeze_net(model: nn.Module):
+ for name, param in model.named_parameters():
+ param.requires_grad = True
+
+
+from .resnet_helper import ResBlock, get_trans_func
+
+
+class ResDecoderBlock(nn.Module):
+ def __init__(
+ self, in_channels, skip_channels, out_channels, use_batchnorm=True,
+ ):
+ super().__init__()
+ trans_func = get_trans_func("bottleneck_transform")
+ self.conv1 = ResBlock(
+ in_channels + skip_channels,
+ out_channels,
+ 3,
+ 1,
+ trans_func,
+ out_channels//2,
+ num_groups=1,
+ stride_1x1=False,
+ inplace_relu=True,
+ eps=1e-5,
+ bn_mmt=0.1,
+ dilation=1,
+ norm_module=nn.BatchNorm3d,
+ )
+
+ self.conv2 = ResBlock(
+ out_channels,
+ out_channels,
+ 3,
+ 1,
+ trans_func,
+ out_channels//2,
+ num_groups=1,
+ stride_1x1=False,
+ inplace_relu=True,
+ eps=1e-5,
+ bn_mmt=0.1,
+ dilation=1,
+ norm_module=nn.BatchNorm3d,
+ )
+
+ def forward(self, x):
+ x = self.conv1(x)
+ x = self.conv2(x)
+ return x
diff --git a/clean/video/pwtf_dvd/inference/slowfast/models/video_model_builder.py b/clean/video/pwtf_dvd/inference/slowfast/models/video_model_builder.py
new file mode 100644
index 0000000000000000000000000000000000000000..9a60fc57ccb7c2a322c2a43bc738650b97d7a958
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/models/video_model_builder.py
@@ -0,0 +1,2772 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Video models."""
+
+import torch
+import torch.nn as nn
+import copy
+
+import slowfast.utils.weight_init_helper as init_helper
+from slowfast.models.batchnorm_helper import get_norm
+
+from . import head_helper, resnet_helper, stem_helper
+from .build import MODEL_REGISTRY
+
+# Number of blocks for different stages given the model depth.
+_MODEL_STAGE_DEPTH = {18:(2,2,2,2),50: (3, 4, 6, 3), 101: (3, 4, 23, 3)}
+
+# Basis of temporal kernel sizes for each of the stage.
+_TEMPORAL_KERNEL_BASIS = {
+ "c2d": [
+ [[1]], # conv1 temporal kernel.
+ [[1]], # res2 temporal kernel.
+ [[1]], # res3 temporal kernel.
+ [[1]], # res4 temporal kernel.
+ [[1]], # res5 temporal kernel.
+ ],
+ "c2d_nopool": [
+ [[1]], # conv1 temporal kernel.
+ [[1]], # res2 temporal kernel.
+ [[1]], # res3 temporal kernel.
+ [[1]], # res4 temporal kernel.
+ [[1]], # res5 temporal kernel.
+ ],
+ "i3d": [
+ [[5]], # conv1 temporal kernel.
+ [[3]], # res2 temporal kernel.
+ [[3, 1]], # res3 temporal kernel.
+ [[3, 1]], # res4 temporal kernel.
+ [[1, 3]], # res5 temporal kernel.
+ ],
+ "r3d_18": [
+ [[3]], # conv1 temporal kernel.
+ [[3]], # res2 temporal kernel.
+ [[3, 1]], # res3 temporal kernel.
+ [[3, 1]], # res4 temporal kernel.
+ [[1, 3]], # res5 temporal kernel.
+ ],
+ "i3d_nopool": [
+ [[5]], # conv1 temporal kernel.
+ [[3]], # res2 temporal kernel.
+ [[3, 1]], # res3 temporal kernel.
+ [[3, 1]], # res4 temporal kernel.
+ [[1, 3]], # res5 temporal kernel.
+ ],
+ "slow": [
+ [[1]], # conv1 temporal kernel.
+ [[1]], # res2 temporal kernel.
+ [[1]], # res3 temporal kernel.
+ [[3]], # res4 temporal kernel.
+ [[3]], # res5 temporal kernel.
+ ],
+ "slowfast": [
+ [[1], [5]], # conv1 temporal kernel for slow and fast pathway.
+ [[1], [3]], # res2 temporal kernel for slow and fast pathway.
+ [[1], [3]], # res3 temporal kernel for slow and fast pathway.
+ [[3], [3]], # res4 temporal kernel for slow and fast pathway.
+ [[3], [3]], # res5 temporal kernel for slow and fast pathway.
+ ],
+}
+
+_POOL1 = {
+ "c2d": [[2, 1, 1]],
+ "c2d_nopool": [[1, 1, 1]],
+ "i3d": [[2, 1, 1]],
+ "r3d_18": [[2, 1, 1]],
+ "i3d_nopool": [[1, 1, 1]],
+ "slow": [[1, 1, 1]],
+ "slowfast": [[1, 1, 1], [1, 1, 1]],
+}
+
+
+
+
+class FuseFastToSlow(nn.Module):
+ """
+ Fuses the information from the Fast pathway to the Slow pathway. Given the
+ tensors from Slow pathway and Fast pathway, fuse information from Fast to
+ Slow, then return the fused tensors from Slow and Fast pathway in order.
+ """
+
+ def __init__(
+ self,
+ dim_in,
+ fusion_conv_channel_ratio,
+ fusion_kernel,
+ alpha,
+ eps=1e-5,
+ bn_mmt=0.1,
+ inplace_relu=True,
+ norm_module=nn.BatchNorm3d,
+ ):
+ """
+ Args:
+ dim_in (int): the channel dimension of the input.
+ fusion_conv_channel_ratio (int): channel ratio for the convolution
+ used to fuse from Fast pathway to Slow pathway.
+ fusion_kernel (int): kernel size of the convolution used to fuse
+ from Fast pathway to Slow pathway.
+ alpha (int): the frame rate ratio between the Fast and Slow pathway.
+ eps (float): epsilon for batch norm.
+ bn_mmt (float): momentum for batch norm. Noted that BN momentum in
+ PyTorch = 1 - BN momentum in Caffe2.
+ inplace_relu (bool): if True, calculate the relu on the original
+ input without allocating new memory.
+ norm_module (nn.Module): nn.Module for the normalization layer. The
+ default is nn.BatchNorm3d.
+ """
+ super(FuseFastToSlow, self).__init__()
+ self.conv_f2s = nn.Conv3d(
+ dim_in,
+ dim_in * fusion_conv_channel_ratio,
+ kernel_size=[fusion_kernel, 1, 1],
+ stride=[alpha, 1, 1],
+ padding=[fusion_kernel // 2, 0, 0],
+ bias=False,
+ )
+ self.bn = norm_module(
+ num_features=dim_in * fusion_conv_channel_ratio,
+ eps=eps,
+ momentum=bn_mmt,
+ )
+ self.relu = nn.ReLU(inplace_relu)
+
+ def forward(self, x):
+ x_s = x[0]
+ x_f = x[1]
+ fuse = self.conv_f2s(x_f)
+ fuse = self.bn(fuse)
+ fuse = self.relu(fuse)
+ x_s_fuse = torch.cat([x_s, fuse], 1)
+ return [x_s_fuse, x_f]
+
+
+
+@MODEL_REGISTRY.register()
+class SlowFast(nn.Module):
+ """
+ SlowFast model builder for SlowFast network.
+
+ Christoph Feichtenhofer, Haoqi Fan, Jitendra Malik, and Kaiming He.
+ "SlowFast networks for video recognition."
+ https://arxiv.org/pdf/1812.03982.pdf
+ """
+
+ def __init__(self, cfg):
+ """
+ The `__init__` method of any subclass should also contain these
+ arguments.
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ super(SlowFast, self).__init__()
+ self.norm_module = get_norm(cfg)
+ self.enable_detection = cfg.DETECTION.ENABLE
+ self.num_pathways = 2
+ self._construct_network(cfg)
+ init_helper.init_weights(
+ self, cfg.MODEL.FC_INIT_STD, cfg.RESNET.ZERO_INIT_FINAL_BN
+ )
+
+ def _construct_network(self, cfg):
+ """
+ Builds a SlowFast model. The first pathway is the Slow pathway and the
+ second pathway is the Fast pathway.
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ assert cfg.MODEL.ARCH in _POOL1.keys()
+ pool_size = _POOL1[cfg.MODEL.ARCH]
+ assert len({len(pool_size), self.num_pathways}) == 1
+ assert cfg.RESNET.DEPTH in _MODEL_STAGE_DEPTH.keys()
+
+ (d2, d3, d4, d5) = _MODEL_STAGE_DEPTH[cfg.RESNET.DEPTH]
+
+ num_groups = cfg.RESNET.NUM_GROUPS
+ width_per_group = cfg.RESNET.WIDTH_PER_GROUP
+ dim_inner = num_groups * width_per_group
+ out_dim_ratio = (
+ cfg.SLOWFAST.BETA_INV // cfg.SLOWFAST.FUSION_CONV_CHANNEL_RATIO
+ )
+
+ temp_kernel = _TEMPORAL_KERNEL_BASIS[cfg.MODEL.ARCH]
+
+ self.s1 = stem_helper.VideoModelStem(
+ dim_in=cfg.DATA.INPUT_CHANNEL_NUM,
+ dim_out=[width_per_group, width_per_group // cfg.SLOWFAST.BETA_INV],
+ kernel=[temp_kernel[0][0] + [7, 7], temp_kernel[0][1] + [7, 7]],
+ stride=[[1, 2, 2]] * 2,
+ padding=[
+ [temp_kernel[0][0][0] // 2, 3, 3],
+ [temp_kernel[0][1][0] // 2, 3, 3],
+ ],
+ norm_module=self.norm_module,
+ )
+ self.s1_fuse = FuseFastToSlow(
+ width_per_group // cfg.SLOWFAST.BETA_INV,
+ cfg.SLOWFAST.FUSION_CONV_CHANNEL_RATIO,
+ cfg.SLOWFAST.FUSION_KERNEL_SZ,
+ cfg.SLOWFAST.ALPHA,
+ norm_module=self.norm_module,
+ )
+
+ self.s2 = resnet_helper.ResStage(
+ dim_in=[
+ width_per_group + width_per_group // out_dim_ratio,
+ width_per_group // cfg.SLOWFAST.BETA_INV,
+ ],
+ dim_out=[
+ width_per_group * 4,
+ width_per_group * 4 // cfg.SLOWFAST.BETA_INV,
+ ],
+ dim_inner=[dim_inner, dim_inner // cfg.SLOWFAST.BETA_INV],
+ temp_kernel_sizes=temp_kernel[1],
+ stride=cfg.RESNET.SPATIAL_STRIDES[0],
+ num_blocks=[d2] * 2,
+ num_groups=[num_groups] * 2,
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[0],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[0],
+ nonlocal_group=cfg.NONLOCAL.GROUP[0],
+ nonlocal_pool=cfg.NONLOCAL.POOL[0],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[0],
+ norm_module=self.norm_module,
+ )
+ self.s2_fuse = FuseFastToSlow(
+ width_per_group * 4 // cfg.SLOWFAST.BETA_INV,
+ cfg.SLOWFAST.FUSION_CONV_CHANNEL_RATIO,
+ cfg.SLOWFAST.FUSION_KERNEL_SZ,
+ cfg.SLOWFAST.ALPHA,
+ norm_module=self.norm_module,
+ )
+
+ for pathway in range(self.num_pathways):
+ pool = nn.MaxPool3d(
+ kernel_size=pool_size[pathway],
+ stride=pool_size[pathway],
+ padding=[0, 0, 0],
+ )
+ self.add_module("pathway{}_pool".format(pathway), pool)
+
+ self.s3 = resnet_helper.ResStage(
+ dim_in=[
+ width_per_group * 4 + width_per_group * 4 // out_dim_ratio,
+ width_per_group * 4 // cfg.SLOWFAST.BETA_INV,
+ ],
+ dim_out=[
+ width_per_group * 8,
+ width_per_group * 8 // cfg.SLOWFAST.BETA_INV,
+ ],
+ dim_inner=[dim_inner * 2, dim_inner * 2 // cfg.SLOWFAST.BETA_INV],
+ temp_kernel_sizes=temp_kernel[2],
+ stride=cfg.RESNET.SPATIAL_STRIDES[1],
+ num_blocks=[d3] * 2,
+ num_groups=[num_groups] * 2,
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[1],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[1],
+ nonlocal_group=cfg.NONLOCAL.GROUP[1],
+ nonlocal_pool=cfg.NONLOCAL.POOL[1],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[1],
+ norm_module=self.norm_module,
+ )
+ self.s3_fuse = FuseFastToSlow(
+ width_per_group * 8 // cfg.SLOWFAST.BETA_INV,
+ cfg.SLOWFAST.FUSION_CONV_CHANNEL_RATIO,
+ cfg.SLOWFAST.FUSION_KERNEL_SZ,
+ cfg.SLOWFAST.ALPHA,
+ norm_module=self.norm_module,
+ )
+
+ self.s4 = resnet_helper.ResStage(
+ dim_in=[
+ width_per_group * 8 + width_per_group * 8 // out_dim_ratio,
+ width_per_group * 8 // cfg.SLOWFAST.BETA_INV,
+ ],
+ dim_out=[
+ width_per_group * 16,
+ width_per_group * 16 // cfg.SLOWFAST.BETA_INV,
+ ],
+ dim_inner=[dim_inner * 4, dim_inner * 4 // cfg.SLOWFAST.BETA_INV],
+ temp_kernel_sizes=temp_kernel[3],
+ stride=cfg.RESNET.SPATIAL_STRIDES[2],
+ num_blocks=[d4] * 2,
+ num_groups=[num_groups] * 2,
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[2],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[2],
+ nonlocal_group=cfg.NONLOCAL.GROUP[2],
+ nonlocal_pool=cfg.NONLOCAL.POOL[2],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[2],
+ norm_module=self.norm_module,
+ )
+ self.s4_fuse = FuseFastToSlow(
+ width_per_group * 16 // cfg.SLOWFAST.BETA_INV,
+ cfg.SLOWFAST.FUSION_CONV_CHANNEL_RATIO,
+ cfg.SLOWFAST.FUSION_KERNEL_SZ,
+ cfg.SLOWFAST.ALPHA,
+ norm_module=self.norm_module,
+ )
+
+ self.s5 = resnet_helper.ResStage(
+ dim_in=[
+ width_per_group * 16 + width_per_group * 16 // out_dim_ratio,
+ width_per_group * 16 // cfg.SLOWFAST.BETA_INV,
+ ],
+ dim_out=[
+ width_per_group * 32,
+ width_per_group * 32 // cfg.SLOWFAST.BETA_INV,
+ ],
+ dim_inner=[dim_inner * 8, dim_inner * 8 // cfg.SLOWFAST.BETA_INV],
+ temp_kernel_sizes=temp_kernel[4],
+ stride=cfg.RESNET.SPATIAL_STRIDES[3],
+ num_blocks=[d5] * 2,
+ num_groups=[num_groups] * 2,
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[3],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[3],
+ nonlocal_group=cfg.NONLOCAL.GROUP[3],
+ nonlocal_pool=cfg.NONLOCAL.POOL[3],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[3],
+ norm_module=self.norm_module,
+ )
+
+ if cfg.DETECTION.ENABLE:
+ raise NotImplementedError
+ else:
+ self.head = head_helper.ResNetBasicHead(
+ dim_in=[
+ width_per_group * 32,
+ width_per_group * 32 // cfg.SLOWFAST.BETA_INV,
+ ],
+ num_classes=cfg.MODEL.NUM_CLASSES,
+ pool_size=[None, None]
+ if cfg.MULTIGRID.SHORT_CYCLE
+ else [
+ [
+ cfg.DATA.NUM_FRAMES
+ // cfg.SLOWFAST.ALPHA
+ // pool_size[0][0],
+ cfg.DATA.CROP_SIZE // 32 // pool_size[0][1],
+ cfg.DATA.CROP_SIZE // 32 // pool_size[0][2],
+ ],
+ [
+ cfg.DATA.NUM_FRAMES // pool_size[1][0],
+ cfg.DATA.CROP_SIZE // 32 // pool_size[1][1],
+ cfg.DATA.CROP_SIZE // 32 // pool_size[1][2],
+ ],
+ ], # None for AdaptiveAvgPool3d((1, 1, 1))
+ dropout_rate=cfg.MODEL.DROPOUT_RATE,
+ act_func=cfg.MODEL.HEAD_ACT,
+ )
+
+ def forward(self, x, bboxes=None):
+ x = self.s1(x)
+ x = self.s1_fuse(x)
+ x = self.s2(x)
+ x = self.s2_fuse(x)
+ for pathway in range(self.num_pathways):
+ pool = getattr(self, "pathway{}_pool".format(pathway))
+ x[pathway] = pool(x[pathway])
+ x = self.s3(x)
+ x = self.s3_fuse(x)
+ x = self.s4(x)
+ x = self.s4_fuse(x)
+ x = self.s5(x)
+ if self.enable_detection:
+ x = self.head(x, bboxes)
+ else:
+ x = self.head(x)
+ return x
+
+#############################
+### ftcn using this Model ###
+#############################
+@MODEL_REGISTRY.register()
+class ResNet(nn.Module):
+ """
+ ResNet model builder. It builds a ResNet like network backbone without
+ lateral connection (C2D, I3D, Slow).
+
+ Christoph Feichtenhofer, Haoqi Fan, Jitendra Malik, and Kaiming He.
+ "SlowFast networks for video recognition."
+ https://arxiv.org/pdf/1812.03982.pdf
+
+ Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He.
+ "Non-local neural networks."
+ https://arxiv.org/pdf/1711.07971.pdf
+ """
+
+ def __init__(self, cfg):
+ """
+ The `__init__` method of any subclass should also contain these
+ arguments.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ super(ResNet, self).__init__()
+ self.norm_module = get_norm(cfg)
+ self.enable_detection = cfg.DETECTION.ENABLE
+ self.num_pathways = 1
+ self._construct_network(cfg)
+ init_helper.init_weights(
+ self, cfg.MODEL.FC_INIT_STD, cfg.RESNET.ZERO_INIT_FINAL_BN
+ )
+
+ def _construct_network(self, cfg):
+ """
+ Builds a single pathway ResNet model.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ assert cfg.MODEL.ARCH in _POOL1.keys()
+ pool_size = _POOL1[cfg.MODEL.ARCH]
+ assert len({len(pool_size), self.num_pathways}) == 1
+ assert cfg.RESNET.DEPTH in _MODEL_STAGE_DEPTH.keys()
+
+ (d2, d3, d4, d5) = _MODEL_STAGE_DEPTH[cfg.RESNET.DEPTH]
+
+ num_groups = cfg.RESNET.NUM_GROUPS
+ width_per_group = cfg.RESNET.WIDTH_PER_GROUP
+ dim_inner = num_groups * width_per_group
+
+ temp_kernel = _TEMPORAL_KERNEL_BASIS[cfg.MODEL.ARCH]
+
+ # fix: inplace_relu=True
+ self.s1 = stem_helper.VideoModelStem(
+ dim_in=cfg.DATA.INPUT_CHANNEL_NUM,
+ dim_out=[width_per_group],
+ kernel=[temp_kernel[0][0] + [7, 7]],
+ stride=[[1, 2, 2]],
+ padding=[[temp_kernel[0][0][0] // 2, 3, 3]],
+ norm_module=self.norm_module,
+ )
+
+ self.s2 = resnet_helper.ResStage(
+ dim_in=[width_per_group],
+ dim_out=[width_per_group * 4],
+ dim_inner=[dim_inner],
+ temp_kernel_sizes=temp_kernel[1],
+ stride=cfg.RESNET.SPATIAL_STRIDES[0],
+ num_blocks=[d2],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[0],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[0],
+ nonlocal_group=cfg.NONLOCAL.GROUP[0],
+ nonlocal_pool=cfg.NONLOCAL.POOL[0],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ # inplace_relu=cfg.RESNET.INPLACE_RELU,
+ inplace_relu=False, # default :: True
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[0],
+ norm_module=self.norm_module,
+ )
+
+ for pathway in range(self.num_pathways):
+ pool = nn.MaxPool3d(
+ kernel_size=pool_size[pathway],
+ stride=pool_size[pathway],
+ padding=[0, 0, 0],
+ )
+ self.add_module("pathway{}_pool".format(pathway), pool)
+
+ self.s3 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 4],
+ dim_out=[width_per_group * 8],
+ dim_inner=[dim_inner * 2],
+ temp_kernel_sizes=temp_kernel[2],
+ stride=cfg.RESNET.SPATIAL_STRIDES[1],
+ num_blocks=[d3],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[1],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[1],
+ nonlocal_group=cfg.NONLOCAL.GROUP[1],
+ nonlocal_pool=cfg.NONLOCAL.POOL[1],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ # inplace_relu=cfg.RESNET.INPLACE_RELU,
+ inplace_relu=False, # default :: True
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[1],
+ norm_module=self.norm_module,
+ )
+
+ self.s4 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 8],
+ dim_out=[width_per_group * 16],
+ dim_inner=[dim_inner * 4],
+ temp_kernel_sizes=temp_kernel[3],
+ stride=cfg.RESNET.SPATIAL_STRIDES[2],
+ num_blocks=[d4],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[2],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[2],
+ nonlocal_group=cfg.NONLOCAL.GROUP[2],
+ nonlocal_pool=cfg.NONLOCAL.POOL[2],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ # inplace_relu=cfg.RESNET.INPLACE_RELU,
+ inplace_relu=False, # default :: True ########### Fixed by TH ############ default :: True
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[2],
+ norm_module=self.norm_module,
+ )
+
+ self.s5 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 16],
+ dim_out=[width_per_group * 32],
+ dim_inner=[dim_inner * 8],
+ temp_kernel_sizes=temp_kernel[4],
+ stride=cfg.RESNET.SPATIAL_STRIDES[3],
+ num_blocks=[d5],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[3],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[3],
+ nonlocal_group=cfg.NONLOCAL.GROUP[3],
+ nonlocal_pool=cfg.NONLOCAL.POOL[3],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ # inplace_relu=cfg.RESNET.INPLACE_RELU,
+ inplace_relu=False, # default :: True
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[3],
+ norm_module=self.norm_module,
+ )
+ ############################# 0918 Head X ###############################################
+ # if self.enable_detection:
+ # raise NotImplementedError
+ # else:
+ # self.head = head_helper.ResNetBasicHead(
+ # dim_in=[width_per_group * 32],
+ # num_classes=cfg.MODEL.NUM_CLASSES,
+ # pool_size=[None, None]
+ # if cfg.MULTIGRID.SHORT_CYCLE
+ # else [
+ # [
+ # cfg.DATA.NUM_FRAMES // pool_size[0][0],
+ # cfg.DATA.CROP_SIZE // 32 // pool_size[0][1],
+ # cfg.DATA.CROP_SIZE // 32 // pool_size[0][2],
+ # ]
+ # ], # None for AdaptiveAvgPool3d((1, 1, 1))
+ # dropout_rate=cfg.MODEL.DROPOUT_RATE,
+ # act_func=cfg.MODEL.HEAD_ACT,
+ # )
+ # x:[images, ]([16, 3, 32, 224, 224])
+ # ft_feats [out1, out2, out3, out4, out5]
+ def forward(self, x,ft_feats, bboxes=None):
+ ### FTCN ###
+ ################# Fourier Transform #################
+ # x shape s1 torch.Size([16, 64, 32, 56, 56])
+ # x shape s2 torch.Size([16, 256, 32, 56, 56])
+ # x shape s3 torch.Size([16, 512, 16, 28, 28])
+ # x shape s4 torch.Size([16, 1024, 16, 14, 14])
+ # x shape s5 torch.Size([16, 2048, 16, 7, 7])
+ #######################################################
+ # FTfeature : (torch.Size([16, 64, 56, 56]),
+ # torch.Size([16, 256, 56, 56]),
+ # torch.Size([16, 512, 28, 28]),
+ # torch.Size([16, 1024, 14, 14]),
+ # torch.Size([16, 2048]))
+ ############## #########################################
+
+ x = self.s1(x)
+ x[0] = x[0].clone() + ft_feats[0].unsqueeze(2).clone()
+ x = self.s2(x)
+ x[0] = x[0].clone() + ft_feats[1].unsqueeze(2).clone()
+ for pathway in range(self.num_pathways):
+ pool = getattr(self, "pathway{}_pool".format(pathway))
+ x[pathway] = pool(x[pathway])
+ x = self.s3(x)
+ x[0] = x[0].clone() + ft_feats[2].unsqueeze(2).clone()
+ x = self.s4(x)
+ x[0] = x[0].clone() + ft_feats[3].unsqueeze(2).clone()
+ x = self.s5(x)
+ ############################# 0918 Head X ###############################################
+ ### FTCN ###
+ # x shape input transformer torch.Size([16, 1024, 16, 14, 14])
+ # x shape input transformer torch.Size([16, 2048, 16, 7, 7])
+ # if self.enable_detection:
+ # x = self.head(x, bboxes)
+ # else:
+ # ### FTCN ###
+ # x = self.head(x,ft_feats[4].clone())
+
+ ###########
+ # x shape output transformer torch.Size([1])
+ return x[0] , ft_feats[4]
+################################################################
+@MODEL_REGISTRY.register()
+class ResNetVar(nn.Module):
+ """
+ ResNet model builder. It builds a ResNet like network backbone without
+ lateral connection (C2D, I3D, Slow).
+
+ Christoph Feichtenhofer, Haoqi Fan, Jitendra Malik, and Kaiming He.
+ "SlowFast networks for video recognition."
+ https://arxiv.org/pdf/1812.03982.pdf
+
+ Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He.
+ "Non-local neural networks."
+ https://arxiv.org/pdf/1711.07971.pdf
+ """
+
+ def __init__(self, cfg):
+ """
+ The `__init__` method of any subclass should also contain these
+ arguments.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ super(ResNetVar, self).__init__()
+ self.norm_module = get_norm(cfg)
+ self.enable_detection = cfg.DETECTION.ENABLE
+ self.num_pathways = 1
+ self._construct_network(cfg)
+ init_helper.init_weights(
+ self, cfg.MODEL.FC_INIT_STD, cfg.RESNET.ZERO_INIT_FINAL_BN
+ )
+
+ def _construct_network(self, cfg):
+ """
+ Builds a single pathway ResNet model.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ assert cfg.MODEL.ARCH in _POOL1.keys()
+ pool_size = _POOL1[cfg.MODEL.ARCH]
+ assert len({len(pool_size), self.num_pathways}) == 1
+ assert cfg.RESNET.DEPTH in _MODEL_STAGE_DEPTH.keys()
+
+ (d2, d3, d4, d5) = _MODEL_STAGE_DEPTH[cfg.RESNET.DEPTH]
+
+ num_groups = cfg.RESNET.NUM_GROUPS
+ width_per_group = cfg.RESNET.WIDTH_PER_GROUP
+ dim_inner = num_groups * width_per_group
+
+ temp_kernel = _TEMPORAL_KERNEL_BASIS[cfg.MODEL.ARCH]
+
+ self.s1 = stem_helper.VideoModelStem(
+ dim_in=cfg.DATA.INPUT_CHANNEL_NUM,
+ dim_out=[width_per_group],
+ kernel=[temp_kernel[0][0] + [7, 7]],
+ stride=[[1, 2, 2]],
+ padding=[[temp_kernel[0][0][0] // 2, 3, 3]],
+ norm_module=self.norm_module,
+ )
+
+ self.s2 = resnet_helper.ResStage(
+ dim_in=[width_per_group],
+ dim_out=[width_per_group * 4],
+ dim_inner=[dim_inner],
+ temp_kernel_sizes=temp_kernel[1],
+ stride=cfg.RESNET.SPATIAL_STRIDES[0],
+ num_blocks=[d2],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[0],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[0],
+ nonlocal_group=cfg.NONLOCAL.GROUP[0],
+ nonlocal_pool=cfg.NONLOCAL.POOL[0],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[0],
+ norm_module=self.norm_module,
+ )
+
+ for pathway in range(self.num_pathways):
+ pool = nn.MaxPool3d(
+ kernel_size=pool_size[pathway],
+ stride=pool_size[pathway],
+ padding=[0, 0, 0],
+ )
+ self.add_module("pathway{}_pool".format(pathway), pool)
+
+ self.s3 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 4],
+ dim_out=[width_per_group * 8],
+ dim_inner=[dim_inner * 2],
+ temp_kernel_sizes=temp_kernel[2],
+ stride=cfg.RESNET.SPATIAL_STRIDES[1],
+ num_blocks=[d3],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[1],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[1],
+ nonlocal_group=cfg.NONLOCAL.GROUP[1],
+ nonlocal_pool=cfg.NONLOCAL.POOL[1],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[1],
+ norm_module=self.norm_module,
+ )
+
+ self.s4 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 8],
+ dim_out=[width_per_group * 16],
+ dim_inner=[dim_inner * 4],
+ temp_kernel_sizes=temp_kernel[3],
+ stride=cfg.RESNET.SPATIAL_STRIDES[2],
+ num_blocks=[d4],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[2],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[2],
+ nonlocal_group=cfg.NONLOCAL.GROUP[2],
+ nonlocal_pool=cfg.NONLOCAL.POOL[2],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[2],
+ norm_module=self.norm_module,
+ )
+
+ self.s5 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 16],
+ dim_out=[width_per_group * 32],
+ dim_inner=[dim_inner * 8],
+ temp_kernel_sizes=temp_kernel[4],
+ stride=cfg.RESNET.SPATIAL_STRIDES[3],
+ num_blocks=[d5],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[3],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[3],
+ nonlocal_group=cfg.NONLOCAL.GROUP[3],
+ nonlocal_pool=cfg.NONLOCAL.POOL[3],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[3],
+ norm_module=self.norm_module,
+ )
+
+ if self.enable_detection:
+ raise NotImplementedError
+ else:
+ self.head = head_helper.ResNetBasicHead(
+ dim_in=[width_per_group * 32],
+ num_classes=cfg.MODEL.NUM_CLASSES,
+ pool_size=[None],
+ dropout_rate=cfg.MODEL.DROPOUT_RATE,
+ act_func=cfg.MODEL.HEAD_ACT,
+ )
+
+ def forward(self, x, bboxes=None):
+ x = self.s1(x)
+ x = self.s2(x)
+ for pathway in range(self.num_pathways):
+ pool = getattr(self, "pathway{}_pool".format(pathway))
+ x[pathway] = pool(x[pathway])
+ x = self.s3(x)
+ x = self.s4(x)
+ x = self.s5(x)
+ if self.enable_detection:
+ x = self.head(x, bboxes)
+ else:
+ x = self.head(x)
+ return x
+
+@MODEL_REGISTRY.register()
+class ResNetBase(nn.Module):
+ """
+ ResNet model builder. It builds a ResNet like network backbone without
+ lateral connection (C2D, I3D, Slow).
+
+ Christoph Feichtenhofer, Haoqi Fan, Jitendra Malik, and Kaiming He.
+ "SlowFast networks for video recognition."
+ https://arxiv.org/pdf/1812.03982.pdf
+
+ Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He.
+ "Non-local neural networks."
+ https://arxiv.org/pdf/1711.07971.pdf
+ """
+
+ def __init__(self, cfg):
+ """
+ The `__init__` method of any subclass should also contain these
+ arguments.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ super(ResNetBase, self).__init__()
+ self.norm_module = get_norm(cfg)
+ self.enable_detection = cfg.DETECTION.ENABLE
+ self.num_pathways = 1
+ self._construct_network(cfg)
+ init_helper.init_weights(
+ self, cfg.MODEL.FC_INIT_STD, cfg.RESNET.ZERO_INIT_FINAL_BN
+ )
+
+ def _construct_network(self, cfg):
+ """
+ Builds a single pathway ResNet model.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ assert cfg.MODEL.ARCH in _POOL1.keys()
+ pool_size = _POOL1[cfg.MODEL.ARCH]
+ assert len({len(pool_size), self.num_pathways}) == 1
+ assert cfg.RESNET.DEPTH in _MODEL_STAGE_DEPTH.keys()
+
+ (d2, d3, d4, d5) = _MODEL_STAGE_DEPTH[cfg.RESNET.DEPTH]
+
+ num_groups = cfg.RESNET.NUM_GROUPS
+ width_per_group = cfg.RESNET.WIDTH_PER_GROUP
+ dim_inner = num_groups * width_per_group
+
+ temp_kernel = _TEMPORAL_KERNEL_BASIS[cfg.MODEL.ARCH]
+
+ self.s1 = stem_helper.VideoModelStem(
+ dim_in=cfg.DATA.INPUT_CHANNEL_NUM,
+ dim_out=[width_per_group],
+ kernel=[temp_kernel[0][0] + [7, 7]],
+ stride=[[1, 2, 2]],
+ padding=[[temp_kernel[0][0][0] // 2, 3, 3]],
+ norm_module=self.norm_module,
+ )
+
+ self.s2 = resnet_helper.ResStage(
+ dim_in=[width_per_group],
+ dim_out=[width_per_group * 4],
+ dim_inner=[dim_inner],
+ temp_kernel_sizes=temp_kernel[1],
+ stride=cfg.RESNET.SPATIAL_STRIDES[0],
+ num_blocks=[d2],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[0],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[0],
+ nonlocal_group=cfg.NONLOCAL.GROUP[0],
+ nonlocal_pool=cfg.NONLOCAL.POOL[0],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[0],
+ norm_module=self.norm_module,
+ )
+
+ for pathway in range(self.num_pathways):
+ pool = nn.MaxPool3d(
+ kernel_size=pool_size[pathway],
+ stride=pool_size[pathway],
+ padding=[0, 0, 0],
+ )
+ self.add_module("pathway{}_pool".format(pathway), pool)
+
+ self.s3 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 4],
+ dim_out=[width_per_group * 8],
+ dim_inner=[dim_inner * 2],
+ temp_kernel_sizes=temp_kernel[2],
+ stride=cfg.RESNET.SPATIAL_STRIDES[1],
+ num_blocks=[d3],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[1],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[1],
+ nonlocal_group=cfg.NONLOCAL.GROUP[1],
+ nonlocal_pool=cfg.NONLOCAL.POOL[1],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[1],
+ norm_module=self.norm_module,
+ )
+
+ self.s4 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 8],
+ dim_out=[width_per_group * 16],
+ dim_inner=[dim_inner * 4],
+ temp_kernel_sizes=temp_kernel[3],
+ stride=cfg.RESNET.SPATIAL_STRIDES[2],
+ num_blocks=[d4],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[2],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[2],
+ nonlocal_group=cfg.NONLOCAL.GROUP[2],
+ nonlocal_pool=cfg.NONLOCAL.POOL[2],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[2],
+ norm_module=self.norm_module,
+ )
+
+ self.s5 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 16],
+ dim_out=[width_per_group * 32],
+ dim_inner=[dim_inner * 8],
+ temp_kernel_sizes=temp_kernel[4],
+ stride=cfg.RESNET.SPATIAL_STRIDES[3],
+ num_blocks=[d5],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[3],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[3],
+ nonlocal_group=cfg.NONLOCAL.GROUP[3],
+ nonlocal_pool=cfg.NONLOCAL.POOL[3],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[3],
+ norm_module=self.norm_module,
+ )
+
+ if self.enable_detection:
+ raise NotImplementedError
+ else:
+ self.head = head_helper.ResNetBasicHead(
+ dim_in=[width_per_group * 32],
+ num_classes=cfg.MODEL.NUM_CLASSES,
+ pool_size=[None, None]
+ if cfg.MULTIGRID.SHORT_CYCLE
+ else [
+ None
+ ], # None for AdaptiveAvgPool3d((1, 1, 1))
+ dropout_rate=cfg.MODEL.DROPOUT_RATE,
+ act_func=cfg.MODEL.HEAD_ACT,
+ )
+
+ def forward(self, x, bboxes=None):
+ x = self.s1(x)
+ x = self.s2(x)
+ for pathway in range(self.num_pathways):
+ pool = getattr(self, "pathway{}_pool".format(pathway))
+ x[pathway] = pool(x[pathway])
+ x = self.s3(x)
+ x = self.s4(x)
+ x = self.s5(x)
+ if self.enable_detection:
+ x = self.head(x, bboxes)
+ else:
+ x = self.head(x)
+ return x
+
+
+@MODEL_REGISTRY.register()
+class ResNetFreeze(nn.Module):
+ """
+ ResNet model builder. It builds a ResNet like network backbone without
+ lateral connection (C2D, I3D, Slow).
+
+ Christoph Feichtenhofer, Haoqi Fan, Jitendra Malik, and Kaiming He.
+ "SlowFast networks for video recognition."
+ https://arxiv.org/pdf/1812.03982.pdf
+
+ Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He.
+ "Non-local neural networks."
+ https://arxiv.org/pdf/1711.07971.pdf
+ """
+
+ def __init__(self, cfg):
+ """
+ The `__init__` method of any subclass should also contain these
+ arguments.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ super(ResNetFreeze, self).__init__()
+ self.norm_module = get_norm(cfg)
+ self.enable_detection = cfg.DETECTION.ENABLE
+ self.num_pathways = 1
+ self._construct_network(cfg)
+ init_helper.init_weights(
+ self, cfg.MODEL.FC_INIT_STD, cfg.RESNET.ZERO_INIT_FINAL_BN
+ )
+
+ def _construct_network(self, cfg):
+ """
+ Builds a single pathway ResNet model.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ assert cfg.MODEL.ARCH in _POOL1.keys()
+ pool_size = _POOL1[cfg.MODEL.ARCH]
+ assert len({len(pool_size), self.num_pathways}) == 1
+ assert cfg.RESNET.DEPTH in _MODEL_STAGE_DEPTH.keys()
+
+ (d2, d3, d4, d5) = _MODEL_STAGE_DEPTH[cfg.RESNET.DEPTH]
+
+ num_groups = cfg.RESNET.NUM_GROUPS
+ width_per_group = cfg.RESNET.WIDTH_PER_GROUP
+ dim_inner = num_groups * width_per_group
+
+ temp_kernel = _TEMPORAL_KERNEL_BASIS[cfg.MODEL.ARCH]
+
+ self.s1 = stem_helper.VideoModelStem(
+ dim_in=cfg.DATA.INPUT_CHANNEL_NUM,
+ dim_out=[width_per_group],
+ kernel=[temp_kernel[0][0] + [7, 7]],
+ stride=[[1, 2, 2]],
+ padding=[[temp_kernel[0][0][0] // 2, 3, 3]],
+ norm_module=self.norm_module,
+ )
+
+ self.s2 = resnet_helper.ResStage(
+ dim_in=[width_per_group],
+ dim_out=[width_per_group * 4],
+ dim_inner=[dim_inner],
+ temp_kernel_sizes=temp_kernel[1],
+ stride=cfg.RESNET.SPATIAL_STRIDES[0],
+ num_blocks=[d2],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[0],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[0],
+ nonlocal_group=cfg.NONLOCAL.GROUP[0],
+ nonlocal_pool=cfg.NONLOCAL.POOL[0],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[0],
+ norm_module=self.norm_module,
+ )
+
+ for pathway in range(self.num_pathways):
+ pool = nn.MaxPool3d(
+ kernel_size=pool_size[pathway],
+ stride=pool_size[pathway],
+ padding=[0, 0, 0],
+ )
+ self.add_module("pathway{}_pool".format(pathway), pool)
+
+ self.s3 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 4],
+ dim_out=[width_per_group * 8],
+ dim_inner=[dim_inner * 2],
+ temp_kernel_sizes=temp_kernel[2],
+ stride=cfg.RESNET.SPATIAL_STRIDES[1],
+ num_blocks=[d3],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[1],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[1],
+ nonlocal_group=cfg.NONLOCAL.GROUP[1],
+ nonlocal_pool=cfg.NONLOCAL.POOL[1],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[1],
+ norm_module=self.norm_module,
+ )
+
+ self.s4 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 8],
+ dim_out=[width_per_group * 16],
+ dim_inner=[dim_inner * 4],
+ temp_kernel_sizes=temp_kernel[3],
+ stride=cfg.RESNET.SPATIAL_STRIDES[2],
+ num_blocks=[d4],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[2],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[2],
+ nonlocal_group=cfg.NONLOCAL.GROUP[2],
+ nonlocal_pool=cfg.NONLOCAL.POOL[2],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[2],
+ norm_module=self.norm_module,
+ )
+
+ self.s5 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 16],
+ dim_out=[width_per_group * 32],
+ dim_inner=[dim_inner * 8],
+ temp_kernel_sizes=temp_kernel[4],
+ stride=cfg.RESNET.SPATIAL_STRIDES[3],
+ num_blocks=[d5],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[3],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[3],
+ nonlocal_group=cfg.NONLOCAL.GROUP[3],
+ nonlocal_pool=cfg.NONLOCAL.POOL[3],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[3],
+ norm_module=self.norm_module,
+ )
+
+ if self.enable_detection:
+ raise NotImplementedError
+ else:
+ self.head = head_helper.ResNetBasicHead(
+ dim_in=[width_per_group * 32],
+ num_classes=cfg.MODEL.NUM_CLASSES,
+ pool_size=[None,None]
+ if cfg.MULTIGRID.SHORT_CYCLE
+ else [
+ None
+ ], # None for AdaptiveAvgPool3d((1, 1, 1))
+ dropout_rate=cfg.MODEL.DROPOUT_RATE,
+ act_func=cfg.MODEL.HEAD_ACT,
+ )
+
+ def forward(self, x, freeze_backbone=False):
+ assert isinstance(freeze_backbone,bool)
+ x = self.s1(x)
+ x = self.s2(x)
+ # for pathway in range(self.num_pathways):
+ # pool = getattr(self, "pathway{}_pool".format(pathway))
+ # x[pathway] = pool(x[pathway])
+ x = self.s3(x)
+ x = self.s4(x)
+ x = self.s5(x)
+ if freeze_backbone:
+ x=[item.detach() for item in x]
+
+ x = self.head(x)
+ return x
+
+
+
+import torch.nn.functional as F
+from .unet_helper import DecoderBlock,LightDecoderBlock,ResDecoderBlock
+
+
+@MODEL_REGISTRY.register()
+class ResUNet(nn.Module):
+ """
+ ResNet model builder. It builds a ResNet like network backbone without
+ lateral connection (C2D, I3D, Slow).
+
+ Christoph Feichtenhofer, Haoqi Fan, Jitendra Malik, and Kaiming He.
+ "SlowFast networks for video recognition."
+ https://arxiv.org/pdf/1812.03982.pdf
+
+ Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He.
+ "Non-local neural networks."
+ https://arxiv.org/pdf/1711.07971.pdf
+ """
+
+ def __init__(self, cfg):
+ """
+ The `__init__` method of any subclass should also contain these
+ arguments.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ super(ResUNet, self).__init__()
+ self.norm_module = get_norm(cfg)
+ self.enable_detection = cfg.DETECTION.ENABLE
+ self.enable_jitter = cfg.JITTER.ENABLE
+ self.num_pathways = 1
+ assert cfg.DATA.TRAIN_CROP_SIZE == cfg.DATA.TEST_CROP_SIZE
+ self.image_size = cfg.DATA.TRAIN_CROP_SIZE
+ self.clip_size = cfg.DATA.NUM_FRAMES
+ self._construct_network(cfg)
+ init_helper.init_weights(
+ self, cfg.MODEL.FC_INIT_STD, cfg.RESNET.ZERO_INIT_FINAL_BN
+ )
+
+ def _construct_network(self, cfg):
+ """
+ Builds a single pathway ResNet model.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ assert cfg.MODEL.ARCH in _POOL1.keys()
+ pool_size = _POOL1[cfg.MODEL.ARCH]
+ self.cfg = cfg
+ assert len({len(pool_size), self.num_pathways}) == 1
+ assert cfg.RESNET.DEPTH in _MODEL_STAGE_DEPTH.keys()
+
+ (d2, d3, d4, d5) = _MODEL_STAGE_DEPTH[cfg.RESNET.DEPTH]
+
+ num_groups = cfg.RESNET.NUM_GROUPS
+ width_per_group = cfg.RESNET.WIDTH_PER_GROUP
+ dim_inner = num_groups * width_per_group
+
+ temp_kernel = _TEMPORAL_KERNEL_BASIS[cfg.MODEL.ARCH]
+
+ self.s1 = stem_helper.VideoModelStem(
+ dim_in=cfg.DATA.INPUT_CHANNEL_NUM,
+ dim_out=[width_per_group],
+ kernel=[temp_kernel[0][0] + [7, 7]],
+ stride=[[1, 2, 2]],
+ padding=[[temp_kernel[0][0][0] // 2, 3, 3]],
+ norm_module=self.norm_module,
+ )
+
+ self.s2 = resnet_helper.ResStage(
+ dim_in=[width_per_group],
+ dim_out=[width_per_group * 4],
+ dim_inner=[dim_inner],
+ temp_kernel_sizes=temp_kernel[1],
+ stride=cfg.RESNET.SPATIAL_STRIDES[0],
+ num_blocks=[d2],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[0],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[0],
+ nonlocal_group=cfg.NONLOCAL.GROUP[0],
+ nonlocal_pool=cfg.NONLOCAL.POOL[0],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[0],
+ norm_module=self.norm_module,
+ )
+
+ for pathway in range(self.num_pathways):
+ pool = nn.MaxPool3d(
+ kernel_size=pool_size[pathway],
+ stride=pool_size[pathway],
+ padding=[0, 0, 0],
+ )
+ self.add_module("pathway{}_pool".format(pathway), pool)
+
+ self.s3 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 4],
+ dim_out=[width_per_group * 8],
+ dim_inner=[dim_inner * 2],
+ temp_kernel_sizes=temp_kernel[2],
+ stride=cfg.RESNET.SPATIAL_STRIDES[1],
+ num_blocks=[d3],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[1],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[1],
+ nonlocal_group=cfg.NONLOCAL.GROUP[1],
+ nonlocal_pool=cfg.NONLOCAL.POOL[1],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[1],
+ norm_module=self.norm_module,
+ )
+
+ self.s4 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 8],
+ dim_out=[width_per_group * 16],
+ dim_inner=[dim_inner * 4],
+ temp_kernel_sizes=temp_kernel[3],
+ stride=cfg.RESNET.SPATIAL_STRIDES[2],
+ num_blocks=[d4],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[2],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[2],
+ nonlocal_group=cfg.NONLOCAL.GROUP[2],
+ nonlocal_pool=cfg.NONLOCAL.POOL[2],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[2],
+ norm_module=self.norm_module,
+ )
+
+ # self.s5 = resnet_helper.ResStage(
+ # dim_in=[width_per_group * 16],
+ # dim_out=[width_per_group * 32],
+ # dim_inner=[dim_inner * 8],
+ # temp_kernel_sizes=temp_kernel[4],
+ # stride=cfg.RESNET.SPATIAL_STRIDES[3],
+ # num_blocks=[d5],
+ # num_groups=[num_groups],
+ # num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[3],
+ # nonlocal_inds=cfg.NONLOCAL.LOCATION[3],
+ # nonlocal_group=cfg.NONLOCAL.GROUP[3],
+ # nonlocal_pool=cfg.NONLOCAL.POOL[3],
+ # instantiation=cfg.NONLOCAL.INSTANTIATION,
+ # trans_func_name=cfg.RESNET.TRANS_FUNC,
+ # stride_1x1=cfg.RESNET.STRIDE_1X1,
+ # inplace_relu=cfg.RESNET.INPLACE_RELU,
+ # dilation=cfg.RESNET.SPATIAL_DILATIONS[3],
+ # norm_module=self.norm_module,
+ # )
+ self.labels=["rotate","light"]
+ self.dual_define("t4",self.labels,DecoderBlock(width_per_group * 16,width_per_group * 8,width_per_group * 8))
+ self.dual_define("t3",self.labels,DecoderBlock(width_per_group * 8,width_per_group * 4, 256))
+ self.dual_define("conv1x1",self.labels,nn.Sequential(
+ nn.Conv3d(width_per_group*4+width_per_group, 1, kernel_size=(1, 1, 1), stride=1, padding=0), nn.Sigmoid()
+ ))
+
+ self.linear = nn.Sequential(nn.Linear(1, 1), nn.Sigmoid())
+
+ def forward_plus(self, x, y, net):
+ return [net(x)[0] + y[0]]
+
+
+ def dual_define(self,name,labels,net):
+ for label in labels:
+ self.add_module(f"{name}_{label}",copy.deepcopy(net))
+
+
+
+ def upsample(self, x, dims=["space"]):
+ ori_size = x[0].shape[2:5]
+ t, h, w = ori_size
+ if "space" in dims:
+ h = 2 * h
+ w = 2 * w
+ if "time" in dims:
+ t = 2 * t
+ size = (t, h, w)
+ return [F.interpolate(x[0], size)]
+
+ def concat(self,x,y):
+ return [torch.cat([x[0],y[0]],1)]
+
+
+
+ # @torchsnooper.snoop()
+ def forward(self, x, bboxes=None):
+ x1 = self.s1(x) # 1,64,8,56,56
+ x2 = self.s2(x1) # 1,256,8,56,56
+ x3 = self.s3(x2) # 1,512,8,28, 28
+ x = self.s4(x3) # 1,1024,8,14,14
+ x = self.upsample(x) # 1,1024, 8, 28, 28
+ x = self.concat(x3,x)# 1,1024+512, 8, 28, 28
+ x=[self.forward_branch(x,x1,x2,label) for label in self.labels]
+ x=torch.cat(x,1)
+ out = x.mean([3, 4]).view(-1, 1)*100
+ out = self.linear(out)
+ out = out.view(x.size(0), -1)
+ return x,out
+
+
+
+ def forward_branch(self,x,x1,x2,label):
+ t4=getattr(self,f"t4_{label}")
+ x = t4(x[0])# 1,512, 8, 28, 28
+ x = self.upsample([x]) # 1,512, 8, 56, 56
+ x = self.concat(x2,x)# 1,256+512, 8, 56, 56
+ t3= getattr(self,f"t3_{label}")
+ x = t3(x[0]) # 1,256, 8, 56, 56
+ x = self.concat(x1,[x]) # 1,320, 8, 56, 56
+ conv1x1=getattr(self,f"conv1x1_{label}")
+ x = conv1x1(x[0]) # 1,2,8,56,56
+ return x
+
+
+
+@MODEL_REGISTRY.register()
+class ResUNetLight(nn.Module):
+ """
+ ResNet model builder. It builds a ResNet like network backbone without
+ lateral connection (C2D, I3D, Slow).
+
+ Christoph Feichtenhofer, Haoqi Fan, Jitendra Malik, and Kaiming He.
+ "SlowFast networks for video recognition."
+ https://arxiv.org/pdf/1812.03982.pdf
+
+ Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He.
+ "Non-local neural networks."
+ https://arxiv.org/pdf/1711.07971.pdf
+ """
+
+ def __init__(self, cfg):
+ """
+ The `__init__` method of any subclass should also contain these
+ arguments.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ super(ResUNetLight, self).__init__()
+ self.norm_module = get_norm(cfg)
+ self.enable_detection = cfg.DETECTION.ENABLE
+ self.enable_jitter = cfg.JITTER.ENABLE
+ self.num_pathways = 1
+ assert cfg.DATA.TRAIN_CROP_SIZE == cfg.DATA.TEST_CROP_SIZE
+ self.image_size = cfg.DATA.TRAIN_CROP_SIZE
+ self.clip_size = cfg.DATA.NUM_FRAMES
+ self._construct_network(cfg)
+ init_helper.init_weights(
+ self, cfg.MODEL.FC_INIT_STD, cfg.RESNET.ZERO_INIT_FINAL_BN
+ )
+
+ def _construct_network(self, cfg):
+ """
+ Builds a single pathway ResNet model.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ assert cfg.MODEL.ARCH in _POOL1.keys()
+ pool_size = _POOL1[cfg.MODEL.ARCH]
+ self.cfg = cfg
+ assert len({len(pool_size), self.num_pathways}) == 1
+ assert cfg.RESNET.DEPTH in _MODEL_STAGE_DEPTH.keys()
+
+ (d2, d3, d4, d5) = _MODEL_STAGE_DEPTH[cfg.RESNET.DEPTH]
+
+ num_groups = cfg.RESNET.NUM_GROUPS
+ width_per_group = cfg.RESNET.WIDTH_PER_GROUP
+ dim_inner = num_groups * width_per_group
+
+ temp_kernel = _TEMPORAL_KERNEL_BASIS[cfg.MODEL.ARCH]
+
+ self.s1 = stem_helper.VideoModelStem(
+ dim_in=cfg.DATA.INPUT_CHANNEL_NUM,
+ dim_out=[width_per_group],
+ kernel=[temp_kernel[0][0] + [7, 7]],
+ stride=[[1, 2, 2]],
+ padding=[[temp_kernel[0][0][0] // 2, 3, 3]],
+ norm_module=self.norm_module,
+ )
+
+ self.s2 = resnet_helper.ResStage(
+ dim_in=[width_per_group],
+ dim_out=[width_per_group * 4],
+ dim_inner=[dim_inner],
+ temp_kernel_sizes=temp_kernel[1],
+ stride=cfg.RESNET.SPATIAL_STRIDES[0],
+ num_blocks=[d2],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[0],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[0],
+ nonlocal_group=cfg.NONLOCAL.GROUP[0],
+ nonlocal_pool=cfg.NONLOCAL.POOL[0],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[0],
+ norm_module=self.norm_module,
+ )
+
+ for pathway in range(self.num_pathways):
+ pool = nn.MaxPool3d(
+ kernel_size=pool_size[pathway],
+ stride=pool_size[pathway],
+ padding=[0, 0, 0],
+ )
+ self.add_module("pathway{}_pool".format(pathway), pool)
+
+ self.s3 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 4],
+ dim_out=[width_per_group * 8],
+ dim_inner=[dim_inner * 2],
+ temp_kernel_sizes=temp_kernel[2],
+ stride=cfg.RESNET.SPATIAL_STRIDES[1],
+ num_blocks=[d3],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[1],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[1],
+ nonlocal_group=cfg.NONLOCAL.GROUP[1],
+ nonlocal_pool=cfg.NONLOCAL.POOL[1],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[1],
+ norm_module=self.norm_module,
+ )
+
+ self.s4 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 8],
+ dim_out=[width_per_group * 16],
+ dim_inner=[dim_inner * 4],
+ temp_kernel_sizes=temp_kernel[3],
+ stride=cfg.RESNET.SPATIAL_STRIDES[2],
+ num_blocks=[d4],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[2],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[2],
+ nonlocal_group=cfg.NONLOCAL.GROUP[2],
+ nonlocal_pool=cfg.NONLOCAL.POOL[2],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[2],
+ norm_module=self.norm_module,
+ )
+
+ # self.s5 = resnet_helper.ResStage(
+ # dim_in=[width_per_group * 16],
+ # dim_out=[width_per_group * 32],
+ # dim_inner=[dim_inner * 8],
+ # temp_kernel_sizes=temp_kernel[4],
+ # stride=cfg.RESNET.SPATIAL_STRIDES[3],
+ # num_blocks=[d5],
+ # num_groups=[num_groups],
+ # num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[3],
+ # nonlocal_inds=cfg.NONLOCAL.LOCATION[3],
+ # nonlocal_group=cfg.NONLOCAL.GROUP[3],
+ # nonlocal_pool=cfg.NONLOCAL.POOL[3],
+ # instantiation=cfg.NONLOCAL.INSTANTIATION,
+ # trans_func_name=cfg.RESNET.TRANS_FUNC,
+ # stride_1x1=cfg.RESNET.STRIDE_1X1,
+ # inplace_relu=cfg.RESNET.INPLACE_RELU,
+ # dilation=cfg.RESNET.SPATIAL_DILATIONS[3],
+ # norm_module=self.norm_module,
+ # )
+ self.labels=["rotate","light"]
+ self.dual_define("t4",self.labels,LightDecoderBlock(width_per_group * 16,width_per_group * 8,width_per_group * 4))
+ self.dual_define("t3",self.labels,LightDecoderBlock(width_per_group * 4,width_per_group * 4, 128))
+ self.dual_define("conv1x1",self.labels,nn.Sequential(
+ nn.Conv3d(128+width_per_group, 1, kernel_size=(1, 1, 1), stride=1, padding=0), nn.Sigmoid()
+ ))
+
+ self.linear = nn.Sequential(nn.Linear(1, 1), nn.Sigmoid())
+
+ def forward_plus(self, x, y, net):
+ return [net(x)[0] + y[0]]
+
+
+ def dual_define(self,name,labels,net):
+ for label in labels:
+ self.add_module(f"{name}_{label}",copy.deepcopy(net))
+
+
+
+ def upsample(self, x, dims=["space"]):
+ ori_size = x[0].shape[2:5]
+ t, h, w = ori_size
+ if "space" in dims:
+ h = 2 * h
+ w = 2 * w
+ if "time" in dims:
+ t = 2 * t
+ size = (t, h, w)
+ return [F.interpolate(x[0], size)]
+
+ def concat(self,x,y):
+ return [torch.cat([x[0],y[0]],1)]
+
+ def get_detach_var(self,x):
+ return [t.detach() for t in x]
+
+ # @torchsnooper.snoop()
+ def forward(self, x, freeze_backbone=False):
+ x1 = self.s1(x) # 1,64,8,56,56
+ x2 = self.s2(x1) # 1,256,8,56,56
+ x3 = self.s3(x2) # 1,512,8,28, 28
+ x = self.s4(x3) # 1,1024,8,14,14
+ assert isinstance(freeze_backbone,bool)
+ if freeze_backbone:
+ x=self.get_detach_var(x)
+ x1=self.get_detach_var(x1)
+ x2=self.get_detach_var(x2)
+ x3=self.get_detach_var(x3)
+
+ x = self.upsample(x) # 1,1024, 8, 28, 28
+ x = self.concat(x3,x)# 1,1024+512, 8, 28, 28
+ x=[self.forward_branch(x,x1,x2,label) for label in self.labels]
+ x=torch.cat(x,1)
+ out = x.mean([3, 4]).view(-1, 1)*100 # 1,2,8,56,56
+ out = self.linear(out)
+ out = out.view(x.size(0), -1)
+ return x,out
+
+
+
+ def forward_branch(self,x,x1,x2,label):
+ t4=getattr(self,f"t4_{label}")
+ x = t4(x[0])# 1,256, 8, 28, 28
+ x = self.upsample([x]) # 1,256, 8, 56, 56
+ x = self.concat(x2,x)# 1,256+256, 8, 56, 56
+ t3= getattr(self,f"t3_{label}")
+ x = t3(x[0]) # 1,128, 8, 56, 56
+ x = self.concat(x1,[x]) # 1,192, 8, 56, 56
+ conv1x1=getattr(self,f"conv1x1_{label}")
+ x = conv1x1(x[0]) # 1,2,8,56,56
+ return x
+
+
+
+@MODEL_REGISTRY.register()
+class ResUNetLightFix(nn.Module):
+ """
+ ResNet model builder. It builds a ResNet like network backbone without
+ lateral connection (C2D, I3D, Slow).
+
+ Christoph Feichtenhofer, Haoqi Fan, Jitendra Malik, and Kaiming He.
+ "SlowFast networks for video recognition."
+ https://arxiv.org/pdf/1812.03982.pdf
+
+ Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He.
+ "Non-local neural networks."
+ https://arxiv.org/pdf/1711.07971.pdf
+ """
+
+ def __init__(self, cfg):
+ """
+ The `__init__` method of any subclass should also contain these
+ arguments.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ super(ResUNetLightFix, self).__init__()
+ self.norm_module = get_norm(cfg)
+ self.enable_detection = cfg.DETECTION.ENABLE
+ self.enable_jitter = cfg.JITTER.ENABLE
+ self.num_pathways = 1
+ assert cfg.DATA.TRAIN_CROP_SIZE == cfg.DATA.TEST_CROP_SIZE
+ self.image_size = cfg.DATA.TRAIN_CROP_SIZE
+ self.clip_size = cfg.DATA.NUM_FRAMES
+ self._construct_network(cfg)
+ init_helper.init_weights(
+ self, cfg.MODEL.FC_INIT_STD, cfg.RESNET.ZERO_INIT_FINAL_BN
+ )
+
+ def _construct_network(self, cfg):
+ """
+ Builds a single pathway ResNet model.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ assert cfg.MODEL.ARCH in _POOL1.keys()
+ pool_size = _POOL1[cfg.MODEL.ARCH]
+ self.cfg = cfg
+ assert len({len(pool_size), self.num_pathways}) == 1
+ assert cfg.RESNET.DEPTH in _MODEL_STAGE_DEPTH.keys()
+
+ (d2, d3, d4, d5) = _MODEL_STAGE_DEPTH[cfg.RESNET.DEPTH]
+
+ num_groups = cfg.RESNET.NUM_GROUPS
+ width_per_group = cfg.RESNET.WIDTH_PER_GROUP
+ dim_inner = num_groups * width_per_group
+
+ temp_kernel = _TEMPORAL_KERNEL_BASIS[cfg.MODEL.ARCH]
+
+ self.s1 = stem_helper.VideoModelStem(
+ dim_in=cfg.DATA.INPUT_CHANNEL_NUM,
+ dim_out=[width_per_group],
+ kernel=[temp_kernel[0][0] + [7, 7]],
+ stride=[[1, 2, 2]],
+ padding=[[temp_kernel[0][0][0] // 2, 3, 3]],
+ norm_module=self.norm_module,
+ )
+
+ self.s2 = resnet_helper.ResStage(
+ dim_in=[width_per_group],
+ dim_out=[width_per_group * 4],
+ dim_inner=[dim_inner],
+ temp_kernel_sizes=temp_kernel[1],
+ stride=cfg.RESNET.SPATIAL_STRIDES[0],
+ num_blocks=[d2],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[0],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[0],
+ nonlocal_group=cfg.NONLOCAL.GROUP[0],
+ nonlocal_pool=cfg.NONLOCAL.POOL[0],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[0],
+ norm_module=self.norm_module,
+ )
+
+ for pathway in range(self.num_pathways):
+ pool = nn.MaxPool3d(
+ kernel_size=pool_size[pathway],
+ stride=pool_size[pathway],
+ padding=[0, 0, 0],
+ )
+ self.add_module("pathway{}_pool".format(pathway), pool)
+
+ self.s3 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 4],
+ dim_out=[width_per_group * 8],
+ dim_inner=[dim_inner * 2],
+ temp_kernel_sizes=temp_kernel[2],
+ stride=cfg.RESNET.SPATIAL_STRIDES[1],
+ num_blocks=[d3],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[1],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[1],
+ nonlocal_group=cfg.NONLOCAL.GROUP[1],
+ nonlocal_pool=cfg.NONLOCAL.POOL[1],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[1],
+ norm_module=self.norm_module,
+ )
+
+ self.s4 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 8],
+ dim_out=[width_per_group * 16],
+ dim_inner=[dim_inner * 4],
+ temp_kernel_sizes=temp_kernel[3],
+ stride=cfg.RESNET.SPATIAL_STRIDES[2],
+ num_blocks=[d4],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[2],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[2],
+ nonlocal_group=cfg.NONLOCAL.GROUP[2],
+ nonlocal_pool=cfg.NONLOCAL.POOL[2],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[2],
+ norm_module=self.norm_module,
+ )
+
+ # self.s5 = resnet_helper.ResStage(
+ # dim_in=[width_per_group * 16],
+ # dim_out=[width_per_group * 32],
+ # dim_inner=[dim_inner * 8],
+ # temp_kernel_sizes=temp_kernel[4],
+ # stride=cfg.RESNET.SPATIAL_STRIDES[3],
+ # num_blocks=[d5],
+ # num_groups=[num_groups],
+ # num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[3],
+ # nonlocal_inds=cfg.NONLOCAL.LOCATION[3],
+ # nonlocal_group=cfg.NONLOCAL.GROUP[3],
+ # nonlocal_pool=cfg.NONLOCAL.POOL[3],
+ # instantiation=cfg.NONLOCAL.INSTANTIATION,
+ # trans_func_name=cfg.RESNET.TRANS_FUNC,
+ # stride_1x1=cfg.RESNET.STRIDE_1X1,
+ # inplace_relu=cfg.RESNET.INPLACE_RELU,
+ # dilation=cfg.RESNET.SPATIAL_DILATIONS[3],
+ # norm_module=self.norm_module,
+ # )
+ self.labels=["rotate","light","skip"]
+ self.dual_define("t4",self.labels,LightDecoderBlock(width_per_group * 16,width_per_group * 8,width_per_group * 4))
+ self.dual_define("t3",self.labels,LightDecoderBlock(width_per_group * 4,width_per_group * 4, 128))
+ self.dual_define("conv1x1",self.labels,nn.Sequential(
+ nn.Conv3d(128+width_per_group, 64, kernel_size=(1, 1, 1), stride=1, padding=0),
+ nn.BatchNorm3d(64),
+ nn.ReLU(),
+ nn.Conv3d(64, 1, kernel_size=(1, 1, 1), stride=1, padding=0),
+ ))
+
+ self.linear = nn.Sequential(nn.Linear(1, 1))
+
+ def forward_plus(self, x, y, net):
+ return [net(x)[0] + y[0]]
+
+
+ def dual_define(self,name,labels,net):
+ for label in labels:
+ self.add_module(f"{name}_{label}",copy.deepcopy(net))
+
+
+
+ def upsample(self, x, dims=["space"]):
+ ori_size = x[0].shape[2:5]
+ t, h, w = ori_size
+ if "space" in dims:
+ h = 2 * h
+ w = 2 * w
+ if "time" in dims:
+ t = 2 * t
+ size = (t, h, w)
+ return [F.interpolate(x[0], size)]
+
+ def concat(self,x,y):
+ return [torch.cat([x[0],y[0]],1)]
+
+ def get_detach_var(self,x):
+ return [t.detach() for t in x]
+
+ # @torchsnooper.snoop()
+ def forward(self, x, freeze_backbone=False):
+ x1 = self.s1(x) # 1,64,8,56,56
+ x2 = self.s2(x1) # 1,256,8,56,56
+ x3 = self.s3(x2) # 1,512,8,28, 28
+ x = self.s4(x3) # 1,1024,8,14,14
+ assert isinstance(freeze_backbone,bool)
+ if freeze_backbone:
+ x=self.get_detach_var(x)
+ x1=self.get_detach_var(x1)
+ x2=self.get_detach_var(x2)
+ x3=self.get_detach_var(x3)
+
+ x = self.upsample(x) # 1,1024, 8, 28, 28
+ x = self.concat(x3,x)# 1,1024+512, 8, 28, 28
+ x=[self.forward_branch(x,x1,x2,label) for label in self.labels]
+ x=torch.cat(x,1)
+ x=torch.sigmoid(x)
+ out = x.mean([3, 4]).view(-1, 1)*100 # 1,2,8,56,56
+ out = self.linear(out)
+ out = out.view(x.size(0), -1)
+ out = torch.sigmoid(out)
+ return x,out
+
+
+
+ def forward_branch(self,x,x1,x2,label):
+ t4=getattr(self,f"t4_{label}")
+ x = t4(x[0])# 1,256, 8, 28, 28
+ x = self.upsample([x]) # 1,256, 8, 56, 56
+ x = self.concat(x2,x)# 1,256+256, 8, 56, 56
+ t3= getattr(self,f"t3_{label}")
+ x = t3(x[0]) # 1,128, 8, 56, 56
+ x = self.concat(x1,[x]) # 1,192, 8, 56, 56
+ conv1x1=getattr(self,f"conv1x1_{label}")
+ x = conv1x1(x[0]) # 1,2,8,56,56
+ return x
+
+
+
+@MODEL_REGISTRY.register()
+class ResUNetContinus(nn.Module):
+ """
+ ResNet model builder. It builds a ResNet like network backbone without
+ lateral connection (C2D, I3D, Slow).
+
+ Christoph Feichtenhofer, Haoqi Fan, Jitendra Malik, and Kaiming He.
+ "SlowFast networks for video recognition."
+ https://arxiv.org/pdf/1812.03982.pdf
+
+ Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He.
+ "Non-local neural networks."
+ https://arxiv.org/pdf/1711.07971.pdf
+ """
+
+ def __init__(self, cfg):
+ """
+ The `__init__` method of any subclass should also contain these
+ arguments.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ super(ResUNetContinus, self).__init__()
+ self.norm_module = get_norm(cfg)
+ self.enable_detection = cfg.DETECTION.ENABLE
+ self.enable_jitter = cfg.JITTER.ENABLE
+ self.num_pathways = 1
+ assert cfg.DATA.TRAIN_CROP_SIZE == cfg.DATA.TEST_CROP_SIZE
+ self.image_size = cfg.DATA.TRAIN_CROP_SIZE
+ self.clip_size = cfg.DATA.NUM_FRAMES
+ self._construct_network(cfg)
+ init_helper.init_weights(
+ self, cfg.MODEL.FC_INIT_STD, cfg.RESNET.ZERO_INIT_FINAL_BN
+ )
+
+ def _construct_network(self, cfg):
+ """
+ Builds a single pathway ResNet model.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ assert cfg.MODEL.ARCH in _POOL1.keys()
+ pool_size = _POOL1[cfg.MODEL.ARCH]
+ self.cfg = cfg
+ assert len({len(pool_size), self.num_pathways}) == 1
+ assert cfg.RESNET.DEPTH in _MODEL_STAGE_DEPTH.keys()
+
+ (d2, d3, d4, d5) = _MODEL_STAGE_DEPTH[cfg.RESNET.DEPTH]
+
+ num_groups = cfg.RESNET.NUM_GROUPS
+ width_per_group = cfg.RESNET.WIDTH_PER_GROUP
+ dim_inner = num_groups * width_per_group
+
+ temp_kernel = _TEMPORAL_KERNEL_BASIS[cfg.MODEL.ARCH]
+
+ self.s1 = stem_helper.VideoModelStem(
+ dim_in=cfg.DATA.INPUT_CHANNEL_NUM,
+ dim_out=[width_per_group],
+ kernel=[temp_kernel[0][0] + [7, 7]],
+ stride=[[1, 2, 2]],
+ padding=[[temp_kernel[0][0][0] // 2, 3, 3]],
+ norm_module=self.norm_module,
+ )
+
+ self.s2 = resnet_helper.ResStage(
+ dim_in=[width_per_group],
+ dim_out=[width_per_group * 4],
+ dim_inner=[dim_inner],
+ temp_kernel_sizes=temp_kernel[1],
+ stride=cfg.RESNET.SPATIAL_STRIDES[0],
+ num_blocks=[d2],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[0],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[0],
+ nonlocal_group=cfg.NONLOCAL.GROUP[0],
+ nonlocal_pool=cfg.NONLOCAL.POOL[0],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[0],
+ norm_module=self.norm_module,
+ )
+
+ for pathway in range(self.num_pathways):
+ pool = nn.MaxPool3d(
+ kernel_size=pool_size[pathway],
+ stride=pool_size[pathway],
+ padding=[0, 0, 0],
+ )
+ self.add_module("pathway{}_pool".format(pathway), pool)
+
+ self.s3 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 4],
+ dim_out=[width_per_group * 8],
+ dim_inner=[dim_inner * 2],
+ temp_kernel_sizes=temp_kernel[2],
+ stride=cfg.RESNET.SPATIAL_STRIDES[1],
+ num_blocks=[d3],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[1],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[1],
+ nonlocal_group=cfg.NONLOCAL.GROUP[1],
+ nonlocal_pool=cfg.NONLOCAL.POOL[1],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[1],
+ norm_module=self.norm_module,
+ )
+
+ self.s4 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 8],
+ dim_out=[width_per_group * 16],
+ dim_inner=[dim_inner * 4],
+ temp_kernel_sizes=temp_kernel[3],
+ stride=cfg.RESNET.SPATIAL_STRIDES[2],
+ num_blocks=[d4],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[2],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[2],
+ nonlocal_group=cfg.NONLOCAL.GROUP[2],
+ nonlocal_pool=cfg.NONLOCAL.POOL[2],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[2],
+ norm_module=self.norm_module,
+ )
+
+ # self.s5 = resnet_helper.ResStage(
+ # dim_in=[width_per_group * 16],
+ # dim_out=[width_per_group * 32],
+ # dim_inner=[dim_inner * 8],
+ # temp_kernel_sizes=temp_kernel[4],
+ # stride=cfg.RESNET.SPATIAL_STRIDES[3],
+ # num_blocks=[d5],
+ # num_groups=[num_groups],
+ # num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[3],
+ # nonlocal_inds=cfg.NONLOCAL.LOCATION[3],
+ # nonlocal_group=cfg.NONLOCAL.GROUP[3],
+ # nonlocal_pool=cfg.NONLOCAL.POOL[3],
+ # instantiation=cfg.NONLOCAL.INSTANTIATION,
+ # trans_func_name=cfg.RESNET.TRANS_FUNC,
+ # stride_1x1=cfg.RESNET.STRIDE_1X1,
+ # inplace_relu=cfg.RESNET.INPLACE_RELU,
+ # dilation=cfg.RESNET.SPATIAL_DILATIONS[3],
+ # norm_module=self.norm_module,
+ # )
+ self.labels=["all"]
+ self.dual_define("t4",self.labels,LightDecoderBlock(width_per_group * 16,width_per_group * 8,width_per_group * 4))
+ self.dual_define("t3",self.labels,LightDecoderBlock(width_per_group * 4,width_per_group * 4, 128))
+ self.dual_define("conv1x1",self.labels,nn.Sequential(
+ nn.Conv3d(128+width_per_group, 64, kernel_size=(1, 1, 1), stride=1, padding=0),
+ nn.BatchNorm3d(64),
+ nn.ReLU(),
+ nn.Conv3d(64, 1, kernel_size=(1, 1, 1), stride=1, padding=0),
+ ))
+
+ self.linear = nn.Sequential(nn.Linear(1, 1))
+
+ def forward_plus(self, x, y, net):
+ return [net(x)[0] + y[0]]
+
+
+ def dual_define(self,name,labels,net):
+ for label in labels:
+ self.add_module(f"{name}_{label}",copy.deepcopy(net))
+
+
+
+ def upsample(self, x, dims=["space"]):
+ ori_size = x[0].shape[2:5]
+ t, h, w = ori_size
+ if "space" in dims:
+ h = 2 * h
+ w = 2 * w
+ if "time" in dims:
+ t = 2 * t
+ size = (t, h, w)
+ return [F.interpolate(x[0], size)]
+
+ def concat(self,x,y):
+ return [torch.cat([x[0],y[0]],1)]
+
+ def get_detach_var(self,x):
+ return [t.detach() for t in x]
+
+ # @torchsnooper.snoop()
+ def forward(self, x, freeze_backbone=False):
+ x1 = self.s1(x) # 1,64,8,56,56
+ x2 = self.s2(x1) # 1,256,8,56,56
+ x3 = self.s3(x2) # 1,512,8,28, 28
+ x = self.s4(x3) # 1,1024,8,14,14
+ assert isinstance(freeze_backbone,bool)
+ if freeze_backbone:
+ x=self.get_detach_var(x)
+ x1=self.get_detach_var(x1)
+ x2=self.get_detach_var(x2)
+ x3=self.get_detach_var(x3)
+
+ x = self.upsample(x) # 1,1024, 8, 28, 28
+ x = self.concat(x3,x)# 1,1024+512, 8, 28, 28
+ x=[self.forward_branch(x,x1,x2,label) for label in self.labels]
+ x=torch.cat(x,1)
+ x=torch.sigmoid(x)
+ out = x.mean([3, 4]).view(-1, 1)*100 # 1,2,8,56,56
+ out = self.linear(out)
+ out = out.view(x.size(0), -1)
+ out = torch.sigmoid(out)
+ return x,out
+
+
+ def forward_branch(self,x,x1,x2,label):
+ t4= getattr(self,f"t4_{label}")
+ x = t4(x[0])# 1,256, 8, 28, 28
+ x = self.upsample([x]) # 1,256, 8, 56, 56
+ x = self.concat(x2,x)# 1,256+256, 8, 56, 56
+ t3= getattr(self,f"t3_{label}")
+ x = t3(x[0]) # 1,128, 8, 56, 56
+ x = self.concat(x1,[x]) # 1,192, 8, 56, 56
+ conv1x1=getattr(self,f"conv1x1_{label}")
+ x = conv1x1(x[0]) # 1,2,8,56,56
+ return x
+
+
+
+
+@MODEL_REGISTRY.register()
+class ResUNetCommon(nn.Module):
+ """
+ ResNet model builder. It builds a ResNet like network backbone without
+ lateral connection (C2D, I3D, Slow).
+
+ Christoph Feichtenhofer, Haoqi Fan, Jitendra Malik, and Kaiming He.
+ "SlowFast networks for video recognition."
+ https://arxiv.org/pdf/1812.03982.pdf
+
+ Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He.
+ "Non-local neural networks."
+ https://arxiv.org/pdf/1711.07971.pdf
+ """
+
+ def __init__(self, cfg):
+ """
+ The `__init__` method of any subclass should also contain these
+ arguments.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ super(ResUNetCommon, self).__init__()
+ self.norm_module = get_norm(cfg)
+ self.enable_detection = cfg.DETECTION.ENABLE
+ self.enable_jitter = cfg.JITTER.ENABLE
+ self.num_pathways = 1
+ assert cfg.DATA.TRAIN_CROP_SIZE == cfg.DATA.TEST_CROP_SIZE
+ self.image_size = cfg.DATA.TRAIN_CROP_SIZE
+ self.clip_size = cfg.DATA.NUM_FRAMES
+ self._construct_network(cfg)
+ init_helper.init_weights(
+ self, cfg.MODEL.FC_INIT_STD, cfg.RESNET.ZERO_INIT_FINAL_BN
+ )
+
+ def _construct_network(self, cfg):
+ """
+ Builds a single pathway ResNet model.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ assert cfg.MODEL.ARCH in _POOL1.keys()
+ pool_size = _POOL1[cfg.MODEL.ARCH]
+ self.cfg = cfg
+ assert len({len(pool_size), self.num_pathways}) == 1
+ assert cfg.RESNET.DEPTH in _MODEL_STAGE_DEPTH.keys()
+
+ (d2, d3, d4, d5) = _MODEL_STAGE_DEPTH[cfg.RESNET.DEPTH]
+
+ num_groups = cfg.RESNET.NUM_GROUPS
+ width_per_group = cfg.RESNET.WIDTH_PER_GROUP
+ dim_inner = num_groups * width_per_group
+
+ temp_kernel = _TEMPORAL_KERNEL_BASIS[cfg.MODEL.ARCH]
+
+ self.s1 = stem_helper.VideoModelStem(
+ dim_in=cfg.DATA.INPUT_CHANNEL_NUM,
+ dim_out=[width_per_group],
+ kernel=[temp_kernel[0][0] + [7, 7]],
+ stride=[[1, 2, 2]],
+ padding=[[temp_kernel[0][0][0] // 2, 3, 3]],
+ norm_module=self.norm_module,
+ )
+
+ self.s2 = resnet_helper.ResStage(
+ dim_in=[width_per_group],
+ dim_out=[width_per_group * 4],
+ dim_inner=[dim_inner],
+ temp_kernel_sizes=temp_kernel[1],
+ stride=cfg.RESNET.SPATIAL_STRIDES[0],
+ num_blocks=[d2],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[0],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[0],
+ nonlocal_group=cfg.NONLOCAL.GROUP[0],
+ nonlocal_pool=cfg.NONLOCAL.POOL[0],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[0],
+ norm_module=self.norm_module,
+ )
+
+ for pathway in range(self.num_pathways):
+ pool = nn.MaxPool3d(
+ kernel_size=pool_size[pathway],
+ stride=pool_size[pathway],
+ padding=[0, 0, 0],
+ )
+ self.add_module("pathway{}_pool".format(pathway), pool)
+
+ self.s3 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 4],
+ dim_out=[width_per_group * 8],
+ dim_inner=[dim_inner * 2],
+ temp_kernel_sizes=temp_kernel[2],
+ stride=cfg.RESNET.SPATIAL_STRIDES[1],
+ num_blocks=[d3],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[1],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[1],
+ nonlocal_group=cfg.NONLOCAL.GROUP[1],
+ nonlocal_pool=cfg.NONLOCAL.POOL[1],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[1],
+ norm_module=self.norm_module,
+ )
+
+ self.s4 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 8],
+ dim_out=[width_per_group * 16],
+ dim_inner=[dim_inner * 4],
+ temp_kernel_sizes=temp_kernel[3],
+ stride=cfg.RESNET.SPATIAL_STRIDES[2],
+ num_blocks=[d4],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[2],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[2],
+ nonlocal_group=cfg.NONLOCAL.GROUP[2],
+ nonlocal_pool=cfg.NONLOCAL.POOL[2],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[2],
+ norm_module=self.norm_module,
+ )
+
+ # self.s5 = resnet_helper.ResStage(
+ # dim_in=[width_per_group * 16],
+ # dim_out=[width_per_group * 32],
+ # dim_inner=[dim_inner * 8],
+ # temp_kernel_sizes=temp_kernel[4],
+ # stride=cfg.RESNET.SPATIAL_STRIDES[3],
+ # num_blocks=[d5],
+ # num_groups=[num_groups],
+ # num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[3],
+ # nonlocal_inds=cfg.NONLOCAL.LOCATION[3],
+ # nonlocal_group=cfg.NONLOCAL.GROUP[3],
+ # nonlocal_pool=cfg.NONLOCAL.POOL[3],
+ # instantiation=cfg.NONLOCAL.INSTANTIATION,
+ # trans_func_name=cfg.RESNET.TRANS_FUNC,
+ # stride_1x1=cfg.RESNET.STRIDE_1X1,
+ # inplace_relu=cfg.RESNET.INPLACE_RELU,
+ # dilation=cfg.RESNET.SPATIAL_DILATIONS[3],
+ # norm_module=self.norm_module,
+ # )
+ self.labels=cfg.RESNET.LABELS
+ self.dual_define("t4",self.labels,LightDecoderBlock(width_per_group * 16,width_per_group * 8,width_per_group * 4))
+ self.dual_define("t3",self.labels,LightDecoderBlock(width_per_group * 4,width_per_group * 4, 128))
+ self.dual_define("conv1x1",self.labels,nn.Sequential(
+ nn.Conv3d(128+width_per_group, 64, kernel_size=(1, 1, 1), stride=1, padding=0),
+ nn.BatchNorm3d(64),
+ nn.ReLU(),
+ nn.Conv3d(64, 1, kernel_size=(1, 1, 1), stride=1, padding=0),
+ ))
+
+ self.linear = nn.Linear(1, 2)
+
+ def forward_plus(self, x, y, net):
+ return [net(x)[0] + y[0]]
+
+
+ def dual_define(self,name,labels,net):
+ for label in labels:
+ self.add_module(f"{name}_{label}",copy.deepcopy(net))
+
+
+ def upsample(self, x, dims=["space"]):
+ ori_size = x[0].shape[2:5]
+ t, h, w = ori_size
+ if "space" in dims:
+ h = 2 * h
+ w = 2 * w
+ if "time" in dims:
+ t = 2 * t
+ size = (t, h, w)
+ return [F.interpolate(x[0], size)]
+
+ def concat(self,x,y):
+ return [torch.cat([x[0],y[0]],1)]
+
+ def get_detach_var(self,x):
+ return [t.detach() for t in x]
+
+ # @torchsnooper.snoop()
+ def forward(self, x, freeze_backbone=False):
+ x = self.get_detach_var(x)
+ x1 = self.s1(x) # 1,64,8,56,56
+ x2 = self.s2(x1) # 1,256,8,56,56
+ x3 = self.s3(x2) # 1,512,8,28, 28
+ feat= self.s4(x3) # 1,1024,8,14,14
+ assert isinstance(freeze_backbone,bool)
+ if freeze_backbone:
+ feat=self.get_detach_var(feat)
+ x1=self.get_detach_var(x1)
+ x2=self.get_detach_var(x2)
+ x3=self.get_detach_var(x3)
+
+ feat = self.upsample(feat) # 1,1024, 8, 28, 28
+ feat = self.concat(x3,feat)# 1,1024+512, 8, 28, 28
+ reg_out=[self.forward_branch(feat,x1,x2,label) for label in self.labels]
+ reg_out=torch.cat(reg_out,1)
+ reg_out=torch.sigmoid(reg_out)
+ class_out = reg_out.mean([3, 4]).view(-1, 1)*100 # 1,2,8,56,56
+ class_out = self.linear(class_out)
+ class_out = class_out.view(reg_out.size(0),len(self.labels),-1)
+ class_out = class_out
+ return reg_out,class_out
+
+
+ def forward_branch(self,feat,x1,x2,label):
+ t4= getattr(self,f"t4_{label}")
+ feat = t4(feat[0])# 1,256, 8, 28, 28
+ feat = self.upsample([feat]) # 1,256, 8, 56, 56
+ feat = self.concat(x2,feat)# 1,256+256, 8, 56, 56
+ t3= getattr(self,f"t3_{label}")
+ feat = t3(feat[0]) # 1,128, 8, 56, 56
+ feat = self.concat(x1,[feat]) # 1,192, 8, 56, 56
+ conv1x1=getattr(self,f"conv1x1_{label}")
+ feat = conv1x1(feat[0]) # 1,2,8,56,56
+ return feat
+
+
+
+
+@MODEL_REGISTRY.register()
+class ResUNetCommon2(nn.Module):
+ """
+ ResNet model builder. It builds a ResNet like network backbone without
+ lateral connection (C2D, I3D, Slow).
+
+ Christoph Feichtenhofer, Haoqi Fan, Jitendra Malik, and Kaiming He.
+ "SlowFast networks for video recognition."
+ https://arxiv.org/pdf/1812.03982.pdf
+
+ Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He.
+ "Non-local neural networks."
+ https://arxiv.org/pdf/1711.07971.pdf
+ """
+
+ def __init__(self, cfg):
+ """
+ The `__init__` method of any subclass should also contain these
+ arguments.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ super(ResUNetCommon2, self).__init__()
+ self.norm_module = get_norm(cfg)
+ self.enable_detection = cfg.DETECTION.ENABLE
+ self.enable_jitter = cfg.JITTER.ENABLE
+ self.num_pathways = 1
+ assert cfg.DATA.TRAIN_CROP_SIZE == cfg.DATA.TEST_CROP_SIZE
+ self.image_size = cfg.DATA.TRAIN_CROP_SIZE
+ self.clip_size = cfg.DATA.NUM_FRAMES
+ self._construct_network(cfg)
+ init_helper.init_weights(
+ self, cfg.MODEL.FC_INIT_STD, cfg.RESNET.ZERO_INIT_FINAL_BN
+ )
+
+ def _construct_network(self, cfg):
+ """
+ Builds a single pathway ResNet model.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ assert cfg.MODEL.ARCH in _POOL1.keys()
+ pool_size = _POOL1[cfg.MODEL.ARCH]
+ self.cfg = cfg
+ assert len({len(pool_size), self.num_pathways}) == 1
+ assert cfg.RESNET.DEPTH in _MODEL_STAGE_DEPTH.keys()
+
+ (d2, d3, d4, d5) = _MODEL_STAGE_DEPTH[cfg.RESNET.DEPTH]
+
+ num_groups = cfg.RESNET.NUM_GROUPS
+ width_per_group = cfg.RESNET.WIDTH_PER_GROUP
+ dim_inner = num_groups * width_per_group
+
+ temp_kernel = _TEMPORAL_KERNEL_BASIS[cfg.MODEL.ARCH]
+
+ self.s1 = stem_helper.VideoModelStem(
+ dim_in=cfg.DATA.INPUT_CHANNEL_NUM,
+ dim_out=[width_per_group],
+ kernel=[temp_kernel[0][0] + [7, 7]],
+ stride=[[1, 2, 2]],
+ padding=[[temp_kernel[0][0][0] // 2, 3, 3]],
+ norm_module=self.norm_module,
+ )
+
+ self.s2 = resnet_helper.ResStage(
+ dim_in=[width_per_group],
+ dim_out=[width_per_group * 4],
+ dim_inner=[dim_inner],
+ temp_kernel_sizes=temp_kernel[1],
+ stride=cfg.RESNET.SPATIAL_STRIDES[0],
+ num_blocks=[d2],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[0],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[0],
+ nonlocal_group=cfg.NONLOCAL.GROUP[0],
+ nonlocal_pool=cfg.NONLOCAL.POOL[0],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[0],
+ norm_module=self.norm_module,
+ )
+
+ for pathway in range(self.num_pathways):
+ pool = nn.MaxPool3d(
+ kernel_size=pool_size[pathway],
+ stride=pool_size[pathway],
+ padding=[0, 0, 0],
+ )
+ self.add_module("pathway{}_pool".format(pathway), pool)
+
+ self.s3 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 4],
+ dim_out=[width_per_group * 8],
+ dim_inner=[dim_inner * 2],
+ temp_kernel_sizes=temp_kernel[2],
+ stride=cfg.RESNET.SPATIAL_STRIDES[1],
+ num_blocks=[d3],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[1],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[1],
+ nonlocal_group=cfg.NONLOCAL.GROUP[1],
+ nonlocal_pool=cfg.NONLOCAL.POOL[1],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[1],
+ norm_module=self.norm_module,
+ )
+
+ self.s4 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 8],
+ dim_out=[width_per_group * 16],
+ dim_inner=[dim_inner * 4],
+ temp_kernel_sizes=temp_kernel[3],
+ stride=cfg.RESNET.SPATIAL_STRIDES[2],
+ num_blocks=[d4],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[2],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[2],
+ nonlocal_group=cfg.NONLOCAL.GROUP[2],
+ nonlocal_pool=cfg.NONLOCAL.POOL[2],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[2],
+ norm_module=self.norm_module,
+ )
+
+ # self.s5 = resnet_helper.ResStage(
+ # dim_in=[width_per_group * 16],
+ # dim_out=[width_per_group * 32],
+ # dim_inner=[dim_inner * 8],
+ # temp_kernel_sizes=temp_kernel[4],
+ # stride=cfg.RESNET.SPATIAL_STRIDES[3],
+ # num_blocks=[d5],
+ # num_groups=[num_groups],
+ # num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[3],
+ # nonlocal_inds=cfg.NONLOCAL.LOCATION[3],
+ # nonlocal_group=cfg.NONLOCAL.GROUP[3],
+ # nonlocal_pool=cfg.NONLOCAL.POOL[3],
+ # instantiation=cfg.NONLOCAL.INSTANTIATION,
+ # trans_func_name=cfg.RESNET.TRANS_FUNC,
+ # stride_1x1=cfg.RESNET.STRIDE_1X1,
+ # inplace_relu=cfg.RESNET.INPLACE_RELU,
+ # dilation=cfg.RESNET.SPATIAL_DILATIONS[3],
+ # norm_module=self.norm_module,
+ # )
+ self.labels=cfg.RESNET.LABELS
+ self.dual_define("t4",self.labels,LightDecoderBlock(width_per_group * 16,width_per_group * 8,width_per_group * 4))
+ self.dual_define("t3",self.labels,LightDecoderBlock(width_per_group * 4,width_per_group * 4, 128))
+ self.dual_define("conv1x1",self.labels,nn.Sequential(
+ nn.Conv3d(128+width_per_group, 64, kernel_size=(1, 1, 1), stride=1, padding=0),
+ nn.BatchNorm3d(64),
+ nn.ReLU(),
+ nn.Conv3d(64, 1, kernel_size=(1, 1, 1), stride=1, padding=0),
+ ))
+
+ self.linear = nn.Linear(1, 1)
+
+ def forward_plus(self, x, y, net):
+ return [net(x)[0] + y[0]]
+
+
+ def dual_define(self,name,labels,net):
+ for label in labels:
+ self.add_module(f"{name}_{label}",copy.deepcopy(net))
+
+
+ def upsample(self, x, dims=["space"]):
+ ori_size = x[0].shape[2:5]
+ t, h, w = ori_size
+ if "space" in dims:
+ h = 2 * h
+ w = 2 * w
+ if "time" in dims:
+ t = 2 * t
+ size = (t, h, w)
+ return [F.interpolate(x[0], size)]
+
+ def concat(self,x,y):
+ return [torch.cat([x[0],y[0]],1)]
+
+ def get_detach_var(self,x):
+ return [t.detach() for t in x]
+
+ # @torchsnooper.snoop()
+ def forward(self, x, freeze_backbone=False):
+ x = self.get_detach_var(x)
+ x1 = self.s1(x) # 1,64,8,56,56
+ x2 = self.s2(x1) # 1,256,8,56,56
+ x3 = self.s3(x2) # 1,512,8,28, 28
+ feat= self.s4(x3) # 1,1024,8,14,14
+ assert isinstance(freeze_backbone,bool)
+ if freeze_backbone:
+ feat=self.get_detach_var(feat)
+ x1=self.get_detach_var(x1)
+ x2=self.get_detach_var(x2)
+ x3=self.get_detach_var(x3)
+
+ feat = self.upsample(feat) # 1,1024, 8, 28, 28
+ feat = self.concat(x3,feat)# 1,1024+512, 8, 28, 28
+ reg_out=[self.forward_branch(feat,x1,x2,label) for label in self.labels]
+ reg_out=torch.cat(reg_out,1)
+ reg_out=torch.sigmoid(reg_out)
+ class_out = reg_out.mean([3, 4]).view(-1, 1)*100 # 1,2,8,56,56
+ class_out = self.linear(class_out)
+ class_out = class_out.view(reg_out.size(0),len(self.labels),-1)
+ class_out = torch.sigmoid(class_out)
+ return reg_out,class_out
+
+
+ def forward_branch(self,feat,x1,x2,label):
+ t4= getattr(self,f"t4_{label}")
+ feat = t4(feat[0])# 1,256, 8, 28, 28
+ feat = self.upsample([feat]) # 1,256, 8, 56, 56
+ feat = self.concat(x2,feat)# 1,256+256, 8, 56, 56
+ t3= getattr(self,f"t3_{label}")
+ feat = t3(feat[0]) # 1,128, 8, 56, 56
+ feat = self.concat(x1,[feat]) # 1,192, 8, 56, 56
+ conv1x1=getattr(self,f"conv1x1_{label}")
+ feat = conv1x1(feat[0]) # 1,2,8,56,56
+ return feat
+
+
+
+@MODEL_REGISTRY.register()
+class ResUNetStrong(nn.Module):
+ """
+ ResNet model builder. It builds a ResNet like network backbone without
+ lateral connection (C2D, I3D, Slow).
+
+ Christoph Feichtenhofer, Haoqi Fan, Jitendra Malik, and Kaiming He.
+ "SlowFast networks for video recognition."
+ https://arxiv.org/pdf/1812.03982.pdf
+
+ Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He.
+ "Non-local neural networks."
+ https://arxiv.org/pdf/1711.07971.pdf
+ """
+
+ def __init__(self, cfg):
+ """
+ The `__init__` method of any subclass should also contain these
+ arguments.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ super(ResUNetStrong, self).__init__()
+ self.norm_module = get_norm(cfg)
+ self.enable_detection = cfg.DETECTION.ENABLE
+ self.enable_jitter = cfg.JITTER.ENABLE
+ self.num_pathways = 1
+ assert cfg.DATA.TRAIN_CROP_SIZE == cfg.DATA.TEST_CROP_SIZE
+ self.image_size = cfg.DATA.TRAIN_CROP_SIZE
+ self.clip_size = cfg.DATA.NUM_FRAMES
+ self._construct_network(cfg)
+ init_helper.init_weights(
+ self, cfg.MODEL.FC_INIT_STD, cfg.RESNET.ZERO_INIT_FINAL_BN
+ )
+
+ def _construct_network(self, cfg):
+ """
+ Builds a single pathway ResNet model.
+
+ Args:
+ cfg (CfgNode): model building configs, details are in the
+ comments of the config file.
+ """
+ assert cfg.MODEL.ARCH in _POOL1.keys()
+ pool_size = _POOL1[cfg.MODEL.ARCH]
+ self.cfg = cfg
+ assert len({len(pool_size), self.num_pathways}) == 1
+ assert cfg.RESNET.DEPTH in _MODEL_STAGE_DEPTH.keys()
+
+ (d2, d3, d4, d5) = _MODEL_STAGE_DEPTH[cfg.RESNET.DEPTH]
+
+ num_groups = cfg.RESNET.NUM_GROUPS
+ width_per_group = cfg.RESNET.WIDTH_PER_GROUP
+ dim_inner = num_groups * width_per_group
+
+ temp_kernel = _TEMPORAL_KERNEL_BASIS[cfg.MODEL.ARCH]
+
+ self.s1 = stem_helper.VideoModelStem(
+ dim_in=cfg.DATA.INPUT_CHANNEL_NUM,
+ dim_out=[width_per_group],
+ kernel=[temp_kernel[0][0] + [7, 7]],
+ stride=[[1, 2, 2]],
+ padding=[[temp_kernel[0][0][0] // 2, 3, 3]],
+ norm_module=self.norm_module,
+ )
+
+ self.s2 = resnet_helper.ResStage(
+ dim_in=[width_per_group],
+ dim_out=[width_per_group * 4],
+ dim_inner=[dim_inner],
+ temp_kernel_sizes=temp_kernel[1],
+ stride=cfg.RESNET.SPATIAL_STRIDES[0],
+ num_blocks=[d2],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[0],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[0],
+ nonlocal_group=cfg.NONLOCAL.GROUP[0],
+ nonlocal_pool=cfg.NONLOCAL.POOL[0],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[0],
+ norm_module=self.norm_module,
+ )
+
+ for pathway in range(self.num_pathways):
+ pool = nn.MaxPool3d(
+ kernel_size=pool_size[pathway],
+ stride=pool_size[pathway],
+ padding=[0, 0, 0],
+ )
+ self.add_module("pathway{}_pool".format(pathway), pool)
+
+ self.s3 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 4],
+ dim_out=[width_per_group * 8],
+ dim_inner=[dim_inner * 2],
+ temp_kernel_sizes=temp_kernel[2],
+ stride=cfg.RESNET.SPATIAL_STRIDES[1],
+ num_blocks=[d3],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[1],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[1],
+ nonlocal_group=cfg.NONLOCAL.GROUP[1],
+ nonlocal_pool=cfg.NONLOCAL.POOL[1],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[1],
+ norm_module=self.norm_module,
+ )
+
+ self.s4 = resnet_helper.ResStage(
+ dim_in=[width_per_group * 8],
+ dim_out=[width_per_group * 16],
+ dim_inner=[dim_inner * 4],
+ temp_kernel_sizes=temp_kernel[3],
+ stride=cfg.RESNET.SPATIAL_STRIDES[2],
+ num_blocks=[d4],
+ num_groups=[num_groups],
+ num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[2],
+ nonlocal_inds=cfg.NONLOCAL.LOCATION[2],
+ nonlocal_group=cfg.NONLOCAL.GROUP[2],
+ nonlocal_pool=cfg.NONLOCAL.POOL[2],
+ instantiation=cfg.NONLOCAL.INSTANTIATION,
+ trans_func_name=cfg.RESNET.TRANS_FUNC,
+ stride_1x1=cfg.RESNET.STRIDE_1X1,
+ inplace_relu=cfg.RESNET.INPLACE_RELU,
+ dilation=cfg.RESNET.SPATIAL_DILATIONS[2],
+ norm_module=self.norm_module,
+ )
+
+ # self.s5 = resnet_helper.ResStage(
+ # dim_in=[width_per_group * 16],
+ # dim_out=[width_per_group * 32],
+ # dim_inner=[dim_inner * 8],
+ # temp_kernel_sizes=temp_kernel[4],
+ # stride=cfg.RESNET.SPATIAL_STRIDES[3],
+ # num_blocks=[d5],
+ # num_groups=[num_groups],
+ # num_block_temp_kernel=cfg.RESNET.NUM_BLOCK_TEMP_KERNEL[3],
+ # nonlocal_inds=cfg.NONLOCAL.LOCATION[3],
+ # nonlocal_group=cfg.NONLOCAL.GROUP[3],
+ # nonlocal_pool=cfg.NONLOCAL.POOL[3],
+ # instantiation=cfg.NONLOCAL.INSTANTIATION,
+ # trans_func_name=cfg.RESNET.TRANS_FUNC,
+ # stride_1x1=cfg.RESNET.STRIDE_1X1,
+ # inplace_relu=cfg.RESNET.INPLACE_RELU,
+ # dilation=cfg.RESNET.SPATIAL_DILATIONS[3],
+ # norm_module=self.norm_module,
+ # )
+
+ self.labels=cfg.RESNET.LABELS
+ self.dual_define("t4",self.labels,ResDecoderBlock(width_per_group * 16,width_per_group * 8,width_per_group * 8))
+ self.dual_define("t3",self.labels,ResDecoderBlock(width_per_group * 8,width_per_group * 4, 256))
+ self.dual_define("conv1x1",self.labels,nn.Sequential(
+ nn.Conv3d(width_per_group*4+width_per_group, 128, kernel_size=(1, 1, 1), stride=1, padding=0),
+ nn.BatchNorm3d(128),
+ nn.ReLU(),
+ nn.Conv3d(128, 1, kernel_size=(1, 1, 1), stride=1, padding=0),
+ ))
+
+ self.linear = nn.Linear(1, 1)
+
+ def forward_plus(self, x, y, net):
+ return [net(x)[0] + y[0]]
+
+
+ def dual_define(self,name,labels,net):
+ for label in labels:
+ self.add_module(f"{name}_{label}",copy.deepcopy(net))
+
+
+ def upsample(self, x, dims=["space"]):
+ ori_size = x[0].shape[2:5]
+ t, h, w = ori_size
+ if "space" in dims:
+ h = 2 * h
+ w = 2 * w
+ if "time" in dims:
+ t = 2 * t
+ size = (t, h, w)
+ return [F.interpolate(x[0], size)]
+
+ def concat(self,x,y):
+ return [torch.cat([x[0],y[0]],1)]
+
+ def get_detach_var(self,x):
+ return [t.detach() for t in x]
+
+ # @torchsnooper.snoop()
+ def forward(self, x, freeze_backbone=False):
+ x = self.get_detach_var(x)
+ x1 = self.s1(x) # 1,64,8,56,56
+ x2 = self.s2(x1) # 1,256,8,56,56
+ x3 = self.s3(x2) # 1,512,8,28, 28
+ feat= self.s4(x3) # 1,1024,8,14,14
+ assert isinstance(freeze_backbone,bool)
+ if freeze_backbone:
+ feat=self.get_detach_var(feat)
+ x1=self.get_detach_var(x1)
+ x2=self.get_detach_var(x2)
+ x3=self.get_detach_var(x3)
+
+ feat = self.upsample(feat) # 1,1024, 8, 28, 28
+ feat = self.concat(x3,feat)# 1,1024+512, 8, 28, 28
+ reg_out=[self.forward_branch(feat,x1,x2,label) for label in self.labels]
+ reg_out=torch.cat(reg_out,1)
+ reg_out=torch.sigmoid(reg_out)
+ class_out = reg_out.mean([3, 4]).view(-1, 1)*100 # 1,2,8,56,56
+ class_out = self.linear(class_out)
+ class_out = class_out.view(reg_out.size(0),len(self.labels),-1)
+ class_out = torch.sigmoid(class_out)
+ return reg_out,class_out
+
+
+ def forward_branch(self,feat,x1,x2,label):
+ t4= getattr(self,f"t4_{label}")
+ feat = t4(feat[0])# 1,256, 8, 28, 28
+ feat = self.upsample([feat]) # 1,256, 8, 56, 56
+ feat = self.concat(x2,feat)# 1,256+256, 8, 56, 56
+ t3= getattr(self,f"t3_{label}")
+ feat = t3(feat[0]) # 1,128, 8, 56, 56
+ feat = self.concat(x1,[feat]) # 1,192, 8, 56, 56
+ conv1x1=getattr(self,f"conv1x1_{label}")
+ feat = conv1x1(feat[0]) # 1,2,8,56,56
+ return feat
diff --git a/clean/video/pwtf_dvd/inference/slowfast/utils/__init__.py b/clean/video/pwtf_dvd/inference/slowfast/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..8dbe96a785072a24a9bcc4841a1934024f2b06a1
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/utils/__init__.py
@@ -0,0 +1,2 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
diff --git a/clean/video/pwtf_dvd/inference/slowfast/utils/ava_eval_helper.py b/clean/video/pwtf_dvd/inference/slowfast/utils/ava_eval_helper.py
new file mode 100644
index 0000000000000000000000000000000000000000..9e8ba5468077053a4dcf1920f25256a1a241f0b9
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/utils/ava_eval_helper.py
@@ -0,0 +1,302 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+##############################################################################
+#
+# Based on:
+# --------------------------------------------------------
+# ActivityNet
+# Copyright (c) 2015 ActivityNet
+# Licensed under The MIT License
+# [see https://github.com/activitynet/ActivityNet/blob/master/LICENSE for details]
+# --------------------------------------------------------
+
+"""Helper functions for AVA evaluation."""
+
+from __future__ import (
+ absolute_import,
+ division,
+ print_function,
+ unicode_literals,
+)
+import csv
+import logging
+import numpy as np
+import pprint
+import time
+from collections import defaultdict
+from fvcore.common.file_io import PathManager
+
+from slowfast.utils.ava_evaluation import (
+ object_detection_evaluation,
+ standard_fields,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def make_image_key(video_id, timestamp):
+ """Returns a unique identifier for a video id & timestamp."""
+ return "%s,%04d" % (video_id, int(timestamp))
+
+
+def read_csv(csv_file, class_whitelist=None, load_score=False):
+ """Loads boxes and class labels from a CSV file in the AVA format.
+ CSV file format described at https://research.google.com/ava/download.html.
+ Args:
+ csv_file: A file object.
+ class_whitelist: If provided, boxes corresponding to (integer) class labels
+ not in this set are skipped.
+ Returns:
+ boxes: A dictionary mapping each unique image key (string) to a list of
+ boxes, given as coordinates [y1, x1, y2, x2].
+ labels: A dictionary mapping each unique image key (string) to a list of
+ integer class lables, matching the corresponding box in `boxes`.
+ scores: A dictionary mapping each unique image key (string) to a list of
+ score values lables, matching the corresponding label in `labels`. If
+ scores are not provided in the csv, then they will default to 1.0.
+ """
+ boxes = defaultdict(list)
+ labels = defaultdict(list)
+ scores = defaultdict(list)
+ with PathManager.open(csv_file, "r") as f:
+ reader = csv.reader(f)
+ for row in reader:
+ assert len(row) in [7, 8], "Wrong number of columns: " + row
+ image_key = make_image_key(row[0], row[1])
+ x1, y1, x2, y2 = [float(n) for n in row[2:6]]
+ action_id = int(row[6])
+ if class_whitelist and action_id not in class_whitelist:
+ continue
+ score = 1.0
+ if load_score:
+ score = float(row[7])
+ boxes[image_key].append([y1, x1, y2, x2])
+ labels[image_key].append(action_id)
+ scores[image_key].append(score)
+ return boxes, labels, scores
+
+
+def read_exclusions(exclusions_file):
+ """Reads a CSV file of excluded timestamps.
+ Args:
+ exclusions_file: A file object containing a csv of video-id,timestamp.
+ Returns:
+ A set of strings containing excluded image keys, e.g. "aaaaaaaaaaa,0904",
+ or an empty set if exclusions file is None.
+ """
+ excluded = set()
+ if exclusions_file:
+ with PathManager.open(exclusions_file, "r") as f:
+ reader = csv.reader(f)
+ for row in reader:
+ assert len(row) == 2, "Expected only 2 columns, got: " + row
+ excluded.add(make_image_key(row[0], row[1]))
+ return excluded
+
+
+def read_labelmap(labelmap_file):
+ """Read label map and class ids."""
+
+ labelmap = []
+ class_ids = set()
+ name = ""
+ class_id = ""
+ with PathManager.open(labelmap_file, "r") as f:
+ for line in f:
+ if line.startswith(" name:"):
+ name = line.split('"')[1]
+ elif line.startswith(" id:") or line.startswith(" label_id:"):
+ class_id = int(line.strip().split(" ")[-1])
+ labelmap.append({"id": class_id, "name": name})
+ class_ids.add(class_id)
+ return labelmap, class_ids
+
+
+def evaluate_ava_from_files(labelmap, groundtruth, detections, exclusions):
+ """Run AVA evaluation given annotation/prediction files."""
+
+ categories, class_whitelist = read_labelmap(labelmap)
+ excluded_keys = read_exclusions(exclusions)
+ groundtruth = read_csv(groundtruth, class_whitelist, load_score=False)
+ detections = read_csv(detections, class_whitelist, load_score=True)
+ run_evaluation(categories, groundtruth, detections, excluded_keys)
+
+
+def evaluate_ava(
+ preds,
+ original_boxes,
+ metadata,
+ excluded_keys,
+ class_whitelist,
+ categories,
+ groundtruth=None,
+ video_idx_to_name=None,
+ name="latest",
+):
+ """Run AVA evaluation given numpy arrays."""
+
+ eval_start = time.time()
+
+ detections = get_ava_eval_data(
+ preds,
+ original_boxes,
+ metadata,
+ class_whitelist,
+ video_idx_to_name=video_idx_to_name,
+ )
+
+ logger.info("Evaluating with %d unique GT frames." % len(groundtruth[0]))
+ logger.info(
+ "Evaluating with %d unique detection frames" % len(detections[0])
+ )
+
+ write_results(detections, "detections_%s.csv" % name)
+ write_results(groundtruth, "groundtruth_%s.csv" % name)
+
+ results = run_evaluation(categories, groundtruth, detections, excluded_keys)
+
+ logger.info("AVA eval done in %f seconds." % (time.time() - eval_start))
+ return results["PascalBoxes_Precision/mAP@0.5IOU"]
+
+
+def run_evaluation(
+ categories, groundtruth, detections, excluded_keys, verbose=True
+):
+ """AVA evaluation main logic."""
+
+ pascal_evaluator = object_detection_evaluation.PascalDetectionEvaluator(
+ categories
+ )
+
+ boxes, labels, _ = groundtruth
+
+ gt_keys = []
+ pred_keys = []
+
+ for image_key in boxes:
+ if image_key in excluded_keys:
+ logging.info(
+ (
+ "Found excluded timestamp in ground truth: %s. "
+ "It will be ignored."
+ ),
+ image_key,
+ )
+ continue
+ pascal_evaluator.add_single_ground_truth_image_info(
+ image_key,
+ {
+ standard_fields.InputDataFields.groundtruth_boxes: np.array(
+ boxes[image_key], dtype=float
+ ),
+ standard_fields.InputDataFields.groundtruth_classes: np.array(
+ labels[image_key], dtype=int
+ ),
+ standard_fields.InputDataFields.groundtruth_difficult: np.zeros(
+ len(boxes[image_key]), dtype=bool
+ ),
+ },
+ )
+
+ gt_keys.append(image_key)
+
+ boxes, labels, scores = detections
+
+ for image_key in boxes:
+ if image_key in excluded_keys:
+ logging.info(
+ (
+ "Found excluded timestamp in detections: %s. "
+ "It will be ignored."
+ ),
+ image_key,
+ )
+ continue
+ pascal_evaluator.add_single_detected_image_info(
+ image_key,
+ {
+ standard_fields.DetectionResultFields.detection_boxes: np.array(
+ boxes[image_key], dtype=float
+ ),
+ standard_fields.DetectionResultFields.detection_classes: np.array(
+ labels[image_key], dtype=int
+ ),
+ standard_fields.DetectionResultFields.detection_scores: np.array(
+ scores[image_key], dtype=float
+ ),
+ },
+ )
+
+ pred_keys.append(image_key)
+
+ metrics = pascal_evaluator.evaluate()
+
+ pprint.pprint(metrics, indent=2)
+ return metrics
+
+
+def get_ava_eval_data(
+ scores,
+ boxes,
+ metadata,
+ class_whitelist,
+ verbose=False,
+ video_idx_to_name=None,
+):
+ """
+ Convert our data format into the data format used in official AVA
+ evaluation.
+ """
+
+ out_scores = defaultdict(list)
+ out_labels = defaultdict(list)
+ out_boxes = defaultdict(list)
+ count = 0
+ for i in range(scores.shape[0]):
+ video_idx = int(np.round(metadata[i][0]))
+ sec = int(np.round(metadata[i][1]))
+
+ video = video_idx_to_name[video_idx]
+
+ key = video + "," + "%04d" % (sec)
+ batch_box = boxes[i].tolist()
+ # The first is batch idx.
+ batch_box = [batch_box[j] for j in [0, 2, 1, 4, 3]]
+
+ one_scores = scores[i].tolist()
+ for cls_idx, score in enumerate(one_scores):
+ if cls_idx + 1 in class_whitelist:
+ out_scores[key].append(score)
+ out_labels[key].append(cls_idx + 1)
+ out_boxes[key].append(batch_box[1:])
+ count += 1
+
+ return out_boxes, out_labels, out_scores
+
+
+def write_results(detections, filename):
+ """Write prediction results into official formats."""
+ start = time.time()
+
+ boxes, labels, scores = detections
+ with PathManager.open(filename, "w") as f:
+ for key in boxes.keys():
+ for box, label, score in zip(boxes[key], labels[key], scores[key]):
+ f.write(
+ "%s,%.03f,%.03f,%.03f,%.03f,%d,%.04f\n"
+ % (key, box[1], box[0], box[3], box[2], label, score)
+ )
+
+ logger.info("AVA results wrote to %s" % filename)
+ logger.info("\ttook %d seconds." % (time.time() - start))
diff --git a/clean/video/pwtf_dvd/inference/slowfast/utils/benchmark.py b/clean/video/pwtf_dvd/inference/slowfast/utils/benchmark.py
new file mode 100644
index 0000000000000000000000000000000000000000..33e5fe9073ad61ecec737d6b4a6a2880eec15cb9
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/utils/benchmark.py
@@ -0,0 +1,103 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+Functions for benchmarks.
+"""
+
+import numpy as np
+import pprint
+import torch
+import tqdm
+from fvcore.common.timer import Timer
+
+import slowfast.utils.logging as logging
+import slowfast.utils.misc as misc
+from slowfast.datasets import loader
+from slowfast.utils.env import setup_environment
+
+logger = logging.get_logger(__name__)
+
+
+def benchmark_data_loading(cfg):
+ """
+ Benchmark the speed of data loading in PySlowFast.
+ Args:
+
+ cfg (CfgNode): configs. Details can be found in
+ slowfast/config/defaults.py
+ """
+ # Set up environment.
+ setup_environment()
+ # Set random seed from configs.
+ np.random.seed(cfg.RNG_SEED)
+ torch.manual_seed(cfg.RNG_SEED)
+
+ # Setup logging format.
+ logging.setup_logging(cfg.OUTPUT_DIR)
+
+ # Print config.
+ logger.info("Benchmark data loading with config:")
+ logger.info(pprint.pformat(cfg))
+
+ timer = Timer()
+ dataloader = loader.construct_loader(cfg, "train")
+ logger.info(
+ "Initialize loader using {:.2f} seconds.".format(timer.seconds())
+ )
+ # Total batch size across different machines.
+ batch_size = cfg.TRAIN.BATCH_SIZE * cfg.NUM_SHARDS
+ log_period = cfg.BENCHMARK.LOG_PERIOD
+ epoch_times = []
+ # Test for a few epochs.
+ for cur_epoch in range(cfg.BENCHMARK.NUM_EPOCHS):
+ timer = Timer()
+ timer_epoch = Timer()
+ iter_times = []
+ if cfg.BENCHMARK.SHUFFLE:
+ loader.shuffle_dataset(dataloader, cur_epoch)
+ for cur_iter, _ in enumerate(tqdm.tqdm(dataloader)):
+ if cur_iter > 0 and cur_iter % log_period == 0:
+ iter_times.append(timer.seconds())
+ ram_usage, ram_total = misc.cpu_mem_usage()
+ logger.info(
+ "Epoch {}: {} iters ({} videos) in {:.2f} seconds. "
+ "RAM Usage: {:.2f}/{:.2f} GB.".format(
+ cur_epoch,
+ log_period,
+ log_period * batch_size,
+ iter_times[-1],
+ ram_usage,
+ ram_total,
+ )
+ )
+ timer.reset()
+ epoch_times.append(timer_epoch.seconds())
+ ram_usage, ram_total = misc.cpu_mem_usage()
+ logger.info(
+ "Epoch {}: in total {} iters ({} videos) in {:.2f} seconds. "
+ "RAM Usage: {:.2f}/{:.2f} GB.".format(
+ cur_epoch,
+ len(dataloader),
+ len(dataloader) * batch_size,
+ epoch_times[-1],
+ ram_usage,
+ ram_total,
+ )
+ )
+ logger.info(
+ "Epoch {}: on average every {} iters ({} videos) take {:.2f}/{:.2f} "
+ "(avg/std) seconds.".format(
+ cur_epoch,
+ log_period,
+ log_period * batch_size,
+ np.mean(iter_times),
+ np.std(iter_times),
+ )
+ )
+ logger.info(
+ "On average every epoch ({} videos) takes {:.2f}/{:.2f} "
+ "(avg/std) seconds.".format(
+ len(dataloader) * batch_size,
+ np.mean(epoch_times),
+ np.std(epoch_times),
+ )
+ )
diff --git a/clean/video/pwtf_dvd/inference/slowfast/utils/bn_helper.py b/clean/video/pwtf_dvd/inference/slowfast/utils/bn_helper.py
new file mode 100644
index 0000000000000000000000000000000000000000..b18d8c76c10d7598db61ba8ca192314140f9ba79
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/utils/bn_helper.py
@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""bn helper."""
+
+import itertools
+import torch
+
+
+@torch.no_grad()
+def compute_and_update_bn_stats(model, data_loader, num_batches=200):
+ """
+ Compute and update the batch norm stats to make it more precise. During
+ training both bn stats and the weight are changing after every iteration,
+ so the bn can not precisely reflect the latest stats of the current model.
+ Here the bn stats is recomputed without change of weights, to make the
+ running mean and running var more precise.
+ Args:
+ model (model): the model using to compute and update the bn stats.
+ data_loader (dataloader): dataloader using to provide inputs.
+ num_batches (int): running iterations using to compute the stats.
+ """
+
+ # Prepares all the bn layers.
+ bn_layers = [
+ m
+ for m in model.modules()
+ if any(
+ (
+ isinstance(m, bn_type)
+ for bn_type in (
+ torch.nn.BatchNorm1d,
+ torch.nn.BatchNorm2d,
+ torch.nn.BatchNorm3d,
+ )
+ )
+ )
+ ]
+
+ # In order to make the running stats only reflect the current batch, the
+ # momentum is disabled.
+ # bn.running_mean = (1 - momentum) * bn.running_mean + momentum * batch_mean
+ # Setting the momentum to 1.0 to compute the stats without momentum.
+ momentum_actual = [bn.momentum for bn in bn_layers]
+ for bn in bn_layers:
+ bn.momentum = 1.0
+
+ # Calculates the running iterations for precise stats computation.
+ running_mean = [torch.zeros_like(bn.running_mean) for bn in bn_layers]
+ running_square_mean = [torch.zeros_like(bn.running_var) for bn in bn_layers]
+
+ for ind, (inputs, _, _) in enumerate(
+ itertools.islice(data_loader, num_batches)
+ ):
+ # Forwards the model to update the bn stats.
+ if isinstance(inputs, (list,)):
+ for i in range(len(inputs)):
+ inputs[i] = inputs[i].float().cuda(non_blocking=True)
+ else:
+ inputs = inputs.cuda(non_blocking=True)
+ model(inputs)
+
+ for i, bn in enumerate(bn_layers):
+ # Accumulates the bn stats.
+ running_mean[i] += (bn.running_mean - running_mean[i]) / (ind + 1)
+ # $E(x^2) = Var(x) + E(x)^2$.
+ cur_square_mean = bn.running_var + bn.running_mean ** 2
+ running_square_mean[i] += (
+ cur_square_mean - running_square_mean[i]
+ ) / (ind + 1)
+
+ for i, bn in enumerate(bn_layers):
+ bn.running_mean = running_mean[i]
+ # Var(x) = $E(x^2) - E(x)^2$.
+ bn.running_var = running_square_mean[i] - bn.running_mean ** 2
+ # Sets the precise bn stats.
+ bn.momentum = momentum_actual[i]
diff --git a/clean/video/pwtf_dvd/inference/slowfast/utils/c2_model_loading.py b/clean/video/pwtf_dvd/inference/slowfast/utils/c2_model_loading.py
new file mode 100644
index 0000000000000000000000000000000000000000..4bcc0759c484fd321917c55e9967835632c2ac54
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/utils/c2_model_loading.py
@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Caffe2 to PyTorch checkpoint name converting utility."""
+
+import re
+
+
+def get_name_convert_func():
+ """
+ Get the function to convert Caffe2 layer names to PyTorch layer names.
+ Returns:
+ (func): function to convert parameter name from Caffe2 format to PyTorch
+ format.
+ """
+ pairs = [
+ # ------------------------------------------------------------
+ # 'nonlocal_conv3_1_theta_w' -> 's3.pathway0_nonlocal3.conv_g.weight'
+ [
+ r"^nonlocal_conv([0-9]+)_([0-9]+)_(.*)",
+ r"s\1.pathway0_nonlocal\2_\3",
+ ],
+ # 'theta' -> 'conv_theta'
+ [r"^(.*)_nonlocal([0-9]+)_(theta)(.*)", r"\1_nonlocal\2.conv_\3\4"],
+ # 'g' -> 'conv_g'
+ [r"^(.*)_nonlocal([0-9]+)_(g)(.*)", r"\1_nonlocal\2.conv_\3\4"],
+ # 'phi' -> 'conv_phi'
+ [r"^(.*)_nonlocal([0-9]+)_(phi)(.*)", r"\1_nonlocal\2.conv_\3\4"],
+ # 'out' -> 'conv_out'
+ [r"^(.*)_nonlocal([0-9]+)_(out)(.*)", r"\1_nonlocal\2.conv_\3\4"],
+ # 'nonlocal_conv4_5_bn_s' -> 's4.pathway0_nonlocal3.bn.weight'
+ [r"^(.*)_nonlocal([0-9]+)_(bn)_(.*)", r"\1_nonlocal\2.\3.\4"],
+ # ------------------------------------------------------------
+ # 't_pool1_subsample_bn' -> 's1_fuse.conv_f2s.bn.running_mean'
+ [r"^t_pool1_subsample_bn_(.*)", r"s1_fuse.bn.\1"],
+ # 't_pool1_subsample' -> 's1_fuse.conv_f2s'
+ [r"^t_pool1_subsample_(.*)", r"s1_fuse.conv_f2s.\1"],
+ # 't_res4_5_branch2c_bn_subsample_bn_rm' -> 's4_fuse.conv_f2s.bias'
+ [
+ r"^t_res([0-9]+)_([0-9]+)_branch2c_bn_subsample_bn_(.*)",
+ r"s\1_fuse.bn.\3",
+ ],
+ # 't_pool1_subsample' -> 's1_fuse.conv_f2s'
+ [
+ r"^t_res([0-9]+)_([0-9]+)_branch2c_bn_subsample_(.*)",
+ r"s\1_fuse.conv_f2s.\3",
+ ],
+ # ------------------------------------------------------------
+ # 'res4_4_branch_2c_bn_b' -> 's4.pathway0_res4.branch2.c_bn_b'
+ [
+ r"^res([0-9]+)_([0-9]+)_branch([0-9]+)([a-z])_(.*)",
+ r"s\1.pathway0_res\2.branch\3.\4_\5",
+ ],
+ # 'res_conv1_bn_' -> 's1.pathway0_stem.bn.'
+ [r"^res_conv1_bn_(.*)", r"s1.pathway0_stem.bn.\1"],
+ # 'conv1_w_momentum' -> 's1.pathway0_stem.conv.'
+ [r"^conv1_(.*)", r"s1.pathway0_stem.conv.\1"],
+ # 'res4_0_branch1_w' -> 'S4.pathway0_res0.branch1.weight'
+ [
+ r"^res([0-9]+)_([0-9]+)_branch([0-9]+)_(.*)",
+ r"s\1.pathway0_res\2.branch\3_\4",
+ ],
+ # 'res_conv1_' -> 's1.pathway0_stem.conv.'
+ [r"^res_conv1_(.*)", r"s1.pathway0_stem.conv.\1"],
+ # ------------------------------------------------------------
+ # 'res4_4_branch_2c_bn_b' -> 's4.pathway0_res4.branch2.c_bn_b'
+ [
+ r"^t_res([0-9]+)_([0-9]+)_branch([0-9]+)([a-z])_(.*)",
+ r"s\1.pathway1_res\2.branch\3.\4_\5",
+ ],
+ # 'res_conv1_bn_' -> 's1.pathway0_stem.bn.'
+ [r"^t_res_conv1_bn_(.*)", r"s1.pathway1_stem.bn.\1"],
+ # 'conv1_w_momentum' -> 's1.pathway0_stem.conv.'
+ [r"^t_conv1_(.*)", r"s1.pathway1_stem.conv.\1"],
+ # 'res4_0_branch1_w' -> 'S4.pathway0_res0.branch1.weight'
+ [
+ r"^t_res([0-9]+)_([0-9]+)_branch([0-9]+)_(.*)",
+ r"s\1.pathway1_res\2.branch\3_\4",
+ ],
+ # 'res_conv1_' -> 's1.pathway0_stem.conv.'
+ [r"^t_res_conv1_(.*)", r"s1.pathway1_stem.conv.\1"],
+ # ------------------------------------------------------------
+ # pred_ -> head.projection.
+ [r"pred_(.*)", r"head.projection.\1"],
+ # '.bn_b' -> '.weight'
+ [r"(.*)bn.b\Z", r"\1bn.bias"],
+ # '.bn_s' -> '.weight'
+ [r"(.*)bn.s\Z", r"\1bn.weight"],
+ # '_bn_rm' -> '.running_mean'
+ [r"(.*)bn.rm\Z", r"\1bn.running_mean"],
+ # '_bn_riv' -> '.running_var'
+ [r"(.*)bn.riv\Z", r"\1bn.running_var"],
+ # '_b' -> '.bias'
+ [r"(.*)[\._]b\Z", r"\1.bias"],
+ # '_w' -> '.weight'
+ [r"(.*)[\._]w\Z", r"\1.weight"],
+ ]
+
+ def convert_caffe2_name_to_pytorch(caffe2_layer_name):
+ """
+ Convert the caffe2_layer_name to pytorch format by apply the list of
+ regular expressions.
+ Args:
+ caffe2_layer_name (str): caffe2 layer name.
+ Returns:
+ (str): pytorch layer name.
+ """
+ for source, dest in pairs:
+ caffe2_layer_name = re.sub(source, dest, caffe2_layer_name)
+ return caffe2_layer_name
+
+ return convert_caffe2_name_to_pytorch
diff --git a/clean/video/pwtf_dvd/inference/slowfast/utils/checkpoint.py b/clean/video/pwtf_dvd/inference/slowfast/utils/checkpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..05d5ac4624feecbbc1f868682f8ed921ac5fef2c
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/utils/checkpoint.py
@@ -0,0 +1,530 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Functions that handle saving and loading of checkpoints."""
+
+import copy
+import numpy as np
+import os
+import pickle
+from collections import OrderedDict
+import torch
+from fvcore.common.file_io import PathManager
+
+import slowfast.utils.distributed as du
+import slowfast.utils.logging as logging
+from slowfast.utils.c2_model_loading import get_name_convert_func
+
+logger = logging.get_logger(__name__)
+
+
+def make_checkpoint_dir(path_to_job):
+ """
+ Creates the checkpoint directory (if not present already).
+ Args:
+ path_to_job (string): the path to the folder of the current job.
+ """
+ checkpoint_dir = os.path.join(path_to_job, "checkpoints")
+ # Create the checkpoint dir from the master process
+ if du.is_master_proc() and not PathManager.exists(checkpoint_dir):
+ try:
+ PathManager.mkdirs(checkpoint_dir)
+ except Exception:
+ pass
+ return checkpoint_dir
+
+
+def get_checkpoint_dir(path_to_job):
+ """
+ Get path for storing checkpoints.
+ Args:
+ path_to_job (string): the path to the folder of the current job.
+ """
+ return os.path.join(path_to_job, "checkpoints")
+
+
+def get_path_to_checkpoint(path_to_job, epoch):
+ """
+ Get the full path to a checkpoint file.
+ Args:
+ path_to_job (string): the path to the folder of the current job.
+ epoch (int): the number of epoch for the checkpoint.
+ """
+ name = "checkpoint_epoch_{:07d}.pyth".format(epoch)
+ return os.path.join(get_checkpoint_dir(path_to_job), name)
+
+
+def get_last_checkpoint(path_to_job):
+ """
+ Get the last checkpoint from the checkpointing folder.
+ Args:
+ path_to_job (string): the path to the folder of the current job.
+ """
+
+ d = get_checkpoint_dir(path_to_job)
+ names = PathManager.ls(d) if PathManager.exists(d) else []
+ names = [f for f in names if "checkpoint" in f]
+ assert len(names), "No checkpoints found in '{}'.".format(d)
+ # Sort the checkpoints by epoch.
+ name = sorted(names)[-1]
+ return os.path.join(d, name)
+
+
+def has_checkpoint(path_to_job):
+ """
+ Determines if the given directory contains a checkpoint.
+ Args:
+ path_to_job (string): the path to the folder of the current job.
+ """
+ d = get_checkpoint_dir(path_to_job)
+ files = PathManager.ls(d) if PathManager.exists(d) else []
+ return any("checkpoint" in f for f in files)
+
+
+def is_checkpoint_epoch(cfg, cur_epoch, multigrid_schedule=None):
+ """
+ Determine if a checkpoint should be saved on current epoch.
+ Args:
+ cfg (CfgNode): configs to save.
+ cur_epoch (int): current number of epoch of the model.
+ multigrid_schedule (List): schedule for multigrid training.
+ """
+ if cur_epoch + 1 == cfg.SOLVER.MAX_EPOCH:
+ return True
+ if multigrid_schedule is not None:
+ prev_epoch = 0
+ for s in multigrid_schedule:
+ if cur_epoch < s[-1]:
+ period = max(
+ (s[-1] - prev_epoch) // cfg.MULTIGRID.EVAL_FREQ + 1, 1
+ )
+ return (s[-1] - 1 - cur_epoch) % period == 0
+ prev_epoch = s[-1]
+
+ return (cur_epoch + 1) % cfg.TRAIN.CHECKPOINT_PERIOD == 0
+
+
+def is_checkpoint_iter(cfg, cur_iter):
+ """
+ Determine if a checkpoint should be saved on current iter.
+ Args:
+ cfg (CfgNode): configs to save.
+ cur_epoch (int): current number of epoch of the model.
+ multigrid_schedule (List): schedule for multigrid training.
+ """
+
+ return (cur_iter+1) % cfg.TRAIN.CHECKPOINT_PERIOD_BY_ITER == 0
+
+
+
+def save_checkpoint_by_iter(path_to_job, model, optimizer, epoch,global_step,cfg):
+ """
+ Save a checkpoint.
+ Args:
+ model (model): model to save the weight to the checkpoint.
+ optimizer (optim): optimizer to save the historical state.
+ epoch (int): current number of epoch of the model.
+ cfg (CfgNode): configs to save.
+ """
+ # Save checkpoints only from the master process.
+ if not du.is_master_proc(cfg.NUM_GPUS * cfg.NUM_SHARDS):
+ return
+ # Ensure that the checkpoint dir exists.
+ PathManager.mkdirs(get_checkpoint_dir(path_to_job))
+ # Omit the DDP wrapper in the multi-gpu setting.
+ sd = model.module.state_dict() if cfg.NUM_GPUS > 1 else model.state_dict()
+ normalized_sd = sub_to_normal_bn(sd)
+
+ # Record the state.
+ checkpoint = {
+ "epoch": epoch,
+ "model_state": normalized_sd,
+ "optimizer_state": optimizer.state_dict(),
+ "global_step": global_step,
+ "cfg": cfg.dump(),
+ }
+ # Write the checkpoint.
+ path_to_checkpoint = get_path_to_checkpoint(path_to_job,global_step+1)
+ with PathManager.open(path_to_checkpoint, "wb") as f:
+ torch.save(checkpoint, f)
+ return path_to_checkpoint
+
+def save_checkpoint(path_to_job, model, optimizer, epoch, cfg):
+ """
+ Save a checkpoint.
+ Args:
+ model (model): model to save the weight to the checkpoint.
+ optimizer (optim): optimizer to save the historical state.
+ epoch (int): current number of epoch of the model.
+ cfg (CfgNode): configs to save.
+ """
+ # Save checkpoints only from the master process.
+ if not du.is_master_proc(cfg.NUM_GPUS * cfg.NUM_SHARDS):
+ return
+ # Ensure that the checkpoint dir exists.
+ PathManager.mkdirs(get_checkpoint_dir(path_to_job))
+ # Omit the DDP wrapper in the multi-gpu setting.
+ sd = model.module.state_dict() if cfg.NUM_GPUS > 1 else model.state_dict()
+ normalized_sd = sub_to_normal_bn(sd)
+
+ # Record the state.
+ checkpoint = {
+ "epoch": epoch,
+ "model_state": normalized_sd,
+ "optimizer_state": optimizer.state_dict(),
+ "cfg": cfg.dump(),
+ }
+ # Write the checkpoint.
+ path_to_checkpoint = get_path_to_checkpoint(path_to_job, epoch + 1)
+ with PathManager.open(path_to_checkpoint, "wb") as f:
+ torch.save(checkpoint, f)
+ return path_to_checkpoint
+
+
+def inflate_weight(state_dict_2d, state_dict_3d):
+ """
+ Inflate 2D model weights in state_dict_2d to the 3D model weights in
+ state_dict_3d. The details can be found in:
+ Joao Carreira, and Andrew Zisserman.
+ "Quo vadis, action recognition? a new model and the kinetics dataset."
+ Args:
+ state_dict_2d (OrderedDict): a dict of parameters from a 2D model.
+ state_dict_3d (OrderedDict): a dict of parameters from a 3D model.
+ Returns:
+ state_dict_inflated (OrderedDict): a dict of inflated parameters.
+ """
+ state_dict_inflated = OrderedDict()
+ for k, v2d in state_dict_2d.items():
+ assert k in state_dict_3d.keys()
+ v3d = state_dict_3d[k]
+ # Inflate the weight of 2D conv to 3D conv.
+ if len(v2d.shape) == 4 and len(v3d.shape) == 5:
+ logger.info(
+ "Inflate {}: {} -> {}: {}".format(k, v2d.shape, k, v3d.shape)
+ )
+ # Dimension need to be match.
+ assert v2d.shape[-2:] == v3d.shape[-2:]
+ assert v2d.shape[:2] == v3d.shape[:2]
+ v3d = (
+ v2d.unsqueeze(2).repeat(1, 1, v3d.shape[2], 1, 1) / v3d.shape[2]
+ )
+ elif v2d.shape == v3d.shape:
+ v3d = v2d
+ else:
+ logger.info(
+ "Unexpected {}: {} -|> {}: {}".format(
+ k, v2d.shape, k, v3d.shape
+ )
+ )
+ state_dict_inflated[k] = v3d.clone()
+ return state_dict_inflated
+
+
+def load_checkpoint(
+ path_to_checkpoint,
+ model,
+ data_parallel=True,
+ optimizer=None,
+ inflation=False,
+ convert_from_caffe2=False,
+):
+ """
+ Load the checkpoint from the given file. If inflation is True, inflate the
+ 2D Conv weights from the checkpoint to 3D Conv.
+ Args:
+ path_to_checkpoint (string): path to the checkpoint to load.
+ model (model): model to load the weights from the checkpoint.
+ data_parallel (bool): if true, model is wrapped by
+ torch.nn.parallel.DistributedDataParallel.
+ optimizer (optim): optimizer to load the historical state.
+ inflation (bool): if True, inflate the weights from the checkpoint.
+ convert_from_caffe2 (bool): if True, load the model from caffe2 and
+ convert it to pytorch.
+ Returns:
+ (int): the number of training epoch of the checkpoint.
+ """
+ assert PathManager.exists(
+ path_to_checkpoint
+ ), "Checkpoint '{}' not found".format(path_to_checkpoint)
+ # Account for the DDP wrapper in the multi-gpu setting.
+ ms = model.module if data_parallel else model
+ if convert_from_caffe2:
+ with PathManager.open(path_to_checkpoint, "rb") as f:
+ caffe2_checkpoint = pickle.load(f, encoding="latin1")
+ state_dict = OrderedDict()
+ name_convert_func = get_name_convert_func()
+ for key in caffe2_checkpoint["blobs"].keys():
+ converted_key = name_convert_func(key)
+ converted_key = c2_normal_to_sub_bn(converted_key, ms.state_dict())
+ if converted_key in ms.state_dict():
+ c2_blob_shape = caffe2_checkpoint["blobs"][key].shape
+ model_blob_shape = ms.state_dict()[converted_key].shape
+ # Load BN stats to Sub-BN.
+ if (
+ len(model_blob_shape) == 1
+ and len(c2_blob_shape) == 1
+ and model_blob_shape[0] > c2_blob_shape[0]
+ and model_blob_shape[0] % c2_blob_shape[0] == 0
+ ):
+ caffe2_checkpoint["blobs"][key] = np.concatenate(
+ [caffe2_checkpoint["blobs"][key]]
+ * (model_blob_shape[0] // c2_blob_shape[0])
+ )
+ c2_blob_shape = caffe2_checkpoint["blobs"][key].shape
+
+ if c2_blob_shape == tuple(model_blob_shape):
+ state_dict[converted_key] = torch.tensor(
+ caffe2_checkpoint["blobs"][key]
+ ).clone()
+ logger.info(
+ "{}: {} => {}: {}".format(
+ key,
+ c2_blob_shape,
+ converted_key,
+ tuple(model_blob_shape),
+ )
+ )
+ else:
+ logger.warn(
+ "!! {}: {} does not match {}: {}".format(
+ key,
+ c2_blob_shape,
+ converted_key,
+ tuple(model_blob_shape),
+ )
+ )
+ else:
+ if not any(
+ prefix in key for prefix in ["momentum", "lr", "model_iter"]
+ ):
+ logger.warn(
+ "!! {}: can not be converted, got {}".format(
+ key, converted_key
+ )
+ )
+ ms.load_state_dict(state_dict, strict=False)
+ epoch = -1
+ global_step=-1
+ else:
+ # Load the checkpoint on CPU to avoid GPU mem spike.
+ with PathManager.open(path_to_checkpoint, "rb") as f:
+ checkpoint = torch.load(f, map_location="cpu")
+ model_state_dict_3d = (
+ model.module.state_dict() if data_parallel else model.state_dict()
+ )
+ checkpoint["model_state"] = normal_to_sub_bn(
+ checkpoint["model_state"], model_state_dict_3d
+ )
+ if inflation:
+ # Try to inflate the model.
+ inflated_model_dict = inflate_weight(
+ checkpoint["model_state"], model_state_dict_3d
+ )
+ ms.load_state_dict(inflated_model_dict, strict=False)
+ else:
+ ms.load_state_dict(checkpoint["model_state"])
+ # Load the optimizer state (commonly not done when fine-tuning)
+ if optimizer:
+ optimizer.load_state_dict(checkpoint["optimizer_state"])
+ if "epoch" in checkpoint.keys():
+ epoch = checkpoint["epoch"]
+ else:
+ epoch = -1
+ if "global_step" in checkpoint.keys():
+ global_step=checkpoint["global_step"]
+ else:
+ global_step=-1
+ return epoch,global_step
+
+
+def sub_to_normal_bn(sd):
+ """
+ Convert the Sub-BN paprameters to normal BN parameters in a state dict.
+ There are two copies of BN layers in a Sub-BN implementation: `bn.bn` and
+ `bn.split_bn`. `bn.split_bn` is used during training and
+ "compute_precise_bn". Before saving or evaluation, its stats are copied to
+ `bn.bn`. We rename `bn.bn` to `bn` and store it to be consistent with normal
+ BN layers.
+ Args:
+ sd (OrderedDict): a dict of parameters whitch might contain Sub-BN
+ parameters.
+ Returns:
+ new_sd (OrderedDict): a dict with Sub-BN parameters reshaped to
+ normal parameters.
+ """
+ new_sd = copy.deepcopy(sd)
+ modifications = [
+ ("bn.bn.running_mean", "bn.running_mean"),
+ ("bn.bn.running_var", "bn.running_var"),
+ ("bn.split_bn.num_batches_tracked", "bn.num_batches_tracked"),
+ ]
+ to_remove = ["bn.bn.", ".split_bn."]
+ for key in sd:
+ for before, after in modifications:
+ if key.endswith(before):
+ new_key = key.split(before)[0] + after
+ new_sd[new_key] = new_sd.pop(key)
+
+ for rm in to_remove:
+ if rm in key and key in new_sd:
+ del new_sd[key]
+
+ for key in new_sd:
+ if key.endswith("bn.weight") or key.endswith("bn.bias"):
+ if len(new_sd[key].size()) == 4:
+ assert all(d == 1 for d in new_sd[key].size()[1:])
+ new_sd[key] = new_sd[key][:, 0, 0, 0]
+
+ return new_sd
+
+
+def c2_normal_to_sub_bn(key, model_keys):
+ """
+ Convert BN parameters to Sub-BN parameters if model contains Sub-BNs.
+ Args:
+ key (OrderedDict): source dict of parameters.
+ mdoel_key (OrderedDict): target dict of parameters.
+ Returns:
+ new_sd (OrderedDict): converted dict of parameters.
+ """
+ if "bn.running_" in key:
+ if key in model_keys:
+ return key
+
+ new_key = key.replace("bn.running_", "bn.split_bn.running_")
+ if new_key in model_keys:
+ return new_key
+ else:
+ return key
+
+
+def normal_to_sub_bn(checkpoint_sd, model_sd):
+ """
+ Convert BN parameters to Sub-BN parameters if model contains Sub-BNs.
+ Args:
+ checkpoint_sd (OrderedDict): source dict of parameters.
+ model_sd (OrderedDict): target dict of parameters.
+ Returns:
+ new_sd (OrderedDict): converted dict of parameters.
+ """
+ for key in model_sd:
+ if key not in checkpoint_sd:
+ if "bn.split_bn." in key:
+ load_key = key.replace("bn.split_bn.", "bn.")
+ bn_key = key.replace("bn.split_bn.", "bn.bn.")
+ checkpoint_sd[key] = checkpoint_sd.pop(load_key)
+ checkpoint_sd[bn_key] = checkpoint_sd[key]
+
+ for key in model_sd:
+ if key in checkpoint_sd:
+ model_blob_shape = model_sd[key].shape
+ c2_blob_shape = checkpoint_sd[key].shape
+
+ if (
+ len(model_blob_shape) == 1
+ and len(c2_blob_shape) == 1
+ and model_blob_shape[0] > c2_blob_shape[0]
+ and model_blob_shape[0] % c2_blob_shape[0] == 0
+ ):
+ before_shape = checkpoint_sd[key].shape
+ checkpoint_sd[key] = torch.cat(
+ [checkpoint_sd[key]]
+ * (model_blob_shape[0] // c2_blob_shape[0])
+ )
+ logger.info(
+ "{} {} -> {}".format(
+ key, before_shape, checkpoint_sd[key].shape
+ )
+ )
+ return checkpoint_sd
+
+
+def load_test_checkpoint(cfg, model):
+ """
+ Loading checkpoint logic for testing.
+ """
+ # Load a checkpoint to test if applicable.
+ if cfg.TEST.CHECKPOINT_FILE_PATH != "":
+ # If no checkpoint found in MODEL_VIS.CHECKPOINT_FILE_PATH or in the current
+ # checkpoint folder, try to load checkpoint from
+ # TEST.CHECKPOINT_FILE_PATH and test it.
+ load_checkpoint(
+ cfg.TEST.CHECKPOINT_FILE_PATH,
+ model,
+ cfg.NUM_GPUS > 1,
+ None,
+ inflation=False,
+ convert_from_caffe2=cfg.TEST.CHECKPOINT_TYPE == "caffe2",
+ )
+ elif has_checkpoint(cfg.OUTPUT_DIR):
+ last_checkpoint = get_last_checkpoint(cfg.OUTPUT_DIR)
+ load_checkpoint(last_checkpoint, model, cfg.NUM_GPUS > 1)
+ elif cfg.TRAIN.CHECKPOINT_FILE_PATH != "":
+ # If no checkpoint found in TEST.CHECKPOINT_FILE_PATH or in the current
+ # checkpoint folder, try to load checkpoint from
+ # TRAIN.CHECKPOINT_FILE_PATH and test it.
+ load_checkpoint(
+ cfg.TRAIN.CHECKPOINT_FILE_PATH,
+ model,
+ cfg.NUM_GPUS > 1,
+ None,
+ inflation=False,
+ convert_from_caffe2=cfg.TRAIN.CHECKPOINT_TYPE == "caffe2",
+ )
+ else:
+ logger.info(
+ "Unknown way of loading checkpoint. Using with random initialization, only for debugging."
+ )
+
+
+def load_train_checkpoint(cfg, model, optimizer):
+ """
+ Loading checkpoint logic for training.
+ """
+ if cfg.TRAIN.AUTO_RESUME and has_checkpoint(cfg.OUTPUT_DIR):
+ last_checkpoint = get_last_checkpoint(cfg.OUTPUT_DIR)
+ logger.info("Load from last checkpoint, {}.".format(last_checkpoint))
+ checkpoint_epoch,global_step = load_checkpoint(
+ last_checkpoint, model, cfg.NUM_GPUS > 1, optimizer
+ )
+ start_epoch = checkpoint_epoch + 1
+ global_step = global_step + 1
+ elif cfg.TRAIN.CHECKPOINT_FILE_PATH != "":
+ if cfg.TRAIN.CHECKPOINT_TYPE=="backbone":
+ logger.info("Load backbone from given checkpoint file.")
+ load_backbone(model,cfg.TRAIN.CHECKPOINT_FILE_PATH)
+ start_epoch = 0
+ global_step = 0
+ else:
+ logger.info("Load from given checkpoint file.")
+ checkpoint_epoch, global_step = load_checkpoint(
+ cfg.TRAIN.CHECKPOINT_FILE_PATH,
+ model,
+ cfg.NUM_GPUS > 1,
+ optimizer,
+ inflation=cfg.TRAIN.CHECKPOINT_INFLATE,
+ convert_from_caffe2=cfg.TRAIN.CHECKPOINT_TYPE == "caffe2",
+ )
+ start_epoch = checkpoint_epoch + 1
+ global_step = global_step + 1
+ else:
+ start_epoch = 0
+ global_step = 0
+
+ return start_epoch, global_step
+
+
+
+def load_backbone(model,file):
+ current_state=model.state_dict()
+ checkpoint=torch.load(file)
+
+ for key in checkpoint:
+ if key in current_state:
+ assert current_state[key].shape==checkpoint[key].shape
+ current_state[key]=checkpoint[key]
+ model.load_state_dict(current_state)
+
+ return model
+
+
diff --git a/clean/video/pwtf_dvd/inference/slowfast/utils/distributed.py b/clean/video/pwtf_dvd/inference/slowfast/utils/distributed.py
new file mode 100644
index 0000000000000000000000000000000000000000..bfbed8e8a4af5fc4b38c1558616fd3f640b587c0
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/utils/distributed.py
@@ -0,0 +1,299 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Distributed helpers."""
+
+import functools
+import logging
+import pickle
+import torch
+import torch.distributed as dist
+
+_LOCAL_PROCESS_GROUP = None
+
+
+def all_gather(tensors):
+ """
+ All gathers the provided tensors from all processes across machines.
+ Args:
+ tensors (list): tensors to perform all gather across all processes in
+ all machines.
+ """
+
+ gather_list = []
+ output_tensor = []
+ world_size = dist.get_world_size()
+ for tensor in tensors:
+ tensor_placeholder = [
+ torch.ones_like(tensor) for _ in range(world_size)
+ ]
+ dist.all_gather(tensor_placeholder, tensor, async_op=False)
+ gather_list.append(tensor_placeholder)
+ for gathered_tensor in gather_list:
+ output_tensor.append(torch.cat(gathered_tensor, dim=0))
+ return output_tensor
+
+
+def all_reduce(tensors, average=True):
+ """
+ All reduce the provided tensors from all processes across machines.
+ Args:
+ tensors (list): tensors to perform all reduce across all processes in
+ all machines.
+ average (bool): scales the reduced tensor by the number of overall
+ processes across all machines.
+ """
+
+ for tensor in tensors:
+ dist.all_reduce(tensor, async_op=False)
+ if average:
+ world_size = dist.get_world_size()
+ for tensor in tensors:
+ tensor.mul_(1.0 / world_size)
+ return tensors
+
+
+def init_process_group(
+ local_rank,
+ local_world_size,
+ shard_id,
+ num_shards,
+ init_method,
+ dist_backend="nccl",
+):
+ """
+ Initializes the default process group.
+ Args:
+ local_rank (int): the rank on the current local machine.
+ local_world_size (int): the world size (number of processes running) on
+ the current local machine.
+ shard_id (int): the shard index (machine rank) of the current machine.
+ num_shards (int): number of shards for distributed training.
+ init_method (string): supporting three different methods for
+ initializing process groups:
+ "file": use shared file system to initialize the groups across
+ different processes.
+ "tcp": use tcp address to initialize the groups across different
+ dist_backend (string): backend to use for distributed training. Options
+ includes gloo, mpi and nccl, the details can be found here:
+ https://pytorch.org/docs/stable/distributed.html
+ """
+ # Sets the GPU to use.
+ torch.cuda.set_device(local_rank)
+ # Initialize the process group.
+ proc_rank = local_rank + shard_id * local_world_size
+ world_size = local_world_size * num_shards
+ dist.init_process_group(
+ backend=dist_backend,
+ init_method=init_method,
+ world_size=world_size,
+ rank=proc_rank,
+ )
+
+
+def is_master_proc(num_gpus=8):
+ """
+ Determines if the current process is the master process.
+ """
+ if torch.distributed.is_initialized():
+ return dist.get_rank() % num_gpus == 0
+ else:
+ return True
+
+
+def get_world_size():
+ """
+ Get the size of the world.
+ """
+ if not dist.is_available():
+ return 1
+ if not dist.is_initialized():
+ return 1
+ return dist.get_world_size()
+
+
+def get_rank():
+ """
+ Get the rank of the current process.
+ """
+ if not dist.is_available():
+ return 0
+ if not dist.is_initialized():
+ return 0
+ return dist.get_rank()
+
+
+def synchronize():
+ """
+ Helper function to synchronize (barrier) among all processes when
+ using distributed training
+ """
+ if not dist.is_available():
+ return
+ if not dist.is_initialized():
+ return
+ world_size = dist.get_world_size()
+ if world_size == 1:
+ return
+ dist.barrier()
+
+
+@functools.lru_cache()
+def _get_global_gloo_group():
+ """
+ Return a process group based on gloo backend, containing all the ranks
+ The result is cached.
+ Returns:
+ (group): pytorch dist group.
+ """
+ if dist.get_backend() == "nccl":
+ return dist.new_group(backend="gloo")
+ else:
+ return dist.group.WORLD
+
+
+def _serialize_to_tensor(data, group):
+ """
+ Seriialize the tensor to ByteTensor. Note that only `gloo` and `nccl`
+ backend is supported.
+ Args:
+ data (data): data to be serialized.
+ group (group): pytorch dist group.
+ Returns:
+ tensor (ByteTensor): tensor that serialized.
+ """
+
+ backend = dist.get_backend(group)
+ assert backend in ["gloo", "nccl"]
+ device = torch.device("cpu" if backend == "gloo" else "cuda")
+
+ buffer = pickle.dumps(data)
+ if len(buffer) > 1024 ** 3:
+ logger = logging.getLogger(__name__)
+ logger.warning(
+ "Rank {} trying to all-gather {:.2f} GB of data on device {}".format(
+ get_rank(), len(buffer) / (1024 ** 3), device
+ )
+ )
+ storage = torch.ByteStorage.from_buffer(buffer)
+ tensor = torch.ByteTensor(storage).to(device=device)
+ return tensor
+
+
+def _pad_to_largest_tensor(tensor, group):
+ """
+ Padding all the tensors from different GPUs to the largest ones.
+ Args:
+ tensor (tensor): tensor to pad.
+ group (group): pytorch dist group.
+ Returns:
+ list[int]: size of the tensor, on each rank
+ Tensor: padded tensor that has the max size
+ """
+ world_size = dist.get_world_size(group=group)
+ assert (
+ world_size >= 1
+ ), "comm.gather/all_gather must be called from ranks within the given group!"
+ local_size = torch.tensor(
+ [tensor.numel()], dtype=torch.int64, device=tensor.device
+ )
+ size_list = [
+ torch.zeros([1], dtype=torch.int64, device=tensor.device)
+ for _ in range(world_size)
+ ]
+ dist.all_gather(size_list, local_size, group=group)
+ size_list = [int(size.item()) for size in size_list]
+
+ max_size = max(size_list)
+
+ # we pad the tensor because torch all_gather does not support
+ # gathering tensors of different shapes
+ if local_size != max_size:
+ padding = torch.zeros(
+ (max_size - local_size,), dtype=torch.uint8, device=tensor.device
+ )
+ tensor = torch.cat((tensor, padding), dim=0)
+ return size_list, tensor
+
+
+def all_gather_unaligned(data, group=None):
+ """
+ Run all_gather on arbitrary picklable data (not necessarily tensors).
+
+ Args:
+ data: any picklable object
+ group: a torch process group. By default, will use a group which
+ contains all ranks on gloo backend.
+
+ Returns:
+ list[data]: list of data gathered from each rank
+ """
+ if get_world_size() == 1:
+ return [data]
+ if group is None:
+ group = _get_global_gloo_group()
+ if dist.get_world_size(group) == 1:
+ return [data]
+
+ tensor = _serialize_to_tensor(data, group)
+
+ size_list, tensor = _pad_to_largest_tensor(tensor, group)
+ max_size = max(size_list)
+
+ # receiving Tensor from all ranks
+ tensor_list = [
+ torch.empty((max_size,), dtype=torch.uint8, device=tensor.device)
+ for _ in size_list
+ ]
+ dist.all_gather(tensor_list, tensor, group=group)
+
+ data_list = []
+ for size, tensor in zip(size_list, tensor_list):
+ buffer = tensor.cpu().numpy().tobytes()[:size]
+ data_list.append(pickle.loads(buffer))
+
+ return data_list
+
+
+def init_distributed_training(cfg):
+ """
+ Initialize variables needed for distributed training.
+ """
+ if cfg.NUM_GPUS <= 1:
+ return
+ num_gpus_per_machine = cfg.NUM_GPUS
+ num_machines = dist.get_world_size() // num_gpus_per_machine
+ for i in range(num_machines):
+ ranks_on_i = list(
+ range(i * num_gpus_per_machine, (i + 1) * num_gpus_per_machine)
+ )
+ pg = dist.new_group(ranks_on_i)
+ if i == cfg.SHARD_ID:
+ global _LOCAL_PROCESS_GROUP
+ _LOCAL_PROCESS_GROUP = pg
+
+
+def get_local_size() -> int:
+ """
+ Returns:
+ The size of the per-machine process group,
+ i.e. the number of processes per machine.
+ """
+ if not dist.is_available():
+ return 1
+ if not dist.is_initialized():
+ return 1
+ return dist.get_world_size(group=_LOCAL_PROCESS_GROUP)
+
+
+def get_local_rank() -> int:
+ """
+ Returns:
+ The rank of the current process within the local (per-machine) process group.
+ """
+ if not dist.is_available():
+ return 0
+ if not dist.is_initialized():
+ return 0
+ assert _LOCAL_PROCESS_GROUP is not None
+ return dist.get_rank(group=_LOCAL_PROCESS_GROUP)
diff --git a/clean/video/pwtf_dvd/inference/slowfast/utils/env.py b/clean/video/pwtf_dvd/inference/slowfast/utils/env.py
new file mode 100644
index 0000000000000000000000000000000000000000..2554915089a6e20c9ce58bba7fa59136ed65c887
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/utils/env.py
@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Set up Environment."""
+
+import slowfast.utils.logging as logging
+
+_ENV_SETUP_DONE = False
+
+
+def setup_environment():
+ global _ENV_SETUP_DONE
+ if _ENV_SETUP_DONE:
+ return
+ _ENV_SETUP_DONE = True
diff --git a/clean/video/pwtf_dvd/inference/slowfast/utils/logging.py b/clean/video/pwtf_dvd/inference/slowfast/utils/logging.py
new file mode 100644
index 0000000000000000000000000000000000000000..f2763b3000c9be4a5499b310a3bc052fd5472d50
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/utils/logging.py
@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Logging."""
+
+import builtins
+import decimal
+import functools
+import logging
+import os
+import sys
+import simplejson
+from fvcore.common.file_io import PathManager
+
+import slowfast.utils.distributed as du
+
+
+def _suppress_print():
+ """
+ Suppresses printing from the current process.
+ """
+
+ def print_pass(*objects, sep=" ", end="\n", file=sys.stdout, flush=False):
+ pass
+
+ builtins.print = print_pass
+
+
+@functools.lru_cache(maxsize=None)
+def _cached_log_stream(filename):
+ return PathManager.open(filename, "a")
+
+
+def setup_logging(output_dir=None):
+ """
+ Sets up the logging for multiple processes. Only enable the logging for the
+ master process, and suppress logging for the non-master processes.
+ """
+ # Set up logging format.
+ _FORMAT = "[%(levelname)s: %(filename)s: %(lineno)4d]: %(message)s"
+
+ if du.is_master_proc():
+ # Enable logging for the master process.
+ logging.root.handlers = []
+ else:
+ # Suppress logging for non-master processes.
+ _suppress_print()
+
+ logger = logging.getLogger()
+ logger.setLevel(logging.DEBUG)
+ logger.propagate = False
+ plain_formatter = logging.Formatter(
+ "[%(asctime)s][%(levelname)s] %(name)s: %(lineno)4d: %(message)s",
+ datefmt="%m/%d %H:%M:%S",
+ )
+
+ if du.is_master_proc():
+ ch = logging.StreamHandler(stream=sys.stdout)
+ ch.setLevel(logging.DEBUG)
+ ch.setFormatter(plain_formatter)
+ logger.addHandler(ch)
+
+ if output_dir is not None and du.is_master_proc(du.get_world_size()):
+ filename = os.path.join(output_dir, "stdout.log")
+ fh = logging.StreamHandler(_cached_log_stream(filename))
+ fh.setLevel(logging.DEBUG)
+ fh.setFormatter(plain_formatter)
+ logger.addHandler(fh)
+
+
+def get_logger(name):
+ """
+ Retrieve the logger with the specified name or, if name is None, return a
+ logger which is the root logger of the hierarchy.
+ Args:
+ name (string): name of the logger.
+ """
+ return logging.getLogger(name)
+
+
+def log_json_stats(stats):
+ """
+ Logs json stats.
+ Args:
+ stats (dict): a dictionary of statistical information to log.
+ """
+ stats = {
+ k: decimal.Decimal("{:.6f}".format(v)) if isinstance(v, float) else v
+ for k, v in stats.items()
+ }
+ json_stats = simplejson.dumps(stats, sort_keys=True, use_decimal=True)
+ logger = get_logger(__name__)
+ logger.info("json_stats: {:s}".format(json_stats))
diff --git a/clean/video/pwtf_dvd/inference/slowfast/utils/lr_policy.py b/clean/video/pwtf_dvd/inference/slowfast/utils/lr_policy.py
new file mode 100644
index 0000000000000000000000000000000000000000..4c67f8e5d9d6d576928986521dbe2eac8b1ca7e5
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/utils/lr_policy.py
@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Learning rate policy."""
+
+import math
+
+
+def get_lr_at_epoch(cfg, cur_epoch):
+ """
+ Retrieve the learning rate of the current epoch with the option to perform
+ warm up in the beginning of the training stage.
+ Args:
+ cfg (CfgNode): configs. Details can be found in
+ slowfast/config/defaults.py
+ cur_epoch (float): the number of epoch of the current training stage.
+ """
+ lr = get_lr_func(cfg.SOLVER.LR_POLICY)(cfg, cur_epoch)
+ # Perform warm up.
+ if cur_epoch < cfg.SOLVER.WARMUP_EPOCHS:
+ lr_start = cfg.SOLVER.WARMUP_START_LR
+ lr_end = get_lr_func(cfg.SOLVER.LR_POLICY)(
+ cfg, cfg.SOLVER.WARMUP_EPOCHS
+ )
+ alpha = (lr_end - lr_start) / cfg.SOLVER.WARMUP_EPOCHS
+ lr = cur_epoch * alpha + lr_start
+ return lr
+
+def get_lr_at_iter(cfg,cur_iter):
+ """LR schedule that should yield 76% converged accuracy with batch size 256"""
+ start_step = cfg.SOLVER.TOTAL_STEP- cfg.SOLVER.LR_STEP
+ duration_step = cfg.SOLVER.LR_STEP
+ base_lr=float(cfg.SOLVER.BASE_LR)
+ if cur_iter <= start_step:
+ return base_lr
+ else:
+ this_step = cur_iter - start_step
+ lr = base_lr * ((this_step / duration_step) ** 2.0)
+ return lr
+
+
+def lr_func_cosine(cfg, cur_epoch):
+ """
+ Retrieve the learning rate to specified values at specified epoch with the
+ cosine learning rate schedule. Details can be found in:
+ Ilya Loshchilov, and Frank Hutter
+ SGDR: Stochastic Gradient Descent With Warm Restarts.
+ Args:
+ cfg (CfgNode): configs. Details can be found in
+ slowfast/config/defaults.py
+ cur_epoch (float): the number of epoch of the current training stage.
+ """
+ return (
+ cfg.SOLVER.BASE_LR
+ * (math.cos(math.pi * cur_epoch / cfg.SOLVER.MAX_EPOCH) + 1.0)
+ * 0.5
+ )
+
+
+def lr_func_steps_with_relative_lrs(cfg, cur_epoch):
+ """
+ Retrieve the learning rate to specified values at specified epoch with the
+ steps with relative learning rate schedule.
+ Args:
+ cfg (CfgNode): configs. Details can be found in
+ slowfast/config/defaults.py
+ cur_epoch (float): the number of epoch of the current training stage.
+ """
+ ind = get_step_index(cfg, cur_epoch)
+ return cfg.SOLVER.LRS[ind] * cfg.SOLVER.BASE_LR
+
+
+def get_step_index(cfg, cur_epoch):
+ """
+ Retrieves the lr step index for the given epoch.
+ Args:
+ cfg (CfgNode): configs. Details can be found in
+ slowfast/config/defaults.py
+ cur_epoch (float): the number of epoch of the current training stage.
+ """
+ steps = cfg.SOLVER.STEPS + [cfg.SOLVER.MAX_EPOCH]
+ for ind, step in enumerate(steps): # NoQA
+ if cur_epoch < step:
+ break
+ return ind - 1
+
+
+def get_lr_func(lr_policy):
+ """
+ Given the configs, retrieve the specified lr policy function.
+ Args:
+ lr_policy (string): the learning rate policy to use for the job.
+ """
+ policy = "lr_func_" + lr_policy
+ if policy not in globals():
+ raise NotImplementedError("Unknown LR policy: {}".format(lr_policy))
+ else:
+ return globals()[policy]
diff --git a/clean/video/pwtf_dvd/inference/slowfast/utils/meters.py b/clean/video/pwtf_dvd/inference/slowfast/utils/meters.py
new file mode 100644
index 0000000000000000000000000000000000000000..2c4e9582a2bb5f6685987ce1f0ce391ac3950419
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/utils/meters.py
@@ -0,0 +1,841 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Meters."""
+
+import datetime
+import numpy as np
+import os
+from collections import defaultdict, deque
+import torch
+from fvcore.common.timer import Timer
+from sklearn.metrics import average_precision_score
+
+import slowfast.datasets.ava_helper as ava_helper
+import slowfast.utils.logging as logging
+import slowfast.utils.metrics as metrics
+import slowfast.utils.misc as misc
+from slowfast.utils.ava_eval_helper import (
+ evaluate_ava,
+ read_csv,
+ read_exclusions,
+ read_labelmap,
+)
+
+logger = logging.get_logger(__name__)
+
+
+def get_ava_mini_groundtruth(full_groundtruth):
+ """
+ Get the groundtruth annotations corresponding the "subset" of AVA val set.
+ We define the subset to be the frames such that (second % 4 == 0).
+ We optionally use subset for faster evaluation during training
+ (in order to track training progress).
+ Args:
+ full_groundtruth(dict): list of groundtruth.
+ """
+ ret = [defaultdict(list), defaultdict(list), defaultdict(list)]
+
+ for i in range(3):
+ for key in full_groundtruth[i].keys():
+ if int(key.split(",")[1]) % 4 == 0:
+ ret[i][key] = full_groundtruth[i][key]
+ return ret
+
+
+class AVAMeter(object):
+ """
+ Measure the AVA train, val, and test stats.
+ """
+
+ def __init__(self, overall_iters, cfg, mode):
+ """
+ overall_iters (int): the overall number of iterations of one epoch.
+ cfg (CfgNode): configs.
+ mode (str): `train`, `val`, or `test` mode.
+ """
+ self.cfg = cfg
+ self.lr = None
+ self.loss = ScalarMeter(cfg.LOG_PERIOD)
+ self.full_ava_test = cfg.AVA.FULL_TEST_ON_VAL
+ self.mode = mode
+ self.iter_timer = Timer()
+ self.all_preds = []
+ self.all_ori_boxes = []
+ self.all_metadata = []
+ self.overall_iters = overall_iters
+ self.excluded_keys = read_exclusions(
+ os.path.join(cfg.AVA.ANNOTATION_DIR, cfg.AVA.EXCLUSION_FILE)
+ )
+ self.categories, self.class_whitelist = read_labelmap(
+ os.path.join(cfg.AVA.ANNOTATION_DIR, cfg.AVA.LABEL_MAP_FILE)
+ )
+ gt_filename = os.path.join(
+ cfg.AVA.ANNOTATION_DIR, cfg.AVA.GROUNDTRUTH_FILE
+ )
+ self.full_groundtruth = read_csv(gt_filename, self.class_whitelist)
+ self.mini_groundtruth = get_ava_mini_groundtruth(self.full_groundtruth)
+
+ _, self.video_idx_to_name = ava_helper.load_image_lists(
+ cfg, mode == "train"
+ )
+
+ def log_iter_stats(self, cur_epoch, cur_iter):
+ """
+ Log the stats.
+ Args:
+ cur_epoch (int): the current epoch.
+ cur_iter (int): the current iteration.
+ """
+
+ if (cur_iter + 1) % self.cfg.LOG_PERIOD != 0:
+ return
+
+ eta_sec = self.iter_timer.seconds() * (self.overall_iters - cur_iter)
+ eta = str(datetime.timedelta(seconds=int(eta_sec)))
+ if self.mode == "train":
+ stats = {
+ "_type": "{}_iter".format(self.mode),
+ "cur_epoch": "{}".format(cur_epoch + 1),
+ "cur_iter": "{}".format(cur_iter + 1),
+ "eta": eta,
+ "time_diff": self.iter_timer.seconds(),
+ "mode": self.mode,
+ "loss": self.loss.get_win_median(),
+ "lr": self.lr,
+ }
+ elif self.mode == "val":
+ stats = {
+ "_type": "{}_iter".format(self.mode),
+ "cur_epoch": "{}".format(cur_epoch + 1),
+ "cur_iter": "{}".format(cur_iter + 1),
+ "eta": eta,
+ "time_diff": self.iter_timer.seconds(),
+ "mode": self.mode,
+ }
+ elif self.mode == "test":
+ stats = {
+ "_type": "{}_iter".format(self.mode),
+ "cur_iter": "{}".format(cur_iter + 1),
+ "eta": eta,
+ "time_diff": self.iter_timer.seconds(),
+ "mode": self.mode,
+ }
+ else:
+ raise NotImplementedError("Unknown mode: {}".format(self.mode))
+
+ logging.log_json_stats(stats)
+
+ def iter_tic(self):
+ """
+ Start to record time.
+ """
+ self.iter_timer.reset()
+
+ def iter_toc(self):
+ """
+ Stop to record time.
+ """
+ self.iter_timer.pause()
+
+ def reset(self):
+ """
+ Reset the Meter.
+ """
+ self.loss.reset()
+
+ self.all_preds = []
+ self.all_ori_boxes = []
+ self.all_metadata = []
+
+ def update_stats(self, preds, ori_boxes, metadata, loss=None, lr=None):
+ """
+ Update the current stats.
+ Args:
+ preds (tensor): prediction embedding.
+ ori_boxes (tensor): original boxes (x1, y1, x2, y2).
+ metadata (tensor): metadata of the AVA data.
+ loss (float): loss value.
+ lr (float): learning rate.
+ """
+ if self.mode in ["val", "test"]:
+ self.all_preds.append(preds)
+ self.all_ori_boxes.append(ori_boxes)
+ self.all_metadata.append(metadata)
+ if loss is not None:
+ self.loss.add_value(loss)
+ if lr is not None:
+ self.lr = lr
+
+ def finalize_metrics(self, log=True):
+ """
+ Calculate and log the final AVA metrics.
+ """
+ all_preds = torch.cat(self.all_preds, dim=0)
+ all_ori_boxes = torch.cat(self.all_ori_boxes, dim=0)
+ all_metadata = torch.cat(self.all_metadata, dim=0)
+
+ if self.mode == "test" or (self.full_ava_test and self.mode == "val"):
+ groundtruth = self.full_groundtruth
+ else:
+ groundtruth = self.mini_groundtruth
+
+ self.full_map = evaluate_ava(
+ all_preds,
+ all_ori_boxes,
+ all_metadata.tolist(),
+ self.excluded_keys,
+ self.class_whitelist,
+ self.categories,
+ groundtruth=groundtruth,
+ video_idx_to_name=self.video_idx_to_name,
+ )
+ if log:
+ stats = {"mode": self.mode, "map": self.full_map}
+ logging.log_json_stats(stats)
+
+ def log_epoch_stats(self, cur_epoch):
+ """
+ Log the stats of the current epoch.
+ Args:
+ cur_epoch (int): the number of current epoch.
+ """
+ if self.mode in ["val", "test"]:
+ self.finalize_metrics(log=False)
+ stats = {
+ "_type": "{}_epoch".format(self.mode),
+ "cur_epoch": "{}".format(cur_epoch + 1),
+ "mode": self.mode,
+ "map": self.full_map,
+ "gpu_mem": "{:.2f} GB".format(misc.gpu_mem_usage()),
+ "RAM": "{:.2f}/{:.2f} GB".format(*misc.cpu_mem_usage()),
+ }
+ logging.log_json_stats(stats)
+
+
+class TestMeter(object):
+ """
+ Perform the multi-view ensemble for testing: each video with an unique index
+ will be sampled with multiple clips, and the predictions of the clips will
+ be aggregated to produce the final prediction for the video.
+ The accuracy is calculated with the given ground truth labels.
+ """
+
+ def __init__(
+ self,
+ num_videos,
+ num_clips,
+ num_cls,
+ overall_iters,
+ multi_label=False,
+ ensemble_method="sum",
+ ):
+ """
+ Construct tensors to store the predictions and labels. Expect to get
+ num_clips predictions from each video, and calculate the metrics on
+ num_videos videos.
+ Args:
+ num_videos (int): number of videos to test.
+ num_clips (int): number of clips sampled from each video for
+ aggregating the final prediction for the video.
+ num_cls (int): number of classes for each prediction.
+ overall_iters (int): overall iterations for testing.
+ multi_label (bool): if True, use map as the metric.
+ ensemble_method (str): method to perform the ensemble, options
+ include "sum", and "max".
+ """
+
+ self.iter_timer = Timer()
+ self.num_clips = num_clips
+ self.overall_iters = overall_iters
+ self.multi_label = multi_label
+ self.ensemble_method = ensemble_method
+ # Initialize tensors.
+ self.video_preds = torch.zeros((num_videos, num_cls))
+ if multi_label:
+ self.video_preds -= 1e10
+
+ self.video_labels = (
+ torch.zeros((num_videos, num_cls))
+ if multi_label
+ else torch.zeros((num_videos)).long()
+ )
+ self.clip_count = torch.zeros((num_videos)).long()
+ # Reset metric.
+ self.reset()
+
+ def reset(self):
+ """
+ Reset the metric.
+ """
+ self.clip_count.zero_()
+ self.video_preds.zero_()
+ if self.multi_label:
+ self.video_preds -= 1e10
+ self.video_labels.zero_()
+
+ def update_stats(self, preds, labels, clip_ids):
+ """
+ Collect the predictions from the current batch and perform on-the-flight
+ summation as ensemble.
+ Args:
+ preds (tensor): predictions from the current batch. Dimension is
+ N x C where N is the batch size and C is the channel size
+ (num_cls).
+ labels (tensor): the corresponding labels of the current batch.
+ Dimension is N.
+ clip_ids (tensor): clip indexes of the current batch, dimension is
+ N.
+ """
+ for ind in range(preds.shape[0]):
+ vid_id = int(clip_ids[ind]) // self.num_clips
+ if self.video_labels[vid_id].sum() > 0:
+ assert torch.equal(
+ self.video_labels[vid_id].type(torch.FloatTensor),
+ labels[ind].type(torch.FloatTensor),
+ )
+ self.video_labels[vid_id] = labels[ind]
+ if self.ensemble_method == "sum":
+ self.video_preds[vid_id] += preds[ind]
+ elif self.ensemble_method == "max":
+ self.video_preds[vid_id] = torch.max(
+ self.video_preds[vid_id], preds[ind]
+ )
+ else:
+ raise NotImplementedError(
+ "Ensemble Method {} is not supported".format(
+ self.ensemble_method
+ )
+ )
+ self.clip_count[vid_id] += 1
+
+ def log_iter_stats(self, cur_iter):
+ """
+ Log the stats.
+ Args:
+ cur_iter (int): the current iteration of testing.
+ """
+ eta_sec = self.iter_timer.seconds() * (self.overall_iters - cur_iter)
+ eta = str(datetime.timedelta(seconds=int(eta_sec)))
+ stats = {
+ "split": "test_iter",
+ "cur_iter": "{}".format(cur_iter + 1),
+ "eta": eta,
+ "time_diff": self.iter_timer.seconds(),
+ }
+ logging.log_json_stats(stats)
+
+ def iter_tic(self):
+ self.iter_timer.reset()
+
+ def iter_toc(self):
+ self.iter_timer.pause()
+
+ def finalize_metrics(self, ks=(1, 5)):
+ """
+ Calculate and log the final ensembled metrics.
+ ks (tuple): list of top-k values for topk_accuracies. For example,
+ ks = (1, 5) correspods to top-1 and top-5 accuracy.
+ """
+ if not all(self.clip_count == self.num_clips):
+ logger.warning(
+ "clip count {} ~= num clips {}".format(
+ ", ".join(
+ [
+ "{}: {}".format(i, k)
+ for i, k in enumerate(self.clip_count.tolist())
+ ]
+ ),
+ self.num_clips,
+ )
+ )
+
+ stats = {"split": "test_final"}
+ if self.multi_label:
+ map = get_map(
+ self.video_preds.cpu().numpy(), self.video_labels.cpu().numpy()
+ )
+ stats["map"] = map
+ else:
+ num_topks_correct = metrics.topks_correct(
+ self.video_preds, self.video_labels, ks
+ )
+ topks = [
+ (x / self.video_preds.size(0)) * 100.0
+ for x in num_topks_correct
+ ]
+ assert len({len(ks), len(topks)}) == 1
+ for k, topk in zip(ks, topks):
+ stats["top{}_acc".format(k)] = "{:.{prec}f}".format(
+ topk, prec=2
+ )
+ logging.log_json_stats(stats)
+
+
+class ScalarMeter(object):
+ """
+ A scalar meter uses a deque to track a series of scaler values with a given
+ window size. It supports calculating the median and average values of the
+ window, and also supports calculating the global average.
+ """
+
+ def __init__(self, window_size):
+ """
+ Args:
+ window_size (int): size of the max length of the deque.
+ """
+ self.deque = deque(maxlen=window_size)
+ self.total = 0.0
+ self.count = 0
+
+ def reset(self):
+ """
+ Reset the deque.
+ """
+ self.deque.clear()
+ self.total = 0.0
+ self.count = 0
+
+ def add_value(self, value):
+ """
+ Add a new scalar value to the deque.
+ """
+ self.deque.append(value)
+ self.count += 1
+ self.total += value
+
+ def get_win_median(self):
+ """
+ Calculate the current median value of the deque.
+ """
+ return np.median(self.deque)
+
+ def get_win_avg(self):
+ """
+ Calculate the current average value of the deque.
+ """
+ return np.mean(self.deque)
+
+ def get_global_avg(self):
+ """
+ Calculate the global mean value.
+ """
+ return self.total / self.count
+
+
+class TrainMeter(object):
+ """
+ Measure training stats.
+ """
+
+ def __init__(self, epoch_iters, cfg):
+ """
+ Args:
+ epoch_iters (int): the overall number of iterations of one epoch.
+ cfg (CfgNode): configs.
+ """
+ self._cfg = cfg
+ self.epoch_iters = epoch_iters
+ self.MAX_EPOCH = cfg.SOLVER.MAX_EPOCH * epoch_iters
+ self.iter_timer = Timer()
+ self.loss = ScalarMeter(cfg.LOG_PERIOD)
+ self.loss_total = 0.0
+ self.lr = None
+ # Current minibatch errors (smoothed over a window).
+ self.mb_top1_err = ScalarMeter(cfg.LOG_PERIOD)
+ self.mb_top5_err = ScalarMeter(cfg.LOG_PERIOD)
+ # Number of misclassified examples.
+ self.num_top1_mis = 0
+ self.num_top5_mis = 0
+ self.num_samples = 0
+
+ def reset(self):
+ """
+ Reset the Meter.
+ """
+ self.loss.reset()
+ self.loss_total = 0.0
+ self.lr = None
+ self.mb_top1_err.reset()
+ self.mb_top5_err.reset()
+ self.num_top1_mis = 0
+ self.num_top5_mis = 0
+ self.num_samples = 0
+
+ def iter_tic(self):
+ """
+ Start to record time.
+ """
+ self.iter_timer.reset()
+
+ def iter_toc(self):
+ """
+ Stop to record time.
+ """
+ self.iter_timer.pause()
+
+ def update_stats(self, top1_err, top5_err, loss, lr, mb_size):
+ """
+ Update the current stats.
+ Args:
+ top1_err (float): top1 error rate.
+ top5_err (float): top5 error rate.
+ loss (float): loss value.
+ lr (float): learning rate.
+ mb_size (int): mini batch size.
+ """
+ self.loss.add_value(loss)
+ self.lr = lr
+ self.loss_total += loss * mb_size
+ self.num_samples += mb_size
+
+ if not self._cfg.DATA.MULTI_LABEL:
+ # Current minibatch stats
+ self.mb_top1_err.add_value(top1_err)
+ self.mb_top5_err.add_value(top5_err)
+ # Aggregate stats
+ self.num_top1_mis += top1_err * mb_size
+ self.num_top5_mis += top5_err * mb_size
+
+ def log_iter_stats(self, cur_epoch, cur_iter):
+ """
+ log the stats of the current iteration.
+ Args:
+ cur_epoch (int): the number of current epoch.
+ cur_iter (int): the number of current iteration.
+ """
+ if (cur_iter + 1) % self._cfg.LOG_PERIOD != 0:
+ return
+ eta_sec = self.iter_timer.seconds() * (
+ self.MAX_EPOCH - (cur_epoch * self.epoch_iters + cur_iter + 1)
+ )
+ eta = str(datetime.timedelta(seconds=int(eta_sec)))
+ stats = {
+ "_type": "train_iter",
+ "epoch": "{}/{}".format(cur_epoch + 1, self._cfg.SOLVER.MAX_EPOCH),
+ "iter": "{}/{}".format(cur_iter + 1, self.epoch_iters),
+ "time_diff": self.iter_timer.seconds(),
+ "eta": eta,
+ "loss": self.loss.get_win_median(),
+ "lr": self.lr,
+ "gpu_mem": "{:.2f} GB".format(misc.gpu_mem_usage()),
+ }
+ if not self._cfg.DATA.MULTI_LABEL:
+ stats["top1_err"] = self.mb_top1_err.get_win_median()
+ stats["top5_err"] = self.mb_top5_err.get_win_median()
+ logging.log_json_stats(stats)
+
+ def log_epoch_stats(self, cur_epoch):
+ """
+ Log the stats of the current epoch.
+ Args:
+ cur_epoch (int): the number of current epoch.
+ """
+ eta_sec = self.iter_timer.seconds() * (
+ self.MAX_EPOCH - (cur_epoch + 1) * self.epoch_iters
+ )
+ eta = str(datetime.timedelta(seconds=int(eta_sec)))
+ stats = {
+ "_type": "train_epoch",
+ "epoch": "{}/{}".format(cur_epoch + 1, self._cfg.SOLVER.MAX_EPOCH),
+ "time_diff": self.iter_timer.seconds(),
+ "eta": eta,
+ "lr": self.lr,
+ "gpu_mem": "{:.2f} GB".format(misc.gpu_mem_usage()),
+ "RAM": "{:.2f}/{:.2f} GB".format(*misc.cpu_mem_usage()),
+ }
+ if not self._cfg.DATA.MULTI_LABEL:
+ top1_err = self.num_top1_mis / self.num_samples
+ top5_err = self.num_top5_mis / self.num_samples
+ avg_loss = self.loss_total / self.num_samples
+ stats["top1_err"] = top1_err
+ stats["top5_err"] = top5_err
+ stats["loss"] = avg_loss
+ logging.log_json_stats(stats)
+
+
+class TrainIterMeter(object):
+ """
+ Measure training stats.
+ """
+
+ def __init__(self, epoch_iters, cfg,extra=[]):
+ """
+ Args:
+ epoch_iters (int): the overall number of iterations of one epoch.
+ cfg (CfgNode): configs.
+ """
+ self._cfg = cfg
+ self.epoch_iters = epoch_iters
+ self.MAX_EPOCH = cfg.SOLVER.MAX_EPOCH * epoch_iters
+ self.iter_timer = Timer()
+ self.loss = ScalarMeter(cfg.LOG_PERIOD)
+ self.loss_total = 0.0
+ self.lr = None
+
+ # Number of misclassified examples.
+ self.num_samples = 0
+
+ self.meters={key:ScalarMeter(cfg.LOG_PERIOD) for key in extra}
+
+ def reset(self):
+ """
+ Reset the Meter.
+ """
+ self.loss.reset()
+ self.loss_total = 0.0
+ self.lr = None
+
+
+ self.num_samples = 0
+
+ for meter in self.meters.values():
+ meter.reset()
+
+ def iter_tic(self):
+ """
+ Start to record time.
+ """
+ self.iter_timer.reset()
+
+ def iter_toc(self):
+ """
+ Stop to record time.
+ """
+ self.iter_timer.pause()
+
+ def update_stats(self, loss, lr, mb_size,extra={}):
+ """
+ Update the current stats.
+ Args:
+ top1_err (float): top1 error rate.
+ top5_err (float): top5 error rate.
+ loss (float): loss value.
+ lr (float): learning rate.
+ mb_size (int): mini batch size.
+ """
+ self.loss.add_value(loss)
+ self.lr = lr
+ self.loss_total += loss * mb_size
+ self.num_samples += mb_size
+
+
+ for key,val in extra.items():
+ self.meters[key].add_value(val)
+
+ def log_iter_stats(self, cur_epoch, cur_iter,extra={}):
+ """
+ log the stats of the current iteration.
+ Args:
+ cur_epoch (int): the number of current epoch.
+ cur_iter (int): the number of current iteration.
+ """
+ if (cur_iter + 1) % self._cfg.LOG_PERIOD != 0:
+ return
+ eta_sec = self.iter_timer.seconds() * (
+ self.MAX_EPOCH - (cur_epoch * self.epoch_iters + cur_iter + 1)
+ )
+ eta = str(datetime.timedelta(seconds=int(eta_sec)))
+ stats = {
+ "_type": "train_iter",
+ "epoch": "{}/{}".format(cur_epoch + 1, self._cfg.SOLVER.MAX_EPOCH),
+ "iter": "{}/{}".format(cur_iter + 1, self.epoch_iters),
+ "time_diff": self.iter_timer.seconds(),
+ "eta": eta,
+ "loss": self.loss.get_win_median(),
+ "lr": self.lr,
+ "gpu_mem": "{:.2f} GB".format(misc.gpu_mem_usage()),
+ }
+
+ for key,meter in self.meters.items():
+ stats[key]=meter.get_win_median()
+ for key,val in extra.items():
+ stats[key]=val
+
+ logging.log_json_stats(stats)
+
+ def log_epoch_stats(self, cur_epoch):
+ """
+ Log the stats of the current epoch.
+ Args:
+ cur_epoch (int): the number of current epoch.
+ """
+ eta_sec = self.iter_timer.seconds() * (
+ self.MAX_EPOCH - (cur_epoch + 1) * self.epoch_iters
+ )
+ eta = str(datetime.timedelta(seconds=int(eta_sec)))
+ stats = {
+ "_type": "train_epoch",
+ "epoch": "{}/{}".format(cur_epoch + 1, self._cfg.SOLVER.MAX_EPOCH),
+ "time_diff": self.iter_timer.seconds(),
+ "eta": eta,
+ "lr": self.lr,
+ "gpu_mem": "{:.2f} GB".format(misc.gpu_mem_usage()),
+ "RAM": "{:.2f}/{:.2f} GB".format(*misc.cpu_mem_usage()),
+ }
+ if not self._cfg.DATA.MULTI_LABEL:
+ avg_loss = self.loss_total / self.num_samples
+ stats["loss"] = avg_loss
+ logging.log_json_stats(stats)
+
+
+
+
+class ValMeter(object):
+ """
+ Measures validation stats.
+ """
+
+ def __init__(self, max_iter, cfg):
+ """
+ Args:
+ max_iter (int): the max number of iteration of the current epoch.
+ cfg (CfgNode): configs.
+ """
+ self._cfg = cfg
+ self.max_iter = max_iter
+ self.iter_timer = Timer()
+ # Current minibatch errors (smoothed over a window).
+ self.mb_top1_err = ScalarMeter(cfg.LOG_PERIOD)
+ self.mb_top5_err = ScalarMeter(cfg.LOG_PERIOD)
+ # Min errors (over the full val set).
+ self.min_top1_err = 100.0
+ self.min_top5_err = 100.0
+ # Number of misclassified examples.
+ self.num_top1_mis = 0
+ self.num_top5_mis = 0
+ self.num_samples = 0
+ self.all_preds = []
+ self.all_labels = []
+
+ def reset(self):
+ """
+ Reset the Meter.
+ """
+ self.iter_timer.reset()
+ self.mb_top1_err.reset()
+ self.mb_top5_err.reset()
+ self.num_top1_mis = 0
+ self.num_top5_mis = 0
+ self.num_samples = 0
+ self.all_preds = []
+ self.all_labels = []
+
+ def iter_tic(self):
+ """
+ Start to record time.
+ """
+ self.iter_timer.reset()
+
+ def iter_toc(self):
+ """
+ Stop to record time.
+ """
+ self.iter_timer.pause()
+
+ def update_stats(self, top1_err, top5_err, mb_size):
+ """
+ Update the current stats.
+ Args:
+ top1_err (float): top1 error rate.
+ top5_err (float): top5 error rate.
+ mb_size (int): mini batch size.
+ """
+ self.mb_top1_err.add_value(top1_err)
+ self.mb_top5_err.add_value(top5_err)
+ self.num_top1_mis += top1_err * mb_size
+ self.num_top5_mis += top5_err * mb_size
+ self.num_samples += mb_size
+
+ def update_predictions(self, preds, labels):
+ """
+ Update predictions and labels.
+ Args:
+ preds (tensor): model output predictions.
+ labels (tensor): labels.
+ """
+ # TODO: merge update_prediction with update_stats.
+ self.all_preds.append(preds)
+ self.all_labels.append(labels)
+
+ def log_iter_stats(self, cur_epoch, cur_iter):
+ """
+ log the stats of the current iteration.
+ Args:
+ cur_epoch (int): the number of current epoch.
+ cur_iter (int): the number of current iteration.
+ """
+ if (cur_iter + 1) % self._cfg.LOG_PERIOD != 0:
+ return
+ eta_sec = self.iter_timer.seconds() * (self.max_iter - cur_iter - 1)
+ eta = str(datetime.timedelta(seconds=int(eta_sec)))
+ stats = {
+ "_type": "val_iter",
+ "epoch": "{}/{}".format(cur_epoch + 1, self._cfg.SOLVER.MAX_EPOCH),
+ "iter": "{}/{}".format(cur_iter + 1, self.max_iter),
+ "time_diff": self.iter_timer.seconds(),
+ "eta": eta,
+ "gpu_mem": "{:.2f} GB".format(misc.gpu_mem_usage()),
+ }
+ if not self._cfg.DATA.MULTI_LABEL:
+ stats["top1_err"] = self.mb_top1_err.get_win_median()
+ stats["top5_err"] = self.mb_top5_err.get_win_median()
+ logging.log_json_stats(stats)
+
+ def log_epoch_stats(self, cur_epoch):
+ """
+ Log the stats of the current epoch.
+ Args:
+ cur_epoch (int): the number of current epoch.
+ """
+ stats = {
+ "_type": "val_epoch",
+ "epoch": "{}/{}".format(cur_epoch + 1, self._cfg.SOLVER.MAX_EPOCH),
+ "time_diff": self.iter_timer.seconds(),
+ "gpu_mem": "{:.2f} GB".format(misc.gpu_mem_usage()),
+ "RAM": "{:.2f}/{:.2f} GB".format(*misc.cpu_mem_usage()),
+ }
+ if self._cfg.DATA.MULTI_LABEL:
+ stats["map"] = get_map(
+ torch.cat(self.all_preds).cpu().numpy(),
+ torch.cat(self.all_labels).cpu().numpy(),
+ )
+ else:
+ top1_err = self.num_top1_mis / self.num_samples
+ top5_err = self.num_top5_mis / self.num_samples
+ self.min_top1_err = min(self.min_top1_err, top1_err)
+ self.min_top5_err = min(self.min_top5_err, top5_err)
+
+ stats["top1_err"] = top1_err
+ stats["top5_err"] = top5_err
+ stats["min_top1_err"] = self.min_top1_err
+ stats["min_top5_err"] = self.min_top5_err
+
+ logging.log_json_stats(stats)
+
+
+def get_map(preds, labels):
+ """
+ Compute mAP for multi-label case.
+ Args:
+ preds (numpy tensor): num_examples x num_classes.
+ labels (numpy tensor): num_examples x num_classes.
+ Returns:
+ mean_ap (int): final mAP score.
+ """
+
+ logger.info("Getting mAP for {} examples".format(preds.shape[0]))
+
+ preds = preds[:, ~(np.all(labels == 0, axis=0))]
+ labels = labels[:, ~(np.all(labels == 0, axis=0))]
+ aps = [0]
+ try:
+ aps = average_precision_score(labels, preds, average=None)
+ except ValueError:
+ print(
+ "Average precision requires a sufficient number of samples \
+ in a batch which are missing in this sample."
+ )
+
+ mean_ap = np.mean(aps)
+ return mean_ap
diff --git a/clean/video/pwtf_dvd/inference/slowfast/utils/metrics.py b/clean/video/pwtf_dvd/inference/slowfast/utils/metrics.py
new file mode 100644
index 0000000000000000000000000000000000000000..0ef01b174aa5c3d54da77923f515f244327c4e80
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/utils/metrics.py
@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Functions for computing metrics."""
+
+import torch
+
+
+def topks_correct(preds, labels, ks):
+ """
+ Given the predictions, labels, and a list of top-k values, compute the
+ number of correct predictions for each top-k value.
+
+ Args:
+ preds (array): array of predictions. Dimension is batchsize
+ N x ClassNum.
+ labels (array): array of labels. Dimension is batchsize N.
+ ks (list): list of top-k values. For example, ks = [1, 5] correspods
+ to top-1 and top-5.
+
+ Returns:
+ topks_correct (list): list of numbers, where the `i`-th entry
+ corresponds to the number of top-`ks[i]` correct predictions.
+ """
+ assert preds.size(0) == labels.size(
+ 0
+ ), "Batch dim of predictions and labels must match"
+ # Find the top max_k predictions for each sample
+ _top_max_k_vals, top_max_k_inds = torch.topk(
+ preds, max(ks), dim=1, largest=True, sorted=True
+ )
+ # (batch_size, max_k) -> (max_k, batch_size).
+ top_max_k_inds = top_max_k_inds.t()
+ # (batch_size, ) -> (max_k, batch_size).
+ rep_max_k_labels = labels.view(1, -1).expand_as(top_max_k_inds)
+ # (i, j) = 1 if top i-th prediction for the j-th sample is correct.
+ top_max_k_correct = top_max_k_inds.eq(rep_max_k_labels)
+ # Compute the number of topk correct predictions for each k.
+ topks_correct = [
+ top_max_k_correct[:k, :].view(-1).float().sum() for k in ks
+ ]
+ return topks_correct
+
+
+def topk_errors(preds, labels, ks):
+ """
+ Computes the top-k error for each k.
+ Args:
+ preds (array): array of predictions. Dimension is N.
+ labels (array): array of labels. Dimension is N.
+ ks (list): list of ks to calculate the top accuracies.
+ """
+ num_topks_correct = topks_correct(preds, labels, ks)
+ return [(1.0 - x / preds.size(0)) * 100.0 for x in num_topks_correct]
+
+
+def topk_accuracies(preds, labels, ks):
+ """
+ Computes the top-k accuracy for each k.
+ Args:
+ preds (array): array of predictions. Dimension is N.
+ labels (array): array of labels. Dimension is N.
+ ks (list): list of ks to calculate the top accuracies.
+ """
+ num_topks_correct = topks_correct(preds, labels, ks)
+ return [(x / preds.size(0)) * 100.0 for x in num_topks_correct]
diff --git a/clean/video/pwtf_dvd/inference/slowfast/utils/misc.py b/clean/video/pwtf_dvd/inference/slowfast/utils/misc.py
new file mode 100644
index 0000000000000000000000000000000000000000..13fea609bca56845d61819707a66cbd807b956d8
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/utils/misc.py
@@ -0,0 +1,359 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+import json
+import logging
+import math
+import numpy as np
+import os
+from datetime import datetime
+import psutil
+import torch
+from fvcore.common.file_io import PathManager
+from fvcore.nn.activation_count import activation_count
+from fvcore.nn.flop_count import flop_count
+from matplotlib import pyplot as plt
+from torch import nn
+
+import slowfast.utils.logging as logging
+import slowfast.utils.multiprocessing as mpu
+from slowfast.datasets.utils import pack_pathway_output
+from slowfast.models.batchnorm_helper import SubBatchNorm3d
+
+logger = logging.get_logger(__name__)
+
+
+def check_nan_losses(loss):
+ """
+ Determine whether the loss is NaN (not a number).
+ Args:
+ loss (loss): loss to check whether is NaN.
+ """
+ if math.isnan(loss):
+ raise RuntimeError("ERROR: Got NaN losses {}".format(datetime.now()))
+
+
+def params_count(model):
+ """
+ Compute the number of parameters.
+ Args:
+ model (model): model to count the number of parameters.
+ """
+ return np.sum([p.numel() for p in model.parameters()]).item()
+
+
+def gpu_mem_usage():
+ """
+ Compute the GPU memory usage for the current device (GB).
+ """
+ if torch.cuda.is_available():
+ mem_usage_bytes = torch.cuda.max_memory_allocated()
+ else:
+ mem_usage_bytes = 0
+ return mem_usage_bytes / 1024 ** 3
+
+
+def cpu_mem_usage():
+ """
+ Compute the system memory (RAM) usage for the current device (GB).
+ Returns:
+ usage (float): used memory (GB).
+ total (float): total memory (GB).
+ """
+ vram = psutil.virtual_memory()
+ usage = (vram.total - vram.available) / 1024 ** 3
+ total = vram.total / 1024 ** 3
+
+ return usage, total
+
+
+def _get_model_analysis_input(cfg, use_train_input):
+ """
+ Return a dummy input for model analysis with batch size 1. The input is
+ used for analyzing the model (counting flops and activations etc.).
+ Args:
+ cfg (CfgNode): configs. Details can be found in
+ slowfast/config/defaults.py
+ use_train_input (bool): if True, return the input for training. Otherwise,
+ return the input for testing.
+
+ Returns:
+ inputs: the input for model analysis.
+ """
+ rgb_dimension = 3
+ if use_train_input:
+ input_tensors = torch.rand(
+ rgb_dimension,
+ cfg.DATA.NUM_FRAMES,
+ cfg.DATA.TRAIN_CROP_SIZE,
+ cfg.DATA.TRAIN_CROP_SIZE,
+ )
+ else:
+ input_tensors = torch.rand(
+ rgb_dimension,
+ cfg.DATA.NUM_FRAMES,
+ cfg.DATA.TEST_CROP_SIZE,
+ cfg.DATA.TEST_CROP_SIZE,
+ )
+ model_inputs = pack_pathway_output(cfg, input_tensors)
+ for i in range(len(model_inputs)):
+ model_inputs[i] = model_inputs[i].unsqueeze(0)
+ if cfg.NUM_GPUS:
+ model_inputs[i] = model_inputs[i].cuda(non_blocking=True)
+
+ # If detection is enabled, count flops for one proposal.
+ if cfg.DETECTION.ENABLE:
+ bbox = torch.tensor([[0, 0, 1.0, 0, 1.0]])
+ if cfg.NUM_GPUS:
+ bbox = bbox.cuda()
+ inputs = (model_inputs, bbox)
+ else:
+ inputs = (model_inputs,)
+ return inputs
+
+
+def get_model_stats(model, cfg, mode, use_train_input):
+ """
+ Compute statistics for the current model given the config.
+ Args:
+ model (model): model to perform analysis.
+ cfg (CfgNode): configs. Details can be found in
+ slowfast/config/defaults.py
+ mode (str): Options include `flop` or `activation`. Compute either flop
+ (gflops) or activation count (mega).
+ use_train_input (bool): if True, compute statistics for training. Otherwise,
+ compute statistics for testing.
+
+ Returns:
+ float: the total number of count of the given model.
+ """
+ assert mode in [
+ "flop",
+ "activation",
+ ], "'{}' not supported for model analysis".format(mode)
+ if mode == "flop":
+ model_stats_fun = flop_count
+ elif mode == "activation":
+ model_stats_fun = activation_count
+
+ # Set model to evaluation mode for analysis.
+ # Evaluation mode can avoid getting stuck with sync batchnorm.
+ model_mode = model.training
+ model.eval()
+ inputs = _get_model_analysis_input(cfg, use_train_input)
+ count_dict, _ = model_stats_fun(model, inputs)
+ count = sum(count_dict.values())
+ model.train(model_mode)
+ return count
+
+
+def log_model_info(model, cfg, use_train_input=True):
+ """
+ Log info, includes number of parameters, gpu usage, gflops and activation count.
+ The model info is computed when the model is in validation mode.
+ Args:
+ model (model): model to log the info.
+ cfg (CfgNode): configs. Details can be found in
+ slowfast/config/defaults.py
+ use_train_input (bool): if True, log info for training. Otherwise,
+ log info for testing.
+ """
+ print("Model:\n{}".format(model))
+ print("Params: {:,}".format(params_count(model)))
+ print("Mem: {:,} MB".format(gpu_mem_usage()))
+ print(
+ "Flops: {:,} G".format(
+ get_model_stats(model, cfg, "flop", use_train_input)
+ )
+ )
+ print(
+ "Activations: {:,} M".format(
+ get_model_stats(model, cfg, "activation", use_train_input)
+ )
+ )
+ logger.info("nvidia-smi")
+ os.system("nvidia-smi")
+
+def is_eval_epoch(cfg, cur_epoch, multigrid_schedule):
+ """
+ Determine if the model should be evaluated at the current epoch.
+ Args:
+ cfg (CfgNode): configs. Details can be found in
+ slowfast/config/defaults.py
+ cur_epoch (int): current epoch.
+ multigrid_schedule (List): schedule for multigrid training.
+ """
+ if cur_epoch + 1 == cfg.SOLVER.MAX_EPOCH:
+ return True
+ if multigrid_schedule is not None:
+ prev_epoch = 0
+ for s in multigrid_schedule:
+ if cur_epoch < s[-1]:
+ period = max(
+ (s[-1] - prev_epoch) // cfg.MULTIGRID.EVAL_FREQ + 1, 1
+ )
+ return (s[-1] - 1 - cur_epoch) % period == 0
+ prev_epoch = s[-1]
+
+ return (cur_epoch + 1) % cfg.TRAIN.EVAL_PERIOD == 0
+
+
+def plot_input(tensor, bboxes=(), texts=(), path="./tmp_vis.png"):
+ """
+ Plot the input tensor with the optional bounding box and save it to disk.
+ Args:
+ tensor (tensor): a tensor with shape of `NxCxHxW`.
+ bboxes (tuple): bounding boxes with format of [[x, y, h, w]].
+ texts (tuple): a tuple of string to plot.
+ path (str): path to the image to save to.
+ """
+ tensor = tensor - tensor.min()
+ tensor = tensor / tensor.max()
+ f, ax = plt.subplots(nrows=1, ncols=tensor.shape[0], figsize=(50, 20))
+ for i in range(tensor.shape[0]):
+ ax[i].axis("off")
+ ax[i].imshow(tensor[i].permute(1, 2, 0))
+ # ax[1][0].axis('off')
+ if bboxes is not None and len(bboxes) > i:
+ for box in bboxes[i]:
+ x1, y1, x2, y2 = box
+ ax[i].vlines(x1, y1, y2, colors="g", linestyles="solid")
+ ax[i].vlines(x2, y1, y2, colors="g", linestyles="solid")
+ ax[i].hlines(y1, x1, x2, colors="g", linestyles="solid")
+ ax[i].hlines(y2, x1, x2, colors="g", linestyles="solid")
+
+ if texts is not None and len(texts) > i:
+ ax[i].text(0, 0, texts[i])
+ f.savefig(path)
+
+
+def frozen_bn_stats(model):
+ """
+ Set all the bn layers to eval mode.
+ Args:
+ model (model): model to set bn layers to eval mode.
+ """
+ for m in model.modules():
+ if isinstance(m, nn.BatchNorm3d):
+ m.eval()
+
+
+def aggregate_sub_bn_stats(module):
+ """
+ Recursively find all SubBN modules and aggregate sub-BN stats.
+ Args:
+ module (nn.Module)
+ Returns:
+ count (int): number of SubBN module found.
+ """
+ count = 0
+ for child in module.children():
+ if isinstance(child, SubBatchNorm3d):
+ child.aggregate_stats()
+ count += 1
+ else:
+ count += aggregate_sub_bn_stats(child)
+ return count
+
+
+def launch_job(cfg, init_method, func, daemon=False):
+ """
+ Run 'func' on one or more GPUs, specified in cfg
+ Args:
+ cfg (CfgNode): configs. Details can be found in
+ slowfast/config/defaults.py
+ init_method (str): initialization method to launch the job with multiple
+ devices.
+ func (function): job to run on GPU(s)
+ daemon (bool): The spawned processes’ daemon flag. If set to True,
+ daemonic processes will be created
+ """
+ if cfg.NUM_GPUS > 1:
+ torch.multiprocessing.spawn(
+ mpu.run,
+ nprocs=cfg.NUM_GPUS,
+ args=(
+ cfg.NUM_GPUS,
+ func,
+ init_method,
+ cfg.SHARD_ID,
+ cfg.NUM_SHARDS,
+ cfg.DIST_BACKEND,
+ cfg,
+ ),
+ daemon=daemon,
+ )
+ else:
+ func(cfg=cfg)
+
+
+def get_class_names(path, parent_path=None, subset_path=None):
+ """
+ Read json file with entries {classname: index} and return
+ an array of class names in order.
+ If parent_path is provided, load and map all children to their ids.
+ Args:
+ path (str): path to class ids json file.
+ File must be in the format {"class1": id1, "class2": id2, ...}
+ parent_path (Optional[str]): path to parent-child json file.
+ File must be in the format {"parent1": ["child1", "child2", ...], ...}
+ subset_path (Optional[str]): path to text file containing a subset
+ of class names, separated by newline characters.
+ Returns:
+ class_names (list of strs): list of class names.
+ class_parents (dict): a dictionary where key is the name of the parent class
+ and value is a list of ids of the children classes.
+ subset_ids (list of ints): list of ids of the classes provided in the
+ subset file.
+ """
+ try:
+ with PathManager.open(path, "r") as f:
+ class2idx = json.load(f)
+ except Exception as err:
+ print("Fail to load file from {} with error {}".format(path, err))
+ return
+
+ max_key = max(class2idx.values())
+ class_names = [None] * (max_key + 1)
+
+ for k, i in class2idx.items():
+ class_names[i] = k
+
+ class_parent = None
+ if parent_path is not None and parent_path != "":
+ try:
+ with PathManager.open(parent_path, "r") as f:
+ d_parent = json.load(f)
+ except EnvironmentError as err:
+ print(
+ "Fail to load file from {} with error {}".format(
+ parent_path, err
+ )
+ )
+ return
+ class_parent = {}
+ for parent, children in d_parent.items():
+ indices = [
+ class2idx[c] for c in children if class2idx.get(c) is not None
+ ]
+ class_parent[parent] = indices
+
+ subset_ids = None
+ if subset_path is not None and subset_path != "":
+ try:
+ with PathManager.open(subset_path, "r") as f:
+ subset = f.read().split("\n")
+ subset_ids = [
+ class2idx[name]
+ for name in subset
+ if class2idx.get(name) is not None
+ ]
+ except EnvironmentError as err:
+ print(
+ "Fail to load file from {} with error {}".format(
+ subset_path, err
+ )
+ )
+ return
+
+ return class_names, class_parent, subset_ids
diff --git a/clean/video/pwtf_dvd/inference/slowfast/utils/multigrid.py b/clean/video/pwtf_dvd/inference/slowfast/utils/multigrid.py
new file mode 100644
index 0000000000000000000000000000000000000000..4aed24bb4889d30960cec5ca94ab20a73b40b9e1
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/utils/multigrid.py
@@ -0,0 +1,240 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Helper functions for multigrid training."""
+
+import numpy as np
+
+import slowfast.utils.logging as logging
+
+logger = logging.get_logger(__name__)
+
+
+class MultigridSchedule(object):
+ """
+ This class defines multigrid training schedule and update cfg accordingly.
+ """
+
+ def init_multigrid(self, cfg):
+ """
+ Update cfg based on multigrid settings.
+ Args:
+ cfg (configs): configs that contains training and multigrid specific
+ hyperparameters. Details can be seen in
+ slowfast/config/defaults.py.
+ Returns:
+ cfg (configs): the updated cfg.
+ """
+ self.schedule = None
+ # We may modify cfg.TRAIN.BATCH_SIZE, cfg.DATA.NUM_FRAMES, and
+ # cfg.DATA.TRAIN_CROP_SIZE during training, so we store their original
+ # value in cfg and use them as global variables.
+ cfg.MULTIGRID.DEFAULT_B = cfg.TRAIN.BATCH_SIZE
+ cfg.MULTIGRID.DEFAULT_T = cfg.DATA.NUM_FRAMES
+ cfg.MULTIGRID.DEFAULT_S = cfg.DATA.TRAIN_CROP_SIZE
+
+ if cfg.MULTIGRID.LONG_CYCLE:
+ self.schedule = self.get_long_cycle_schedule(cfg)
+ cfg.SOLVER.STEPS = [0] + [s[-1] for s in self.schedule]
+ # Fine-tuning phase.
+ cfg.SOLVER.STEPS[-1] = (
+ cfg.SOLVER.STEPS[-2] + cfg.SOLVER.STEPS[-1]
+ ) // 2
+ cfg.SOLVER.LRS = [
+ cfg.SOLVER.GAMMA ** s[0] * s[1][0] for s in self.schedule
+ ]
+ # Fine-tuning phase.
+ cfg.SOLVER.LRS = cfg.SOLVER.LRS[:-1] + [
+ cfg.SOLVER.LRS[-2],
+ cfg.SOLVER.LRS[-1],
+ ]
+
+ cfg.SOLVER.MAX_EPOCH = self.schedule[-1][-1]
+
+ elif cfg.MULTIGRID.SHORT_CYCLE:
+ cfg.SOLVER.STEPS = [
+ int(s * cfg.MULTIGRID.EPOCH_FACTOR) for s in cfg.SOLVER.STEPS
+ ]
+ cfg.SOLVER.MAX_EPOCH = int(
+ cfg.SOLVER.MAX_EPOCH * cfg.MULTIGRID.EPOCH_FACTOR
+ )
+ return cfg
+
+ def update_long_cycle(self, cfg, cur_epoch):
+ """
+ Before every epoch, check if long cycle shape should change. If it
+ should, update cfg accordingly.
+ Args:
+ cfg (configs): configs that contains training and multigrid specific
+ hyperparameters. Details can be seen in
+ slowfast/config/defaults.py.
+ cur_epoch (int): current epoch index.
+ Returns:
+ cfg (configs): the updated cfg.
+ changed (bool): do we change long cycle shape at this epoch?
+ """
+ base_b, base_t, base_s = get_current_long_cycle_shape(
+ self.schedule, cur_epoch
+ )
+ if base_s != cfg.DATA.TRAIN_CROP_SIZE or base_t != cfg.DATA.NUM_FRAMES:
+
+ cfg.DATA.NUM_FRAMES = base_t
+ cfg.DATA.TRAIN_CROP_SIZE = base_s
+ cfg.TRAIN.BATCH_SIZE = base_b * cfg.MULTIGRID.DEFAULT_B
+
+ bs_factor = (
+ float(cfg.TRAIN.BATCH_SIZE / cfg.NUM_GPUS)
+ / cfg.MULTIGRID.BN_BASE_SIZE
+ )
+
+ if bs_factor < 1:
+ cfg.BN.NORM_TYPE = "sync_batchnorm"
+ cfg.BN.NUM_SYNC_DEVICES = int(1.0 / bs_factor)
+ elif bs_factor > 1:
+ cfg.BN.NORM_TYPE = "sub_batchnorm"
+ cfg.BN.NUM_SPLITS = int(bs_factor)
+ else:
+ cfg.BN.NORM_TYPE = "batchnorm"
+
+ cfg.MULTIGRID.LONG_CYCLE_SAMPLING_RATE = cfg.DATA.SAMPLING_RATE * (
+ cfg.MULTIGRID.DEFAULT_T // cfg.DATA.NUM_FRAMES
+ )
+ logger.info("Long cycle updates:")
+ logger.info("\tBN.NORM_TYPE: {}".format(cfg.BN.NORM_TYPE))
+ if cfg.BN.NORM_TYPE == "sync_batchnorm":
+ logger.info(
+ "\tBN.NUM_SYNC_DEVICES: {}".format(cfg.BN.NUM_SYNC_DEVICES)
+ )
+ elif cfg.BN.NORM_TYPE == "sub_batchnorm":
+ logger.info("\tBN.NUM_SPLITS: {}".format(cfg.BN.NUM_SPLITS))
+ logger.info("\tTRAIN.BATCH_SIZE: {}".format(cfg.TRAIN.BATCH_SIZE))
+ logger.info(
+ "\tDATA.NUM_FRAMES x LONG_CYCLE_SAMPLING_RATE: {}x{}".format(
+ cfg.DATA.NUM_FRAMES, cfg.MULTIGRID.LONG_CYCLE_SAMPLING_RATE
+ )
+ )
+ logger.info(
+ "\tDATA.TRAIN_CROP_SIZE: {}".format(cfg.DATA.TRAIN_CROP_SIZE)
+ )
+ return cfg, True
+ else:
+ return cfg, False
+
+ def get_long_cycle_schedule(self, cfg):
+ """
+ Based on multigrid hyperparameters, define the schedule of a long cycle.
+ Args:
+ cfg (configs): configs that contains training and multigrid specific
+ hyperparameters. Details can be seen in
+ slowfast/config/defaults.py.
+ Returns:
+ schedule (list): Specifies a list long cycle base shapes and their
+ corresponding training epochs.
+ """
+
+ steps = cfg.SOLVER.STEPS
+
+ default_size = float(
+ cfg.DATA.NUM_FRAMES * cfg.DATA.TRAIN_CROP_SIZE ** 2
+ )
+ default_iters = steps[-1]
+
+ # Get shapes and average batch size for each long cycle shape.
+ avg_bs = []
+ all_shapes = []
+ for t_factor, s_factor in cfg.MULTIGRID.LONG_CYCLE_FACTORS:
+ base_t = int(round(cfg.DATA.NUM_FRAMES * t_factor))
+ base_s = int(round(cfg.DATA.TRAIN_CROP_SIZE * s_factor))
+ if cfg.MULTIGRID.SHORT_CYCLE:
+ shapes = [
+ [
+ base_t,
+ cfg.MULTIGRID.DEFAULT_S
+ * cfg.MULTIGRID.SHORT_CYCLE_FACTORS[0],
+ ],
+ [
+ base_t,
+ cfg.MULTIGRID.DEFAULT_S
+ * cfg.MULTIGRID.SHORT_CYCLE_FACTORS[1],
+ ],
+ [base_t, base_s],
+ ]
+ else:
+ shapes = [[base_t, base_s]]
+
+ # (T, S) -> (B, T, S)
+ shapes = [
+ [int(round(default_size / (s[0] * s[1] * s[1]))), s[0], s[1]]
+ for s in shapes
+ ]
+ avg_bs.append(np.mean([s[0] for s in shapes]))
+ all_shapes.append(shapes)
+
+ # Get schedule regardless of cfg.MULTIGRID.EPOCH_FACTOR.
+ total_iters = 0
+ schedule = []
+ for step_index in range(len(steps) - 1):
+ step_epochs = steps[step_index + 1] - steps[step_index]
+
+ for long_cycle_index, shapes in enumerate(all_shapes):
+ cur_epochs = (
+ step_epochs * avg_bs[long_cycle_index] / sum(avg_bs)
+ )
+
+ cur_iters = cur_epochs / avg_bs[long_cycle_index]
+ total_iters += cur_iters
+ schedule.append((step_index, shapes[-1], cur_epochs))
+
+ iter_saving = default_iters / total_iters
+
+ final_step_epochs = cfg.SOLVER.MAX_EPOCH - steps[-1]
+
+ # We define the fine-tuning phase to have the same amount of iteration
+ # saving as the rest of the training.
+ ft_epochs = final_step_epochs / iter_saving * avg_bs[-1]
+
+ schedule.append((step_index + 1, all_shapes[-1][2], ft_epochs))
+
+ # Obtrain final schedule given desired cfg.MULTIGRID.EPOCH_FACTOR.
+ x = (
+ cfg.SOLVER.MAX_EPOCH
+ * cfg.MULTIGRID.EPOCH_FACTOR
+ / sum(s[-1] for s in schedule)
+ )
+
+ final_schedule = []
+ total_epochs = 0
+ for s in schedule:
+ epochs = s[2] * x
+ total_epochs += epochs
+ final_schedule.append((s[0], s[1], int(round(total_epochs))))
+ print_schedule(final_schedule)
+ return final_schedule
+
+
+def print_schedule(schedule):
+ """
+ Log schedule.
+ """
+ logger.info("Long cycle index\tBase shape\tEpochs")
+ for s in schedule:
+ logger.info("{}\t{}\t{}".format(s[0], s[1], s[2]))
+
+
+def get_current_long_cycle_shape(schedule, epoch):
+ """
+ Given a schedule and epoch index, return the long cycle base shape.
+ Args:
+ schedule (configs): configs that contains training and multigrid specific
+ hyperparameters. Details can be seen in
+ slowfast/config/defaults.py.
+ cur_epoch (int): current epoch index.
+ Returns:
+ shapes (list): A list describing the base shape in a long cycle:
+ [batch size relative to default,
+ number of frames, spatial dimension].
+ """
+ for s in schedule:
+ if epoch < s[-1]:
+ return s[1]
+ return schedule[-1][1]
diff --git a/clean/video/pwtf_dvd/inference/slowfast/utils/multiprocessing.py b/clean/video/pwtf_dvd/inference/slowfast/utils/multiprocessing.py
new file mode 100644
index 0000000000000000000000000000000000000000..a56aa603697dd3a6a37871ddd4407523aeebbb8c
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/utils/multiprocessing.py
@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Multiprocessing helpers."""
+
+import torch
+
+
+def run(
+ local_rank, num_proc, func, init_method, shard_id, num_shards, backend, cfg
+):
+ """
+ Runs a function from a child process.
+ Args:
+ local_rank (int): rank of the current process on the current machine.
+ num_proc (int): number of processes per machine.
+ func (function): function to execute on each of the process.
+ init_method (string): method to initialize the distributed training.
+ TCP initialization: equiring a network address reachable from all
+ processes followed by the port.
+ Shared file-system initialization: makes use of a file system that
+ is shared and visible from all machines. The URL should start with
+ file:// and contain a path to a non-existent file on a shared file
+ system.
+ shard_id (int): the rank of the current machine.
+ num_shards (int): number of overall machines for the distributed
+ training job.
+ backend (string): three distributed backends ('nccl', 'gloo', 'mpi') are
+ supports, each with different capabilities. Details can be found
+ here:
+ https://pytorch.org/docs/stable/distributed.html
+ cfg (CfgNode): configs. Details can be found in
+ slowfast/config/defaults.py
+ """
+ # Initialize the process group.
+ world_size = num_proc * num_shards
+ rank = shard_id * num_proc + local_rank
+
+ try:
+ torch.distributed.init_process_group(
+ backend=backend,
+ init_method=init_method,
+ world_size=world_size,
+ rank=rank,
+ )
+ except Exception as e:
+ raise e
+
+ torch.cuda.set_device(local_rank)
+ func(cfg)
diff --git a/clean/video/pwtf_dvd/inference/slowfast/utils/parser.py b/clean/video/pwtf_dvd/inference/slowfast/utils/parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..06b4373e3b3736ceb310eed465fbc75e3bae1eb5
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/utils/parser.py
@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Argument parser functions."""
+
+import argparse
+import sys
+
+import slowfast.utils.checkpoint as cu
+from slowfast.config.defaults import get_cfg
+
+
+def parse_args():
+ """
+ Parse the following arguments for a default parser for PySlowFast users.
+ Args:
+ shard_id (int): shard id for the current machine. Starts from 0 to
+ num_shards - 1. If single machine is used, then set shard id to 0.
+ num_shards (int): number of shards using by the job.
+ init_method (str): initialization method to launch the job with multiple
+ devices. Options includes TCP or shared file-system for
+ initialization. details can be find in
+ https://pytorch.org/docs/stable/distributed.html#tcp-initialization
+ cfg (str): path to the config file.
+ opts (argument): provide addtional options from the command line, it
+ overwrites the config loaded from file.
+ """
+ parser = argparse.ArgumentParser(
+ description="Provide SlowFast video training and testing pipeline."
+ )
+ parser.add_argument(
+ "--shard_id",
+ help="The shard id of current node, Starts from 0 to num_shards - 1",
+ default=0,
+ type=int,
+ )
+ parser.add_argument(
+ "--num_shards",
+ help="Number of shards using by the job",
+ default=1,
+ type=int,
+ )
+ parser.add_argument(
+ "--init_method",
+ help="Initialization method, includes TCP or shared file-system",
+ default="tcp://localhost:9999",
+ type=str,
+ )
+ parser.add_argument(
+ "--cfg",
+ dest="cfg_file",
+ help="Path to the config file",
+ default="configs/Kinetics/SLOWFAST_4x16_R50.yaml",
+ type=str,
+ )
+ parser.add_argument(
+ "opts",
+ help="See slowfast/config/defaults.py for all options",
+ default=None,
+ nargs=argparse.REMAINDER,
+ )
+ if len(sys.argv) == 1:
+ parser.print_help()
+ return parser.parse_args()
+
+
+def load_config(args):
+ """
+ Given the arguemnts, load and initialize the configs.
+ Args:
+ args (argument): arguments includes `shard_id`, `num_shards`,
+ `init_method`, `cfg_file`, and `opts`.
+ """
+ # Setup cfg.
+ cfg = get_cfg()
+ # Load config from cfg.
+ if args.cfg_file is not None:
+ cfg.merge_from_file(args.cfg_file)
+ # Load config from command line, overwrite config from opts.
+ if args.opts is not None:
+ cfg.merge_from_list(args.opts)
+
+ # Inherit parameters from args.
+ if hasattr(args, "num_shards") and hasattr(args, "shard_id"):
+ cfg.NUM_SHARDS = args.num_shards
+ cfg.SHARD_ID = args.shard_id
+ if hasattr(args, "rng_seed"):
+ cfg.RNG_SEED = args.rng_seed
+ if hasattr(args, "output_dir"):
+ cfg.OUTPUT_DIR = args.output_dir
+
+ # Create the checkpoint dir.
+ cu.make_checkpoint_dir(cfg.OUTPUT_DIR)
+ return cfg
diff --git a/clean/video/pwtf_dvd/inference/slowfast/utils/weight_init_helper.py b/clean/video/pwtf_dvd/inference/slowfast/utils/weight_init_helper.py
new file mode 100644
index 0000000000000000000000000000000000000000..0b5544a70529f5dd1b06ba05a6aca4c7f508bdf3
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/slowfast/utils/weight_init_helper.py
@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+"""Utility function for weight initialization"""
+
+import torch.nn as nn
+from fvcore.nn.weight_init import c2_msra_fill
+
+
+def init_weights(model, fc_init_std=0.01, zero_init_final_bn=True):
+ """
+ Performs ResNet style weight initialization.
+ Args:
+ fc_init_std (float): the expected standard deviation for fc layer.
+ zero_init_final_bn (bool): if True, zero initialize the final bn for
+ every bottleneck.
+ """
+ for m in model.modules():
+ if isinstance(m, nn.Conv3d):
+ """
+ Follow the initialization method proposed in:
+ {He, Kaiming, et al.
+ "Delving deep into rectifiers: Surpassing human-level
+ performance on imagenet classification."
+ arXiv preprint arXiv:1502.01852 (2015)}
+ """
+ c2_msra_fill(m)
+ elif isinstance(m, nn.BatchNorm3d):
+ if (
+ hasattr(m, "transform_final_bn")
+ and m.transform_final_bn
+ and zero_init_final_bn
+ ):
+ batchnorm_weight = 0.0
+ else:
+ batchnorm_weight = 1.0
+ if m.weight is not None:
+ m.weight.data.fill_(batchnorm_weight)
+ if m.bias is not None:
+ m.bias.data.zero_()
+ if isinstance(m, nn.Linear):
+ m.weight.data.normal_(mean=0.0, std=fc_init_std)
+ m.bias.data.zero_()
diff --git a/clean/video/pwtf_dvd/inference/test_on_raw_video.py b/clean/video/pwtf_dvd/inference/test_on_raw_video.py
new file mode 100644
index 0000000000000000000000000000000000000000..9074dcd8ab946d4db33467594722f9e3e83199a7
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_on_raw_video.py
@@ -0,0 +1,251 @@
+
+import torch
+from torch.nn import functional as F
+import os
+import numpy as np
+from test_tools.common import detect_all, grab_all_frames
+from test_tools.utils import get_crop_box
+from test_tools.ct.operations import find_longest, multiple_tracking
+from test_tools.faster_crop_align_xray import FasterCropAlignXRay
+from test_tools.supply_writer import SupplyWriter
+import argparse
+from tqdm import tqdm
+from torchvision.transforms import Compose, ToTensor, Normalize
+from model.framework import get_model
+import cv2
+from PIL import Image
+
+mean = torch.tensor([0.485 * 255, 0.456 * 255, 0.406 * 255,]).cuda().view(1, 3, 1, 1, 1)
+std = torch.tensor([0.229 * 255, 0.224 * 255, 0.225 * 255,]).cuda().view(1, 3, 1, 1, 1)
+
+def main():
+ """Main inference function"""
+ parser = argparse.ArgumentParser(description="FTCN Face Depth Estimation Inference")
+ parser.add_argument("--video", type=str, help="Input video path", default="./examples/shining.mp4")
+ parser.add_argument("--out_dir", type=str, help="Output directory", default="./examples")
+ parser.add_argument("--model_path", type=str, help="Model checkpoint path", default="./model/for_deploy_model3.pth")
+ parser.add_argument("--max_frame", type=int, help="Maximum number of frames to process", default=768)
+ args = parser.parse_args()
+
+ # Load model
+ print("Loading model...")
+ model = get_model()
+ model.load_state_dict(torch.load(args.model_path, map_location='cpu'))
+ model.eval()
+ model.cuda()
+ print("Model loaded successfully!")
+
+ # Initialize preprocessing functions
+ crop_align_func = FasterCropAlignXRay(224)
+
+ # Setup input/output paths
+ input_file = args.video
+ os.makedirs(args.out_dir, exist_ok=True)
+ basename = os.path.splitext(os.path.basename(input_file))[0] +"_detect.mp4"
+ out_file = os.path.join(args.out_dir, basename)
+
+ # Process video frames
+ max_frame = args.max_frame
+ cache_file = f"{input_file}_{str(max_frame)}.pth"
+
+ if os.path.exists(cache_file):
+ print("Loading cached detection results...")
+ detect_res, all_lm68 = torch.load(cache_file)
+ frames = grab_all_frames(input_file, max_size=max_frame, cvt=True)
+ print("Detection results loaded from cache")
+ else:
+ print("Performing face detection...")
+ detect_res, all_lm68, frames = detect_all(
+ input_file, return_frames=True, max_size=max_frame
+ )
+ torch.save((detect_res, all_lm68), cache_file)
+ print("Face detection completed")
+
+ print(f"Processing {len(frames)} frames")
+
+
+ # Process detection results
+ shape = frames[0].shape[:2]
+ all_detect_res = []
+
+ assert len(all_lm68) == len(detect_res)
+
+ for faces, faces_lm68 in zip(detect_res, all_lm68):
+ new_faces = []
+ for (box, lm5, score), face_lm68 in zip(faces, faces_lm68):
+ new_face = (box, lm5, face_lm68, score)
+ new_faces.append(new_face)
+ all_detect_res.append(new_faces)
+
+ detect_res = all_detect_res
+
+ # Track faces across frames
+ print("Tracking faces across frames...")
+ tracks = multiple_tracking(detect_res)
+ tuples = [(0, len(detect_res))] * len(tracks)
+
+ print(f"Found {len(tracks)} face tracks")
+
+ if len(tracks) == 0:
+ print("No tracks found, using longest sequence...")
+ tuples, tracks = find_longest(detect_res)
+
+ # Extract face crops and landmarks
+ data_storage = {}
+ frame_boxes = {}
+ super_clips = []
+
+ for track_i, ((start, end), track) in enumerate(zip(tuples, tracks)):
+ print(f"Processing track {track_i}: frames {start}-{end}")
+ assert len(detect_res[start:end]) == len(track)
+
+ super_clips.append(len(track))
+
+ for face, frame_idx, j in zip(track, range(start, end), range(len(track))):
+ box, lm5, lm68 = face[:3]
+ big_box = get_crop_box(shape, box, scale=0.5)
+
+ top_left = big_box[:2][None, :]
+
+ new_lm5 = lm5 - top_left
+ new_lm68 = lm68 - top_left
+ new_box = (box.reshape(2, 2) - top_left).reshape(-1)
+
+ info = (new_box, new_lm5, new_lm68, big_box)
+
+ x1, y1, x2, y2 = big_box
+ cropped = frames[frame_idx][y1:y2, x1:x2]
+
+ base_key = f"{track_i}_{j}_"
+ data_storage[base_key + "img"] = cropped
+ data_storage[base_key + "ldm"] = info
+ data_storage[base_key + "idx"] = frame_idx
+
+ frame_boxes[frame_idx] = np.rint(box).astype(np.int32)
+
+ print(f"Sampling clips from super clips: {super_clips}")
+
+ # Generate clips for temporal analysis
+ clips_for_video = []
+ clip_size = 32
+ pad_length = clip_size - 1
+
+ for super_clip_idx, super_clip_size in enumerate(super_clips):
+ inner_index = list(range(super_clip_size))
+
+ if super_clip_size < clip_size: # Need padding for short sequences
+ post_module = inner_index[1:-1][::-1] + inner_index
+ l_post = len(post_module)
+ post_module = post_module * (pad_length // l_post + 1)
+ post_module = post_module[:pad_length]
+
+ if len(post_module) != pad_length:
+ continue # Skip sequences that are too short
+
+ pre_module = inner_index + inner_index[1:-1][::-1]
+ l_pre = len(post_module)
+ pre_module = pre_module * (pad_length // l_pre + 1)
+ pre_module = pre_module[-pad_length:]
+
+ if len(pre_module) != pad_length:
+ continue # Skip sequences that are too short
+
+ inner_index = pre_module + inner_index + post_module
+
+ super_clip_size = len(inner_index)
+
+ # Generate sliding window clips
+ frame_range = [
+ inner_index[i : i + clip_size]
+ for i in range(super_clip_size)
+ if i + clip_size <= super_clip_size
+ ]
+
+ for indices in frame_range:
+ clip = [(super_clip_idx, t) for t in indices]
+ clips_for_video.append(clip)
+
+ # Run inference on clips
+ preds = []
+ frame_res = {}
+ test_transform = Compose([
+ ToTensor(),
+ Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+ ])
+
+ print(f"Running inference on {len(clips_for_video)} clips...")
+ for clip in tqdm(clips_for_video, desc="Processing clips"):
+ # Prepare data for this clip
+ images = [data_storage[f"{i}_{j}_img"] for i, j in clip]
+ landmarks = [data_storage[f"{i}_{j}_ldm"] for i, j in clip]
+ frame_ids = [data_storage[f"{i}_{j}_idx"] for i, j in clip]
+
+ # Align and crop faces
+ landmarks, images = crop_align_func(landmarks, images)
+
+ # Convert to temporal frequency domain
+ images_tensor = []
+ ft_images = []
+ for image in images:
+ image = np.array(image)
+ img_pil = Image.fromarray(image)
+ img_pil = test_transform(img_pil)
+ images_tensor.append(img_pil)
+
+ # Apply median filter and compute frequency domain
+ img_filtered = cv2.medianBlur(image.copy(), 5)
+ ft_images.append(cv2.cvtColor((image - img_filtered), cv2.COLOR_RGB2GRAY))
+
+ # Compute FFT for temporal analysis
+ ft_images = np.array(ft_images)
+ ft_images = np.absolute(np.fft.fft(ft_images, axis=0)[:clip_size//2] * 1/clip_size)
+ ft_images = torch.from_numpy(ft_images).cuda()
+ ft_images = ft_images.unsqueeze(0)
+
+ # Prepare image tensor
+ images = torch.stack(images_tensor, dim=1)
+ images = images.unsqueeze(0)
+ images = images.cuda()
+
+ # Run model inference
+ with torch.no_grad():
+ output = model(images, ft_images)
+ output = F.sigmoid(output)
+ output = output.squeeze(0)
+
+ pred = float(output.item())
+
+ # Store predictions for each frame
+ for f_id in frame_ids:
+ if f_id not in frame_res:
+ frame_res[f_id] = []
+ frame_res[f_id].append(pred)
+ preds.append(pred)
+
+
+ # Aggregate results
+ mean_pred = np.mean(preds)
+ print(f"Average prediction score: {mean_pred:.4f}")
+
+ # Prepare final results
+ boxes = []
+ scores = []
+
+ for frame_idx in range(len(frames)):
+ if frame_idx in frame_res:
+ pred_prob = np.mean(frame_res[frame_idx])
+ rect = frame_boxes[frame_idx]
+ else:
+ pred_prob = None
+ rect = None
+ scores.append(pred_prob)
+ boxes.append(rect)
+
+ # Save results to video
+ print(f"Saving results to {out_file}")
+ SupplyWriter(args.video, out_file, 0.002584857167676091).run(frames, scores, boxes)
+ print("Inference completed successfully!")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/clean/video/pwtf_dvd/inference/test_tools/__init__.py b/clean/video/pwtf_dvd/inference/test_tools/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/clean/video/pwtf_dvd/inference/test_tools/common.py b/clean/video/pwtf_dvd/inference/test_tools/common.py
new file mode 100644
index 0000000000000000000000000000000000000000..e76ac698b8fe9a32dfbdf4e294d5afa521250e71
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_tools/common.py
@@ -0,0 +1,122 @@
+import os
+
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+
+from .ct.detection.utils import grab_all_frames, get_valid_faces, sample_chunks
+from .ct.operations import multiple_tracking
+import numpy as np
+from .ct.face_alignment import LandmarkPredictor
+from .ct.detection import FaceDetector
+import cv2
+from .utils import flatten,partition
+
+
+detector = FaceDetector(0)
+predictor = LandmarkPredictor(0)
+
+
+def get_five(ldm68):
+ groups = [range(36, 42), range(42, 48), [30], [48], [54]]
+ points = []
+ for group in groups:
+ points.append(ldm68[group].mean(0))
+ return np.array(points)
+
+
+def get_bbox(mask):
+ try:
+ y, x = np.nonzero(mask[..., 0])
+ return x.min() - 1, y.min() - 1, x.max() + 1, y.max() + 1
+ except:
+ return None
+
+
+def get_bigger_box(image, box, scale=0.5):
+ height, width = image.shape[:2]
+ box = np.rint(box).astype(np.int)
+ new_box = box.reshape(2, 2)
+ size = new_box[1] - new_box[0]
+ diff = scale * size
+ diff = diff[None, :] * np.array([-1, 1])[:, None]
+ new_box = new_box + diff
+ new_box[:, 0] = np.clip(new_box[:, 0], 0, width - 1)
+ new_box[:, 1] = np.clip(new_box[:, 1], 0, height - 1)
+ new_box = np.rint(new_box).astype(np.int)
+ return new_box.reshape(-1)
+
+
+def process_bigger_clips(clips, dete_res, clip_size, step, scale=0.5):
+ assert len(clips) % clip_size == 0
+ detect_results = sample_chunks(dete_res, clip_size, step)
+ clips = sample_chunks(clips, clip_size, step)
+ new_clips = []
+ for i, (frame_clip, record_clip) in enumerate(zip(clips, detect_results)):
+ tracks = multiple_tracking(record_clip)
+ for j, track in enumerate(tracks):
+ new_images = []
+ for (box, ldm, _), frame in zip(track, frame_clip):
+ big_box = get_bigger_box(frame, box, scale)
+ x1, y1, x2, y2 = big_box
+ top_left = big_box[:2][None, :]
+ new_ldm5 = ldm - top_left
+ box = np.rint(box).astype(np.int)
+ new_box = (box.reshape(2, 2) - top_left).reshape(-1)
+ feed = LandmarkPredictor.prepare_feed(frame, box)
+ ldm68 = predictor(feed) - top_left
+ new_images.append(
+ (frame[y1:y2, x1:x2], big_box, new_box, new_ldm5, ldm68)
+ )
+ new_clips.append(new_images)
+ return new_clips
+
+
+def post(detected_faces):
+ return [[face[:4], None, face[-1]] for face in detected_faces]
+
+
+def check(detect_res):
+ return min([len(faces) for faces in detect_res]) != 0
+
+
+def detect_all(file, sfd_only=False, return_frames=False, max_size=None):
+ frames = grab_all_frames(file, max_size=max_size, cvt=True)
+ if not sfd_only:
+ detect_res = flatten(
+ [detector.detect(item) for item in partition(frames, 50)]
+ )
+ detect_res = get_valid_faces(detect_res, thres=0.5)
+ else:
+ raise NotImplementedError
+
+ all_68 = get_lm68(frames, detect_res)
+ if not return_frames:
+ return detect_res, all_68
+ else:
+ return detect_res, all_68, frames
+
+
+def get_lm68(frames, detect_res):
+ assert len(frames) == len(detect_res)
+ frame_count = len(frames)
+ all_68 = []
+ for i in range(frame_count):
+ frame = frames[i]
+ faces = detect_res[i]
+ if len(faces) == 0:
+ res_68 = []
+ else:
+ feeds = []
+ for face in faces:
+ assert len(face) == 3
+ box = face[0]
+ feed = LandmarkPredictor.prepare_feed(frame, box)
+ feeds.append(feed)
+ res_68 = predictor(feeds)
+ assert len(res_68) == len(faces)
+ for face, l_68 in zip(faces, res_68):
+ if face[1] is None:
+ face[1] = get_five(l_68)
+ all_68.append(res_68)
+
+ assert len(all_68) == len(detect_res)
+ return all_68
diff --git a/clean/video/pwtf_dvd/inference/test_tools/ct/detection/__init__.py b/clean/video/pwtf_dvd/inference/test_tools/ct/detection/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..23447e84f4f5f23a4c4818df2f26cc19a9200363
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_tools/ct/detection/__init__.py
@@ -0,0 +1,56 @@
+import cv2
+from .detector import RetinaFace
+from .utils import *
+
+
+def assert_bounded(val, low, up):
+ return val >= low and val < up
+
+
+def check_valid(face, w, h):
+ box = face[0]
+ if box[0] > box[2]:
+ return False
+ if box[1] > box[3]:
+ return False
+ for idx, bound in zip([0, 1, 2, 3], [w, h, w, h]):
+ if not assert_bounded(box[idx], 0, bound):
+ return False
+ pts = face[1]
+ for p in pts:
+ for idx, bound in zip([0, 1], [w, h]):
+ if not assert_bounded(p[idx], 0, bound):
+ return False
+ return True
+
+
+def post_detect(detect_results, scale, w, h):
+ new_results = []
+ for frame_faces in detect_results:
+ new_frame_faces = []
+ for box, ldm, score in frame_faces:
+ box = box * scale
+ ldm = ldm * scale
+ face = (box, ldm, score)
+ if check_valid(face, w=w, h=h):
+ new_frame_faces.append(face)
+ new_results.append(new_frame_faces)
+ return new_results
+
+
+class FaceDetector(RetinaFace):
+ def scale_detect(self, images):
+ max_res = 1920
+ h, w = images[0].shape[:2]
+ if max(h, w) > max_res:
+ init_scale = max(h, w) / max_res
+ else:
+ init_scale = 1
+ resize_scale = 2 * init_scale
+ resize_w = int(w / resize_scale)
+ resize_h = int(h / resize_scale)
+ detect_input = [cv2.resize(frame, (resize_w, resize_h)) for frame in images]
+ detect_results = post_detect(
+ self.detect(detect_input), scale=resize_scale, w=w, h=h,
+ )
+ return detect_results
diff --git a/clean/video/pwtf_dvd/inference/test_tools/ct/detection/alignment.py b/clean/video/pwtf_dvd/inference/test_tools/ct/detection/alignment.py
new file mode 100644
index 0000000000000000000000000000000000000000..64692a3490c7e8e29e399bd174c5906716b3e6fc
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_tools/ct/detection/alignment.py
@@ -0,0 +1,608 @@
+from itertools import product as product
+from math import ceil
+
+import numpy as np
+import torch
+import torch.backends.cudnn as cudnn
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models._utils as _utils
+
+
+def conv_bn(inp, oup, stride=1, leaky=0):
+ return nn.Sequential(
+ nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
+ nn.BatchNorm2d(oup),
+ nn.LeakyReLU(negative_slope=leaky, inplace=True),
+ )
+
+
+def conv_bn_no_relu(inp, oup, stride):
+ return nn.Sequential(
+ nn.Conv2d(inp, oup, 3, stride, 1, bias=False), nn.BatchNorm2d(oup),
+ )
+
+
+def conv_bn1X1(inp, oup, stride, leaky=0):
+ return nn.Sequential(
+ nn.Conv2d(inp, oup, 1, stride, padding=0, bias=False),
+ nn.BatchNorm2d(oup),
+ nn.LeakyReLU(negative_slope=leaky, inplace=True),
+ )
+
+
+def conv_dw(inp, oup, stride, leaky=0.1):
+ return nn.Sequential(
+ nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
+ nn.BatchNorm2d(inp),
+ nn.LeakyReLU(negative_slope=leaky, inplace=True),
+ nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
+ nn.BatchNorm2d(oup),
+ nn.LeakyReLU(negative_slope=leaky, inplace=True),
+ )
+
+
+class SSH(nn.Module):
+ def __init__(self, in_channel, out_channel):
+ super(SSH, self).__init__()
+ assert out_channel % 4 == 0
+ leaky = 0
+ if out_channel <= 64:
+ leaky = 0.1
+ self.conv3X3 = conv_bn_no_relu(in_channel, out_channel // 2, stride=1)
+
+ self.conv5X5_1 = conv_bn(in_channel, out_channel // 4, stride=1, leaky=leaky)
+ self.conv5X5_2 = conv_bn_no_relu(out_channel // 4, out_channel // 4, stride=1)
+
+ self.conv7X7_2 = conv_bn(
+ out_channel // 4, out_channel // 4, stride=1, leaky=leaky
+ )
+ self.conv7x7_3 = conv_bn_no_relu(out_channel // 4, out_channel // 4, stride=1)
+
+ def forward(self, input):
+ conv3X3 = self.conv3X3(input)
+
+ conv5X5_1 = self.conv5X5_1(input)
+ conv5X5 = self.conv5X5_2(conv5X5_1)
+
+ conv7X7_2 = self.conv7X7_2(conv5X5_1)
+ conv7X7 = self.conv7x7_3(conv7X7_2)
+
+ out = torch.cat([conv3X3, conv5X5, conv7X7], dim=1)
+ out = F.relu(out)
+ return out
+
+
+class FPN(nn.Module):
+ def __init__(self, in_channels_list, out_channels):
+ super(FPN, self).__init__()
+ leaky = 0
+ if out_channels <= 64:
+ leaky = 0.1
+ self.output1 = conv_bn1X1(
+ in_channels_list[0], out_channels, stride=1, leaky=leaky
+ )
+ self.output2 = conv_bn1X1(
+ in_channels_list[1], out_channels, stride=1, leaky=leaky
+ )
+ self.output3 = conv_bn1X1(
+ in_channels_list[2], out_channels, stride=1, leaky=leaky
+ )
+
+ self.merge1 = conv_bn(out_channels, out_channels, leaky=leaky)
+ self.merge2 = conv_bn(out_channels, out_channels, leaky=leaky)
+
+ def forward(self, input):
+ # names = list(input.keys())
+ input = list(input.values())
+
+ output1 = self.output1(input[0])
+ output2 = self.output2(input[1])
+ output3 = self.output3(input[2])
+
+ up3 = F.interpolate(
+ output3, size=[output2.size(2), output2.size(3)], mode="nearest"
+ )
+ output2 = output2 + up3
+ output2 = self.merge2(output2)
+
+ up2 = F.interpolate(
+ output2, size=[output1.size(2), output1.size(3)], mode="nearest"
+ )
+ output1 = output1 + up2
+ output1 = self.merge1(output1)
+
+ out = [output1, output2, output3]
+ return out
+
+
+class MobileNetV1(nn.Module):
+ def __init__(self):
+ super(MobileNetV1, self).__init__()
+ self.stage1 = nn.Sequential(
+ conv_bn(3, 8, 2, leaky=0.1), # 3
+ conv_dw(8, 16, 1), # 7
+ conv_dw(16, 32, 2), # 11
+ conv_dw(32, 32, 1), # 19
+ conv_dw(32, 64, 2), # 27
+ conv_dw(64, 64, 1), # 43
+ )
+ self.stage2 = nn.Sequential(
+ conv_dw(64, 128, 2), # 43 + 16 = 59
+ conv_dw(128, 128, 1), # 59 + 32 = 91
+ conv_dw(128, 128, 1), # 91 + 32 = 123
+ conv_dw(128, 128, 1), # 123 + 32 = 155
+ conv_dw(128, 128, 1), # 155 + 32 = 187
+ conv_dw(128, 128, 1), # 187 + 32 = 219
+ )
+ self.stage3 = nn.Sequential(
+ conv_dw(128, 256, 2), # 219 +3 2 = 241
+ conv_dw(256, 256, 1), # 241 + 64 = 301
+ )
+ self.avg = nn.AdaptiveAvgPool2d((1, 1))
+ self.fc = nn.Linear(256, 1000)
+
+ def forward(self, x):
+ x = self.stage1(x)
+ x = self.stage2(x)
+ x = self.stage3(x)
+ x = self.avg(x)
+ # x = self.model(x)
+ x = x.view(-1, 256)
+ x = self.fc(x)
+ return x
+
+
+class ClassHead(nn.Module):
+ def __init__(self, inchannels=512, num_anchors=3):
+ super(ClassHead, self).__init__()
+ self.num_anchors = num_anchors
+ self.conv1x1 = nn.Conv2d(
+ inchannels, self.num_anchors * 2, kernel_size=(1, 1), stride=1, padding=0
+ )
+
+ def forward(self, x):
+ out = self.conv1x1(x)
+ out = out.permute(0, 2, 3, 1).contiguous()
+
+ return out.view(out.shape[0], -1, 2)
+
+
+class BboxHead(nn.Module):
+ def __init__(self, inchannels=512, num_anchors=3):
+ super(BboxHead, self).__init__()
+ self.conv1x1 = nn.Conv2d(
+ inchannels, num_anchors * 4, kernel_size=(1, 1), stride=1, padding=0
+ )
+
+ def forward(self, x):
+ out = self.conv1x1(x)
+ out = out.permute(0, 2, 3, 1).contiguous()
+
+ return out.view(out.shape[0], -1, 4)
+
+
+class LandmarkHead(nn.Module):
+ def __init__(self, inchannels=512, num_anchors=3):
+ super(LandmarkHead, self).__init__()
+ self.conv1x1 = nn.Conv2d(
+ inchannels, num_anchors * 10, kernel_size=(1, 1), stride=1, padding=0
+ )
+
+ def forward(self, x):
+ out = self.conv1x1(x)
+ out = out.permute(0, 2, 3, 1).contiguous()
+
+ return out.view(out.shape[0], -1, 10)
+
+
+class RetinaFace(nn.Module):
+ def __init__(self, cfg=None, phase="train"):
+ """
+ :param cfg: Network related settings.
+ :param phase: train or test.
+ """
+ super(RetinaFace, self).__init__()
+ self.phase = phase
+ backbone = None
+ if cfg["name"] == "mobilenet0.25":
+ backbone = MobileNetV1()
+ elif cfg["name"] == "Resnet50":
+ import torchvision.models as models
+
+ backbone = models.resnet50(pretrained=cfg["pretrain"])
+
+ self.body = _utils.IntermediateLayerGetter(backbone, cfg["return_layers"])
+ in_channels_stage2 = cfg["in_channel"]
+ in_channels_list = [
+ in_channels_stage2 * 2,
+ in_channels_stage2 * 4,
+ in_channels_stage2 * 8,
+ ]
+ out_channels = cfg["out_channel"]
+ self.fpn = FPN(in_channels_list, out_channels)
+ self.ssh1 = SSH(out_channels, out_channels)
+ self.ssh2 = SSH(out_channels, out_channels)
+ self.ssh3 = SSH(out_channels, out_channels)
+
+ self.ClassHead = self._make_class_head(fpn_num=3, inchannels=cfg["out_channel"])
+ self.BboxHead = self._make_bbox_head(fpn_num=3, inchannels=cfg["out_channel"])
+ self.LandmarkHead = self._make_landmark_head(
+ fpn_num=3, inchannels=cfg["out_channel"]
+ )
+
+ def _make_class_head(self, fpn_num=3, inchannels=64, anchor_num=2):
+ classhead = nn.ModuleList()
+ for i in range(fpn_num):
+ classhead.append(ClassHead(inchannels, anchor_num))
+ return classhead
+
+ def _make_bbox_head(self, fpn_num=3, inchannels=64, anchor_num=2):
+ bboxhead = nn.ModuleList()
+ for i in range(fpn_num):
+ bboxhead.append(BboxHead(inchannels, anchor_num))
+ return bboxhead
+
+ def _make_landmark_head(self, fpn_num=3, inchannels=64, anchor_num=2):
+ landmarkhead = nn.ModuleList()
+ for i in range(fpn_num):
+ landmarkhead.append(LandmarkHead(inchannels, anchor_num))
+ return landmarkhead
+
+ def forward(self, inputs):
+ out = self.body(inputs)
+
+ # FPN
+ fpn = self.fpn(out)
+
+ # SSH
+ feature1 = self.ssh1(fpn[0])
+ feature2 = self.ssh2(fpn[1])
+ feature3 = self.ssh3(fpn[2])
+ features = [feature1, feature2, feature3]
+
+ bbox_regressions = torch.cat(
+ [self.BboxHead[i](feature) for i, feature in enumerate(features)], dim=1
+ )
+ classifications = torch.cat(
+ [self.ClassHead[i](feature) for i, feature in enumerate(features)], dim=1
+ )
+ ldm_regressions = torch.cat(
+ [self.LandmarkHead[i](feature) for i, feature in enumerate(features)], dim=1
+ )
+
+ if self.phase == "train":
+ output = (bbox_regressions, classifications, ldm_regressions)
+ else:
+ output = (
+ bbox_regressions,
+ F.softmax(classifications, dim=-1),
+ ldm_regressions,
+ )
+ return output
+
+
+# Adapted from https://github.com/Hakuyume/chainer-ssd
+def decode(loc, priors, variances):
+ boxes = torch.cat(
+ (
+ priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
+ priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1]),
+ ),
+ 1,
+ )
+ boxes[:, :2] -= boxes[:, 2:] / 2
+ boxes[:, 2:] += boxes[:, :2]
+ return boxes
+
+
+def decode_landm(pre, priors, variances):
+ landms = torch.cat(
+ (
+ priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],
+ priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:],
+ priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:],
+ priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:],
+ priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:],
+ ),
+ dim=1,
+ )
+ return landms
+
+
+def py_cpu_nms(dets, thresh):
+ """Pure Python NMS baseline."""
+ x1 = dets[:, 0]
+ y1 = dets[:, 1]
+ x2 = dets[:, 2]
+ y2 = dets[:, 3]
+ scores = dets[:, 4]
+
+ areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+ order = scores.argsort()[::-1]
+
+ keep = []
+ while order.size > 0:
+ i = order[0]
+ keep.append(i)
+ xx1 = np.maximum(x1[i], x1[order[1:]])
+ yy1 = np.maximum(y1[i], y1[order[1:]])
+ xx2 = np.minimum(x2[i], x2[order[1:]])
+ yy2 = np.minimum(y2[i], y2[order[1:]])
+
+ w = np.maximum(0.0, xx2 - xx1 + 1)
+ h = np.maximum(0.0, yy2 - yy1 + 1)
+ inter = w * h
+ ovr = inter / (areas[i] + areas[order[1:]] - inter)
+
+ inds = np.where(ovr <= thresh)[0]
+ order = order[inds + 1]
+
+ return keep
+
+
+class PriorBox(object):
+ def __init__(self, cfg, image_size=None, phase="train"):
+ super(PriorBox, self).__init__()
+ self.min_sizes = cfg["min_sizes"]
+ self.steps = cfg["steps"]
+ self.clip = cfg["clip"]
+ self.image_size = image_size
+ self.feature_maps = [
+ [ceil(self.image_size[0] / step), ceil(self.image_size[1] / step)]
+ for step in self.steps
+ ]
+ self.name = "s"
+
+ def forward(self):
+ anchors = []
+ for k, f in enumerate(self.feature_maps):
+ min_sizes = self.min_sizes[k]
+ for i, j in product(range(f[0]), range(f[1])):
+ for min_size in min_sizes:
+ s_kx = min_size / self.image_size[1]
+ s_ky = min_size / self.image_size[0]
+ dense_cx = [
+ x * self.steps[k] / self.image_size[1] for x in [j + 0.5]
+ ]
+ dense_cy = [
+ y * self.steps[k] / self.image_size[0] for y in [i + 0.5]
+ ]
+ for cy, cx in product(dense_cy, dense_cx):
+ anchors += [cx, cy, s_kx, s_ky]
+
+ # back to torch land
+ output = torch.Tensor(anchors).view(-1, 4)
+ if self.clip:
+ output.clamp_(max=1, min=0)
+ return output
+
+
+cfg_mnet = {
+ "name": "mobilenet0.25",
+ "min_sizes": [[16, 32], [64, 128], [256, 512]],
+ "steps": [8, 16, 32],
+ "variance": [0.1, 0.2],
+ "clip": False,
+ "loc_weight": 2.0,
+ "gpu_train": True,
+ "batch_size": 32,
+ "ngpu": 1,
+ "epoch": 250,
+ "decay1": 190,
+ "decay2": 220,
+ "image_size": 640,
+ "pretrain": True,
+ "return_layers": {"stage1": 1, "stage2": 2, "stage3": 3},
+ "in_channel": 32,
+ "out_channel": 64,
+}
+
+cfg_re50 = {
+ "name": "Resnet50",
+ "min_sizes": [[16, 32], [64, 128], [256, 512]],
+ "steps": [8, 16, 32],
+ "variance": [0.1, 0.2],
+ "clip": False,
+ "loc_weight": 2.0,
+ "gpu_train": True,
+ "batch_size": 24,
+ "ngpu": 4,
+ "epoch": 100,
+ "decay1": 70,
+ "decay2": 90,
+ "image_size": 840,
+ "pretrain": False,
+ "return_layers": {"layer2": 1, "layer3": 2, "layer4": 3},
+ "in_channel": 256,
+ "out_channel": 256,
+}
+
+
+def check_keys(model, pretrained_state_dict):
+ ckpt_keys = set(pretrained_state_dict.keys())
+ model_keys = set(model.state_dict().keys())
+ used_pretrained_keys = model_keys & ckpt_keys
+ assert len(used_pretrained_keys) > 0, "load NONE from pretrained checkpoint"
+ return True
+
+
+def remove_prefix(state_dict, prefix):
+ """ Old style model is stored with all names of parameters sharing common prefix 'module.' """
+ f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else x
+ return {f(key): value for key, value in state_dict.items()}
+
+
+def load_model(model, pretrained_path, load_to_cpu):
+ if load_to_cpu:
+ if pretrained_path is None:
+ url = "https://github.com/yinglinzheng/face_weights/releases/download/v1/mobilenet0.25_Final.pth"
+ pretrained_dict = torch.utils.model_zoo.load_url(url)
+ else:
+ pretrained_dict = torch.load(
+ pretrained_path, map_location=lambda storage, loc: storage
+ )
+ else:
+ device = torch.cuda.current_device()
+ pretrained_dict = torch.load(
+ pretrained_path, map_location=lambda storage, loc: storage.cuda(device)
+ )
+ if "state_dict" in pretrained_dict.keys():
+ pretrained_dict = remove_prefix(pretrained_dict["state_dict"], "module.")
+ else:
+ pretrained_dict = remove_prefix(pretrained_dict, "module.")
+ check_keys(model, pretrained_dict)
+ model.load_state_dict(pretrained_dict, strict=False)
+ return model
+
+
+def load_net(model_path, device, network="mobilenet"):
+ if network == "mobilenet":
+ cfg = cfg_mnet
+ elif network == "resnet50":
+ cfg = cfg_re50
+ # net and model
+ net = RetinaFace(cfg=cfg, phase="test")
+ net = load_model(net, model_path, True)
+ net.eval()
+ cudnn.benchmark = True
+ net = net.to(device)
+ return net
+
+
+def parse_det(det):
+ landmarks = det[5:].reshape(5, 2)
+ box = det[:4]
+ score = det[4]
+ return box, landmarks, score
+
+
+def post_process(
+ loc,
+ conf,
+ landms,
+ prior_data,
+ cfg,
+ scale,
+ scale1,
+ resize,
+ confidence_threshold,
+ top_k,
+ nms_threshold,
+ keep_top_k,
+):
+ boxes = decode(loc, prior_data, cfg["variance"])
+ boxes = boxes * scale / resize
+ boxes = boxes.cpu().numpy()
+ scores = conf.cpu().numpy()[:, 1]
+ landms_copy = decode_landm(landms, prior_data, cfg["variance"])
+
+ landms_copy = landms_copy * scale1 / resize
+ landms_copy = landms_copy.cpu().numpy()
+
+ # ignore low scores
+ inds = np.where(scores > confidence_threshold)[0]
+ boxes = boxes[inds]
+ landms_copy = landms_copy[inds]
+ scores = scores[inds]
+
+ # keep top-K before NMS
+ order = scores.argsort()[::-1][:top_k]
+ boxes = boxes[order]
+ landms_copy = landms_copy[order]
+ scores = scores[order]
+
+ # do NMS
+ dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
+ keep = py_cpu_nms(dets, nms_threshold)
+ # keep = nms(dets, args.nms_threshold,force_cpu=args.cpu)
+ dets = dets[keep, :]
+ landms_copy = landms_copy[keep]
+
+ # keep top-K faster NMS
+ dets = dets[:keep_top_k, :]
+ landms_copy = landms_copy[:keep_top_k, :]
+
+ dets = np.concatenate((dets, landms_copy), axis=1)
+ # show image
+ dets = sorted(dets, key=lambda x: x[4], reverse=True)
+ dets = [parse_det(x) for x in dets]
+
+ return dets
+
+
+def batch_detect(net, images, device, is_tensor=False, normalized=False):
+ with torch.no_grad():
+ confidence_threshold = 0.02
+ cfg = cfg_mnet
+ top_k = 5000
+ nms_threshold = 0.4
+ keep_top_k = 750
+ resize = 1
+ if not is_tensor:
+ try:
+ img = np.float32(images)
+ except ValueError:
+ raise NotImplementedError("Input images must of same size")
+ img = torch.from_numpy(img)
+ else:
+ img = images.float()
+ img = img.to(device)
+ mean = (
+ torch.as_tensor([104, 117, 123], dtype=img.dtype, device=img.device)
+ .unsqueeze(0)
+ .unsqueeze(0)
+ .unsqueeze(0)
+ )
+ img -= mean
+ img = img.permute(0, 3, 1, 2)
+ (batch_size, _, im_height, im_width,) = img.shape
+ scale = torch.as_tensor(
+ [im_width, im_height, im_width, im_height],
+ dtype=img.dtype,
+ device=img.device,
+ )
+ scale = scale.to(device)
+
+ loc, conf, landms = net(img) # forward pass
+
+ priorbox = PriorBox(cfg, image_size=(im_height, im_width))
+ priors = priorbox.forward()
+ prior_data = priors.to(device)
+ scale1 = torch.as_tensor(
+ [
+ img.shape[3],
+ img.shape[2],
+ img.shape[3],
+ img.shape[2],
+ img.shape[3],
+ img.shape[2],
+ img.shape[3],
+ img.shape[2],
+ img.shape[3],
+ img.shape[2],
+ ],
+ dtype=img.dtype,
+ device=img.device,
+ )
+ scale1 = scale1.to(device)
+
+ all_dets = [
+ post_process(
+ loc_i,
+ conf_i,
+ landms_i,
+ prior_data,
+ cfg,
+ scale,
+ scale1,
+ resize,
+ confidence_threshold,
+ top_k,
+ nms_threshold,
+ keep_top_k,
+ )
+ for loc_i, conf_i, landms_i in zip(loc, conf, landms)
+ ]
+
+ return all_dets
diff --git a/clean/video/pwtf_dvd/inference/test_tools/ct/detection/detector.py b/clean/video/pwtf_dvd/inference/test_tools/ct/detection/detector.py
new file mode 100644
index 0000000000000000000000000000000000000000..f38050e7a050eeb05323e7d710a63c2bf1a12455
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_tools/ct/detection/detector.py
@@ -0,0 +1,46 @@
+import os
+
+import numpy as np
+import torch
+
+from .alignment import load_net, batch_detect
+
+
+def get_project_dir():
+ current_path = os.path.abspath(os.path.join(__file__, "../"))
+ return current_path
+
+
+def relative(path):
+ path = os.path.join(get_project_dir(), path)
+ return os.path.abspath(path)
+
+
+class RetinaFace:
+ def __init__(
+ self, gpu_id=-1, model_path=None, network="mobilenet",
+ ):
+ self.gpu_id = gpu_id
+ self.device = (
+ torch.device("cpu") if gpu_id == -1 else torch.device("cuda", gpu_id)
+ )
+ self.model = load_net(model_path, self.device, network)
+
+ def detect(self, images):
+ if isinstance(images, np.ndarray):
+ if len(images.shape) == 3:
+ return batch_detect(self.model, [images], self.device)[0]
+ elif len(images.shape) == 4:
+ return batch_detect(self.model, images, self.device)
+ elif isinstance(images, list):
+ return batch_detect(self.model, np.array(images), self.device)
+ elif isinstance(images, torch.Tensor):
+ if len(images.shape) == 3:
+ return batch_detect(self.model, images.unsqueeze(0), self.device)[0]
+ elif len(images.shape) == 4:
+ return batch_detect(self.model, images, self.device)
+ else:
+ raise NotImplementedError()
+
+ def __call__(self, images):
+ return self.detect(images)
diff --git a/clean/video/pwtf_dvd/inference/test_tools/ct/detection/utils.py b/clean/video/pwtf_dvd/inference/test_tools/ct/detection/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..512d824f02645ae9473331a553b18850e642b08f
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_tools/ct/detection/utils.py
@@ -0,0 +1,146 @@
+import cv2
+from test_tools.utils import flatten
+import numpy as np
+
+
+def chunks(l, n, step=None):
+ if step is None:
+ step = n
+ return [l[i : i + n] for i in range(0, len(l), step)]
+
+
+def sample_chunks(l, n, step=None):
+ return [l[i : i + n] for i in range(0, len(l), step) if i + n <= len(l)]
+
+
+def grab_all_frames(path, max_size, cvt=False):
+ capture = cv2.VideoCapture(path)
+ ret = True
+ frames = []
+ while ret:
+ ret, frame = capture.read()
+ if ret:
+ if cvt:
+ frame = frame[..., ::-1]
+ frames.append(frame)
+ if len(frames) == max_size:
+ break
+ capture.release()
+ return frames
+
+
+def get_clips_uniform(path, count, clip_size):
+ capture = cv2.VideoCapture(path)
+ n_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
+ max_clip_available = n_frames + 1 - clip_size
+ if count > max_clip_available:
+ count = max_clip_available
+ final_start = max_clip_available - 1
+ start_indices = np.linspace(0, final_start, count, endpoint=True, dtype=np.int)
+ all_clip_idx = [list(range(start, start + clip_size)) for start in start_indices]
+ valid = set(flatten(all_clip_idx))
+ max_idx = max(valid)
+
+ frames = {}
+ for idx in range(max_idx + 1):
+ # Get the next frame, but don't decode if we're not using it.
+ ret = capture.grab()
+ if not ret:
+ continue
+
+ if idx in valid:
+ ret, frame = capture.retrieve()
+ if not ret or frame is None:
+ continue
+ else:
+ # frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+ frames[idx] = frame
+
+ capture.release()
+ clips = []
+ for clip_idx in all_clip_idx:
+ clip = []
+ flag = True
+ for idx in clip_idx:
+ if idx not in frames:
+ flag = False
+ break
+ clip.append(frames[idx])
+ if flag:
+ clips.append(clip)
+ return clips
+
+
+def get_valid_faces(detect_results, max_count=10, thres=0.5, at_least=False):
+ new_results = []
+ for i, faces in enumerate(detect_results):
+ if len(faces) > max_count:
+ faces = faces[:max_count]
+ l = []
+ for j, face in enumerate(faces):
+ if face[-1] < thres and not (j == 0 and at_least):
+ continue
+ box, lm, score = face
+ box = box.astype(np.float32)
+ lm = lm.astype(np.float32)
+ l.append((box, lm, score))
+ new_results.append(l)
+ return new_results
+
+
+def scale_box(box, scale_h, scale_w, h, w):
+ x1, y1, x2, y2 = box.astype(np.int32)
+ center_x = (x1 + x2) // 2
+ center_y = (y1 + y2) // 2
+ box_h = int((y2 - y1) * scale_h)
+ box_w = int((x2 - x1) * scale_w)
+ new_x1 = center_x - box_w // 2
+ new_x2 = new_x1 + box_w
+ new_y1 = center_y - box_h // 2
+ new_y2 = new_y1 + box_h
+ new_x1 = max(new_x1, 0)
+ new_y1 = max(new_y1, 0)
+ new_y2 = min(new_y2, h)
+ new_x2 = min(new_x2, w)
+ return new_x1, new_y1, new_x2, new_y2
+
+
+def get_bbox(detect_res):
+ tmp_detect_res = get_valid_faces(detect_res, max_count=4, thres=0.5)
+ all_face_bboxs = []
+ for faces in tmp_detect_res:
+ all_face_bboxs.extend([face[0] for face in faces])
+ all_face_bboxs = np.array(all_face_bboxs).astype(np.int)
+ x1 = all_face_bboxs[:, 0].min()
+ x2 = all_face_bboxs[:, 2].max()
+ y1 = all_face_bboxs[:, 1].min()
+ y2 = all_face_bboxs[:, 3].max()
+
+ return x1, y1, x2, y2
+
+
+def delta_detect_res(detect_res, x1, y1):
+ diff = np.array([[x1, y1]])
+ new_detect_res = []
+ for faces in detect_res:
+ f = []
+ for face in faces:
+ box, lm, score = face
+ box = box.astype(np.float)
+ box[[0, 2]] -= x1
+ box[[1, 3]] -= y1
+ lm = lm.astype(np.float) - diff
+ f.append((box, lm, score))
+ new_detect_res.append(f)
+ return new_detect_res
+
+
+def pre_crop(clips, detect_res):
+ box = np.array(get_bbox(detect_res))
+ w = box[2] - box[0]
+ h = box[3] - box[1]
+ x1, y1, x2, y2 = scale_box(
+ box, 1.5, 1.2 if w > 2 * h else 1.5, clips[0].shape[0], clips[0].shape[1]
+ )
+ clips = np.array(clips)
+ return clips[:, y1:y2, x1:x2], delta_detect_res(detect_res, x1, y1)
diff --git a/clean/video/pwtf_dvd/inference/test_tools/ct/face_alignment/__init__.py b/clean/video/pwtf_dvd/inference/test_tools/ct/face_alignment/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..b0545f2f8c2aa872a127bf4e815c1e614a971723
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_tools/ct/face_alignment/__init__.py
@@ -0,0 +1 @@
+from .predictor import LandmarkPredictor
\ No newline at end of file
diff --git a/clean/video/pwtf_dvd/inference/test_tools/ct/face_alignment/basenet.py b/clean/video/pwtf_dvd/inference/test_tools/ct/face_alignment/basenet.py
new file mode 100644
index 0000000000000000000000000000000000000000..699b163e0584d56cc6d8f916b171f7ba86326efa
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_tools/ct/face_alignment/basenet.py
@@ -0,0 +1,107 @@
+# Backbone networks used for face landmark detection
+# Cunjian Chen (cunjian@msu.edu)
+
+import torch.nn as nn
+import torchvision.models as models
+
+
+class ConvBlock(nn.Module):
+ def __init__(self, inp, oup, k, s, p, dw=False, linear=False):
+ super(ConvBlock, self).__init__()
+ self.linear = linear
+ if dw:
+ self.conv = nn.Conv2d(inp, oup, k, s, p, groups=inp, bias=False)
+ else:
+ self.conv = nn.Conv2d(inp, oup, k, s, p, bias=False)
+ self.bn = nn.BatchNorm2d(oup)
+ if not linear:
+ self.prelu = nn.PReLU(oup)
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.bn(x)
+ if self.linear:
+ return x
+ else:
+ return self.prelu(x)
+
+
+# SE module
+# https://github.com/wujiyang/Face_Pytorch/blob/master/backbone/cbam.py
+class SEModule(nn.Module):
+ """Squeeze and Excitation Module"""
+
+ def __init__(self, channels, reduction):
+ super(SEModule, self).__init__()
+ self.avg_pool = nn.AdaptiveAvgPool2d(1)
+ self.fc1 = nn.Conv2d(
+ channels, channels // reduction, kernel_size=1, padding=0, bias=False
+ )
+ self.relu = nn.ReLU(inplace=True)
+ self.fc2 = nn.Conv2d(
+ channels // reduction, channels, kernel_size=1, padding=0, bias=False
+ )
+ self.sigmoid = nn.Sigmoid()
+
+ def forward(self, x):
+ input = x
+ x = self.avg_pool(x)
+ x = self.fc1(x)
+ x = self.relu(x)
+ x = self.fc2(x)
+ x = self.sigmoid(x)
+
+ return input * x
+
+
+# USE global depthwise convolution layer. Compatible with MobileNetV2 (224×224), MobileNetV2_ExternalData (224×224)
+class MobileNet_GDConv(nn.Module):
+ def __init__(self, num_classes):
+ super(MobileNet_GDConv, self).__init__()
+ self.pretrain_net = models.mobilenet_v2(pretrained=False)
+ self.base_net = nn.Sequential(*list(self.pretrain_net.children())[:-1])
+ self.linear7 = ConvBlock(1280, 1280, (7, 7), 1, 0, dw=True, linear=True)
+ self.linear1 = ConvBlock(1280, num_classes, 1, 1, 0, linear=True)
+
+ def forward(self, x):
+ x = self.base_net(x)
+ x = self.linear7(x)
+ x = self.linear1(x)
+ x = x.view(x.size(0), -1)
+ return x
+
+
+# USE global depthwise convolution layer. Compatible with MobileNetV2 (56×56)
+class MobileNet_GDConv_56(nn.Module):
+ def __init__(self, num_classes):
+ super(MobileNet_GDConv_56, self).__init__()
+ self.pretrain_net = models.mobilenet_v2(pretrained=False)
+ self.base_net = nn.Sequential(*list(self.pretrain_net.children())[:-1])
+ self.linear7 = ConvBlock(1280, 1280, (2, 2), 1, 0, dw=True, linear=True)
+ self.linear1 = ConvBlock(1280, num_classes, 1, 1, 0, linear=True)
+
+ def forward(self, x):
+ x = self.base_net(x)
+ x = self.linear7(x)
+ x = self.linear1(x)
+ x = x.view(x.size(0), -1)
+ return x
+
+
+# MobileNetV2 with SE; Compatible with MobileNetV2_SE (224×224) and MobileNetV2_SE_RE (224×224)
+class MobileNet_GDConv_SE(nn.Module):
+ def __init__(self, num_classes):
+ super(MobileNet_GDConv_SE, self).__init__()
+ self.pretrain_net = models.mobilenet_v2(pretrained=True)
+ self.base_net = nn.Sequential(*list(self.pretrain_net.children())[:-1])
+ self.linear7 = ConvBlock(1280, 1280, (7, 7), 1, 0, dw=True, linear=True)
+ self.linear1 = ConvBlock(1280, num_classes, 1, 1, 0, linear=True)
+ self.attention = SEModule(1280, 8)
+
+ def forward(self, x):
+ x = self.base_net(x)
+ x = self.attention(x)
+ x = self.linear7(x)
+ x = self.linear1(x)
+ x = x.view(x.size(0), -1)
+ return x
diff --git a/clean/video/pwtf_dvd/inference/test_tools/ct/face_alignment/predictor.py b/clean/video/pwtf_dvd/inference/test_tools/ct/face_alignment/predictor.py
new file mode 100644
index 0000000000000000000000000000000000000000..b210b673ab3af90d805e7fff9b0b66b00826fd1f
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_tools/ct/face_alignment/predictor.py
@@ -0,0 +1,143 @@
+# Face alignment demo
+# Uses MTCNN as face detector
+# Cunjian Chen (ccunjian@gmail.com)
+import torch
+import cv2
+import numpy as np
+from torch.utils.data import DataLoader
+from .basenet import MobileNet_GDConv
+
+
+def get_device(gpu_id):
+ if gpu_id > -1:
+ return torch.device(f"cuda:{str(gpu_id)}")
+ else:
+ return torch.device("cpu")
+
+
+def load_model(file):
+ model = MobileNet_GDConv(136)
+ if file is not None:
+ model.load_state_dict(torch.load(file, map_location="cpu"))
+ else:
+ url = "https://github.com/yinglinzheng/face_weights/releases/download/v1/mobilenet_224_model_best_gdconv_external.pth"
+ model.load_state_dict(torch.utils.model_zoo.load_url(url))
+ return model
+
+
+# landmark of (5L, 2L) from [0,1] to real range
+def reproject(bbox, landmark):
+ landmark_ = landmark.clone()
+ x1, y1, x2, y2 = bbox
+ w = x2 - x1
+ h = y2 - y1
+ landmark_[:, 0] *= w
+ landmark_[:, 0] += x1
+ landmark_[:, 1] *= h
+ landmark_[:, 1] += y1
+ return landmark_
+
+
+def prepare_feed(img, face):
+ height, width, _ = img.shape
+ mean = np.asarray([0.485, 0.456, 0.406])
+ std = np.asarray([0.229, 0.224, 0.225])
+ out_size = 224
+ x1, y1, x2, y2 = face[:4]
+
+ w = x2 - x1 + 1
+ h = y2 - y1 + 1
+ size = int(min([w, h]) * 1.2)
+ cx = x1 + w // 2
+ cy = y1 + h // 2
+ x1 = cx - size // 2
+ x2 = x1 + size
+ y1 = cy - size // 2
+ y2 = y1 + size
+
+ dx = max(0, -x1)
+ dy = max(0, -y1)
+ x1 = max(0, x1)
+ y1 = max(0, y1)
+
+ edx = max(0, x2 - width)
+ edy = max(0, y2 - height)
+ x2 = min(width, x2)
+ y2 = min(height, y2)
+ new_bbox = torch.Tensor([x1, y1, x2, y2]).int()
+ x1, y1, x2, y2 = new_bbox
+ cropped = img[y1:y2, x1:x2]
+ if dx > 0 or dy > 0 or edx > 0 or edy > 0:
+ cropped = cv2.copyMakeBorder(
+ cropped, int(dy), int(edy), int(dx), int(edx), cv2.BORDER_CONSTANT, 0
+ )
+ cropped_face = cv2.resize(cropped, (out_size, out_size))
+
+ if cropped_face.shape[0] <= 0 or cropped_face.shape[1] <= 0:
+ return None
+ test_face = cropped_face.copy()
+ test_face = test_face / 255.0
+ test_face = (test_face - mean) / std
+ test_face = test_face.transpose((2, 0, 1))
+ data = torch.from_numpy(test_face).float()
+ return dict(data=data, bbox=new_bbox)
+
+
+@torch.no_grad()
+def single_predict(model, feed, device):
+ landmark = model(feed["data"].unsqueeze(0).to(device)).cpu()
+ landmark = landmark.reshape(-1, 2)
+ landmark = reproject(feed["bbox"], landmark)
+ return landmark.numpy()
+
+
+@torch.no_grad()
+def batch_predict(model, feeds, device):
+ if not isinstance(feeds, list):
+ feeds = [feeds]
+ # loader = DataLoader(FeedDataset(feeds), batch_size=50, shuffle=False)
+ data = []
+ for feed in feeds:
+ data.append(feed["data"].unsqueeze(0))
+ data = torch.cat(data, 0).to(device)
+ results = []
+
+ landmarks = model(data).cpu()
+ for landmark, feed in zip(landmarks, feeds):
+ landmark = landmark.reshape(-1, 2)
+ landmark = reproject(feed["bbox"], landmark)
+ results.append(landmark.numpy())
+ return results
+
+
+@torch.no_grad()
+def batch_predict2(model, feeds, device, batch_size=None):
+ if not isinstance(feeds, list):
+ feeds = [feeds]
+ if batch_size is None:
+ batch_size = len(feeds)
+ loader = DataLoader(feeds, batch_size=len(feeds), shuffle=False)
+ results = []
+ for feed in loader:
+ landmarks = model(feed["data"].to(device)).cpu()
+ for landmark, bbox in zip(landmarks, feed["bbox"]):
+ landmark = landmark.reshape(-1, 2)
+ landmark = reproject(bbox, landmark)
+ results.append(landmark.numpy())
+ return results
+
+
+class LandmarkPredictor:
+ def __init__(self, gpu_id=0, file=None):
+ self.device = get_device(gpu_id)
+ self.model = load_model(file).to(self.device).eval()
+
+ def __call__(self, feeds):
+ results = batch_predict2(self.model, feeds, self.device)
+ if not isinstance(feeds, list):
+ results = results[0]
+ return results
+
+ @staticmethod
+ def prepare_feed(img, face):
+ return prepare_feed(img, face)
diff --git a/clean/video/pwtf_dvd/inference/test_tools/ct/face_alignment/utils.py b/clean/video/pwtf_dvd/inference/test_tools/ct/face_alignment/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..e871b155963b0d1ce6891589d32615c5fc97bb66
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_tools/ct/face_alignment/utils.py
@@ -0,0 +1,17 @@
+import cv2
+
+
+def drawLandmark_multiple(img, bbox, landmark):
+ """
+ Input:
+ - img: gray or RGB
+ - bbox: type of BBox
+ - landmark: reproject landmark of (5L, 2L)
+ Output:
+ - img marked with landmark and bbox
+ """
+ x1, y1, x2, y2 = bbox
+ cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2)
+ for x, y in landmark:
+ cv2.circle(img, (int(x), int(y)), 2, (0, 255, 0), -1)
+ return img
diff --git a/clean/video/pwtf_dvd/inference/test_tools/ct/operations.py b/clean/video/pwtf_dvd/inference/test_tools/ct/operations.py
new file mode 100644
index 0000000000000000000000000000000000000000..68375fcc4efd768abcfddf39521af21d574c4aac
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_tools/ct/operations.py
@@ -0,0 +1,79 @@
+import os
+
+import os
+import cv2
+import numpy as np
+from .tracking.sort import iou
+
+
+def face_iou(f1, f2):
+ return iou(f1[0], f2[0])
+
+
+def simple_tracking(batch_landmarks, index=0, thres=0.5):
+ track = []
+
+ for i, faces in enumerate(batch_landmarks):
+ if i == 0:
+ if len(faces) <= index or faces[index][-1] < 0.8:
+ return None
+ if index != 0:
+ for idx in range(index):
+ if face_iou(faces[idx], faces[index]) > thres:
+ return None
+ track.append(faces[index])
+ else:
+ last = track[i - 1]
+ if len(faces) == 0:
+ return None
+ sorted_faces = sorted(faces, key=lambda x: face_iou(x, last), reverse=True)
+ if face_iou(sorted_faces[0], last) < thres:
+ return None
+ track.append(sorted_faces[0])
+ return track
+
+
+def multiple_tracking(batch_landmarks):
+ tracks = []
+ for i in range(len(batch_landmarks[0])):
+ track = simple_tracking(batch_landmarks, index=i)
+ if track is None:
+ continue
+ tracks.append(track)
+ return tracks
+
+def find_longest(detect_res):
+ fc = len(detect_res)
+ tuples = []
+ start = 0
+ end = 0
+ previous_count = -1
+ all_tracks = []
+ # start 取得到，end 取不到
+ while start < (fc - 1):
+ for end in range(start + 2, fc + 1):
+ tracks = multiple_tracking(detect_res[start:end])
+ if (len(tracks) != previous_count and previous_count != -1) or len(
+ tracks
+ ) == 0:
+ break
+ previous_count = len(tracks)
+ if end - start > 2:
+ if end != fc:
+ un_reach_end = end - 1
+ else:
+ un_reach_end = end
+ sub_tracks = multiple_tracking(detect_res[start:un_reach_end])
+ if end == fc and len(sub_tracks) == 0:
+ un_reach_end = end - 1
+ sub_tracks = multiple_tracking(detect_res[start:un_reach_end])
+ if len(sub_tracks) > 0:
+ tpl = (start, un_reach_end)
+ tuples.append(tpl)
+ all_tracks.append(sub_tracks[0])
+ else:
+ raise NotImplementedError
+ previous_count = -1
+ end = un_reach_end
+ start = end
+ return tuples, all_tracks
\ No newline at end of file
diff --git a/clean/video/pwtf_dvd/inference/test_tools/ct/tracking/__init__.py b/clean/video/pwtf_dvd/inference/test_tools/ct/tracking/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/clean/video/pwtf_dvd/inference/test_tools/ct/tracking/sort.py b/clean/video/pwtf_dvd/inference/test_tools/ct/tracking/sort.py
new file mode 100644
index 0000000000000000000000000000000000000000..dc7b0838e7110a2d3521c6d4cfefbb40cc266c23
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_tools/ct/tracking/sort.py
@@ -0,0 +1,285 @@
+"""
+ SORT: A Simple, Online and Realtime Tracker
+ Copyright (C) 2016 Alex Bewley alex@dynamicdetection.com
+
+ This program is free software: you can redistribute it and/or modify
+ it under the terms of the GNU General Public License as published by
+ the Free Software Foundation, either version 3 of the License, or
+ (at your option) any later version.
+
+ This program is distributed in the hope that it will be useful,
+ but WITHOUT ANY WARRANTY; without even the implied warranty of
+ MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ GNU General Public License for more details.
+
+ You should have received a copy of the GNU General Public License
+ along with this program. If not, see .
+"""
+from __future__ import print_function
+import os.path
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from scipy.optimize import linear_sum_assignment
+import glob
+import time
+import argparse
+from filterpy.kalman import KalmanFilter
+
+
+def iou(bb_test, bb_gt):
+ """
+ Computes IUO between two bboxes in the form [x1,y1,x2,y2]
+ """
+ xx1 = np.maximum(bb_test[0], bb_gt[0])
+ yy1 = np.maximum(bb_test[1], bb_gt[1])
+ xx2 = np.minimum(bb_test[2], bb_gt[2])
+ yy2 = np.minimum(bb_test[3], bb_gt[3])
+ w = np.maximum(0.0, xx2 - xx1)
+ h = np.maximum(0.0, yy2 - yy1)
+ wh = w * h
+ o = wh / (
+ (bb_test[2] - bb_test[0]) * (bb_test[3] - bb_test[1])
+ + (bb_gt[2] - bb_gt[0]) * (bb_gt[3] - bb_gt[1])
+ - wh
+ )
+ return o
+
+
+def convert_bbox_to_z(bbox):
+ """
+ Takes a bounding box in the form [x1,y1,x2,y2] and returns z in the form
+ [x,y,s,r] where x,y is the centre of the box and s is the scale/area and r is
+ the aspect ratio
+ """
+ w = bbox[2] - bbox[0]
+ h = bbox[3] - bbox[1]
+ x = bbox[0] + w / 2.0
+ y = bbox[1] + h / 2.0
+ s = w * h # scale is just area
+ r = w / float(h)
+ return np.array([x, y, s, r]).reshape((4, 1))
+
+
+def convert_x_to_bbox(x, score=None):
+ """
+ Takes a bounding box in the centre form [x,y,s,r] and returns it in the form
+ [x1,y1,x2,y2] where x1,y1 is the top left and x2,y2 is the bottom right
+ """
+ w = np.sqrt(x[2] * x[3])
+ h = x[2] / w
+ if score == None:
+ return np.array(
+ [x[0] - w / 2.0, x[1] - h / 2.0, x[0] + w / 2.0, x[1] + h / 2.0]
+ ).reshape((1, 4))
+ else:
+ return np.array(
+ [x[0] - w / 2.0, x[1] - h / 2.0, x[0] + w / 2.0, x[1] + h / 2.0, score]
+ ).reshape((1, 5))
+
+
+class KalmanBoxTracker(object):
+ """
+ This class represents the internel state of individual tracked objects observed as bbox.
+ """
+
+ count = 0
+
+ def __init__(self, bbox):
+ """
+ Initialises a tracker using initial bounding box.
+ """
+ # define constant velocity model
+ self.kf = KalmanFilter(dim_x=7, dim_z=4)
+ self.kf.F = np.array(
+ [
+ [1, 0, 0, 0, 1, 0, 0],
+ [0, 1, 0, 0, 0, 1, 0],
+ [0, 0, 1, 0, 0, 0, 1],
+ [0, 0, 0, 1, 0, 0, 0],
+ [0, 0, 0, 0, 1, 0, 0],
+ [0, 0, 0, 0, 0, 1, 0],
+ [0, 0, 0, 0, 0, 0, 1],
+ ]
+ )
+ self.kf.H = np.array(
+ [
+ [1, 0, 0, 0, 0, 0, 0],
+ [0, 1, 0, 0, 0, 0, 0],
+ [0, 0, 1, 0, 0, 0, 0],
+ [0, 0, 0, 1, 0, 0, 0],
+ ]
+ )
+
+ self.kf.R[2:, 2:] *= 10.0
+ self.kf.P[
+ 4:, 4:
+ ] *= 1000.0 # give high uncertainty to the unobservable initial velocities
+ self.kf.P *= 10.0
+ self.kf.Q[-1, -1] *= 0.01
+ self.kf.Q[4:, 4:] *= 0.01
+
+ self.kf.x[:4] = convert_bbox_to_z(bbox)
+ self.time_since_update = 0
+ self.id = KalmanBoxTracker.count
+ KalmanBoxTracker.count += 1
+ self.history = []
+ self.hits = 0
+ self.hit_streak = 0
+ self.age = 0
+
+ def update(self, bbox):
+ """
+ Updates the state vector with observed bbox.
+ """
+ self.time_since_update = 0
+ self.history = []
+ self.hits += 1
+ self.hit_streak += 1
+ self.kf.update(convert_bbox_to_z(bbox))
+
+ def predict(self):
+ """
+ Advances the state vector and returns the predicted bounding box estimate.
+ """
+ if (self.kf.x[6] + self.kf.x[2]) <= 0:
+ self.kf.x[6] *= 0.0
+ self.kf.predict()
+ self.age += 1
+ if self.time_since_update > 0:
+ self.hit_streak = 0
+ self.time_since_update += 1
+ self.history.append(convert_x_to_bbox(self.kf.x))
+ return self.history[-1]
+
+ def get_state(self):
+ """
+ Returns the current bounding box estimate.
+ """
+ return convert_x_to_bbox(self.kf.x)
+
+
+def associate_detections_to_trackers(detections, trackers, iou_threshold=0.3):
+ """
+ Assigns detections to tracked object (both represented as bounding boxes)
+
+ Returns 3 lists of matches, unmatched_detections and unmatched_trackers
+ """
+ if len(trackers) == 0:
+ return (
+ np.empty((0, 2), dtype=int),
+ np.arange(len(detections)),
+ np.empty((0, 5), dtype=int),
+ )
+ iou_matrix = np.zeros((len(detections), len(trackers)), dtype=np.float32)
+
+ for d, det in enumerate(detections):
+ for t, trk in enumerate(trackers):
+ iou_matrix[d, t] = iou(det, trk)
+
+ matched_indices = linear_sum_assignment(-iou_matrix)
+ matched_indices = np.array(list(zip(*matched_indices)), dtype=np.int)
+ matched_indices.shape = (-1, 2)
+ # print(matched_indices)
+ # print(type(matched_indices))
+
+ unmatched_detections = []
+ for d, det in enumerate(detections):
+ if d not in matched_indices[:, 0]:
+ unmatched_detections.append(d)
+ unmatched_trackers = []
+ for t, trk in enumerate(trackers):
+ if t not in matched_indices[:, 1]:
+ unmatched_trackers.append(t)
+
+ # filter out matched with low IOU
+ matches = []
+ for m in matched_indices:
+ if iou_matrix[m[0], m[1]] < iou_threshold:
+ unmatched_detections.append(m[0])
+ unmatched_trackers.append(m[1])
+ else:
+ matches.append(m.reshape(1, 2))
+ if len(matches) == 0:
+ matches = np.empty((0, 2), dtype=int)
+ else:
+ matches = np.concatenate(matches, axis=0)
+
+ return matches, np.array(unmatched_detections), np.array(unmatched_trackers)
+
+
+class Sort(object):
+ def __init__(self, max_age=1, min_hits=3):
+ """
+ Sets key parameters for SORT
+ """
+ self.max_age = max_age
+ self.min_hits = min_hits
+ self.trackers = []
+ self.frame_count = 0
+
+ def update(self, dets):
+ """
+ Params:
+ dets - a numpy array of detections in the format [[x1,y1,x2,y2,score],[x1,y1,x2,y2,score],...]
+ Requires: this method must be called once for each frame even with empty detections.
+ Returns the a similar array, where the last column is the object ID.
+
+ NOTE: The number of objects returned may differ from the number of detections provided.
+ """
+ self.frame_count += 1
+ # get predicted locations from existing trackers.
+ trks = np.zeros((len(self.trackers), 5))
+ to_del = []
+ ret = []
+ for t, trk in enumerate(trks):
+ pos = self.trackers[t].predict()[0]
+ trk[:] = [pos[0], pos[1], pos[2], pos[3], 0]
+ if np.any(np.isnan(pos)):
+ to_del.append(t)
+ trks = np.ma.compress_rows(np.ma.masked_invalid(trks))
+ for t in reversed(to_del):
+ self.trackers.pop(t)
+ matched, unmatched_dets, unmatched_trks = associate_detections_to_trackers(
+ dets, trks
+ )
+
+ # update matched trackers with assigned detections
+ for t, trk in enumerate(self.trackers):
+ if t not in unmatched_trks:
+ d = matched[np.where(matched[:, 1] == t)[0], 0]
+ trk.update(dets[d, :][0])
+
+ # create and initialise new trackers for unmatched detections
+ for i in unmatched_dets:
+ trk = KalmanBoxTracker(dets[i, :])
+ self.trackers.append(trk)
+ i = len(self.trackers)
+ for trk in reversed(self.trackers):
+ d = trk.get_state()[0]
+ if (trk.time_since_update < 1) and (
+ trk.hit_streak >= self.min_hits or self.frame_count <= self.min_hits
+ ):
+ ret.append(
+ np.concatenate((d, [trk.id + 1])).reshape(1, -1)
+ ) # +1 as MOT benchmark requires positive
+ i -= 1
+ # remove dead tracklet
+ if trk.time_since_update > self.max_age:
+ self.trackers.pop(i)
+ if len(ret) > 0:
+ return np.concatenate(ret)
+ return np.empty((0, 5))
+
+
+def parse_args():
+ """Parse input arguments."""
+ parser = argparse.ArgumentParser(description="SORT demo")
+ parser.add_argument(
+ "--display",
+ dest="display",
+ help="Display online tracker output (slow) [False]",
+ action="store_true",
+ )
+ args = parser.parse_args()
+ return args
diff --git a/clean/video/pwtf_dvd/inference/test_tools/ct/tracking/tracker.py b/clean/video/pwtf_dvd/inference/test_tools/ct/tracking/tracker.py
new file mode 100644
index 0000000000000000000000000000000000000000..20dd79f41a56bcb4873e83bb9c76db9ffbf0f627
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_tools/ct/tracking/tracker.py
@@ -0,0 +1,27 @@
+from .sort import Sort
+import numpy as np
+
+
+def get_detections(faces):
+ detections = []
+ for face in faces:
+ x1, y1, x2, y2 = face[0]
+ detections.append((x1, y1, x2, y2, face[-1]))
+ return np.array(detections)
+
+
+def get_tracks(detect_results):
+ tracks = {}
+ mot_tracker = Sort()
+ for faces in detect_results:
+ detections = get_detections(faces)
+ track_bbs_ids = mot_tracker.update(detections)
+ for track in track_bbs_ids: # 单独框出每一张人脸
+ id = int(track[-1])
+ box = track[:4]
+ if id in tracks:
+ tracks[id].append(box)
+ else:
+ tracks[id] = [box]
+
+ return [track for id, track in tracks.items() if len(track) == len(detect_results)]
diff --git a/clean/video/pwtf_dvd/inference/test_tools/ct/utils.py b/clean/video/pwtf_dvd/inference/test_tools/ct/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..5ff180fb2a24b9e8bd673d3428e3b101da16a6b0
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_tools/ct/utils.py
@@ -0,0 +1,5 @@
+import cv2
+
+
+def write_img(file, img):
+ cv2.imwrite(file, img, [cv2.IMWRITE_PNG_COMPRESSION, 0])
diff --git a/clean/video/pwtf_dvd/inference/test_tools/faster_crop_align_xray.py b/clean/video/pwtf_dvd/inference/test_tools/faster_crop_align_xray.py
new file mode 100644
index 0000000000000000000000000000000000000000..3e99f2f0930e26539793e3d76d74e9df7afc186b
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_tools/faster_crop_align_xray.py
@@ -0,0 +1,73 @@
+import numpy as np
+import cv2
+from .warp_for_xray import (
+ estimiate_batch_transform,
+ transform_landmarks,
+ std_points_256,
+)
+import numpy as np
+
+
+class FasterCropAlignXRay:
+ """
+ 修正到统一坐标系，统一图像大小到标准尺寸
+ """
+
+ def __init__(self, size=256):
+ self.image_size = size
+ self.std_points = std_points_256 * size / 256.0
+
+ def __call__(self, landmarks, images=None, jitter=False):
+ landmarks = [landmark[:4] for landmark in landmarks]
+ ori_boxes = np.array([ori_box for _, _, _, ori_box in landmarks])
+ five_landmarks = np.array([ldm5 for _, ldm5, _, _ in landmarks])
+ landmarks68 = np.array([ldm68 for _, _, ldm68, _ in landmarks])
+ # assert landmarks68.min() > 0
+
+ left_top = ori_boxes[:, :2].min(0)
+
+ right_bottom = ori_boxes[:, 2:].max(0)
+
+ size = right_bottom - left_top
+
+ w, h = size
+
+ diff = ori_boxes[:, :2] - left_top[None, ...]
+
+ new_five_landmarks = five_landmarks + diff[:, None, :]
+ new_landmarks68 = landmarks68 + diff[:, None, :]
+
+ landmark_for_estimiate = new_five_landmarks.copy()
+ if jitter:
+ landmark_for_estimiate += np.random.uniform(
+ -4, 4, landmark_for_estimiate.shape
+ )
+
+ tfm, trans = estimiate_batch_transform(
+ landmark_for_estimiate, tgt_pts=self.std_points
+ )
+
+ transformed_landmarks68 = np.array(
+ [transform_landmarks(ldm68, trans) for ldm68 in new_landmarks68]
+ )
+
+ if images is not None:
+ transformed_images = [
+ self.process_sinlge(tfm, image, d, h, w)
+ for image, d in zip(images, diff)
+ ] # 拼接 func 的参数
+ transformed_images = np.stack(transformed_images)
+ return transformed_landmarks68, transformed_images
+ else:
+ return transformed_landmarks68
+
+ def process_sinlge(self, tfm, image, d, h, w):
+ assert isinstance(image, np.ndarray)
+ new_image = np.zeros((h, w, 3), dtype=np.uint8)
+ x, y = d
+ ih, iw, _ = image.shape
+ new_image[y : y + ih, x : x + iw] = image
+ transformed_image = cv2.warpAffine(
+ new_image, tfm, (self.image_size, self.image_size)
+ )
+ return transformed_image
diff --git a/clean/video/pwtf_dvd/inference/test_tools/supply_writer.py b/clean/video/pwtf_dvd/inference/test_tools/supply_writer.py
new file mode 100644
index 0000000000000000000000000000000000000000..0394dfbb8fc50d47984eaa937537de144e51ee81
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_tools/supply_writer.py
@@ -0,0 +1,49 @@
+import cv2
+
+class SupplyWriter:
+ def __init__(self, intput_video, output_video, opt_thres, rgb_input=True):
+ reader = cv2.VideoCapture(intput_video)
+ fourcc = cv2.VideoWriter_fourcc(*"XVID")
+ fps = reader.get(cv2.CAP_PROP_FPS)
+ width = int(reader.get(3))
+ height = int(reader.get(4))
+ reader.release()
+ self.padding = 40
+
+ self.writer = cv2.VideoWriter(output_video, fourcc, fps, (height, width)[::-1])
+ self.rgb_input = rgb_input
+ self.opt_thres = opt_thres
+
+ def run(self, images, scores, boxes):
+ # Text variables
+ font_face = cv2.FONT_HERSHEY_SIMPLEX
+ thickness = 5
+ font_scale = 3
+
+ for image, score, box in zip(images, scores, boxes):
+ if self.rgb_input:
+ image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+ if box is not None:
+ label = "fake" if score > self.opt_thres else "real"
+ x1, y1, x2, y2 = box
+ x = int(x1)
+ y = int(y1)
+ w = int(x2 - x1)
+ h = int(y2 - y1)
+ color = (
+ (255, 255, 0) if label == "real" else (0, 255, 255)
+ ) # BGR 255 0
+ cv2.putText(
+ image,
+ label,
+ (x, y + h + 68),
+ font_face,
+ font_scale,
+ color,
+ thickness,
+ 2,
+ )
+ # draw box over face
+ cv2.rectangle(image, (x, y), (x + w, y + h), color, 10)
+ self.writer.write(image)
+ self.writer.release()
diff --git a/clean/video/pwtf_dvd/inference/test_tools/utils.py b/clean/video/pwtf_dvd/inference/test_tools/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..8069d72b1db1bab8155f88782b5fb6a15fb9563c
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_tools/utils.py
@@ -0,0 +1,115 @@
+import numpy as np
+import cv2
+import os
+import platform
+import json
+import errno
+
+
+
+def weak_check(detect_res):
+ return sum([len(faces) for faces in detect_res]) > len(detect_res) * 0.75
+
+
+def get_crop_box(shape, box, scale=0.5):
+ height, width = shape
+ box = np.rint(box).astype(np.int32)
+ new_box = box.reshape(2, 2)
+ size = new_box[1] - new_box[0]
+ diff = scale * size
+ diff = diff[None, :] * np.array([-1, 1])[:, None]
+ new_box = new_box + diff
+ new_box[:, 0] = np.clip(new_box[:, 0], 0, width - 1)
+ new_box[:, 1] = np.clip(new_box[:, 1], 0, height - 1)
+ new_box = np.rint(new_box).astype(np.int32)
+ return new_box.reshape(-1)
+
+
+def get_fps(input_file):
+ reader = cv2.VideoCapture(input_file)
+ fps = reader.get(cv2.CAP_PROP_FPS)
+ reader.release()
+ return fps
+
+
+
+def mkdir_p(dirname):
+ """Like "mkdir -p", make a dir recursively, but do nothing if the dir exists
+ 这个是线程安全的, from Lingzhi Li
+ Args:
+ dirname(str):
+ """
+ assert dirname is not None
+ if dirname == "" or os.path.isdir(dirname):
+ return
+ try:
+ os.makedirs(dirname)
+ except OSError as e:
+ if e.errno != errno.EEXIST:
+ raise e
+
+
+def mkdir(*args):
+ for folder in args:
+ if not os.path.isdir(folder):
+ mkdir_p(folder)
+
+
+def make_join(*args):
+ folder = os.path.join(*args)
+ mkdir(folder)
+ return folder
+
+
+def list_dir(folder, condition=None, key=lambda x: x, reverse=False, co_join=[]):
+ files = os.listdir(folder)
+ if condition is not None:
+ files = filter(condition, files)
+ co_join = [folder] + co_join
+ if key is not None:
+ files = sorted(files, key=key, reverse=reverse)
+ files = [(file, *[os.path.join(fold, file) for fold in co_join]) for file in files]
+ return files
+
+def get_jointer(file):
+ def jointer(folder):
+ return os.path.join(folder, file)
+
+ return jointer
+
+def flatten(l):
+ return [item for sublist in l for item in sublist]
+
+
+def is_win():
+ return platform.system() == "Windows"
+
+
+def get_postfix(post_fix):
+ return lambda x: x.endswith(post_fix)
+
+
+def partition(images, size):
+ """
+ Returns a new list with elements
+ of which is a list of certain size.
+
+ >>> partition([1, 2, 3, 4], 3)
+ [[1, 2, 3], [4]]
+ """
+ return [
+ images[i : i + size] if i + size <= len(images) else images[i:]
+ for i in range(0, len(images), size)
+ ]
+
+
+def load_json(file):
+ with open(file, "r") as f:
+ res = json.load(f)
+ return res
+
+
+def save_json(file, obj):
+ with open(file, "w", encoding="utf-8") as f:
+ json.dump(obj, f, indent=4, ensure_ascii=False)
+
diff --git a/clean/video/pwtf_dvd/inference/test_tools/warp_for_xray.py b/clean/video/pwtf_dvd/inference/test_tools/warp_for_xray.py
new file mode 100644
index 0000000000000000000000000000000000000000..ea68def0bcc6b384f1198f21fb0d72253f8085d6
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/test_tools/warp_for_xray.py
@@ -0,0 +1,574 @@
+import numpy as np
+import cv2
+
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Jul 11 06:54:28 2017
+
+@author: zhaoyafei
+"""
+
+import numpy as np
+from numpy.linalg import inv, norm, lstsq
+from numpy.linalg import matrix_rank as rank
+
+"""
+Introduction:
+----------
+numpy implemetation form matlab function CP2TFORM(...)
+with 'transformtype':
+ 1) 'nonreflective similarity'
+ 2) 'similarity'
+
+
+MATLAB code:
+----------
+%--------------------------------------
+% Function findNonreflectiveSimilarity
+%
+function [trans, output] = findNonreflectiveSimilarity(uv,xy,options)
+%
+% For a nonreflective similarity:
+%
+% let sc = s*cos(theta)
+% let ss = s*sin(theta)
+%
+% [ sc -ss
+% [u v] = [x y 1] * ss sc
+% tx ty]
+%
+% There are 4 unknowns: sc,ss,tx,ty.
+%
+% Another way to write this is:
+%
+% u = [x y 1 0] * [sc
+% ss
+% tx
+% ty]
+%
+% v = [y -x 0 1] * [sc
+% ss
+% tx
+% ty]
+%
+% With 2 or more correspondence points we can combine the u equations and
+% the v equations for one linear system to solve for sc,ss,tx,ty.
+%
+% [ u1 ] = [ x1 y1 1 0 ] * [sc]
+% [ u2 ] [ x2 y2 1 0 ] [ss]
+% [ ... ] [ ... ] [tx]
+% [ un ] [ xn yn 1 0 ] [ty]
+% [ v1 ] [ y1 -x1 0 1 ]
+% [ v2 ] [ y2 -x2 0 1 ]
+% [ ... ] [ ... ]
+% [ vn ] [ yn -xn 0 1 ]
+%
+% Or rewriting the above matrix equation:
+% U = X * r, where r = [sc ss tx ty]'
+% so r = X\ U.
+%
+
+K = options.K;
+M = size(xy,1);
+x = xy(:,1);
+y = xy(:,2);
+X = [x y ones(M,1) zeros(M,1);
+ y -x zeros(M,1) ones(M,1) ];
+
+u = uv(:,1);
+v = uv(:,2);
+U = [u; v];
+
+% We know that X * r = U
+if rank(X) >= 2*K
+ r = X \ U;
+else
+ error(message('images:cp2tform:twoUniquePointsReq'))
+end
+
+sc = r(1);
+ss = r(2);
+tx = r(3);
+ty = r(4);
+
+Tinv = [sc -ss 0;
+ ss sc 0;
+ tx ty 1];
+
+T = inv(Tinv);
+T(:,3) = [0 0 1]';
+
+trans = maketform('affine', T);
+output = [];
+
+%-------------------------
+% Function findSimilarity
+%
+function [trans, output] = findSimilarity(uv,xy,options)
+%
+% The similarities are a superset of the nonreflective similarities as they may
+% also include reflection.
+%
+% let sc = s*cos(theta)
+% let ss = s*sin(theta)
+%
+% [ sc -ss
+% [u v] = [x y 1] * ss sc
+% tx ty]
+%
+% OR
+%
+% [ sc ss
+% [u v] = [x y 1] * ss -sc
+% tx ty]
+%
+% Algorithm:
+% 1) Solve for trans1, a nonreflective similarity.
+% 2) Reflect the xy data across the Y-axis,
+% and solve for trans2r, also a nonreflective similarity.
+% 3) Transform trans2r to trans2, undoing the reflection done in step 2.
+% 4) Use TFORMFWD to transform uv using both trans1 and trans2,
+% and compare the results, Returnsing the transformation corresponding
+% to the smaller L2 norm.
+
+% Need to reset options.K to prepare for calls to findNonreflectiveSimilarity.
+% This is safe because we already checked that there are enough point pairs.
+options.K = 2;
+
+% Solve for trans1
+[trans1, output] = findNonreflectiveSimilarity(uv,xy,options);
+
+
+% Solve for trans2
+
+% manually reflect the xy data across the Y-axis
+xyR = xy;
+xyR(:,1) = -1*xyR(:,1);
+
+trans2r = findNonreflectiveSimilarity(uv,xyR,options);
+
+% manually reflect the tform to undo the reflection done on xyR
+TreflectY = [-1 0 0;
+ 0 1 0;
+ 0 0 1];
+trans2 = maketform('affine', trans2r.tdata.T * TreflectY);
+
+
+% Figure out if trans1 or trans2 is better
+xy1 = tformfwd(trans1,uv);
+norm1 = norm(xy1-xy);
+
+xy2 = tformfwd(trans2,uv);
+norm2 = norm(xy2-xy);
+
+if norm1 <= norm2
+ trans = trans1;
+else
+ trans = trans2;
+end
+"""
+
+
+class MatlabCp2tormException(Exception):
+ def __str__(self):
+ return "In File {}:{}".format(__file__, super.__str__(self))
+
+
+def tformfwd(trans, uv):
+ """
+ Function:
+ ----------
+ apply affine transform 'trans' to uv
+
+ Parameters:
+ ----------
+ @trans: 3x3 np.array
+ transform matrix
+ @uv: Kx2 np.array
+ each row is a pair of coordinates (x, y)
+
+ Returns:
+ ----------
+ @xy: Kx2 np.array
+ each row is a pair of transformed coordinates (x, y)
+ """
+ uv = np.hstack((uv, np.ones((uv.shape[0], 1))))
+ xy = np.dot(uv, trans)
+ xy = xy[:, 0:-1]
+ return xy
+
+
+def tforminv(trans, uv):
+ """
+ Function:
+ ----------
+ apply the inverse of affine transform 'trans' to uv
+
+ Parameters:
+ ----------
+ @trans: 3x3 np.array
+ transform matrix
+ @uv: Kx2 np.array
+ each row is a pair of coordinates (x, y)
+
+ Returns:
+ ----------
+ @xy: Kx2 np.array
+ each row is a pair of inverse-transformed coordinates (x, y)
+ """
+ Tinv = inv(trans)
+ xy = tformfwd(Tinv, uv)
+ return xy
+
+
+def findNonreflectiveSimilarity(uv, xy, options=None):
+ """
+ Function:
+ ----------
+ Find Non-reflective Similarity Transform Matrix 'trans':
+ u = uv[:, 0]
+ v = uv[:, 1]
+ x = xy[:, 0]
+ y = xy[:, 1]
+ [x, y, 1] = [u, v, 1] * trans
+
+ Parameters:
+ ----------
+ @uv: Kx2 np.array
+ source points each row is a pair of coordinates (x, y)
+ @xy: Kx2 np.array
+ each row is a pair of inverse-transformed
+ @option: not used, keep it as None
+
+ Returns:
+ @trans: 3x3 np.array
+ transform matrix from uv to xy
+ @trans_inv: 3x3 np.array
+ inverse of trans, transform matrix from xy to uv
+
+ Matlab:
+ ----------
+ % For a nonreflective similarity:
+ %
+ % let sc = s*cos(theta)
+ % let ss = s*sin(theta)
+ %
+ % [ sc -ss
+ % [u v] = [x y 1] * ss sc
+ % tx ty]
+ %
+ % There are 4 unknowns: sc,ss,tx,ty.
+ %
+ % Another way to write this is:
+ %
+ % u = [x y 1 0] * [sc
+ % ss
+ % tx
+ % ty]
+ %
+ % v = [y -x 0 1] * [sc
+ % ss
+ % tx
+ % ty]
+ %
+ % With 2 or more correspondence points we can combine the u equations and
+ % the v equations for one linear system to solve for sc,ss,tx,ty.
+ %
+ % [ u1 ] = [ x1 y1 1 0 ] * [sc]
+ % [ u2 ] [ x2 y2 1 0 ] [ss]
+ % [ ... ] [ ... ] [tx]
+ % [ un ] [ xn yn 1 0 ] [ty]
+ % [ v1 ] [ y1 -x1 0 1 ]
+ % [ v2 ] [ y2 -x2 0 1 ]
+ % [ ... ] [ ... ]
+ % [ vn ] [ yn -xn 0 1 ]
+ %
+ % Or rewriting the above matrix equation:
+ % U = X * r, where r = [sc ss tx ty]'
+ % so r = X\ U.
+ %
+ """
+ options = {"K": 2}
+
+ K = options["K"]
+ M = xy.shape[0]
+ x = xy[:, 0].reshape((-1, 1)) # use reshape to keep a column vector
+ y = xy[:, 1].reshape((-1, 1)) # use reshape to keep a column vector
+ # print '--->x, y:\n', x, y
+
+ tmp1 = np.hstack((x, y, np.ones((M, 1)), np.zeros((M, 1))))
+ tmp2 = np.hstack((y, -x, np.zeros((M, 1)), np.ones((M, 1))))
+ X = np.vstack((tmp1, tmp2))
+ # print '--->X.shape: ', X.shape
+ # print 'X:\n', X
+
+ u = uv[:, 0].reshape((-1, 1)) # use reshape to keep a column vector
+ v = uv[:, 1].reshape((-1, 1)) # use reshape to keep a column vector
+ U = np.vstack((u, v))
+ # print '--->U.shape: ', U.shape
+ # print 'U:\n', U
+
+ # We know that X * r = U
+ if rank(X) >= 2 * K:
+ r, _, _, _ = lstsq(X, U, rcond=-1)
+ r = np.squeeze(r)
+ else:
+ raise Exception("cp2tform:twoUniquePointsReq")
+
+ # print '--->r:\n', r
+
+ sc = r[0]
+ ss = r[1]
+ tx = r[2]
+ ty = r[3]
+
+ Tinv = np.array([[sc, -ss, 0], [ss, sc, 0], [tx, ty, 1]])
+
+ # print '--->Tinv:\n', Tinv
+
+ T = inv(Tinv)
+ # print '--->T:\n', T
+
+ T[:, 2] = np.array([0, 0, 1])
+
+ return T, Tinv
+
+
+def findSimilarity(uv, xy, options=None):
+ """
+ Function:
+ ----------
+ Find Reflective Similarity Transform Matrix 'trans':
+ u = uv[:, 0]
+ v = uv[:, 1]
+ x = xy[:, 0]
+ y = xy[:, 1]
+ [x, y, 1] = [u, v, 1] * trans
+
+ Parameters:
+ ----------
+ @uv: Kx2 np.array
+ source points each row is a pair of coordinates (x, y)
+ @xy: Kx2 np.array
+ each row is a pair of inverse-transformed
+ @option: not used, keep it as None
+
+ Returns:
+ ----------
+ @trans: 3x3 np.array
+ transform matrix from uv to xy
+ @trans_inv: 3x3 np.array
+ inverse of trans, transform matrix from xy to uv
+
+ Matlab:
+ ----------
+ % The similarities are a superset of the nonreflective similarities as they may
+ % also include reflection.
+ %
+ % let sc = s*cos(theta)
+ % let ss = s*sin(theta)
+ %
+ % [ sc -ss
+ % [u v] = [x y 1] * ss sc
+ % tx ty]
+ %
+ % OR
+ %
+ % [ sc ss
+ % [u v] = [x y 1] * ss -sc
+ % tx ty]
+ %
+ % Algorithm:
+ % 1) Solve for trans1, a nonreflective similarity.
+ % 2) Reflect the xy data across the Y-axis,
+ % and solve for trans2r, also a nonreflective similarity.
+ % 3) Transform trans2r to trans2, undoing the reflection done in step 2.
+ % 4) Use TFORMFWD to transform uv using both trans1 and trans2,
+ % and compare the results, Returnsing the transformation corresponding
+ % to the smaller L2 norm.
+
+ % Need to reset options.K to prepare for calls to findNonreflectiveSimilarity.
+ % This is safe because we already checked that there are enough point pairs.
+ """
+ options = {"K": 2}
+
+ # uv = np.array(uv)
+ # xy = np.array(xy)
+
+ # Solve for trans1
+ trans1, trans1_inv = findNonreflectiveSimilarity(uv, xy, options)
+
+ # Solve for trans2
+
+ # manually reflect the xy data across the Y-axis
+ xyR = xy
+ xyR[:, 0] = -1 * xyR[:, 0]
+
+ trans2r, trans2r_inv = findNonreflectiveSimilarity(uv, xyR, options)
+
+ # manually reflect the tform to undo the reflection done on xyR
+ TreflectY = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, 1]])
+
+ trans2 = np.dot(trans2r, TreflectY)
+
+ # Figure out if trans1 or trans2 is better
+ xy1 = tformfwd(trans1, uv)
+ norm1 = norm(xy1 - xy)
+
+ xy2 = tformfwd(trans2, uv)
+ norm2 = norm(xy2 - xy)
+
+ if norm1 <= norm2:
+ return trans1, trans1_inv
+ else:
+ trans2_inv = inv(trans2)
+ return trans2, trans2_inv
+
+
+def get_similarity_transform(src_pts, dst_pts, reflective=True):
+ """
+ Function:
+ ----------
+ Find Similarity Transform Matrix 'trans':
+ u = src_pts[:, 0]
+ v = src_pts[:, 1]
+ x = dst_pts[:, 0]
+ y = dst_pts[:, 1]
+ [x, y, 1] = [u, v, 1] * trans
+
+ Parameters:
+ ----------
+ @src_pts: Kx2 np.array
+ source points, each row is a pair of coordinates (x, y)
+ @dst_pts: Kx2 np.array
+ destination points, each row is a pair of transformed
+ coordinates (x, y)
+ @reflective: True or False
+ if True:
+ use reflective similarity transform
+ else:
+ use non-reflective similarity transform
+
+ Returns:
+ ----------
+ @trans: 3x3 np.array
+ transform matrix from uv to xy
+ trans_inv: 3x3 np.array
+ inverse of trans, transform matrix from xy to uv
+ """
+
+ if reflective:
+ trans, trans_inv = findSimilarity(src_pts, dst_pts)
+ else:
+ trans, trans_inv = findNonreflectiveSimilarity(src_pts, dst_pts)
+
+ return trans, trans_inv
+
+
+def cvt_tform_mat_for_cv2(trans):
+ """
+ Function:
+ ----------
+ Convert Transform Matrix 'trans' into 'cv2_trans' which could be
+ directly used by cv2.warpAffine():
+ u = src_pts[:, 0]
+ v = src_pts[:, 1]
+ x = dst_pts[:, 0]
+ y = dst_pts[:, 1]
+ [x, y].T = cv_trans * [u, v, 1].T
+
+ Parameters:
+ ----------
+ @trans: 3x3 np.array
+ transform matrix from uv to xy
+
+ Returns:
+ ----------
+ @cv2_trans: 2x3 np.array
+ transform matrix from src_pts to dst_pts, could be directly used
+ for cv2.warpAffine()
+ """
+ cv2_trans = trans[:, 0:2].T
+
+ return cv2_trans
+
+
+def get_similarity_transform_for_cv2(src_pts, dst_pts, reflective=True):
+ """
+ Function:
+ ----------
+ Find Similarity Transform Matrix 'cv2_trans' which could be
+ directly used by cv2.warpAffine():
+ u = src_pts[:, 0]
+ v = src_pts[:, 1]
+ x = dst_pts[:, 0]
+ y = dst_pts[:, 1]
+ [x, y].T = cv_trans * [u, v, 1].T
+
+ Parameters:
+ ----------
+ @src_pts: Kx2 np.array
+ source points, each row is a pair of coordinates (x, y)
+ @dst_pts: Kx2 np.array
+ destination points, each row is a pair of transformed
+ coordinates (x, y)
+ reflective: True or False
+ if True:
+ use reflective similarity transform
+ else:
+ use non-reflective similarity transform
+
+ Returns:
+ ----------
+ @cv2_trans: 2x3 np.array
+ transform matrix from src_pts to dst_pts, could be directly used
+ for cv2.warpAffine()
+ """
+ trans, trans_inv = get_similarity_transform(src_pts, dst_pts, reflective)
+ cv2_trans = cvt_tform_mat_for_cv2(trans)
+ return cv2_trans, trans
+
+
+std_points_317 = np.array(
+ [
+ [85.82991, 115.7792],
+ [169.0532, 114.3381],
+ [127.574, 167.0006],
+ [90.6964, 204.7014],
+ [167.3069, 203.3733],
+ ]
+)
+
+
+padding = 30
+
+std_points_317 = std_points_317 + padding
+
+std_points_256=std_points_317.copy()
+std_points_256[..., 0] -= 30
+std_points_256[..., 1] -= 60
+
+def warp_as_face_x_ray(img, src_pts, tgt_pts=std_points_317):
+ tfm, trans = get_similarity_transform_for_cv2(src_pts.copy(), tgt_pts.copy())
+ return cv2.warpAffine(img, tfm, (317, 317)), trans
+
+
+def estimiate_batch_transform(all_src_pts, tgt_pts=std_points_317):
+ tgt_pts = np.repeat(tgt_pts[None, ...], len(all_src_pts), 0).reshape(-1, 2)
+ src_pts = np.array(all_src_pts).reshape(-1, 2)
+ tfm, trans = get_similarity_transform_for_cv2(src_pts, tgt_pts)
+ return tfm, trans
+
+
+def batch_warp_as_face_x_ray(images, all_src_pts, tgt_pts=std_points_317):
+ tfm, trans = estimiate_batch_transform(all_src_pts, tgt_pts)
+ return [cv2.warpAffine(img, tfm, (317, 317)) for img in images], trans
+
+
+def transform_landmarks(landmarks, trans):
+ transformed = np.hstack((landmarks, np.ones((landmarks.shape[0], 1))))
+ transformed = np.dot(transformed, trans)
+ return transformed[:, :2]
+
+def compute_reverse_trans(trans):
+ return np.linalg.inv(trans)
\ No newline at end of file
diff --git a/clean/video/pwtf_dvd/inference/utils/__init__.py b/clean/video/pwtf_dvd/inference/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..7437239283d4494444c4b0993d34c90f6891e0c8
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/utils/__init__.py
@@ -0,0 +1,7 @@
+
+__all__ = [] # do not use ' from utils import *'
+
+from .common import *
+from .plugin_loader import PluginLoader
+#from .plugin_loaderv2 import PluginLoader as PluginLoaderV2
+from .model_loader import add_loader
\ No newline at end of file
diff --git a/clean/video/pwtf_dvd/inference/utils/common.py b/clean/video/pwtf_dvd/inference/utils/common.py
new file mode 100644
index 0000000000000000000000000000000000000000..7c975bf6a08077296351bd10c03ea719538e48a3
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/utils/common.py
@@ -0,0 +1,80 @@
+#!/usr/bin/python
+# -*- coding: UTF-8 -*-
+
+
+import os
+import torch
+from torch.autograd import Variable
+import errno
+import torch.distributed as dist
+import math
+from functools import reduce
+def make_folder(path, version):
+ if not os.path.exists(os.path.join(path, version)):
+ print(os.path.join(path, version))
+ os.makedirs(os.path.join(path, version))
+
+
+def tensor2var(x, grad=False):
+ if torch.cuda.is_available():
+ x = x.cuda()
+ return Variable(x, requires_grad=grad)
+
+def var2tensor(x):
+ return x.data.cpu()
+
+def var2numpy(x):
+ return x.data.cpu().numpy()
+
+def denorm(x):
+ out = (x + 1) / 2
+ return out.clamp_(0, 1)
+
+def mkdir_p(dirname):
+ """ Like "mkdir -p", make a dir recursively, but do nothing if the dir exists
+ Args:
+ dirname(str):
+ """
+ assert dirname is not None
+ if dirname == '' or os.path.isdir(dirname):
+ return
+ try:
+ os.makedirs(dirname)
+ except OSError as e:
+ if e.errno != errno.EEXIST:
+ raise e
+
+
+def skipShardSplit(aList, drop_last=False, num_replicas=None, rank=None):
+ if not isinstance(aList, list) and not isinstance(aList, tuple):
+ aList = List
+
+ if num_replicas is None:
+ num_replicas = dist.get_world_size() if dist.is_initialized() else 1
+ if rank is None:
+ rank = dist.get_rank() if dist.is_initialized() else 0
+
+ num_replicas = num_replicas
+ rank = rank
+ drop_last = drop_last
+
+ if drop_last:
+ aList = aList[0: (len(aList) // num_replicas) * num_replicas]
+
+ # subsample
+ aList = aList[rank::num_replicas]
+
+ return aList
+
+def mixb2a(a,b):
+ if len(b) > len(a):
+ a,b = b,a
+ if len(b) == 0:
+ return a
+ chunk_num = (len(b))
+ a_chunk = splitIntoChunk(a, chunk_num)
+ b_chunk = list(map(lambda x:[x],b))
+ return reduce(lambda x, y: x+y, [_a+_b for _a,_b in zip(a_chunk, b_chunk)])
+
+def splitIntoChunk(aList, chunk_num):
+ return [aList[math.ceil(k * (len(aList) / chunk_num)):math.ceil((k + 1) * (len(aList) / chunk_num)):] for k in range(chunk_num)]
\ No newline at end of file
diff --git a/clean/video/pwtf_dvd/inference/utils/logger.py b/clean/video/pwtf_dvd/inference/utils/logger.py
new file mode 100644
index 0000000000000000000000000000000000000000..b5e3c52c3ae53a698ba67581a43bf20d832eca0d
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/utils/logger.py
@@ -0,0 +1,182 @@
+#!/usr/bin/python
+# -*- coding: UTF-8 -*-
+# Modified by: algohunt
+# Microsoft Research & Peking University
+# lilingzhi@pku.edu.cn
+# Copyright (c) 2019
+
+
+# -*- coding: utf-8 -*-
+
+"""
+Borrow from tensorpack credit goes to yuxin wu
+The logger module itself has the common logging functions of Python's
+:class:`logging.Logger`. For example:
+
+.. code-block:: python
+
+ from utils import logger
+ logger.set_logger_dir('train_log/test')
+ logger.info("Test")
+ logger.error("Error happened!")
+"""
+
+
+import logging
+import os
+import os.path
+import shutil
+import sys
+from datetime import datetime, timedelta
+from six.moves import input
+from termcolor import colored
+import time
+
+__all__ = ['set_logger_dir', 'auto_set_dir', 'get_logger_dir']
+
+
+class _MyFormatter(logging.Formatter):
+ def format(self, record):
+ date = colored('[%(asctime)s @%(filename)s:%(lineno)d]', 'green')
+ msg = '%(message)s'
+ if record.levelno == logging.WARNING:
+ fmt = date + ' ' + colored('WRN', 'red', attrs=['blink']) + ' ' + msg
+ elif record.levelno == logging.ERROR or record.levelno == logging.CRITICAL:
+ fmt = date + ' ' + colored('ERR', 'red', attrs=['blink', 'underline']) + ' ' + msg
+ elif record.levelno == logging.DEBUG:
+ fmt = date + ' ' + colored('DBG', 'yellow', attrs=['blink']) + ' ' + msg
+ else:
+ fmt = date + ' ' + msg
+ if hasattr(self, '_style'):
+ # Python3 compatibility
+ self._style._fmt = fmt
+ self._fmt = fmt
+ return super(_MyFormatter, self).format(record)
+
+
+def _getlogger():
+ logger = logging.getLogger('tensorpack')
+ logger.propagate = False
+ logger.setLevel(logging.INFO)
+ handler = logging.StreamHandler(sys.stdout)
+ handler.setFormatter(_MyFormatter(datefmt='%m%d %H:%M:%S'))
+ logger.addHandler(handler)
+ return logger
+
+
+_logger = _getlogger()
+_LOGGING_METHOD = ['info', 'warning', 'error', 'critical', 'exception', 'debug', 'setLevel']
+# export logger functions
+for func in _LOGGING_METHOD:
+ locals()[func] = getattr(_logger, func)
+ __all__.append(func)
+# 'warn' is deprecated in logging module
+warn = _logger.warning
+__all__.append('warn')
+
+
+def _get_time_str():
+ utc_time = datetime.utcfromtimestamp(time.time())
+ beijing_time = utc_time- timedelta(hours=8)
+ return beijing_time.strftime('%m%d-%H%M%S')
+
+
+# globals: logger file and directory:
+LOG_DIR = None
+_FILE_HANDLER = None
+
+
+def _set_file(path):
+ global _FILE_HANDLER
+ if os.path.isfile(path):
+ backup_name = path + '.' + _get_time_str()
+ shutil.move(path, backup_name)
+ _logger.info("Existing log file '{}' backuped to '{}'".format(path, backup_name)) # noqa: F821
+ hdl = logging.FileHandler(
+ filename=path, encoding='utf-8', mode='w')
+ hdl.setFormatter(_MyFormatter(datefmt='%m%d %H:%M:%S'))
+
+ _FILE_HANDLER = hdl
+ _logger.addHandler(hdl)
+ _logger.info("Argv: " + ' '.join(sys.argv))
+
+
+def set_logger_dir(dirname, action=None):
+ """
+ Set the directory for global logging.
+
+ Args:
+ dirname(str): log directory
+ action(str): an action of ["k","d","q"] to be performed
+ when the directory exists. Will ask user by default.
+
+ "d": delete the directory. Note that the deletion may fail when
+ the directory is used by tensorboard.
+
+ "k": keep the directory. This is useful when you resume from a
+ previous training and want the directory to look as if the
+ training was not interrupted.
+ Note that this option does not load old models or any other
+ old states for you. It simply does nothing.
+
+ """
+ global LOG_DIR, _FILE_HANDLER
+ if _FILE_HANDLER:
+ # unload and close the old file handler, so that we may safely delete the logger directory
+ _logger.removeHandler(_FILE_HANDLER)
+ del _FILE_HANDLER
+
+ def dir_nonempty(dirname):
+ # If directory exists and nonempty (ignore hidden files), prompt for action
+ return os.path.isdir(dirname) and len([x for x in os.listdir(dirname) if x[0] != '.'])
+
+ if dir_nonempty(dirname):
+ if not action:
+ _logger.warn("""\
+Log directory {} exists! Use 'd' to delete it. """.format(dirname))
+ _logger.warn("""\
+If you're resuming from a previous run, you can choose to keep it.
+Press any other key to exit. """)
+ while not action:
+ action = input("Select Action: k (keep) / d (delete) / q (quit):").lower().strip()
+ act = action
+ if act == 'b':
+ backup_name = dirname + _get_time_str()
+ shutil.move(dirname, backup_name)
+ info("Directory '{}' backuped to '{}'".format(dirname, backup_name)) # noqa: F821
+ elif act == 'd':
+ shutil.rmtree(dirname, ignore_errors=True)
+ if dir_nonempty(dirname):
+ shutil.rmtree(dirname, ignore_errors=False)
+ elif act == 'n':
+ dirname = dirname + _get_time_str()
+ info("Use a new log directory {}".format(dirname)) # noqa: F821
+ elif act == 'k':
+ pass
+ else:
+ raise OSError("Directory {} exits!".format(dirname))
+ LOG_DIR = dirname
+ from . import mkdir_p
+ mkdir_p(dirname)
+ _set_file(os.path.join(dirname, 'log.log'))
+
+
+def auto_set_dir(action=None, name=None):
+ """
+ Use :func:`logger.set_logger_dir` to set log directory to
+ "./train_log/{scriptname}:{name}". "scriptname" is the name of the main python file currently running"""
+ mod = sys.modules['__main__']
+ basename = os.path.basename(mod.__file__)
+ auto_dirname = os.path.join('train_log', basename[:basename.rfind('.')])
+ if name:
+ auto_dirname += '_%s' % name if os.name == 'nt' else ':%s' % name
+ set_logger_dir(auto_dirname, action=action)
+
+
+def get_logger_dir():
+ """
+ Returns:
+ The logger directory, or None if not set.
+ The directory is used for general logging, tensorboard events, checkpoints, etc.
+ """
+ return LOG_DIR
diff --git a/clean/video/pwtf_dvd/inference/utils/model_loader.py b/clean/video/pwtf_dvd/inference/utils/model_loader.py
new file mode 100644
index 0000000000000000000000000000000000000000..9be6e723022cd03a23018a72181230caf9525282
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/utils/model_loader.py
@@ -0,0 +1,117 @@
+#!/usr/bin/python
+# -*- coding: UTF-8 -*-
+
+
+import types
+from utils import logger
+from config import config as cfg
+import os
+import sys
+import glob
+import torch
+import traceback
+import types
+import torch.distributed as dist
+import copy
+from .torch_save import torch_save
+
+def add_loader(target, name,max_to_keep=2):
+
+ def get_rank(self):
+ return dist.get_rank() if dist.is_initialized() else 0
+
+ def save_models(self, epoch):
+ """ Backup and save the models """
+ if self.get_rank() == 0:
+ logger.debug("Backing up and saving models")
+ if not os.path.exists(self.model_dir):
+ os.mkdir(self.model_dir)
+
+ torch_save(self.state_dict(), self.get_checkpoint_path(epoch))
+ if os.path.exists(self.get_checkpoint_path(epoch - self.max_to_keep)):
+ os.remove(self.get_checkpoint_path(epoch - self.max_to_keep))
+ logger.info("{} models saved".format(self.name))
+
+ def load(self, fullpath=None, epoch=-1):
+ """ Force Loading a model, or load the latest model"""
+ if fullpath is None:
+ fullpath, loaded_epoch = self.find_last(epoch)
+ else:
+ loaded_epoch = epoch
+
+ if fullpath is None:
+ logger.info("No existing {} model found".format(self.name))
+ return False, -1
+ logger.debug("Loading model: '%s'", fullpath)
+ try:
+ saved_state_dict = torch.load(fullpath, map_location='cpu')
+ self.load_state_dict(saved_state_dict)
+ logger.info(" consume training from {}".format(fullpath))
+ except ValueError as err:
+ logger.warning("Failed loading existing training data for {}. Generating new models".format(self.name))
+ logger.debug("Exception: %s", str(err))
+ return False, -1
+ except OSError as err:
+ logger.warning("Failed loading existing training data for {}. Generating new models".format(self.name))
+ logger.debug("Exception: %s", str(err))
+ return False, -1
+ except RuntimeError as err:
+ logger.warning("{} model has corrupted, try to load earlier one".format(self.name))
+ logger.debug("Exception: %s", str(err))
+ return False, -1
+ except:
+ logger.error(traceback.format_exc())
+ raise
+
+ return True, loaded_epoch
+
+ def get_checkpoint_path(self, epoch):
+ """" returning the checkpoint path w.r.t epoch which should be {name}_{epoch}.pth"""
+ return os.path.join(self.model_dir, self.name + '_' +str(epoch) + '.pth')
+
+
+ def find_last(self, epoch=-1, model_dir=None):
+ """Finds the last checkpoint file of the last trained model in the
+ model directory.
+ Returns:
+ checkpoint :The path of the last checkpoint file
+
+ """
+ if model_dir is None:
+ model_dir = self.model_dir
+ if not os.path.exists(model_dir):
+ logger.info("model dir not exists {} ".format(model_dir))
+ return None, -1
+ #assert os.path.exists(self.model_dir), "model dir not exists {}".format(self.model_dir)
+ checkpoints = glob.glob(os.path.join(model_dir, '*.pth'))
+
+
+ checkpoints = list(filter(lambda x: os.path.basename(x).startswith(self.name), checkpoints))
+ if len(checkpoints) == 0:
+ return None, -1
+ checkpoints = {int(os.path.basename(x).split('.')[0].split('_')[-1]):x for x in checkpoints}
+
+ start = min(checkpoints.keys())
+ end = max(checkpoints.keys())
+
+ if epoch == -1:
+ return checkpoints[end], end
+ elif epoch < start :
+ raise RuntimeError(
+ "model for epoch {} has been deleted as we only keep {} models".format(epoch,self.max_to_keep))
+ elif epoch > end:
+ raise RuntimeError(
+ "epoch {} is bigger than all exist checkpoints".format(epoch))
+ else:
+ return checkpoints[epoch], epoch
+
+ target.find_last = types.MethodType(find_last, target)
+ target.get_checkpoint_path = types.MethodType(get_checkpoint_path, target)
+ target.load = types.MethodType(load, target)
+ target.save_models = types.MethodType(save_models, target)
+ target.get_rank = types.MethodType(get_rank, target)
+
+ target.max_to_keep = max_to_keep
+ target.name = name
+ target.model_dir = os.path.join(cfg.path.model_dir, cfg.setting_name)
+ return target
\ No newline at end of file
diff --git a/clean/video/pwtf_dvd/inference/utils/plugin_loader.py b/clean/video/pwtf_dvd/inference/utils/plugin_loader.py
new file mode 100644
index 0000000000000000000000000000000000000000..bcd5ababfeb1bb358ef15fe56f35f58d164b7c62
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/utils/plugin_loader.py
@@ -0,0 +1,69 @@
+#!/usr/bin/python
+# -*- coding: UTF-8 -*-
+
+
+
+""" Plugin loader for extract, training and model tasks """
+
+from utils import logger
+import os
+from importlib import import_module
+from typing import Type
+from trainer._base import TrainerBase
+from torch.utils.data import Dataset
+from model._base import ModelBase
+
+class PluginLoader():
+ """
+ Plugin loader for extract, training and model tasks
+ function: get_{model_type}
+ args: {model_name}
+ will return the a class named model_type under model_type.model_name.py
+
+ as it return a class you should also annotate the returning classtype to make
+ code linting avaliable in some IDE
+ """
+ @staticmethod
+ def get_classifier(name) -> Type[ModelBase]:
+ """ Return requested attribute encoder plugin """
+ return PluginLoader._import("model.classifier", name)
+
+ @staticmethod
+ def get_trainer(name) -> Type[TrainerBase]:
+ """ Return requested trainer plugin """
+ return PluginLoader._import("trainer", name)
+
+ @staticmethod
+ def get_dataset(name) -> Type[Dataset]:
+ """ Return requested trainer plugin """
+ return PluginLoader._import("dataset", name)
+
+ @staticmethod
+ def _import(attr, name):
+ """ Import the plugin's module """
+ name = name.replace("-", "_")
+ ttl = attr.split(".")[-1].title()
+ logger.info("Loading %s from %s plugin...", ttl, name.title())
+ attr = "model" if attr == "Trainer" else attr.lower()
+ mod = ".".join((attr, name))
+ module = import_module(mod)
+ logger.info(str(module) + str(ttl))
+ return getattr(module, ttl)
+
+ @staticmethod
+ def get_available_trainer():
+ """ Return a list of available models """
+ modelpath = os.path.join(os.path.dirname(__file__), "trainer")
+ models = sorted(item.name.replace(".py", "").replace("_", "-")
+ for item in os.scandir(modelpath)
+ if not item.name.startswith("_")
+ and item.name.endswith(".py"))
+ return models
+
+ @staticmethod
+ def get_default_model():
+ """ Return the default model """
+ models = PluginLoader.get_available_models()
+ return 'original' if 'original' in models else models[0]
+
+
diff --git a/clean/video/pwtf_dvd/inference/utils/torch_save.py b/clean/video/pwtf_dvd/inference/utils/torch_save.py
new file mode 100644
index 0000000000000000000000000000000000000000..88fce6ba32ac66b3607a8496a03ae6e8f6321866
--- /dev/null
+++ b/clean/video/pwtf_dvd/inference/utils/torch_save.py
@@ -0,0 +1,9 @@
+import torch
+
+def torch_save(arr,file):
+ if torch.__version__>="1.6.0":
+ torch.save(arr, file, _use_new_zipfile_serialization=False)
+ else:
+ torch.save(arr, file)
+
+
diff --git a/clean/video/pwtf_dvd/preprocessing/preprocess.py b/clean/video/pwtf_dvd/preprocessing/preprocess.py
new file mode 100644
index 0000000000000000000000000000000000000000..f0ffc0860ecc92f717b459d3e299ec6313848db6
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/preprocess.py
@@ -0,0 +1,258 @@
+import os
+from os.path import join
+import argparse
+import glob
+import subprocess
+import cv2
+from tqdm import tqdm
+import numpy as np
+import logging
+import torch
+from test_tools.common import detect_all, grab_all_frames
+from test_tools.faster_crop_align_xray import FasterCropAlignXRay
+from test_tools.warp_for_xray import (
+ estimiate_batch_transform,
+ transform_landmarks,
+ std_points_256,
+)
+from test_tools.ct.operations import find_longest, multiple_tracking
+from test_tools.utils import get_crop_box
+import datetime
+# from FaceForensics.face_detection_save import get_boundingbox
+
+os.environ['CUDA_LAUNCH_BLOCKING'] = "1"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+device=torch.device('cuda')
+#Date
+now = datetime.datetime.now()
+logger = logging.getLogger("main") #Logger 선언
+stream_handler = logging.StreamHandler() # Logger output 방법 선언
+formatter = logging.Formatter('[%(asctime)s][%(levelname)s|%(filename)s:%(lineno)s] >> %(message)s')
+stream_handler.setFormatter(formatter)
+logger.addHandler(stream_handler)
+logger.setLevel(logging.DEBUG)
+
+crop_align_func = FasterCropAlignXRay(256)
+max_frame= 10000
+
+### video path를 받으면 crop된 face를 저장하는 함수 ###
+def crop_face_from_video(video_path,cache_path,crop_path,clip_size):
+ # mp4 파일이 아니면 return
+ if 'mp4' not in video_path : return
+ # video name
+ video_name = video_path.split('/')[-1].replace('.mp4','')
+ # 만약 crop image path에 crop된 이미지가 110개 이상이면 return
+
+ if os.path.exists(crop_path):
+ if len(os.listdir(crop_path))>clip_limit:
+ logger.info(f'{video_name} already exists')
+ return
+
+ ##########################################
+ # detect_res : list, 전체 frame, whole frame
+ # detect_res [] : list, len = 사람 수로 예상 the number of detected face in a frame
+ # detect_res [] [] : tuple, length = 3
+ # detect_res [] [] 의 각 요소는 각각 box, lm5 : landmark (5,2) , score
+ ##########################################
+ # all_lm68 : list, 전체 frame, whole frame
+ # all_lm68 : list, len = 사람 수로 예상, the number of detected face in a frame
+ # all_lm68 : np.array : landmark 68개 (68,2)
+ ##########################################
+ # frames : each frame's np.array
+
+
+ # cache_file : cache file path
+ # landmark와 box를 저장하는 cache file
+ cache_file = f"{cache_path}.pth"
+
+ if os.path.exists(cache_file):
+ # cache file이 존재하면 load하고 frame만 불러옴
+ detect_res, all_lm68 = torch.load(cache_file)
+ frames = grab_all_frames(video_path, max_size=max_frame, cvt=True)
+ logger.info("detection result loaded from cache")
+ else:
+ # cache file이 존재하지 않으면 detect_all 함수를 통해 detect_res, all_lm68, frames를 불러옴
+ # detection_all 함수는 retina_face를 이용해서 box와 landmark를 찾는 함수
+ detect_res, all_lm68, frames = detect_all(
+ video_path, return_frames=True, max_size=10000
+ )
+ torch.save((detect_res, all_lm68), cache_file)
+ try:
+ shape = frames[0].shape[:2]
+ except IndexError: # if there is no frame in the video, error list에 저장
+ f = open("./indexerror.txt", 'a')
+ f.write("{}\n".format(video_path))
+ f.close()
+ return
+
+ # 모든 detect_res
+ all_detect_res = []
+
+ assert len(all_lm68) == len(detect_res)
+ # in each frame, save the detected face's bounding box, landmark(5, 68), score as a tuple and save it in a list
+ for faces, faces_lm68 in zip(detect_res, all_lm68):
+ new_faces = []
+ for (box, lm5, score), face_lm68 in zip(faces, faces_lm68):
+ new_face = (box, lm5, face_lm68, score)
+ new_faces.append(new_face)
+ all_detect_res.append(new_faces)
+ detect_res = all_detect_res
+ # SORT tracking
+ # tracks : list, len = 사람 수로 예상, the number of detected face in a frame
+ # tracks [] : list, len = 프레임 수, the number of frames
+ # tracks [] [] : tuple, length = 4, 각각 box, lm5 : landmark (5,2) , lm68 : landmark (68,2), score
+ tracks = multiple_tracking(detect_res)
+ # tuples : list, len = 사람 수로 예상, the number of detected face in a frame
+ # tuples [] : tuple, length = 2, 각각 0, 프레임 수 the number of frames
+ tuples = [(0, len(detect_res))] * len(tracks)
+ # if there is no face detected, find the longest face in the video
+ if len(tracks) == 0:
+ tuples, tracks = find_longest(detect_res)
+ data_storage = {}
+ frame_boxes = {}
+ super_clips = []
+ frame_res = {}
+ super_clips_start_end = []
+ # super_clips : tracking된 face들을 의미하는 것으로 보임
+ for track_i, ((start, end), track) in enumerate(zip(tuples, tracks)): # each track(=face)
+
+ # if detect_res's length is not equal to track's length, raise error
+ assert len(detect_res[start:end]) == len(track)
+
+ super_clips.append(len(track))
+ super_clips_start_end.append((start, end))
+ for face, frame_idx, j in zip(track, range(start, end), range(len(track))): # frame에서 각각의 face
+ box,lm5,lm68 = face[:3] # box, lm5, lm68
+ big_box = get_crop_box(shape, box, scale=0.5) # get crop box
+
+ top_left = big_box[:2][None, :] # top left point
+
+ new_lm5 = lm5 - top_left
+ new_lm68 = lm68 - top_left
+
+ new_box = (box.reshape(2, 2) - top_left).reshape(-1)
+
+ info = (new_box, new_lm5, new_lm68, big_box) # face info
+
+
+ x1, y1, x2, y2 = big_box
+ cropped = frames[frame_idx][y1:y2, x1:x2]
+ # cropped = cv2.resize(cropped, (512, 512))
+ # face들을 tracking한 박스들로 crop함
+ # landmark들도 box에 맞게 변환
+ # data_storage에 저장 i는 face id, j는 frame을 의미
+ base_key = f"{track_i}_{j}_" # i : face, j : frame
+ data_storage[base_key + "img"] = cropped
+ data_storage[base_key + "ldm"] = info
+ data_storage[base_key + "idx"] = frame_idx
+ frame_boxes[frame_idx] = np.rint(box).astype(np.int64)
+ # 총 crop된 face들과 그 face들의 frame 수를 알려줌
+ logger.info(f"{crop_path} : sampling clips from super clips {super_clips}")
+ clips_for_video = []
+ clip_size = clip_size
+ pad_length = clip_size - 1
+
+ # 각 face id 별로 clip을 만듦
+ # 아래의 영어 표기로는 8clip을 의미하지만 정확하겐 clip size 만큼 함
+ for super_clip_idx, super_clip_size in enumerate(super_clips): # cut the super clip into clips, overlap 7frames, 8frames per clip
+ inner_index = list(range(super_clip_size))
+
+ if super_clip_size < clip_size: # if there is not enough frames to make a clip, pad the frames
+ # to do : how to operate the padding
+ # 정확하게 이 코드가 어떻게 동작하는지 모르겠지만
+ # 대략적으로 frame들을 clipsize로 나눌때 부족하면
+ # clip size만큼의 frame이 되도록 padding을 함
+ if super_clip_size < clip_size//2 : continue
+ post_module = inner_index[1:-1][::-1] + inner_index
+
+ l_post = len(post_module)
+ post_module = post_module * (pad_length // l_post + 1)
+ post_module = post_module[:pad_length]
+ assert len(post_module) == pad_length
+
+ pre_module = inner_index + inner_index[1:-1][::-1]
+ l_pre = len(post_module)
+ pre_module = pre_module * (pad_length // l_pre + 1)
+ pre_module = pre_module[-pad_length:]
+ assert len(pre_module) == pad_length
+
+ inner_index = pre_module + inner_index + post_module
+
+ super_clip_size = len(inner_index)
+
+ frame_range = [
+ inner_index[i : i + clip_size] for i in range(super_clip_size) if i + clip_size <= super_clip_size
+ ]
+ for indices in frame_range:
+ clip = [(super_clip_idx, t) for t in indices]
+ clips_for_video.append(clip)
+
+ # landmarks, images = crop_align_func(landmarks, images) # i : face, j : frame
+ processed_clips = 0 # Track number of processed clips
+ for clip in clips_for_video:
+ # Check if we've reached the clip limit
+ if processed_clips >= clip_limit:
+ logger.info(f"Reached clip limit of {clip_limit}, stopping processing")
+ break
+
+ # 각 자른 clip에 대해서 진행
+ images = [data_storage[f"{i}_{j}_img"] for i, j in clip] # call cropped face images from data_storage, i : face, j : frame
+ landmarks = [data_storage[f"{i}_{j}_ldm"] for i, j in clip] # call landmarks from data_storage, i : face, j : frame
+ # landmark를 기준으로 crop align func을 진행
+ # 해당 함수가 clip에 있는 얼굴들의 landmark 평균을 기준으로 박스를 설정하고
+ # 박스를 기준으로 crop align을 진행
+ # 다르게 말하면, landmark 평균을 기준으로 박스의 geometry를 설정하고
+ # 박스의 geometry는 고정한체로 얼굴이 움직이는 걸 찍었다 생각하면 됨
+ # 다시 또 말하면, 카메라를 고정하고 사람이 움직이는 것을 찍은것처럼
+ # PPT 참조
+ landmarks, images = crop_align_func(landmarks, images) # align the face images by landmarks in the clip
+ i, j = clip[-1]
+ k = super_clips[i]%clip_size
+
+ ##########################################################################
+ # 코드 변경시 이 함수에서는 이부분만 변경할 것을 권고 !!!!!!!!!!!!!!!!!!!!!!!!
+ # 특히, cv2.imwrite함수만 변경할 것을 추천
+ ##########################################################################
+ if (j+1)%clip_size==0: # if last frame number of the clip is multiple of clip_size, save all images in the clip
+ # it means save face alignments in 8 frames in the video so that they don't overlap
+ for f, (i,j) in enumerate(clip) :
+ cv2.imwrite(join(crop_path, f'{i:02}_{j:04}.png'), cv2.cvtColor(images[f], cv2.COLOR_BGR2RGB))
+ if j == super_clips[i]-1: # if the clip have last frame image, save all images in the clip
+ if k!=0 : # if the clip is not multiple of clip_size, save the last k images in the clip
+ # k is the number of frames that are not overlapped
+ for l in range(clip_size-k,clip_size):
+ ci,cj = clip[l]
+ cv2.imwrite(join(crop_path, f'{ci:02}_{cj:04}.png'), cv2.cvtColor(images[l], cv2.COLOR_BGR2RGB)) # clip means face alignment images in 8 frames, non-overlap
+
+ processed_clips += 1 # Increment processed clip counter
+ ##########################################################################
+
+
+if __name__ == '__main__':
+ p = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+ p.add_argument('--video_path','-i', type=str, default='/videos.mp4', help='path to input video')
+ p.add_argument('--save_path','-s', type=str, default='/data/crop_face', help='path to save cropped faces')
+ p.add_argument('--cachepath', '-c', type=str, default='/data/cache', help='path to cache detection results')
+ p.add_argument('--clipsize','-l',type=int,default=32, help='number of frames in a clip')
+ args = p.parse_args()
+ video_path = args.video_path
+ save_path = args.save_path
+ cache_path = args.cachepath
+ clip_size = args.clipsize
+
+ crop_face_from_video(video_path, cache_path, crop_path, clip_size)
+
+
+
+
+###################### reference ######################
+
+# this code reference from FTCN Official Code in git hub
+# link is https://github.com/yinglinzheng/FTCN
+
+# - Zheng, Y., Bao, J., Chen, D., Zeng, M., & Wen, F. (2021). Exploring Temporal Coherence for More General Video Face Forgery Detection. In Proceedings of the IEEE/CVF International Conference on Computer Vision (pp. 15044–15054).
+
+
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/__init__.py b/clean/video/pwtf_dvd/preprocessing/test_tools/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/common.py b/clean/video/pwtf_dvd/preprocessing/test_tools/common.py
new file mode 100644
index 0000000000000000000000000000000000000000..e76ac698b8fe9a32dfbdf4e294d5afa521250e71
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/test_tools/common.py
@@ -0,0 +1,122 @@
+import os
+
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+
+from .ct.detection.utils import grab_all_frames, get_valid_faces, sample_chunks
+from .ct.operations import multiple_tracking
+import numpy as np
+from .ct.face_alignment import LandmarkPredictor
+from .ct.detection import FaceDetector
+import cv2
+from .utils import flatten,partition
+
+
+detector = FaceDetector(0)
+predictor = LandmarkPredictor(0)
+
+
+def get_five(ldm68):
+ groups = [range(36, 42), range(42, 48), [30], [48], [54]]
+ points = []
+ for group in groups:
+ points.append(ldm68[group].mean(0))
+ return np.array(points)
+
+
+def get_bbox(mask):
+ try:
+ y, x = np.nonzero(mask[..., 0])
+ return x.min() - 1, y.min() - 1, x.max() + 1, y.max() + 1
+ except:
+ return None
+
+
+def get_bigger_box(image, box, scale=0.5):
+ height, width = image.shape[:2]
+ box = np.rint(box).astype(np.int)
+ new_box = box.reshape(2, 2)
+ size = new_box[1] - new_box[0]
+ diff = scale * size
+ diff = diff[None, :] * np.array([-1, 1])[:, None]
+ new_box = new_box + diff
+ new_box[:, 0] = np.clip(new_box[:, 0], 0, width - 1)
+ new_box[:, 1] = np.clip(new_box[:, 1], 0, height - 1)
+ new_box = np.rint(new_box).astype(np.int)
+ return new_box.reshape(-1)
+
+
+def process_bigger_clips(clips, dete_res, clip_size, step, scale=0.5):
+ assert len(clips) % clip_size == 0
+ detect_results = sample_chunks(dete_res, clip_size, step)
+ clips = sample_chunks(clips, clip_size, step)
+ new_clips = []
+ for i, (frame_clip, record_clip) in enumerate(zip(clips, detect_results)):
+ tracks = multiple_tracking(record_clip)
+ for j, track in enumerate(tracks):
+ new_images = []
+ for (box, ldm, _), frame in zip(track, frame_clip):
+ big_box = get_bigger_box(frame, box, scale)
+ x1, y1, x2, y2 = big_box
+ top_left = big_box[:2][None, :]
+ new_ldm5 = ldm - top_left
+ box = np.rint(box).astype(np.int)
+ new_box = (box.reshape(2, 2) - top_left).reshape(-1)
+ feed = LandmarkPredictor.prepare_feed(frame, box)
+ ldm68 = predictor(feed) - top_left
+ new_images.append(
+ (frame[y1:y2, x1:x2], big_box, new_box, new_ldm5, ldm68)
+ )
+ new_clips.append(new_images)
+ return new_clips
+
+
+def post(detected_faces):
+ return [[face[:4], None, face[-1]] for face in detected_faces]
+
+
+def check(detect_res):
+ return min([len(faces) for faces in detect_res]) != 0
+
+
+def detect_all(file, sfd_only=False, return_frames=False, max_size=None):
+ frames = grab_all_frames(file, max_size=max_size, cvt=True)
+ if not sfd_only:
+ detect_res = flatten(
+ [detector.detect(item) for item in partition(frames, 50)]
+ )
+ detect_res = get_valid_faces(detect_res, thres=0.5)
+ else:
+ raise NotImplementedError
+
+ all_68 = get_lm68(frames, detect_res)
+ if not return_frames:
+ return detect_res, all_68
+ else:
+ return detect_res, all_68, frames
+
+
+def get_lm68(frames, detect_res):
+ assert len(frames) == len(detect_res)
+ frame_count = len(frames)
+ all_68 = []
+ for i in range(frame_count):
+ frame = frames[i]
+ faces = detect_res[i]
+ if len(faces) == 0:
+ res_68 = []
+ else:
+ feeds = []
+ for face in faces:
+ assert len(face) == 3
+ box = face[0]
+ feed = LandmarkPredictor.prepare_feed(frame, box)
+ feeds.append(feed)
+ res_68 = predictor(feeds)
+ assert len(res_68) == len(faces)
+ for face, l_68 in zip(faces, res_68):
+ if face[1] is None:
+ face[1] = get_five(l_68)
+ all_68.append(res_68)
+
+ assert len(all_68) == len(detect_res)
+ return all_68
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/ct/detection/__init__.py b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/detection/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..23447e84f4f5f23a4c4818df2f26cc19a9200363
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/detection/__init__.py
@@ -0,0 +1,56 @@
+import cv2
+from .detector import RetinaFace
+from .utils import *
+
+
+def assert_bounded(val, low, up):
+ return val >= low and val < up
+
+
+def check_valid(face, w, h):
+ box = face[0]
+ if box[0] > box[2]:
+ return False
+ if box[1] > box[3]:
+ return False
+ for idx, bound in zip([0, 1, 2, 3], [w, h, w, h]):
+ if not assert_bounded(box[idx], 0, bound):
+ return False
+ pts = face[1]
+ for p in pts:
+ for idx, bound in zip([0, 1], [w, h]):
+ if not assert_bounded(p[idx], 0, bound):
+ return False
+ return True
+
+
+def post_detect(detect_results, scale, w, h):
+ new_results = []
+ for frame_faces in detect_results:
+ new_frame_faces = []
+ for box, ldm, score in frame_faces:
+ box = box * scale
+ ldm = ldm * scale
+ face = (box, ldm, score)
+ if check_valid(face, w=w, h=h):
+ new_frame_faces.append(face)
+ new_results.append(new_frame_faces)
+ return new_results
+
+
+class FaceDetector(RetinaFace):
+ def scale_detect(self, images):
+ max_res = 1920
+ h, w = images[0].shape[:2]
+ if max(h, w) > max_res:
+ init_scale = max(h, w) / max_res
+ else:
+ init_scale = 1
+ resize_scale = 2 * init_scale
+ resize_w = int(w / resize_scale)
+ resize_h = int(h / resize_scale)
+ detect_input = [cv2.resize(frame, (resize_w, resize_h)) for frame in images]
+ detect_results = post_detect(
+ self.detect(detect_input), scale=resize_scale, w=w, h=h,
+ )
+ return detect_results
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/ct/detection/alignment.py b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/detection/alignment.py
new file mode 100644
index 0000000000000000000000000000000000000000..64692a3490c7e8e29e399bd174c5906716b3e6fc
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/detection/alignment.py
@@ -0,0 +1,608 @@
+from itertools import product as product
+from math import ceil
+
+import numpy as np
+import torch
+import torch.backends.cudnn as cudnn
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models._utils as _utils
+
+
+def conv_bn(inp, oup, stride=1, leaky=0):
+ return nn.Sequential(
+ nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
+ nn.BatchNorm2d(oup),
+ nn.LeakyReLU(negative_slope=leaky, inplace=True),
+ )
+
+
+def conv_bn_no_relu(inp, oup, stride):
+ return nn.Sequential(
+ nn.Conv2d(inp, oup, 3, stride, 1, bias=False), nn.BatchNorm2d(oup),
+ )
+
+
+def conv_bn1X1(inp, oup, stride, leaky=0):
+ return nn.Sequential(
+ nn.Conv2d(inp, oup, 1, stride, padding=0, bias=False),
+ nn.BatchNorm2d(oup),
+ nn.LeakyReLU(negative_slope=leaky, inplace=True),
+ )
+
+
+def conv_dw(inp, oup, stride, leaky=0.1):
+ return nn.Sequential(
+ nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
+ nn.BatchNorm2d(inp),
+ nn.LeakyReLU(negative_slope=leaky, inplace=True),
+ nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
+ nn.BatchNorm2d(oup),
+ nn.LeakyReLU(negative_slope=leaky, inplace=True),
+ )
+
+
+class SSH(nn.Module):
+ def __init__(self, in_channel, out_channel):
+ super(SSH, self).__init__()
+ assert out_channel % 4 == 0
+ leaky = 0
+ if out_channel <= 64:
+ leaky = 0.1
+ self.conv3X3 = conv_bn_no_relu(in_channel, out_channel // 2, stride=1)
+
+ self.conv5X5_1 = conv_bn(in_channel, out_channel // 4, stride=1, leaky=leaky)
+ self.conv5X5_2 = conv_bn_no_relu(out_channel // 4, out_channel // 4, stride=1)
+
+ self.conv7X7_2 = conv_bn(
+ out_channel // 4, out_channel // 4, stride=1, leaky=leaky
+ )
+ self.conv7x7_3 = conv_bn_no_relu(out_channel // 4, out_channel // 4, stride=1)
+
+ def forward(self, input):
+ conv3X3 = self.conv3X3(input)
+
+ conv5X5_1 = self.conv5X5_1(input)
+ conv5X5 = self.conv5X5_2(conv5X5_1)
+
+ conv7X7_2 = self.conv7X7_2(conv5X5_1)
+ conv7X7 = self.conv7x7_3(conv7X7_2)
+
+ out = torch.cat([conv3X3, conv5X5, conv7X7], dim=1)
+ out = F.relu(out)
+ return out
+
+
+class FPN(nn.Module):
+ def __init__(self, in_channels_list, out_channels):
+ super(FPN, self).__init__()
+ leaky = 0
+ if out_channels <= 64:
+ leaky = 0.1
+ self.output1 = conv_bn1X1(
+ in_channels_list[0], out_channels, stride=1, leaky=leaky
+ )
+ self.output2 = conv_bn1X1(
+ in_channels_list[1], out_channels, stride=1, leaky=leaky
+ )
+ self.output3 = conv_bn1X1(
+ in_channels_list[2], out_channels, stride=1, leaky=leaky
+ )
+
+ self.merge1 = conv_bn(out_channels, out_channels, leaky=leaky)
+ self.merge2 = conv_bn(out_channels, out_channels, leaky=leaky)
+
+ def forward(self, input):
+ # names = list(input.keys())
+ input = list(input.values())
+
+ output1 = self.output1(input[0])
+ output2 = self.output2(input[1])
+ output3 = self.output3(input[2])
+
+ up3 = F.interpolate(
+ output3, size=[output2.size(2), output2.size(3)], mode="nearest"
+ )
+ output2 = output2 + up3
+ output2 = self.merge2(output2)
+
+ up2 = F.interpolate(
+ output2, size=[output1.size(2), output1.size(3)], mode="nearest"
+ )
+ output1 = output1 + up2
+ output1 = self.merge1(output1)
+
+ out = [output1, output2, output3]
+ return out
+
+
+class MobileNetV1(nn.Module):
+ def __init__(self):
+ super(MobileNetV1, self).__init__()
+ self.stage1 = nn.Sequential(
+ conv_bn(3, 8, 2, leaky=0.1), # 3
+ conv_dw(8, 16, 1), # 7
+ conv_dw(16, 32, 2), # 11
+ conv_dw(32, 32, 1), # 19
+ conv_dw(32, 64, 2), # 27
+ conv_dw(64, 64, 1), # 43
+ )
+ self.stage2 = nn.Sequential(
+ conv_dw(64, 128, 2), # 43 + 16 = 59
+ conv_dw(128, 128, 1), # 59 + 32 = 91
+ conv_dw(128, 128, 1), # 91 + 32 = 123
+ conv_dw(128, 128, 1), # 123 + 32 = 155
+ conv_dw(128, 128, 1), # 155 + 32 = 187
+ conv_dw(128, 128, 1), # 187 + 32 = 219
+ )
+ self.stage3 = nn.Sequential(
+ conv_dw(128, 256, 2), # 219 +3 2 = 241
+ conv_dw(256, 256, 1), # 241 + 64 = 301
+ )
+ self.avg = nn.AdaptiveAvgPool2d((1, 1))
+ self.fc = nn.Linear(256, 1000)
+
+ def forward(self, x):
+ x = self.stage1(x)
+ x = self.stage2(x)
+ x = self.stage3(x)
+ x = self.avg(x)
+ # x = self.model(x)
+ x = x.view(-1, 256)
+ x = self.fc(x)
+ return x
+
+
+class ClassHead(nn.Module):
+ def __init__(self, inchannels=512, num_anchors=3):
+ super(ClassHead, self).__init__()
+ self.num_anchors = num_anchors
+ self.conv1x1 = nn.Conv2d(
+ inchannels, self.num_anchors * 2, kernel_size=(1, 1), stride=1, padding=0
+ )
+
+ def forward(self, x):
+ out = self.conv1x1(x)
+ out = out.permute(0, 2, 3, 1).contiguous()
+
+ return out.view(out.shape[0], -1, 2)
+
+
+class BboxHead(nn.Module):
+ def __init__(self, inchannels=512, num_anchors=3):
+ super(BboxHead, self).__init__()
+ self.conv1x1 = nn.Conv2d(
+ inchannels, num_anchors * 4, kernel_size=(1, 1), stride=1, padding=0
+ )
+
+ def forward(self, x):
+ out = self.conv1x1(x)
+ out = out.permute(0, 2, 3, 1).contiguous()
+
+ return out.view(out.shape[0], -1, 4)
+
+
+class LandmarkHead(nn.Module):
+ def __init__(self, inchannels=512, num_anchors=3):
+ super(LandmarkHead, self).__init__()
+ self.conv1x1 = nn.Conv2d(
+ inchannels, num_anchors * 10, kernel_size=(1, 1), stride=1, padding=0
+ )
+
+ def forward(self, x):
+ out = self.conv1x1(x)
+ out = out.permute(0, 2, 3, 1).contiguous()
+
+ return out.view(out.shape[0], -1, 10)
+
+
+class RetinaFace(nn.Module):
+ def __init__(self, cfg=None, phase="train"):
+ """
+ :param cfg: Network related settings.
+ :param phase: train or test.
+ """
+ super(RetinaFace, self).__init__()
+ self.phase = phase
+ backbone = None
+ if cfg["name"] == "mobilenet0.25":
+ backbone = MobileNetV1()
+ elif cfg["name"] == "Resnet50":
+ import torchvision.models as models
+
+ backbone = models.resnet50(pretrained=cfg["pretrain"])
+
+ self.body = _utils.IntermediateLayerGetter(backbone, cfg["return_layers"])
+ in_channels_stage2 = cfg["in_channel"]
+ in_channels_list = [
+ in_channels_stage2 * 2,
+ in_channels_stage2 * 4,
+ in_channels_stage2 * 8,
+ ]
+ out_channels = cfg["out_channel"]
+ self.fpn = FPN(in_channels_list, out_channels)
+ self.ssh1 = SSH(out_channels, out_channels)
+ self.ssh2 = SSH(out_channels, out_channels)
+ self.ssh3 = SSH(out_channels, out_channels)
+
+ self.ClassHead = self._make_class_head(fpn_num=3, inchannels=cfg["out_channel"])
+ self.BboxHead = self._make_bbox_head(fpn_num=3, inchannels=cfg["out_channel"])
+ self.LandmarkHead = self._make_landmark_head(
+ fpn_num=3, inchannels=cfg["out_channel"]
+ )
+
+ def _make_class_head(self, fpn_num=3, inchannels=64, anchor_num=2):
+ classhead = nn.ModuleList()
+ for i in range(fpn_num):
+ classhead.append(ClassHead(inchannels, anchor_num))
+ return classhead
+
+ def _make_bbox_head(self, fpn_num=3, inchannels=64, anchor_num=2):
+ bboxhead = nn.ModuleList()
+ for i in range(fpn_num):
+ bboxhead.append(BboxHead(inchannels, anchor_num))
+ return bboxhead
+
+ def _make_landmark_head(self, fpn_num=3, inchannels=64, anchor_num=2):
+ landmarkhead = nn.ModuleList()
+ for i in range(fpn_num):
+ landmarkhead.append(LandmarkHead(inchannels, anchor_num))
+ return landmarkhead
+
+ def forward(self, inputs):
+ out = self.body(inputs)
+
+ # FPN
+ fpn = self.fpn(out)
+
+ # SSH
+ feature1 = self.ssh1(fpn[0])
+ feature2 = self.ssh2(fpn[1])
+ feature3 = self.ssh3(fpn[2])
+ features = [feature1, feature2, feature3]
+
+ bbox_regressions = torch.cat(
+ [self.BboxHead[i](feature) for i, feature in enumerate(features)], dim=1
+ )
+ classifications = torch.cat(
+ [self.ClassHead[i](feature) for i, feature in enumerate(features)], dim=1
+ )
+ ldm_regressions = torch.cat(
+ [self.LandmarkHead[i](feature) for i, feature in enumerate(features)], dim=1
+ )
+
+ if self.phase == "train":
+ output = (bbox_regressions, classifications, ldm_regressions)
+ else:
+ output = (
+ bbox_regressions,
+ F.softmax(classifications, dim=-1),
+ ldm_regressions,
+ )
+ return output
+
+
+# Adapted from https://github.com/Hakuyume/chainer-ssd
+def decode(loc, priors, variances):
+ boxes = torch.cat(
+ (
+ priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
+ priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1]),
+ ),
+ 1,
+ )
+ boxes[:, :2] -= boxes[:, 2:] / 2
+ boxes[:, 2:] += boxes[:, :2]
+ return boxes
+
+
+def decode_landm(pre, priors, variances):
+ landms = torch.cat(
+ (
+ priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],
+ priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:],
+ priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:],
+ priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:],
+ priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:],
+ ),
+ dim=1,
+ )
+ return landms
+
+
+def py_cpu_nms(dets, thresh):
+ """Pure Python NMS baseline."""
+ x1 = dets[:, 0]
+ y1 = dets[:, 1]
+ x2 = dets[:, 2]
+ y2 = dets[:, 3]
+ scores = dets[:, 4]
+
+ areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+ order = scores.argsort()[::-1]
+
+ keep = []
+ while order.size > 0:
+ i = order[0]
+ keep.append(i)
+ xx1 = np.maximum(x1[i], x1[order[1:]])
+ yy1 = np.maximum(y1[i], y1[order[1:]])
+ xx2 = np.minimum(x2[i], x2[order[1:]])
+ yy2 = np.minimum(y2[i], y2[order[1:]])
+
+ w = np.maximum(0.0, xx2 - xx1 + 1)
+ h = np.maximum(0.0, yy2 - yy1 + 1)
+ inter = w * h
+ ovr = inter / (areas[i] + areas[order[1:]] - inter)
+
+ inds = np.where(ovr <= thresh)[0]
+ order = order[inds + 1]
+
+ return keep
+
+
+class PriorBox(object):
+ def __init__(self, cfg, image_size=None, phase="train"):
+ super(PriorBox, self).__init__()
+ self.min_sizes = cfg["min_sizes"]
+ self.steps = cfg["steps"]
+ self.clip = cfg["clip"]
+ self.image_size = image_size
+ self.feature_maps = [
+ [ceil(self.image_size[0] / step), ceil(self.image_size[1] / step)]
+ for step in self.steps
+ ]
+ self.name = "s"
+
+ def forward(self):
+ anchors = []
+ for k, f in enumerate(self.feature_maps):
+ min_sizes = self.min_sizes[k]
+ for i, j in product(range(f[0]), range(f[1])):
+ for min_size in min_sizes:
+ s_kx = min_size / self.image_size[1]
+ s_ky = min_size / self.image_size[0]
+ dense_cx = [
+ x * self.steps[k] / self.image_size[1] for x in [j + 0.5]
+ ]
+ dense_cy = [
+ y * self.steps[k] / self.image_size[0] for y in [i + 0.5]
+ ]
+ for cy, cx in product(dense_cy, dense_cx):
+ anchors += [cx, cy, s_kx, s_ky]
+
+ # back to torch land
+ output = torch.Tensor(anchors).view(-1, 4)
+ if self.clip:
+ output.clamp_(max=1, min=0)
+ return output
+
+
+cfg_mnet = {
+ "name": "mobilenet0.25",
+ "min_sizes": [[16, 32], [64, 128], [256, 512]],
+ "steps": [8, 16, 32],
+ "variance": [0.1, 0.2],
+ "clip": False,
+ "loc_weight": 2.0,
+ "gpu_train": True,
+ "batch_size": 32,
+ "ngpu": 1,
+ "epoch": 250,
+ "decay1": 190,
+ "decay2": 220,
+ "image_size": 640,
+ "pretrain": True,
+ "return_layers": {"stage1": 1, "stage2": 2, "stage3": 3},
+ "in_channel": 32,
+ "out_channel": 64,
+}
+
+cfg_re50 = {
+ "name": "Resnet50",
+ "min_sizes": [[16, 32], [64, 128], [256, 512]],
+ "steps": [8, 16, 32],
+ "variance": [0.1, 0.2],
+ "clip": False,
+ "loc_weight": 2.0,
+ "gpu_train": True,
+ "batch_size": 24,
+ "ngpu": 4,
+ "epoch": 100,
+ "decay1": 70,
+ "decay2": 90,
+ "image_size": 840,
+ "pretrain": False,
+ "return_layers": {"layer2": 1, "layer3": 2, "layer4": 3},
+ "in_channel": 256,
+ "out_channel": 256,
+}
+
+
+def check_keys(model, pretrained_state_dict):
+ ckpt_keys = set(pretrained_state_dict.keys())
+ model_keys = set(model.state_dict().keys())
+ used_pretrained_keys = model_keys & ckpt_keys
+ assert len(used_pretrained_keys) > 0, "load NONE from pretrained checkpoint"
+ return True
+
+
+def remove_prefix(state_dict, prefix):
+ """ Old style model is stored with all names of parameters sharing common prefix 'module.' """
+ f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else x
+ return {f(key): value for key, value in state_dict.items()}
+
+
+def load_model(model, pretrained_path, load_to_cpu):
+ if load_to_cpu:
+ if pretrained_path is None:
+ url = "https://github.com/yinglinzheng/face_weights/releases/download/v1/mobilenet0.25_Final.pth"
+ pretrained_dict = torch.utils.model_zoo.load_url(url)
+ else:
+ pretrained_dict = torch.load(
+ pretrained_path, map_location=lambda storage, loc: storage
+ )
+ else:
+ device = torch.cuda.current_device()
+ pretrained_dict = torch.load(
+ pretrained_path, map_location=lambda storage, loc: storage.cuda(device)
+ )
+ if "state_dict" in pretrained_dict.keys():
+ pretrained_dict = remove_prefix(pretrained_dict["state_dict"], "module.")
+ else:
+ pretrained_dict = remove_prefix(pretrained_dict, "module.")
+ check_keys(model, pretrained_dict)
+ model.load_state_dict(pretrained_dict, strict=False)
+ return model
+
+
+def load_net(model_path, device, network="mobilenet"):
+ if network == "mobilenet":
+ cfg = cfg_mnet
+ elif network == "resnet50":
+ cfg = cfg_re50
+ # net and model
+ net = RetinaFace(cfg=cfg, phase="test")
+ net = load_model(net, model_path, True)
+ net.eval()
+ cudnn.benchmark = True
+ net = net.to(device)
+ return net
+
+
+def parse_det(det):
+ landmarks = det[5:].reshape(5, 2)
+ box = det[:4]
+ score = det[4]
+ return box, landmarks, score
+
+
+def post_process(
+ loc,
+ conf,
+ landms,
+ prior_data,
+ cfg,
+ scale,
+ scale1,
+ resize,
+ confidence_threshold,
+ top_k,
+ nms_threshold,
+ keep_top_k,
+):
+ boxes = decode(loc, prior_data, cfg["variance"])
+ boxes = boxes * scale / resize
+ boxes = boxes.cpu().numpy()
+ scores = conf.cpu().numpy()[:, 1]
+ landms_copy = decode_landm(landms, prior_data, cfg["variance"])
+
+ landms_copy = landms_copy * scale1 / resize
+ landms_copy = landms_copy.cpu().numpy()
+
+ # ignore low scores
+ inds = np.where(scores > confidence_threshold)[0]
+ boxes = boxes[inds]
+ landms_copy = landms_copy[inds]
+ scores = scores[inds]
+
+ # keep top-K before NMS
+ order = scores.argsort()[::-1][:top_k]
+ boxes = boxes[order]
+ landms_copy = landms_copy[order]
+ scores = scores[order]
+
+ # do NMS
+ dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
+ keep = py_cpu_nms(dets, nms_threshold)
+ # keep = nms(dets, args.nms_threshold,force_cpu=args.cpu)
+ dets = dets[keep, :]
+ landms_copy = landms_copy[keep]
+
+ # keep top-K faster NMS
+ dets = dets[:keep_top_k, :]
+ landms_copy = landms_copy[:keep_top_k, :]
+
+ dets = np.concatenate((dets, landms_copy), axis=1)
+ # show image
+ dets = sorted(dets, key=lambda x: x[4], reverse=True)
+ dets = [parse_det(x) for x in dets]
+
+ return dets
+
+
+def batch_detect(net, images, device, is_tensor=False, normalized=False):
+ with torch.no_grad():
+ confidence_threshold = 0.02
+ cfg = cfg_mnet
+ top_k = 5000
+ nms_threshold = 0.4
+ keep_top_k = 750
+ resize = 1
+ if not is_tensor:
+ try:
+ img = np.float32(images)
+ except ValueError:
+ raise NotImplementedError("Input images must of same size")
+ img = torch.from_numpy(img)
+ else:
+ img = images.float()
+ img = img.to(device)
+ mean = (
+ torch.as_tensor([104, 117, 123], dtype=img.dtype, device=img.device)
+ .unsqueeze(0)
+ .unsqueeze(0)
+ .unsqueeze(0)
+ )
+ img -= mean
+ img = img.permute(0, 3, 1, 2)
+ (batch_size, _, im_height, im_width,) = img.shape
+ scale = torch.as_tensor(
+ [im_width, im_height, im_width, im_height],
+ dtype=img.dtype,
+ device=img.device,
+ )
+ scale = scale.to(device)
+
+ loc, conf, landms = net(img) # forward pass
+
+ priorbox = PriorBox(cfg, image_size=(im_height, im_width))
+ priors = priorbox.forward()
+ prior_data = priors.to(device)
+ scale1 = torch.as_tensor(
+ [
+ img.shape[3],
+ img.shape[2],
+ img.shape[3],
+ img.shape[2],
+ img.shape[3],
+ img.shape[2],
+ img.shape[3],
+ img.shape[2],
+ img.shape[3],
+ img.shape[2],
+ ],
+ dtype=img.dtype,
+ device=img.device,
+ )
+ scale1 = scale1.to(device)
+
+ all_dets = [
+ post_process(
+ loc_i,
+ conf_i,
+ landms_i,
+ prior_data,
+ cfg,
+ scale,
+ scale1,
+ resize,
+ confidence_threshold,
+ top_k,
+ nms_threshold,
+ keep_top_k,
+ )
+ for loc_i, conf_i, landms_i in zip(loc, conf, landms)
+ ]
+
+ return all_dets
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/ct/detection/detector.py b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/detection/detector.py
new file mode 100644
index 0000000000000000000000000000000000000000..f38050e7a050eeb05323e7d710a63c2bf1a12455
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/detection/detector.py
@@ -0,0 +1,46 @@
+import os
+
+import numpy as np
+import torch
+
+from .alignment import load_net, batch_detect
+
+
+def get_project_dir():
+ current_path = os.path.abspath(os.path.join(__file__, "../"))
+ return current_path
+
+
+def relative(path):
+ path = os.path.join(get_project_dir(), path)
+ return os.path.abspath(path)
+
+
+class RetinaFace:
+ def __init__(
+ self, gpu_id=-1, model_path=None, network="mobilenet",
+ ):
+ self.gpu_id = gpu_id
+ self.device = (
+ torch.device("cpu") if gpu_id == -1 else torch.device("cuda", gpu_id)
+ )
+ self.model = load_net(model_path, self.device, network)
+
+ def detect(self, images):
+ if isinstance(images, np.ndarray):
+ if len(images.shape) == 3:
+ return batch_detect(self.model, [images], self.device)[0]
+ elif len(images.shape) == 4:
+ return batch_detect(self.model, images, self.device)
+ elif isinstance(images, list):
+ return batch_detect(self.model, np.array(images), self.device)
+ elif isinstance(images, torch.Tensor):
+ if len(images.shape) == 3:
+ return batch_detect(self.model, images.unsqueeze(0), self.device)[0]
+ elif len(images.shape) == 4:
+ return batch_detect(self.model, images, self.device)
+ else:
+ raise NotImplementedError()
+
+ def __call__(self, images):
+ return self.detect(images)
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/ct/detection/utils.py b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/detection/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..381fe3c717daeced48806cd273aecd971a16e84c
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/detection/utils.py
@@ -0,0 +1,147 @@
+import cv2
+# from test_tools.utils import flatten
+import numpy as np
+def flatten(l):
+ return [item for sublist in l for item in sublist]
+
+def chunks(l, n, step=None):
+ if step is None:
+ step = n
+ return [l[i : i + n] for i in range(0, len(l), step)]
+
+
+def sample_chunks(l, n, step=None):
+ return [l[i : i + n] for i in range(0, len(l), step) if i + n <= len(l)]
+
+
+def grab_all_frames(path, max_size, cvt=False):
+ capture = cv2.VideoCapture(path)
+ ret = True
+ frames = []
+ while ret:
+ ret, frame = capture.read()
+ if ret:
+ if cvt:
+ frame = frame[..., ::-1]
+ frames.append(frame)
+ if len(frames) == max_size:
+ break
+ capture.release()
+ return frames
+
+
+def get_clips_uniform(path, count, clip_size):
+ capture = cv2.VideoCapture(path)
+ n_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
+ max_clip_available = n_frames + 1 - clip_size
+ if count > max_clip_available:
+ count = max_clip_available
+ final_start = max_clip_available - 1
+ start_indices = np.linspace(0, final_start, count, endpoint=True, dtype=np.int)
+ all_clip_idx = [list(range(start, start + clip_size)) for start in start_indices]
+ valid = set(flatten(all_clip_idx))
+ max_idx = max(valid)
+
+ frames = {}
+ for idx in range(max_idx + 1):
+ # Get the next frame, but don't decode if we're not using it.
+ ret = capture.grab()
+ if not ret:
+ continue
+
+ if idx in valid:
+ ret, frame = capture.retrieve()
+ if not ret or frame is None:
+ continue
+ else:
+ # frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+ frames[idx] = frame
+
+ capture.release()
+ clips = []
+ for clip_idx in all_clip_idx:
+ clip = []
+ flag = True
+ for idx in clip_idx:
+ if idx not in frames:
+ flag = False
+ break
+ clip.append(frames[idx])
+ if flag:
+ clips.append(clip)
+ return clips
+
+
+def get_valid_faces(detect_results, max_count=10, thres=0.5, at_least=False):
+ new_results = []
+ for i, faces in enumerate(detect_results):
+ if len(faces) > max_count:
+ faces = faces[:max_count]
+ l = []
+ for j, face in enumerate(faces):
+ if face[-1] < thres and not (j == 0 and at_least):
+ continue
+ box, lm, score = face
+ box = box.astype(np.float64)
+ lm = lm.astype(np.float64)
+ l.append((box, lm, score))
+ new_results.append(l)
+ return new_results
+
+
+def scale_box(box, scale_h, scale_w, h, w):
+ x1, y1, x2, y2 = box.astype(np.int32)
+ center_x = (x1 + x2) // 2
+ center_y = (y1 + y2) // 2
+ box_h = int((y2 - y1) * scale_h)
+ box_w = int((x2 - x1) * scale_w)
+ new_x1 = center_x - box_w // 2
+ new_x2 = new_x1 + box_w
+ new_y1 = center_y - box_h // 2
+ new_y2 = new_y1 + box_h
+ new_x1 = max(new_x1, 0)
+ new_y1 = max(new_y1, 0)
+ new_y2 = min(new_y2, h)
+ new_x2 = min(new_x2, w)
+ return new_x1, new_y1, new_x2, new_y2
+
+
+def get_bbox(detect_res):
+ tmp_detect_res = get_valid_faces(detect_res, max_count=4, thres=0.5)
+ all_face_bboxs = []
+ for faces in tmp_detect_res:
+ all_face_bboxs.extend([face[0] for face in faces])
+ all_face_bboxs = np.array(all_face_bboxs).astype(np.int)
+ x1 = all_face_bboxs[:, 0].min()
+ x2 = all_face_bboxs[:, 2].max()
+ y1 = all_face_bboxs[:, 1].min()
+ y2 = all_face_bboxs[:, 3].max()
+
+ return x1, y1, x2, y2
+
+
+def delta_detect_res(detect_res, x1, y1):
+ diff = np.array([[x1, y1]])
+ new_detect_res = []
+ for faces in detect_res:
+ f = []
+ for face in faces:
+ box, lm, score = face
+ box = box.astype(np.float64)
+ box[[0, 2]] -= x1
+ box[[1, 3]] -= y1
+ lm = lm.astype(np.float64) - diff
+ f.append((box, lm, score))
+ new_detect_res.append(f)
+ return new_detect_res
+
+
+def pre_crop(clips, detect_res):
+ box = np.array(get_bbox(detect_res))
+ w = box[2] - box[0]
+ h = box[3] - box[1]
+ x1, y1, x2, y2 = scale_box(
+ box, 1.5, 1.2 if w > 2 * h else 1.5, clips[0].shape[0], clips[0].shape[1]
+ )
+ clips = np.array(clips)
+ return clips[:, y1:y2, x1:x2], delta_detect_res(detect_res, x1, y1)
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/ct/face_alignment/__init__.py b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/face_alignment/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..b0545f2f8c2aa872a127bf4e815c1e614a971723
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/face_alignment/__init__.py
@@ -0,0 +1 @@
+from .predictor import LandmarkPredictor
\ No newline at end of file
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/ct/face_alignment/basenet.py b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/face_alignment/basenet.py
new file mode 100644
index 0000000000000000000000000000000000000000..699b163e0584d56cc6d8f916b171f7ba86326efa
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/face_alignment/basenet.py
@@ -0,0 +1,107 @@
+# Backbone networks used for face landmark detection
+# Cunjian Chen (cunjian@msu.edu)
+
+import torch.nn as nn
+import torchvision.models as models
+
+
+class ConvBlock(nn.Module):
+ def __init__(self, inp, oup, k, s, p, dw=False, linear=False):
+ super(ConvBlock, self).__init__()
+ self.linear = linear
+ if dw:
+ self.conv = nn.Conv2d(inp, oup, k, s, p, groups=inp, bias=False)
+ else:
+ self.conv = nn.Conv2d(inp, oup, k, s, p, bias=False)
+ self.bn = nn.BatchNorm2d(oup)
+ if not linear:
+ self.prelu = nn.PReLU(oup)
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.bn(x)
+ if self.linear:
+ return x
+ else:
+ return self.prelu(x)
+
+
+# SE module
+# https://github.com/wujiyang/Face_Pytorch/blob/master/backbone/cbam.py
+class SEModule(nn.Module):
+ """Squeeze and Excitation Module"""
+
+ def __init__(self, channels, reduction):
+ super(SEModule, self).__init__()
+ self.avg_pool = nn.AdaptiveAvgPool2d(1)
+ self.fc1 = nn.Conv2d(
+ channels, channels // reduction, kernel_size=1, padding=0, bias=False
+ )
+ self.relu = nn.ReLU(inplace=True)
+ self.fc2 = nn.Conv2d(
+ channels // reduction, channels, kernel_size=1, padding=0, bias=False
+ )
+ self.sigmoid = nn.Sigmoid()
+
+ def forward(self, x):
+ input = x
+ x = self.avg_pool(x)
+ x = self.fc1(x)
+ x = self.relu(x)
+ x = self.fc2(x)
+ x = self.sigmoid(x)
+
+ return input * x
+
+
+# USE global depthwise convolution layer. Compatible with MobileNetV2 (224×224), MobileNetV2_ExternalData (224×224)
+class MobileNet_GDConv(nn.Module):
+ def __init__(self, num_classes):
+ super(MobileNet_GDConv, self).__init__()
+ self.pretrain_net = models.mobilenet_v2(pretrained=False)
+ self.base_net = nn.Sequential(*list(self.pretrain_net.children())[:-1])
+ self.linear7 = ConvBlock(1280, 1280, (7, 7), 1, 0, dw=True, linear=True)
+ self.linear1 = ConvBlock(1280, num_classes, 1, 1, 0, linear=True)
+
+ def forward(self, x):
+ x = self.base_net(x)
+ x = self.linear7(x)
+ x = self.linear1(x)
+ x = x.view(x.size(0), -1)
+ return x
+
+
+# USE global depthwise convolution layer. Compatible with MobileNetV2 (56×56)
+class MobileNet_GDConv_56(nn.Module):
+ def __init__(self, num_classes):
+ super(MobileNet_GDConv_56, self).__init__()
+ self.pretrain_net = models.mobilenet_v2(pretrained=False)
+ self.base_net = nn.Sequential(*list(self.pretrain_net.children())[:-1])
+ self.linear7 = ConvBlock(1280, 1280, (2, 2), 1, 0, dw=True, linear=True)
+ self.linear1 = ConvBlock(1280, num_classes, 1, 1, 0, linear=True)
+
+ def forward(self, x):
+ x = self.base_net(x)
+ x = self.linear7(x)
+ x = self.linear1(x)
+ x = x.view(x.size(0), -1)
+ return x
+
+
+# MobileNetV2 with SE; Compatible with MobileNetV2_SE (224×224) and MobileNetV2_SE_RE (224×224)
+class MobileNet_GDConv_SE(nn.Module):
+ def __init__(self, num_classes):
+ super(MobileNet_GDConv_SE, self).__init__()
+ self.pretrain_net = models.mobilenet_v2(pretrained=True)
+ self.base_net = nn.Sequential(*list(self.pretrain_net.children())[:-1])
+ self.linear7 = ConvBlock(1280, 1280, (7, 7), 1, 0, dw=True, linear=True)
+ self.linear1 = ConvBlock(1280, num_classes, 1, 1, 0, linear=True)
+ self.attention = SEModule(1280, 8)
+
+ def forward(self, x):
+ x = self.base_net(x)
+ x = self.attention(x)
+ x = self.linear7(x)
+ x = self.linear1(x)
+ x = x.view(x.size(0), -1)
+ return x
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/ct/face_alignment/predictor.py b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/face_alignment/predictor.py
new file mode 100644
index 0000000000000000000000000000000000000000..b210b673ab3af90d805e7fff9b0b66b00826fd1f
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/face_alignment/predictor.py
@@ -0,0 +1,143 @@
+# Face alignment demo
+# Uses MTCNN as face detector
+# Cunjian Chen (ccunjian@gmail.com)
+import torch
+import cv2
+import numpy as np
+from torch.utils.data import DataLoader
+from .basenet import MobileNet_GDConv
+
+
+def get_device(gpu_id):
+ if gpu_id > -1:
+ return torch.device(f"cuda:{str(gpu_id)}")
+ else:
+ return torch.device("cpu")
+
+
+def load_model(file):
+ model = MobileNet_GDConv(136)
+ if file is not None:
+ model.load_state_dict(torch.load(file, map_location="cpu"))
+ else:
+ url = "https://github.com/yinglinzheng/face_weights/releases/download/v1/mobilenet_224_model_best_gdconv_external.pth"
+ model.load_state_dict(torch.utils.model_zoo.load_url(url))
+ return model
+
+
+# landmark of (5L, 2L) from [0,1] to real range
+def reproject(bbox, landmark):
+ landmark_ = landmark.clone()
+ x1, y1, x2, y2 = bbox
+ w = x2 - x1
+ h = y2 - y1
+ landmark_[:, 0] *= w
+ landmark_[:, 0] += x1
+ landmark_[:, 1] *= h
+ landmark_[:, 1] += y1
+ return landmark_
+
+
+def prepare_feed(img, face):
+ height, width, _ = img.shape
+ mean = np.asarray([0.485, 0.456, 0.406])
+ std = np.asarray([0.229, 0.224, 0.225])
+ out_size = 224
+ x1, y1, x2, y2 = face[:4]
+
+ w = x2 - x1 + 1
+ h = y2 - y1 + 1
+ size = int(min([w, h]) * 1.2)
+ cx = x1 + w // 2
+ cy = y1 + h // 2
+ x1 = cx - size // 2
+ x2 = x1 + size
+ y1 = cy - size // 2
+ y2 = y1 + size
+
+ dx = max(0, -x1)
+ dy = max(0, -y1)
+ x1 = max(0, x1)
+ y1 = max(0, y1)
+
+ edx = max(0, x2 - width)
+ edy = max(0, y2 - height)
+ x2 = min(width, x2)
+ y2 = min(height, y2)
+ new_bbox = torch.Tensor([x1, y1, x2, y2]).int()
+ x1, y1, x2, y2 = new_bbox
+ cropped = img[y1:y2, x1:x2]
+ if dx > 0 or dy > 0 or edx > 0 or edy > 0:
+ cropped = cv2.copyMakeBorder(
+ cropped, int(dy), int(edy), int(dx), int(edx), cv2.BORDER_CONSTANT, 0
+ )
+ cropped_face = cv2.resize(cropped, (out_size, out_size))
+
+ if cropped_face.shape[0] <= 0 or cropped_face.shape[1] <= 0:
+ return None
+ test_face = cropped_face.copy()
+ test_face = test_face / 255.0
+ test_face = (test_face - mean) / std
+ test_face = test_face.transpose((2, 0, 1))
+ data = torch.from_numpy(test_face).float()
+ return dict(data=data, bbox=new_bbox)
+
+
+@torch.no_grad()
+def single_predict(model, feed, device):
+ landmark = model(feed["data"].unsqueeze(0).to(device)).cpu()
+ landmark = landmark.reshape(-1, 2)
+ landmark = reproject(feed["bbox"], landmark)
+ return landmark.numpy()
+
+
+@torch.no_grad()
+def batch_predict(model, feeds, device):
+ if not isinstance(feeds, list):
+ feeds = [feeds]
+ # loader = DataLoader(FeedDataset(feeds), batch_size=50, shuffle=False)
+ data = []
+ for feed in feeds:
+ data.append(feed["data"].unsqueeze(0))
+ data = torch.cat(data, 0).to(device)
+ results = []
+
+ landmarks = model(data).cpu()
+ for landmark, feed in zip(landmarks, feeds):
+ landmark = landmark.reshape(-1, 2)
+ landmark = reproject(feed["bbox"], landmark)
+ results.append(landmark.numpy())
+ return results
+
+
+@torch.no_grad()
+def batch_predict2(model, feeds, device, batch_size=None):
+ if not isinstance(feeds, list):
+ feeds = [feeds]
+ if batch_size is None:
+ batch_size = len(feeds)
+ loader = DataLoader(feeds, batch_size=len(feeds), shuffle=False)
+ results = []
+ for feed in loader:
+ landmarks = model(feed["data"].to(device)).cpu()
+ for landmark, bbox in zip(landmarks, feed["bbox"]):
+ landmark = landmark.reshape(-1, 2)
+ landmark = reproject(bbox, landmark)
+ results.append(landmark.numpy())
+ return results
+
+
+class LandmarkPredictor:
+ def __init__(self, gpu_id=0, file=None):
+ self.device = get_device(gpu_id)
+ self.model = load_model(file).to(self.device).eval()
+
+ def __call__(self, feeds):
+ results = batch_predict2(self.model, feeds, self.device)
+ if not isinstance(feeds, list):
+ results = results[0]
+ return results
+
+ @staticmethod
+ def prepare_feed(img, face):
+ return prepare_feed(img, face)
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/ct/face_alignment/utils.py b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/face_alignment/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..e871b155963b0d1ce6891589d32615c5fc97bb66
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/face_alignment/utils.py
@@ -0,0 +1,17 @@
+import cv2
+
+
+def drawLandmark_multiple(img, bbox, landmark):
+ """
+ Input:
+ - img: gray or RGB
+ - bbox: type of BBox
+ - landmark: reproject landmark of (5L, 2L)
+ Output:
+ - img marked with landmark and bbox
+ """
+ x1, y1, x2, y2 = bbox
+ cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2)
+ for x, y in landmark:
+ cv2.circle(img, (int(x), int(y)), 2, (0, 255, 0), -1)
+ return img
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/ct/operations.py b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/operations.py
new file mode 100644
index 0000000000000000000000000000000000000000..68375fcc4efd768abcfddf39521af21d574c4aac
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/operations.py
@@ -0,0 +1,79 @@
+import os
+
+import os
+import cv2
+import numpy as np
+from .tracking.sort import iou
+
+
+def face_iou(f1, f2):
+ return iou(f1[0], f2[0])
+
+
+def simple_tracking(batch_landmarks, index=0, thres=0.5):
+ track = []
+
+ for i, faces in enumerate(batch_landmarks):
+ if i == 0:
+ if len(faces) <= index or faces[index][-1] < 0.8:
+ return None
+ if index != 0:
+ for idx in range(index):
+ if face_iou(faces[idx], faces[index]) > thres:
+ return None
+ track.append(faces[index])
+ else:
+ last = track[i - 1]
+ if len(faces) == 0:
+ return None
+ sorted_faces = sorted(faces, key=lambda x: face_iou(x, last), reverse=True)
+ if face_iou(sorted_faces[0], last) < thres:
+ return None
+ track.append(sorted_faces[0])
+ return track
+
+
+def multiple_tracking(batch_landmarks):
+ tracks = []
+ for i in range(len(batch_landmarks[0])):
+ track = simple_tracking(batch_landmarks, index=i)
+ if track is None:
+ continue
+ tracks.append(track)
+ return tracks
+
+def find_longest(detect_res):
+ fc = len(detect_res)
+ tuples = []
+ start = 0
+ end = 0
+ previous_count = -1
+ all_tracks = []
+ # start 取得到，end 取不到
+ while start < (fc - 1):
+ for end in range(start + 2, fc + 1):
+ tracks = multiple_tracking(detect_res[start:end])
+ if (len(tracks) != previous_count and previous_count != -1) or len(
+ tracks
+ ) == 0:
+ break
+ previous_count = len(tracks)
+ if end - start > 2:
+ if end != fc:
+ un_reach_end = end - 1
+ else:
+ un_reach_end = end
+ sub_tracks = multiple_tracking(detect_res[start:un_reach_end])
+ if end == fc and len(sub_tracks) == 0:
+ un_reach_end = end - 1
+ sub_tracks = multiple_tracking(detect_res[start:un_reach_end])
+ if len(sub_tracks) > 0:
+ tpl = (start, un_reach_end)
+ tuples.append(tpl)
+ all_tracks.append(sub_tracks[0])
+ else:
+ raise NotImplementedError
+ previous_count = -1
+ end = un_reach_end
+ start = end
+ return tuples, all_tracks
\ No newline at end of file
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/ct/tracking/__init__.py b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/tracking/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/ct/tracking/sort.py b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/tracking/sort.py
new file mode 100644
index 0000000000000000000000000000000000000000..dc7b0838e7110a2d3521c6d4cfefbb40cc266c23
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/tracking/sort.py
@@ -0,0 +1,285 @@
+"""
+ SORT: A Simple, Online and Realtime Tracker
+ Copyright (C) 2016 Alex Bewley alex@dynamicdetection.com
+
+ This program is free software: you can redistribute it and/or modify
+ it under the terms of the GNU General Public License as published by
+ the Free Software Foundation, either version 3 of the License, or
+ (at your option) any later version.
+
+ This program is distributed in the hope that it will be useful,
+ but WITHOUT ANY WARRANTY; without even the implied warranty of
+ MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ GNU General Public License for more details.
+
+ You should have received a copy of the GNU General Public License
+ along with this program. If not, see .
+"""
+from __future__ import print_function
+import os.path
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from scipy.optimize import linear_sum_assignment
+import glob
+import time
+import argparse
+from filterpy.kalman import KalmanFilter
+
+
+def iou(bb_test, bb_gt):
+ """
+ Computes IUO between two bboxes in the form [x1,y1,x2,y2]
+ """
+ xx1 = np.maximum(bb_test[0], bb_gt[0])
+ yy1 = np.maximum(bb_test[1], bb_gt[1])
+ xx2 = np.minimum(bb_test[2], bb_gt[2])
+ yy2 = np.minimum(bb_test[3], bb_gt[3])
+ w = np.maximum(0.0, xx2 - xx1)
+ h = np.maximum(0.0, yy2 - yy1)
+ wh = w * h
+ o = wh / (
+ (bb_test[2] - bb_test[0]) * (bb_test[3] - bb_test[1])
+ + (bb_gt[2] - bb_gt[0]) * (bb_gt[3] - bb_gt[1])
+ - wh
+ )
+ return o
+
+
+def convert_bbox_to_z(bbox):
+ """
+ Takes a bounding box in the form [x1,y1,x2,y2] and returns z in the form
+ [x,y,s,r] where x,y is the centre of the box and s is the scale/area and r is
+ the aspect ratio
+ """
+ w = bbox[2] - bbox[0]
+ h = bbox[3] - bbox[1]
+ x = bbox[0] + w / 2.0
+ y = bbox[1] + h / 2.0
+ s = w * h # scale is just area
+ r = w / float(h)
+ return np.array([x, y, s, r]).reshape((4, 1))
+
+
+def convert_x_to_bbox(x, score=None):
+ """
+ Takes a bounding box in the centre form [x,y,s,r] and returns it in the form
+ [x1,y1,x2,y2] where x1,y1 is the top left and x2,y2 is the bottom right
+ """
+ w = np.sqrt(x[2] * x[3])
+ h = x[2] / w
+ if score == None:
+ return np.array(
+ [x[0] - w / 2.0, x[1] - h / 2.0, x[0] + w / 2.0, x[1] + h / 2.0]
+ ).reshape((1, 4))
+ else:
+ return np.array(
+ [x[0] - w / 2.0, x[1] - h / 2.0, x[0] + w / 2.0, x[1] + h / 2.0, score]
+ ).reshape((1, 5))
+
+
+class KalmanBoxTracker(object):
+ """
+ This class represents the internel state of individual tracked objects observed as bbox.
+ """
+
+ count = 0
+
+ def __init__(self, bbox):
+ """
+ Initialises a tracker using initial bounding box.
+ """
+ # define constant velocity model
+ self.kf = KalmanFilter(dim_x=7, dim_z=4)
+ self.kf.F = np.array(
+ [
+ [1, 0, 0, 0, 1, 0, 0],
+ [0, 1, 0, 0, 0, 1, 0],
+ [0, 0, 1, 0, 0, 0, 1],
+ [0, 0, 0, 1, 0, 0, 0],
+ [0, 0, 0, 0, 1, 0, 0],
+ [0, 0, 0, 0, 0, 1, 0],
+ [0, 0, 0, 0, 0, 0, 1],
+ ]
+ )
+ self.kf.H = np.array(
+ [
+ [1, 0, 0, 0, 0, 0, 0],
+ [0, 1, 0, 0, 0, 0, 0],
+ [0, 0, 1, 0, 0, 0, 0],
+ [0, 0, 0, 1, 0, 0, 0],
+ ]
+ )
+
+ self.kf.R[2:, 2:] *= 10.0
+ self.kf.P[
+ 4:, 4:
+ ] *= 1000.0 # give high uncertainty to the unobservable initial velocities
+ self.kf.P *= 10.0
+ self.kf.Q[-1, -1] *= 0.01
+ self.kf.Q[4:, 4:] *= 0.01
+
+ self.kf.x[:4] = convert_bbox_to_z(bbox)
+ self.time_since_update = 0
+ self.id = KalmanBoxTracker.count
+ KalmanBoxTracker.count += 1
+ self.history = []
+ self.hits = 0
+ self.hit_streak = 0
+ self.age = 0
+
+ def update(self, bbox):
+ """
+ Updates the state vector with observed bbox.
+ """
+ self.time_since_update = 0
+ self.history = []
+ self.hits += 1
+ self.hit_streak += 1
+ self.kf.update(convert_bbox_to_z(bbox))
+
+ def predict(self):
+ """
+ Advances the state vector and returns the predicted bounding box estimate.
+ """
+ if (self.kf.x[6] + self.kf.x[2]) <= 0:
+ self.kf.x[6] *= 0.0
+ self.kf.predict()
+ self.age += 1
+ if self.time_since_update > 0:
+ self.hit_streak = 0
+ self.time_since_update += 1
+ self.history.append(convert_x_to_bbox(self.kf.x))
+ return self.history[-1]
+
+ def get_state(self):
+ """
+ Returns the current bounding box estimate.
+ """
+ return convert_x_to_bbox(self.kf.x)
+
+
+def associate_detections_to_trackers(detections, trackers, iou_threshold=0.3):
+ """
+ Assigns detections to tracked object (both represented as bounding boxes)
+
+ Returns 3 lists of matches, unmatched_detections and unmatched_trackers
+ """
+ if len(trackers) == 0:
+ return (
+ np.empty((0, 2), dtype=int),
+ np.arange(len(detections)),
+ np.empty((0, 5), dtype=int),
+ )
+ iou_matrix = np.zeros((len(detections), len(trackers)), dtype=np.float32)
+
+ for d, det in enumerate(detections):
+ for t, trk in enumerate(trackers):
+ iou_matrix[d, t] = iou(det, trk)
+
+ matched_indices = linear_sum_assignment(-iou_matrix)
+ matched_indices = np.array(list(zip(*matched_indices)), dtype=np.int)
+ matched_indices.shape = (-1, 2)
+ # print(matched_indices)
+ # print(type(matched_indices))
+
+ unmatched_detections = []
+ for d, det in enumerate(detections):
+ if d not in matched_indices[:, 0]:
+ unmatched_detections.append(d)
+ unmatched_trackers = []
+ for t, trk in enumerate(trackers):
+ if t not in matched_indices[:, 1]:
+ unmatched_trackers.append(t)
+
+ # filter out matched with low IOU
+ matches = []
+ for m in matched_indices:
+ if iou_matrix[m[0], m[1]] < iou_threshold:
+ unmatched_detections.append(m[0])
+ unmatched_trackers.append(m[1])
+ else:
+ matches.append(m.reshape(1, 2))
+ if len(matches) == 0:
+ matches = np.empty((0, 2), dtype=int)
+ else:
+ matches = np.concatenate(matches, axis=0)
+
+ return matches, np.array(unmatched_detections), np.array(unmatched_trackers)
+
+
+class Sort(object):
+ def __init__(self, max_age=1, min_hits=3):
+ """
+ Sets key parameters for SORT
+ """
+ self.max_age = max_age
+ self.min_hits = min_hits
+ self.trackers = []
+ self.frame_count = 0
+
+ def update(self, dets):
+ """
+ Params:
+ dets - a numpy array of detections in the format [[x1,y1,x2,y2,score],[x1,y1,x2,y2,score],...]
+ Requires: this method must be called once for each frame even with empty detections.
+ Returns the a similar array, where the last column is the object ID.
+
+ NOTE: The number of objects returned may differ from the number of detections provided.
+ """
+ self.frame_count += 1
+ # get predicted locations from existing trackers.
+ trks = np.zeros((len(self.trackers), 5))
+ to_del = []
+ ret = []
+ for t, trk in enumerate(trks):
+ pos = self.trackers[t].predict()[0]
+ trk[:] = [pos[0], pos[1], pos[2], pos[3], 0]
+ if np.any(np.isnan(pos)):
+ to_del.append(t)
+ trks = np.ma.compress_rows(np.ma.masked_invalid(trks))
+ for t in reversed(to_del):
+ self.trackers.pop(t)
+ matched, unmatched_dets, unmatched_trks = associate_detections_to_trackers(
+ dets, trks
+ )
+
+ # update matched trackers with assigned detections
+ for t, trk in enumerate(self.trackers):
+ if t not in unmatched_trks:
+ d = matched[np.where(matched[:, 1] == t)[0], 0]
+ trk.update(dets[d, :][0])
+
+ # create and initialise new trackers for unmatched detections
+ for i in unmatched_dets:
+ trk = KalmanBoxTracker(dets[i, :])
+ self.trackers.append(trk)
+ i = len(self.trackers)
+ for trk in reversed(self.trackers):
+ d = trk.get_state()[0]
+ if (trk.time_since_update < 1) and (
+ trk.hit_streak >= self.min_hits or self.frame_count <= self.min_hits
+ ):
+ ret.append(
+ np.concatenate((d, [trk.id + 1])).reshape(1, -1)
+ ) # +1 as MOT benchmark requires positive
+ i -= 1
+ # remove dead tracklet
+ if trk.time_since_update > self.max_age:
+ self.trackers.pop(i)
+ if len(ret) > 0:
+ return np.concatenate(ret)
+ return np.empty((0, 5))
+
+
+def parse_args():
+ """Parse input arguments."""
+ parser = argparse.ArgumentParser(description="SORT demo")
+ parser.add_argument(
+ "--display",
+ dest="display",
+ help="Display online tracker output (slow) [False]",
+ action="store_true",
+ )
+ args = parser.parse_args()
+ return args
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/ct/tracking/tracker.py b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/tracking/tracker.py
new file mode 100644
index 0000000000000000000000000000000000000000..20dd79f41a56bcb4873e83bb9c76db9ffbf0f627
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/tracking/tracker.py
@@ -0,0 +1,27 @@
+from .sort import Sort
+import numpy as np
+
+
+def get_detections(faces):
+ detections = []
+ for face in faces:
+ x1, y1, x2, y2 = face[0]
+ detections.append((x1, y1, x2, y2, face[-1]))
+ return np.array(detections)
+
+
+def get_tracks(detect_results):
+ tracks = {}
+ mot_tracker = Sort()
+ for faces in detect_results:
+ detections = get_detections(faces)
+ track_bbs_ids = mot_tracker.update(detections)
+ for track in track_bbs_ids: # 单独框出每一张人脸
+ id = int(track[-1])
+ box = track[:4]
+ if id in tracks:
+ tracks[id].append(box)
+ else:
+ tracks[id] = [box]
+
+ return [track for id, track in tracks.items() if len(track) == len(detect_results)]
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/ct/utils.py b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..5ff180fb2a24b9e8bd673d3428e3b101da16a6b0
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/test_tools/ct/utils.py
@@ -0,0 +1,5 @@
+import cv2
+
+
+def write_img(file, img):
+ cv2.imwrite(file, img, [cv2.IMWRITE_PNG_COMPRESSION, 0])
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/faster_crop_align_xray.py b/clean/video/pwtf_dvd/preprocessing/test_tools/faster_crop_align_xray.py
new file mode 100644
index 0000000000000000000000000000000000000000..3e99f2f0930e26539793e3d76d74e9df7afc186b
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/test_tools/faster_crop_align_xray.py
@@ -0,0 +1,73 @@
+import numpy as np
+import cv2
+from .warp_for_xray import (
+ estimiate_batch_transform,
+ transform_landmarks,
+ std_points_256,
+)
+import numpy as np
+
+
+class FasterCropAlignXRay:
+ """
+ 修正到统一坐标系，统一图像大小到标准尺寸
+ """
+
+ def __init__(self, size=256):
+ self.image_size = size
+ self.std_points = std_points_256 * size / 256.0
+
+ def __call__(self, landmarks, images=None, jitter=False):
+ landmarks = [landmark[:4] for landmark in landmarks]
+ ori_boxes = np.array([ori_box for _, _, _, ori_box in landmarks])
+ five_landmarks = np.array([ldm5 for _, ldm5, _, _ in landmarks])
+ landmarks68 = np.array([ldm68 for _, _, ldm68, _ in landmarks])
+ # assert landmarks68.min() > 0
+
+ left_top = ori_boxes[:, :2].min(0)
+
+ right_bottom = ori_boxes[:, 2:].max(0)
+
+ size = right_bottom - left_top
+
+ w, h = size
+
+ diff = ori_boxes[:, :2] - left_top[None, ...]
+
+ new_five_landmarks = five_landmarks + diff[:, None, :]
+ new_landmarks68 = landmarks68 + diff[:, None, :]
+
+ landmark_for_estimiate = new_five_landmarks.copy()
+ if jitter:
+ landmark_for_estimiate += np.random.uniform(
+ -4, 4, landmark_for_estimiate.shape
+ )
+
+ tfm, trans = estimiate_batch_transform(
+ landmark_for_estimiate, tgt_pts=self.std_points
+ )
+
+ transformed_landmarks68 = np.array(
+ [transform_landmarks(ldm68, trans) for ldm68 in new_landmarks68]
+ )
+
+ if images is not None:
+ transformed_images = [
+ self.process_sinlge(tfm, image, d, h, w)
+ for image, d in zip(images, diff)
+ ] # 拼接 func 的参数
+ transformed_images = np.stack(transformed_images)
+ return transformed_landmarks68, transformed_images
+ else:
+ return transformed_landmarks68
+
+ def process_sinlge(self, tfm, image, d, h, w):
+ assert isinstance(image, np.ndarray)
+ new_image = np.zeros((h, w, 3), dtype=np.uint8)
+ x, y = d
+ ih, iw, _ = image.shape
+ new_image[y : y + ih, x : x + iw] = image
+ transformed_image = cv2.warpAffine(
+ new_image, tfm, (self.image_size, self.image_size)
+ )
+ return transformed_image
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/supply_writer.py b/clean/video/pwtf_dvd/preprocessing/test_tools/supply_writer.py
new file mode 100644
index 0000000000000000000000000000000000000000..0394dfbb8fc50d47984eaa937537de144e51ee81
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/test_tools/supply_writer.py
@@ -0,0 +1,49 @@
+import cv2
+
+class SupplyWriter:
+ def __init__(self, intput_video, output_video, opt_thres, rgb_input=True):
+ reader = cv2.VideoCapture(intput_video)
+ fourcc = cv2.VideoWriter_fourcc(*"XVID")
+ fps = reader.get(cv2.CAP_PROP_FPS)
+ width = int(reader.get(3))
+ height = int(reader.get(4))
+ reader.release()
+ self.padding = 40
+
+ self.writer = cv2.VideoWriter(output_video, fourcc, fps, (height, width)[::-1])
+ self.rgb_input = rgb_input
+ self.opt_thres = opt_thres
+
+ def run(self, images, scores, boxes):
+ # Text variables
+ font_face = cv2.FONT_HERSHEY_SIMPLEX
+ thickness = 5
+ font_scale = 3
+
+ for image, score, box in zip(images, scores, boxes):
+ if self.rgb_input:
+ image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+ if box is not None:
+ label = "fake" if score > self.opt_thres else "real"
+ x1, y1, x2, y2 = box
+ x = int(x1)
+ y = int(y1)
+ w = int(x2 - x1)
+ h = int(y2 - y1)
+ color = (
+ (255, 255, 0) if label == "real" else (0, 255, 255)
+ ) # BGR 255 0
+ cv2.putText(
+ image,
+ label,
+ (x, y + h + 68),
+ font_face,
+ font_scale,
+ color,
+ thickness,
+ 2,
+ )
+ # draw box over face
+ cv2.rectangle(image, (x, y), (x + w, y + h), color, 10)
+ self.writer.write(image)
+ self.writer.release()
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/utils.py b/clean/video/pwtf_dvd/preprocessing/test_tools/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..051f29c7c5808fd08aae14b6a44c54490986ce3a
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/test_tools/utils.py
@@ -0,0 +1,115 @@
+import numpy as np
+import cv2
+import os
+import platform
+import json
+import errno
+
+
+
+def weak_check(detect_res):
+ return sum([len(faces) for faces in detect_res]) > len(detect_res) * 0.75
+
+
+def get_crop_box(shape, box, scale=0.5):
+ height, width = shape
+ box = np.rint(box).astype(np.int64)
+ new_box = box.reshape(2, 2)
+ size = new_box[1] - new_box[0]
+ diff = scale * size
+ diff = diff[None, :] * np.array([-1, 1])[:, None]
+ new_box = new_box + diff
+ new_box[:, 0] = np.clip(new_box[:, 0], 0, width - 1)
+ new_box[:, 1] = np.clip(new_box[:, 1], 0, height - 1)
+ new_box = np.rint(new_box).astype(np.int64)
+ return new_box.reshape(-1)
+
+
+def get_fps(input_file):
+ reader = cv2.VideoCapture(input_file)
+ fps = reader.get(cv2.CAP_PROP_FPS)
+ reader.release()
+ return fps
+
+
+
+def mkdir_p(dirname):
+ """Like "mkdir -p", make a dir recursively, but do nothing if the dir exists
+ 这个是线程安全的, from Lingzhi Li
+ Args:
+ dirname(str):
+ """
+ assert dirname is not None
+ if dirname == "" or os.path.isdir(dirname):
+ return
+ try:
+ os.makedirs(dirname)
+ except OSError as e:
+ if e.errno != errno.EEXIST:
+ raise e
+
+
+def mkdir(*args):
+ for folder in args:
+ if not os.path.isdir(folder):
+ mkdir_p(folder)
+
+
+def make_join(*args):
+ folder = os.path.join(*args)
+ mkdir(folder)
+ return folder
+
+
+def list_dir(folder, condition=None, key=lambda x: x, reverse=False, co_join=[]):
+ files = os.listdir(folder)
+ if condition is not None:
+ files = filter(condition, files)
+ co_join = [folder] + co_join
+ if key is not None:
+ files = sorted(files, key=key, reverse=reverse)
+ files = [(file, *[os.path.join(fold, file) for fold in co_join]) for file in files]
+ return files
+
+def get_jointer(file):
+ def jointer(folder):
+ return os.path.join(folder, file)
+
+ return jointer
+
+def flatten(l):
+ return [item for sublist in l for item in sublist]
+
+
+def is_win():
+ return platform.system() == "Windows"
+
+
+def get_postfix(post_fix):
+ return lambda x: x.endswith(post_fix)
+
+
+def partition(images, size):
+ """
+ Returns a new list with elements
+ of which is a list of certain size.
+
+ >>> partition([1, 2, 3, 4], 3)
+ [[1, 2, 3], [4]]
+ """
+ return [
+ images[i : i + size] if i + size <= len(images) else images[i:]
+ for i in range(0, len(images), size)
+ ]
+
+
+def load_json(file):
+ with open(file, "r") as f:
+ res = json.load(f)
+ return res
+
+
+def save_json(file, obj):
+ with open(file, "w", encoding="utf-8") as f:
+ json.dump(obj, f, indent=4, ensure_ascii=False)
+
diff --git a/clean/video/pwtf_dvd/preprocessing/test_tools/warp_for_xray.py b/clean/video/pwtf_dvd/preprocessing/test_tools/warp_for_xray.py
new file mode 100644
index 0000000000000000000000000000000000000000..ea68def0bcc6b384f1198f21fb0d72253f8085d6
--- /dev/null
+++ b/clean/video/pwtf_dvd/preprocessing/test_tools/warp_for_xray.py
@@ -0,0 +1,574 @@
+import numpy as np
+import cv2
+
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Jul 11 06:54:28 2017
+
+@author: zhaoyafei
+"""
+
+import numpy as np
+from numpy.linalg import inv, norm, lstsq
+from numpy.linalg import matrix_rank as rank
+
+"""
+Introduction:
+----------
+numpy implemetation form matlab function CP2TFORM(...)
+with 'transformtype':
+ 1) 'nonreflective similarity'
+ 2) 'similarity'
+
+
+MATLAB code:
+----------
+%--------------------------------------
+% Function findNonreflectiveSimilarity
+%
+function [trans, output] = findNonreflectiveSimilarity(uv,xy,options)
+%
+% For a nonreflective similarity:
+%
+% let sc = s*cos(theta)
+% let ss = s*sin(theta)
+%
+% [ sc -ss
+% [u v] = [x y 1] * ss sc
+% tx ty]
+%
+% There are 4 unknowns: sc,ss,tx,ty.
+%
+% Another way to write this is:
+%
+% u = [x y 1 0] * [sc
+% ss
+% tx
+% ty]
+%
+% v = [y -x 0 1] * [sc
+% ss
+% tx
+% ty]
+%
+% With 2 or more correspondence points we can combine the u equations and
+% the v equations for one linear system to solve for sc,ss,tx,ty.
+%
+% [ u1 ] = [ x1 y1 1 0 ] * [sc]
+% [ u2 ] [ x2 y2 1 0 ] [ss]
+% [ ... ] [ ... ] [tx]
+% [ un ] [ xn yn 1 0 ] [ty]
+% [ v1 ] [ y1 -x1 0 1 ]
+% [ v2 ] [ y2 -x2 0 1 ]
+% [ ... ] [ ... ]
+% [ vn ] [ yn -xn 0 1 ]
+%
+% Or rewriting the above matrix equation:
+% U = X * r, where r = [sc ss tx ty]'
+% so r = X\ U.
+%
+
+K = options.K;
+M = size(xy,1);
+x = xy(:,1);
+y = xy(:,2);
+X = [x y ones(M,1) zeros(M,1);
+ y -x zeros(M,1) ones(M,1) ];
+
+u = uv(:,1);
+v = uv(:,2);
+U = [u; v];
+
+% We know that X * r = U
+if rank(X) >= 2*K
+ r = X \ U;
+else
+ error(message('images:cp2tform:twoUniquePointsReq'))
+end
+
+sc = r(1);
+ss = r(2);
+tx = r(3);
+ty = r(4);
+
+Tinv = [sc -ss 0;
+ ss sc 0;
+ tx ty 1];
+
+T = inv(Tinv);
+T(:,3) = [0 0 1]';
+
+trans = maketform('affine', T);
+output = [];
+
+%-------------------------
+% Function findSimilarity
+%
+function [trans, output] = findSimilarity(uv,xy,options)
+%
+% The similarities are a superset of the nonreflective similarities as they may
+% also include reflection.
+%
+% let sc = s*cos(theta)
+% let ss = s*sin(theta)
+%
+% [ sc -ss
+% [u v] = [x y 1] * ss sc
+% tx ty]
+%
+% OR
+%
+% [ sc ss
+% [u v] = [x y 1] * ss -sc
+% tx ty]
+%
+% Algorithm:
+% 1) Solve for trans1, a nonreflective similarity.
+% 2) Reflect the xy data across the Y-axis,
+% and solve for trans2r, also a nonreflective similarity.
+% 3) Transform trans2r to trans2, undoing the reflection done in step 2.
+% 4) Use TFORMFWD to transform uv using both trans1 and trans2,
+% and compare the results, Returnsing the transformation corresponding
+% to the smaller L2 norm.
+
+% Need to reset options.K to prepare for calls to findNonreflectiveSimilarity.
+% This is safe because we already checked that there are enough point pairs.
+options.K = 2;
+
+% Solve for trans1
+[trans1, output] = findNonreflectiveSimilarity(uv,xy,options);
+
+
+% Solve for trans2
+
+% manually reflect the xy data across the Y-axis
+xyR = xy;
+xyR(:,1) = -1*xyR(:,1);
+
+trans2r = findNonreflectiveSimilarity(uv,xyR,options);
+
+% manually reflect the tform to undo the reflection done on xyR
+TreflectY = [-1 0 0;
+ 0 1 0;
+ 0 0 1];
+trans2 = maketform('affine', trans2r.tdata.T * TreflectY);
+
+
+% Figure out if trans1 or trans2 is better
+xy1 = tformfwd(trans1,uv);
+norm1 = norm(xy1-xy);
+
+xy2 = tformfwd(trans2,uv);
+norm2 = norm(xy2-xy);
+
+if norm1 <= norm2
+ trans = trans1;
+else
+ trans = trans2;
+end
+"""
+
+
+class MatlabCp2tormException(Exception):
+ def __str__(self):
+ return "In File {}:{}".format(__file__, super.__str__(self))
+
+
+def tformfwd(trans, uv):
+ """
+ Function:
+ ----------
+ apply affine transform 'trans' to uv
+
+ Parameters:
+ ----------
+ @trans: 3x3 np.array
+ transform matrix
+ @uv: Kx2 np.array
+ each row is a pair of coordinates (x, y)
+
+ Returns:
+ ----------
+ @xy: Kx2 np.array
+ each row is a pair of transformed coordinates (x, y)
+ """
+ uv = np.hstack((uv, np.ones((uv.shape[0], 1))))
+ xy = np.dot(uv, trans)
+ xy = xy[:, 0:-1]
+ return xy
+
+
+def tforminv(trans, uv):
+ """
+ Function:
+ ----------
+ apply the inverse of affine transform 'trans' to uv
+
+ Parameters:
+ ----------
+ @trans: 3x3 np.array
+ transform matrix
+ @uv: Kx2 np.array
+ each row is a pair of coordinates (x, y)
+
+ Returns:
+ ----------
+ @xy: Kx2 np.array
+ each row is a pair of inverse-transformed coordinates (x, y)
+ """
+ Tinv = inv(trans)
+ xy = tformfwd(Tinv, uv)
+ return xy
+
+
+def findNonreflectiveSimilarity(uv, xy, options=None):
+ """
+ Function:
+ ----------
+ Find Non-reflective Similarity Transform Matrix 'trans':
+ u = uv[:, 0]
+ v = uv[:, 1]
+ x = xy[:, 0]
+ y = xy[:, 1]
+ [x, y, 1] = [u, v, 1] * trans
+
+ Parameters:
+ ----------
+ @uv: Kx2 np.array
+ source points each row is a pair of coordinates (x, y)
+ @xy: Kx2 np.array
+ each row is a pair of inverse-transformed
+ @option: not used, keep it as None
+
+ Returns:
+ @trans: 3x3 np.array
+ transform matrix from uv to xy
+ @trans_inv: 3x3 np.array
+ inverse of trans, transform matrix from xy to uv
+
+ Matlab:
+ ----------
+ % For a nonreflective similarity:
+ %
+ % let sc = s*cos(theta)
+ % let ss = s*sin(theta)
+ %
+ % [ sc -ss
+ % [u v] = [x y 1] * ss sc
+ % tx ty]
+ %
+ % There are 4 unknowns: sc,ss,tx,ty.
+ %
+ % Another way to write this is:
+ %
+ % u = [x y 1 0] * [sc
+ % ss
+ % tx
+ % ty]
+ %
+ % v = [y -x 0 1] * [sc
+ % ss
+ % tx
+ % ty]
+ %
+ % With 2 or more correspondence points we can combine the u equations and
+ % the v equations for one linear system to solve for sc,ss,tx,ty.
+ %
+ % [ u1 ] = [ x1 y1 1 0 ] * [sc]
+ % [ u2 ] [ x2 y2 1 0 ] [ss]
+ % [ ... ] [ ... ] [tx]
+ % [ un ] [ xn yn 1 0 ] [ty]
+ % [ v1 ] [ y1 -x1 0 1 ]
+ % [ v2 ] [ y2 -x2 0 1 ]
+ % [ ... ] [ ... ]
+ % [ vn ] [ yn -xn 0 1 ]
+ %
+ % Or rewriting the above matrix equation:
+ % U = X * r, where r = [sc ss tx ty]'
+ % so r = X\ U.
+ %
+ """
+ options = {"K": 2}
+
+ K = options["K"]
+ M = xy.shape[0]
+ x = xy[:, 0].reshape((-1, 1)) # use reshape to keep a column vector
+ y = xy[:, 1].reshape((-1, 1)) # use reshape to keep a column vector
+ # print '--->x, y:\n', x, y
+
+ tmp1 = np.hstack((x, y, np.ones((M, 1)), np.zeros((M, 1))))
+ tmp2 = np.hstack((y, -x, np.zeros((M, 1)), np.ones((M, 1))))
+ X = np.vstack((tmp1, tmp2))
+ # print '--->X.shape: ', X.shape
+ # print 'X:\n', X
+
+ u = uv[:, 0].reshape((-1, 1)) # use reshape to keep a column vector
+ v = uv[:, 1].reshape((-1, 1)) # use reshape to keep a column vector
+ U = np.vstack((u, v))
+ # print '--->U.shape: ', U.shape
+ # print 'U:\n', U
+
+ # We know that X * r = U
+ if rank(X) >= 2 * K:
+ r, _, _, _ = lstsq(X, U, rcond=-1)
+ r = np.squeeze(r)
+ else:
+ raise Exception("cp2tform:twoUniquePointsReq")
+
+ # print '--->r:\n', r
+
+ sc = r[0]
+ ss = r[1]
+ tx = r[2]
+ ty = r[3]
+
+ Tinv = np.array([[sc, -ss, 0], [ss, sc, 0], [tx, ty, 1]])
+
+ # print '--->Tinv:\n', Tinv
+
+ T = inv(Tinv)
+ # print '--->T:\n', T
+
+ T[:, 2] = np.array([0, 0, 1])
+
+ return T, Tinv
+
+
+def findSimilarity(uv, xy, options=None):
+ """
+ Function:
+ ----------
+ Find Reflective Similarity Transform Matrix 'trans':
+ u = uv[:, 0]
+ v = uv[:, 1]
+ x = xy[:, 0]
+ y = xy[:, 1]
+ [x, y, 1] = [u, v, 1] * trans
+
+ Parameters:
+ ----------
+ @uv: Kx2 np.array
+ source points each row is a pair of coordinates (x, y)
+ @xy: Kx2 np.array
+ each row is a pair of inverse-transformed
+ @option: not used, keep it as None
+
+ Returns:
+ ----------
+ @trans: 3x3 np.array
+ transform matrix from uv to xy
+ @trans_inv: 3x3 np.array
+ inverse of trans, transform matrix from xy to uv
+
+ Matlab:
+ ----------
+ % The similarities are a superset of the nonreflective similarities as they may
+ % also include reflection.
+ %
+ % let sc = s*cos(theta)
+ % let ss = s*sin(theta)
+ %
+ % [ sc -ss
+ % [u v] = [x y 1] * ss sc
+ % tx ty]
+ %
+ % OR
+ %
+ % [ sc ss
+ % [u v] = [x y 1] * ss -sc
+ % tx ty]
+ %
+ % Algorithm:
+ % 1) Solve for trans1, a nonreflective similarity.
+ % 2) Reflect the xy data across the Y-axis,
+ % and solve for trans2r, also a nonreflective similarity.
+ % 3) Transform trans2r to trans2, undoing the reflection done in step 2.
+ % 4) Use TFORMFWD to transform uv using both trans1 and trans2,
+ % and compare the results, Returnsing the transformation corresponding
+ % to the smaller L2 norm.
+
+ % Need to reset options.K to prepare for calls to findNonreflectiveSimilarity.
+ % This is safe because we already checked that there are enough point pairs.
+ """
+ options = {"K": 2}
+
+ # uv = np.array(uv)
+ # xy = np.array(xy)
+
+ # Solve for trans1
+ trans1, trans1_inv = findNonreflectiveSimilarity(uv, xy, options)
+
+ # Solve for trans2
+
+ # manually reflect the xy data across the Y-axis
+ xyR = xy
+ xyR[:, 0] = -1 * xyR[:, 0]
+
+ trans2r, trans2r_inv = findNonreflectiveSimilarity(uv, xyR, options)
+
+ # manually reflect the tform to undo the reflection done on xyR
+ TreflectY = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, 1]])
+
+ trans2 = np.dot(trans2r, TreflectY)
+
+ # Figure out if trans1 or trans2 is better
+ xy1 = tformfwd(trans1, uv)
+ norm1 = norm(xy1 - xy)
+
+ xy2 = tformfwd(trans2, uv)
+ norm2 = norm(xy2 - xy)
+
+ if norm1 <= norm2:
+ return trans1, trans1_inv
+ else:
+ trans2_inv = inv(trans2)
+ return trans2, trans2_inv
+
+
+def get_similarity_transform(src_pts, dst_pts, reflective=True):
+ """
+ Function:
+ ----------
+ Find Similarity Transform Matrix 'trans':
+ u = src_pts[:, 0]
+ v = src_pts[:, 1]
+ x = dst_pts[:, 0]
+ y = dst_pts[:, 1]
+ [x, y, 1] = [u, v, 1] * trans
+
+ Parameters:
+ ----------
+ @src_pts: Kx2 np.array
+ source points, each row is a pair of coordinates (x, y)
+ @dst_pts: Kx2 np.array
+ destination points, each row is a pair of transformed
+ coordinates (x, y)
+ @reflective: True or False
+ if True:
+ use reflective similarity transform
+ else:
+ use non-reflective similarity transform
+
+ Returns:
+ ----------
+ @trans: 3x3 np.array
+ transform matrix from uv to xy
+ trans_inv: 3x3 np.array
+ inverse of trans, transform matrix from xy to uv
+ """
+
+ if reflective:
+ trans, trans_inv = findSimilarity(src_pts, dst_pts)
+ else:
+ trans, trans_inv = findNonreflectiveSimilarity(src_pts, dst_pts)
+
+ return trans, trans_inv
+
+
+def cvt_tform_mat_for_cv2(trans):
+ """
+ Function:
+ ----------
+ Convert Transform Matrix 'trans' into 'cv2_trans' which could be
+ directly used by cv2.warpAffine():
+ u = src_pts[:, 0]
+ v = src_pts[:, 1]
+ x = dst_pts[:, 0]
+ y = dst_pts[:, 1]
+ [x, y].T = cv_trans * [u, v, 1].T
+
+ Parameters:
+ ----------
+ @trans: 3x3 np.array
+ transform matrix from uv to xy
+
+ Returns:
+ ----------
+ @cv2_trans: 2x3 np.array
+ transform matrix from src_pts to dst_pts, could be directly used
+ for cv2.warpAffine()
+ """
+ cv2_trans = trans[:, 0:2].T
+
+ return cv2_trans
+
+
+def get_similarity_transform_for_cv2(src_pts, dst_pts, reflective=True):
+ """
+ Function:
+ ----------
+ Find Similarity Transform Matrix 'cv2_trans' which could be
+ directly used by cv2.warpAffine():
+ u = src_pts[:, 0]
+ v = src_pts[:, 1]
+ x = dst_pts[:, 0]
+ y = dst_pts[:, 1]
+ [x, y].T = cv_trans * [u, v, 1].T
+
+ Parameters:
+ ----------
+ @src_pts: Kx2 np.array
+ source points, each row is a pair of coordinates (x, y)
+ @dst_pts: Kx2 np.array
+ destination points, each row is a pair of transformed
+ coordinates (x, y)
+ reflective: True or False
+ if True:
+ use reflective similarity transform
+ else:
+ use non-reflective similarity transform
+
+ Returns:
+ ----------
+ @cv2_trans: 2x3 np.array
+ transform matrix from src_pts to dst_pts, could be directly used
+ for cv2.warpAffine()
+ """
+ trans, trans_inv = get_similarity_transform(src_pts, dst_pts, reflective)
+ cv2_trans = cvt_tform_mat_for_cv2(trans)
+ return cv2_trans, trans
+
+
+std_points_317 = np.array(
+ [
+ [85.82991, 115.7792],
+ [169.0532, 114.3381],
+ [127.574, 167.0006],
+ [90.6964, 204.7014],
+ [167.3069, 203.3733],
+ ]
+)
+
+
+padding = 30
+
+std_points_317 = std_points_317 + padding
+
+std_points_256=std_points_317.copy()
+std_points_256[..., 0] -= 30
+std_points_256[..., 1] -= 60
+
+def warp_as_face_x_ray(img, src_pts, tgt_pts=std_points_317):
+ tfm, trans = get_similarity_transform_for_cv2(src_pts.copy(), tgt_pts.copy())
+ return cv2.warpAffine(img, tfm, (317, 317)), trans
+
+
+def estimiate_batch_transform(all_src_pts, tgt_pts=std_points_317):
+ tgt_pts = np.repeat(tgt_pts[None, ...], len(all_src_pts), 0).reshape(-1, 2)
+ src_pts = np.array(all_src_pts).reshape(-1, 2)
+ tfm, trans = get_similarity_transform_for_cv2(src_pts, tgt_pts)
+ return tfm, trans
+
+
+def batch_warp_as_face_x_ray(images, all_src_pts, tgt_pts=std_points_317):
+ tfm, trans = estimiate_batch_transform(all_src_pts, tgt_pts)
+ return [cv2.warpAffine(img, tfm, (317, 317)) for img in images], trans
+
+
+def transform_landmarks(landmarks, trans):
+ transformed = np.hstack((landmarks, np.ones((landmarks.shape[0], 1))))
+ transformed = np.dot(transformed, trans)
+ return transformed[:, :2]
+
+def compute_reverse_trans(trans):
+ return np.linalg.inv(trans)
\ No newline at end of file
diff --git a/clean/video/recce/.gitignore b/clean/video/recce/.gitignore
new file mode 100644
index 0000000000000000000000000000000000000000..4c59502f6d5d0ac94d4206f836bb4bb32f975057
--- /dev/null
+++ b/clean/video/recce/.gitignore
@@ -0,0 +1,134 @@
+# ignore directory
+runs/
+.idea/
+
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+# Usually these files are written by a python script from a template
+# before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+# However, in case of collaboration, if having platform-specific dependencies or dependencies
+# having no cross-platform support, pipenv may install dependencies that don't work, or not
+# install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
diff --git a/clean/video/recce/LICENSE b/clean/video/recce/LICENSE
new file mode 100644
index 0000000000000000000000000000000000000000..0db8340366af7e7f31596692ad9de4e746e1e3da
--- /dev/null
+++ b/clean/video/recce/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 SJTU Vision and Learning Group
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/clean/video/recce/README.md b/clean/video/recce/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..1547c5c58097d5a9fc718458e3a8c4fcf9ddf5a8
--- /dev/null
+++ b/clean/video/recce/README.md
@@ -0,0 +1,82 @@
+# RECCE CVPR 2022
+
+:page_facing_up: End-to-End Reconstruction-Classification Learning for Face Forgery Detection
+
+:boy: Junyi Cao, Chao Ma, Taiping Yao, Shen Chen, Shouhong Ding, Xiaokang Yang
+
+**Please consider citing our paper if you find it interesting or helpful to your research.**
+```
+@InProceedings{Cao_2022_CVPR,
+ author = {Cao, Junyi and Ma, Chao and Yao, Taiping and Chen, Shen and Ding, Shouhong and Yang, Xiaokang},
+ title = {End-to-End Reconstruction-Classification Learning for Face Forgery Detection},
+ booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+ month = {June},
+ year = {2022},
+ pages = {4113-4122}
+}
+```
+
+----
+
+### Introduction
+
+This repository is an implementation for *End-to-End Reconstruction-Classification Learning for Face Forgery Detection* presented in CVPR 2022. In the paper, we propose a novel **REC**onstruction-**C**lassification l**E**arning framework called **RECCE** to detect face forgeries. The code is based on Pytorch. Please follow the instructions below to get started.
+
+
+### Motivation
+
+Briefly, we train a reconstruction network over genuine images only and use the output of the latent feature by the encoder to perform binary classification. Due to the discrepancy in the data distribution between genuine and forged faces, the reconstruction differences of forged faces are obvious and also indicate the probably forged regions.
+
+
+### Basic Requirements
+Please ensure that you have already installed the following packages.
+- [Pytorch](https://pytorch.org/get-started/previous-versions/) 1.7.1
+- [Torchvision](https://pytorch.org/get-started/previous-versions/) 0.8.2
+- [Albumentations](https://github.com/albumentations-team/albumentations#spatial-level-transforms) 1.0.3
+- [Timm](https://github.com/rwightman/pytorch-image-models) 0.3.4
+- [TensorboardX](https://pypi.org/project/tensorboardX/#history) 2.1
+- [Scipy](https://pypi.org/project/scipy/#history) 1.5.2
+- [PyYaml](https://pypi.org/project/PyYAML/#history) 5.3.1
+
+### Dataset Preparation
+- We include the dataset loaders for several commonly-used face forgery datasets, *i.e.,* [FaceForensics++](https://github.com/ondyari/FaceForensics), [Celeb-DF](https://www.cs.albany.edu/~lsw/celeb-deepfakeforensics.html), [WildDeepfake](https://github.com/deepfakeinthewild/deepfake-in-the-wild), and [DFDC](https://ai.facebook.com/datasets/dfdc). You can enter the dataset website to download the original data.
+- For FaceForensics++, Celeb-DF, and DFDC, since the original data are in video format, you should first extract the facial images from the sequences and store them. We use [RetinaFace](https://github.com/biubug6/Pytorch_Retinaface) to do this.
+
+### Config Files
+- We have already provided the config templates in `config/`. You can adjust the parameters in the yaml files to specify a training process. More information is presented in [config/README.md](./config/README.md).
+
+### Training
+- We use `torch.distributed` package to train the models, for more information, please refer to [PyTorch Distributed Overview](https://pytorch.org/tutorials/beginner/dist_overview.html).
+- To train a model, run the following script in your console.
+```{bash}
+CUDA_VISIBLE_DEVICES=0 python -m torch.distributed.launch --nproc_per_node=1 --master_port 12345 train.py --config path/to/config.yaml
+```
+- `--config`: Specify the path of the config file.
+
+### Testing
+- To test a model, run the following script in your console.
+```{bash}
+python test.py --config path/to/config.yaml
+```
+- `--config`: Specify the path of the config file.
+
+### Inference
+- We provide the script in `inference.py` to help you do inference using custom data.
+- To do inference, run the following script in your console.
+```{bash}
+python inference.py --bin path/to/model.bin --image_folder path/to/image_folder --device $DEVICE --image_size $IMAGE_SIZE
+```
+- `--bin`: Specify the path of the model bin generated by the training script of this project.
+- `--image_folder`: Specify the directory of custom facial images. The script accepts images end with `.jpg` or `.png`.
+- `--device`: Specify the device to run the experiment, e.g., `cpu`, `cuda:0`.
+- `--image_size`: Specify the spatial size of input images.
+- The program will output the fake probability for each input image like this:
+ ```
+ path: path/to/image1.jpg | fake probability: 0.1296 | prediction: real
+ path: path/to/image2.jpg | fake probability: 0.9146 | prediction: fake
+ ```
+- Type `python inference.py -h` in your console for more information about available arguments.
+
+
+### Acknowledgement
+- We thank Qiqi Gu for helping plot the schematic diagram of the proposed method in the manuscript.
diff --git a/clean/video/recce/SOURCE.md b/clean/video/recce/SOURCE.md
new file mode 100644
index 0000000000000000000000000000000000000000..f33f4e7b7717bf51e16d577505cfbc92f5208112
--- /dev/null
+++ b/clean/video/recce/SOURCE.md
@@ -0,0 +1,17 @@
+# Source: video/recce
+
+| Field | Value |
+|---|---|
+| Upstream | **UNVERIFIED** -- provenance was lost when this code was vendored |
+| Paper | not recorded |
+| Commit SHA | **not recorded** -- the vendoring step did not preserve it |
+| Mirrored on | 2026-09-15 |
+| Upstream license | LICENSE |
+
+This is a **mirror**, stripped to the files needed for inference. The full
+untouched snapshot is at `archive/video__recce.tar.gz`.
+
+This code is the work of its original authors and is **not** covered by the
+DeepSafe project license. If you are an author and want this removed, open an
+issue on https://github.com/deepsafehq/deepsafe-bench and it will be taken down
+within 48 hours, no questions asked.
diff --git a/clean/video/recce/config/README.md b/clean/video/recce/config/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..dfb996189a40279bef13deb836f4b1b87ef0e57f
--- /dev/null
+++ b/clean/video/recce/config/README.md
@@ -0,0 +1,50 @@
+## Configuration Files
+
+#### Model Configuration
+- We use a yaml file to specify the hyperparameters of a model. All the training logs will be placed in `${project_root}/runs/${model_name}/${experiment_id}`. An example are shown below.
+
+```yaml
+model:
+ name: Recce # Model Name
+ num_classes: 1
+config:
+ lambda_1: 0.1 # balancing weight for L_r
+ lambda_2: 0.1 # balancing weight for L_m
+ distribute:
+ backend: nccl
+ optimizer:
+ name: adam
+ lr: 0.0002
+ weight_decay: 0.00001
+ scheduler:
+ name: StepLR
+ step_size: 22500
+ gamma: 0.5
+ resume: False
+ resume_best: False
+ id: FF++c40 # Specify a unique experiment id.
+ loss: binary_ce # Loss type, either 'binary_ce' or 'cross_entropy'.
+ metric: Acc # Main metric, either 'Acc', 'AUC', or 'LogLoss'.
+ debug: False
+ device: "cuda:1" # NOTE: Used only when testing, annotation this line when training.
+ ckpt: best_model_1000 # NOTE: Used only when testing to specify a checkpoint id, annotating this line when training.
+data:
+ train_batch_size: 32
+ val_batch_size: 64
+ test_batch_size: 64
+ name: FaceForensics
+ file: "./config/dataset/faceforensics.yml" # config file for a dataset
+ train_branch: "train_cfg"
+ val_branch: "test_cfg"
+ test_branch: "test_cfg"
+```
+
+- We set different hyper-parameters for the learning rate scheduler according to the used dataset as follows:
+ - FaceForensics++: The learning rate is decayed by 0.5 every 10 epochs.
+ - Celeb-DF: The learning rate is decayed by 0.5 every 10 epochs.
+ - WildDeepfake: The learning rate is decayed by 0.9 every 3000 iterations.
+ - DFDC: The learning rate is decayed by 0.5 every 3 epochs.
+
+#### Dataset Configuration
+- We also use a yaml file to specify the dataset to load for the experiment. These files are placed under `config/dataset/` subfold.
+- Briefly, you should change the `root` parameter according to your storage path.
\ No newline at end of file
diff --git a/clean/video/recce/config/Recce.yml b/clean/video/recce/config/Recce.yml
new file mode 100644
index 0000000000000000000000000000000000000000..beb2f60001712ef0ceab1c03bfd0d2b46aa832b9
--- /dev/null
+++ b/clean/video/recce/config/Recce.yml
@@ -0,0 +1,33 @@
+model:
+ name: Recce
+ num_classes: 1
+config:
+ lambda_1: 0.1
+ lambda_2: 0.1
+ distribute:
+ backend: nccl
+ optimizer:
+ name: adam
+ lr: 0.0002
+ weight_decay: 0.00001
+ scheduler:
+ name: StepLR
+ step_size: 22500
+ gamma: 0.5
+ resume: False
+ resume_best: False
+ id: FF++c40
+ loss: binary_ce
+ metric: Acc
+ debug: False
+# device: "cuda:1"
+# ckpt: best_model_1000
+data:
+ train_batch_size: 32
+ val_batch_size: 64
+ test_batch_size: 64
+ name: FaceForensics
+ file: "./config/dataset/faceforensics.yml"
+ train_branch: "train_cfg"
+ val_branch: "test_cfg"
+ test_branch: "test_cfg"
diff --git a/clean/video/recce/inference.py b/clean/video/recce/inference.py
new file mode 100644
index 0000000000000000000000000000000000000000..02c6d70219b2cf02dad6551d1fda5d9946014a4f
--- /dev/null
+++ b/clean/video/recce/inference.py
@@ -0,0 +1,129 @@
+import cv2
+import torch
+import random
+import argparse
+from glob import glob
+from os.path import join
+from model.network import Recce
+from model.common import freeze_weights
+from albumentations import Compose, Normalize, Resize
+from albumentations.pytorch.transforms import ToTensorV2
+
+# fix random seed
+seed = 0
+random.seed(seed)
+torch.manual_seed(seed)
+torch.cuda.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+
+parser = argparse.ArgumentParser(description="This code helps you use a trained model to "
+ "do inference.")
+parser.add_argument("--weight", "-w",
+ type=str,
+ default=None,
+ help="Specify the path to the model weight (the state dict file). "
+ "Do not use this argument when '--bin' is set.")
+parser.add_argument("--bin", "-b",
+ type=str,
+ default=None,
+ help="Specify the path to the model bin which ends up with '.bin' "
+ "(which is generated by the trainer of this project). "
+ "Do not use this argument when '--weight' is set.")
+parser.add_argument("--image", "-i",
+ type=str,
+ default=None,
+ help="Specify the path to the input image. "
+ "Do not use this argument when '--image_folder' is set.")
+parser.add_argument("--image_folder", "-f",
+ type=str,
+ default=None,
+ help="Specify the directory to evaluate all the images. "
+ "Do not use this argument when '--image' is set.")
+parser.add_argument('--device', '-d', type=str,
+ default="cpu",
+ help="Specify the device to load the model. Default: 'cpu'.")
+parser.add_argument('--image_size', '-s', type=int,
+ default=299,
+ help="Specify the spatial size of the input image(s). Default: 299.")
+parser.add_argument('--visualize', '-v', action="store_true",
+ default=False, help='Visualize images.')
+
+
+def preprocess(file_path):
+ img = cv2.imread(file_path)
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+ compose = Compose([Resize(height=args.image_size, width=args.image_size),
+ Normalize(mean=[0.5] * 3, std=[0.5] * 3),
+ ToTensorV2()])
+ img = compose(image=img)['image'].unsqueeze(0)
+ return img
+
+
+def prepare_data():
+ paths = list()
+ images = list()
+ # check the console arguments
+ if args.image and args.image_folder:
+ raise ValueError("Only one of '--image' or '--image_folder' can be set.")
+ elif args.image:
+ images.append(preprocess(args.image))
+ paths.append(args.image)
+ elif args.image_folder:
+ image_paths = glob(join(args.image_folder, "*.jpg"))
+ image_paths.extend(glob(join(args.image_folder, "*.png")))
+ for _ in image_paths:
+ images.append(preprocess(_))
+ paths.append(_)
+ else:
+ raise ValueError("Neither of '--image' nor '--image_folder' is set. Please specify either "
+ "one of these two arguments to load input image(s) properly.")
+ return paths, images
+
+
+def inference(model, images, paths, device):
+ for img, pt in zip(images, paths):
+ img = img.to(device)
+ prediction = model(img)
+ prediction = torch.sigmoid(prediction).cpu()
+ fake = True if prediction >= 0.5 else False
+ print(f"path: {pt} \t\t| fake probability: {prediction.item():.4f} \t| "
+ f"prediction: {'fake' if fake else 'real'}")
+ if args.visualize:
+ cvimg = cv2.imread(pt)
+ cvimg = cv2.putText(cvimg, f'p: {prediction.item():.2f}, ' + f"{'fake' if fake else 'real'}",
+ (5, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
+ (0, 0, 255) if fake else (255, 0, 0), 2)
+ cv2.imshow("image", cvimg)
+ cv2.waitKey(0)
+ cv2.destroyWindow("image")
+
+
+def main():
+ print("Arguments:\n", args, end="\n\n")
+ # set device
+ device = torch.device(args.device)
+ # load model
+ model = eval("Recce")(num_classes=1)
+ # check the console arguments
+ if args.weight and args.bin:
+ raise ValueError("Only one of '--weight' or '--bin' can be set.")
+ elif args.weight:
+ weights = torch.load(args.weight, map_location="cpu")
+ elif args.bin:
+ weights = torch.load(args.bin, map_location="cpu")["model"]
+ else:
+ raise ValueError("Neither of '--weight' nor '--bin' is set. Please specify either "
+ "one of these two arguments to load model's weight properly.")
+ model.load_state_dict(weights)
+ model = model.to(device)
+ freeze_weights(model)
+ model.eval()
+
+ paths, images = prepare_data()
+ print("Inference:")
+ inference(model, images=images, paths=paths, device=device)
+
+
+if __name__ == '__main__':
+ args = parser.parse_args()
+ main()
diff --git a/clean/video/recce/loss/__init__.py b/clean/video/recce/loss/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..fb8d2007cb4a3f852c7bde93fe84b0e162408d71
--- /dev/null
+++ b/clean/video/recce/loss/__init__.py
@@ -0,0 +1,12 @@
+import torch.nn as nn
+
+
+def get_loss(name="cross_entropy", device="cuda:0"):
+ print(f"Using loss: '{LOSSES[name]}'")
+ return LOSSES[name].to(device)
+
+
+LOSSES = {
+ "binary_ce": nn.BCEWithLogitsLoss(),
+ "cross_entropy": nn.CrossEntropyLoss()
+}
diff --git a/clean/video/recce/model/__init__.py b/clean/video/recce/model/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c8028d946b3ef30ba6b9ade62b503c3455c49ed6
--- /dev/null
+++ b/clean/video/recce/model/__init__.py
@@ -0,0 +1,12 @@
+from .network import *
+from .common import *
+
+MODELS = {
+ "Recce": Recce
+}
+
+
+def load_model(name="Recce"):
+ assert name in MODELS.keys(), f"Model name can only be one of {MODELS.keys()}."
+ print(f"Using model: '{name}'")
+ return MODELS[name]
diff --git a/clean/video/recce/model/common.py b/clean/video/recce/model/common.py
new file mode 100644
index 0000000000000000000000000000000000000000..1c8a1b962ac277a2164d4387d60083632ce9096b
--- /dev/null
+++ b/clean/video/recce/model/common.py
@@ -0,0 +1,200 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def freeze_weights(module):
+ for param in module.parameters():
+ param.requires_grad = False
+
+
+def l1_regularize(module):
+ reg_loss = 0.
+ for key, param in module.reg_params.items():
+ if "weight" in key and param.requires_grad:
+ reg_loss += torch.sum(torch.abs(param))
+ return reg_loss
+
+
+class SeparableConv2d(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False):
+ super(SeparableConv2d, self).__init__()
+
+ self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size, stride, padding, dilation,
+ groups=in_channels, bias=bias)
+ self.pointwise = nn.Conv2d(in_channels, out_channels, 1, 1, 0, 1, 1, bias=bias)
+
+ def forward(self, x):
+ x = self.conv1(x)
+ x = self.pointwise(x)
+ return x
+
+
+class Block(nn.Module):
+ def __init__(self, in_channels, out_channels, reps, strides=1,
+ start_with_relu=True, grow_first=True, with_bn=True):
+ super(Block, self).__init__()
+
+ self.with_bn = with_bn
+
+ if out_channels != in_channels or strides != 1:
+ self.skip = nn.Conv2d(in_channels, out_channels, 1, stride=strides, bias=False)
+ if with_bn:
+ self.skipbn = nn.BatchNorm2d(out_channels)
+ else:
+ self.skip = None
+
+ rep = []
+ for i in range(reps):
+ if grow_first:
+ inc = in_channels if i == 0 else out_channels
+ outc = out_channels
+ else:
+ inc = in_channels
+ outc = in_channels if i < (reps - 1) else out_channels
+ rep.append(nn.ReLU(inplace=True))
+ rep.append(SeparableConv2d(inc, outc, 3, stride=1, padding=1))
+ if with_bn:
+ rep.append(nn.BatchNorm2d(outc))
+
+ if not start_with_relu:
+ rep = rep[1:]
+ else:
+ rep[0] = nn.ReLU(inplace=False)
+
+ if strides != 1:
+ rep.append(nn.MaxPool2d(3, strides, 1))
+ self.rep = nn.Sequential(*rep)
+
+ def forward(self, inp):
+ x = self.rep(inp)
+
+ if self.skip is not None:
+ skip = self.skip(inp)
+ if self.with_bn:
+ skip = self.skipbn(skip)
+ else:
+ skip = inp
+
+ x += skip
+ return x
+
+
+class GraphReasoning(nn.Module):
+ """ Graph Reasoning Module for information aggregation. """
+
+ def __init__(self, va_in, va_out, vb_in, vb_out, vc_in, vc_out, spatial_ratio, drop_rate):
+ super(GraphReasoning, self).__init__()
+ self.ratio = spatial_ratio
+ self.va_embedding = nn.Sequential(
+ nn.Conv2d(va_in, va_out, 1, bias=False),
+ nn.ReLU(True),
+ nn.Conv2d(va_out, va_out, 1, bias=False),
+ )
+ self.va_gated_b = nn.Sequential(
+ nn.Conv2d(va_in, va_out, 1, bias=False),
+ nn.Sigmoid()
+ )
+ self.va_gated_c = nn.Sequential(
+ nn.Conv2d(va_in, va_out, 1, bias=False),
+ nn.Sigmoid()
+ )
+ self.vb_embedding = nn.Sequential(
+ nn.Linear(vb_in, vb_out, bias=False),
+ nn.ReLU(True),
+ nn.Linear(vb_out, vb_out, bias=False),
+ )
+ self.vc_embedding = nn.Sequential(
+ nn.Linear(vc_in, vc_out, bias=False),
+ nn.ReLU(True),
+ nn.Linear(vc_out, vc_out, bias=False),
+ )
+ self.unfold_b = nn.Unfold(kernel_size=spatial_ratio[0], stride=spatial_ratio[0])
+ self.unfold_c = nn.Unfold(kernel_size=spatial_ratio[1], stride=spatial_ratio[1])
+ self.reweight_ab = nn.Sequential(
+ nn.Linear(va_out + vb_out, 1, bias=False),
+ nn.ReLU(True),
+ nn.Softmax(dim=1)
+ )
+ self.reweight_ac = nn.Sequential(
+ nn.Linear(va_out + vc_out, 1, bias=False),
+ nn.ReLU(True),
+ nn.Softmax(dim=1)
+ )
+ self.reproject = nn.Sequential(
+ nn.Conv2d(va_out + vb_out + vc_out, va_in, kernel_size=1, bias=False),
+ nn.ReLU(True),
+ nn.Conv2d(va_in, va_in, kernel_size=1, bias=False),
+ nn.Dropout(drop_rate) if drop_rate is not None else nn.Identity(),
+ )
+
+ def forward(self, vert_a, vert_b, vert_c):
+ emb_vert_a = self.va_embedding(vert_a)
+ emb_vert_a = emb_vert_a.reshape([emb_vert_a.shape[0], emb_vert_a.shape[1], -1])
+
+ gate_vert_b = 1 - self.va_gated_b(vert_a)
+ gate_vert_b = gate_vert_b.reshape(*emb_vert_a.shape)
+ gate_vert_c = 1 - self.va_gated_c(vert_a)
+ gate_vert_c = gate_vert_c.reshape(*emb_vert_a.shape)
+
+ vert_b = self.unfold_b(vert_b).reshape(
+ [vert_b.shape[0], vert_b.shape[1], self.ratio[0] * self.ratio[0], -1])
+ vert_b = vert_b.permute([0, 2, 3, 1])
+ emb_vert_b = self.vb_embedding(vert_b)
+
+ vert_c = self.unfold_c(vert_c).reshape(
+ [vert_c.shape[0], vert_c.shape[1], self.ratio[1] * self.ratio[1], -1])
+ vert_c = vert_c.permute([0, 2, 3, 1])
+ emb_vert_c = self.vc_embedding(vert_c)
+
+ agg_vb = list()
+ agg_vc = list()
+ for j in range(emb_vert_a.shape[-1]):
+ # ab propagating
+ emb_v_a = torch.stack([emb_vert_a[:, :, j]] * (self.ratio[0] ** 2), dim=1)
+ emb_v_b = emb_vert_b[:, :, j, :]
+ emb_v_ab = torch.cat([emb_v_a, emb_v_b], dim=-1)
+ w = self.reweight_ab(emb_v_ab)
+ agg_vb.append(torch.bmm(emb_v_b.transpose(1, 2), w).squeeze() * gate_vert_b[:, :, j])
+
+ # ac propagating
+ emb_v_a = torch.stack([emb_vert_a[:, :, j]] * (self.ratio[1] ** 2), dim=1)
+ emb_v_c = emb_vert_c[:, :, j, :]
+ emb_v_ac = torch.cat([emb_v_a, emb_v_c], dim=-1)
+ w = self.reweight_ac(emb_v_ac)
+ agg_vc.append(torch.bmm(emb_v_c.transpose(1, 2), w).squeeze() * gate_vert_c[:, :, j])
+
+ agg_vert_b = torch.stack(agg_vb, dim=-1)
+ agg_vert_c = torch.stack(agg_vc, dim=-1)
+ agg_vert_bc = torch.cat([agg_vert_b, agg_vert_c], dim=1)
+ agg_vert_abc = torch.cat([agg_vert_bc, emb_vert_a], dim=1)
+ agg_vert_abc = torch.sigmoid(agg_vert_abc)
+ agg_vert_abc = agg_vert_abc.reshape(vert_a.shape[0], -1, vert_a.shape[2], vert_a.shape[3])
+ return self.reproject(agg_vert_abc)
+
+
+class GuidedAttention(nn.Module):
+ """ Reconstruction Guided Attention. """
+
+ def __init__(self, depth=728, drop_rate=0.2):
+ super(GuidedAttention, self).__init__()
+ self.depth = depth
+ self.gated = nn.Sequential(
+ nn.Conv2d(3, 3, kernel_size=3, stride=1, padding=1, bias=False),
+ nn.ReLU(True),
+ nn.Conv2d(3, 1, 1, bias=False),
+ nn.Sigmoid()
+ )
+ self.h = nn.Sequential(
+ nn.Conv2d(depth, depth, 1, 1, bias=False),
+ nn.BatchNorm2d(depth),
+ nn.ReLU(True),
+ )
+ self.dropout = nn.Dropout(drop_rate)
+
+ def forward(self, x, pred_x, embedding):
+ residual_full = torch.abs(x - pred_x)
+ residual_x = F.interpolate(residual_full, size=embedding.shape[-2:],
+ mode='bilinear', align_corners=True)
+ res_map = self.gated(residual_x)
+ return res_map * self.h(embedding) + self.dropout(embedding)
diff --git a/clean/video/recce/model/network/Recce.py b/clean/video/recce/model/network/Recce.py
new file mode 100644
index 0000000000000000000000000000000000000000..eb647cfc245542108bf4de3450ee3db05c55fb6d
--- /dev/null
+++ b/clean/video/recce/model/network/Recce.py
@@ -0,0 +1,133 @@
+from functools import partial
+from timm.models import xception
+from model.common import SeparableConv2d, Block
+from model.common import GuidedAttention, GraphReasoning
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+encoder_params = {
+ "xception": {
+ "features": 2048,
+ "init_op": partial(xception, pretrained=True)
+ }
+}
+
+
+class Recce(nn.Module):
+ """ End-to-End Reconstruction-Classification Learning for Face Forgery Detection """
+
+ def __init__(self, num_classes, drop_rate=0.2):
+ super(Recce, self).__init__()
+ self.name = "xception"
+ self.loss_inputs = dict()
+ self.encoder = encoder_params[self.name]["init_op"]()
+ self.global_pool = nn.AdaptiveAvgPool2d((1, 1))
+ self.dropout = nn.Dropout(drop_rate)
+ self.fc = nn.Linear(encoder_params[self.name]["features"], num_classes)
+
+ self.attention = GuidedAttention(depth=728, drop_rate=drop_rate)
+ self.reasoning = GraphReasoning(728, 256, 256, 256, 128, 256, [2, 4], drop_rate)
+
+ self.decoder1 = nn.Sequential(
+ nn.UpsamplingNearest2d(scale_factor=2),
+ SeparableConv2d(728, 256, 3, 1, 1, bias=False),
+ nn.BatchNorm2d(256),
+ nn.ReLU(inplace=True)
+ )
+ self.decoder2 = Block(256, 256, 3, 1)
+ self.decoder3 = nn.Sequential(
+ nn.UpsamplingNearest2d(scale_factor=2),
+ SeparableConv2d(256, 128, 3, 1, 1, bias=False),
+ nn.BatchNorm2d(128),
+ nn.ReLU(inplace=True)
+ )
+ self.decoder4 = Block(128, 128, 3, 1)
+ self.decoder5 = nn.Sequential(
+ nn.UpsamplingNearest2d(scale_factor=2),
+ SeparableConv2d(128, 64, 3, 1, 1, bias=False),
+ nn.BatchNorm2d(64),
+ nn.ReLU(inplace=True)
+ )
+ self.decoder6 = nn.Sequential(
+ nn.Conv2d(64, 3, 1, 1, bias=False),
+ nn.Tanh()
+ )
+
+ def norm_n_corr(self, x):
+ norm_embed = F.normalize(self.global_pool(x), p=2, dim=1)
+ corr = (torch.matmul(norm_embed.squeeze(), norm_embed.squeeze().T) + 1.) / 2.
+ return norm_embed, corr
+
+ @staticmethod
+ def add_white_noise(tensor, mean=0., std=1e-6):
+ rand = torch.rand([tensor.shape[0], 1, 1, 1])
+ rand = torch.where(rand > 0.5, 1., 0.).to(tensor.device)
+ white_noise = torch.normal(mean, std, size=tensor.shape, device=tensor.device)
+ noise_t = tensor + white_noise * rand
+ noise_t = torch.clip(noise_t, -1., 1.)
+ return noise_t
+
+ def forward(self, x):
+ # clear the loss inputs
+ self.loss_inputs = dict(recons=[], contra=[])
+ noise_x = self.add_white_noise(x) if self.training else x
+ out = self.encoder.conv1(noise_x)
+ out = self.encoder.bn1(out)
+ out = self.encoder.act1(out)
+ out = self.encoder.conv2(out)
+ out = self.encoder.bn2(out)
+ out = self.encoder.act2(out)
+ out = self.encoder.block1(out)
+ out = self.encoder.block2(out)
+ out = self.encoder.block3(out)
+ embedding = self.encoder.block4(out)
+
+ norm_embed, corr = self.norm_n_corr(embedding)
+ self.loss_inputs['contra'].append(corr)
+
+ out = self.dropout(embedding)
+ out = self.decoder1(out)
+ out_d2 = self.decoder2(out)
+
+ norm_embed, corr = self.norm_n_corr(out_d2)
+ self.loss_inputs['contra'].append(corr)
+
+ out = self.decoder3(out_d2)
+ out_d4 = self.decoder4(out)
+
+ norm_embed, corr = self.norm_n_corr(out_d4)
+ self.loss_inputs['contra'].append(corr)
+
+ out = self.decoder5(out_d4)
+ pred = self.decoder6(out)
+
+ recons_x = F.interpolate(pred, size=x.shape[-2:], mode='bilinear', align_corners=True)
+ self.loss_inputs['recons'].append(recons_x)
+
+ embedding = self.encoder.block5(embedding)
+ embedding = self.encoder.block6(embedding)
+ embedding = self.encoder.block7(embedding)
+
+ fusion = self.reasoning(embedding, out_d2, out_d4) + embedding
+
+ embedding = self.encoder.block8(fusion)
+ img_att = self.attention(x, recons_x, embedding)
+
+ embedding = self.encoder.block9(img_att)
+ embedding = self.encoder.block10(embedding)
+ embedding = self.encoder.block11(embedding)
+ embedding = self.encoder.block12(embedding)
+
+ embedding = self.encoder.conv3(embedding)
+ embedding = self.encoder.bn3(embedding)
+ embedding = self.encoder.act3(embedding)
+ embedding = self.encoder.conv4(embedding)
+ embedding = self.encoder.bn4(embedding)
+ embedding = self.encoder.act4(embedding)
+
+ embedding = self.global_pool(embedding).squeeze()
+
+ out = self.dropout(embedding)
+ return self.fc(out)
diff --git a/clean/video/recce/model/network/__init__.py b/clean/video/recce/model/network/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e552eb84c758d2b1ac6542725d7faed4d1fca3bd
--- /dev/null
+++ b/clean/video/recce/model/network/__init__.py
@@ -0,0 +1 @@
+from .Recce import Recce
diff --git a/clean/video/recce/optimizer/__init__.py b/clean/video/recce/optimizer/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..74db9744a1f35e9f082e3a61997025ad4c22a8e7
--- /dev/null
+++ b/clean/video/recce/optimizer/__init__.py
@@ -0,0 +1,30 @@
+from torch.optim import SGD
+from torch.optim import Adam
+from torch.optim import ASGD
+from torch.optim import Adamax
+from torch.optim import Adadelta
+from torch.optim import Adagrad
+from torch.optim import RMSprop
+
+key2opt = {
+ 'sgd': SGD,
+ 'adam': Adam,
+ 'asgd': ASGD,
+ 'adamax': Adamax,
+ 'adadelta': Adadelta,
+ 'adagrad': Adagrad,
+ 'rmsprop': RMSprop,
+}
+
+
+def get_optimizer(optimizer_name=None):
+ if optimizer_name is None:
+ print("Using default 'SGD' optimizer")
+ return SGD
+
+ else:
+ if optimizer_name not in key2opt:
+ raise NotImplementedError(f"Optimizer '{optimizer_name}' not implemented")
+
+ print(f"Using optimizer: '{optimizer_name}'")
+ return key2opt[optimizer_name]
diff --git a/clean/video/recce/scheduler/__init__.py b/clean/video/recce/scheduler/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d02f3edf959d8e10b66c08e276ced619ecfd15dd
--- /dev/null
+++ b/clean/video/recce/scheduler/__init__.py
@@ -0,0 +1,36 @@
+from torch.optim.lr_scheduler import _LRScheduler
+from torch.optim.lr_scheduler import StepLR
+from torch.optim.lr_scheduler import MultiStepLR
+from torch.optim.lr_scheduler import ExponentialLR
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+
+
+class ConstantLR(_LRScheduler):
+ def __init__(self, optimizer, last_epoch=-1):
+ super(ConstantLR, self).__init__(optimizer, last_epoch)
+
+ def get_lr(self):
+ return [base_lr for base_lr in self.base_lrs]
+
+
+SCHEDULERS = {
+ 'ConstantLR': ConstantLR,
+ "StepLR": StepLR,
+ "MultiStepLR": MultiStepLR,
+ "CosineAnnealingLR": CosineAnnealingLR,
+ "CosineAnnealingWarmRestarts": CosineAnnealingWarmRestarts,
+ "ExponentialLR": ExponentialLR,
+ "ReduceLROnPlateau": ReduceLROnPlateau
+}
+
+
+def get_scheduler(optimizer, kwargs):
+ if kwargs is None:
+ print("No lr scheduler is used.")
+ return ConstantLR(optimizer)
+ name = kwargs["name"]
+ kwargs.pop("name")
+ print("Using scheduler: '%s' with params: %s" % (name, kwargs))
+ return SCHEDULERS[name](optimizer, **kwargs)
diff --git a/clean/video/recce/test.py b/clean/video/recce/test.py
new file mode 100644
index 0000000000000000000000000000000000000000..7467842f9affda299fbd5a06bb44cc6d691600f9
--- /dev/null
+++ b/clean/video/recce/test.py
@@ -0,0 +1,31 @@
+import yaml
+import argparse
+
+from trainer import ExpTester
+
+
+def arg_parser():
+ parser = argparse.ArgumentParser(description="config")
+ parser.add_argument("--config",
+ type=str,
+ default="config/Recce.yml",
+ help="Specify the path of configuration file to be used.")
+ parser.add_argument('--display', '-d', action="store_true",
+ default=False, help='Display some images.')
+ return parser.parse_args()
+
+
+if __name__ == '__main__':
+ import torch
+
+ torch.backends.cudnn.benchmark = True
+ torch.backends.cudnn.enabled = True
+
+ arg = arg_parser()
+ config = arg.config
+
+ with open(config) as config_file:
+ config = yaml.load(config_file, Loader=yaml.FullLoader)
+
+ trainer = ExpTester(config, stage="Test")
+ trainer.test(display_images=arg.display)
diff --git a/clean/video/recce/train.py b/clean/video/recce/train.py
new file mode 100644
index 0000000000000000000000000000000000000000..cae7ad48eddf9d945142cc77c818df3dd8b30a71
--- /dev/null
+++ b/clean/video/recce/train.py
@@ -0,0 +1,33 @@
+import yaml
+import argparse
+
+from trainer import ExpMultiGpuTrainer
+
+
+def arg_parser():
+ parser = argparse.ArgumentParser(description="config")
+ parser.add_argument("--config",
+ type=str,
+ default="config/Recce.yml",
+ help="Specified the path of configuration file to be used.")
+ parser.add_argument("--local_rank", default=0,
+ type=int,
+ help="Specified the node rank for distributed training.")
+ return parser.parse_args()
+
+
+if __name__ == '__main__':
+ import torch
+
+ torch.backends.cudnn.benchmark = True
+ torch.backends.cudnn.enabled = True
+
+ arg = arg_parser()
+ config = arg.config
+
+ with open(config) as config_file:
+ config = yaml.load(config_file, Loader=yaml.FullLoader)
+ config["config"]["local_rank"] = arg.local_rank
+
+ trainer = ExpMultiGpuTrainer(config, stage="Train")
+ trainer.train()
diff --git a/clean/video/recce/trainer/__init__.py b/clean/video/recce/trainer/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d7f4da6b97a20cf6c6dae6761dffea145a336cea
--- /dev/null
+++ b/clean/video/recce/trainer/__init__.py
@@ -0,0 +1,5 @@
+from .abstract_trainer import AbstractTrainer, LEGAL_METRIC
+from .exp_mgpu_trainer import ExpMultiGpuTrainer
+from .exp_tester import ExpTester
+from .utils import center_print, reduce_tensor
+from .utils import exp_recons_loss
diff --git a/clean/video/recce/trainer/abstract_trainer.py b/clean/video/recce/trainer/abstract_trainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..4e1354a94ec55db2fcad0d951ecaca2f7b804fbe
--- /dev/null
+++ b/clean/video/recce/trainer/abstract_trainer.py
@@ -0,0 +1,100 @@
+import os
+import torch
+import random
+from collections import OrderedDict
+from torchvision.utils import make_grid
+
+LEGAL_METRIC = ['Acc', 'AUC', 'LogLoss']
+
+
+class AbstractTrainer(object):
+ def __init__(self, config, stage="Train"):
+ feasible_stage = ["Train", "Test"]
+ if stage not in feasible_stage:
+ raise ValueError(f"stage should be in {feasible_stage}, but found '{stage}'")
+
+ self.config = config
+ model_cfg = config.get("model", None)
+ data_cfg = config.get("data", None)
+ config_cfg = config.get("config", None)
+
+ self.model_name = model_cfg.pop("name")
+
+ self.gpu = None
+ self.dir = None
+ self.debug = None
+ self.device = None
+ self.resume = None
+ self.local_rank = None
+ self.num_classes = None
+
+ self.best_metric = 0.0
+ self.best_step = 1
+ self.start_step = 1
+
+ self._initiated_settings(model_cfg, data_cfg, config_cfg)
+
+ if stage == 'Train':
+ self._train_settings(model_cfg, data_cfg, config_cfg)
+ if stage == 'Test':
+ self._test_settings(model_cfg, data_cfg, config_cfg)
+
+ def _initiated_settings(self, model_cfg, data_cfg, config_cfg):
+ raise NotImplementedError("Not implemented in abstract class.")
+
+ def _train_settings(self, model_cfg, data_cfg, config_cfg):
+ raise NotImplementedError("Not implemented in abstract class.")
+
+ def _test_settings(self, model_cfg, data_cfg, config_cfg):
+ raise NotImplementedError("Not implemented in abstract class.")
+
+ def _save_ckpt(self, step, best=False):
+ raise NotImplementedError("Not implemented in abstract class.")
+
+ def _load_ckpt(self, best=False, train=False):
+ raise NotImplementedError("Not implemented in abstract class.")
+
+ def to_device(self, items):
+ return [obj.to(self.device) for obj in items]
+
+ @staticmethod
+ def fixed_randomness():
+ random.seed(0)
+ torch.manual_seed(0)
+ torch.cuda.manual_seed(0)
+ torch.cuda.manual_seed_all(0)
+
+ def train(self):
+ raise NotImplementedError("Not implemented in abstract class.")
+
+ def validate(self, epoch, step, timer, writer):
+ raise NotImplementedError("Not implemented in abstract class.")
+
+ def test(self):
+ raise NotImplementedError("Not implemented in abstract class.")
+
+ def plot_figure(self, images, pred, gt, nrow, categories=None, show=True):
+ import matplotlib.pyplot as plt
+ plot = make_grid(
+ images, nrow, padding=4, normalize=True, scale_each=True, pad_value=1)
+ if self.num_classes == 1:
+ pred = (pred >= 0.5).cpu().numpy()
+ else:
+ pred = pred.argmax(1).cpu().numpy()
+ gt = gt.cpu().numpy()
+ if categories is not None:
+ pred = [categories[i] for i in pred]
+ gt = [categories[i] for i in gt]
+ plot = plot.permute([1, 2, 0])
+ plot = plot.cpu().numpy()
+ ret = plt.figure()
+ plt.imshow(plot)
+ plt.title("pred: %s\ngt: %s" % (pred, gt))
+ plt.axis("off")
+ if show:
+ plt.savefig(os.path.join(self.dir, "test_image.png"), dpi=300)
+ plt.show()
+ plt.close()
+ else:
+ plt.close()
+ return ret
diff --git a/clean/video/recce/trainer/exp_mgpu_trainer.py b/clean/video/recce/trainer/exp_mgpu_trainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..f9e2d43c4e040b4390870b7ca5af07809f5e30c3
--- /dev/null
+++ b/clean/video/recce/trainer/exp_mgpu_trainer.py
@@ -0,0 +1,370 @@
+import os
+import sys
+import time
+import math
+import yaml
+import torch
+import random
+import numpy as np
+
+from tqdm import tqdm
+from pprint import pprint
+from torch.utils import data
+import torch.distributed as dist
+from torch.cuda.amp import autocast, GradScaler
+from tensorboardX import SummaryWriter
+
+from dataset import load_dataset
+from loss import get_loss
+from model import load_model
+from optimizer import get_optimizer
+from scheduler import get_scheduler
+from trainer import AbstractTrainer, LEGAL_METRIC
+from trainer.utils import exp_recons_loss, MLLoss, reduce_tensor, center_print
+from trainer.utils import MODELS_PATH, AccMeter, AUCMeter, AverageMeter, Logger, Timer
+
+
+class ExpMultiGpuTrainer(AbstractTrainer):
+ def __init__(self, config, stage="Train"):
+ super(ExpMultiGpuTrainer, self).__init__(config, stage)
+ np.random.seed(2021)
+
+ def _mprint(self, content=""):
+ if self.local_rank == 0:
+ print(content)
+
+ def _initiated_settings(self, model_cfg=None, data_cfg=None, config_cfg=None):
+ self.local_rank = config_cfg["local_rank"]
+
+ def _train_settings(self, model_cfg, data_cfg, config_cfg):
+ # debug mode: no log dir, no train_val operation.
+ self.debug = config_cfg["debug"]
+ self._mprint(f"Using debug mode: {self.debug}.")
+ self._mprint("*" * 20)
+
+ self.eval_metric = config_cfg["metric"]
+ if self.eval_metric not in LEGAL_METRIC:
+ raise ValueError(f"Evaluation metric must be in {LEGAL_METRIC}, but found "
+ f"{self.eval_metric}.")
+ if self.eval_metric == LEGAL_METRIC[-1]:
+ self.best_metric = 1.0e8
+
+ # distribution
+ dist.init_process_group(config_cfg["distribute"]["backend"])
+
+ # load training dataset
+ train_dataset = data_cfg["file"]
+ branch = data_cfg["train_branch"]
+ name = data_cfg["name"]
+ with open(train_dataset, "r") as f:
+ options = yaml.load(f, Loader=yaml.FullLoader)
+ train_options = options[branch]
+ self.train_set = load_dataset(name)(train_options)
+ # define training sampler
+ self.train_sampler = data.distributed.DistributedSampler(self.train_set)
+ # wrapped with data loader
+ self.train_loader = data.DataLoader(self.train_set, shuffle=False,
+ sampler=self.train_sampler,
+ num_workers=data_cfg.get("num_workers", 4),
+ batch_size=data_cfg["train_batch_size"])
+
+ if self.local_rank == 0:
+ # load validation dataset
+ val_options = options[data_cfg["val_branch"]]
+ self.val_set = load_dataset(name)(val_options)
+ # wrapped with data loader
+ self.val_loader = data.DataLoader(self.val_set, shuffle=True,
+ num_workers=data_cfg.get("num_workers", 4),
+ batch_size=data_cfg["val_batch_size"])
+
+ self.resume = config_cfg.get("resume", False)
+
+ if not self.debug:
+ time_format = "%Y-%m-%d...%H.%M.%S"
+ run_id = time.strftime(time_format, time.localtime(time.time()))
+ self.run_id = config_cfg.get("id", run_id)
+ self.dir = os.path.join("runs", self.model_name, self.run_id)
+
+ if self.local_rank == 0:
+ if not self.resume:
+ if os.path.exists(self.dir):
+ raise ValueError("Error: given id '%s' already exists." % self.run_id)
+ os.makedirs(self.dir, exist_ok=True)
+ print(f"Writing config file to file directory: {self.dir}.")
+ yaml.dump({"config": self.config,
+ "train_data": train_options,
+ "val_data": val_options},
+ open(os.path.join(self.dir, 'train_config.yml'), 'w'))
+ # copy the script for the training model
+ model_file = MODELS_PATH[self.model_name]
+ os.system("cp " + model_file + " " + self.dir)
+ else:
+ print(f"Resuming the history in file directory: {self.dir}.")
+
+ print(f"Logging directory: {self.dir}.")
+
+ # redirect the std out stream
+ sys.stdout = Logger(os.path.join(self.dir, 'records.txt'))
+ center_print('Train configurations begins.')
+ pprint(self.config)
+ pprint(train_options)
+ pprint(val_options)
+ center_print('Train configurations ends.')
+
+ # load model
+ self.num_classes = model_cfg["num_classes"]
+ self.device = "cuda:" + str(self.local_rank)
+ self.model = load_model(self.model_name)(**model_cfg)
+ self.model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.model).to(self.device)
+ self._mprint(f"Using SyncBatchNorm.")
+ self.model = torch.nn.parallel.DistributedDataParallel(
+ self.model, device_ids=[self.local_rank], find_unused_parameters=True)
+
+ # load optimizer
+ optim_cfg = config_cfg.get("optimizer", None)
+ optim_name = optim_cfg.pop("name")
+ self.optimizer = get_optimizer(optim_name)(self.model.parameters(), **optim_cfg)
+ # load scheduler
+ self.scheduler = get_scheduler(self.optimizer, config_cfg.get("scheduler", None))
+ # load loss
+ self.loss_criterion = get_loss(config_cfg.get("loss", None), device=self.device)
+
+ # total number of steps (or epoch) to train
+ self.num_steps = train_options["num_steps"]
+ self.num_epoch = math.ceil(self.num_steps / len(self.train_loader))
+
+ # the number of steps to write down a log
+ self.log_steps = train_options["log_steps"]
+ # the number of steps to validate on val dataset once
+ self.val_steps = train_options["val_steps"]
+
+ # balance coefficients
+ self.lambda_1 = config_cfg["lambda_1"]
+ self.lambda_2 = config_cfg["lambda_2"]
+ self.warmup_step = config_cfg.get('warmup_step', 0)
+
+ self.contra_loss = MLLoss()
+ self.acc_meter = AccMeter()
+ self.loss_meter = AverageMeter()
+ self.recons_loss_meter = AverageMeter()
+ self.contra_loss_meter = AverageMeter()
+
+ if self.resume and self.local_rank == 0:
+ self._load_ckpt(best=config_cfg.get("resume_best", False), train=True)
+
+ def _test_settings(self, model_cfg, data_cfg, config_cfg):
+ # Not used.
+ raise NotImplementedError("The function is not intended to be used here.")
+
+ def _load_ckpt(self, best=False, train=False):
+ # Not used.
+ raise NotImplementedError("The function is not intended to be used here.")
+
+ def _save_ckpt(self, step, best=False):
+ save_dir = os.path.join(self.dir, f"best_model_{step}.bin" if best else "latest_model.bin")
+ torch.save({
+ "step": step,
+ "best_step": self.best_step,
+ "best_metric": self.best_metric,
+ "eval_metric": self.eval_metric,
+ "model": self.model.module.state_dict(),
+ "optimizer": self.optimizer.state_dict(),
+ "scheduler": self.scheduler.state_dict(),
+ }, save_dir)
+
+ def train(self):
+ try:
+ timer = Timer()
+ grad_scalar = GradScaler(2 ** 10)
+ if self.local_rank == 0:
+ writer = None if self.debug else SummaryWriter(log_dir=self.dir)
+ center_print("Training begins......")
+ else:
+ writer = None
+ start_epoch = self.start_step // len(self.train_loader) + 1
+ for epoch_idx in range(start_epoch, self.num_epoch + 1):
+ # set sampler
+ self.train_sampler.set_epoch(epoch_idx)
+
+ # reset meter
+ self.acc_meter.reset()
+ self.loss_meter.reset()
+ self.recons_loss_meter.reset()
+ self.contra_loss_meter.reset()
+ self.optimizer.step()
+
+ train_generator = enumerate(self.train_loader, 1)
+ # wrap train generator with tqdm for process 0
+ if self.local_rank == 0:
+ train_generator = tqdm(train_generator, position=0, leave=True)
+
+ for batch_idx, train_data in train_generator:
+ global_step = (epoch_idx - 1) * len(self.train_loader) + batch_idx
+ self.model.train()
+ I, Y = train_data
+ I = self.train_loader.dataset.load_item(I)
+ in_I, Y = self.to_device((I, Y))
+
+ # warm-up lr
+ if self.warmup_step != 0 and global_step <= self.warmup_step:
+ lr = self.config['config']['optimizer']['lr'] * float(global_step) / self.warmup_step
+ for param_group in self.optimizer.param_groups:
+ param_group['lr'] = lr
+
+ self.optimizer.zero_grad()
+ with autocast():
+ Y_pre = self.model(in_I)
+
+ # for BCE Setting:
+ if self.num_classes == 1:
+ Y_pre = Y_pre.squeeze()
+ loss = self.loss_criterion(Y_pre, Y.float())
+ Y_pre = torch.sigmoid(Y_pre)
+ else:
+ loss = self.loss_criterion(Y_pre, Y)
+
+ # flood
+ loss = (loss - 0.04).abs() + 0.04
+ recons_loss = exp_recons_loss(self.model.module.loss_inputs['recons'], (in_I, Y))
+ contra_loss = self.contra_loss(self.model.module.loss_inputs['contra'], Y)
+ loss += self.lambda_1 * recons_loss + self.lambda_2 * contra_loss
+
+ grad_scalar.scale(loss).backward()
+ grad_scalar.step(self.optimizer)
+ grad_scalar.update()
+ if self.warmup_step == 0 or global_step > self.warmup_step:
+ self.scheduler.step()
+
+ self.acc_meter.update(Y_pre, Y, self.num_classes == 1)
+ self.loss_meter.update(reduce_tensor(loss).item())
+ self.recons_loss_meter.update(reduce_tensor(recons_loss).item())
+ self.contra_loss_meter.update(reduce_tensor(contra_loss).item())
+ iter_acc = reduce_tensor(self.acc_meter.mean_acc()).item()
+
+ if self.local_rank == 0:
+ if global_step % self.log_steps == 0 and writer is not None:
+ writer.add_scalar("train/Acc", iter_acc, global_step)
+ writer.add_scalar("train/Loss", self.loss_meter.avg, global_step)
+ writer.add_scalar("train/Recons_Loss",
+ self.recons_loss_meter.avg if self.lambda_1 != 0 else 0.,
+ global_step)
+ writer.add_scalar("train/Contra_Loss", self.contra_loss_meter.avg, global_step)
+ writer.add_scalar("train/LR", self.scheduler.get_last_lr()[0], global_step)
+
+ # log training step
+ train_generator.set_description(
+ "Train Epoch %d (%d/%d), Global Step %d, Loss %.4f, Recons %.4f, con %.4f, "
+ "ACC %.4f, LR %.6f" % (
+ epoch_idx, batch_idx, len(self.train_loader), global_step,
+ self.loss_meter.avg, self.recons_loss_meter.avg, self.contra_loss_meter.avg,
+ iter_acc, self.scheduler.get_last_lr()[0])
+ )
+
+ # validating process
+ if global_step % self.val_steps == 0 and not self.debug:
+ print()
+ self.validate(epoch_idx, global_step, timer, writer)
+
+ # when num_steps has been set and the training process will
+ # be stopped earlier than the specified num_epochs, then stop.
+ if self.num_steps is not None and global_step == self.num_steps:
+ if writer is not None:
+ writer.close()
+ if self.local_rank == 0:
+ print()
+ center_print("Training process ends.")
+ dist.destroy_process_group()
+ return
+ # close the tqdm bar when one epoch ends
+ if self.local_rank == 0:
+ train_generator.close()
+ print()
+ # training ends with integer epochs
+ if self.local_rank == 0:
+ if writer is not None:
+ writer.close()
+ center_print("Training process ends.")
+ dist.destroy_process_group()
+ except Exception as e:
+ dist.destroy_process_group()
+ raise e
+
+ def validate(self, epoch, step, timer, writer):
+ v_idx = random.randint(1, len(self.val_loader) + 1)
+ categories = self.val_loader.dataset.categories
+ self.model.eval()
+ with torch.no_grad():
+ acc = AccMeter()
+ auc = AUCMeter()
+ loss_meter = AverageMeter()
+ cur_acc = 0.0 # Higher is better
+ cur_auc = 0.0 # Higher is better
+ cur_loss = 1e8 # Lower is better
+ val_generator = tqdm(enumerate(self.val_loader, 1), position=0, leave=True)
+ for val_idx, val_data in val_generator:
+ I, Y = val_data
+ I = self.val_loader.dataset.load_item(I)
+ in_I, Y = self.to_device((I, Y))
+ Y_pre = self.model(in_I)
+
+ # for BCE Setting:
+ if self.num_classes == 1:
+ Y_pre = Y_pre.squeeze()
+ loss = self.loss_criterion(Y_pre, Y.float())
+ Y_pre = torch.sigmoid(Y_pre)
+ else:
+ loss = self.loss_criterion(Y_pre, Y)
+
+ acc.update(Y_pre, Y, self.num_classes == 1)
+ auc.update(Y_pre, Y, self.num_classes == 1)
+ loss_meter.update(loss.item())
+
+ cur_acc = acc.mean_acc()
+ cur_loss = loss_meter.avg
+
+ val_generator.set_description(
+ "Eval Epoch %d (%d/%d), Global Step %d, Loss %.4f, ACC %.4f" % (
+ epoch, val_idx, len(self.val_loader), step,
+ cur_loss, cur_acc)
+ )
+
+ if val_idx == v_idx or val_idx == 1:
+ sample_recons = list()
+ for _ in self.model.module.loss_inputs['recons']:
+ sample_recons.append(_[:4].to("cpu"))
+ # show images
+ images = I[:4]
+ images = torch.cat([images, *sample_recons], dim=0)
+ pred = Y_pre[:4]
+ gt = Y[:4]
+ figure = self.plot_figure(images, pred, gt, 4, categories, show=False)
+
+ cur_auc = auc.mean_auc()
+ print("Eval Epoch %d, Loss %.4f, ACC %.4f, AUC %.4f" % (epoch, cur_loss, cur_acc, cur_auc))
+ if writer is not None:
+ writer.add_scalar("val/Loss", cur_loss, step)
+ writer.add_scalar("val/Acc", cur_acc, step)
+ writer.add_scalar("val/AUC", cur_auc, step)
+ writer.add_figure("val/Figures", figure, step)
+ # record the best acc and the corresponding step
+ if self.eval_metric == 'Acc' and cur_acc >= self.best_metric:
+ self.best_metric = cur_acc
+ self.best_step = step
+ self._save_ckpt(step, best=True)
+ elif self.eval_metric == 'AUC' and cur_auc >= self.best_metric:
+ self.best_metric = cur_auc
+ self.best_step = step
+ self._save_ckpt(step, best=True)
+ elif self.eval_metric == 'LogLoss' and cur_loss <= self.best_metric:
+ self.best_metric = cur_loss
+ self.best_step = step
+ self._save_ckpt(step, best=True)
+ print("Best Step %d, Best %s %.4f, Running Time: %s, Estimated Time: %s" % (
+ self.best_step, self.eval_metric, self.best_metric,
+ timer.measure(), timer.measure(step / self.num_steps)
+ ))
+ self._save_ckpt(step, best=False)
+
+ def test(self):
+ # Not used.
+ raise NotImplementedError("The function is not intended to be used here.")
diff --git a/clean/video/recce/trainer/exp_tester.py b/clean/video/recce/trainer/exp_tester.py
new file mode 100644
index 0000000000000000000000000000000000000000..ee4f219f0ba60c6ded44e3b7c95eb7b21fc0be46
--- /dev/null
+++ b/clean/video/recce/trainer/exp_tester.py
@@ -0,0 +1,144 @@
+import os
+import sys
+import yaml
+import torch
+import random
+
+from tqdm import tqdm
+from pprint import pprint
+from torch.utils import data
+
+from dataset import load_dataset
+from loss import get_loss
+from model import load_model
+from model.common import freeze_weights
+from trainer import AbstractTrainer
+from trainer.utils import AccMeter, AUCMeter, AverageMeter, Logger, center_print
+
+
+class ExpTester(AbstractTrainer):
+ def __init__(self, config, stage="Test"):
+ super(ExpTester, self).__init__(config, stage)
+
+ if torch.cuda.is_available() and self.device is not None:
+ print(f"Using cuda device: {self.device}.")
+ self.gpu = True
+ self.model = self.model.to(self.device)
+ else:
+ print("Using cpu device.")
+ self.device = torch.device("cpu")
+
+ def _initiated_settings(self, model_cfg=None, data_cfg=None, config_cfg=None):
+ self.gpu = False
+ self.device = config_cfg.get("device", None)
+
+ def _train_settings(self, model_cfg=None, data_cfg=None, config_cfg=None):
+ # Not used.
+ raise NotImplementedError("The function is not intended to be used here.")
+
+ def _test_settings(self, model_cfg=None, data_cfg=None, config_cfg=None):
+ # load test dataset
+ test_dataset = data_cfg["file"]
+ branch = data_cfg["test_branch"]
+ name = data_cfg["name"]
+ with open(test_dataset, "r") as f:
+ options = yaml.load(f, Loader=yaml.FullLoader)
+ test_options = options[branch]
+ self.test_set = load_dataset(name)(test_options)
+ # wrapped with data loader
+ self.test_batch_size = data_cfg["test_batch_size"]
+ self.test_loader = data.DataLoader(self.test_set, shuffle=False,
+ batch_size=self.test_batch_size)
+ self.run_id = config_cfg["id"]
+ self.ckpt_fold = config_cfg.get("ckpt_fold", "runs")
+ self.dir = os.path.join(self.ckpt_fold, self.model_name, self.run_id)
+
+ # load model
+ self.num_classes = model_cfg["num_classes"]
+ self.model = load_model(self.model_name)(**model_cfg)
+
+ # load loss
+ self.loss_criterion = get_loss(config_cfg.get("loss", None))
+
+ # redirect the std out stream
+ sys.stdout = Logger(os.path.join(self.dir, "test_result.txt"))
+ print('Run dir: {}'.format(self.dir))
+
+ center_print('Test configurations begins')
+ pprint(self.config)
+ pprint(test_options)
+ center_print('Test configurations ends')
+
+ self.ckpt = config_cfg.get("ckpt", "best_model")
+ self._load_ckpt(best=True, train=False)
+
+ def _save_ckpt(self, step, best=False):
+ # Not used.
+ raise NotImplementedError("The function is not intended to be used here.")
+
+ def _load_ckpt(self, best=False, train=False):
+ load_dir = os.path.join(self.dir, self.ckpt + ".bin" if best else "latest_model.bin")
+ load_dict = torch.load(load_dir, map_location=self.device)
+ self.start_step = load_dict["step"]
+ self.best_step = load_dict["best_step"]
+ self.best_metric = load_dict.get("best_metric", None)
+ if self.best_metric is None:
+ self.best_metric = load_dict.get("best_acc")
+ self.eval_metric = load_dict.get("eval_metric", None)
+ if self.eval_metric is None:
+ self.eval_metric = load_dict.get("Acc")
+ self.model.load_state_dict(load_dict["model"])
+ print(f"Loading checkpoint from {load_dir}, best step: {self.best_step}, "
+ f"best {self.eval_metric}: {round(self.best_metric.item(), 4)}.")
+
+ def train(self):
+ # Not used.
+ raise NotImplementedError("The function is not intended to be used here.")
+
+ def validate(self, epoch, step, timer, writer):
+ # Not used.
+ raise NotImplementedError("The function is not intended to be used here.")
+
+ def test(self, display_images=False):
+ freeze_weights(self.model)
+ t_idx = random.randint(1, len(self.test_loader) + 1)
+ self.fixed_randomness() # for reproduction
+
+ acc = AccMeter()
+ auc = AUCMeter()
+ logloss = AverageMeter()
+ test_generator = tqdm(enumerate(self.test_loader, 1))
+ categories = self.test_loader.dataset.categories
+ for idx, test_data in test_generator:
+ self.model.eval()
+ I, Y = test_data
+ I = self.test_loader.dataset.load_item(I)
+ if self.gpu:
+ in_I, Y = self.to_device((I, Y))
+ else:
+ in_I, Y = (I, Y)
+ Y_pre = self.model(in_I)
+
+ # for BCE Setting:
+ if self.num_classes == 1:
+ Y_pre = Y_pre.squeeze()
+ loss = self.loss_criterion(Y_pre, Y.float())
+ Y_pre = torch.sigmoid(Y_pre)
+ else:
+ loss = self.loss_criterion(Y_pre, Y)
+
+ acc.update(Y_pre, Y, use_bce=self.num_classes == 1)
+ auc.update(Y_pre, Y, use_bce=self.num_classes == 1)
+ logloss.update(loss.item())
+
+ test_generator.set_description("Test %d/%d" % (idx, len(self.test_loader)))
+ if display_images and idx == t_idx:
+ # show images
+ images = I[:4]
+ pred = Y_pre[:4]
+ gt = Y[:4]
+ self.plot_figure(images, pred, gt, 2, categories)
+
+ print("Test, FINAL LOSS %.4f, FINAL ACC %.4f, FINAL AUC %.4f" %
+ (logloss.avg, acc.mean_acc(), auc.mean_auc()))
+ auc.curve(self.dir)
diff --git a/clean/video/recce/trainer/utils.py b/clean/video/recce/trainer/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..fcf1c4a864d34a18133667293c9dc76215af1011
--- /dev/null
+++ b/clean/video/recce/trainer/utils.py
@@ -0,0 +1,183 @@
+import os
+import sys
+import time
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+from collections import OrderedDict
+
+import numpy as np
+from sklearn.metrics import roc_auc_score, roc_curve
+from scipy.optimize import brentq
+from scipy.interpolate import interp1d
+
+# Tracking the path to the definition of the model.
+MODELS_PATH = {
+ "Recce": "model/network/Recce.py"
+}
+
+
+def exp_recons_loss(recons, x):
+ x, y = x
+ loss = torch.tensor(0., device=y.device)
+ real_index = torch.where(1 - y)[0]
+ for r in recons:
+ if real_index.numel() > 0:
+ real_x = torch.index_select(x, dim=0, index=real_index)
+ real_rec = torch.index_select(r, dim=0, index=real_index)
+ real_rec = F.interpolate(real_rec, size=x.shape[-2:], mode='bilinear', align_corners=True)
+ loss += torch.mean(torch.abs(real_rec - real_x))
+ return loss
+
+
+def center_print(content, around='*', repeat_around=10):
+ num = repeat_around
+ s = around
+ print(num * s + ' %s ' % content + num * s)
+
+
+def reduce_tensor(t):
+ rt = t.clone()
+ dist.all_reduce(rt)
+ rt /= float(dist.get_world_size())
+ return rt
+
+
+def tensor2image(tensor):
+ image = tensor.permute([1, 2, 0]).cpu().detach().numpy()
+ return (image - np.min(image)) / (np.max(image) - np.min(image))
+
+
+def state_dict(state_dict):
+ """ Remove 'module' keyword in state dictionary. """
+ weights = OrderedDict()
+ for k, v in state_dict.items():
+ weights.update({k.replace("module.", ""): v})
+ return weights
+
+
+class Logger(object):
+ def __init__(self, filename):
+ self.terminal = sys.stdout
+ self.log = open(filename, "a")
+
+ def write(self, message):
+ self.terminal.write(message)
+ self.log.write(message)
+ self.log.flush()
+
+ def flush(self):
+ pass
+
+
+class Timer(object):
+ """The class for timer."""
+
+ def __init__(self):
+ self.o = time.time()
+
+ def measure(self, p=1):
+ x = (time.time() - self.o) / p
+ x = int(x)
+ if x >= 3600:
+ return '{:.1f}h'.format(x / 3600)
+ if x >= 60:
+ return '{}m'.format(round(x / 60))
+ return '{}s'.format(x)
+
+
+class MLLoss(nn.Module):
+ def __init__(self):
+ super(MLLoss, self).__init__()
+
+ def forward(self, input, target, eps=1e-6):
+ # 0 - real; 1 - fake.
+ loss = torch.tensor(0., device=target.device)
+ batch_size = target.shape[0]
+ mat_1 = torch.hstack([target.unsqueeze(-1)] * batch_size)
+ mat_2 = torch.vstack([target] * batch_size)
+ diff_mat = torch.logical_xor(mat_1, mat_2).float()
+ or_mat = torch.logical_or(mat_1, mat_2)
+ eye = torch.eye(batch_size, device=target.device)
+ or_mat = torch.logical_or(or_mat, eye).float()
+ sim_mat = 1. - or_mat
+ for _ in input:
+ diff = torch.sum(_ * diff_mat, dim=[0, 1]) / (torch.sum(diff_mat, dim=[0, 1]) + eps)
+ sim = torch.sum(_ * sim_mat, dim=[0, 1]) / (torch.sum(sim_mat, dim=[0, 1]) + eps)
+ partial_loss = 1. - sim + diff
+ loss += max(partial_loss, torch.zeros_like(partial_loss))
+ return loss
+
+
+class AccMeter(object):
+ def __init__(self):
+ self.nums = 0
+ self.acc = 0
+
+ def reset(self):
+ self.nums = 0
+ self.acc = 0
+
+ def update(self, pred, target, use_bce=False):
+ if use_bce:
+ pred = (pred >= 0.5).int()
+ else:
+ pred = pred.argmax(1)
+ self.nums += target.shape[0]
+ self.acc += torch.sum(pred == target)
+
+ def mean_acc(self):
+ return self.acc / self.nums
+
+
+class AUCMeter(object):
+ def __init__(self):
+ self.score = None
+ self.true = None
+
+ def reset(self):
+ self.score = None
+ self.true = None
+
+ def update(self, score, true, use_bce=False):
+ if use_bce:
+ score = score.detach().cpu().numpy()
+ else:
+ score = torch.softmax(score.detach(), dim=-1)
+ score = torch.select(score, 1, 1).cpu().numpy()
+ true = true.flatten().cpu().numpy()
+ self.score = score if self.score is None else np.concatenate([self.score, score])
+ self.true = true if self.true is None else np.concatenate([self.true, true])
+
+ def mean_auc(self):
+ return roc_auc_score(self.true, self.score)
+
+ def curve(self, prefix):
+ fpr, tpr, thresholds = roc_curve(self.true, self.score, pos_label=1)
+ eer = brentq(lambda x: 1. - x - interp1d(fpr, tpr)(x), 0., 1.)
+ thresh = interp1d(fpr, thresholds)(eer)
+ print(f"# EER: {eer:.4f}(thresh: {thresh:.4f})")
+ torch.save([fpr, tpr, thresholds], os.path.join(prefix, "roc_curve.pickle"))
+
+
+class AverageMeter(object):
+ """Computes and stores the average and current value"""
+
+ def __init__(self):
+ self.val = 0
+ self.avg = 0
+ self.sum = 0
+ self.count = 0
+
+ def reset(self):
+ self.val = 0
+ self.avg = 0
+ self.sum = 0
+ self.count = 0
+
+ def update(self, val, n=1):
+ self.val = val
+ self.sum += val * n
+ self.count += n
+ self.avg = self.sum / self.count
diff --git a/clean/video/univfd_video/LICENSE b/clean/video/univfd_video/LICENSE
new file mode 100644
index 0000000000000000000000000000000000000000..7bcf6bcca73c769ede1f597bb73794e567039fb8
--- /dev/null
+++ b/clean/video/univfd_video/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Wisconsin AI and Vision Lab (WAIV)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/clean/video/univfd_video/README.md b/clean/video/univfd_video/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..393d091108d187521780a0c53f1260fa8967cad6
--- /dev/null
+++ b/clean/video/univfd_video/README.md
@@ -0,0 +1,110 @@
+# Detecting fake images
+
+**Towards Universal Fake Image Detectors that Generalize Across Generative Models**
+[Utkarsh Ojha*](https://utkarshojha.github.io/), [Yuheng Li*](https://yuheng-li.github.io/), [Yong Jae Lee](https://pages.cs.wisc.edu/~yongjaelee/)
+(*Equal contribution)
+CVPR 2023
+
+[[Project Page](https://utkarshojha.github.io/universal-fake-detection/)] [[Paper](https://arxiv.org/abs/2302.10174)]
+
+
+ >
+ Using images from one type of generative model (e.g., GAN), detect fake images from other breeds (e.g., Diffusion models)
+
+
+## Contents
+
+- [Setup](#setup)
+- [Pretrained model](#weights)
+- [Data](#data)
+- [Evaluation](#evaluation)
+- [Training](#training)
+
+
+## Setup
+
+1. Clone this repository
+```bash
+git clone https://github.com/Yuheng-Li/UniversalFakeDetect
+cd UniversalFakeDetect
+```
+
+2. Install the necessary libraries
+```bash
+pip install torch torchvision
+```
+
+## Data
+
+- Of the 19 models studied overall (Table 1/2 in the main paper), 11 are taken from a [previous work](https://arxiv.org/abs/1912.11035). Download the test set, i.e., real/fake images for those 11 models given by the authors from [here](https://drive.google.com/file/d/1z_fD3UKgWQyOTZIBbYSaQ-hz4AzUrLC1/view) (dataset size ~19GB).
+- Download the file and unzip it in `datasets/test`. You could also use the bash scripts provided by the authors, as described [here](https://github.com/PeterWang512/CNNDetection#download-the-dataset) in their code repository.
+- This should create a directory structure as follows:
+```
+
+datasets
+└── test
+ ├── progan
+ │── cyclegan
+ │── biggan
+ │ .
+ │ .
+
+```
+- Each directory (e.g., progan) will contain real/fake images under `0_real` and `1_fake` folders respectively.
+- Dataset for the diffusion models (e.g., LDM/Glide) can be found [here](https://drive.google.com/file/d/1FXlGIRh_Ud3cScMgSVDbEWmPDmjcrm1t/view?usp=drive_link). Note that in the paper (Table 2/3), we had reported the results over 10k randomly sampled images. Since providing that many images for all the domains will take up too much space, we are only releasing 1k images for each domain; i.e., 1k images fake images and 1k real images for each domain (e.g., LDM-200).
+- Download and unzip the file into `./diffusion_datasets` directory.
+
+
+## Evaluation
+
+- You can evaluate the model on all the dataset at once by running:
+```bash
+python validate.py --arch=CLIP:ViT-L/14 --ckpt=pretrained_weights/fc_weights.pth --result_folder=clip_vitl14
+```
+
+- You can also evaluate the model on one generative model by specifying the paths of real and fake datasets
+```bash
+python validate.py --arch=CLIP:ViT-L/14 --ckpt=pretrained_weights/fc_weights.pth --result_folder=clip_vitl14 --real_path datasets/test/progan/0_real --fake_path datasets/test/progan/1_fake
+```
+
+Note that if no arguments are provided for `real_path` and `fake_path`, the script will perform the evaluation on all the domains specified in `dataset_paths.py`.
+
+- The results will be stored in `results/` in two files: `ap.txt` stores the Average Prevision for each of the test domains, and `acc.txt` stores the accuracy (with 0.5 as the threshold) for the same domains.
+
+## Training
+
+- Our main model is trained on the same dataset used by the authors of [this work](https://arxiv.org/abs/1912.11035). Download the official training dataset provided [here](https://drive.google.com/file/d/1iVNBV0glknyTYGA9bCxT_d0CVTOgGcKh/view) (dataset size ~ 72GB).
+
+- Download and unzip the dataset in `datasets/train` directory. The overall structure should look like the following:
+```
+datasets
+└── train
+ └── progan
+ ├── airplane
+ │── bird
+ │── boat
+ │ .
+ │ .
+```
+- A total of 20 different object categories, with each folder containing the corresponding real and fake images in `0_real` and `1_fake` folders.
+- The model can then be trained with the following command:
+```bash
+python train.py --name=clip_vitl14 --wang2020_data_path=datasets/ --data_mode=wang2020 --arch=CLIP:ViT-L/14 --fix_backbone
+```
+- **Important**: do not forget to use the `--fix_backbone` argument during training, which makes sure that the only the linear layer's parameters will be trained.
+
+## Acknowledgement
+
+We would like to thank [Sheng-Yu Wang](https://github.com/PeterWang512) for releasing the real/fake images from different generative models. Our training pipeline is also inspired by his [open-source code](https://github.com/PeterWang512/CNNDetection). We would also like to thank [CompVis](https://github.com/CompVis) for releasing the pre-trained [LDMs](https://github.com/CompVis/latent-diffusion) and [LAION](https://laion.ai/) for open-sourcing [LAION-400M dataset](https://laion.ai/blog/laion-400-open-dataset/).
+
+## Citation
+
+If you find our work helpful in your research, please cite it using the following:
+```bibtex
+@inproceedings{ojha2023fakedetect,
+ title={Towards Universal Fake Image Detectors that Generalize Across Generative Models},
+ author={Ojha, Utkarsh and Li, Yuheng and Lee, Yong Jae},
+ booktitle={CVPR},
+ year={2023},
+}
+```
diff --git a/clean/video/univfd_video/SOURCE.md b/clean/video/univfd_video/SOURCE.md
new file mode 100644
index 0000000000000000000000000000000000000000..616eb7fb3256d53790d32235fdf82c9da84153c0
--- /dev/null
+++ b/clean/video/univfd_video/SOURCE.md
@@ -0,0 +1,20 @@
+# Source: video/univfd_video
+
+| Field | Value |
+|---|---|
+| Upstream | **UNVERIFIED** -- provenance was lost when this code was vendored |
+| Paper | not recorded |
+| Commit SHA | **not recorded** -- the vendoring step did not preserve it |
+| Mirrored on | 2026-09-15 |
+| Upstream license | LICENSE |
+
+This is a **mirror**, stripped to the files needed for inference. The full
+untouched snapshot is at `archive/video__univfd_video.tar.gz`.
+
+This code is the work of its original authors and is **not** covered by the
+DeepSafe project license. If you are an author and want this removed, open an
+issue on https://github.com/deepsafehq/deepsafe-bench and it will be taken down
+within 48 hours, no questions asked.
+
+**Alias.** This model shares its code with `image/universal`; only the
+weights differ. No separate archive is published.
diff --git a/clean/video/univfd_video/dataset_paths.py b/clean/video/univfd_video/dataset_paths.py
new file mode 100644
index 0000000000000000000000000000000000000000..3e1c8acfecfda1f209e29cb109e1ad26c64a0ab4
--- /dev/null
+++ b/clean/video/univfd_video/dataset_paths.py
@@ -0,0 +1,153 @@
+DATASET_PATHS = [
+
+
+ dict(
+ real_path='../FAKE_IMAGES/CNN/test/progan',
+ fake_path='../FAKE_IMAGES/CNN/test/progan',
+ data_mode='wang2020',
+ key='progan'
+ ),
+
+ dict(
+ real_path='../FAKE_IMAGES/CNN/test/cyclegan',
+ fake_path='../FAKE_IMAGES/CNN/test/cyclegan',
+ data_mode='wang2020',
+ key='cyclegan'
+ ),
+
+ dict(
+ real_path='../FAKE_IMAGES/CNN/test/biggan/', # Imagenet
+ fake_path='../FAKE_IMAGES/CNN/test/biggan/',
+ data_mode='wang2020',
+ key='biggan'
+ ),
+
+
+ dict(
+ real_path='../FAKE_IMAGES/CNN/test/stylegan',
+ fake_path='../FAKE_IMAGES/CNN/test/stylegan',
+ data_mode='wang2020',
+ key='stylegan'
+ ),
+
+
+ dict(
+ real_path='../FAKE_IMAGES/CNN/test/gaugan', # It is COCO
+ fake_path='../FAKE_IMAGES/CNN/test/gaugan',
+ data_mode='wang2020',
+ key='gaugan'
+ ),
+
+
+ dict(
+ real_path='../FAKE_IMAGES/CNN/test/stargan',
+ fake_path='../FAKE_IMAGES/CNN/test/stargan',
+ data_mode='wang2020',
+ key='stargan'
+ ),
+
+
+ dict(
+ real_path='../FAKE_IMAGES/CNN/test/deepfake',
+ fake_path='../FAKE_IMAGES/CNN/test/deepfake',
+ data_mode='wang2020',
+ key='deepfake'
+ ),
+
+
+ dict(
+ real_path='../FAKE_IMAGES/CNN/test/seeingdark',
+ fake_path='../FAKE_IMAGES/CNN/test/seeingdark',
+ data_mode='wang2020',
+ key='sitd'
+ ),
+
+
+ dict(
+ real_path='../FAKE_IMAGES/CNN/test/san',
+ fake_path='../FAKE_IMAGES/CNN/test/san',
+ data_mode='wang2020',
+ key='san'
+ ),
+
+
+ dict(
+ real_path='../FAKE_IMAGES/CNN/test/crn', # Images from some video games
+ fake_path='../FAKE_IMAGES/CNN/test/crn',
+ data_mode='wang2020',
+ key='crn'
+ ),
+
+
+ dict(
+ real_path='../FAKE_IMAGES/CNN/test/imle', # Images from some video games
+ fake_path='../FAKE_IMAGES/CNN/test/imle',
+ data_mode='wang2020',
+ key='imle'
+ ),
+
+
+ dict(
+ real_path='./diffusion_datasets/imagenet',
+ fake_path='./diffusion_datasets/guided',
+ data_mode='wang2020',
+ key='guided'
+ ),
+
+
+ dict(
+ real_path='./diffusion_datasets/laion',
+ fake_path='./diffusion_datasets/ldm_200',
+ data_mode='wang2020',
+ key='ldm_200'
+ ),
+
+ dict(
+ real_path='./diffusion_datasets/laion',
+ fake_path='./diffusion_datasets/ldm_200_cfg',
+ data_mode='wang2020',
+ key='ldm_200_cfg'
+ ),
+
+ dict(
+ real_path='./diffusion_datasets/laion',
+ fake_path='./diffusion_datasets/ldm_100',
+ data_mode='wang2020',
+ key='ldm_100'
+ ),
+
+
+ dict(
+ real_path='./diffusion_datasets/laion',
+ fake_path='./diffusion_datasets/glide_100_27',
+ data_mode='wang2020',
+ key='glide_100_27'
+ ),
+
+
+ dict(
+ real_path='./diffusion_datasets/laion',
+ fake_path='./diffusion_datasets/glide_50_27',
+ data_mode='wang2020',
+ key='glide_50_27'
+ ),
+
+
+ dict(
+ real_path='./diffusion_datasets/laion',
+ fake_path='./diffusion_datasets/glide_100_10',
+ data_mode='wang2020',
+ key='glide_100_10'
+ ),
+
+
+ dict(
+ real_path='./diffusion_datasets/laion',
+ fake_path='./diffusion_datasets/dalle',
+ data_mode='wang2020',
+ key='dalle'
+ ),
+
+
+
+]
diff --git a/clean/video/univfd_video/models/__init__.py b/clean/video/univfd_video/models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d7b790b40d29d1f3bf02f398f3522eea8e4c2c22
--- /dev/null
+++ b/clean/video/univfd_video/models/__init__.py
@@ -0,0 +1,43 @@
+from .clip_models import CLIPModel
+from .imagenet_models import ImagenetModel
+
+
+VALID_NAMES = [
+ 'Imagenet:resnet18',
+ 'Imagenet:resnet34',
+ 'Imagenet:resnet50',
+ 'Imagenet:resnet101',
+ 'Imagenet:resnet152',
+ 'Imagenet:vgg11',
+ 'Imagenet:vgg19',
+ 'Imagenet:swin-b',
+ 'Imagenet:swin-s',
+ 'Imagenet:swin-t',
+ 'Imagenet:vit_b_16',
+ 'Imagenet:vit_b_32',
+ 'Imagenet:vit_l_16',
+ 'Imagenet:vit_l_32',
+
+ 'CLIP:RN50',
+ 'CLIP:RN101',
+ 'CLIP:RN50x4',
+ 'CLIP:RN50x16',
+ 'CLIP:RN50x64',
+ 'CLIP:ViT-B/32',
+ 'CLIP:ViT-B/16',
+ 'CLIP:ViT-L/14',
+ 'CLIP:ViT-L/14@336px',
+]
+
+
+
+
+
+def get_model(name):
+ assert name in VALID_NAMES
+ if name.startswith("Imagenet:"):
+ return ImagenetModel(name[9:])
+ elif name.startswith("CLIP:"):
+ return CLIPModel(name[5:])
+ else:
+ assert False
diff --git a/clean/video/univfd_video/models/clip/__init__.py b/clean/video/univfd_video/models/clip/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..dcc5619538c0f7c782508bdbd9587259d805e0d9
--- /dev/null
+++ b/clean/video/univfd_video/models/clip/__init__.py
@@ -0,0 +1 @@
+from .clip import *
diff --git a/clean/video/univfd_video/models/clip/bpe_simple_vocab_16e6.txt.gz b/clean/video/univfd_video/models/clip/bpe_simple_vocab_16e6.txt.gz
new file mode 100644
index 0000000000000000000000000000000000000000..36a15856e00a06a9fbed8cdd34d2393fea4a3113
--- /dev/null
+++ b/clean/video/univfd_video/models/clip/bpe_simple_vocab_16e6.txt.gz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a
+size 1356917
diff --git a/clean/video/univfd_video/models/clip/clip.py b/clean/video/univfd_video/models/clip/clip.py
new file mode 100644
index 0000000000000000000000000000000000000000..257511e1d40c120e0d64a0f1562d44b2b8a40a17
--- /dev/null
+++ b/clean/video/univfd_video/models/clip/clip.py
@@ -0,0 +1,237 @@
+import hashlib
+import os
+import urllib
+import warnings
+from typing import Any, Union, List
+from pkg_resources import packaging
+
+import torch
+from PIL import Image
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+from tqdm import tqdm
+
+from .model import build_model
+from .simple_tokenizer import SimpleTokenizer as _Tokenizer
+
+try:
+ from torchvision.transforms import InterpolationMode
+ BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+ BICUBIC = Image.BICUBIC
+
+
+if packaging.version.parse(torch.__version__) < packaging.version.parse("1.7.1"):
+ warnings.warn("PyTorch version 1.7.1 or higher is recommended")
+
+
+__all__ = ["available_models", "load", "tokenize"]
+_tokenizer = _Tokenizer()
+
+_MODELS = {
+ "RN50": "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt",
+ "RN101": "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt",
+ "RN50x4": "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt",
+ "RN50x16": "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt",
+ "RN50x64": "https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt",
+ "ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt",
+ "ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt",
+ "ViT-L/14": "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt",
+ "ViT-L/14@336px": "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt",
+}
+
+
+def _download(url: str, root: str):
+ os.makedirs(root, exist_ok=True)
+ filename = os.path.basename(url)
+
+ expected_sha256 = url.split("/")[-2]
+ download_target = os.path.join(root, filename)
+
+ if os.path.exists(download_target) and not os.path.isfile(download_target):
+ raise RuntimeError(f"{download_target} exists and is not a regular file")
+
+ if os.path.isfile(download_target):
+ if hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256:
+ return download_target
+ else:
+ warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
+
+ with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+ with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True, unit_divisor=1024) as loop:
+ while True:
+ buffer = source.read(8192)
+ if not buffer:
+ break
+
+ output.write(buffer)
+ loop.update(len(buffer))
+
+ if hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256:
+ raise RuntimeError("Model has been downloaded but the SHA256 checksum does not not match")
+
+ return download_target
+
+
+def _convert_image_to_rgb(image):
+ return image.convert("RGB")
+
+
+def _transform(n_px):
+ return Compose([
+ Resize(n_px, interpolation=BICUBIC),
+ CenterCrop(n_px),
+ _convert_image_to_rgb,
+ ToTensor(),
+ Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+ ])
+
+
+def available_models() -> List[str]:
+ """Returns the names of available CLIP models"""
+ return list(_MODELS.keys())
+
+
+def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit: bool = False, download_root: str = None):
+ """Load a CLIP model
+
+ Parameters
+ ----------
+ name : str
+ A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
+
+ device : Union[str, torch.device]
+ The device to put the loaded model
+
+ jit : bool
+ Whether to load the optimized JIT model or more hackable non-JIT model (default).
+
+ download_root: str
+ path to download the model files; by default, it uses "~/.cache/clip"
+
+ Returns
+ -------
+ model : torch.nn.Module
+ The CLIP model
+
+ preprocess : Callable[[PIL.Image], torch.Tensor]
+ A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
+ """
+ if name in _MODELS:
+ model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/clip"))
+ elif os.path.isfile(name):
+ model_path = name
+ else:
+ raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
+
+ with open(model_path, 'rb') as opened_file:
+ try:
+ # loading JIT archive
+ model = torch.jit.load(opened_file, map_location=device if jit else "cpu").eval()
+ state_dict = None
+ except RuntimeError:
+ # loading saved state dict
+ if jit:
+ warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
+ jit = False
+ state_dict = torch.load(opened_file, map_location="cpu")
+
+ if not jit:
+ model = build_model(state_dict or model.state_dict()).to(device)
+ if str(device) == "cpu":
+ model.float()
+ return model, _transform(model.visual.input_resolution)
+
+ # patch the device names
+ device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])
+ device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]
+
+ def patch_device(module):
+ try:
+ graphs = [module.graph] if hasattr(module, "graph") else []
+ except RuntimeError:
+ graphs = []
+
+ if hasattr(module, "forward1"):
+ graphs.append(module.forward1.graph)
+
+ for graph in graphs:
+ for node in graph.findAllNodes("prim::Constant"):
+ if "value" in node.attributeNames() and str(node["value"]).startswith("cuda"):
+ node.copyAttributes(device_node)
+
+ model.apply(patch_device)
+ patch_device(model.encode_image)
+ patch_device(model.encode_text)
+
+ # patch dtype to float32 on CPU
+ if str(device) == "cpu":
+ float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])
+ float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
+ float_node = float_input.node()
+
+ def patch_float(module):
+ try:
+ graphs = [module.graph] if hasattr(module, "graph") else []
+ except RuntimeError:
+ graphs = []
+
+ if hasattr(module, "forward1"):
+ graphs.append(module.forward1.graph)
+
+ for graph in graphs:
+ for node in graph.findAllNodes("aten::to"):
+ inputs = list(node.inputs())
+ for i in [1, 2]: # dtype can be the second or third argument to aten::to()
+ if inputs[i].node()["value"] == 5:
+ inputs[i].node().copyAttributes(float_node)
+
+ model.apply(patch_float)
+ patch_float(model.encode_image)
+ patch_float(model.encode_text)
+
+ model.float()
+
+ return model, _transform(model.input_resolution.item())
+
+
+def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> Union[torch.IntTensor, torch.LongTensor]:
+ """
+ Returns the tokenized representation of given input string(s)
+
+ Parameters
+ ----------
+ texts : Union[str, List[str]]
+ An input string or a list of input strings to tokenize
+
+ context_length : int
+ The context length to use; all CLIP models use 77 as the context length
+
+ truncate: bool
+ Whether to truncate the text in case its encoding is longer than the context length
+
+ Returns
+ -------
+ A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length].
+ We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long.
+ """
+ if isinstance(texts, str):
+ texts = [texts]
+
+ sot_token = _tokenizer.encoder["<|startoftext|>"]
+ eot_token = _tokenizer.encoder["<|endoftext|>"]
+ all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
+ if packaging.version.parse(torch.__version__) < packaging.version.parse("1.8.0"):
+ result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+ else:
+ result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)
+
+ for i, tokens in enumerate(all_tokens):
+ if len(tokens) > context_length:
+ if truncate:
+ tokens = tokens[:context_length]
+ tokens[-1] = eot_token
+ else:
+ raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}")
+ result[i, :len(tokens)] = torch.tensor(tokens)
+
+ return result
diff --git a/clean/video/univfd_video/models/clip/model.py b/clean/video/univfd_video/models/clip/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..c60bda69ab0d35af0b64517d32595d9c03f8721c
--- /dev/null
+++ b/clean/video/univfd_video/models/clip/model.py
@@ -0,0 +1,452 @@
+from collections import OrderedDict
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Bottleneck(nn.Module):
+ expansion = 4
+
+ def __init__(self, inplanes, planes, stride=1):
+ super().__init__()
+
+ # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+ self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+ self.bn1 = nn.BatchNorm2d(planes)
+ self.relu1 = nn.ReLU(inplace=True)
+
+ self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+ self.bn2 = nn.BatchNorm2d(planes)
+ self.relu2 = nn.ReLU(inplace=True)
+
+ self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+ self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+ self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+ self.relu3 = nn.ReLU(inplace=True)
+
+ self.downsample = None
+ self.stride = stride
+
+ if stride > 1 or inplanes != planes * Bottleneck.expansion:
+ # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+ self.downsample = nn.Sequential(OrderedDict([
+ ("-1", nn.AvgPool2d(stride)),
+ ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+ ("1", nn.BatchNorm2d(planes * self.expansion))
+ ]))
+
+ def forward(self, x: torch.Tensor):
+ identity = x
+
+ out = self.relu1(self.bn1(self.conv1(x)))
+ out = self.relu2(self.bn2(self.conv2(out)))
+ out = self.avgpool(out)
+ out = self.bn3(self.conv3(out))
+
+ if self.downsample is not None:
+ identity = self.downsample(x)
+
+ out += identity
+ out = self.relu3(out)
+ return out
+
+
+class AttentionPool2d(nn.Module):
+ def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+ super().__init__()
+ self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
+ self.k_proj = nn.Linear(embed_dim, embed_dim)
+ self.q_proj = nn.Linear(embed_dim, embed_dim)
+ self.v_proj = nn.Linear(embed_dim, embed_dim)
+ self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+ self.num_heads = num_heads
+
+ def forward(self, x):
+ x = x.flatten(start_dim=2).permute(2, 0, 1) # NCHW -> (HW)NC
+ x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0) # (HW+1)NC
+ x = x + self.positional_embedding[:, None, :].to(x.dtype) # (HW+1)NC
+ x, _ = F.multi_head_attention_forward(
+ query=x[:1], key=x, value=x,
+ embed_dim_to_check=x.shape[-1],
+ num_heads=self.num_heads,
+ q_proj_weight=self.q_proj.weight,
+ k_proj_weight=self.k_proj.weight,
+ v_proj_weight=self.v_proj.weight,
+ in_proj_weight=None,
+ in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+ bias_k=None,
+ bias_v=None,
+ add_zero_attn=False,
+ dropout_p=0,
+ out_proj_weight=self.c_proj.weight,
+ out_proj_bias=self.c_proj.bias,
+ use_separate_proj_weight=True,
+ training=self.training,
+ need_weights=False
+ )
+ return x.squeeze(0)
+
+
+class ModifiedResNet(nn.Module):
+ """
+ A ResNet class that is similar to torchvision's but contains the following changes:
+ - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+ - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+ - The final pooling layer is a QKV attention instead of an average pool
+ """
+
+ def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
+ super().__init__()
+ self.output_dim = output_dim
+ self.input_resolution = input_resolution
+
+ # the 3-layer stem
+ self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(width // 2)
+ self.relu1 = nn.ReLU(inplace=True)
+ self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+ self.bn2 = nn.BatchNorm2d(width // 2)
+ self.relu2 = nn.ReLU(inplace=True)
+ self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+ self.bn3 = nn.BatchNorm2d(width)
+ self.relu3 = nn.ReLU(inplace=True)
+ self.avgpool = nn.AvgPool2d(2)
+
+ # residual layers
+ self._inplanes = width # this is a *mutable* variable used during construction
+ self.layer1 = self._make_layer(width, layers[0])
+ self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+ self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+ self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+
+ embed_dim = width * 32 # the ResNet feature dimension
+ self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim)
+
+ def _make_layer(self, planes, blocks, stride=1):
+ layers = [Bottleneck(self._inplanes, planes, stride)]
+
+ self._inplanes = planes * Bottleneck.expansion
+ for _ in range(1, blocks):
+ layers.append(Bottleneck(self._inplanes, planes))
+
+ return nn.Sequential(*layers)
+
+ def forward(self, x):
+ def stem(x):
+ x = self.relu1(self.bn1(self.conv1(x)))
+ x = self.relu2(self.bn2(self.conv2(x)))
+ x = self.relu3(self.bn3(self.conv3(x)))
+ x = self.avgpool(x)
+ return x
+
+ x = x.type(self.conv1.weight.dtype)
+ x = stem(x)
+ x = self.layer1(x)
+ x = self.layer2(x)
+ x = self.layer3(x)
+ x = self.layer4(x)
+ x = self.attnpool(x)
+
+ return x
+
+
+class LayerNorm(nn.LayerNorm):
+ """Subclass torch's LayerNorm to handle fp16."""
+
+ def forward(self, x: torch.Tensor):
+ orig_type = x.dtype
+ ret = super().forward(x.type(torch.float32))
+ return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+ def forward(self, x: torch.Tensor):
+ return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+ def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
+ super().__init__()
+
+ self.attn = nn.MultiheadAttention(d_model, n_head)
+ self.ln_1 = LayerNorm(d_model)
+ self.mlp = nn.Sequential(OrderedDict([
+ ("c_fc", nn.Linear(d_model, d_model * 4)),
+ ("gelu", QuickGELU()),
+ ("c_proj", nn.Linear(d_model * 4, d_model))
+ ]))
+ self.ln_2 = LayerNorm(d_model)
+ self.attn_mask = attn_mask
+
+ def attention(self, x: torch.Tensor):
+ self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+ return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+ def forward(self, x: torch.Tensor):
+ x = x + self.attention(self.ln_1(x))
+ x = x + self.mlp(self.ln_2(x))
+ return x
+
+
+class Transformer(nn.Module):
+ def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
+ super().__init__()
+ self.width = width
+ self.layers = layers
+ self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])
+
+ def forward(self, x: torch.Tensor):
+ out = {}
+ for idx, layer in enumerate(self.resblocks.children()):
+ x = layer(x)
+ out['layer'+str(idx)] = x[0] # shape:LND. choose cls token feature
+ return out, x
+
+ # return self.resblocks(x) # This is the original code
+
+
+class VisionTransformer(nn.Module):
+ def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int):
+ super().__init__()
+ self.input_resolution = input_resolution
+ self.output_dim = output_dim
+ self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
+
+ scale = width ** -0.5
+ self.class_embedding = nn.Parameter(scale * torch.randn(width))
+ self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
+ self.ln_pre = LayerNorm(width)
+
+ self.transformer = Transformer(width, layers, heads)
+
+ self.ln_post = LayerNorm(width)
+ self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+
+
+
+ def forward(self, x: torch.Tensor):
+ x = self.conv1(x) # shape = [*, width, grid, grid]
+ x = x.reshape(x.shape[0], x.shape[1], -1) # shape = [*, width, grid ** 2]
+ x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width]
+ x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1) # shape = [*, grid ** 2 + 1, width]
+ x = x + self.positional_embedding.to(x.dtype)
+ x = self.ln_pre(x)
+
+ x = x.permute(1, 0, 2) # NLD -> LND
+ out, x = self.transformer(x)
+ x = x.permute(1, 0, 2) # LND -> NLD
+
+ x = self.ln_post(x[:, 0, :])
+
+
+ out['before_projection'] = x
+
+ if self.proj is not None:
+ x = x @ self.proj
+ out['after_projection'] = x
+
+ # Return both intermediate features and final clip feature
+ # return out
+
+ # This only returns CLIP features
+ return x
+
+
+class CLIP(nn.Module):
+ def __init__(self,
+ embed_dim: int,
+ # vision
+ image_resolution: int,
+ vision_layers: Union[Tuple[int, int, int, int], int],
+ vision_width: int,
+ vision_patch_size: int,
+ # text
+ context_length: int,
+ vocab_size: int,
+ transformer_width: int,
+ transformer_heads: int,
+ transformer_layers: int
+ ):
+ super().__init__()
+
+ self.context_length = context_length
+
+ if isinstance(vision_layers, (tuple, list)):
+ vision_heads = vision_width * 32 // 64
+ self.visual = ModifiedResNet(
+ layers=vision_layers,
+ output_dim=embed_dim,
+ heads=vision_heads,
+ input_resolution=image_resolution,
+ width=vision_width
+ )
+ else:
+ vision_heads = vision_width // 64
+ self.visual = VisionTransformer(
+ input_resolution=image_resolution,
+ patch_size=vision_patch_size,
+ width=vision_width,
+ layers=vision_layers,
+ heads=vision_heads,
+ output_dim=embed_dim
+ )
+
+ self.transformer = Transformer(
+ width=transformer_width,
+ layers=transformer_layers,
+ heads=transformer_heads,
+ attn_mask=self.build_attention_mask()
+ )
+
+ self.vocab_size = vocab_size
+ self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+ self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
+ self.ln_final = LayerNorm(transformer_width)
+
+ self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
+ self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+ self.initialize_parameters()
+
+ def initialize_parameters(self):
+ nn.init.normal_(self.token_embedding.weight, std=0.02)
+ nn.init.normal_(self.positional_embedding, std=0.01)
+
+ if isinstance(self.visual, ModifiedResNet):
+ if self.visual.attnpool is not None:
+ std = self.visual.attnpool.c_proj.in_features ** -0.5
+ nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
+ nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
+ nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
+ nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
+
+ for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]:
+ for name, param in resnet_block.named_parameters():
+ if name.endswith("bn3.weight"):
+ nn.init.zeros_(param)
+
+ proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+ attn_std = self.transformer.width ** -0.5
+ fc_std = (2 * self.transformer.width) ** -0.5
+ for block in self.transformer.resblocks:
+ nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+ nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+ nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+ nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+ if self.text_projection is not None:
+ nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+
+ def build_attention_mask(self):
+ # lazily create causal attention mask, with full attention between the vision tokens
+ # pytorch uses additive attention mask; fill with -inf
+ mask = torch.empty(self.context_length, self.context_length)
+ mask.fill_(float("-inf"))
+ mask.triu_(1) # zero out the lower diagonal
+ return mask
+
+ @property
+ def dtype(self):
+ return self.visual.conv1.weight.dtype
+
+ def encode_image(self, image):
+ return self.visual(image.type(self.dtype))
+
+ def encode_text(self, text):
+ x = self.token_embedding(text).type(self.dtype) # [batch_size, n_ctx, d_model]
+
+ x = x + self.positional_embedding.type(self.dtype)
+ x = x.permute(1, 0, 2) # NLD -> LND
+ x = self.transformer(x)
+ x = x.permute(1, 0, 2) # LND -> NLD
+ x = self.ln_final(x).type(self.dtype)
+
+ # x.shape = [batch_size, n_ctx, transformer.width]
+ # take features from the eot embedding (eot_token is the highest number in each sequence)
+ x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+
+ return x
+
+ def forward(self, image, text):
+ image_features = self.encode_image(image)
+ text_features = self.encode_text(text)
+
+ # normalized features
+ image_features = image_features / image_features.norm(dim=1, keepdim=True)
+ text_features = text_features / text_features.norm(dim=1, keepdim=True)
+
+ # cosine similarity as logits
+ logit_scale = self.logit_scale.exp()
+ logits_per_image = logit_scale * image_features @ text_features.t()
+ logits_per_text = logits_per_image.t()
+
+ # shape = [global_batch_size, global_batch_size]
+ return logits_per_image, logits_per_text
+
+
+def convert_weights(model: nn.Module):
+ """Convert applicable model parameters to fp16"""
+
+ def _convert_weights_to_fp16(l):
+ if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+ l.weight.data = l.weight.data.half()
+ if l.bias is not None:
+ l.bias.data = l.bias.data.half()
+
+ if isinstance(l, nn.MultiheadAttention):
+ for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+ tensor = getattr(l, attr)
+ if tensor is not None:
+ tensor.data = tensor.data.half()
+
+ for name in ["text_projection", "proj"]:
+ if hasattr(l, name):
+ attr = getattr(l, name)
+ if attr is not None:
+ attr.data = attr.data.half()
+
+ model.apply(_convert_weights_to_fp16)
+
+
+def build_model(state_dict: dict):
+ vit = "visual.proj" in state_dict
+
+ if vit:
+ vision_width = state_dict["visual.conv1.weight"].shape[0]
+ vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+ vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+ grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+ image_resolution = vision_patch_size * grid_size
+ else:
+ counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+ vision_layers = tuple(counts)
+ vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+ output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+ vision_patch_size = None
+ assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+ image_resolution = output_width * 32
+
+ embed_dim = state_dict["text_projection"].shape[1]
+ context_length = state_dict["positional_embedding"].shape[0]
+ vocab_size = state_dict["token_embedding.weight"].shape[0]
+ transformer_width = state_dict["ln_final.weight"].shape[0]
+ transformer_heads = transformer_width // 64
+ transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks")))
+
+ model = CLIP(
+ embed_dim,
+ image_resolution, vision_layers, vision_width, vision_patch_size,
+ context_length, vocab_size, transformer_width, transformer_heads, transformer_layers
+ )
+
+ for key in ["input_resolution", "context_length", "vocab_size"]:
+ if key in state_dict:
+ del state_dict[key]
+
+ convert_weights(model)
+ model.load_state_dict(state_dict)
+ return model.eval()
diff --git a/clean/video/univfd_video/models/clip/simple_tokenizer.py b/clean/video/univfd_video/models/clip/simple_tokenizer.py
new file mode 100644
index 0000000000000000000000000000000000000000..0a66286b7d5019c6e221932a813768038f839c91
--- /dev/null
+++ b/clean/video/univfd_video/models/clip/simple_tokenizer.py
@@ -0,0 +1,132 @@
+import gzip
+import html
+import os
+from functools import lru_cache
+
+import ftfy
+import regex as re
+
+
+@lru_cache()
+def default_bpe():
+ return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz")
+
+
+@lru_cache()
+def bytes_to_unicode():
+ """
+ Returns list of utf-8 byte and a corresponding list of unicode strings.
+ The reversible bpe codes work on unicode strings.
+ This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
+ When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
+ This is a signficant percentage of your normal, say, 32K bpe vocab.
+ To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
+ And avoids mapping to whitespace/control characters the bpe code barfs on.
+ """
+ bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
+ cs = bs[:]
+ n = 0
+ for b in range(2**8):
+ if b not in bs:
+ bs.append(b)
+ cs.append(2**8+n)
+ n += 1
+ cs = [chr(n) for n in cs]
+ return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+ """Return set of symbol pairs in a word.
+ Word is represented as tuple of symbols (symbols being variable-length strings).
+ """
+ pairs = set()
+ prev_char = word[0]
+ for char in word[1:]:
+ pairs.add((prev_char, char))
+ prev_char = char
+ return pairs
+
+
+def basic_clean(text):
+ text = ftfy.fix_text(text)
+ text = html.unescape(html.unescape(text))
+ return text.strip()
+
+
+def whitespace_clean(text):
+ text = re.sub(r'\s+', ' ', text)
+ text = text.strip()
+ return text
+
+
+class SimpleTokenizer(object):
+ def __init__(self, bpe_path: str = default_bpe()):
+ self.byte_encoder = bytes_to_unicode()
+ self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+ merges = gzip.open(bpe_path).read().decode("utf-8").split('\n')
+ merges = merges[1:49152-256-2+1]
+ merges = [tuple(merge.split()) for merge in merges]
+ vocab = list(bytes_to_unicode().values())
+ vocab = vocab + [v+'' for v in vocab]
+ for merge in merges:
+ vocab.append(''.join(merge))
+ vocab.extend(['<|startoftext|>', '<|endoftext|>'])
+ self.encoder = dict(zip(vocab, range(len(vocab))))
+ self.decoder = {v: k for k, v in self.encoder.items()}
+ self.bpe_ranks = dict(zip(merges, range(len(merges))))
+ self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
+ self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE)
+
+ def bpe(self, token):
+ if token in self.cache:
+ return self.cache[token]
+ word = tuple(token[:-1]) + ( token[-1] + '',)
+ pairs = get_pairs(word)
+
+ if not pairs:
+ return token+''
+
+ while True:
+ bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+ if bigram not in self.bpe_ranks:
+ break
+ first, second = bigram
+ new_word = []
+ i = 0
+ while i < len(word):
+ try:
+ j = word.index(first, i)
+ new_word.extend(word[i:j])
+ i = j
+ except:
+ new_word.extend(word[i:])
+ break
+
+ if word[i] == first and i < len(word)-1 and word[i+1] == second:
+ new_word.append(first+second)
+ i += 2
+ else:
+ new_word.append(word[i])
+ i += 1
+ new_word = tuple(new_word)
+ word = new_word
+ if len(word) == 1:
+ break
+ else:
+ pairs = get_pairs(word)
+ word = ' '.join(word)
+ self.cache[token] = word
+ return word
+
+ def encode(self, text):
+ bpe_tokens = []
+ text = whitespace_clean(basic_clean(text)).lower()
+ for token in re.findall(self.pat, text):
+ token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
+ bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
+ return bpe_tokens
+
+ def decode(self, tokens):
+ text = ''.join([self.decoder[token] for token in tokens])
+ text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('', ' ')
+ return text
diff --git a/clean/video/univfd_video/models/clip_models.py b/clean/video/univfd_video/models/clip_models.py
new file mode 100644
index 0000000000000000000000000000000000000000..44de4db8ee4dd77690970fc76eb671b1b3a43882
--- /dev/null
+++ b/clean/video/univfd_video/models/clip_models.py
@@ -0,0 +1,24 @@
+from .clip import clip
+from PIL import Image
+import torch.nn as nn
+
+
+CHANNELS = {
+ "RN50" : 1024,
+ "ViT-L/14" : 768
+}
+
+class CLIPModel(nn.Module):
+ def __init__(self, name, num_classes=1):
+ super(CLIPModel, self).__init__()
+
+ self.model, self.preprocess = clip.load(name, device="cpu") # self.preprecess will not be used during training, which is handled in Dataset class
+ self.fc = nn.Linear( CHANNELS[name], num_classes )
+
+
+ def forward(self, x, return_feature=False):
+ features = self.model.encode_image(x)
+ if return_feature:
+ return features
+ return self.fc(features)
+
diff --git a/clean/video/univfd_video/models/imagenet_models.py b/clean/video/univfd_video/models/imagenet_models.py
new file mode 100644
index 0000000000000000000000000000000000000000..20a40b916793d926c915aa2f62602651613fec04
--- /dev/null
+++ b/clean/video/univfd_video/models/imagenet_models.py
@@ -0,0 +1,40 @@
+from .resnet import resnet18, resnet34, resnet50, resnet101, resnet152
+from .vision_transformer import vit_b_16, vit_b_32, vit_l_16, vit_l_32
+
+from torchvision import transforms
+from PIL import Image
+import torch
+import torch.nn as nn
+
+
+model_dict = {
+ 'resnet18': resnet18,
+ 'resnet34': resnet34,
+ 'resnet50': resnet50,
+ 'resnet101': resnet101,
+ 'resnet152': resnet152,
+ 'vit_b_16': vit_b_16,
+ 'vit_b_32': vit_b_32,
+ 'vit_l_16': vit_l_16,
+ 'vit_l_32': vit_l_32
+}
+
+
+CHANNELS = {
+ "resnet50" : 2048,
+ "vit_b_16" : 768,
+}
+
+
+
+class ImagenetModel(nn.Module):
+ def __init__(self, name, num_classes=1):
+ super(ImagenetModel, self).__init__()
+
+ self.model = model_dict[name](pretrained=True)
+ self.fc = nn.Linear(CHANNELS[name], num_classes) #manually define a fc layer here
+
+
+ def forward(self, x):
+ feature = self.model(x)["penultimate"]
+ return self.fc(feature)
diff --git a/clean/video/univfd_video/models/resnet.py b/clean/video/univfd_video/models/resnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..a78e3d65e263cb9dbd1afa0e1a88dba9f5ddd164
--- /dev/null
+++ b/clean/video/univfd_video/models/resnet.py
@@ -0,0 +1,337 @@
+import torch
+from torch import Tensor
+import torch.nn as nn
+from typing import Type, Any, Callable, Union, List, Optional
+
+try:
+ from torch.hub import load_state_dict_from_url
+except ImportError:
+ from torch.utils.model_zoo import load_url as load_state_dict_from_url
+
+
+model_urls = {
+ 'resnet18': 'https://download.pytorch.org/models/resnet18-f37072fd.pth',
+ 'resnet34': 'https://download.pytorch.org/models/resnet34-b627a593.pth',
+ 'resnet50': 'https://download.pytorch.org/models/resnet50-0676ba61.pth',
+ 'resnet101': 'https://download.pytorch.org/models/resnet101-63fe2227.pth',
+ 'resnet152': 'https://download.pytorch.org/models/resnet152-394f9c45.pth',
+ 'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
+ 'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
+ 'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
+ 'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
+}
+
+
+
+
+def conv3x3(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d:
+ """3x3 convolution with padding"""
+ return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+ padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+
+def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:
+ """1x1 convolution"""
+ return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+ expansion: int = 1
+
+ def __init__(
+ self,
+ inplanes: int,
+ planes: int,
+ stride: int = 1,
+ downsample: Optional[nn.Module] = None,
+ groups: int = 1,
+ base_width: int = 64,
+ dilation: int = 1,
+ norm_layer: Optional[Callable[..., nn.Module]] = None
+ ) -> None:
+ super(BasicBlock, self).__init__()
+ if norm_layer is None:
+ norm_layer = nn.BatchNorm2d
+ if groups != 1 or base_width != 64:
+ raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+ if dilation > 1:
+ raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+ # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+ self.conv1 = conv3x3(inplanes, planes, stride)
+ self.bn1 = norm_layer(planes)
+ self.relu = nn.ReLU(inplace=True)
+ self.conv2 = conv3x3(planes, planes)
+ self.bn2 = norm_layer(planes)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x: Tensor) -> Tensor:
+ identity = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+
+ if self.downsample is not None:
+ identity = self.downsample(x)
+
+ out += identity
+ out = self.relu(out)
+
+ return out
+
+
+class Bottleneck(nn.Module):
+ # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+ # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+ # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
+ # This variant is also known as ResNet V1.5 and improves accuracy according to
+ # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+
+ expansion: int = 4
+
+ def __init__(
+ self,
+ inplanes: int,
+ planes: int,
+ stride: int = 1,
+ downsample: Optional[nn.Module] = None,
+ groups: int = 1,
+ base_width: int = 64,
+ dilation: int = 1,
+ norm_layer: Optional[Callable[..., nn.Module]] = None
+ ) -> None:
+ super(Bottleneck, self).__init__()
+ if norm_layer is None:
+ norm_layer = nn.BatchNorm2d
+ width = int(planes * (base_width / 64.)) * groups
+ # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+ self.conv1 = conv1x1(inplanes, width)
+ self.bn1 = norm_layer(width)
+ self.conv2 = conv3x3(width, width, stride, groups, dilation)
+ self.bn2 = norm_layer(width)
+ self.conv3 = conv1x1(width, planes * self.expansion)
+ self.bn3 = norm_layer(planes * self.expansion)
+ self.relu = nn.ReLU(inplace=True)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x: Tensor) -> Tensor:
+ identity = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+ out = self.relu(out)
+
+ out = self.conv3(out)
+ out = self.bn3(out)
+
+ if self.downsample is not None:
+ identity = self.downsample(x)
+
+ out += identity
+ out = self.relu(out)
+
+ return out
+
+
+class ResNet(nn.Module):
+
+ def __init__(
+ self,
+ block: Type[Union[BasicBlock, Bottleneck]],
+ layers: List[int],
+ num_classes: int = 1000,
+ zero_init_residual: bool = False,
+ groups: int = 1,
+ width_per_group: int = 64,
+ replace_stride_with_dilation: Optional[List[bool]] = None,
+ norm_layer: Optional[Callable[..., nn.Module]] = None
+ ) -> None:
+ super(ResNet, self).__init__()
+ if norm_layer is None:
+ norm_layer = nn.BatchNorm2d
+ self._norm_layer = norm_layer
+
+ self.inplanes = 64
+ self.dilation = 1
+ if replace_stride_with_dilation is None:
+ # each element in the tuple indicates if we should replace
+ # the 2x2 stride with a dilated convolution instead
+ replace_stride_with_dilation = [False, False, False]
+ if len(replace_stride_with_dilation) != 3:
+ raise ValueError("replace_stride_with_dilation should be None "
+ "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+ self.groups = groups
+ self.base_width = width_per_group
+ self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,
+ bias=False)
+ self.bn1 = norm_layer(self.inplanes)
+ self.relu = nn.ReLU(inplace=True)
+ self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+ self.layer1 = self._make_layer(block, 64, layers[0])
+ self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
+ dilate=replace_stride_with_dilation[0])
+ self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
+ dilate=replace_stride_with_dilation[1])
+ self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
+ dilate=replace_stride_with_dilation[2])
+ self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+ self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+ for m in self.modules():
+ if isinstance(m, nn.Conv2d):
+ nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+ elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+ nn.init.constant_(m.weight, 1)
+ nn.init.constant_(m.bias, 0)
+
+ # Zero-initialize the last BN in each residual branch,
+ # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+ # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+ if zero_init_residual:
+ for m in self.modules():
+ if isinstance(m, Bottleneck):
+ nn.init.constant_(m.bn3.weight, 0) # type: ignore[arg-type]
+ elif isinstance(m, BasicBlock):
+ nn.init.constant_(m.bn2.weight, 0) # type: ignore[arg-type]
+
+ def _make_layer(self, block: Type[Union[BasicBlock, Bottleneck]], planes: int, blocks: int,
+ stride: int = 1, dilate: bool = False) -> nn.Sequential:
+ norm_layer = self._norm_layer
+ downsample = None
+ previous_dilation = self.dilation
+ if dilate:
+ self.dilation *= stride
+ stride = 1
+ if stride != 1 or self.inplanes != planes * block.expansion:
+ downsample = nn.Sequential(
+ conv1x1(self.inplanes, planes * block.expansion, stride),
+ norm_layer(planes * block.expansion),
+ )
+
+ layers = []
+ layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
+ self.base_width, previous_dilation, norm_layer))
+ self.inplanes = planes * block.expansion
+ for _ in range(1, blocks):
+ layers.append(block(self.inplanes, planes, groups=self.groups,
+ base_width=self.base_width, dilation=self.dilation,
+ norm_layer=norm_layer))
+
+ return nn.Sequential(*layers)
+
+ def _forward_impl(self, x):
+ # The comment resolution is based on input size is 224*224 imagenet
+ out = {}
+ x = self.conv1(x)
+ x = self.bn1(x)
+ x = self.relu(x)
+ x = self.maxpool(x)
+ out['f0'] = x # N*64*56*56
+
+ x = self.layer1(x)
+ out['f1'] = x # N*64*56*56
+
+ x = self.layer2(x)
+ out['f2'] = x # N*128*28*28
+
+ x = self.layer3(x)
+ out['f3'] = x # N*256*14*14
+
+ x = self.layer4(x)
+ out['f4'] = x # N*512*7*7
+
+ x = self.avgpool(x)
+ x = torch.flatten(x, 1)
+ out['penultimate'] = x # N*512
+
+ x = self.fc(x)
+ out['logits'] = x # N*1000
+
+ # return all features
+ return out
+
+ # return final classification result
+ # return x
+
+ def forward(self, x):
+ return self._forward_impl(x)
+
+
+def _resnet(
+ arch: str,
+ block: Type[Union[BasicBlock, Bottleneck]],
+ layers: List[int],
+ pretrained: bool,
+ progress: bool,
+ **kwargs: Any
+) -> ResNet:
+ model = ResNet(block, layers, **kwargs)
+ if pretrained:
+ state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
+ model.load_state_dict(state_dict)
+ return model
+
+
+def resnet18(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+ r"""ResNet-18 model from
+ `"Deep Residual Learning for Image Recognition" `_.
+
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ progress (bool): If True, displays a progress bar of the download to stderr
+ """
+ return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress, **kwargs)
+
+
+def resnet34(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+ r"""ResNet-34 model from
+ `"Deep Residual Learning for Image Recognition" `_.
+
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ progress (bool): If True, displays a progress bar of the download to stderr
+ """
+ return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress, **kwargs)
+
+
+def resnet50(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+ r"""ResNet-50 model from
+ `"Deep Residual Learning for Image Recognition" `_.
+
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ progress (bool): If True, displays a progress bar of the download to stderr
+ """
+ return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs)
+
+
+def resnet101(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+ r"""ResNet-101 model from
+ `"Deep Residual Learning for Image Recognition" `_.
+
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ progress (bool): If True, displays a progress bar of the download to stderr
+ """
+ return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs)
+
+
+def resnet152(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+ r"""ResNet-152 model from
+ `"Deep Residual Learning for Image Recognition" `_.
+
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ progress (bool): If True, displays a progress bar of the download to stderr
+ """
+ return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress, **kwargs)
+
diff --git a/clean/video/univfd_video/models/vgg.py b/clean/video/univfd_video/models/vgg.py
new file mode 100644
index 0000000000000000000000000000000000000000..a30a1df18a64f9ab2ca309b264cd4e8409b0cf64
--- /dev/null
+++ b/clean/video/univfd_video/models/vgg.py
@@ -0,0 +1,120 @@
+import torch
+import torch.nn as nn
+from typing import Union, List, Dict, Any, cast
+import torchvision
+import torch.nn.functional as F
+
+
+
+
+
+class VGG(torch.nn.Module):
+ def __init__(self, arch_type, pretrained, progress):
+ super().__init__()
+
+ self.layer1 = torch.nn.Sequential()
+ self.layer2 = torch.nn.Sequential()
+ self.layer3 = torch.nn.Sequential()
+ self.layer4 = torch.nn.Sequential()
+ self.layer5 = torch.nn.Sequential()
+
+ if arch_type == 'vgg11':
+ official_vgg = torchvision.models.vgg11(pretrained=pretrained, progress=progress)
+ blocks = [ [0,2], [2,5], [5,10], [10,15], [15,20] ]
+ last_idx = 20
+ elif arch_type == 'vgg19':
+ official_vgg = torchvision.models.vgg19(pretrained=pretrained, progress=progress)
+ blocks = [ [0,4], [4,9], [9,18], [18,27], [27,36] ]
+ last_idx = 36
+ else:
+ raise NotImplementedError
+
+
+ for x in range( *blocks[0] ):
+ self.layer1.add_module(str(x), official_vgg.features[x])
+ for x in range( *blocks[1] ):
+ self.layer2.add_module(str(x), official_vgg.features[x])
+ for x in range( *blocks[2] ):
+ self.layer3.add_module(str(x), official_vgg.features[x])
+ for x in range( *blocks[3] ):
+ self.layer4.add_module(str(x), official_vgg.features[x])
+ for x in range( *blocks[4] ):
+ self.layer5.add_module(str(x), official_vgg.features[x])
+
+ self.max_pool = official_vgg.features[last_idx]
+ self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
+
+ self.fc1 = official_vgg.classifier[0]
+ self.fc2 = official_vgg.classifier[3]
+ self.fc3 = official_vgg.classifier[6]
+ self.dropout = nn.Dropout()
+
+
+ def forward(self, x):
+ out = {}
+
+ x = self.layer1(x)
+ out['f0'] = x
+
+ x = self.layer2(x)
+ out['f1'] = x
+
+ x = self.layer3(x)
+ out['f2'] = x
+
+ x = self.layer4(x)
+ out['f3'] = x
+
+ x = self.layer5(x)
+ out['f4'] = x
+
+ x = self.max_pool(x)
+ x = self.avgpool(x)
+ x = x.view(-1,512*7*7)
+
+ x = self.fc1(x)
+ x = F.relu(x)
+ x = self.dropout(x)
+ x = self.fc2(x)
+ x = F.relu(x)
+ out['penultimate'] = x
+ x = self.dropout(x)
+ x = self.fc3(x)
+ out['logits'] = x
+
+ return out
+
+
+
+
+
+
+
+
+
+
+def vgg11(pretrained=False, progress=True):
+ r"""VGG 11-layer model (configuration "A") from
+ `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_.
+
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ progress (bool): If True, displays a progress bar of the download to stderr
+ """
+ return VGG('vgg11', pretrained, progress)
+
+
+
+def vgg19(pretrained=False, progress=True):
+ r"""VGG 19-layer model (configuration "E")
+ `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_.
+
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ progress (bool): If True, displays a progress bar of the download to stderr
+ """
+ return VGG('vgg19', pretrained, progress)
+
+
+
+
diff --git a/clean/video/univfd_video/models/vision_transformer.py b/clean/video/univfd_video/models/vision_transformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..618e9626ca43f1afdb3419e19be11f3a3048f81e
--- /dev/null
+++ b/clean/video/univfd_video/models/vision_transformer.py
@@ -0,0 +1,481 @@
+import math
+from collections import OrderedDict
+from functools import partial
+from typing import Any, Callable, List, NamedTuple, Optional
+
+import torch
+import torch.nn as nn
+
+# from .._internally_replaced_utils import load_state_dict_from_url
+from .vision_transformer_misc import ConvNormActivation
+from .vision_transformer_utils import _log_api_usage_once
+
+try:
+ from torch.hub import load_state_dict_from_url
+except ImportError:
+ from torch.utils.model_zoo import load_url as load_state_dict_from_url
+
+# __all__ = [
+# "VisionTransformer",
+# "vit_b_16",
+# "vit_b_32",
+# "vit_l_16",
+# "vit_l_32",
+# ]
+
+model_urls = {
+ "vit_b_16": "https://download.pytorch.org/models/vit_b_16-c867db91.pth",
+ "vit_b_32": "https://download.pytorch.org/models/vit_b_32-d86f8d99.pth",
+ "vit_l_16": "https://download.pytorch.org/models/vit_l_16-852ce7e3.pth",
+ "vit_l_32": "https://download.pytorch.org/models/vit_l_32-c7638314.pth",
+}
+
+
+class ConvStemConfig(NamedTuple):
+ out_channels: int
+ kernel_size: int
+ stride: int
+ norm_layer: Callable[..., nn.Module] = nn.BatchNorm2d
+ activation_layer: Callable[..., nn.Module] = nn.ReLU
+
+
+class MLPBlock(nn.Sequential):
+ """Transformer MLP block."""
+
+ def __init__(self, in_dim: int, mlp_dim: int, dropout: float):
+ super().__init__()
+ self.linear_1 = nn.Linear(in_dim, mlp_dim)
+ self.act = nn.GELU()
+ self.dropout_1 = nn.Dropout(dropout)
+ self.linear_2 = nn.Linear(mlp_dim, in_dim)
+ self.dropout_2 = nn.Dropout(dropout)
+
+ nn.init.xavier_uniform_(self.linear_1.weight)
+ nn.init.xavier_uniform_(self.linear_2.weight)
+ nn.init.normal_(self.linear_1.bias, std=1e-6)
+ nn.init.normal_(self.linear_2.bias, std=1e-6)
+
+
+class EncoderBlock(nn.Module):
+ """Transformer encoder block."""
+
+ def __init__(
+ self,
+ num_heads: int,
+ hidden_dim: int,
+ mlp_dim: int,
+ dropout: float,
+ attention_dropout: float,
+ norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
+ ):
+ super().__init__()
+ self.num_heads = num_heads
+
+ # Attention block
+ self.ln_1 = norm_layer(hidden_dim)
+ self.self_attention = nn.MultiheadAttention(hidden_dim, num_heads, dropout=attention_dropout, batch_first=True)
+ self.dropout = nn.Dropout(dropout)
+
+ # MLP block
+ self.ln_2 = norm_layer(hidden_dim)
+ self.mlp = MLPBlock(hidden_dim, mlp_dim, dropout)
+
+ def forward(self, input: torch.Tensor):
+ torch._assert(input.dim() == 3, f"Expected (seq_length, batch_size, hidden_dim) got {input.shape}")
+ x = self.ln_1(input)
+ x, _ = self.self_attention(query=x, key=x, value=x, need_weights=False)
+ x = self.dropout(x)
+ x = x + input
+
+ y = self.ln_2(x)
+ y = self.mlp(y)
+ return x + y
+
+
+class Encoder(nn.Module):
+ """Transformer Model Encoder for sequence to sequence translation."""
+
+ def __init__(
+ self,
+ seq_length: int,
+ num_layers: int,
+ num_heads: int,
+ hidden_dim: int,
+ mlp_dim: int,
+ dropout: float,
+ attention_dropout: float,
+ norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
+ ):
+ super().__init__()
+ # Note that batch_size is on the first dim because
+ # we have batch_first=True in nn.MultiAttention() by default
+ self.pos_embedding = nn.Parameter(torch.empty(1, seq_length, hidden_dim).normal_(std=0.02)) # from BERT
+ self.dropout = nn.Dropout(dropout)
+ layers: OrderedDict[str, nn.Module] = OrderedDict()
+ for i in range(num_layers):
+ layers[f"encoder_layer_{i}"] = EncoderBlock(
+ num_heads,
+ hidden_dim,
+ mlp_dim,
+ dropout,
+ attention_dropout,
+ norm_layer,
+ )
+ self.layers = nn.Sequential(layers)
+ self.ln = norm_layer(hidden_dim)
+
+ def forward(self, input: torch.Tensor):
+ torch._assert(input.dim() == 3, f"Expected (batch_size, seq_length, hidden_dim) got {input.shape}")
+ input = input + self.pos_embedding
+ return self.ln(self.layers(self.dropout(input)))
+
+
+class VisionTransformer(nn.Module):
+ """Vision Transformer as per https://arxiv.org/abs/2010.11929."""
+
+ def __init__(
+ self,
+ image_size: int,
+ patch_size: int,
+ num_layers: int,
+ num_heads: int,
+ hidden_dim: int,
+ mlp_dim: int,
+ dropout: float = 0.0,
+ attention_dropout: float = 0.0,
+ num_classes: int = 1000,
+ representation_size: Optional[int] = None,
+ norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
+ conv_stem_configs: Optional[List[ConvStemConfig]] = None,
+ ):
+ super().__init__()
+ _log_api_usage_once(self)
+ torch._assert(image_size % patch_size == 0, "Input shape indivisible by patch size!")
+ self.image_size = image_size
+ self.patch_size = patch_size
+ self.hidden_dim = hidden_dim
+ self.mlp_dim = mlp_dim
+ self.attention_dropout = attention_dropout
+ self.dropout = dropout
+ self.num_classes = num_classes
+ self.representation_size = representation_size
+ self.norm_layer = norm_layer
+
+ if conv_stem_configs is not None:
+ # As per https://arxiv.org/abs/2106.14881
+ seq_proj = nn.Sequential()
+ prev_channels = 3
+ for i, conv_stem_layer_config in enumerate(conv_stem_configs):
+ seq_proj.add_module(
+ f"conv_bn_relu_{i}",
+ ConvNormActivation(
+ in_channels=prev_channels,
+ out_channels=conv_stem_layer_config.out_channels,
+ kernel_size=conv_stem_layer_config.kernel_size,
+ stride=conv_stem_layer_config.stride,
+ norm_layer=conv_stem_layer_config.norm_layer,
+ activation_layer=conv_stem_layer_config.activation_layer,
+ ),
+ )
+ prev_channels = conv_stem_layer_config.out_channels
+ seq_proj.add_module(
+ "conv_last", nn.Conv2d(in_channels=prev_channels, out_channels=hidden_dim, kernel_size=1)
+ )
+ self.conv_proj: nn.Module = seq_proj
+ else:
+ self.conv_proj = nn.Conv2d(
+ in_channels=3, out_channels=hidden_dim, kernel_size=patch_size, stride=patch_size
+ )
+
+ seq_length = (image_size // patch_size) ** 2
+
+ # Add a class token
+ self.class_token = nn.Parameter(torch.zeros(1, 1, hidden_dim))
+ seq_length += 1
+
+ self.encoder = Encoder(
+ seq_length,
+ num_layers,
+ num_heads,
+ hidden_dim,
+ mlp_dim,
+ dropout,
+ attention_dropout,
+ norm_layer,
+ )
+ self.seq_length = seq_length
+
+ heads_layers: OrderedDict[str, nn.Module] = OrderedDict()
+ if representation_size is None:
+ heads_layers["head"] = nn.Linear(hidden_dim, num_classes)
+ else:
+ heads_layers["pre_logits"] = nn.Linear(hidden_dim, representation_size)
+ heads_layers["act"] = nn.Tanh()
+ heads_layers["head"] = nn.Linear(representation_size, num_classes)
+
+ self.heads = nn.Sequential(heads_layers)
+
+ if isinstance(self.conv_proj, nn.Conv2d):
+ # Init the patchify stem
+ fan_in = self.conv_proj.in_channels * self.conv_proj.kernel_size[0] * self.conv_proj.kernel_size[1]
+ nn.init.trunc_normal_(self.conv_proj.weight, std=math.sqrt(1 / fan_in))
+ if self.conv_proj.bias is not None:
+ nn.init.zeros_(self.conv_proj.bias)
+ elif self.conv_proj.conv_last is not None and isinstance(self.conv_proj.conv_last, nn.Conv2d):
+ # Init the last 1x1 conv of the conv stem
+ nn.init.normal_(
+ self.conv_proj.conv_last.weight, mean=0.0, std=math.sqrt(2.0 / self.conv_proj.conv_last.out_channels)
+ )
+ if self.conv_proj.conv_last.bias is not None:
+ nn.init.zeros_(self.conv_proj.conv_last.bias)
+
+ if hasattr(self.heads, "pre_logits") and isinstance(self.heads.pre_logits, nn.Linear):
+ fan_in = self.heads.pre_logits.in_features
+ nn.init.trunc_normal_(self.heads.pre_logits.weight, std=math.sqrt(1 / fan_in))
+ nn.init.zeros_(self.heads.pre_logits.bias)
+
+ if isinstance(self.heads.head, nn.Linear):
+ nn.init.zeros_(self.heads.head.weight)
+ nn.init.zeros_(self.heads.head.bias)
+
+ def _process_input(self, x: torch.Tensor) -> torch.Tensor:
+ n, c, h, w = x.shape
+ p = self.patch_size
+ torch._assert(h == self.image_size, "Wrong image height!")
+ torch._assert(w == self.image_size, "Wrong image width!")
+ n_h = h // p
+ n_w = w // p
+
+ # (n, c, h, w) -> (n, hidden_dim, n_h, n_w)
+ x = self.conv_proj(x)
+ # (n, hidden_dim, n_h, n_w) -> (n, hidden_dim, (n_h * n_w))
+ x = x.reshape(n, self.hidden_dim, n_h * n_w)
+
+ # (n, hidden_dim, (n_h * n_w)) -> (n, (n_h * n_w), hidden_dim)
+ # The self attention layer expects inputs in the format (N, S, E)
+ # where S is the source sequence length, N is the batch size, E is the
+ # embedding dimension
+ x = x.permute(0, 2, 1)
+
+ return x
+
+ def forward(self, x: torch.Tensor):
+ out = {}
+
+ # Reshape and permute the input tensor
+ x = self._process_input(x)
+ n = x.shape[0]
+
+ # Expand the class token to the full batch
+ batch_class_token = self.class_token.expand(n, -1, -1)
+ x = torch.cat([batch_class_token, x], dim=1)
+
+
+ x = self.encoder(x)
+ img_feature = x[:,1:]
+ H = W = int(self.image_size / self.patch_size)
+ out['f4'] = img_feature.view(n, H, W, self.hidden_dim).permute(0,3,1,2)
+
+ # Classifier "token" as used by standard language architectures
+ x = x[:, 0]
+ out['penultimate'] = x
+
+ x = self.heads(x) # I checked that for all pretrained ViT, this is just a fc
+ out['logits'] = x
+
+ return out
+
+
+def _vision_transformer(
+ arch: str,
+ patch_size: int,
+ num_layers: int,
+ num_heads: int,
+ hidden_dim: int,
+ mlp_dim: int,
+ pretrained: bool,
+ progress: bool,
+ **kwargs: Any,
+) -> VisionTransformer:
+ image_size = kwargs.pop("image_size", 224)
+
+ model = VisionTransformer(
+ image_size=image_size,
+ patch_size=patch_size,
+ num_layers=num_layers,
+ num_heads=num_heads,
+ hidden_dim=hidden_dim,
+ mlp_dim=mlp_dim,
+ **kwargs,
+ )
+
+ if pretrained:
+ if arch not in model_urls:
+ raise ValueError(f"No checkpoint is available for model type '{arch}'!")
+ state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
+ model.load_state_dict(state_dict)
+
+ return model
+
+
+def vit_b_16(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VisionTransformer:
+ """
+ Constructs a vit_b_16 architecture from
+ `"An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale" `_.
+
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ progress (bool): If True, displays a progress bar of the download to stderr
+ """
+ return _vision_transformer(
+ arch="vit_b_16",
+ patch_size=16,
+ num_layers=12,
+ num_heads=12,
+ hidden_dim=768,
+ mlp_dim=3072,
+ pretrained=pretrained,
+ progress=progress,
+ **kwargs,
+ )
+
+
+def vit_b_32(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VisionTransformer:
+ """
+ Constructs a vit_b_32 architecture from
+ `"An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale" `_.
+
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ progress (bool): If True, displays a progress bar of the download to stderr
+ """
+ return _vision_transformer(
+ arch="vit_b_32",
+ patch_size=32,
+ num_layers=12,
+ num_heads=12,
+ hidden_dim=768,
+ mlp_dim=3072,
+ pretrained=pretrained,
+ progress=progress,
+ **kwargs,
+ )
+
+
+def vit_l_16(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VisionTransformer:
+ """
+ Constructs a vit_l_16 architecture from
+ `"An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale" `_.
+
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ progress (bool): If True, displays a progress bar of the download to stderr
+ """
+ return _vision_transformer(
+ arch="vit_l_16",
+ patch_size=16,
+ num_layers=24,
+ num_heads=16,
+ hidden_dim=1024,
+ mlp_dim=4096,
+ pretrained=pretrained,
+ progress=progress,
+ **kwargs,
+ )
+
+
+def vit_l_32(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VisionTransformer:
+ """
+ Constructs a vit_l_32 architecture from
+ `"An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale" `_.
+
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ progress (bool): If True, displays a progress bar of the download to stderr
+ """
+ return _vision_transformer(
+ arch="vit_l_32",
+ patch_size=32,
+ num_layers=24,
+ num_heads=16,
+ hidden_dim=1024,
+ mlp_dim=4096,
+ pretrained=pretrained,
+ progress=progress,
+ **kwargs,
+ )
+
+
+def interpolate_embeddings(
+ image_size: int,
+ patch_size: int,
+ model_state: "OrderedDict[str, torch.Tensor]",
+ interpolation_mode: str = "bicubic",
+ reset_heads: bool = False,
+) -> "OrderedDict[str, torch.Tensor]":
+ """This function helps interpolating positional embeddings during checkpoint loading,
+ especially when you want to apply a pre-trained model on images with different resolution.
+
+ Args:
+ image_size (int): Image size of the new model.
+ patch_size (int): Patch size of the new model.
+ model_state (OrderedDict[str, torch.Tensor]): State dict of the pre-trained model.
+ interpolation_mode (str): The algorithm used for upsampling. Default: bicubic.
+ reset_heads (bool): If true, not copying the state of heads. Default: False.
+
+ Returns:
+ OrderedDict[str, torch.Tensor]: A state dict which can be loaded into the new model.
+ """
+ # Shape of pos_embedding is (1, seq_length, hidden_dim)
+ pos_embedding = model_state["encoder.pos_embedding"]
+ n, seq_length, hidden_dim = pos_embedding.shape
+ if n != 1:
+ raise ValueError(f"Unexpected position embedding shape: {pos_embedding.shape}")
+
+ new_seq_length = (image_size // patch_size) ** 2 + 1
+
+ # Need to interpolate the weights for the position embedding.
+ # We do this by reshaping the positions embeddings to a 2d grid, performing
+ # an interpolation in the (h, w) space and then reshaping back to a 1d grid.
+ if new_seq_length != seq_length:
+ # The class token embedding shouldn't be interpolated so we split it up.
+ seq_length -= 1
+ new_seq_length -= 1
+ pos_embedding_token = pos_embedding[:, :1, :]
+ pos_embedding_img = pos_embedding[:, 1:, :]
+
+ # (1, seq_length, hidden_dim) -> (1, hidden_dim, seq_length)
+ pos_embedding_img = pos_embedding_img.permute(0, 2, 1)
+ seq_length_1d = int(math.sqrt(seq_length))
+ torch._assert(seq_length_1d * seq_length_1d == seq_length, "seq_length is not a perfect square!")
+
+ # (1, hidden_dim, seq_length) -> (1, hidden_dim, seq_l_1d, seq_l_1d)
+ pos_embedding_img = pos_embedding_img.reshape(1, hidden_dim, seq_length_1d, seq_length_1d)
+ new_seq_length_1d = image_size // patch_size
+
+ # Perform interpolation.
+ # (1, hidden_dim, seq_l_1d, seq_l_1d) -> (1, hidden_dim, new_seq_l_1d, new_seq_l_1d)
+ new_pos_embedding_img = nn.functional.interpolate(
+ pos_embedding_img,
+ size=new_seq_length_1d,
+ mode=interpolation_mode,
+ align_corners=True,
+ )
+
+ # (1, hidden_dim, new_seq_l_1d, new_seq_l_1d) -> (1, hidden_dim, new_seq_length)
+ new_pos_embedding_img = new_pos_embedding_img.reshape(1, hidden_dim, new_seq_length)
+
+ # (1, hidden_dim, new_seq_length) -> (1, new_seq_length, hidden_dim)
+ new_pos_embedding_img = new_pos_embedding_img.permute(0, 2, 1)
+ new_pos_embedding = torch.cat([pos_embedding_token, new_pos_embedding_img], dim=1)
+
+ model_state["encoder.pos_embedding"] = new_pos_embedding
+
+ if reset_heads:
+ model_state_copy: "OrderedDict[str, torch.Tensor]" = OrderedDict()
+ for k, v in model_state.items():
+ if not k.startswith("heads"):
+ model_state_copy[k] = v
+ model_state = model_state_copy
+
+ return model_state
diff --git a/clean/video/univfd_video/models/vision_transformer_misc.py b/clean/video/univfd_video/models/vision_transformer_misc.py
new file mode 100644
index 0000000000000000000000000000000000000000..7915f036c00f0d9c57c176e621afc9f1e69dcb30
--- /dev/null
+++ b/clean/video/univfd_video/models/vision_transformer_misc.py
@@ -0,0 +1,163 @@
+from typing import Callable, List, Optional
+
+import torch
+from torch import Tensor
+
+from .vision_transformer_utils import _log_api_usage_once
+
+
+interpolate = torch.nn.functional.interpolate
+
+
+# This is not in nn
+class FrozenBatchNorm2d(torch.nn.Module):
+ """
+ BatchNorm2d where the batch statistics and the affine parameters are fixed
+
+ Args:
+ num_features (int): Number of features ``C`` from an expected input of size ``(N, C, H, W)``
+ eps (float): a value added to the denominator for numerical stability. Default: 1e-5
+ """
+
+ def __init__(
+ self,
+ num_features: int,
+ eps: float = 1e-5,
+ ):
+ super().__init__()
+ _log_api_usage_once(self)
+ self.eps = eps
+ self.register_buffer("weight", torch.ones(num_features))
+ self.register_buffer("bias", torch.zeros(num_features))
+ self.register_buffer("running_mean", torch.zeros(num_features))
+ self.register_buffer("running_var", torch.ones(num_features))
+
+ def _load_from_state_dict(
+ self,
+ state_dict: dict,
+ prefix: str,
+ local_metadata: dict,
+ strict: bool,
+ missing_keys: List[str],
+ unexpected_keys: List[str],
+ error_msgs: List[str],
+ ):
+ num_batches_tracked_key = prefix + "num_batches_tracked"
+ if num_batches_tracked_key in state_dict:
+ del state_dict[num_batches_tracked_key]
+
+ super()._load_from_state_dict(
+ state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+ )
+
+ def forward(self, x: Tensor) -> Tensor:
+ # move reshapes to the beginning
+ # to make it fuser-friendly
+ w = self.weight.reshape(1, -1, 1, 1)
+ b = self.bias.reshape(1, -1, 1, 1)
+ rv = self.running_var.reshape(1, -1, 1, 1)
+ rm = self.running_mean.reshape(1, -1, 1, 1)
+ scale = w * (rv + self.eps).rsqrt()
+ bias = b - rm * scale
+ return x * scale + bias
+
+ def __repr__(self) -> str:
+ return f"{self.__class__.__name__}({self.weight.shape[0]}, eps={self.eps})"
+
+
+class ConvNormActivation(torch.nn.Sequential):
+ """
+ Configurable block used for Convolution-Normalzation-Activation blocks.
+
+ Args:
+ in_channels (int): Number of channels in the input image
+ out_channels (int): Number of channels produced by the Convolution-Normalzation-Activation block
+ kernel_size: (int, optional): Size of the convolving kernel. Default: 3
+ stride (int, optional): Stride of the convolution. Default: 1
+ padding (int, tuple or str, optional): Padding added to all four sides of the input. Default: None, in wich case it will calculated as ``padding = (kernel_size - 1) // 2 * dilation``
+ groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+ norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the convolutiuon layer. If ``None`` this layer wont be used. Default: ``torch.nn.BatchNorm2d``
+ activation_layer (Callable[..., torch.nn.Module], optinal): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the conv layer. If ``None`` this layer wont be used. Default: ``torch.nn.ReLU``
+ dilation (int): Spacing between kernel elements. Default: 1
+ inplace (bool): Parameter for the activation layer, which can optionally do the operation in-place. Default ``True``
+ bias (bool, optional): Whether to use bias in the convolution layer. By default, biases are included if ``norm_layer is None``.
+
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int = 3,
+ stride: int = 1,
+ padding: Optional[int] = None,
+ groups: int = 1,
+ norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d,
+ activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
+ dilation: int = 1,
+ inplace: Optional[bool] = True,
+ bias: Optional[bool] = None,
+ ) -> None:
+ if padding is None:
+ padding = (kernel_size - 1) // 2 * dilation
+ if bias is None:
+ bias = norm_layer is None
+ layers = [
+ torch.nn.Conv2d(
+ in_channels,
+ out_channels,
+ kernel_size,
+ stride,
+ padding,
+ dilation=dilation,
+ groups=groups,
+ bias=bias,
+ )
+ ]
+ if norm_layer is not None:
+ layers.append(norm_layer(out_channels))
+ if activation_layer is not None:
+ params = {} if inplace is None else {"inplace": inplace}
+ layers.append(activation_layer(**params))
+ super().__init__(*layers)
+ _log_api_usage_once(self)
+ self.out_channels = out_channels
+
+
+class SqueezeExcitation(torch.nn.Module):
+ """
+ This block implements the Squeeze-and-Excitation block from https://arxiv.org/abs/1709.01507 (see Fig. 1).
+ Parameters ``activation``, and ``scale_activation`` correspond to ``delta`` and ``sigma`` in in eq. 3.
+
+ Args:
+ input_channels (int): Number of channels in the input image
+ squeeze_channels (int): Number of squeeze channels
+ activation (Callable[..., torch.nn.Module], optional): ``delta`` activation. Default: ``torch.nn.ReLU``
+ scale_activation (Callable[..., torch.nn.Module]): ``sigma`` activation. Default: ``torch.nn.Sigmoid``
+ """
+
+ def __init__(
+ self,
+ input_channels: int,
+ squeeze_channels: int,
+ activation: Callable[..., torch.nn.Module] = torch.nn.ReLU,
+ scale_activation: Callable[..., torch.nn.Module] = torch.nn.Sigmoid,
+ ) -> None:
+ super().__init__()
+ _log_api_usage_once(self)
+ self.avgpool = torch.nn.AdaptiveAvgPool2d(1)
+ self.fc1 = torch.nn.Conv2d(input_channels, squeeze_channels, 1)
+ self.fc2 = torch.nn.Conv2d(squeeze_channels, input_channels, 1)
+ self.activation = activation()
+ self.scale_activation = scale_activation()
+
+ def _scale(self, input: Tensor) -> Tensor:
+ scale = self.avgpool(input)
+ scale = self.fc1(scale)
+ scale = self.activation(scale)
+ scale = self.fc2(scale)
+ return self.scale_activation(scale)
+
+ def forward(self, input: Tensor) -> Tensor:
+ scale = self._scale(input)
+ return scale * input
diff --git a/clean/video/univfd_video/models/vision_transformer_utils.py b/clean/video/univfd_video/models/vision_transformer_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..6d3293d103d0e186a1244e7cc0c6e3bde63d1df3
--- /dev/null
+++ b/clean/video/univfd_video/models/vision_transformer_utils.py
@@ -0,0 +1,549 @@
+import math
+import pathlib
+import warnings
+from types import FunctionType
+from typing import Any, BinaryIO, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from PIL import Image, ImageColor, ImageDraw, ImageFont
+
+__all__ = [
+ "make_grid",
+ "save_image",
+ "draw_bounding_boxes",
+ "draw_segmentation_masks",
+ "draw_keypoints",
+ "flow_to_image",
+]
+
+
+@torch.no_grad()
+def make_grid(
+ tensor: Union[torch.Tensor, List[torch.Tensor]],
+ nrow: int = 8,
+ padding: int = 2,
+ normalize: bool = False,
+ value_range: Optional[Tuple[int, int]] = None,
+ scale_each: bool = False,
+ pad_value: float = 0.0,
+ **kwargs,
+) -> torch.Tensor:
+ """
+ Make a grid of images.
+
+ Args:
+ tensor (Tensor or list): 4D mini-batch Tensor of shape (B x C x H x W)
+ or a list of images all of the same size.
+ nrow (int, optional): Number of images displayed in each row of the grid.
+ The final grid size is ``(B / nrow, nrow)``. Default: ``8``.
+ padding (int, optional): amount of padding. Default: ``2``.
+ normalize (bool, optional): If True, shift the image to the range (0, 1),
+ by the min and max values specified by ``value_range``. Default: ``False``.
+ value_range (tuple, optional): tuple (min, max) where min and max are numbers,
+ then these numbers are used to normalize the image. By default, min and max
+ are computed from the tensor.
+ range (tuple. optional):
+ .. warning::
+ This parameter was deprecated in ``0.12`` and will be removed in ``0.14``. Please use ``value_range``
+ instead.
+ scale_each (bool, optional): If ``True``, scale each image in the batch of
+ images separately rather than the (min, max) over all images. Default: ``False``.
+ pad_value (float, optional): Value for the padded pixels. Default: ``0``.
+
+ Returns:
+ grid (Tensor): the tensor containing grid of images.
+ """
+ if not torch.jit.is_scripting() and not torch.jit.is_tracing():
+ _log_api_usage_once(make_grid)
+ if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))):
+ raise TypeError(f"tensor or list of tensors expected, got {type(tensor)}")
+
+ if "range" in kwargs.keys():
+ warnings.warn(
+ "The parameter 'range' is deprecated since 0.12 and will be removed in 0.14. "
+ "Please use 'value_range' instead."
+ )
+ value_range = kwargs["range"]
+
+ # if list of tensors, convert to a 4D mini-batch Tensor
+ if isinstance(tensor, list):
+ tensor = torch.stack(tensor, dim=0)
+
+ if tensor.dim() == 2: # single image H x W
+ tensor = tensor.unsqueeze(0)
+ if tensor.dim() == 3: # single image
+ if tensor.size(0) == 1: # if single-channel, convert to 3-channel
+ tensor = torch.cat((tensor, tensor, tensor), 0)
+ tensor = tensor.unsqueeze(0)
+
+ if tensor.dim() == 4 and tensor.size(1) == 1: # single-channel images
+ tensor = torch.cat((tensor, tensor, tensor), 1)
+
+ if normalize is True:
+ tensor = tensor.clone() # avoid modifying tensor in-place
+ if value_range is not None:
+ assert isinstance(
+ value_range, tuple
+ ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"
+
+ def norm_ip(img, low, high):
+ img.clamp_(min=low, max=high)
+ img.sub_(low).div_(max(high - low, 1e-5))
+
+ def norm_range(t, value_range):
+ if value_range is not None:
+ norm_ip(t, value_range[0], value_range[1])
+ else:
+ norm_ip(t, float(t.min()), float(t.max()))
+
+ if scale_each is True:
+ for t in tensor: # loop over mini-batch dimension
+ norm_range(t, value_range)
+ else:
+ norm_range(tensor, value_range)
+
+ assert isinstance(tensor, torch.Tensor)
+ if tensor.size(0) == 1:
+ return tensor.squeeze(0)
+
+ # make the mini-batch of images into a grid
+ nmaps = tensor.size(0)
+ xmaps = min(nrow, nmaps)
+ ymaps = int(math.ceil(float(nmaps) / xmaps))
+ height, width = int(tensor.size(2) + padding), int(tensor.size(3) + padding)
+ num_channels = tensor.size(1)
+ grid = tensor.new_full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
+ k = 0
+ for y in range(ymaps):
+ for x in range(xmaps):
+ if k >= nmaps:
+ break
+ # Tensor.copy_() is a valid method but seems to be missing from the stubs
+ # https://pytorch.org/docs/stable/tensors.html#torch.Tensor.copy_
+ grid.narrow(1, y * height + padding, height - padding).narrow( # type: ignore[attr-defined]
+ 2, x * width + padding, width - padding
+ ).copy_(tensor[k])
+ k = k + 1
+ return grid
+
+
+@torch.no_grad()
+def save_image(
+ tensor: Union[torch.Tensor, List[torch.Tensor]],
+ fp: Union[str, pathlib.Path, BinaryIO],
+ format: Optional[str] = None,
+ **kwargs,
+) -> None:
+ """
+ Save a given Tensor into an image file.
+
+ Args:
+ tensor (Tensor or list): Image to be saved. If given a mini-batch tensor,
+ saves the tensor as a grid of images by calling ``make_grid``.
+ fp (string or file object): A filename or a file object
+ format(Optional): If omitted, the format to use is determined from the filename extension.
+ If a file object was used instead of a filename, this parameter should always be used.
+ **kwargs: Other arguments are documented in ``make_grid``.
+ """
+
+ if not torch.jit.is_scripting() and not torch.jit.is_tracing():
+ _log_api_usage_once(save_image)
+ grid = make_grid(tensor, **kwargs)
+ # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
+ ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()
+ im = Image.fromarray(ndarr)
+ im.save(fp, format=format)
+
+
+@torch.no_grad()
+def draw_bounding_boxes(
+ image: torch.Tensor,
+ boxes: torch.Tensor,
+ labels: Optional[List[str]] = None,
+ colors: Optional[Union[List[Union[str, Tuple[int, int, int]]], str, Tuple[int, int, int]]] = None,
+ fill: Optional[bool] = False,
+ width: int = 1,
+ font: Optional[str] = None,
+ font_size: int = 10,
+) -> torch.Tensor:
+
+ """
+ Draws bounding boxes on given image.
+ The values of the input image should be uint8 between 0 and 255.
+ If fill is True, Resulting Tensor should be saved as PNG image.
+
+ Args:
+ image (Tensor): Tensor of shape (C x H x W) and dtype uint8.
+ boxes (Tensor): Tensor of size (N, 4) containing bounding boxes in (xmin, ymin, xmax, ymax) format. Note that
+ the boxes are absolute coordinates with respect to the image. In other words: `0 <= xmin < xmax < W` and
+ `0 <= ymin < ymax < H`.
+ labels (List[str]): List containing the labels of bounding boxes.
+ colors (color or list of colors, optional): List containing the colors
+ of the boxes or single color for all boxes. The color can be represented as
+ PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``.
+ By default, random colors are generated for boxes.
+ fill (bool): If `True` fills the bounding box with specified color.
+ width (int): Width of bounding box.
+ font (str): A filename containing a TrueType font. If the file is not found in this filename, the loader may
+ also search in other directories, such as the `fonts/` directory on Windows or `/Library/Fonts/`,
+ `/System/Library/Fonts/` and `~/Library/Fonts/` on macOS.
+ font_size (int): The requested font size in points.
+
+ Returns:
+ img (Tensor[C, H, W]): Image Tensor of dtype uint8 with bounding boxes plotted.
+ """
+
+ if not torch.jit.is_scripting() and not torch.jit.is_tracing():
+ _log_api_usage_once(draw_bounding_boxes)
+ if not isinstance(image, torch.Tensor):
+ raise TypeError(f"Tensor expected, got {type(image)}")
+ elif image.dtype != torch.uint8:
+ raise ValueError(f"Tensor uint8 expected, got {image.dtype}")
+ elif image.dim() != 3:
+ raise ValueError("Pass individual images, not batches")
+ elif image.size(0) not in {1, 3}:
+ raise ValueError("Only grayscale and RGB images are supported")
+
+ num_boxes = boxes.shape[0]
+
+ if labels is None:
+ labels: Union[List[str], List[None]] = [None] * num_boxes # type: ignore[no-redef]
+ elif len(labels) != num_boxes:
+ raise ValueError(
+ f"Number of boxes ({num_boxes}) and labels ({len(labels)}) mismatch. Please specify labels for each box."
+ )
+
+ if colors is None:
+ colors = _generate_color_palette(num_boxes)
+ elif isinstance(colors, list):
+ if len(colors) < num_boxes:
+ raise ValueError(f"Number of colors ({len(colors)}) is less than number of boxes ({num_boxes}). ")
+ else: # colors specifies a single color for all boxes
+ colors = [colors] * num_boxes
+
+ colors = [(ImageColor.getrgb(color) if isinstance(color, str) else color) for color in colors]
+
+ # Handle Grayscale images
+ if image.size(0) == 1:
+ image = torch.tile(image, (3, 1, 1))
+
+ ndarr = image.permute(1, 2, 0).cpu().numpy()
+ img_to_draw = Image.fromarray(ndarr)
+ img_boxes = boxes.to(torch.int64).tolist()
+
+ if fill:
+ draw = ImageDraw.Draw(img_to_draw, "RGBA")
+ else:
+ draw = ImageDraw.Draw(img_to_draw)
+
+ txt_font = ImageFont.load_default() if font is None else ImageFont.truetype(font=font, size=font_size)
+
+ for bbox, color, label in zip(img_boxes, colors, labels): # type: ignore[arg-type]
+ if fill:
+ fill_color = color + (100,)
+ draw.rectangle(bbox, width=width, outline=color, fill=fill_color)
+ else:
+ draw.rectangle(bbox, width=width, outline=color)
+
+ if label is not None:
+ margin = width + 1
+ draw.text((bbox[0] + margin, bbox[1] + margin), label, fill=color, font=txt_font)
+
+ return torch.from_numpy(np.array(img_to_draw)).permute(2, 0, 1).to(dtype=torch.uint8)
+
+
+@torch.no_grad()
+def draw_segmentation_masks(
+ image: torch.Tensor,
+ masks: torch.Tensor,
+ alpha: float = 0.8,
+ colors: Optional[Union[List[Union[str, Tuple[int, int, int]]], str, Tuple[int, int, int]]] = None,
+) -> torch.Tensor:
+
+ """
+ Draws segmentation masks on given RGB image.
+ The values of the input image should be uint8 between 0 and 255.
+
+ Args:
+ image (Tensor): Tensor of shape (3, H, W) and dtype uint8.
+ masks (Tensor): Tensor of shape (num_masks, H, W) or (H, W) and dtype bool.
+ alpha (float): Float number between 0 and 1 denoting the transparency of the masks.
+ 0 means full transparency, 1 means no transparency.
+ colors (color or list of colors, optional): List containing the colors
+ of the masks or single color for all masks. The color can be represented as
+ PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``.
+ By default, random colors are generated for each mask.
+
+ Returns:
+ img (Tensor[C, H, W]): Image Tensor, with segmentation masks drawn on top.
+ """
+
+ if not torch.jit.is_scripting() and not torch.jit.is_tracing():
+ _log_api_usage_once(draw_segmentation_masks)
+ if not isinstance(image, torch.Tensor):
+ raise TypeError(f"The image must be a tensor, got {type(image)}")
+ elif image.dtype != torch.uint8:
+ raise ValueError(f"The image dtype must be uint8, got {image.dtype}")
+ elif image.dim() != 3:
+ raise ValueError("Pass individual images, not batches")
+ elif image.size()[0] != 3:
+ raise ValueError("Pass an RGB image. Other Image formats are not supported")
+ if masks.ndim == 2:
+ masks = masks[None, :, :]
+ if masks.ndim != 3:
+ raise ValueError("masks must be of shape (H, W) or (batch_size, H, W)")
+ if masks.dtype != torch.bool:
+ raise ValueError(f"The masks must be of dtype bool. Got {masks.dtype}")
+ if masks.shape[-2:] != image.shape[-2:]:
+ raise ValueError("The image and the masks must have the same height and width")
+
+ num_masks = masks.size()[0]
+ if colors is not None and num_masks > len(colors):
+ raise ValueError(f"There are more masks ({num_masks}) than colors ({len(colors)})")
+
+ if colors is None:
+ colors = _generate_color_palette(num_masks)
+
+ if not isinstance(colors, list):
+ colors = [colors]
+ if not isinstance(colors[0], (tuple, str)):
+ raise ValueError("colors must be a tuple or a string, or a list thereof")
+ if isinstance(colors[0], tuple) and len(colors[0]) != 3:
+ raise ValueError("It seems that you passed a tuple of colors instead of a list of colors")
+
+ out_dtype = torch.uint8
+
+ colors_ = []
+ for color in colors:
+ if isinstance(color, str):
+ color = ImageColor.getrgb(color)
+ colors_.append(torch.tensor(color, dtype=out_dtype))
+
+ img_to_draw = image.detach().clone()
+ # TODO: There might be a way to vectorize this
+ for mask, color in zip(masks, colors_):
+ img_to_draw[:, mask] = color[:, None]
+
+ out = image * (1 - alpha) + img_to_draw * alpha
+ return out.to(out_dtype)
+
+
+@torch.no_grad()
+def draw_keypoints(
+ image: torch.Tensor,
+ keypoints: torch.Tensor,
+ connectivity: Optional[List[Tuple[int, int]]] = None,
+ colors: Optional[Union[str, Tuple[int, int, int]]] = None,
+ radius: int = 2,
+ width: int = 3,
+) -> torch.Tensor:
+
+ """
+ Draws Keypoints on given RGB image.
+ The values of the input image should be uint8 between 0 and 255.
+
+ Args:
+ image (Tensor): Tensor of shape (3, H, W) and dtype uint8.
+ keypoints (Tensor): Tensor of shape (num_instances, K, 2) the K keypoints location for each of the N instances,
+ in the format [x, y].
+ connectivity (List[Tuple[int, int]]]): A List of tuple where,
+ each tuple contains pair of keypoints to be connected.
+ colors (str, Tuple): The color can be represented as
+ PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``.
+ radius (int): Integer denoting radius of keypoint.
+ width (int): Integer denoting width of line connecting keypoints.
+
+ Returns:
+ img (Tensor[C, H, W]): Image Tensor of dtype uint8 with keypoints drawn.
+ """
+
+ if not torch.jit.is_scripting() and not torch.jit.is_tracing():
+ _log_api_usage_once(draw_keypoints)
+ if not isinstance(image, torch.Tensor):
+ raise TypeError(f"The image must be a tensor, got {type(image)}")
+ elif image.dtype != torch.uint8:
+ raise ValueError(f"The image dtype must be uint8, got {image.dtype}")
+ elif image.dim() != 3:
+ raise ValueError("Pass individual images, not batches")
+ elif image.size()[0] != 3:
+ raise ValueError("Pass an RGB image. Other Image formats are not supported")
+
+ if keypoints.ndim != 3:
+ raise ValueError("keypoints must be of shape (num_instances, K, 2)")
+
+ ndarr = image.permute(1, 2, 0).cpu().numpy()
+ img_to_draw = Image.fromarray(ndarr)
+ draw = ImageDraw.Draw(img_to_draw)
+ img_kpts = keypoints.to(torch.int64).tolist()
+
+ for kpt_id, kpt_inst in enumerate(img_kpts):
+ for inst_id, kpt in enumerate(kpt_inst):
+ x1 = kpt[0] - radius
+ x2 = kpt[0] + radius
+ y1 = kpt[1] - radius
+ y2 = kpt[1] + radius
+ draw.ellipse([x1, y1, x2, y2], fill=colors, outline=None, width=0)
+
+ if connectivity:
+ for connection in connectivity:
+ start_pt_x = kpt_inst[connection[0]][0]
+ start_pt_y = kpt_inst[connection[0]][1]
+
+ end_pt_x = kpt_inst[connection[1]][0]
+ end_pt_y = kpt_inst[connection[1]][1]
+
+ draw.line(
+ ((start_pt_x, start_pt_y), (end_pt_x, end_pt_y)),
+ width=width,
+ )
+
+ return torch.from_numpy(np.array(img_to_draw)).permute(2, 0, 1).to(dtype=torch.uint8)
+
+
+# Flow visualization code adapted from https://github.com/tomrunia/OpticalFlow_Visualization
+@torch.no_grad()
+def flow_to_image(flow: torch.Tensor) -> torch.Tensor:
+
+ """
+ Converts a flow to an RGB image.
+
+ Args:
+ flow (Tensor): Flow of shape (N, 2, H, W) or (2, H, W) and dtype torch.float.
+
+ Returns:
+ img (Tensor): Image Tensor of dtype uint8 where each color corresponds
+ to a given flow direction. Shape is (N, 3, H, W) or (3, H, W) depending on the input.
+ """
+
+ if flow.dtype != torch.float:
+ raise ValueError(f"Flow should be of dtype torch.float, got {flow.dtype}.")
+
+ orig_shape = flow.shape
+ if flow.ndim == 3:
+ flow = flow[None] # Add batch dim
+
+ if flow.ndim != 4 or flow.shape[1] != 2:
+ raise ValueError(f"Input flow should have shape (2, H, W) or (N, 2, H, W), got {orig_shape}.")
+
+ max_norm = torch.sum(flow ** 2, dim=1).sqrt().max()
+ epsilon = torch.finfo((flow).dtype).eps
+ normalized_flow = flow / (max_norm + epsilon)
+ img = _normalized_flow_to_image(normalized_flow)
+
+ if len(orig_shape) == 3:
+ img = img[0] # Remove batch dim
+ return img
+
+
+@torch.no_grad()
+def _normalized_flow_to_image(normalized_flow: torch.Tensor) -> torch.Tensor:
+
+ """
+ Converts a batch of normalized flow to an RGB image.
+
+ Args:
+ normalized_flow (torch.Tensor): Normalized flow tensor of shape (N, 2, H, W)
+ Returns:
+ img (Tensor(N, 3, H, W)): Flow visualization image of dtype uint8.
+ """
+
+ N, _, H, W = normalized_flow.shape
+ device = normalized_flow.device
+ flow_image = torch.zeros((N, 3, H, W), dtype=torch.uint8, device=device)
+ colorwheel = _make_colorwheel().to(device) # shape [55x3]
+ num_cols = colorwheel.shape[0]
+ norm = torch.sum(normalized_flow ** 2, dim=1).sqrt()
+ a = torch.atan2(-normalized_flow[:, 1, :, :], -normalized_flow[:, 0, :, :]) / torch.pi
+ fk = (a + 1) / 2 * (num_cols - 1)
+ k0 = torch.floor(fk).to(torch.long)
+ k1 = k0 + 1
+ k1[k1 == num_cols] = 0
+ f = fk - k0
+
+ for c in range(colorwheel.shape[1]):
+ tmp = colorwheel[:, c]
+ col0 = tmp[k0] / 255.0
+ col1 = tmp[k1] / 255.0
+ col = (1 - f) * col0 + f * col1
+ col = 1 - norm * (1 - col)
+ flow_image[:, c, :, :] = torch.floor(255 * col)
+ return flow_image
+
+
+def _make_colorwheel() -> torch.Tensor:
+ """
+ Generates a color wheel for optical flow visualization as presented in:
+ Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007)
+ URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf.
+
+ Returns:
+ colorwheel (Tensor[55, 3]): Colorwheel Tensor.
+ """
+
+ RY = 15
+ YG = 6
+ GC = 4
+ CB = 11
+ BM = 13
+ MR = 6
+
+ ncols = RY + YG + GC + CB + BM + MR
+ colorwheel = torch.zeros((ncols, 3))
+ col = 0
+
+ # RY
+ colorwheel[0:RY, 0] = 255
+ colorwheel[0:RY, 1] = torch.floor(255 * torch.arange(0, RY) / RY)
+ col = col + RY
+ # YG
+ colorwheel[col : col + YG, 0] = 255 - torch.floor(255 * torch.arange(0, YG) / YG)
+ colorwheel[col : col + YG, 1] = 255
+ col = col + YG
+ # GC
+ colorwheel[col : col + GC, 1] = 255
+ colorwheel[col : col + GC, 2] = torch.floor(255 * torch.arange(0, GC) / GC)
+ col = col + GC
+ # CB
+ colorwheel[col : col + CB, 1] = 255 - torch.floor(255 * torch.arange(CB) / CB)
+ colorwheel[col : col + CB, 2] = 255
+ col = col + CB
+ # BM
+ colorwheel[col : col + BM, 2] = 255
+ colorwheel[col : col + BM, 0] = torch.floor(255 * torch.arange(0, BM) / BM)
+ col = col + BM
+ # MR
+ colorwheel[col : col + MR, 2] = 255 - torch.floor(255 * torch.arange(MR) / MR)
+ colorwheel[col : col + MR, 0] = 255
+ return colorwheel
+
+
+def _generate_color_palette(num_objects: int):
+ palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
+ return [tuple((i * palette) % 255) for i in range(num_objects)]
+
+
+def _log_api_usage_once(obj: Any) -> None:
+
+ """
+ Logs API usage(module and name) within an organization.
+ In a large ecosystem, it's often useful to track the PyTorch and
+ TorchVision APIs usage. This API provides the similar functionality to the
+ logging module in the Python stdlib. It can be used for debugging purpose
+ to log which methods are used and by default it is inactive, unless the user
+ manually subscribes a logger via the `SetAPIUsageLogger method `_.
+ Please note it is triggered only once for the same API call within a process.
+ It does not collect any data from open-source users since it is no-op by default.
+ For more information, please refer to
+ * PyTorch note: https://pytorch.org/docs/stable/notes/large_scale_deployments.html#api-usage-logging;
+ * Logging policy: https://github.com/pytorch/vision/issues/5052;
+
+ Args:
+ obj (class instance or method): an object to extract info from.
+ """
+ if not obj.__module__.startswith("torchvision"):
+ return
+ name = obj.__class__.__name__
+ if isinstance(obj, FunctionType):
+ name = obj.__name__
+ torch._C._log_api_usage_once(f"{obj.__module__}.{name}")
diff --git a/clean/video/univfd_video/networks/__init__.py b/clean/video/univfd_video/networks/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/clean/video/univfd_video/networks/base_model.py b/clean/video/univfd_video/networks/base_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..684bdd31004eb9d5664da1aba08dc3ba3b7c4d80
--- /dev/null
+++ b/clean/video/univfd_video/networks/base_model.py
@@ -0,0 +1,58 @@
+import os
+import torch
+import torch.nn as nn
+from torch.nn import init
+from torch.optim import lr_scheduler
+
+
+class BaseModel(nn.Module):
+ def __init__(self, opt):
+ super(BaseModel, self).__init__()
+ self.opt = opt
+ self.total_steps = 0
+ self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
+ self.device = torch.device('cuda:{}'.format(opt.gpu_ids[0])) if opt.gpu_ids else torch.device('cpu')
+
+ def save_networks(self, save_filename):
+ save_path = os.path.join(self.save_dir, save_filename)
+
+ # serialize model and optimizer to dict
+ state_dict = {
+ 'model': self.model.state_dict(),
+ 'optimizer' : self.optimizer.state_dict(),
+ 'total_steps' : self.total_steps,
+ }
+
+ torch.save(state_dict, save_path)
+
+
+ def eval(self):
+ self.model.eval()
+
+ def test(self):
+ with torch.no_grad():
+ self.forward()
+
+
+def init_weights(net, init_type='normal', gain=0.02):
+ def init_func(m):
+ classname = m.__class__.__name__
+ if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
+ if init_type == 'normal':
+ init.normal_(m.weight.data, 0.0, gain)
+ elif init_type == 'xavier':
+ init.xavier_normal_(m.weight.data, gain=gain)
+ elif init_type == 'kaiming':
+ init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+ elif init_type == 'orthogonal':
+ init.orthogonal_(m.weight.data, gain=gain)
+ else:
+ raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
+ if hasattr(m, 'bias') and m.bias is not None:
+ init.constant_(m.bias.data, 0.0)
+ elif classname.find('BatchNorm2d') != -1:
+ init.normal_(m.weight.data, 1.0, gain)
+ init.constant_(m.bias.data, 0.0)
+
+ print('initialize network with %s' % init_type)
+ net.apply(init_func)
diff --git a/clean/video/univfd_video/networks/lpf.py b/clean/video/univfd_video/networks/lpf.py
new file mode 100644
index 0000000000000000000000000000000000000000..f64030bd9a73786f249e03b4d6ce02b32d5ecf92
--- /dev/null
+++ b/clean/video/univfd_video/networks/lpf.py
@@ -0,0 +1,120 @@
+# Copyright (c) 2019, Adobe Inc. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial-ShareAlike
+# 4.0 International Public License. To view a copy of this license, visit
+# https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode.
+
+import torch
+import torch.nn.parallel
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+from IPython import embed
+
+class Downsample(nn.Module):
+ def __init__(self, pad_type='reflect', filt_size=3, stride=2, channels=None, pad_off=0):
+ super(Downsample, self).__init__()
+ self.filt_size = filt_size
+ self.pad_off = pad_off
+ self.pad_sizes = [int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)), int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2))]
+ self.pad_sizes = [pad_size+pad_off for pad_size in self.pad_sizes]
+ self.stride = stride
+ self.off = int((self.stride-1)/2.)
+ self.channels = channels
+
+ # print('Filter size [%i]'%filt_size)
+ if(self.filt_size==1):
+ a = np.array([1.,])
+ elif(self.filt_size==2):
+ a = np.array([1., 1.])
+ elif(self.filt_size==3):
+ a = np.array([1., 2., 1.])
+ elif(self.filt_size==4):
+ a = np.array([1., 3., 3., 1.])
+ elif(self.filt_size==5):
+ a = np.array([1., 4., 6., 4., 1.])
+ elif(self.filt_size==6):
+ a = np.array([1., 5., 10., 10., 5., 1.])
+ elif(self.filt_size==7):
+ a = np.array([1., 6., 15., 20., 15., 6., 1.])
+
+ filt = torch.Tensor(a[:,None]*a[None,:])
+ filt = filt/torch.sum(filt)
+ self.register_buffer('filt', filt[None,None,:,:].repeat((self.channels,1,1,1)))
+
+ self.pad = get_pad_layer(pad_type)(self.pad_sizes)
+
+ def forward(self, inp):
+ if(self.filt_size==1):
+ if(self.pad_off==0):
+ return inp[:,:,::self.stride,::self.stride]
+ else:
+ return self.pad(inp)[:,:,::self.stride,::self.stride]
+ else:
+ return F.conv2d(self.pad(inp), self.filt, stride=self.stride, groups=inp.shape[1])
+
+def get_pad_layer(pad_type):
+ if(pad_type in ['refl','reflect']):
+ PadLayer = nn.ReflectionPad2d
+ elif(pad_type in ['repl','replicate']):
+ PadLayer = nn.ReplicationPad2d
+ elif(pad_type=='zero'):
+ PadLayer = nn.ZeroPad2d
+ else:
+ print('Pad type [%s] not recognized'%pad_type)
+ return PadLayer
+
+
+class Downsample1D(nn.Module):
+ def __init__(self, pad_type='reflect', filt_size=3, stride=2, channels=None, pad_off=0):
+ super(Downsample1D, self).__init__()
+ self.filt_size = filt_size
+ self.pad_off = pad_off
+ self.pad_sizes = [int(1. * (filt_size - 1) / 2), int(np.ceil(1. * (filt_size - 1) / 2))]
+ self.pad_sizes = [pad_size + pad_off for pad_size in self.pad_sizes]
+ self.stride = stride
+ self.off = int((self.stride - 1) / 2.)
+ self.channels = channels
+
+ # print('Filter size [%i]' % filt_size)
+ if(self.filt_size == 1):
+ a = np.array([1., ])
+ elif(self.filt_size == 2):
+ a = np.array([1., 1.])
+ elif(self.filt_size == 3):
+ a = np.array([1., 2., 1.])
+ elif(self.filt_size == 4):
+ a = np.array([1., 3., 3., 1.])
+ elif(self.filt_size == 5):
+ a = np.array([1., 4., 6., 4., 1.])
+ elif(self.filt_size == 6):
+ a = np.array([1., 5., 10., 10., 5., 1.])
+ elif(self.filt_size == 7):
+ a = np.array([1., 6., 15., 20., 15., 6., 1.])
+
+ filt = torch.Tensor(a)
+ filt = filt / torch.sum(filt)
+ self.register_buffer('filt', filt[None, None, :].repeat((self.channels, 1, 1)))
+
+ self.pad = get_pad_layer_1d(pad_type)(self.pad_sizes)
+
+ def forward(self, inp):
+ if(self.filt_size == 1):
+ if(self.pad_off == 0):
+ return inp[:, :, ::self.stride]
+ else:
+ return self.pad(inp)[:, :, ::self.stride]
+ else:
+ return F.conv1d(self.pad(inp), self.filt, stride=self.stride, groups=inp.shape[1])
+
+
+def get_pad_layer_1d(pad_type):
+ if(pad_type in ['refl', 'reflect']):
+ PadLayer = nn.ReflectionPad1d
+ elif(pad_type in ['repl', 'replicate']):
+ PadLayer = nn.ReplicationPad1d
+ elif(pad_type == 'zero'):
+ PadLayer = nn.ZeroPad1d
+ else:
+ print('Pad type [%s] not recognized' % pad_type)
+ return PadLayer
diff --git a/clean/video/univfd_video/networks/resnet_lpf.py b/clean/video/univfd_video/networks/resnet_lpf.py
new file mode 100644
index 0000000000000000000000000000000000000000..f9e34254eadc00e701245d03ffd86c25c114ce3f
--- /dev/null
+++ b/clean/video/univfd_video/networks/resnet_lpf.py
@@ -0,0 +1,313 @@
+# This code is built from the PyTorch examples repository: https://github.com/pytorch/vision/tree/master/torchvision/models.
+# Copyright (c) 2017 Torch Contributors.
+# The Pytorch examples are available under the BSD 3-Clause License.
+#
+# ==========================================================================================
+#
+# Adobe’s modifications are Copyright 2019 Adobe. All rights reserved.
+# Adobe’s modifications are licensed under the Creative Commons Attribution-NonCommercial-ShareAlike
+# 4.0 International Public License (CC-NC-SA-4.0). To view a copy of the license, visit
+# https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode.
+#
+# ==========================================================================================
+#
+# BSD-3 License
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+from .lpf import *
+
+__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
+ 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d']
+
+
+# model_urls = {
+# 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+# 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+# 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+# 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+# 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+# }
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1):
+ """3x3 convolution with padding"""
+ return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+ padding=1, groups=groups, bias=False)
+
+def conv1x1(in_planes, out_planes, stride=1):
+ """1x1 convolution"""
+ return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+class BasicBlock(nn.Module):
+ expansion = 1
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, norm_layer=None, filter_size=1):
+ super(BasicBlock, self).__init__()
+ if norm_layer is None:
+ norm_layer = nn.BatchNorm2d
+ if groups != 1:
+ raise ValueError('BasicBlock only supports groups=1')
+ # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+ self.conv1 = conv3x3(inplanes, planes)
+ self.bn1 = norm_layer(planes)
+ self.relu = nn.ReLU(inplace=True)
+ if(stride==1):
+ self.conv2 = conv3x3(planes,planes)
+ else:
+ self.conv2 = nn.Sequential(Downsample(filt_size=filter_size, stride=stride, channels=planes),
+ conv3x3(planes, planes),)
+ self.bn2 = norm_layer(planes)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x):
+ identity = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+
+ if self.downsample is not None:
+ identity = self.downsample(x)
+
+ out += identity
+ out = self.relu(out)
+
+ return out
+
+
+class Bottleneck(nn.Module):
+ expansion = 4
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, norm_layer=None, filter_size=1):
+ super(Bottleneck, self).__init__()
+ if norm_layer is None:
+ norm_layer = nn.BatchNorm2d
+ # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+ self.conv1 = conv1x1(inplanes, planes)
+ self.bn1 = norm_layer(planes)
+ self.conv2 = conv3x3(planes, planes, groups) # stride moved
+ self.bn2 = norm_layer(planes)
+ if(stride==1):
+ self.conv3 = conv1x1(planes, planes * self.expansion)
+ else:
+ self.conv3 = nn.Sequential(Downsample(filt_size=filter_size, stride=stride, channels=planes),
+ conv1x1(planes, planes * self.expansion))
+ self.bn3 = norm_layer(planes * self.expansion)
+ self.relu = nn.ReLU(inplace=True)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x):
+ identity = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+ out = self.relu(out)
+
+ out = self.conv3(out)
+ out = self.bn3(out)
+
+ if self.downsample is not None:
+ identity = self.downsample(x)
+
+ out += identity
+ out = self.relu(out)
+
+ return out
+
+
+class ResNet(nn.Module):
+
+ def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
+ groups=1, width_per_group=64, norm_layer=None, filter_size=1, pool_only=True):
+ super(ResNet, self).__init__()
+ if norm_layer is None:
+ norm_layer = nn.BatchNorm2d
+ planes = [int(width_per_group * groups * 2 ** i) for i in range(4)]
+ self.inplanes = planes[0]
+
+ if(pool_only):
+ self.conv1 = nn.Conv2d(3, planes[0], kernel_size=7, stride=2, padding=3, bias=False)
+ else:
+ self.conv1 = nn.Conv2d(3, planes[0], kernel_size=7, stride=1, padding=3, bias=False)
+ self.bn1 = norm_layer(planes[0])
+ self.relu = nn.ReLU(inplace=True)
+
+ if(pool_only):
+ self.maxpool = nn.Sequential(*[nn.MaxPool2d(kernel_size=2, stride=1),
+ Downsample(filt_size=filter_size, stride=2, channels=planes[0])])
+ else:
+ self.maxpool = nn.Sequential(*[Downsample(filt_size=filter_size, stride=2, channels=planes[0]),
+ nn.MaxPool2d(kernel_size=2, stride=1),
+ Downsample(filt_size=filter_size, stride=2, channels=planes[0])])
+
+ self.layer1 = self._make_layer(block, planes[0], layers[0], groups=groups, norm_layer=norm_layer)
+ self.layer2 = self._make_layer(block, planes[1], layers[1], stride=2, groups=groups, norm_layer=norm_layer, filter_size=filter_size)
+ self.layer3 = self._make_layer(block, planes[2], layers[2], stride=2, groups=groups, norm_layer=norm_layer, filter_size=filter_size)
+ self.layer4 = self._make_layer(block, planes[3], layers[3], stride=2, groups=groups, norm_layer=norm_layer, filter_size=filter_size)
+ self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+ self.fc = nn.Linear(planes[3] * block.expansion, num_classes)
+
+ for m in self.modules():
+ if isinstance(m, nn.Conv2d):
+ if(m.in_channels!=m.out_channels or m.out_channels!=m.groups or m.bias is not None):
+ # don't want to reinitialize downsample layers, code assuming normal conv layers will not have these characteristics
+ nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+ else:
+ print('Not initializing')
+ elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+ nn.init.constant_(m.weight, 1)
+ nn.init.constant_(m.bias, 0)
+
+ # Zero-initialize the last BN in each residual branch,
+ # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+ # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+ if zero_init_residual:
+ for m in self.modules():
+ if isinstance(m, Bottleneck):
+ nn.init.constant_(m.bn3.weight, 0)
+ elif isinstance(m, BasicBlock):
+ nn.init.constant_(m.bn2.weight, 0)
+
+ def _make_layer(self, block, planes, blocks, stride=1, groups=1, norm_layer=None, filter_size=1):
+ if norm_layer is None:
+ norm_layer = nn.BatchNorm2d
+ downsample = None
+ if stride != 1 or self.inplanes != planes * block.expansion:
+ # downsample = nn.Sequential(
+ # conv1x1(self.inplanes, planes * block.expansion, stride, filter_size=filter_size),
+ # norm_layer(planes * block.expansion),
+ # )
+
+ downsample = [Downsample(filt_size=filter_size, stride=stride, channels=self.inplanes),] if(stride !=1) else []
+ downsample += [conv1x1(self.inplanes, planes * block.expansion, 1),
+ norm_layer(planes * block.expansion)]
+ # print(downsample)
+ downsample = nn.Sequential(*downsample)
+
+ layers = []
+ layers.append(block(self.inplanes, planes, stride, downsample, groups, norm_layer, filter_size=filter_size))
+ self.inplanes = planes * block.expansion
+ for _ in range(1, blocks):
+ layers.append(block(self.inplanes, planes, groups=groups, norm_layer=norm_layer, filter_size=filter_size))
+
+ return nn.Sequential(*layers)
+
+ def forward(self, x):
+ x = self.conv1(x)
+ x = self.bn1(x)
+ x = self.relu(x)
+ x = self.maxpool(x)
+
+ x = self.layer1(x)
+ x = self.layer2(x)
+ x = self.layer3(x)
+ x = self.layer4(x)
+
+ x = self.avgpool(x)
+ x = x.view(x.size(0), -1)
+ x = self.fc(x)
+
+ return x
+
+
+def resnet18(pretrained=False, filter_size=1, pool_only=True, **kwargs):
+ """Constructs a ResNet-18 model.
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ """
+ model = ResNet(BasicBlock, [2, 2, 2, 2], filter_size=filter_size, pool_only=pool_only, **kwargs)
+ if pretrained:
+ model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
+ return model
+
+
+def resnet34(pretrained=False, filter_size=1, pool_only=True, **kwargs):
+ """Constructs a ResNet-34 model.
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ """
+ model = ResNet(BasicBlock, [3, 4, 6, 3], filter_size=filter_size, pool_only=pool_only, **kwargs)
+ if pretrained:
+ model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
+ return model
+
+
+def resnet50(pretrained=False, filter_size=1, pool_only=True, **kwargs):
+ """Constructs a ResNet-50 model.
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ """
+ model = ResNet(Bottleneck, [3, 4, 6, 3], filter_size=filter_size, pool_only=pool_only, **kwargs)
+ if pretrained:
+ model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
+ return model
+
+
+def resnet101(pretrained=False, filter_size=1, pool_only=True, **kwargs):
+ """Constructs a ResNet-101 model.
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ """
+ model = ResNet(Bottleneck, [3, 4, 23, 3], filter_size=filter_size, pool_only=pool_only, **kwargs)
+ if pretrained:
+ model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
+ return model
+
+
+def resnet152(pretrained=False, filter_size=1, pool_only=True, **kwargs):
+ """Constructs a ResNet-152 model.
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ """
+ model = ResNet(Bottleneck, [3, 8, 36, 3], filter_size=filter_size, pool_only=pool_only, **kwargs)
+ if pretrained:
+ model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
+ return model
+
+
+def resnext50_32x4d(pretrained=False, filter_size=1, pool_only=True, **kwargs):
+ model = ResNet(Bottleneck, [3, 4, 6, 3], groups=4, width_per_group=32, filter_size=filter_size, pool_only=pool_only, **kwargs)
+ # if pretrained:
+ # model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
+ return model
+
+
+def resnext101_32x8d(pretrained=False, filter_size=1, pool_only=True, **kwargs):
+ model = ResNet(Bottleneck, [3, 4, 23, 3], groups=8, width_per_group=32, filter_size=filter_size, pool_only=pool_only, **kwargs)
+ # if pretrained:
+ # model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
+ return model
diff --git a/clean/video/univfd_video/networks/trainer.py b/clean/video/univfd_video/networks/trainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..73d5bc0bcf5d48117e473900b38088c738f9c555
--- /dev/null
+++ b/clean/video/univfd_video/networks/trainer.py
@@ -0,0 +1,74 @@
+import functools
+import torch
+import torch.nn as nn
+from networks.base_model import BaseModel, init_weights
+import sys
+from models import get_model
+
+class Trainer(BaseModel):
+ def name(self):
+ return 'Trainer'
+
+ def __init__(self, opt):
+ super(Trainer, self).__init__(opt)
+ self.opt = opt
+ self.model = get_model(opt.arch)
+ torch.nn.init.normal_(self.model.fc.weight.data, 0.0, opt.init_gain)
+
+ if opt.fix_backbone:
+ params = []
+ for name, p in self.model.named_parameters():
+ if name=="fc.weight" or name=="fc.bias":
+ params.append(p)
+ else:
+ p.requires_grad = False
+ else:
+ print("Your backbone is not fixed. Are you sure you want to proceed? If this is a mistake, enable the --fix_backbone command during training and rerun")
+ import time
+ time.sleep(3)
+ params = self.model.parameters()
+
+
+
+ if opt.optim == 'adam':
+ self.optimizer = torch.optim.AdamW(params, lr=opt.lr, betas=(opt.beta1, 0.999), weight_decay=opt.weight_decay)
+ elif opt.optim == 'sgd':
+ self.optimizer = torch.optim.SGD(params, lr=opt.lr, momentum=0.0, weight_decay=opt.weight_decay)
+ else:
+ raise ValueError("optim should be [adam, sgd]")
+
+ self.loss_fn = nn.BCEWithLogitsLoss()
+
+ self.model.to(opt.gpu_ids[0])
+
+
+ def adjust_learning_rate(self, min_lr=1e-6):
+ for param_group in self.optimizer.param_groups:
+ param_group['lr'] /= 10.
+ if param_group['lr'] < min_lr:
+ return False
+ return True
+
+
+ def set_input(self, input):
+ self.input = input[0].to(self.device)
+ self.label = input[1].to(self.device).float()
+
+
+ def forward(self):
+ self.output = self.model(self.input)
+ self.output = self.output.view(-1).unsqueeze(1)
+
+
+ def get_loss(self):
+ return self.loss_fn(self.output.squeeze(1), self.label)
+
+ def optimize_parameters(self):
+ self.forward()
+ self.loss = self.loss_fn(self.output.squeeze(1), self.label)
+ self.optimizer.zero_grad()
+ self.loss.backward()
+ self.optimizer.step()
+
+
+
diff --git a/clean/video/univfd_video/options/__init__.py b/clean/video/univfd_video/options/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/clean/video/univfd_video/options/base_options.py b/clean/video/univfd_video/options/base_options.py
new file mode 100644
index 0000000000000000000000000000000000000000..69d35b17cf409eac0a55a74ed187f23028cb93c5
--- /dev/null
+++ b/clean/video/univfd_video/options/base_options.py
@@ -0,0 +1,117 @@
+import argparse
+import os
+import util
+import torch
+
+
+class BaseOptions():
+ def __init__(self):
+ self.initialized = False
+
+ def initialize(self, parser):
+ parser.add_argument('--mode', default='binary')
+ parser.add_argument('--arch', type=str, default='res50', help='see my_models/__init__.py')
+ parser.add_argument('--fix_backbone', action='store_true')
+
+ # data augmentation
+ parser.add_argument('--rz_interp', default='bilinear')
+ parser.add_argument('--blur_prob', type=float, default=0.5)
+ parser.add_argument('--blur_sig', default='0.0,3.0')
+ parser.add_argument('--jpg_prob', type=float, default=0.5)
+ parser.add_argument('--jpg_method', default='cv2,pil')
+ parser.add_argument('--jpg_qual', default='30,100')
+
+
+ parser.add_argument('--real_list_path', default=None, help='only used if data_mode==ours: path for the list of real images, which should contain train.pickle and val.pickle')
+ parser.add_argument('--fake_list_path', default=None, help='only used if data_mode==ours: path for the list of fake images, which should contain train.pickle and val.pickle')
+ parser.add_argument('--wang2020_data_path', default=None, help='only used if data_mode==wang2020 it should contain train and test folders')
+ parser.add_argument('--data_mode', default='ours', help='wang2020 or ours')
+ parser.add_argument('--data_label', default='train', help='label to decide whether train or validation dataset')
+ parser.add_argument('--weight_decay', type=float, default=0.0, help='loss weight for l2 reg')
+
+ parser.add_argument('--class_bal', action='store_true') # what is this ?
+ parser.add_argument('--batch_size', type=int, default=256, help='input batch size')
+ parser.add_argument('--loadSize', type=int, default=256, help='scale images to this size')
+ parser.add_argument('--cropSize', type=int, default=224, help='then crop to this size')
+ parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU')
+ parser.add_argument('--name', type=str, default='experiment_name', help='name of the experiment. It decides where to store samples and models')
+ parser.add_argument('--num_threads', default=4, type=int, help='# threads for loading data')
+ parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
+ parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')
+ parser.add_argument('--resize_or_crop', type=str, default='scale_and_crop', help='scaling and cropping of images at load time [resize_and_crop|crop|scale_width|scale_width_and_crop|none]')
+ parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data augmentation')
+ parser.add_argument('--init_type', type=str, default='normal', help='network initialization [normal|xavier|kaiming|orthogonal]')
+ parser.add_argument('--init_gain', type=float, default=0.02, help='scaling factor for normal, xavier and orthogonal.')
+ parser.add_argument('--suffix', default='', type=str, help='customized suffix: opt.name = opt.name + suffix: e.g., {model}_{netG}_size{loadSize}')
+ self.initialized = True
+ return parser
+
+ def gather_options(self):
+ # initialize parser with basic options
+ if not self.initialized:
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+ parser = self.initialize(parser)
+
+ # get the basic options
+ opt, _ = parser.parse_known_args()
+ self.parser = parser
+
+ return parser.parse_args()
+
+ def print_options(self, opt):
+ message = ''
+ message += '----------------- Options ---------------\n'
+ for k, v in sorted(vars(opt).items()):
+ comment = ''
+ default = self.parser.get_default(k)
+ if v != default:
+ comment = '\t[default: %s]' % str(default)
+ message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
+ message += '----------------- End -------------------'
+ print(message)
+
+ # save to the disk
+ expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
+ util.mkdirs(expr_dir)
+ file_name = os.path.join(expr_dir, 'opt.txt')
+ with open(file_name, 'wt') as opt_file:
+ opt_file.write(message)
+ opt_file.write('\n')
+
+ def parse(self, print_options=True):
+
+ opt = self.gather_options()
+ opt.isTrain = self.isTrain # train or test
+
+ # process opt.suffix
+ if opt.suffix:
+ suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else ''
+ opt.name = opt.name + suffix
+
+ if print_options:
+ self.print_options(opt)
+
+ # set gpu ids
+ str_ids = opt.gpu_ids.split(',')
+ opt.gpu_ids = []
+ for str_id in str_ids:
+ id = int(str_id)
+ if id >= 0:
+ opt.gpu_ids.append(id)
+ if len(opt.gpu_ids) > 0:
+ torch.cuda.set_device(opt.gpu_ids[0])
+
+ # additional
+ #opt.classes = opt.classes.split(',')
+ opt.rz_interp = opt.rz_interp.split(',')
+ opt.blur_sig = [float(s) for s in opt.blur_sig.split(',')]
+ opt.jpg_method = opt.jpg_method.split(',')
+ opt.jpg_qual = [int(s) for s in opt.jpg_qual.split(',')]
+ if len(opt.jpg_qual) == 2:
+ opt.jpg_qual = list(range(opt.jpg_qual[0], opt.jpg_qual[1] + 1))
+ elif len(opt.jpg_qual) > 2:
+ raise ValueError("Shouldn't have more than 2 values for --jpg_qual.")
+
+ self.opt = opt
+ return self.opt
diff --git a/clean/video/univfd_video/options/test_options.py b/clean/video/univfd_video/options/test_options.py
new file mode 100644
index 0000000000000000000000000000000000000000..f824c7aae81d325bf81e16c5a5e2c1931ccd5b34
--- /dev/null
+++ b/clean/video/univfd_video/options/test_options.py
@@ -0,0 +1,13 @@
+from .base_options import BaseOptions
+
+
+class TestOptions(BaseOptions):
+ def initialize(self, parser):
+ parser = BaseOptions.initialize(self, parser)
+ parser.add_argument('--model_path')
+ parser.add_argument('--no_resize', action='store_true')
+ parser.add_argument('--no_crop', action='store_true')
+ parser.add_argument('--eval', action='store_true', help='use eval mode during test time.')
+
+ self.isTrain = False
+ return parser
diff --git a/clean/video/univfd_video/options/train_options.py b/clean/video/univfd_video/options/train_options.py
new file mode 100644
index 0000000000000000000000000000000000000000..e7f1e54cea63ef424ad84d1ccaa1538e74119466
--- /dev/null
+++ b/clean/video/univfd_video/options/train_options.py
@@ -0,0 +1,22 @@
+from .base_options import BaseOptions
+
+
+class TrainOptions(BaseOptions):
+ def initialize(self, parser):
+ parser = BaseOptions.initialize(self, parser)
+ parser.add_argument('--earlystop_epoch', type=int, default=5)
+ parser.add_argument('--data_aug', action='store_true', help='if specified, perform additional data augmentation (photometric, blurring, jpegging)')
+ parser.add_argument('--optim', type=str, default='adam', help='optim to use [sgd, adam]')
+ parser.add_argument('--new_optim', action='store_true', help='new optimizer instead of loading the optim state')
+ parser.add_argument('--loss_freq', type=int, default=400, help='frequency of showing loss on tensorboard')
+ parser.add_argument('--save_epoch_freq', type=int, default=1, help='frequency of saving checkpoints at the end of epochs')
+ parser.add_argument('--epoch_count', type=int, default=1, help='the starting epoch count, we save the model by , +, ...')
+ parser.add_argument('--last_epoch', type=int, default=-1, help='starting epoch count for scheduler intialization')
+ parser.add_argument('--train_split', type=str, default='train', help='train, val, test, etc')
+ parser.add_argument('--val_split', type=str, default='val', help='train, val, test, etc')
+ parser.add_argument('--niter', type=int, default=100, help='total epoches')
+ parser.add_argument('--beta1', type=float, default=0.9, help='momentum term of adam')
+ parser.add_argument('--lr', type=float, default=0.0001, help='initial learning rate for adam')
+
+ self.isTrain = True
+ return parser
diff --git a/clean/video/univfd_video/pretrained_weights/fc_weights.pth b/clean/video/univfd_video/pretrained_weights/fc_weights.pth
new file mode 100644
index 0000000000000000000000000000000000000000..989708188fa14aa3f7ddfcddb579d7b9426d5e8e
--- /dev/null
+++ b/clean/video/univfd_video/pretrained_weights/fc_weights.pth
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:477100745713bcc957beb2b40859536859b6483fd6301b3b9293151b194c7847
+size 4083
diff --git a/clean/video/univfd_video/test.sh b/clean/video/univfd_video/test.sh
new file mode 100644
index 0000000000000000000000000000000000000000..983d05b7c4cf0e8b0a977b835849a102fa8b75fc
--- /dev/null
+++ b/clean/video/univfd_video/test.sh
@@ -0,0 +1 @@
+CUDA_VISIBLE_DEVICES=0 python3 validate.py --arch=CLIP:ViT-L/14 --ckpt=pretrained_weights/fc_weights.pth --result_folder=clip_vitl14
diff --git a/clean/video/univfd_video/train.py b/clean/video/univfd_video/train.py
new file mode 100644
index 0000000000000000000000000000000000000000..c24917730db61c90f62015ca68572b83219e2050
--- /dev/null
+++ b/clean/video/univfd_video/train.py
@@ -0,0 +1,85 @@
+import os
+import time
+from tensorboardX import SummaryWriter
+
+from validate import validate
+from data import create_dataloader
+from earlystop import EarlyStopping
+from networks.trainer import Trainer
+from options.train_options import TrainOptions
+
+
+"""Currently assumes jpg_prob, blur_prob 0 or 1"""
+def get_val_opt():
+ val_opt = TrainOptions().parse(print_options=False)
+ val_opt.isTrain = False
+ val_opt.no_resize = False
+ val_opt.no_crop = False
+ val_opt.serial_batches = True
+ val_opt.data_label = 'val'
+ val_opt.jpg_method = ['pil']
+ if len(val_opt.blur_sig) == 2:
+ b_sig = val_opt.blur_sig
+ val_opt.blur_sig = [(b_sig[0] + b_sig[1]) / 2]
+ if len(val_opt.jpg_qual) != 1:
+ j_qual = val_opt.jpg_qual
+ val_opt.jpg_qual = [int((j_qual[0] + j_qual[-1]) / 2)]
+
+ return val_opt
+
+
+
+if __name__ == '__main__':
+ opt = TrainOptions().parse()
+ val_opt = get_val_opt()
+
+ model = Trainer(opt)
+
+ data_loader = create_dataloader(opt)
+ val_loader = create_dataloader(val_opt)
+
+ train_writer = SummaryWriter(os.path.join(opt.checkpoints_dir, opt.name, "train"))
+ val_writer = SummaryWriter(os.path.join(opt.checkpoints_dir, opt.name, "val"))
+
+ early_stopping = EarlyStopping(patience=opt.earlystop_epoch, delta=-0.001, verbose=True)
+ start_time = time.time()
+ print ("Length of data loader: %d" %(len(data_loader)))
+ for epoch in range(opt.niter):
+
+ for i, data in enumerate(data_loader):
+ model.total_steps += 1
+
+ model.set_input(data)
+ model.optimize_parameters()
+
+ if model.total_steps % opt.loss_freq == 0:
+ print("Train loss: {} at step: {}".format(model.loss, model.total_steps))
+ train_writer.add_scalar('loss', model.loss, model.total_steps)
+ print("Iter time: ", ((time.time()-start_time)/model.total_steps) )
+
+ if model.total_steps in [10,30,50,100,1000,5000,10000] and False: # save models at these iters
+ model.save_networks('model_iters_%s.pth' % model.total_steps)
+
+ if epoch % opt.save_epoch_freq == 0:
+ print('saving the model at the end of epoch %d' % (epoch))
+ model.save_networks( 'model_epoch_best.pth' )
+ model.save_networks( 'model_epoch_%s.pth' % epoch )
+
+ # Validation
+ model.eval()
+ ap, r_acc, f_acc, acc = validate(model.model, val_loader)
+ val_writer.add_scalar('accuracy', acc, model.total_steps)
+ val_writer.add_scalar('ap', ap, model.total_steps)
+ print("(Val @ epoch {}) acc: {}; ap: {}".format(epoch, acc, ap))
+
+ early_stopping(acc, model)
+ if early_stopping.early_stop:
+ cont_train = model.adjust_learning_rate()
+ if cont_train:
+ print("Learning rate dropped by 10, continue training...")
+ early_stopping = EarlyStopping(patience=opt.earlystop_epoch, delta=-0.002, verbose=True)
+ else:
+ print("Early stopping.")
+ break
+ model.train()
+
diff --git a/clean/video/univfd_video/validate.py b/clean/video/univfd_video/validate.py
new file mode 100644
index 0000000000000000000000000000000000000000..4887014457e584b5801e868a229db98bbd70117c
--- /dev/null
+++ b/clean/video/univfd_video/validate.py
@@ -0,0 +1,312 @@
+import argparse
+from ast import arg
+import os
+import csv
+import torch
+import torchvision.transforms as transforms
+import torch.utils.data
+import numpy as np
+from sklearn.metrics import average_precision_score, precision_recall_curve, accuracy_score
+from torch.utils.data import Dataset
+import sys
+from models import get_model
+from PIL import Image
+import pickle
+from tqdm import tqdm
+from io import BytesIO
+from copy import deepcopy
+from dataset_paths import DATASET_PATHS
+import random
+import shutil
+from scipy.ndimage.filters import gaussian_filter
+
+SEED = 0
+def set_seed():
+ torch.manual_seed(SEED)
+ torch.cuda.manual_seed(SEED)
+ np.random.seed(SEED)
+ random.seed(SEED)
+
+
+MEAN = {
+ "imagenet":[0.485, 0.456, 0.406],
+ "clip":[0.48145466, 0.4578275, 0.40821073]
+}
+
+STD = {
+ "imagenet":[0.229, 0.224, 0.225],
+ "clip":[0.26862954, 0.26130258, 0.27577711]
+}
+
+
+
+
+
+def find_best_threshold(y_true, y_pred):
+ "We assume first half is real 0, and the second half is fake 1"
+
+ N = y_true.shape[0]
+
+ if y_pred[0:N//2].max() <= y_pred[N//2:N].min(): # perfectly separable case
+ return (y_pred[0:N//2].max() + y_pred[N//2:N].min()) / 2
+
+ best_acc = 0
+ best_thres = 0
+ for thres in y_pred:
+ temp = deepcopy(y_pred)
+ temp[temp>=thres] = 1
+ temp[temp= best_acc:
+ best_thres = thres
+ best_acc = acc
+
+ return best_thres
+
+
+
+def png2jpg(img, quality):
+ out = BytesIO()
+ img.save(out, format='jpeg', quality=quality) # ranging from 0-95, 75 is default
+ img = Image.open(out)
+ # load from memory before ByteIO closes
+ img = np.array(img)
+ out.close()
+ return Image.fromarray(img)
+
+
+def gaussian_blur(img, sigma):
+ img = np.array(img)
+
+ gaussian_filter(img[:,:,0], output=img[:,:,0], sigma=sigma)
+ gaussian_filter(img[:,:,1], output=img[:,:,1], sigma=sigma)
+ gaussian_filter(img[:,:,2], output=img[:,:,2], sigma=sigma)
+
+ return Image.fromarray(img)
+
+
+
+def calculate_acc(y_true, y_pred, thres):
+ r_acc = accuracy_score(y_true[y_true==0], y_pred[y_true==0] > thres)
+ f_acc = accuracy_score(y_true[y_true==1], y_pred[y_true==1] > thres)
+ acc = accuracy_score(y_true, y_pred > thres)
+ return r_acc, f_acc, acc
+
+
+def validate(model, loader, find_thres=False):
+
+ with torch.no_grad():
+ y_true, y_pred = [], []
+ print ("Length of dataset: %d" %(len(loader)))
+ for img, label in loader:
+ in_tens = img.cuda()
+
+ y_pred.extend(model(in_tens).sigmoid().flatten().tolist())
+ y_true.extend(label.flatten().tolist())
+
+ y_true, y_pred = np.array(y_true), np.array(y_pred)
+
+ # ================== save this if you want to plot the curves =========== #
+ # torch.save( torch.stack( [torch.tensor(y_true), torch.tensor(y_pred)] ), 'baseline_predication_for_pr_roc_curve.pth' )
+ # exit()
+ # =================================================================== #
+
+ # Get AP
+ ap = average_precision_score(y_true, y_pred)
+
+ # Acc based on 0.5
+ r_acc0, f_acc0, acc0 = calculate_acc(y_true, y_pred, 0.5)
+ if not find_thres:
+ return ap, r_acc0, f_acc0, acc0
+
+
+ # Acc based on the best thres
+ best_thres = find_best_threshold(y_true, y_pred)
+ r_acc1, f_acc1, acc1 = calculate_acc(y_true, y_pred, best_thres)
+
+ return ap, r_acc0, f_acc0, acc0, r_acc1, f_acc1, acc1, best_thres
+
+
+
+
+
+
+# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = #
+
+
+
+
+def recursively_read(rootdir, must_contain, exts=["png", "jpg", "JPEG", "jpeg", "bmp"]):
+ out = []
+ for r, d, f in os.walk(rootdir):
+ for file in f:
+ if (file.split('.')[1] in exts) and (must_contain in os.path.join(r, file)):
+ out.append(os.path.join(r, file))
+ return out
+
+
+def get_list(path, must_contain=''):
+ if ".pickle" in path:
+ with open(path, 'rb') as f:
+ image_list = pickle.load(f)
+ image_list = [ item for item in image_list if must_contain in item ]
+ else:
+ image_list = recursively_read(path, must_contain)
+ return image_list
+
+
+
+
+
+class RealFakeDataset(Dataset):
+ def __init__(self, real_path,
+ fake_path,
+ data_mode,
+ max_sample,
+ arch,
+ jpeg_quality=None,
+ gaussian_sigma=None):
+
+ assert data_mode in ["wang2020", "ours"]
+ self.jpeg_quality = jpeg_quality
+ self.gaussian_sigma = gaussian_sigma
+
+ # = = = = = = data path = = = = = = = = = #
+ if type(real_path) == str and type(fake_path) == str:
+ real_list, fake_list = self.read_path(real_path, fake_path, data_mode, max_sample)
+ else:
+ real_list = []
+ fake_list = []
+ for real_p, fake_p in zip(real_path, fake_path):
+ real_l, fake_l = self.read_path(real_p, fake_p, data_mode, max_sample)
+ real_list += real_l
+ fake_list += fake_l
+
+ self.total_list = real_list + fake_list
+
+
+ # = = = = = = label = = = = = = = = = #
+
+ self.labels_dict = {}
+ for i in real_list:
+ self.labels_dict[i] = 0
+ for i in fake_list:
+ self.labels_dict[i] = 1
+
+ stat_from = "imagenet" if arch.lower().startswith("imagenet") else "clip"
+ self.transform = transforms.Compose([
+ transforms.CenterCrop(224),
+ transforms.ToTensor(),
+ transforms.Normalize( mean=MEAN[stat_from], std=STD[stat_from] ),
+ ])
+
+
+ def read_path(self, real_path, fake_path, data_mode, max_sample):
+
+ if data_mode == 'wang2020':
+ real_list = get_list(real_path, must_contain='0_real')
+ fake_list = get_list(fake_path, must_contain='1_fake')
+ else:
+ real_list = get_list(real_path)
+ fake_list = get_list(fake_path)
+
+
+ if max_sample is not None:
+ if (max_sample > len(real_list)) or (max_sample > len(fake_list)):
+ max_sample = 100
+ print("not enough images, max_sample falling to 100")
+ random.shuffle(real_list)
+ random.shuffle(fake_list)
+ real_list = real_list[0:max_sample]
+ fake_list = fake_list[0:max_sample]
+
+ assert len(real_list) == len(fake_list)
+
+ return real_list, fake_list
+
+
+
+ def __len__(self):
+ return len(self.total_list)
+
+ def __getitem__(self, idx):
+
+ img_path = self.total_list[idx]
+
+ label = self.labels_dict[img_path]
+ img = Image.open(img_path).convert("RGB")
+
+ if self.gaussian_sigma is not None:
+ img = gaussian_blur(img, self.gaussian_sigma)
+ if self.jpeg_quality is not None:
+ img = png2jpg(img, self.jpeg_quality)
+
+ img = self.transform(img)
+ return img, label
+
+
+
+
+
+if __name__ == '__main__':
+
+
+ parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+ parser.add_argument('--real_path', type=str, default=None, help='dir name or a pickle')
+ parser.add_argument('--fake_path', type=str, default=None, help='dir name or a pickle')
+ parser.add_argument('--data_mode', type=str, default=None, help='wang2020 or ours')
+ parser.add_argument('--max_sample', type=int, default=1000, help='only check this number of images for both fake/real')
+
+ parser.add_argument('--arch', type=str, default='res50')
+ parser.add_argument('--ckpt', type=str, default='./pretrained_weights/fc_weights.pth')
+
+ parser.add_argument('--result_folder', type=str, default='result', help='')
+ parser.add_argument('--batch_size', type=int, default=128)
+
+ parser.add_argument('--jpeg_quality', type=int, default=None, help="100, 90, 80, ... 30. Used to test robustness of our model. Not apply if None")
+ parser.add_argument('--gaussian_sigma', type=int, default=None, help="0,1,2,3,4. Used to test robustness of our model. Not apply if None")
+
+
+ opt = parser.parse_args()
+
+
+ if os.path.exists(opt.result_folder):
+ shutil.rmtree(opt.result_folder)
+ os.makedirs(opt.result_folder)
+
+ model = get_model(opt.arch)
+ state_dict = torch.load(opt.ckpt, map_location='cpu')
+ model.fc.load_state_dict(state_dict)
+ print ("Model loaded..")
+ model.eval()
+ model.cuda()
+
+ if (opt.real_path == None) or (opt.fake_path == None) or (opt.data_mode == None):
+ dataset_paths = DATASET_PATHS
+ else:
+ dataset_paths = [ dict(real_path=opt.real_path, fake_path=opt.fake_path, data_mode=opt.data_mode) ]
+
+
+
+ for dataset_path in (dataset_paths):
+ set_seed()
+
+ dataset = RealFakeDataset( dataset_path['real_path'],
+ dataset_path['fake_path'],
+ dataset_path['data_mode'],
+ opt.max_sample,
+ opt.arch,
+ jpeg_quality=opt.jpeg_quality,
+ gaussian_sigma=opt.gaussian_sigma,
+ )
+
+ loader = torch.utils.data.DataLoader(dataset, batch_size=opt.batch_size, shuffle=False, num_workers=4)
+ ap, r_acc0, f_acc0, acc0, r_acc1, f_acc1, acc1, best_thres = validate(model, loader, find_thres=True)
+
+ with open( os.path.join(opt.result_folder,'ap.txt'), 'a') as f:
+ f.write(dataset_path['key']+': ' + str(round(ap*100, 2))+'\n' )
+
+ with open( os.path.join(opt.result_folder,'acc0.txt'), 'a') as f:
+ f.write(dataset_path['key']+': ' + str(round(r_acc0*100, 2))+' '+str(round(f_acc0*100, 2))+' '+str(round(acc0*100, 2))+'\n' )
+
: Content Privacy-Preserving Audio Deepfake Detection
+
+[](https://arxiv.org/abs/2409.09272)
+[](https://makeapullrequest.com)
+[](https://creativecommons.org/licenses/by/4.0/)
+
+
+
+
+
+By [1] Zhejiang University, [2] Tsinghua University.
+* [Xinfeng Li](https://letterligo.github.io)* [1], [Kai Li](https://cslikai.cn)* [2], Yifan Zheng [1], Chen Yan† [1], Xiaoyu Ji [1], Wenyuan Xu [1].
+
+This repository is an official implementation of the SafeEar accepted to **ACM CCS 2024** (Core-A*, CCF-A, Big4) .
+
+Please also visit our 
+ 
+ 
+ 
+
+
+## CO-SPY-Bench
+We have released our benchmark on [Huggingface](https://huggingface.co/datasets/ruojiruoli/Co-Spy-Bench), designed to offer diverse and comprehensive coverage of the latest generative models:
+- Captions are sourced from five real-world datasets: MSCOCO2017, CC3M, Flickr, TextCaps, and SBU.
+- Synthetic images are generated using 22 different models, covering a wide range of architectures.
+- Diverse generation parameters, such as diffusion steps and guidance scales, are used to enrich variability.
+
+## Main Code Architecture
+ .
+ ├── data # Dataset folder
+ │ ├── in_the_wild # CO-SPY-Bench in-the-wild synthetic samples
+ │ ├── test # Test dataset (CO-SPY-Bench & AIGCDetectionBenchMark)
+ │ └── train # Training dataset (DRCT-2M & CNNDet)
+ ├── dataSets # Various dataset classes
+ ├── detectors # Various detector classes
+ │ ├── progan # Detectors for CNNDet training set
+ │ └── sd-v1_4 # Detectors for DRCT-2M training set
+ ├── pretrained # Pre-trained weights
+ ├── main.py # Main function
+ ├── evaluate.py # Evaluation function
+ ├── main.py # Main function (Entry point)
+ ├── train.py # Training function
+ ├── train.sh # Recommended training script
+ └── utils.py # Utility functions
+
+## Environments
+```bash
+# Create python environment (optional)
+conda env create -f environment.yml
+source activate cospy
+```
+
+## Prerequisites
+Please download the required datasets and pre-trained weights for full evaluation.
+
+### Download Datasets
+Make sure you have `7z` an `unzip` installed. You can install them via conda:
+```bash
+conda install p7zip
+conda install unzip
+```
+
+To download the training and test datasets, run the following commands respectively:
+```bash
+###############################
+# Download the training dataset
+###############################
+# Download CNNDet (ProGAN) training set
+cd data/train/progan
+sh download.sh
+
+# Download DRCT-2M (Stable Diffusion v1.4) training set
+cd data/train/sd-v1_4
+sh download.sh
+
+###############################
+# Download the test dataset
+###############################
+# Download AIGCDetectionBenchMark test set
+cd data/test/AIGCDetectionBenchMark
+python download.py
+
+# Download CO-SPY-Bench test set
+cd data/test/Co-Spy-Bench
+sh download.sh
+```
+
+Finally, the directory structure should look like this:
+```
+.
+├── data
+│ ├── test
+│ │ ├── AIGCDetectionBenchMark
+│ │ │ └── test
+│ │ └── Co-Spy-Bench
+│ │ ├── real_image_examples
+│ │ └── synthetic
+│ └── train
+│ ├── progan
+│ │ ├── train
+│ │ └── val
+│ └── sd-v1_4
+│ ├── mscoco2017
+│ └── stable_diffusion_v1-4
+```
+*Note: Please ensure the use of these datasets complies with the original licenses.*
+
+
+> Please refer to `data/in_the_wild/README.md` for detailed instructions on accessing the **CO-SPY-Bench in-the-wild synthetic samples**.
+>
+> *Note: Some samples are temporarily unavailable (e.g., [instavibe.ai](https://www.instavibe.ai/)). As we cannot confirm whether redistributing these images would violate the original sources’ intellectual property rights, we choose not to release them at this time. This decision and its rationale are discussed in [Issue #6](https://github.com/Megum1/CO-SPY/issues/6). We apologize for any inconvenience this may cause.*
+
+
+### Download Pre-trained Weights
+```bash
+# Download the pre-trained weights
+cd pretrained
+sh download.sh
+# It contains the pre-trained weights on CNNDet (`progan`) and DRCT-2M (`sd-v1_4`) training sets.
+```
+
+
+
+## Experiments
+We provide the source code for training and evaluating the CO-SPY detector.
+
+### Evaluation of Pre-trained Detectors
+To evaluate the pre-trained detector (trained on DRCT-2M (`sd-v1_4`)) on [CO-SPY-Bench](https://huggingface.co/datasets/ruojiruoli/Co-Spy-Bench), run
+```bash
+python main.py --gpu 0 --phase eval --train_dataset sd-v1_4 --pretrain
+```
+When finished, evaluation results will be saved to `ckpt/sd-v1_4/fusion/pretrain_Co-Spy-Bench`.
+
+
+
+
+
+
+
+
+Reference github repository for the paper [Rethinking the Up-Sampling Operations in CNN-based Generative Network for Generalizable Deepfake Detection](https://arxiv.org/abs/2312.10461).
+```
+@misc{tan2023rethinking,
+ title={Rethinking the Up-Sampling Operations in CNN-based Generative Network for Generalizable Deepfake Detection},
+ author={Chuangchuang Tan and Huan Liu and Yao Zhao and Shikui Wei and Guanghua Gu and Ping Liu and Yunchao Wei},
+ year={2023},
+ eprint={2312.10461},
+ archivePrefix={arXiv},
+ primaryClass={cs.CV}
+}
+```
+
+## News 🆕
+- `2024/02`: NPR is accepted by CVPR 2024! Congratulations and thanks to my all co-authors!
+- `2024/05`: [🤗Online Demo](https://huggingface.co/spaces/tancc/Generalizable_Deepfake_Detection-NPR-CVPR2024)
+
+
+
+## Environment setup
+**Classification environment:**
+We recommend installing the required packages by running the command:
+```sh
+pip install -r requirements.txt
+```
+In order to ensure the reproducibility of the results, we provide the following suggestions：
+- Docker image: nvcr.io/nvidia/tensorflow:21.02-tf1-py3
+- Conda environment: [./pytorch18/bin/python](https://drive.google.com/file/d/16MK7KnPebBZx5yeN6jqJ49k7VWbEYQPr/view)
+- Random seed during testing period: [Random seed](https://github.com/chuangchuangtan/NPR-DeepfakeDetection/blob/b4e1bfa59ec58542ab5b1e78a3b75b54df67f3b8/test.py#L14)
+
+## Getting the data
+
+| | paper | Url |
+|:----------------------:|:-----:|:-----:|
+| Train set | [CNNDetection CVPR2020](https://github.com/PeterWang512/CNNDetection) | [Baidudrive](https://pan.baidu.com/s/1l-rXoVhoc8xJDl20Cdwy4Q?pwd=ft8b) |
+| Val set | [CNNDetection CVPR2020](https://github.com/PeterWang512/CNNDetection) | [Baidudrive](https://pan.baidu.com/s/1l-rXoVhoc8xJDl20Cdwy4Q?pwd=ft8b) |
+| Table1 Test | [CNNDetection CVPR2020](https://github.com/PeterWang512/CNNDetection) | [Baidudrive](https://pan.baidu.com/s/1l-rXoVhoc8xJDl20Cdwy4Q?pwd=ft8b) |
+| Table2 Test | [FreqNet AAAI2024](https://github.com/chuangchuangtan/FreqNet-DeepfakeDetection) | [googledrive](https://drive.google.com/drive/folders/11E0Knf9J1qlv2UuTnJSOFUjIIi90czSj?usp=sharing) |
+| Table3 Test | [DIRE ICCV2023](https://github.com/ZhendongWang6/DIRE) | [googledrive](https://drive.google.com/drive/folders/1jZE4hg6SxRvKaPYO_yyMeJN_DOcqGMEf?usp=sharing) |
+| Table4 Test | [UniversalFakeDetect CVPR2023](https://github.com/Yuheng-Li/UniversalFakeDetect) | [googledrive](https://drive.google.com/drive/folders/1nkCXClC7kFM01_fqmLrVNtnOYEFPtWO-?usp=sharing)|
+| Table5 Test | Diffusion1kStep | [googledrive](https://drive.google.com/drive/folders/14f0vApTLiukiPvIHukHDzLujrvJpDpRq?usp=sharing) |
+
+```
+pip install gdown==4.7.1
+
+chmod 777 ./download_dataset.sh
+
+./download_dataset.sh
+```
+## Directory structure
+
> 
+ 
+
+
+
+
+
+
+
+> **Abstract.** In recent years, DeepFake technology has achieved unprecedented success in high-quality video synthesis, but these methods also pose potential and severe security threats to humanity. DeepFake can be bifurcated into entertainment applications like face swapping and illicit uses such as lip-syncing fraud. However, lip-forgery videos, which neither change identity nor have discernible visual artifacts, present a formidable challenge to existing DeepFake detection methods. Our preliminary experiments have shown that the effectiveness of the existing methods often drastically decrease or even fail when tackling lip-syncing videos.
+> In this paper, for the first time, we propose a novel approach dedicated to lip-forgery identification that exploits the inconsistency between lip movements and audio signals. We also mimic human natural cognition by capturing subtle biological links between lips and head regions to boost accuracy. To better illustrate the effectiveness and advances of our proposed method, we create a high-quality LipSync dataset, AVLips, by employing the state-of-the-art lip generators. We hope this high-quality and diverse dataset could be well served the further research on this challenging and interesting field. Experimental results show that our approach gives an average accuracy of more than 95.3% in spotting lip-syncing videos, significantly outperforming the baselines. Extensive experiments demonstrate the capability to tackle deepfakes and the robustness in surviving diverse input transformations. Our method achieves an accuracy of up to 90.2% in real-world scenarios (e.g., WeChat video call) and shows its powerful capabilities in real scenario deployment.
+
+
+
+
+
+## 🔥 AVLips: A high-quality audio-visual dataset for LipSync detection
+
+To the best of our knowledge, the majority of public DeepFake datasets consist solely of videos or images, with no specialized one specifically dedicated to LipSync detection available. To fill this gap, we construct a high-quality **A**udio-**V**isual **Lip**-syncing Dataset, **AVLips**, which contains up to 340,000 audio-visual samples generated by several SOTA LipSync methods. The workflow is demonstrated below.
+
+**High quality.** We employed a combination of static MakeItTalk and dynamic Wav2Lip, TalkLip, SadTalker generation methods to simulate realistic lip movements. These methods are widely recognized as high-quality work, capable of generating high-resolution videos while ensuring accurate lip movements. We applied a noise reduction algorithm to all audio samples before synthesis to reduce irrelevant background noise, ensuring the models can focus on speech content.
+
+**Diversity.** Our dataset encompasses a wide range of scenarios, covering not only well-known public datasets but also real-world data. Our aim is for this collection to act as a catalyst for advancing real-time forgery detection. To better simulate the nuances of real-world conditions, we have employed six perturbation techniques — saturation, contrast, compression, Gaussian noise, Gaussian blur, and pixelation — at various degrees, thus ensuring the dataset's realism and practical relevance.
+
+**Download Link: [AVLips v1.0](https://drive.google.com/file/d/1fEiUo22GBSnWD7nfEwDW86Eiza-pOEJm/view?usp=share_link)**
+
+
+
+
+#### Update (Aug 25, 2020)
+
+This update adds:
+ * A new `include_top` (default: `True`) option ([#208](https://github.com/lukemelas/EfficientNet-PyTorch/pull/208))
+ * Continuous testing with [sotabench](https://sotabench.com/)
+ * Code quality improvements and fixes ([#215](https://github.com/lukemelas/EfficientNet-PyTorch/pull/215) [#223](https://github.com/lukemelas/EfficientNet-PyTorch/pull/223))
+
+#### Update (May 14, 2020)
+
+This update adds comprehensive comments and documentation (thanks to @workingcoder).
+
+#### Update (January 23, 2020)
+
+This update adds a new category of pre-trained model based on adversarial training, called _advprop_. It is important to note that the preprocessing required for the advprop pretrained models is slightly different from normal ImageNet preprocessing. As a result, by default, advprop models are not used. To load a model with advprop, use:
+```python
+model = EfficientNet.from_pretrained("efficientnet-b0", advprop=True)
+```
+There is also a new, large `efficientnet-b8` pretrained model that is only available in advprop form. When using these models, replace ImageNet preprocessing code as follows:
+```python
+if advprop: # for models using advprop pretrained weights
+ normalize = transforms.Lambda(lambda img: img * 2.0 - 1.0)
+else:
+ normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
+ std=[0.229, 0.224, 0.225])
+```
+This update also addresses multiple other issues ([#115](https://github.com/lukemelas/EfficientNet-PyTorch/issues/115), [#128](https://github.com/lukemelas/EfficientNet-PyTorch/issues/128)).
+
+#### Update (October 15, 2019)
+
+This update allows you to choose whether to use a memory-efficient Swish activation. The memory-efficient version is chosen by default, but it cannot be used when exporting using PyTorch JIT. For this purpose, we have also included a standard (export-friendly) swish activation function. To switch to the export-friendly version, simply call `model.set_swish(memory_efficient=False)` after loading your desired model. This update addresses issues [#88](https://github.com/lukemelas/EfficientNet-PyTorch/pull/88) and [#89](https://github.com/lukemelas/EfficientNet-PyTorch/pull/89).
+
+#### Update (October 12, 2019)
+
+This update makes the Swish activation function more memory-efficient. It also addresses pull requests [#72](https://github.com/lukemelas/EfficientNet-PyTorch/pull/72), [#73](https://github.com/lukemelas/EfficientNet-PyTorch/pull/73), [#85](https://github.com/lukemelas/EfficientNet-PyTorch/pull/85), and [#86](https://github.com/lukemelas/EfficientNet-PyTorch/pull/86). Thanks to the authors of all the pull requests!
+
+#### Update (July 31, 2019)
+
+_Upgrade the pip package with_ `pip install --upgrade efficientnet-pytorch`
+
+The B6 and B7 models are now available. Additionally, _all_ pretrained models have been updated to use AutoAugment preprocessing, which translates to better performance across the board. Usage is the same as before:
+```python
+from efficientnet_pytorch import EfficientNet
+model = EfficientNet.from_pretrained('efficientnet-b7')
+```
+
+#### Update (June 29, 2019)
+
+This update adds easy model exporting ([#20](https://github.com/lukemelas/EfficientNet-PyTorch/issues/20)) and feature extraction ([#38](https://github.com/lukemelas/EfficientNet-PyTorch/issues/38)).
+
+ * [Example: Export to ONNX](#example-export)
+ * [Example: Extract features](#example-feature-extraction)
+ * Also: fixed a CUDA/CPU bug ([#32](https://github.com/lukemelas/EfficientNet-PyTorch/issues/32))
+
+It is also now incredibly simple to load a pretrained model with a new number of classes for transfer learning:
+```python
+model = EfficientNet.from_pretrained('efficientnet-b1', num_classes=23)
+```
+
+
+#### Update (June 23, 2019)
+
+The B4 and B5 models are now available. Their usage is identical to the other models:
+```python
+from efficientnet_pytorch import EfficientNet
+model = EfficientNet.from_pretrained('efficientnet-b4')
+```
+
+### Overview
+This repository contains an op-for-op PyTorch reimplementation of [EfficientNet](https://arxiv.org/abs/1905.11946), along with pre-trained models and examples.
+
+The goal of this implementation is to be simple, highly extensible, and easy to integrate into your own projects. This implementation is a work in progress -- new features are currently being implemented.
+
+At the moment, you can easily:
+ * Load pretrained EfficientNet models
+ * Use EfficientNet models for classification or feature extraction
+ * Evaluate EfficientNet models on ImageNet or your own images
+
+_Upcoming features_: In the next few days, you will be able to:
+ * Train new models from scratch on ImageNet with a simple command
+ * Quickly finetune an EfficientNet on your own dataset
+ * Export EfficientNet models for production
+
+### Table of contents
+1. [About EfficientNet](#about-efficientnet)
+2. [About EfficientNet-PyTorch](#about-efficientnet-pytorch)
+3. [Installation](#installation)
+4. [Usage](#usage)
+ * [Load pretrained models](#loading-pretrained-models)
+ * [Example: Classify](#example-classification)
+ * [Example: Extract features](#example-feature-extraction)
+ * [Example: Export to ONNX](#example-export)
+6. [Contributing](#contributing)
+
+### About EfficientNet
+
+If you're new to EfficientNets, here is an explanation straight from the official TensorFlow implementation:
+
+EfficientNets are a family of image classification models, which achieve state-of-the-art accuracy, yet being an order-of-magnitude smaller and faster than previous models. We develop EfficientNets based on AutoML and Compound Scaling. In particular, we first use [AutoML Mobile framework](https://ai.googleblog.com/2018/08/mnasnet-towards-automating-design-of.html) to develop a mobile-size baseline network, named as EfficientNet-B0; Then, we use the compound scaling method to scale up this baseline to obtain EfficientNet-B1 to B7.
+
+
+ 
+ 
+ 
+ 
+ 
+ 
+ 
+ 
+ 
+ 
+