diff --git a/audio/aasist3/.gitignore b/audio/aasist3/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..d285b0dd1d6c257ec0dea379665e40a74f8e8759 --- /dev/null +++ b/audio/aasist3/.gitignore @@ -0,0 +1 @@ +weights/ diff --git a/audio/aasist3/Dockerfile b/audio/aasist3/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..2539f7771b2cc49301a55c41599384eb68e391f0 --- /dev/null +++ b/audio/aasist3/Dockerfile @@ -0,0 +1,35 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3 python3-pip \ + ffmpeg libsndfile1 \ + && rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +WORKDIR /app + +# Install PyTorch with CUDA 12.1 +RUN pip install --no-cache-dir \ + torch==2.5.1 torchaudio==2.5.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +RUN python -c "from transformers import Wav2Vec2Config; Wav2Vec2Config.from_pretrained('facebook/wav2vec2-large-xlsr-53', cache_dir='/app/w2v_cache')" + +COPY model/ /app/model/ +COPY api.py . +RUN mkdir -p /app/weights +COPY weights/ /app/weights/ + +EXPOSE 8005 + +RUN adduser --disabled-password --gecos '' appuser +USER appuser + +CMD ["python", "api.py"] diff --git a/audio/aasist3/api.py b/audio/aasist3/api.py new file mode 100644 index 0000000000000000000000000000000000000000..b1fc8a868dd8686a3eb1725eb76aa67e543c5e8e --- /dev/null +++ b/audio/aasist3/api.py @@ -0,0 +1,345 @@ +"""AASIST3 Audio Deepfake Detection API. + +Detects synthetic speech using the AASIST3 model architecture: +- Frontend: XLSR wav2vec 2.0 (HuggingFace Transformers) +- Backend: AASIST with KAN (Kolmogorov-Arnold Network) linear + layers and Graph Attention Networks + +Reference: https://github.com/AI4Bharat/AASIST3 +""" + +import base64 +import io +import logging +import os +import platform +import time +from typing import Optional + +import numpy as np +import soundfile as sf +import torch +import torch.nn.functional as F +import uvicorn +from fastapi import FastAPI, HTTPException +from pydantic import BaseModel, Field + +# Point transformers cache to pre-cached wav2vec2 config +# (must be set before importing model code) +os.environ["TRANSFORMERS_CACHE"] = "/app/w2v_cache" + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", +) +logger = logging.getLogger("aasist3_api") + +# Import model class +try: + from model import aasist3 as AASIST3Model +except ImportError as e: + logger.error(f"Failed to import AASIST3 model: {e}") + AASIST3Model = None + +# Constants +MODEL_NAME = "aasist3" +MODEL_ID = "aasist3_kan_mlaad" +WEIGHTS_DIR = "/app/weights" + + +def _get_device(): + """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if override: + pass # Invalid override, fall through to auto-detect + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + return torch.device("cpu") + + +DEVICE = _get_device() + +if DEVICE.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + +if DEVICE.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) +else: + logger.warning( + "Device: %s (no CUDA available -- check nvidia-container-toolkit)", + DEVICE, + ) + +SAMPLE_RATE = 16000 +TARGET_SAMPLES = 64600 # ~4.04 seconds at 16kHz + +# Global model instance +model = None + + +class AudioInput(BaseModel): + """Request schema for audio deepfake detection.""" + + audio_data: str = Field( + ..., description="Base64 encoded audio string (WAV/MP3/etc)" + ) + threshold: Optional[float] = Field( + 0.5, ge=0.0, le=1.0, description="Classification threshold" + ) + + +app = FastAPI( + title="AASIST3 Audio Deepfake Detection API", + description=( + "Service for detecting synthetic speech using the " + "AASIST3 model (HuggingFace wav2vec 2.0 + AASIST " + "with KAN layers)." + ), + version="1.0.0", +) + + +def load_model(): + """Load the AASIST3 model from pretrained weights. + + Returns: + The loaded model, or None if loading fails. + """ + global model + if model is not None: + return model + + logger.info(f"Loading AASIST3 model onto {DEVICE}...") + + if AASIST3Model is None: + logger.error("AASIST3 model class not available.") + return None + + weights_safetensors = os.path.join(WEIGHTS_DIR, "model.safetensors") + weights_config = os.path.join(WEIGHTS_DIR, "config.json") + + if not os.path.exists(weights_safetensors): + logger.error(f"Model weights not found at {weights_safetensors}") + return None + + if not os.path.exists(weights_config): + logger.error(f"Model config not found at {weights_config}") + return None + + try: + model = AASIST3Model.from_pretrained(WEIGHTS_DIR) + model.to(DEVICE) + # Set model to inference mode (disables dropout, batchnorm) + model.train(False) + + logger.info("AASIST3 model loaded successfully.") + return model + except Exception as e: + logger.exception(f"Failed to load AASIST3 model: {e}") + model = None + return None + + +@app.on_event("startup") +async def startup_event(): + """Load model on service startup.""" + load_model() + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and DEVICE.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +async def health(): + """Health check endpoint.""" + return { + "status": "healthy" if model is not None else "degraded", + "model": MODEL_NAME, + "model_id": MODEL_ID, + "device": str(DEVICE), + "weights_found": os.path.exists(os.path.join(WEIGHTS_DIR, "model.safetensors")), + **_gpu_health_info(), + } + + +def _load_audio_bytes(audio_bytes: bytes) -> tuple: + """Load audio from raw bytes using soundfile with torchaudio fallback. + + Args: + audio_bytes: Raw audio file bytes. + + Returns: + Tuple of (audio_numpy_array, sample_rate). + + Raises: + ValueError: If audio cannot be loaded by any backend. + """ + # Try soundfile first (handles WAV, FLAC natively) + sf_error = None + try: + audio, sr = sf.read(io.BytesIO(audio_bytes), dtype="float32") + if audio.ndim > 1: + audio = audio.mean(axis=1) # Convert to mono + return audio, sr + except Exception as sf_err: + sf_error = sf_err + logger.debug(f"soundfile failed, trying torchaudio: {sf_err}") + + # Fallback to torchaudio (handles MP3, compressed formats) + try: + import torchaudio + + buf = io.BytesIO(audio_bytes) + waveform, sr = torchaudio.load(buf) + if waveform.shape[0] > 1: + waveform = waveform.mean(dim=0, keepdim=True) + return waveform.squeeze(0).numpy(), sr + except Exception as ta_err: + raise ValueError( + f"Failed to load audio with soundfile and torchaudio: " + f"sf={sf_error}, ta={ta_err}" + ) + + +def preprocess_audio(audio_bytes: bytes) -> torch.Tensor: + """Preprocess audio for AASIST3 inference. + + Loads audio, resamples to 16kHz mono, and zero-pads or + truncates to TARGET_SAMPLES. + + Args: + audio_bytes: Raw audio file bytes. + + Returns: + Audio tensor of shape (1, TARGET_SAMPLES). + + Raises: + ValueError: If audio preprocessing fails. + """ + try: + logger.info("Starting audio preprocessing...") + audio, sr = _load_audio_bytes(audio_bytes) + logger.info(f"Audio loaded. Length: {len(audio)} samples at {sr}Hz") + + # Resample to 16kHz if needed + if sr != SAMPLE_RATE: + import torchaudio + + resampler = torchaudio.transforms.Resample( + orig_freq=sr, new_freq=SAMPLE_RATE + ) + audio_tensor = torch.FloatTensor(audio).unsqueeze(0) + audio_tensor = resampler(audio_tensor).squeeze(0) + audio = audio_tensor.numpy() + logger.info( + f"Resampled from {sr}Hz to {SAMPLE_RATE}Hz. " + f"New length: {len(audio)} samples" + ) + + # Zero-pad or truncate to TARGET_SAMPLES + if len(audio) >= TARGET_SAMPLES: + audio = audio[:TARGET_SAMPLES] + else: + pad_length = TARGET_SAMPLES - len(audio) + audio = np.pad(audio, (0, pad_length), mode="constant") + + logger.info(f"Audio padded/trimmed to {TARGET_SAMPLES} samples") + + audio_tensor = torch.FloatTensor(audio).unsqueeze(0).to(DEVICE) + return audio_tensor + except Exception as e: + logger.error(f"Error preprocessing audio: {e}") + raise ValueError(f"Audio preprocessing failed: {str(e)}") + + +@app.post("/predict") +async def predict(input_data: AudioInput): + """Run deepfake detection on base64-encoded audio. + + The model outputs 2 logits: [bonafide_score, spoof_score]. + Class 0 = bonafide (real), Class 1 = spoof (fake). + The returned probability is the spoof/fake probability. + """ + if model is None: + if load_model() is None: + raise HTTPException(status_code=503, detail="Model not loaded") + + try: + start_time = time.time() + logger.info( + f"Prediction request. Data size: " f"{len(input_data.audio_data)} chars" + ) + + # Decode base64 audio + audio_bytes = base64.b64decode(input_data.audio_data) + + # Preprocess + audio_tensor = preprocess_audio(audio_bytes) + + # Inference + logger.info("Starting model inference...") + with torch.no_grad(): + output = model(audio_tensor) + + # output shape: [batch, 2] + # Index 0 = bonafide logit, Index 1 = spoof logit + probs = torch.softmax(output, dim=1) + prob_fake = probs[0, 1].item() + + prediction = 1 if prob_fake >= input_data.threshold else 0 + verdict = "fake" if prediction == 1 else "real" + inference_time = time.time() - start_time + + logger.info( + f"Prediction: {verdict} (prob_fake={prob_fake:.4f}, " + f"time={inference_time:.3f}s)" + ) + + return { + "model": MODEL_NAME, + "probability": float(prob_fake), + "prediction": int(prediction), + "class": verdict, + "inference_time": float(inference_time), + } + + except Exception as e: + logger.exception(f"Error during prediction: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +if __name__ == "__main__": + port = int(os.environ.get("MODEL_PORT", 8005)) + uvicorn.run(app, host="0.0.0.0", port=port) diff --git a/audio/aasist3/model/__init__.py b/audio/aasist3/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..71783b17169a638280c0985f13af7df6b89a4c05 --- /dev/null +++ b/audio/aasist3/model/__init__.py @@ -0,0 +1 @@ +from .full_model import aasist3 \ No newline at end of file diff --git a/audio/aasist3/model/branch.py b/audio/aasist3/model/branch.py new file mode 100644 index 0000000000000000000000000000000000000000..e35e26f0367cd3cbca62102e6f8ab52532268701 --- /dev/null +++ b/audio/aasist3/model/branch.py @@ -0,0 +1,34 @@ +import torch.nn as nn + +from .hs_gal import HtrgGraphAttentionLayer +from .pool import GraphPool + +class InferenceBranch(nn.Module): + def __init__(self, gat_dims, temperature, pool_ratio, size): + super().__init__() + self.htrg_gat1 = HtrgGraphAttentionLayer( + gat_dims[0], gat_dims[1], temperature=temperature, size=size + ) + self.htrg_gat2 = HtrgGraphAttentionLayer( + gat_dims[1], gat_dims[1], temperature=temperature, size=size + ) + + self.pool_hS = GraphPool(pool_ratio, gat_dims[1], 0.3, size=size) + self.pool_hT = GraphPool(pool_ratio, gat_dims[1], 0.3, size=size) + + def forward(self, out_T, out_S, master): + # Первая стадия + out_T_res, out_S_res, master_res = self.htrg_gat1(out_T, out_S, master=master) + + # Пулинг + out_S_res = self.pool_hS(out_S_res) + out_T_res = self.pool_hT(out_T_res) + + # Вторая стадия с residual connection + out_T_aug, out_S_aug, master_aug = self.htrg_gat2(out_T_res, out_S_res, master=master_res) + + out_T_final = out_T_res + out_T_aug + out_S_final = out_S_res + out_S_aug + master_final = master_res + master_aug + + return out_T_final, out_S_final, master_final \ No newline at end of file diff --git a/audio/aasist3/model/full_model.py b/audio/aasist3/model/full_model.py new file mode 100644 index 0000000000000000000000000000000000000000..dddccd136c9e5576157067bd9f2666af63149684 --- /dev/null +++ b/audio/aasist3/model/full_model.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python + +import torch +import torch.nn.functional as F +import torch.nn as nn +from huggingface_hub import PyTorchModelHubMixin + +from .kan import KANLinear +from .gat import GraphAttentionLayer +from .pool import GraphPool +from .branch import InferenceBranch +from .residual import Residual_block +from .wav2vec import Wav2Vec2Encoder + + +class aasist3(nn.Module, PyTorchModelHubMixin): + def __init__(self, d_args={ + "architecture": "AASIST", + "nb_samp": 64600, + "first_conv": 128, + "filts": [70, [1, 32], [32, 32], [32, 64], [64, 64]], + "gat_dims": [64, 32], + "pool_ratios": [0.5, 0.7, 0.5, 0.5], + "temperatures": [2.0, 2.0, 100.0, 100.0], + }, size=200, w2v_cache_dir="weights/", load_pretrained=True): + super().__init__() + + self.w2v_encoder = Wav2Vec2Encoder(cache_dir=w2v_cache_dir, load_pretrained=load_pretrained) + self.bridge = KANLinear(1024, 128) + + self.d_args = d_args + filts = d_args["filts"] + gat_dims = d_args["gat_dims"] + pool_ratios = d_args["pool_ratios"] + temperatures = d_args["temperatures"] + + self.first_bn = nn.BatchNorm2d(num_features=1) + self.selu = nn.SELU(inplace=True) + self.drop = nn.Dropout(0.5, inplace=True) + self.drop_way = nn.Dropout(0.2, inplace=True) + + self.encoder = nn.Sequential( + nn.Sequential(Residual_block(nb_filts=filts[1], first=True)), + nn.Sequential(Residual_block(nb_filts=filts[2])), + nn.Sequential(Residual_block(nb_filts=filts[3])), + nn.Sequential(Residual_block(nb_filts=filts[4])), + nn.Sequential(Residual_block(nb_filts=filts[4])), + nn.Sequential(Residual_block(nb_filts=filts[4]))) + + self.pos_S = nn.Parameter(torch.randn(1, 42, filts[-1][-1])) + self.pos_T = nn.Parameter(torch.randn(1, 67, filts[-1][-1])) + + self.GAT_layer_S = GraphAttentionLayer(filts[-1][-1], gat_dims[0], temperature=temperatures[0], size=size) + self.GAT_layer_T = GraphAttentionLayer(filts[-1][-1], gat_dims[0], temperature=temperatures[1], size=size) + + self.pool_S = GraphPool(pool_ratios[0], gat_dims[0], 0.3, size=size) + self.pool_T = GraphPool(pool_ratios[1], gat_dims[0], 0.3, size=size) + + self.master1 = nn.Parameter(torch.randn(1, 1, gat_dims[0])) + self.master2 = nn.Parameter(torch.randn(1, 1, gat_dims[0])) + self.master3 = nn.Parameter(torch.randn(1, 1, gat_dims[0])) + self.master4 = nn.Parameter(torch.randn(1, 1, gat_dims[0])) + + self.inference_branch1 = InferenceBranch( + gat_dims=gat_dims, + temperature=temperatures[2], + pool_ratio=pool_ratios[2], + size=size + ) + self.inference_branch2 = InferenceBranch( + gat_dims=gat_dims, + temperature=temperatures[2], + pool_ratio=pool_ratios[2], + size=size + ) + self.inference_branch3 = InferenceBranch( + gat_dims=gat_dims, + temperature=temperatures[2], + pool_ratio=pool_ratios[2], + size=size + ) + self.inference_branch4 = InferenceBranch( + gat_dims=gat_dims, + temperature=temperatures[2], + pool_ratio=pool_ratios[2], + size=size + ) + + self.out_layer = KANLinear(5 * gat_dims[1], 2) + + def forward(self, x, Freq_aug=False): + x = self.w2v_encoder(x) + x = self.bridge(x) + x = x.transpose(1, 2) + x = x.unsqueeze(dim=1) + x = F.max_pool2d(torch.abs(x), (3, 3)) + x = self.first_bn(x) + x = self.selu(x) + + e = self.encoder(x) + + # GAT-S + e_S, _ = torch.max(torch.abs(e), dim=3) + e_S = e_S.transpose(1, 2) + self.pos_S + gat_S = self.GAT_layer_S(e_S) + out_S = self.pool_S(gat_S) + + # GAT-T + e_T, _ = torch.max(torch.abs(e), dim=2) + e_T = e_T.transpose(1, 2) + self.pos_T + gat_T = self.GAT_layer_T(e_T) + out_T = self.pool_T(gat_T) + + out_T1, out_S1, master1 = self.inference_branch1(out_T, out_S, self.master1) + out_T2, out_S2, master2 = self.inference_branch2(out_T, out_S, self.master2) + out_T3, out_S3, master3 = self.inference_branch3(out_T, out_S, self.master3) + out_T4, out_S4, master4 = self.inference_branch4(out_T, out_S, self.master4) + + out_T1, out_T2 = self.drop_way(out_T1), self.drop_way(out_T2) + out_T3, out_T4 = self.drop_way(out_T3), self.drop_way(out_T4) + out_S1, out_S2 = self.drop_way(out_S1), self.drop_way(out_S2) + out_S3, out_S4 = self.drop_way(out_S3), self.drop_way(out_S4) + master1, master2 = self.drop_way(master1), self.drop_way(master2) + master3, master4 = self.drop_way(master3), self.drop_way(master4) + + out_T = torch.stack([out_T1, out_T2, out_T3, out_T4]).max(dim=0)[0] + out_S = torch.stack([out_S1, out_S2, out_S3, out_S4]).max(dim=0)[0] + master = torch.stack([master1, master2, master3, master4]).max(dim=0)[0] + + T_max, _ = torch.max(torch.abs(out_T), dim=1) + T_avg = torch.mean(out_T, dim=1) + S_max, _ = torch.max(torch.abs(out_S), dim=1) + S_avg = torch.mean(out_S, dim=1) + + last_hidden = torch.cat([T_max, T_avg, S_max, S_avg, master.squeeze(1)], dim=1) + last_hidden = self.drop(last_hidden) + output = self.out_layer(last_hidden) + + return output diff --git a/audio/aasist3/model/gat.py b/audio/aasist3/model/gat.py new file mode 100644 index 0000000000000000000000000000000000000000..9acbceffffaeebbaaee5d68607b7dd3e12fe46ff --- /dev/null +++ b/audio/aasist3/model/gat.py @@ -0,0 +1,99 @@ +import torch, torch.nn as nn, torch.nn.functional as F + +from .kan import KANLinear + + +class GraphAttentionLayer(nn.Module): + def __init__(self, in_dim, out_dim, size, **kwargs): + super().__init__() + + # attention map + self.att_proj = KANLinear(in_dim, out_dim) + self.att_weight = self._init_new_params(out_dim, 1) + + # project + self.proj_with_att = KANLinear(in_dim, out_dim) + self.proj_without_att = KANLinear(in_dim, out_dim) + + # batch norm + self.bn = nn.BatchNorm1d(out_dim) + + # dropout for inputs + self.input_drop = nn.Dropout(p=0.2) + + # activate + self.act = nn.SELU(inplace=True) + + # temperature + self.temp = 1. + if "temperature" in kwargs: + self.temp = kwargs["temperature"] + + def forward(self, x): + ''' + x :(#bs, #node, #dim) + ''' + # apply input dropout + x = self.input_drop(x) + + # derive attention map + att_map = self._derive_att_map(x) + + # projection + x = self._project(x, att_map) + + # apply batch norm + x = self._apply_BN(x) + x = self.act(x) + return x + + def _pairwise_mul_nodes(self, x): + ''' + Calculates pairwise multiplication of nodes. + - for attention map + x :(#bs, #node, #dim) + out_shape :(#bs, #node, #node, #dim) + ''' + + nb_nodes = x.size(1) + x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1) + x_mirror = x.transpose(1, 2) + + return x * x_mirror + + def _derive_att_map(self, x): + ''' + x :(#bs, #node, #dim) + out_shape :(#bs, #node, #node, 1) + ''' + att_map = self._pairwise_mul_nodes(x) + # size: (#bs, #node, #node, #dim_out) + att_map = torch.tanh(self.att_proj(att_map)) + # size: (#bs, #node, #node, 1) + att_map = torch.matmul(att_map, self.att_weight) + + # apply temperature + att_map = att_map / self.temp + + att_map = F.softmax(att_map, dim=-2) + + return att_map + + def _project(self, x, att_map): + x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x)) + x2 = self.proj_without_att(x) + + return x1 + x2 + + def _apply_BN(self, x): + org_size = x.size() + x = x.view(-1, org_size[-1]) + x = self.bn(x) + x = x.view(org_size) + + return x + + def _init_new_params(self, *size): + out = nn.Parameter(torch.FloatTensor(*size)) + nn.init.xavier_normal_(out) + return out diff --git a/audio/aasist3/model/hs_gal.py b/audio/aasist3/model/hs_gal.py new file mode 100644 index 0000000000000000000000000000000000000000..99edcae0769edf34ca835b0128355002b8ba1ad5 --- /dev/null +++ b/audio/aasist3/model/hs_gal.py @@ -0,0 +1,176 @@ +import torch, torch.nn as nn, torch.nn.functional as F + +from .kan import KANLinear + + +class HtrgGraphAttentionLayer(nn.Module): + def __init__(self, in_dim, out_dim, size, **kwargs): + super().__init__() + + self.proj_type1 = KANLinear(in_dim, in_dim) + self.proj_type2 = KANLinear(in_dim, in_dim) + + # attention map + self.att_proj = KANLinear(in_dim, out_dim) + self.att_projM = KANLinear(in_dim, out_dim) + + self.att_weight11 = self._init_new_params(out_dim, 1) + self.att_weight22 = self._init_new_params(out_dim, 1) + self.att_weight12 = self._init_new_params(out_dim, 1) + self.att_weightM = self._init_new_params(out_dim, 1) + + # project + self.proj_with_att = KANLinear(in_dim, out_dim) + self.proj_without_att = KANLinear(in_dim, out_dim) + + self.proj_with_attM = KANLinear(in_dim, out_dim) + self.proj_without_attM = KANLinear(in_dim, out_dim) + + # batch norm + self.bn = nn.BatchNorm1d(out_dim) + + # dropout for inputs + self.input_drop = nn.Dropout(p=0.2) + + # activate + self.act = nn.SELU(inplace=True) + + # temperature + self.temp = 1. + if "temperature" in kwargs: + self.temp = kwargs["temperature"] + + def forward(self, x1, x2, master=None): + ''' + x1 :(#bs, #node, #dim) + x2 :(#bs, #node, #dim) + ''' + num_type1 = x1.size(1) + num_type2 = x2.size(1) + + x1 = self.proj_type1(x1) + x2 = self.proj_type2(x2) + + x = torch.cat([x1, x2], dim=1) + + if master is None: + master = torch.mean(x, dim=1, keepdim=True) + + # apply input dropout + x = self.input_drop(x) + + # derive attention map + att_map = self._derive_att_map(x, num_type1, num_type2) + + # directional edge for master node + master = self._update_master(x, master) + + # projection + x = self._project(x, att_map) + + # apply batch norm + x = self._apply_BN(x) + # x = self.act(x) + + x1 = x.narrow(1, 0, num_type1) + x2 = x.narrow(1, num_type1, num_type2) + + return x1, x2, master + + def _update_master(self, x, master): + + att_map = self._derive_att_map_master(x, master) + master = self._project_master(x, master, att_map) + + return master + + def _pairwise_mul_nodes(self, x): + ''' + Calculates pairwise multiplication of nodes. + - for attention map + x :(#bs, #node, #dim) + out_shape :(#bs, #node, #node, #dim) + ''' + + nb_nodes = x.size(1) + x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1) + x_mirror = x.transpose(1, 2) + + return x * x_mirror + + def _derive_att_map_master(self, x, master): + ''' + x :(#bs, #node, #dim) + out_shape :(#bs, #node, #node, 1) + ''' + att_map = x * master + att_map = torch.tanh(self.att_projM(att_map)) + + att_map = torch.matmul(att_map, self.att_weightM) + + # apply temperature + att_map = att_map / self.temp + + att_map = F.softmax(att_map, dim=-2) + + return att_map + + def _derive_att_map(self, x, num_type1, num_type2): + ''' + x :(#bs, #node, #dim) + out_shape :(#bs, #node, #node, 1) + ''' + att_map = self._pairwise_mul_nodes(x) + # size: (#bs, #node, #node, #dim_out) + att_map = torch.tanh(self.att_proj(att_map)) + # size: (#bs, #node, #node, 1) + + att_board = torch.zeros_like(att_map[:, :, :, 0]).unsqueeze(-1) + + att_board[:, :num_type1, :num_type1, :] = torch.matmul( + att_map[:, :num_type1, :num_type1, :], self.att_weight11) + att_board[:, num_type1:, num_type1:, :] = torch.matmul( + att_map[:, num_type1:, num_type1:, :], self.att_weight22) + att_board[:, :num_type1, num_type1:, :] = torch.matmul( + att_map[:, :num_type1, num_type1:, :], self.att_weight12) + att_board[:, num_type1:, :num_type1, :] = torch.matmul( + att_map[:, num_type1:, :num_type1, :], self.att_weight12) + + att_map = att_board + + # att_map = torch.matmul(att_map, self.att_weight12) + + # apply temperature + att_map = att_map / self.temp + + att_map = F.softmax(att_map, dim=-2) + + return att_map + + def _project(self, x, att_map): + x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x)) + x2 = self.proj_without_att(x) + + return x1 + x2 + + def _project_master(self, x, master, att_map): + + x1 = self.proj_with_attM(torch.matmul( + att_map.squeeze(-1).unsqueeze(1), x)) + x2 = self.proj_without_attM(master) + + return x1 + x2 + + def _apply_BN(self, x): + org_size = x.size() + x = x.view(-1, org_size[-1]) + x = self.bn(x) + x = x.view(org_size) + + return x + + def _init_new_params(self, *size): + out = nn.Parameter(torch.FloatTensor(*size)) + nn.init.xavier_normal_(out) + return out + diff --git a/audio/aasist3/model/kan.py b/audio/aasist3/model/kan.py new file mode 100644 index 0000000000000000000000000000000000000000..9ceb844778c85c7abc5b3b0eb71ed5c346457917 --- /dev/null +++ b/audio/aasist3/model/kan.py @@ -0,0 +1,213 @@ +import torch, math, torch.nn.functional as F + + +class KANLinear(torch.nn.Module): + def __init__( + self, + in_features, + out_features, + grid_size=16, + spline_order=4, + scale_noise=0.1, + scale_base=1.0, + scale_spline=1.0, + enable_standalone_scale_spline=True, + base_activation=torch.nn.PReLU, + grid_eps=0.02, + grid_range=[-1, 1], + ): + super(KANLinear, self).__init__() + self.in_features = in_features + self.out_features = out_features + self.grid_size = grid_size + self.spline_order = spline_order + + h = (grid_range[1] - grid_range[0]) / grid_size + grid = ( + ( + torch.arange(-spline_order, grid_size + spline_order + 1) * h + + grid_range[0] + ) + .expand(in_features, -1) + .contiguous() + ) + self.register_buffer("grid", grid) + + self.base_weight = torch.nn.Parameter(torch.Tensor(out_features, in_features)) + self.spline_weight = torch.nn.Parameter( + torch.Tensor(out_features, in_features, grid_size + spline_order) + ) + if enable_standalone_scale_spline: + self.spline_scaler = torch.nn.Parameter( + torch.Tensor(out_features, in_features) + ) + + self.scale_noise = scale_noise + self.scale_base = scale_base + self.scale_spline = scale_spline + self.enable_standalone_scale_spline = enable_standalone_scale_spline + self.base_activation = base_activation() + self.grid_eps = grid_eps + + self.reset_parameters() + + def reset_parameters(self): + torch.nn.init.kaiming_uniform_(self.base_weight, a=math.sqrt(5) * self.scale_base) + with torch.no_grad(): + noise = ( + ( + torch.rand(self.grid_size + 1, self.in_features, self.out_features) + - 1 / 2 + ) + * self.scale_noise + / self.grid_size + ) + self.spline_weight.data.copy_( + (self.scale_spline if not self.enable_standalone_scale_spline else 1.0) + * self.curve2coeff( + self.grid.T[self.spline_order : -self.spline_order], + noise, + ) + ) + if self.enable_standalone_scale_spline: + # torch.nn.init.constant_(self.spline_scaler, self.scale_spline) + torch.nn.init.kaiming_uniform_(self.spline_scaler, a=math.sqrt(5) * self.scale_spline) + + def b_splines(self, x: torch.Tensor): + """ + Compute the B-spline bases for the given input tensor. + + Args: + x (torch.Tensor): Input tensor of shape (batch_size, in_features). + + Returns: + torch.Tensor: B-spline bases tensor of shape (batch_size, in_features, grid_size + spline_order). + """ + assert x.dim() == 2 and x.size(1) == self.in_features + + grid: torch.Tensor = ( + self.grid + ) # (in_features, grid_size + 2 * spline_order + 1) + x = x.unsqueeze(-1) + bases = ((x >= grid[:, :-1]) & (x < grid[:, 1:])).to(x.dtype) + for k in range(1, self.spline_order + 1): + bases = ( + (x - grid[:, : -(k + 1)]) + / (grid[:, k:-1] - grid[:, : -(k + 1)]) + * bases[:, :, :-1] + ) + ( + (grid[:, k + 1 :] - x) + / (grid[:, k + 1 :] - grid[:, 1:(-k)]) + * bases[:, :, 1:] + ) + + assert bases.size() == ( + x.size(0), + self.in_features, + self.grid_size + self.spline_order, + ) + return bases.contiguous() + + def curve2coeff(self, x: torch.Tensor, y: torch.Tensor): + """ + Compute the coefficients of the curve that interpolates the given points. + + Args: + x (torch.Tensor): Input tensor of shape (batch_size, in_features). + y (torch.Tensor): Output tensor of shape (batch_size, in_features, out_features). + + Returns: + torch.Tensor: Coefficients tensor of shape (out_features, in_features, grid_size + spline_order). + """ + assert x.dim() == 2 and x.size(1) == self.in_features + assert y.size() == (x.size(0), self.in_features, self.out_features) + + A = self.b_splines(x).transpose( + 0, 1 + ) # (in_features, batch_size, grid_size + spline_order) + B = y.transpose(0, 1) # (in_features, batch_size, out_features) + solution = torch.linalg.lstsq( + A, B + ).solution # (in_features, grid_size + spline_order, out_features) + result = solution.permute( + 2, 0, 1 + ) # (out_features, in_features, grid_size + spline_order) + + assert result.size() == ( + self.out_features, + self.in_features, + self.grid_size + self.spline_order, + ) + return result.contiguous() + + @property + def scaled_spline_weight(self): + return self.spline_weight * ( + self.spline_scaler.unsqueeze(-1) + if self.enable_standalone_scale_spline + else 1.0 + ) + + def forward(self, x: torch.Tensor): + assert x.size(-1) == self.in_features + original_shape = x.shape + x = x.reshape(-1, self.in_features) + + base_output = F.linear(self.base_activation(x), self.base_weight) + spline_output = F.linear( + self.b_splines(x).view(x.size(0), -1), + self.scaled_spline_weight.reshape(self.out_features, -1), + ) + output = base_output + spline_output + # print(*original_shape[:-1], output.shape) + output = output.view(*original_shape[:-1], self.out_features) + return output + + @torch.no_grad() + def update_grid(self, x: torch.Tensor, margin=0.01): + assert x.dim() == 2 and x.size(1) == self.in_features + batch = x.size(0) + + splines = self.b_splines(x) # (batch, in, coeff) + splines = splines.permute(1, 0, 2) # (in, batch, coeff) + orig_coeff = self.scaled_spline_weight # (out, in, coeff) + orig_coeff = orig_coeff.permute(1, 2, 0) # (in, coeff, out) + unreduced_spline_output = torch.bmm(splines, orig_coeff) # (in, batch, out) + unreduced_spline_output = unreduced_spline_output.permute( + 1, 0, 2 + ) # (batch, in, out) + + # sort each channel individually to collect data distribution + x_sorted = torch.sort(x, dim=0)[0] + grid_adaptive = x_sorted[ + torch.linspace( + 0, batch - 1, self.grid_size + 1, dtype=torch.int64, device=x.device + ) + ] + + uniform_step = (x_sorted[-1] - x_sorted[0] + 2 * margin) / self.grid_size + grid_uniform = ( + torch.arange( + self.grid_size + 1, dtype=torch.float32, device=x.device + ).unsqueeze(1) + * uniform_step + + x_sorted[0] + - margin + ) + + grid = self.grid_eps * grid_uniform + (1 - self.grid_eps) * grid_adaptive + grid = torch.concatenate( + [ + grid[:1] + - uniform_step + * torch.arange(self.spline_order, 0, -1, device=x.device).unsqueeze(1), + grid, + grid[-1:] + + uniform_step + * torch.arange(1, self.spline_order + 1, device=x.device).unsqueeze(1), + ], + dim=0, + ) + + self.grid.copy_(grid.T) + self.spline_weight.data.copy_(self.curve2coeff(x, unreduced_spline_output)) diff --git a/audio/aasist3/model/pool.py b/audio/aasist3/model/pool.py new file mode 100644 index 0000000000000000000000000000000000000000..b9d5c3b79e3c4b67d106bb62c4e7bb6e51c1fa3f --- /dev/null +++ b/audio/aasist3/model/pool.py @@ -0,0 +1,45 @@ +import torch, torch.nn as nn +from typing import Union + +from .kan import KANLinear + + +class GraphPool(nn.Module): + def __init__(self, k: float, in_dim: int, p: Union[float, int], size): + super().__init__() + self.k = k + self.sigmoid = nn.Sigmoid() + self.proj = KANLinear(in_dim, 1) + self.drop = nn.Dropout(p=p) if p > 0 else nn.Identity() + self.in_dim = in_dim + + def forward(self, h): + Z = self.drop(h) + weights = self.proj(Z) + scores = self.sigmoid(weights) + new_h = self.top_k_graph(scores, h, self.k) + + return new_h + + def top_k_graph(self, scores, h, k): + """ + args + ===== + scores: attention-based weights (#bs, #node, 1) + h: graph data (#bs, #node, #dim) + k: ratio of remaining nodes, (float) + + returns + ===== + h: graph pool applied data (#bs, #node', #dim) + """ + _, n_nodes, n_feat = h.size() + n_nodes = max(int(n_nodes * k), 1) + _, idx = torch.topk(scores, n_nodes, dim=1) + idx = idx.expand(-1, -1, n_feat) + + h = h * scores + h = torch.gather(h, 1, idx) + + return h + diff --git a/audio/aasist3/model/residual.py b/audio/aasist3/model/residual.py new file mode 100644 index 0000000000000000000000000000000000000000..6fd300a30e8eaca4d5c59135b3c1f2990c7cfcd7 --- /dev/null +++ b/audio/aasist3/model/residual.py @@ -0,0 +1,56 @@ +import torch.nn as nn + +class Residual_block(nn.Module): + def __init__(self, nb_filts, first=False): + super().__init__() + self.first = first + + if not self.first: + self.bn1 = nn.BatchNorm2d(num_features=nb_filts[0]) + self.conv1 = nn.Conv2d(in_channels=nb_filts[0], + out_channels=nb_filts[1], + kernel_size=(2, 3), + padding=(1, 1), + stride=1) + self.selu = nn.SELU() + + self.bn2 = nn.BatchNorm2d(num_features=nb_filts[1]) + self.conv2 = nn.Conv2d(in_channels=nb_filts[1], + out_channels=nb_filts[1], + kernel_size=(2, 3), + padding=(0, 1), + stride=1) + + if nb_filts[0] != nb_filts[1]: + self.downsample = True + self.conv_downsample = nn.Conv2d(in_channels=nb_filts[0], + out_channels=nb_filts[1], + padding=(0, 1), + kernel_size=(1, 3), + stride=1) + + else: + self.downsample = False + # self.mp = nn.MaxPool2d((1, 3)) # self.mp = nn.MaxPool2d((1,4)) + + def forward(self, x): + identity = x + if not self.first: + out = self.bn1(x) + out = self.selu(out) + else: + out = x + out = self.conv1(x) + + # print('out',out.shape) + out = self.bn2(out) + out = self.selu(out) + # print('out',out.shape) + out = self.conv2(out) + #print('conv2 out',out.shape) + if self.downsample: + identity = self.conv_downsample(identity) + + out += identity + # out = self.mp(out) + return out \ No newline at end of file diff --git a/audio/aasist3/model/wav2vec.py b/audio/aasist3/model/wav2vec.py new file mode 100644 index 0000000000000000000000000000000000000000..72a787e3a0a235709f1cbf599dae5954df4b6621 --- /dev/null +++ b/audio/aasist3/model/wav2vec.py @@ -0,0 +1,82 @@ +import torch, torch.nn as nn +from transformers import Wav2Vec2Model, Wav2Vec2Config + + +class Wav2Vec2Encoder(nn.Module): + """SSL encoder based on Hugging Face's Wav2Vec2 model.""" + + def __init__(self, + model_name_or_path: str = "facebook/wav2vec2-large-xlsr-53", + ssl_out_dim: int = 768, + use_ssl_n_layers: int = None, + freeze_ssl_n_layers: int = 0, + output_attentions: bool = False, + output_hidden_states: bool = False, + normalize_waveform: bool = True, + cache_dir: str = "weights", + load_pretrained: bool = True): + """Initialize the Wav2Vec2 encoder. + + Args: + model_name_or_path: HuggingFace model name or path to local model. + ssl_out_dim: Output dimension of the Wav2Vec2 encoder. + use_ssl_n_layers: Number of Wav2Vec2 layers to use. If None, use all layers. + freeze_ssl_n_layers: Number of Wav2Vec2 layers to freeze during training. + output_attentions: Whether to output attentions. + output_hidden_states: Whether to output hidden states. + normalize_waveform: Whether to normalize the waveform input. + cache_dir: Directory to cache pretrained models. + load_pretrained: Whether to load pretrained weights. If False, initializes with random weights. + """ + super().__init__() + + self.model_name_or_path = model_name_or_path + self.ssl_out_dim = ssl_out_dim + self.use_ssl_n_layers = use_ssl_n_layers + self.freeze_ssl_n_layers = freeze_ssl_n_layers + self.output_attentions = output_attentions + self.output_hidden_states = output_hidden_states + self.normalize_waveform = normalize_waveform + + if load_pretrained: + self.model = Wav2Vec2Model.from_pretrained(model_name_or_path, cache_dir=cache_dir) + else: + config = Wav2Vec2Config.from_pretrained( + model_name_or_path, + cache_dir=cache_dir, + local_files_only=False + ) + self.model = Wav2Vec2Model(config) + self.model.init_weights() + + def forward(self, x): + """Forward pass through the Wav2Vec2 encoder. + + Args: + x: Input tensor of shape (batch_size, sequence_length, channels) + + Returns: + Extracted features of shape (batch_size, sequence_length, ssl_out_dim) + """ + # Handle shape: convert (batch_size, sequence_length, channels) to (batch_size, sequence_length) + if x.ndim == 3: + x = x.squeeze(-1) # Remove channel dimension if present + + if self.normalize_waveform: + x = x / (torch.max(torch.abs(x), dim=1, keepdim=True)[0] + 1e-8) + + outputs = self.model( + x, + output_attentions=self.output_attentions, + output_hidden_states=self.output_hidden_states, + return_dict=True + ) + + last_hidden_state = outputs.last_hidden_state + + if self.use_ssl_n_layers is not None and self.output_hidden_states and outputs.hidden_states is not None: + selected = outputs.hidden_states[-self.use_ssl_n_layers:] + last_hidden_state = torch.mean(torch.stack(selected, dim=0), dim=0) + del outputs + + return last_hidden_state \ No newline at end of file diff --git a/audio/aasist3/requirements.txt b/audio/aasist3/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..e77b3643d0b1ad3e86b7d3b0f79b85c8425e4af4 --- /dev/null +++ b/audio/aasist3/requirements.txt @@ -0,0 +1,11 @@ +torch==2.5.1 +torchaudio==2.5.1 +transformers>=4.40.0 +huggingface-hub>=0.20.0 +safetensors>=0.4.0 +numpy<2.0 +soundfile +fastapi +uvicorn[standard] +pydantic +python-multipart diff --git a/audio/nes2net/.gitignore b/audio/nes2net/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..d285b0dd1d6c257ec0dea379665e40a74f8e8759 --- /dev/null +++ b/audio/nes2net/.gitignore @@ -0,0 +1 @@ +weights/ diff --git a/audio/nes2net/Dockerfile b/audio/nes2net/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..d2a777d1882caa3eb548a65dd14f343de8190b92 --- /dev/null +++ b/audio/nes2net/Dockerfile @@ -0,0 +1,41 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3 python3-pip python3-dev \ + git ffmpeg libsndfile1 wget \ + build-essential g++ \ + && rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +WORKDIR /app + +# Install PyTorch with CUDA 12.1 +RUN pip install --no-cache-dir \ + torch==2.5.1 torchaudio==2.5.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +# Clone fairseq with patched C extensions (same as ShiftySpeech) +RUN git clone https://github.com/facebookresearch/fairseq.git /app/fairseq_repo && \ + cd /app/fairseq_repo && \ + git checkout a54021305d6b3c4c5959ac9395135f63202db8f1 && \ + sed -i 's/ext_modules=extensions/ext_modules=[]/' setup.py && \ + pip install --no-cache-dir --no-deps -e . + +COPY model_scripts /app/model_scripts +COPY api.py . +RUN mkdir -p /app/weights +COPY weights/ /app/weights/ + +EXPOSE 8004 + +RUN adduser --disabled-password --gecos '' appuser +USER appuser + +CMD ["python", "api.py"] diff --git a/audio/nes2net/api.py b/audio/nes2net/api.py new file mode 100644 index 0000000000000000000000000000000000000000..d085da9d55e6041cb882565c2077fbf39868e0ba --- /dev/null +++ b/audio/nes2net/api.py @@ -0,0 +1,315 @@ +"""Nes2Net (XLSR + Nested Res2Net TDNN) Audio Deepfake Detection API. + +Detects synthetic speech using the Nes2Net model architecture: +- Frontend: XLSR wav2vec 2.0 (Self-Supervised Learning) +- Backend: Nested Res2Net TDNN with SE modules + +Reference: https://github.com/TianchiLiu/Nes2Net +""" + +import argparse +import base64 +import io +import logging +import os +import platform +import sys +import time +import warnings +from typing import Optional + +import librosa +import numpy as np +import torch +import uvicorn +from fastapi import FastAPI, HTTPException +from pydantic import BaseModel, Field + +# Suppress deprecation warnings from fairseq/omegaconf compatibility +warnings.filterwarnings("ignore", category=DeprecationWarning) + +# Monkey-patch omegaconf for fairseq compatibility (older fairseq +# expects is_primitive_type which was removed in newer omegaconf). +import omegaconf._utils as _omegaconf_utils + +if not hasattr(_omegaconf_utils, "is_primitive_type"): + _omegaconf_utils.is_primitive_type = lambda t: t in (int, float, bool, str, bytes) + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", +) +logger = logging.getLogger("nes2net_api") + +# Add the model code to the path +if "/app" not in sys.path: + sys.path.insert(0, "/app") + +# Import model class (deferred to allow path setup) +try: + from model_scripts.wav2vec2_Nes2Net_X import ( + wav2vec2_Nes2Net_no_Res_w_allT as Nes2NetModel, + ) +except ImportError as e: + logger.error(f"Failed to import Nes2Net model: {e}") + Nes2NetModel = None + +# Constants +MODEL_NAME = "nes2net" +MODEL_ID = "nes2net_xlsr_itw_valaug" +WEIGHTS_PATH = "/app/weights/nes2net_itw_valaug.pt" + + +def _get_device(): + """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if override: + pass # Invalid override, fall through to auto-detect + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + return torch.device("cpu") + + +DEVICE = _get_device() + +if DEVICE.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + +if DEVICE.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) +else: + logger.warning( + "Device: %s (no CUDA available -- check nvidia-container-toolkit)", + DEVICE, + ) + +SAMPLE_RATE = 16000 +TARGET_SAMPLES = 64600 # ~4.04 seconds at 16kHz + +# Global model instance +model = None + + +class AudioInput(BaseModel): + """Request schema for audio deepfake detection.""" + + audio_data: str = Field( + ..., description="Base64 encoded audio string (WAV/MP3/etc)" + ) + threshold: Optional[float] = Field( + 0.5, ge=0.0, le=1.0, description="Classification threshold" + ) + + +app = FastAPI( + title="Nes2Net Audio Deepfake Detection API", + description=( + "Service for detecting synthetic speech using the " + "Nes2Net model (XLSR wav2vec 2.0 + Nested Res2Net TDNN)." + ), + version="1.0.0", +) + + +def load_model(): + """Load the Nes2Net model with fine-tuned weights. + + Returns: + The loaded model, or None if loading fails. + """ + global model + if model is not None: + return model + + logger.info(f"Loading Nes2Net model onto {DEVICE}...") + + if Nes2NetModel is None: + logger.error("Nes2Net model class not available.") + return None + + if not os.path.exists(WEIGHTS_PATH): + logger.error(f"Model weights not found at {WEIGHTS_PATH}") + return None + + try: + args = argparse.Namespace( + n_output_logits=2, + dilation=2, + pool_func="mean", + SE_ratio=[1], + Nes_ratio=[8, 8], + ) + model = Nes2NetModel(args, str(DEVICE)) + + # Load fine-tuned weights + try: + state_dict = torch.load( + WEIGHTS_PATH, + map_location=DEVICE, + weights_only=False, + ) + except TypeError: + state_dict = torch.load(WEIGHTS_PATH, map_location=DEVICE) + + model.load_state_dict(state_dict) + model.to(DEVICE) + model.eval() + + logger.info("Nes2Net model loaded successfully.") + return model + except Exception as e: + logger.exception(f"Failed to load Nes2Net model: {e}") + model = None + return None + + +@app.on_event("startup") +async def startup_event(): + """Load model on service startup.""" + load_model() + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and DEVICE.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +async def health(): + """Health check endpoint.""" + return { + "status": "healthy" if model is not None else "degraded", + "model": MODEL_NAME, + "model_id": MODEL_ID, + "device": str(DEVICE), + "weights_found": os.path.exists(WEIGHTS_PATH), + **_gpu_health_info(), + } + + +def preprocess_audio(audio_bytes: bytes) -> torch.Tensor: + """Preprocess audio for Nes2Net inference. + + Loads audio, resamples to 16kHz mono, and pads/trims + to TARGET_SAMPLES using tiling (matching original training + preprocessing). + + Args: + audio_bytes: Raw audio file bytes. + + Returns: + Audio tensor of shape (1, TARGET_SAMPLES). + + Raises: + ValueError: If audio preprocessing fails. + """ + try: + logger.info("Starting audio preprocessing...") + audio, sr = librosa.load(io.BytesIO(audio_bytes), sr=SAMPLE_RATE, mono=True) + logger.info(f"Audio loaded. Length: {len(audio)} samples at {sr}Hz") + + # Pad/trim to TARGET_SAMPLES using tiling + if len(audio) >= TARGET_SAMPLES: + audio = audio[:TARGET_SAMPLES] + else: + num_repeats = TARGET_SAMPLES // len(audio) + 1 + audio = np.tile(audio, num_repeats)[:TARGET_SAMPLES] + + logger.info(f"Audio padded/trimmed to {TARGET_SAMPLES} samples") + + audio_tensor = torch.FloatTensor(audio).unsqueeze(0).to(DEVICE) + return audio_tensor + except Exception as e: + logger.error(f"Error preprocessing audio: {e}") + raise ValueError(f"Audio preprocessing failed: {str(e)}") + + +@app.post("/predict") +async def predict(input_data: AudioInput): + """Run deepfake detection on base64-encoded audio. + + The model outputs 2 logits: [spoof_score, bonafide_score]. + Class 0 = spoof (fake), Class 1 = bonafide (real). + The returned probability is the spoof/fake probability. + """ + if model is None: + if load_model() is None: + raise HTTPException(status_code=503, detail="Model not loaded") + + try: + start_time = time.time() + logger.info( + f"Prediction request. Data size: " f"{len(input_data.audio_data)} chars" + ) + + # Decode base64 audio + audio_bytes = base64.b64decode(input_data.audio_data) + + # Preprocess + audio_tensor = preprocess_audio(audio_bytes) + + # Inference + logger.info("Starting model inference...") + with torch.no_grad(): + output = model(audio_tensor) + + # output shape: [batch, 2] + # Index 0 = spoof logit, Index 1 = bonafide logit + probs = torch.softmax(output, dim=1) + prob_fake = probs[0, 0].item() + + prediction = 1 if prob_fake >= input_data.threshold else 0 + verdict = "fake" if prediction == 1 else "real" + inference_time = time.time() - start_time + + logger.info( + f"Prediction: {verdict} (prob_fake={prob_fake:.4f}, " + f"time={inference_time:.3f}s)" + ) + + return { + "model": MODEL_NAME, + "probability": float(prob_fake), + "prediction": int(prediction), + "class": verdict, + "inference_time": float(inference_time), + } + + except Exception as e: + logger.exception(f"Error during prediction: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +if __name__ == "__main__": + port = int(os.environ.get("MODEL_PORT", 8004)) + uvicorn.run(app, host="0.0.0.0", port=port) diff --git a/audio/nes2net/model_scripts/__init__.py b/audio/nes2net/model_scripts/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/audio/nes2net/model_scripts/wav2vec2_Nes2Net_X.py b/audio/nes2net/model_scripts/wav2vec2_Nes2Net_X.py new file mode 100644 index 0000000000000000000000000000000000000000..4d8b5917d96c3217fe457828e619f1c1abbc69a3 --- /dev/null +++ b/audio/nes2net/model_scripts/wav2vec2_Nes2Net_X.py @@ -0,0 +1,317 @@ +import math + +import fairseq +import torch +import torch.nn as nn + +___author__ = "Tianchi Liu" +__email__ = "tianchi_liu@u.nus.edu" +# modified from the model script from Hemlata Tak + + +class SSLModel(nn.Module): + def __init__(self, device): + super(SSLModel, self).__init__() + cp_path = ( + "/app/weights/xlsr2_300m.pt" # Change the pre-trained XLSR model path. + ) + model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task( + [cp_path] + ) + self.model = model[0] + self.device = device + self.out_dim = 1024 + return + + def extract_feat(self, input_data): + # put the model to GPU if it not there + if ( + next(self.model.parameters()).device != input_data.device + or next(self.model.parameters()).dtype != input_data.dtype + ): + self.model.to(input_data.device, dtype=input_data.dtype) + self.model.train() + if True: + # input should be in shape (batch, length) + if input_data.ndim == 3: + input_tmp = input_data[:, :, 0] + else: + input_tmp = input_data + # [batch, length, dim] + emb = self.model(input_tmp, mask=False, features_only=True)["x"] + return emb + + +class SEModule(nn.Module): + def __init__(self, channels, SE_ratio=8): + super(SEModule, self).__init__() + self.se = nn.Sequential( + nn.AdaptiveAvgPool1d(1), + nn.Conv1d(channels, channels // SE_ratio, kernel_size=1, padding=0), + nn.ReLU(), + nn.Conv1d(channels // SE_ratio, channels, kernel_size=1, padding=0), + nn.Sigmoid(), + ) + + def forward(self, input): + x = self.se(input) + return input * x + + +class Bottle2neck(nn.Module): + + def __init__( + self, inplanes, planes, kernel_size=None, dilation=None, scale=8, SE_ratio=8 + ): + super(Bottle2neck, self).__init__() + width = int(math.floor(planes / scale)) + self.conv1 = nn.Conv1d(inplanes, width * scale, kernel_size=1) + self.bn1 = nn.BatchNorm1d(width * scale) + self.nums = scale - 1 + convs = [] + bns = [] + weighted_sum = [] + num_pad = math.floor(kernel_size / 2) * dilation + for i in range(self.nums): + convs.append( + nn.Conv2d( + width, + width, + kernel_size=(kernel_size, 1), + dilation=(dilation, 1), + padding=(num_pad, 0), + ) + ) + bns.append(nn.BatchNorm2d(width)) + initial_value = torch.ones(1, 1, 1, i + 2) * (1 / (i + 2)) + weighted_sum.append(nn.Parameter(initial_value, requires_grad=True)) + self.weighted_sum = nn.ParameterList(weighted_sum) + self.convs = nn.ModuleList(convs) + self.bns = nn.ModuleList(bns) + self.conv3 = nn.Conv1d(width * scale, planes, kernel_size=1) + self.bn3 = nn.BatchNorm1d(planes) + self.relu = nn.ReLU() + self.width = width + self.se = SEModule(planes, SE_ratio) + + def forward(self, x): + residual = x + out = self.conv1(x) + out = self.relu(out) + out = self.bn1(out).unsqueeze(-1) # bz c T 1 + + spx = torch.split(out, self.width, 1) + sp = spx[self.nums] + for i in range(self.nums): + sp = torch.cat((sp, spx[i]), -1) + + sp = self.bns[i](self.relu(self.convs[i](sp))) + sp_s = sp * self.weighted_sum[i] + sp_s = torch.sum(sp_s, dim=-1, keepdim=False) + + if i == 0: + out = sp_s + else: + out = torch.cat((out, sp_s), 1) + out = torch.cat((out, spx[self.nums].squeeze(-1)), 1) + out = self.conv3(out) + out = self.relu(out) + out = self.bn3(out) + out = self.se(out) + out += residual + return out + + +class ASTP(nn.Module): + """Attentive statistics pooling: Channel- and context-dependent + statistics pooling, first used in ECAPA_TDNN. + """ + + def __init__(self, in_dim, bottleneck_dim=128, global_context_att=False): + super(ASTP, self).__init__() + self.global_context_att = global_context_att + + # Use Conv1d with stride == 1 rather than Linear, then we don't + # need to transpose inputs. + if global_context_att: + self.linear1 = nn.Conv1d( + in_dim * 3, bottleneck_dim, kernel_size=1 + ) # equals W and b in the paper + else: + self.linear1 = nn.Conv1d( + in_dim, bottleneck_dim, kernel_size=1 + ) # equals W and b in the paper + self.linear2 = nn.Conv1d( + bottleneck_dim, in_dim, kernel_size=1 + ) # equals V and k in the paper + + def forward(self, x): + """ + x: a 3-dimensional tensor in tdnn-based architecture (B,F,T) + or a 4-dimensional tensor in resnet architecture (B,C,F,T) + 0-dim: batch-dimension, last-dim: time-dimension (frame-dimension) + """ + if len(x.shape) == 4: + x = x.reshape(x.shape[0], x.shape[1] * x.shape[2], x.shape[3]) + assert len(x.shape) == 3 + + if self.global_context_att: + context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x) + context_std = torch.sqrt( + torch.var(x, dim=-1, keepdim=True) + 1e-10 + ).expand_as(x) + x_in = torch.cat((x, context_mean, context_std), dim=1) + else: + x_in = x + + # DON'T use ReLU here! ReLU may be hard to converge. + alpha = torch.tanh(self.linear1(x_in)) # alpha = F.relu(self.linear1(x_in)) + alpha = torch.softmax(self.linear2(alpha), dim=2) + mean = torch.sum(alpha * x, dim=2) + var = torch.sum(alpha * (x**2), dim=2) - mean**2 + std = torch.sqrt(var.clamp(min=1e-10)) + return torch.cat([mean, std], dim=1) + + +class Nested_Res2Net_TDNN(nn.Module): + + def __init__( + self, + Nes_ratio=[8, 8], + input_channel=1024, + n_output_logits=2, + dilation=2, + pool_func="mean", + SE_ratio=[8], + ): + + super(Nested_Res2Net_TDNN, self).__init__() + self.Nes_ratio = Nes_ratio[0] + assert input_channel % Nes_ratio[0] == 0 + C = input_channel // Nes_ratio[0] + self.C = C + Build_in_Res2Nets = [] + bns = [] + for i in range(Nes_ratio[0] - 1): + Build_in_Res2Nets.append( + Bottle2neck( + C, + C, + kernel_size=3, + dilation=dilation, + scale=Nes_ratio[1], + SE_ratio=SE_ratio[0], + ) + ) + bns.append(nn.BatchNorm1d(C)) + self.Build_in_Res2Nets = nn.ModuleList(Build_in_Res2Nets) + self.bns = nn.ModuleList(bns) + self.bn = nn.BatchNorm1d(1024) + self.relu = nn.ReLU() + self.pool_func = pool_func + if pool_func == "mean": + self.fc = nn.Linear(1024, n_output_logits) + elif pool_func == "ASTP": + self.pooling = ASTP( + in_dim=input_channel, bottleneck_dim=128, global_context_att=False + ) + self.fc = nn.Linear(2048, n_output_logits) + + def forward(self, x): + spx = torch.split(x, self.C, 1) + for i in range(self.Nes_ratio - 1): + if i == 0: + sp = spx[i] + else: + sp = sp + spx[i] + sp = self.Build_in_Res2Nets[i](sp) + sp = self.relu(sp) + sp = self.bns[i](sp) + if i == 0: + out = sp + else: + out = torch.cat((out, sp), 1) + out = torch.cat((out, spx[-1]), 1) + out = self.bn(out) + out = self.relu(out) + if self.pool_func == "mean": + out = torch.mean(out, dim=-1) + elif self.pool_func == "ASTP": + out = self.pooling(out) + out = self.fc(out) + return out + + +class wav2vec2_Nes2Net_no_Res_w_allT(nn.Module): + def __init__(self, args, device): + super().__init__() + self.device = device + + self.n_output_logits = args.n_output_logits + + #### + # create network wav2vec 2.0 + #### + self.ssl_model = SSLModel(self.device) + self.Nested_Res2Net_TDNN = Nested_Res2Net_TDNN( + Nes_ratio=args.Nes_ratio, + input_channel=1024, + n_output_logits=self.n_output_logits, + dilation=args.dilation, + pool_func=args.pool_func, + SE_ratio=args.SE_ratio, + ) + + def forward(self, x): + # -------pre-trained Wav2vec model fine tunning ------------------------## + x_ssl_feat = self.ssl_model.extract_feat(x.squeeze(-1)) + x_ssl_feat = x_ssl_feat.permute(0, 2, 1) + output = self.Nested_Res2Net_TDNN(x_ssl_feat) + + return output + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("--n_output_logits", type=int, default=2) + parser.add_argument("--dilation", type=int, default=2) # not important + parser.add_argument( + "--pool_func", + type=str, + default="mean", + choices=["mean", "ASTP"], + help="pooling function, choose from mean and ASTP", + ) + parser.add_argument( + "--Nes_ratio", + type=int, + nargs="+", + default=[8, 8], + help="Nes_ratio, from outer to inner", + ) + parser.add_argument( + "--SE_ratio", + type=int, + nargs="+", + default=[1], + help="SE downsampling ratio in the bottleneck", + ) + args = parser.parse_args() + + model = wav2vec2_Nes2Net_no_Res_w_allT(args=args, device="cpu") + x = torch.rand((4, 32000)).to("cpu") + model = model.to("cpu") + y = model(x) + print(y) + trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + print("all:", trainable_params) + trainable_params = sum( + p.numel() for p in model.ssl_model.parameters() if p.requires_grad + ) + print("SSL:", trainable_params) + trainable_params = sum( + p.numel() for p in model.Nested_Res2Net_TDNN.parameters() if p.requires_grad + ) + print("Backend:", trainable_params) diff --git a/audio/nes2net/requirements.txt b/audio/nes2net/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..9ba4009e0d0d24b612425b63f22112303c5e49cc --- /dev/null +++ b/audio/nes2net/requirements.txt @@ -0,0 +1,13 @@ +torch==2.5.1 +torchaudio==2.5.1 +numpy==1.23.5 +librosa==0.9.1 +soundfile +scipy +omegaconf +hydra-core +bitarray +fastapi +uvicorn[standard] +pydantic +python-multipart diff --git a/audio/safeear/Dockerfile b/audio/safeear/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..414d643ddcd0a1e946dd376377e7727ca01af359 --- /dev/null +++ b/audio/safeear/Dockerfile @@ -0,0 +1,49 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3 python3-pip python3-dev \ + git ffmpeg libsndfile1 wget \ + build-essential \ + && rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +WORKDIR /app + +# Install PyTorch with CUDA 12.1 +RUN pip install --no-cache-dir \ + torch==2.5.1 torchaudio==2.5.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +# Clone SafeEar repository (for model code imports) +RUN git clone --depth 1 https://github.com/LetterLiGo/SafeEar.git /app/safeear_repo + +# Install the fairseq fork with C extensions PATCHED OUT +# (same proven patch used by ShiftySpeech and Nes2Net -- +# C extensions are not needed for checkpoint loading) +WORKDIR /app/safeear_repo/fairseq_ours +RUN sed -i 's/ext_modules=extensions/ext_modules=[]/' setup.py && \ + pip install --no-cache-dir --no-deps -e . +WORKDIR /app + +# Download model weights from HuggingFace +RUN mkdir -p /app/weights && \ + wget -q -O /app/weights/SpeechTokenizer.pt \ + "https://huggingface.co/TEC2004/SafeEar-ASV19-spoof-detection/resolve/main/SpeechTokenizer.pt" && \ + wget -q -O /app/weights/model.ckpt \ + "https://huggingface.co/TEC2004/SafeEar-ASV19-spoof-detection/resolve/main/model.ckpt" + +COPY api.py . + +EXPOSE 8002 + +RUN adduser --disabled-password --gecos '' appuser +USER appuser + +CMD ["python", "api.py"] diff --git a/audio/safeear/api.py b/audio/safeear/api.py new file mode 100644 index 0000000000000000000000000000000000000000..0ea75779e64d3e101dd30afcbcd81f982ea4e697 --- /dev/null +++ b/audio/safeear/api.py @@ -0,0 +1,321 @@ +"""SafeEar audio deepfake detection API service. + +Uses the SafeEar content privacy-preserving model (CCS 2024) to detect +synthetic speech. Two-stage pipeline: + 1. SpeechTokenizer (neural audio codec) decouples acoustic features + 2. SafeEar1s (transformer classifier) detects spoofing from acoustic tokens + +Weights: HuggingFace TEC2004/SafeEar-ASV19-spoof-detection +""" + +import base64 +import logging +import os +import sys +import tempfile +import time +from typing import Optional + +import uvicorn +from fastapi import FastAPI, HTTPException +from pydantic import BaseModel, Field + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", +) +logger = logging.getLogger("safeear_api") + +import platform + +import librosa +import numpy as np +import torch + +# Add SafeEar repo to path for model imports +SAFEEAR_REPO_PATH = os.environ.get( + "SAFEEAR_REPO_PATH", + os.path.join(os.path.dirname(__file__), "safeear_repo"), +) +if SAFEEAR_REPO_PATH not in sys.path: + sys.path.insert(0, SAFEEAR_REPO_PATH) + + +def _get_device(): + """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + return torch.device("cpu") + + +# Constants +MODEL_NAME = "safeear" +WEIGHTS_DIR = os.environ.get( + "WEIGHTS_DIR", + os.path.join(os.path.dirname(__file__), "weights"), +) +DEVICE = _get_device() + +if DEVICE.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + +if DEVICE.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) +else: + logger.warning( + "Device: %s (no CUDA available -- check nvidia-container-toolkit)", + DEVICE, + ) + +SAMPLE_RATE = 16000 +MAX_AUDIO_LENGTH = 64600 # ~4 seconds at 16kHz (ASVspoof standard) +SOFTMAX_TEMPERATURE = 5.0 # Calibration temperature for out-of-distribution data +NUM_INFERENCE_PASSES = 5 # Monte Carlo passes for stable predictions + +# Global model instances +decouple_model = None +detect_model = None + + +class AudioInput(BaseModel): + """Schema for audio prediction requests.""" + + audio_data: str = Field( + ..., description="Base64 encoded audio string (WAV/MP3/etc)" + ) + threshold: Optional[float] = Field( + 0.5, ge=0.0, le=1.0, description="Classification threshold" + ) + + +app = FastAPI( + title="SafeEar Audio Deepfake Detection API", + description="Content privacy-preserving deepfake detection using SafeEar.", + version="1.0.0", +) + + +def load_models(): + """Load both the decouple model (SpeechTokenizer) and detect model.""" + global decouple_model, detect_model + + if decouple_model is not None and detect_model is not None: + return True + + speech_tokenizer_path = os.path.join(WEIGHTS_DIR, "SpeechTokenizer.pt") + checkpoint_path = os.path.join(WEIGHTS_DIR, "model.ckpt") + + if not os.path.exists(speech_tokenizer_path): + logger.error(f"SpeechTokenizer weights not found: {speech_tokenizer_path}") + return False + if not os.path.exists(checkpoint_path): + logger.error(f"Model checkpoint not found: {checkpoint_path}") + return False + + try: + # --- Load SpeechTokenizer (decouple model) --- + from safeear.models.decouple import SpeechTokenizer + + logger.info("Loading SpeechTokenizer...") + decouple_model = SpeechTokenizer( + n_filters=64, + strides=[8, 5, 4, 2], + dimension=1024, + semantic_dimension=768, + bidirectional=True, + dilation_base=2, + residual_kernel_size=3, + n_residual_layers=1, + lstm_layers=2, + activation="ELU", + codebook_size=1024, + n_q=8, + sample_rate=16000, + ) + st_state = torch.load(speech_tokenizer_path, map_location="cpu") + decouple_model.load_state_dict(st_state) + decouple_model.to(DEVICE) + decouple_model.eval() + logger.info("SpeechTokenizer loaded.") + + # --- Load SafeEar1s (detect model) from Lightning checkpoint --- + from safeear.models.safeear import SafeEar1s, SE_Rawformer_front + + logger.info("Loading SafeEar1s detect model...") + detect_model = SafeEar1s( + front=SE_Rawformer_front(), + embedding_dim=1024, + dropout_rate=0.1, + attention_dropout=0.1, + stochastic_depth=0.1, + num_layers=2, + num_heads=8, + num_classes=2, + positional_embedding="sine", + mlp_ratio=1.0, + ) + + # The .ckpt is a PyTorch Lightning checkpoint + ckpt = torch.load(checkpoint_path, map_location="cpu") + state_dict = ckpt.get("state_dict", ckpt) + + # Lightning prefixes keys with "detect_model." + detect_state = {} + for k, v in state_dict.items(): + if k.startswith("detect_model."): + detect_state[k.replace("detect_model.", "", 1)] = v + + detect_model.load_state_dict(detect_state) + detect_model.to(DEVICE) + detect_model.eval() + logger.info("SafeEar1s detect model loaded.") + return True + + except Exception as e: + logger.exception(f"Failed to load SafeEar models: {e}") + decouple_model = None + detect_model = None + return False + + +def preprocess_audio(audio_bytes: bytes) -> torch.Tensor: + """Load audio bytes, resample to 16kHz mono, pad/trim.""" + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: + tmp.write(audio_bytes) + tmp_path = tmp.name + + try: + waveform, _ = librosa.load(tmp_path, sr=SAMPLE_RATE, mono=True) + finally: + os.unlink(tmp_path) + + if len(waveform) < MAX_AUDIO_LENGTH: + waveform = np.pad(waveform, (0, MAX_AUDIO_LENGTH - len(waveform))) + else: + waveform = waveform[:MAX_AUDIO_LENGTH] + + # Shape: (1, 1, samples) -- batch=1, channels=1, time + tensor = torch.FloatTensor(waveform).unsqueeze(0).unsqueeze(0).to(DEVICE) + return tensor + + +@app.on_event("startup") +async def startup_event(): + """Attempt to load models at startup.""" + load_models() + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and DEVICE.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +async def health(): + """Return service health status and model availability.""" + models_loaded = decouple_model is not None and detect_model is not None + return { + "status": "healthy" if models_loaded else "degraded", + "model": MODEL_NAME, + "device": str(DEVICE), + "weights_found": ( + os.path.exists(os.path.join(WEIGHTS_DIR, "SpeechTokenizer.pt")) + and os.path.exists(os.path.join(WEIGHTS_DIR, "model.ckpt")) + ), + **_gpu_health_info(), + } + + +@app.post("/predict") +async def predict(input_data: AudioInput): + """Run SafeEar inference on base64-encoded audio data.""" + if decouple_model is None or detect_model is None: + if not load_models(): + raise HTTPException(status_code=503, detail="Models not loaded") + + try: + start_time = time.time() + logger.info( + "Received prediction request. " + f"Data size: {len(input_data.audio_data)} chars" + ) + + audio_bytes = base64.b64decode(input_data.audio_data) + x_wav = preprocess_audio(audio_bytes) + + with torch.no_grad(): + # Step 1: Extract acoustic tokens via SpeechTokenizer + # forward() returns: + # (reconstructed, commit_loss, semantic_feature, acoustic_tokens) + # layers=[0,1,2,3,4,5,6,7] means layer 0 goes to + # semantic_feature; layers 1-7 go to acoustic_tokens list + _, _, _, acoustic_tokens = decouple_model( + x_wav, layers=[0, 1, 2, 3, 4, 5, 6, 7] + ) + + # Step 2: Run detection model with Monte Carlo averaging + # SafeEar1s uses torch.randperm() in forward, so we average + # multiple passes for stable predictions + logit_sum = torch.zeros(1, 2, device=DEVICE) + for _ in range(NUM_INFERENCE_PASSES): + raw_logits, _ = detect_model(acoustic_tokens) + logit_sum += raw_logits + avg_logits = logit_sum / NUM_INFERENCE_PASSES + + # Step 3: Get fake probability with temperature-scaled softmax + # The model produces extreme logits that saturate standard + # softmax. Temperature scaling preserves discrimination while + # giving more interpretable probabilities. + probs = torch.softmax(avg_logits / SOFTMAX_TEMPERATURE, dim=-1) + prob_fake = probs[0, 1].item() + + prediction = 1 if prob_fake >= input_data.threshold else 0 + verdict = "fake" if prediction == 1 else "real" + inference_time = time.time() - start_time + + return { + "model": MODEL_NAME, + "probability": float(prob_fake), + "prediction": int(prediction), + "class": verdict, + "inference_time": float(inference_time), + } + + except Exception as e: + logger.exception(f"Error during prediction: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +if __name__ == "__main__": + port = int(os.environ.get("MODEL_PORT", 8002)) + uvicorn.run(app, host="0.0.0.0", port=port) diff --git a/audio/safeear/download_weights.sh b/audio/safeear/download_weights.sh new file mode 100644 index 0000000000000000000000000000000000000000..beb0406d96ce1f1b333965c429952c57325e2222 --- /dev/null +++ b/audio/safeear/download_weights.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +set -euo pipefail + +WEIGHTS_DIR="${WEIGHTS_DIR:-/app/weights}" +REPO_DIR="${REPO_DIR:-/app/safeear_repo}" + +mkdir -p "$WEIGHTS_DIR" + +echo "==> Cloning SafeEar source repository..." +if [ ! -d "$REPO_DIR/.git" ]; then + git clone --depth 1 https://github.com/LetterLiGo/SafeEar.git "$REPO_DIR" +fi + +echo "==> Downloading SpeechTokenizer.pt from HuggingFace..." +wget -q --show-progress -O "$WEIGHTS_DIR/SpeechTokenizer.pt" \ + "https://huggingface.co/TEC2004/SafeEar-ASV19-spoof-detection/resolve/main/SpeechTokenizer.pt" + +echo "==> Downloading model.ckpt from HuggingFace..." +wget -q --show-progress -O "$WEIGHTS_DIR/model.ckpt" \ + "https://huggingface.co/TEC2004/SafeEar-ASV19-spoof-detection/resolve/main/model.ckpt" + +echo "==> Weights downloaded to $WEIGHTS_DIR" +ls -lh "$WEIGHTS_DIR" diff --git a/audio/safeear/requirements.txt b/audio/safeear/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..2c3414f7989d54b3981bddd947a17d53b3b5e394 --- /dev/null +++ b/audio/safeear/requirements.txt @@ -0,0 +1,15 @@ +torch==2.5.1 +torchaudio==2.5.1 +librosa>=0.10.0 +soundfile>=0.11.0 +numpy>=1.23.0 +einops>=0.7.0 +timm>=0.9.0 +hydra-core>=1.0.7 +omegaconf>=2.1.0 +pytorch-lightning>=1.6.0 +scipy>=1.11.0 +fastapi>=0.100.0 +uvicorn[standard]>=0.20.0 +python-multipart>=0.0.5 +pydantic>=2.0.0 diff --git a/audio/shiftyspeech/Dockerfile b/audio/shiftyspeech/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..ce862b0ac4f1d9ab15157a56061777ec2f968c7e --- /dev/null +++ b/audio/shiftyspeech/Dockerfile @@ -0,0 +1,43 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3 python3-pip python3-dev \ + git ffmpeg libsndfile1 \ + build-essential g++ \ + && rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +WORKDIR /app + +# Install PyTorch with CUDA 12.1 +RUN pip install --no-cache-dir \ + torch==2.5.1 torchaudio==2.5.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +# Clone fairseq with patched C extensions +RUN git clone https://github.com/facebookresearch/fairseq.git /app/fairseq_repo && \ + cd /app/fairseq_repo && \ + git checkout a54021305d6b3c4c5959ac9395135f63202db8f1 && \ + sed -i 's/ext_modules=extensions/ext_modules=[]/' setup.py && \ + pip install --no-cache-dir --no-deps -e . + +COPY synthetic_speech_detection /app/synthetic_speech_detection +RUN mkdir -p /app/models +COPY api.py . +RUN mkdir -p /app/weights +COPY weights/ /app/weights/ +RUN ln -sf /app/weights/xlsr_53_56k.pt /app/models/xlsr2_300m.pt + +EXPOSE 8001 + +RUN adduser --disabled-password --gecos '' appuser +USER appuser + +CMD ["python", "api.py"] diff --git a/audio/shiftyspeech/api.py b/audio/shiftyspeech/api.py new file mode 100644 index 0000000000000000000000000000000000000000..729e5449a82830f0f8be9b51c7c4ae8dd3bea8a9 --- /dev/null +++ b/audio/shiftyspeech/api.py @@ -0,0 +1,315 @@ +"""ShiftySpeech (SSL-AASIST) Audio Deepfake Detection API. + +Detects synthetic speech using the SSL-AASIST model architecture: +- Frontend: XLSR wav2vec 2.0 (Self-Supervised Learning) +- Backend: AASIST (Audio Anti-Spoofing using Integrated + Spectro-Temporal Graph Attention Networks) + +Reference: https://github.com/Ashigarg123/ShiftySpeech +""" + +import base64 +import io +import logging +import os +import platform +import sys +import time +import warnings +from typing import Optional + +import librosa +import numpy as np +import torch +import uvicorn +from fastapi import FastAPI, HTTPException +from pydantic import BaseModel, Field + +# Suppress deprecation warnings from fairseq/omegaconf compatibility +warnings.filterwarnings("ignore", category=DeprecationWarning) + +# Monkey-patch omegaconf for fairseq compatibility (older fairseq +# expects is_primitive_type which was removed in newer omegaconf). +import omegaconf._utils as _omegaconf_utils + +if not hasattr(_omegaconf_utils, "is_primitive_type"): + _omegaconf_utils.is_primitive_type = lambda t: t in (int, float, bool, str, bytes) + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", +) +logger = logging.getLogger("shiftyspeech_api") + +# Add the SSL_Anti-spoofing model code to the path +MODEL_CODE_PATH = "/app/synthetic_speech_detection/SSL_Anti-spoofing" +if MODEL_CODE_PATH not in sys.path: + sys.path.insert(0, MODEL_CODE_PATH) + +# Import model class (deferred to allow path setup) +try: + from model import Model as SSLAASISTModel +except ImportError as e: + logger.error(f"Failed to import SSL-AASIST model: {e}") + SSLAASISTModel = None + +# Constants +MODEL_NAME = "shiftyspeech" +MODEL_ID = "ssl_aasist_augmented" +WEIGHTS_PATH = "/app/weights/hfg_aug_1_2.pt" +XLSR_DIR = "/app/models" + + +def _get_device(): + """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if override: + pass # Invalid override, fall through to auto-detect + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + return torch.device("cpu") + + +DEVICE = _get_device() + +if DEVICE.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + +if DEVICE.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) +else: + logger.warning( + "Device: %s (no CUDA available -- check nvidia-container-toolkit)", + DEVICE, + ) + +SAMPLE_RATE = 16000 +TARGET_SAMPLES = 64600 # ~4.04 seconds at 16kHz + +# Global model instance +model = None + + +class AudioInput(BaseModel): + """Request schema for audio deepfake detection.""" + + audio_data: str = Field( + ..., description="Base64 encoded audio string (WAV/MP3/etc)" + ) + threshold: Optional[float] = Field( + 0.5, ge=0.0, le=1.0, description="Classification threshold" + ) + + +app = FastAPI( + title="ShiftySpeech Audio Deepfake Detection API", + description=( + "Service for detecting synthetic speech using the " + "SSL-AASIST model (XLSR wav2vec 2.0 + AASIST backend)." + ), + version="1.0.0", +) + + +def load_model(): + """Load the SSL-AASIST model with augmented weights. + + Returns: + The loaded model, or None if loading fails. + """ + global model + if model is not None: + return model + + logger.info(f"Loading SSL-AASIST model onto {DEVICE}...") + + if SSLAASISTModel is None: + logger.error("SSL-AASIST model class not available.") + return None + + if not os.path.exists(WEIGHTS_PATH): + logger.error(f"Model weights not found at {WEIGHTS_PATH}") + return None + + try: + # Ensure XLSR model directory exists for architecture init + os.makedirs(XLSR_DIR, exist_ok=True) + + import argparse + + args = argparse.Namespace() + model = SSLAASISTModel(args, str(DEVICE)) + + # Load fine-tuned weights (includes XLSR weights) + try: + state_dict = torch.load( + WEIGHTS_PATH, + map_location=DEVICE, + weights_only=False, + ) + except TypeError: + state_dict = torch.load(WEIGHTS_PATH, map_location=DEVICE) + + model.load_state_dict(state_dict) + model.to(DEVICE) + model.eval() + + logger.info("SSL-AASIST model loaded successfully.") + return model + except Exception as e: + logger.exception(f"Failed to load SSL-AASIST model: {e}") + model = None + return None + + +@app.on_event("startup") +async def startup_event(): + """Load model on service startup.""" + load_model() + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and DEVICE.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +async def health(): + """Health check endpoint.""" + return { + "status": "healthy" if model is not None else "degraded", + "model": MODEL_NAME, + "model_id": MODEL_ID, + "device": str(DEVICE), + "weights_found": os.path.exists(WEIGHTS_PATH), + **_gpu_health_info(), + } + + +def preprocess_audio(audio_bytes: bytes) -> torch.Tensor: + """Preprocess audio for SSL-AASIST inference. + + Loads audio, resamples to 16kHz mono, and pads/trims + to TARGET_SAMPLES using tiling (matching original training + preprocessing from data_utils.py). + + Args: + audio_bytes: Raw audio file bytes. + + Returns: + Audio tensor of shape (1, TARGET_SAMPLES). + + Raises: + ValueError: If audio preprocessing fails. + """ + try: + logger.info("Starting audio preprocessing...") + audio, sr = librosa.load(io.BytesIO(audio_bytes), sr=SAMPLE_RATE, mono=True) + logger.info(f"Audio loaded. Length: {len(audio)} samples at {sr}Hz") + + # Pad/trim to TARGET_SAMPLES using tiling + # (matches original data_utils.pad function) + if len(audio) >= TARGET_SAMPLES: + audio = audio[:TARGET_SAMPLES] + else: + num_repeats = TARGET_SAMPLES // len(audio) + 1 + audio = np.tile(audio, num_repeats)[:TARGET_SAMPLES] + + logger.info(f"Audio padded/trimmed to {TARGET_SAMPLES} samples") + + audio_tensor = torch.FloatTensor(audio).unsqueeze(0).to(DEVICE) + return audio_tensor + except Exception as e: + logger.error(f"Error preprocessing audio: {e}") + raise ValueError(f"Audio preprocessing failed: {str(e)}") + + +@app.post("/predict") +async def predict(input_data: AudioInput): + """Run deepfake detection on base64-encoded audio. + + The model outputs 2 logits: [spoof_score, bonafide_score]. + Class 0 = spoof (fake), Class 1 = bonafide (real). + The returned probability is the spoof/fake probability. + """ + if model is None: + if load_model() is None: + raise HTTPException(status_code=503, detail="Model not loaded") + + try: + start_time = time.time() + logger.info( + f"Prediction request. Data size: " f"{len(input_data.audio_data)} chars" + ) + + # Decode base64 audio + audio_bytes = base64.b64decode(input_data.audio_data) + + # Preprocess + audio_tensor = preprocess_audio(audio_bytes) + + # Inference + logger.info("Starting model inference...") + with torch.no_grad(): + output = model(audio_tensor) + + # output shape: [batch, 2] + # Index 0 = spoof logit, Index 1 = bonafide logit + probs = torch.softmax(output, dim=1) + prob_fake = probs[0, 0].item() + + prediction = 1 if prob_fake >= input_data.threshold else 0 + verdict = "fake" if prediction == 1 else "real" + inference_time = time.time() - start_time + + logger.info( + f"Prediction: {verdict} (prob_fake={prob_fake:.4f}, " + f"time={inference_time:.3f}s)" + ) + + return { + "model": MODEL_NAME, + "probability": float(prob_fake), + "prediction": int(prediction), + "class": verdict, + "inference_time": float(inference_time), + } + + except Exception as e: + logger.exception(f"Error during prediction: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +if __name__ == "__main__": + port = int(os.environ.get("MODEL_PORT", 8001)) + uvicorn.run(app, host="0.0.0.0", port=port) diff --git a/audio/shiftyspeech/evaluate.py b/audio/shiftyspeech/evaluate.py new file mode 100644 index 0000000000000000000000000000000000000000..3d7dc12e408ba25195a342f377684a143ac7114d --- /dev/null +++ b/audio/shiftyspeech/evaluate.py @@ -0,0 +1,210 @@ +"""Evaluate ShiftySpeech SSL-AASIST on the DeepSafe audio dataset. + +Reports accuracy, precision, recall, F1, EER, and per-file results. +""" + +import os +import sys +import time +import warnings + +warnings.filterwarnings("ignore", category=DeprecationWarning) + +# Monkey-patch omegaconf for fairseq compatibility +import omegaconf._utils as _omegaconf_utils + +if not hasattr(_omegaconf_utils, "is_primitive_type"): + _omegaconf_utils.is_primitive_type = lambda t: t in (int, float, bool, str, bytes) + +import argparse + +import librosa +import numpy as np +import torch + +# Add model code to path +SERVICE_DIR = os.path.dirname(os.path.abspath(__file__)) +MODEL_CODE_PATH = os.path.join( + SERVICE_DIR, "synthetic_speech_detection", "SSL_Anti-spoofing" +) +sys.path.insert(0, MODEL_CODE_PATH) + +from model import Model as SSLAASISTModel + +SAMPLE_RATE = 16000 +TARGET_SAMPLES = 64600 +DATASET_DIR = os.path.join( + SERVICE_DIR, os.pardir, os.pardir, os.pardir, "dataset", "audio" +) +DATASET_DIR = os.path.normpath(DATASET_DIR) + + +def pad_audio(audio, target=TARGET_SAMPLES): + """Pad/trim audio to target length using tiling.""" + if len(audio) >= target: + return audio[:target] + num_repeats = target // len(audio) + 1 + return np.tile(audio, num_repeats)[:target] + + +def compute_eer(target_scores, nontarget_scores): + """Compute Equal Error Rate.""" + n_scores = target_scores.size + nontarget_scores.size + all_scores = np.concatenate((target_scores, nontarget_scores)) + labels = np.concatenate( + (np.ones(target_scores.size), np.zeros(nontarget_scores.size)) + ) + indices = np.argsort(all_scores, kind="mergesort") + labels = labels[indices] + tar_trial_sums = np.cumsum(labels) + nontarget_trial_sums = nontarget_scores.size - ( + np.arange(1, n_scores + 1) - tar_trial_sums + ) + frr = np.concatenate((np.atleast_1d(0), tar_trial_sums / target_scores.size)) + far = np.concatenate( + ( + np.atleast_1d(1), + nontarget_trial_sums / nontarget_scores.size, + ) + ) + abs_diffs = np.abs(frr - far) + min_index = np.argmin(abs_diffs) + eer = np.mean((frr[min_index], far[min_index])) + return eer + + +def main(): + weights_path = os.path.join(SERVICE_DIR, "weights", "hfg_aug_1_2.pt") + os.makedirs(os.path.join(SERVICE_DIR, "models"), exist_ok=True) + + print("=" * 70) + print("ShiftySpeech (SSL-AASIST) - DeepSafe Dataset Evaluation") + print("=" * 70) + print(f"Weights: {weights_path}") + print(f"Dataset: {DATASET_DIR}") + print(f"Device: cpu") + print() + + # Load model + print("Loading model...") + start = time.time() + args_ns = argparse.Namespace() + ssl_model = SSLAASISTModel(args_ns, "cpu") + state_dict = torch.load(weights_path, map_location="cpu", weights_only=False) + ssl_model.load_state_dict(state_dict) + ssl_model.eval() + print(f"Model loaded in {time.time() - start:.1f}s") + print() + + # Collect audio files + real_dir = os.path.join(DATASET_DIR, "real") + fake_dir = os.path.join(DATASET_DIR, "fake") + + files = [] + for fname in sorted(os.listdir(real_dir)): + if fname.endswith(".wav"): + files.append((os.path.join(real_dir, fname), 0, fname)) + for fname in sorted(os.listdir(fake_dir)): + if fname.endswith(".wav"): + files.append((os.path.join(fake_dir, fname), 1, fname)) + + n_real = sum(1 for _, label, _ in files if label == 0) + n_fake = sum(1 for _, label, _ in files if label == 1) + print(f"Total files: {len(files)} (real: {n_real}, fake: {n_fake})") + print() + + # Run inference + results = [] + total_time = 0.0 + + print( + f"{'File':<20} {'True':>5} {'Pred':>5} {'P(fake)':>8} " + f"{'P(real)':>8} {'Time':>6}" + ) + print("-" * 60) + + for path, true_label, fname in files: + audio, sr = librosa.load(path, sr=SAMPLE_RATE, mono=True) + audio = pad_audio(audio) + x = torch.FloatTensor(audio).unsqueeze(0) + + t0 = time.time() + with torch.no_grad(): + out = ssl_model(x) + elapsed = time.time() - t0 + total_time += elapsed + + probs = torch.softmax(out, dim=1) + p_fake = probs[0, 0].item() + p_real = probs[0, 1].item() + pred = 1 if p_fake >= 0.5 else 0 + + results.append( + { + "file": fname, + "true_label": true_label, + "pred_label": pred, + "prob_fake": p_fake, + "prob_real": p_real, + } + ) + + true_str = "FAKE" if true_label == 1 else "REAL" + pred_str = "FAKE" if pred == 1 else "REAL" + correct = "ok" if pred == true_label else "XX" + print( + f"{fname:<20} {true_str:>5} {pred_str:>5} " + f"{p_fake:>8.4f} {p_real:>8.4f} {elapsed:>5.2f}s " + f"[{correct}]" + ) + + print() + print("=" * 70) + print("METRICS") + print("=" * 70) + + # Compute metrics + true_labels = np.array([r["true_label"] for r in results]) + pred_labels = np.array([r["pred_label"] for r in results]) + + tp = int(np.sum((pred_labels == 1) & (true_labels == 1))) + tn = int(np.sum((pred_labels == 0) & (true_labels == 0))) + fp = int(np.sum((pred_labels == 1) & (true_labels == 0))) + fn = int(np.sum((pred_labels == 0) & (true_labels == 1))) + + accuracy = (tp + tn) / len(results) if len(results) > 0 else 0 + precision = tp / (tp + fp) if (tp + fp) > 0 else 0 + recall = tp / (tp + fn) if (tp + fn) > 0 else 0 + f1 = ( + 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0 + ) + specificity = tn / (tn + fp) if (tn + fp) > 0 else 0 + + # EER using bonafide scores (prob_real: higher = more real) + bonafide_scores = np.array( + [r["prob_real"] for r in results if r["true_label"] == 0] + ) + spoof_scores = np.array([r["prob_real"] for r in results if r["true_label"] == 1]) + if len(bonafide_scores) > 0 and len(spoof_scores) > 0: + eer = compute_eer(bonafide_scores, spoof_scores) + else: + eer = float("nan") + + print(f"Accuracy: {accuracy:.4f} ({accuracy * 100:.1f}%)") + print(f"Precision: {precision:.4f}") + print(f"Recall: {recall:.4f}") + print(f"F1 Score: {f1:.4f}") + print(f"Specificity: {specificity:.4f}") + print(f"EER: {eer:.4f} ({eer * 100:.1f}%)") + print() + print(f"Confusion Matrix:") + print(f" TP={tp:>3d} FP={fp:>3d}") + print(f" FN={fn:>3d} TN={tn:>3d}") + print() + print(f"Total inference time: {total_time:.1f}s") + print(f"Avg per file: {total_time / len(results):.3f}s") + print(f"Total files: {len(results)}") + + +if __name__ == "__main__": + main() diff --git a/audio/shiftyspeech/requirements.txt b/audio/shiftyspeech/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..6e888dbd201cd7df0c63ebc4ec044044266ec90d --- /dev/null +++ b/audio/shiftyspeech/requirements.txt @@ -0,0 +1,14 @@ +fastapi +uvicorn +python-multipart +torch==2.5.1 +torchaudio==2.5.1 +numpy==1.23.5 +scipy +librosa==0.9.1 +soundfile +pydantic +omegaconf +hydra-core +scikit-learn +bitarray diff --git a/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/.env b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/.env new file mode 100644 index 0000000000000000000000000000000000000000..3d3f446144d3d3463b50d3f7a6946eda40b00b51 --- /dev/null +++ b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/.env @@ -0,0 +1,2 @@ +WANDB_API_KEY="" +WANDB_PROJECT_NAME="SSL-AASIST" \ No newline at end of file diff --git a/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/LICENSE b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..24853134ee78c4a7d7c21398fd5546193b8ffb43 --- /dev/null +++ b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2022 Hemlata + +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/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/RawBoost.py b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/RawBoost.py new file mode 100644 index 0000000000000000000000000000000000000000..2c48cdc299024b4e1227e6392dd530edb4cda159 --- /dev/null +++ b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/RawBoost.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +import copy + +import numpy as np +from scipy import signal + +""" + Hemlata Tak, Madhu Kamble, Jose Patino, Massimiliano Todisco, Nicholas Evans. + RawBoost: A Raw Data Boosting and Augmentation Method applied to Automatic Speaker Verification Anti-Spoofing. + In Proc. ICASSP 2022, pp:6382--6386. +""" + + +def randRange(x1, x2, integer): + y = np.random.uniform(low=x1, high=x2, size=(1,)) + if integer: + y = int(y) + return y + + +def normWav(x, always): + if always: + x = x / np.amax(abs(x)) + elif np.amax(abs(x)) > 1: + x = x / np.amax(abs(x)) + return x + + +def genNotchCoeffs( + nBands, minF, maxF, minBW, maxBW, minCoeff, maxCoeff, minG, maxG, fs +): + b = 1 + for i in range(0, nBands): + fc = randRange(minF, maxF, 0) + bw = randRange(minBW, maxBW, 0) + c = randRange(minCoeff, maxCoeff, 1) + + if c / 2 == int(c / 2): + c = c + 1 + f1 = fc - bw / 2 + f2 = fc + bw / 2 + if f1 <= 0: + f1 = 1 / 1000 + if f2 >= fs / 2: + f2 = fs / 2 - 1 / 1000 + b = np.convolve( + signal.firwin(c, [float(f1), float(f2)], window="hamming", fs=fs), b + ) + + G = randRange(minG, maxG, 0) + _, h = signal.freqz(b, 1, fs=fs) + b = pow(10, G / 20) * b / np.amax(abs(h)) + return b + + +def filterFIR(x, b): + N = b.shape[0] + 1 + xpad = np.pad(x, (0, N), "constant") + y = signal.lfilter(b, 1, xpad) + y = y[int(N / 2) : int(y.shape[0] - N / 2)] + return y + + +# Linear and non-linear convolutive noise +def LnL_convolutive_noise( + x, + N_f, + nBands, + minF, + maxF, + minBW, + maxBW, + minCoeff, + maxCoeff, + minG, + maxG, + minBiasLinNonLin, + maxBiasLinNonLin, + fs, +): + y = [0] * x.shape[0] + for i in range(0, N_f): + if i == 1: + minG = minG - minBiasLinNonLin + maxG = maxG - maxBiasLinNonLin + b = genNotchCoeffs( + nBands, minF, maxF, minBW, maxBW, minCoeff, maxCoeff, minG, maxG, fs + ) + y = y + filterFIR(np.power(x, (i + 1)), b) + y = y - np.mean(y) + y = normWav(y, 0) + return y + + +# Impulsive signal dependent noise +def ISD_additive_noise(x, P, g_sd): + beta = randRange(0, P, 0) + + y = copy.deepcopy(x) + x_len = x.shape[0] + n = int(x_len * (beta / 100)) + p = np.random.permutation(x_len)[:n] + f_r = np.multiply( + ((2 * np.random.rand(p.shape[0])) - 1), ((2 * np.random.rand(p.shape[0])) - 1) + ) + r = g_sd * x[p] * f_r + y[p] = x[p] + r + y = normWav(y, 0) + return y + + +# Stationary signal independent noise + + +def SSI_additive_noise( + x, + SNRmin, + SNRmax, + nBands, + minF, + maxF, + minBW, + maxBW, + minCoeff, + maxCoeff, + minG, + maxG, + fs, +): + noise = np.random.normal(0, 1, x.shape[0]) + b = genNotchCoeffs( + nBands, minF, maxF, minBW, maxBW, minCoeff, maxCoeff, minG, maxG, fs + ) + noise = filterFIR(noise, b) + noise = normWav(noise, 1) + SNR = randRange(SNRmin, SNRmax, 0) + noise = ( + noise / np.linalg.norm(noise, 2) * np.linalg.norm(x, 2) / 10.0 ** (0.05 * SNR) + ) + x = x + noise + return x diff --git a/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/Simplified_CM_solution.py b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/Simplified_CM_solution.py new file mode 100644 index 0000000000000000000000000000000000000000..cf247e1e0ab8c1781b743dbee28302c71a880301 --- /dev/null +++ b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/Simplified_CM_solution.py @@ -0,0 +1,227 @@ +import math +from collections import OrderedDict + +import fairseq +import numpy as np +import scipy.io as sio +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.autograd import Variable +from torch.nn.parameter import Parameter +from torch.utils import data + +___author__ = "Hemlata Tak" +__email__ = "tak@eurecom.fr" + +# from losses_anti_spoofing import AMSoftmax + +############################ +## FOR fine-tuning SSL MODEL +############################ + + +class SSLModel(nn.Module): + def __init__(self, device): + super(SSLModel, self).__init__() + + cp_path = "/change_to_path_to_pre_trained_model_XLR_300M/xlsr2_300m.pt" + model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task( + [cp_path] + ) + self.model = model[0] + self.device = device + self.out_dim = 1024 + return + + def extract_feat(self, input_data): + + # put the model to GPU if it not there + if ( + next(self.model.parameters()).device != input_data.device + or next(self.model.parameters()).dtype != input_data.dtype + ): + self.model.to(input_data.device, dtype=input_data.dtype) + self.model.train() + + if True: + # input should be in shape (batch, length) + if input_data.ndim == 3: + input_tmp = input_data[:, :, 0] + else: + input_tmp = input_data + + # [batch, length, dim] + emb = self.model(input_tmp, mask=False, features_only=True)["x"] + return emb + + +# ---------Graph attention simple back-end------------------------# +""" + Hemlata Tak, Jee-weon Jung, Jose Patino, Madhu Kamble, Massimiliano Todisco, Nicholas Evans. + End-to-end spectro-temporal graph attention networks for speaker verification anti-spoofing and speech deepfake detection. + In Proc. Automatic Speaker Verification and Spoofing Countermeasures Challenge 2021 Interspeech 2021 satellite workshop. +""" + + +class GraphAttentionLayer(nn.Module): + def __init__(self, in_dim, out_dim, **kwargs): + super(GraphAttentionLayer, self).__init__() + + # attention map + self.att_proj = nn.Linear(in_dim, out_dim) + self.att_weight = self._init_new_params(out_dim, 1) + + # project + self.proj_with_att = nn.Linear(in_dim, out_dim) + self.proj_without_att = nn.Linear(in_dim, out_dim) + + # batch norm + self.bn = nn.BatchNorm1d(out_dim) + + # dropout for inputs + self.input_drop = nn.Dropout(p=0.2) + + self.act = nn.SELU(inplace=True) + + def forward(self, x): + """ + x :(#bs, #node, #dim) + """ + # apply input dropout + x = self.input_drop(x) + + # derive attention map + att_map = self._derive_att_map(x) + + # projection + x = self._project(x, att_map) + + # apply batch norm + x = self._apply_BN(x) + x = self.act(x) + + return x + + def _pairwise_mul_nodes(self, x): + """ + Calculates pairwise multiplication of nodes. + - for attention map + x :(#bs, #node, #dim) + out_shape :(#bs, #node, #node, #dim) + """ + + nb_nodes = x.size(1) + x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1) + x_mirror = x.transpose(1, 2) + + return x * x_mirror + + def _derive_att_map(self, x): + """ + x :(#bs, #node, #dim) + out_shape :(#bs, #node, #node, 1) + """ + att_map = self._pairwise_mul_nodes(x) + att_map = torch.tanh( + self.att_proj(att_map) + ) # size: (#bs, #node, #node, #dim_out) + att_map = torch.matmul(att_map, self.att_weight) # size: (#bs, #node, #node, 1) + att_map = F.softmax(att_map, dim=-2) + + return att_map + + def _project(self, x, att_map): + x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x)) + x2 = self.proj_without_att(x) + + return x1 + x2 + + def _apply_BN(self, x): + org_size = x.size() + x = x.view(-1, org_size[-1]) + x = self.bn(x) + x = x.view(org_size) + + return x + + def _init_new_params(self, *size): + out = nn.Parameter(torch.FloatTensor(*size)) + nn.init.xavier_normal_(out) + return out + + +class GraphPool(nn.Module): + def __init__(self, k: float, in_dim: int, p): + super().__init__() + self.k = k + self.sigmoid = nn.Sigmoid() + self.proj = nn.Linear(in_dim, 1) + self.drop = nn.Dropout(p=p) if p > 0 else nn.Identity() + self.in_dim = in_dim + + def forward(self, h): + Z = self.drop(h) + weights = self.proj(Z) + scores = self.sigmoid(weights) + new_h = self.top_k_graph(scores, h, self.k) + + return new_h + + def top_k_graph(self, scores, h, k): + """ + args + ===== + scores: attention-based weights (#bs, #node, 1) + h: graph data (#bs, #node, #dim) + k: ratio of remaining nodes, (float) + returns + ===== + h: graph pool applied data (#bs, #node', #dim) + """ + _, n_nodes, n_feat = h.size() + n_nodes = max(int(n_nodes * k), 1) + _, idx = torch.topk(scores, n_nodes, dim=1) + idx = idx.expand(-1, -1, n_feat) + + h = h * scores + h = torch.gather(h, 1, idx) + + return h + + +class Model(nn.Module): + def __init__(self, d_args, device): + super(Model, self).__init__() + + # SSL model + self.device = device + self.ssl_model = SSLModel(self.device) + self.LL = nn.Linear(self.ssl_model.out_dim, 128) + self.first_bn = nn.BatchNorm1d(num_features=128) + self.selu = nn.SELU(inplace=True) + + # graph module layer + self.GAT_layer = GraphAttentionLayer(128, 64) + self.proj = nn.Linear(64, 1) + self.pool = GraphPool(0.8, 64, 0.3) + + # classifier head + self.proj_node = nn.Linear(53, 2) + + def forward(self, x_inp, Freq_aug=False): + # SSL wav2vec 2.0 model + x_ssl_feat = self.ssl_model.extract_feat(x_inp.squeeze(-1)) + x_SSL = self.LL(x_ssl_feat) # (bs,frame_number,feat_out_dim) + x_SSL = x_SSL.transpose(1, 2) # (bs,feat_out_dim,frame_number) + + x = F.max_pool1d(x_SSL, (3)) + x = self.first_bn(x) + x = self.selu(x) + + x = self.GAT_layer(x.transpose(1, 2)) + x = self.pool(x) + x = self.proj(x).flatten(1) + output = self.proj_node(x) + return output diff --git a/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/data_utils.py b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/data_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5ca02a8d1d7cafdea37a291d568dd39ccf6ca3f6 --- /dev/null +++ b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/data_utils.py @@ -0,0 +1,292 @@ +import os +import random +from random import randrange + +import librosa +import numpy as np +import torch +import torch.nn as nn +from RawBoost import ( + ISD_additive_noise, + LnL_convolutive_noise, + SSI_additive_noise, + normWav, +) +from torch import Tensor +from torch.utils.data import Dataset + +__author__ = "Hemlata Tak" +__email__ = "tak@eurecom.fr" + + +def genSpoof_list(dir_meta, is_train=False, is_eval=False): + d_meta = {} + file_list = [] + with open(dir_meta, "r") as f: + l_meta = f.readlines() + + if is_train: + for line in l_meta: + key, label = line.strip().split() + file_list.append(key) + d_meta[key] = 1 if label == "bonafide" else 0 + return d_meta, file_list + + elif is_eval: + for line in l_meta: + key, _ = line.strip().split(" ") + file_list.append(key) + return file_list + else: + for line in l_meta: + key, label = line.strip().split() + file_list.append(key) + d_meta[key] = 1 if label == "bonafide" else 0 + return d_meta, file_list + + +def pad(x, max_len=64600): + x_len = x.shape[0] + if x_len >= max_len: + return x[:max_len] + # need to pad + num_repeats = int(max_len / x_len) + 1 + padded_x = np.tile(x, (1, num_repeats))[:, :max_len][0] + return padded_x + + +class Dataset_ASVspoof2019_train(Dataset): + def __init__(self, args, metafile, algo): + """self.list_IDs : list of strings (each string: utt key), + self.labels: dictionary (key: utt key, value: label integer)""" + + self.uttpath_labels = [] + with open(metafile, "r") as f: + for line in f: + items = line.strip().split() + lb = 1 if items[-1] == "bonafide" else 0 + self.uttpath_labels.append((items[0], lb)) + + self.algo = algo + self.args = args + self.cut = 64600 # take ~4 sec audio (64600 samples) + + def __len__(self): + return len(self.uttpath_labels) + + def __getitem__(self, index): + path, target = self.uttpath_labels[index] + X, fs = librosa.load(path, sr=16000) + Y = process_Rawboost_feature(X, fs, self.args, self.algo) + X_pad = pad(Y, self.cut) + x_inp = Tensor(X_pad) + return x_inp, target + + +class Dataset_ASVspoof2021_eval(Dataset): + def __init__(self, list_IDs): + """self.list_IDs : list of strings (each string: utt key),""" + + self.list_IDs = list_IDs + self.cut = 64600 # take ~4 sec audio (64600 samples) + + def __len__(self): + return len(self.list_IDs) + + def __getitem__(self, index): + utt_id = self.list_IDs[index] + X, fs = librosa.load(utt_id, sr=16000) + X_pad = pad(X, self.cut) + x_inp = Tensor(X_pad) + return x_inp, utt_id + + +# --------------RawBoost data augmentation algorithms---------------------------## +def process_Rawboost_feature(feature, sr, args, algo): + + # Data process by Convolutive noise (1st algo) + if algo == 1: + + feature = LnL_convolutive_noise( + feature, + args.N_f, + args.nBands, + args.minF, + args.maxF, + args.minBW, + args.maxBW, + args.minCoeff, + args.maxCoeff, + args.minG, + args.maxG, + args.minBiasLinNonLin, + args.maxBiasLinNonLin, + sr, + ) + + # Data process by Impulsive noise (2nd algo) + elif algo == 2: + + feature = ISD_additive_noise(feature, args.P, args.g_sd) + + # Data process by coloured additive noise (3rd algo) + elif algo == 3: + + feature = SSI_additive_noise( + feature, + args.SNRmin, + args.SNRmax, + args.nBands, + args.minF, + args.maxF, + args.minBW, + args.maxBW, + args.minCoeff, + args.maxCoeff, + args.minG, + args.maxG, + sr, + ) + + # Data process by all 3 algo. together in series (1+2+3) + elif algo == 4: + + feature = LnL_convolutive_noise( + feature, + args.N_f, + args.nBands, + args.minF, + args.maxF, + args.minBW, + args.maxBW, + args.minCoeff, + args.maxCoeff, + args.minG, + args.maxG, + args.minBiasLinNonLin, + args.maxBiasLinNonLin, + sr, + ) + feature = ISD_additive_noise(feature, args.P, args.g_sd) + feature = SSI_additive_noise( + feature, + args.SNRmin, + args.SNRmax, + args.nBands, + args.minF, + args.maxF, + args.minBW, + args.maxBW, + args.minCoeff, + args.maxCoeff, + args.minG, + args.maxG, + sr, + ) + + # Data process by 1st two algo. together in series (1+2) + elif algo == 5: + + feature = LnL_convolutive_noise( + feature, + args.N_f, + args.nBands, + args.minF, + args.maxF, + args.minBW, + args.maxBW, + args.minCoeff, + args.maxCoeff, + args.minG, + args.maxG, + args.minBiasLinNonLin, + args.maxBiasLinNonLin, + sr, + ) + feature = ISD_additive_noise(feature, args.P, args.g_sd) + + # Data process by 1st and 3rd algo. together in series (1+3) + elif algo == 6: + + feature = LnL_convolutive_noise( + feature, + args.N_f, + args.nBands, + args.minF, + args.maxF, + args.minBW, + args.maxBW, + args.minCoeff, + args.maxCoeff, + args.minG, + args.maxG, + args.minBiasLinNonLin, + args.maxBiasLinNonLin, + sr, + ) + feature = SSI_additive_noise( + feature, + args.SNRmin, + args.SNRmax, + args.nBands, + args.minF, + args.maxF, + args.minBW, + args.maxBW, + args.minCoeff, + args.maxCoeff, + args.minG, + args.maxG, + sr, + ) + + # Data process by 2nd and 3rd algo. together in series (2+3) + elif algo == 7: + + feature = ISD_additive_noise(feature, args.P, args.g_sd) + feature = SSI_additive_noise( + feature, + args.SNRmin, + args.SNRmax, + args.nBands, + args.minF, + args.maxF, + args.minBW, + args.maxBW, + args.minCoeff, + args.maxCoeff, + args.minG, + args.maxG, + sr, + ) + + # Data process by 1st two algo. together in Parallel (1||2) + elif algo == 8: + + feature1 = LnL_convolutive_noise( + feature, + args.N_f, + args.nBands, + args.minF, + args.maxF, + args.minBW, + args.maxBW, + args.minCoeff, + args.maxCoeff, + args.minG, + args.maxG, + args.minBiasLinNonLin, + args.maxBiasLinNonLin, + sr, + ) + feature2 = ISD_additive_noise(feature, args.P, args.g_sd) + + feature_para = feature1 + feature2 + feature = normWav(feature_para, 0) # normalized resultant waveform + + # original data without Rawboost processing + else: + + feature = feature + + return feature diff --git a/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/model.py b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/model.py new file mode 100644 index 0000000000000000000000000000000000000000..8a84f2122255f07e858fcc9d97f643d7da3bce2d --- /dev/null +++ b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/model.py @@ -0,0 +1,603 @@ +import random +from typing import Union + +import fairseq +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + +___author__ = "Hemlata Tak" +__email__ = "tak@eurecom.fr" + +############################ +## FOR fine-tuned SSL MODEL +############################ + + +class SSLModel(nn.Module): + def __init__(self, device): + super(SSLModel, self).__init__() + + cp_path = "models/xlsr2_300m.pt" # Change the pre-trained XLSR model path. + model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task( + [cp_path] + ) + self.model = model[0] + self.device = device + self.out_dim = 1024 + return + + def extract_feat(self, input_data): + + # put the model to GPU if it not there + if ( + next(self.model.parameters()).device != input_data.device + or next(self.model.parameters()).dtype != input_data.dtype + ): + self.model.to(input_data.device, dtype=input_data.dtype) + self.model.train() + + if True: + # input should be in shape (batch, length) + if input_data.ndim == 3: + input_tmp = input_data[:, :, 0] + else: + input_tmp = input_data + + # [batch, length, dim] + emb = self.model(input_tmp, mask=False, features_only=True)["x"] + return emb + + +# ---------AASIST back-end------------------------# +""" Jee-weon Jung, Hee-Soo Heo, Hemlata Tak, Hye-jin Shim, Joon Son Chung, Bong-Jin Lee, Ha-Jin Yu and Nicholas Evans. + AASIST: Audio Anti-Spoofing Using Integrated Spectro-Temporal Graph Attention Networks. + In Proc. ICASSP 2022, pp: 6367--6371.""" + + +class GraphAttentionLayer(nn.Module): + def __init__(self, in_dim, out_dim, **kwargs): + super().__init__() + + # attention map + self.att_proj = nn.Linear(in_dim, out_dim) + self.att_weight = self._init_new_params(out_dim, 1) + + # project + self.proj_with_att = nn.Linear(in_dim, out_dim) + self.proj_without_att = nn.Linear(in_dim, out_dim) + + # batch norm + self.bn = nn.BatchNorm1d(out_dim) + + # dropout for inputs + self.input_drop = nn.Dropout(p=0.2) + + # activate + self.act = nn.SELU(inplace=True) + + # temperature + self.temp = 1.0 + if "temperature" in kwargs: + self.temp = kwargs["temperature"] + + def forward(self, x): + """ + x :(#bs, #node, #dim) + """ + # apply input dropout + x = self.input_drop(x) + + # derive attention map + att_map = self._derive_att_map(x) + + # projection + x = self._project(x, att_map) + + # apply batch norm + x = self._apply_BN(x) + x = self.act(x) + return x + + def _pairwise_mul_nodes(self, x): + """ + Calculates pairwise multiplication of nodes. + - for attention map + x :(#bs, #node, #dim) + out_shape :(#bs, #node, #node, #dim) + """ + + nb_nodes = x.size(1) + x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1) + x_mirror = x.transpose(1, 2) + + return x * x_mirror + + def _derive_att_map(self, x): + """ + x :(#bs, #node, #dim) + out_shape :(#bs, #node, #node, 1) + """ + att_map = self._pairwise_mul_nodes(x) + # size: (#bs, #node, #node, #dim_out) + att_map = torch.tanh(self.att_proj(att_map)) + # size: (#bs, #node, #node, 1) + att_map = torch.matmul(att_map, self.att_weight) + + # apply temperature + att_map = att_map / self.temp + + att_map = F.softmax(att_map, dim=-2) + + return att_map + + def _project(self, x, att_map): + x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x)) + x2 = self.proj_without_att(x) + + return x1 + x2 + + def _apply_BN(self, x): + org_size = x.size() + x = x.view(-1, org_size[-1]) + x = self.bn(x) + x = x.view(org_size) + + return x + + def _init_new_params(self, *size): + out = nn.Parameter(torch.FloatTensor(*size)) + nn.init.xavier_normal_(out) + return out + + +class HtrgGraphAttentionLayer(nn.Module): + def __init__(self, in_dim, out_dim, **kwargs): + super().__init__() + + self.proj_type1 = nn.Linear(in_dim, in_dim) + self.proj_type2 = nn.Linear(in_dim, in_dim) + + # attention map + self.att_proj = nn.Linear(in_dim, out_dim) + self.att_projM = nn.Linear(in_dim, out_dim) + + self.att_weight11 = self._init_new_params(out_dim, 1) + self.att_weight22 = self._init_new_params(out_dim, 1) + self.att_weight12 = self._init_new_params(out_dim, 1) + self.att_weightM = self._init_new_params(out_dim, 1) + + # project + self.proj_with_att = nn.Linear(in_dim, out_dim) + self.proj_without_att = nn.Linear(in_dim, out_dim) + + self.proj_with_attM = nn.Linear(in_dim, out_dim) + self.proj_without_attM = nn.Linear(in_dim, out_dim) + + # batch norm + self.bn = nn.BatchNorm1d(out_dim) + + # dropout for inputs + self.input_drop = nn.Dropout(p=0.2) + + # activate + self.act = nn.SELU(inplace=True) + + # temperature + self.temp = 1.0 + if "temperature" in kwargs: + self.temp = kwargs["temperature"] + + def forward(self, x1, x2, master=None): + """ + x1 :(#bs, #node, #dim) + x2 :(#bs, #node, #dim) + """ + + num_type1 = x1.size(1) + num_type2 = x2.size(1) + + x1 = self.proj_type1(x1) + + x2 = self.proj_type2(x2) + + x = torch.cat([x1, x2], dim=1) + + if master is None: + master = torch.mean(x, dim=1, keepdim=True) + + # apply input dropout + x = self.input_drop(x) + + # derive attention map + att_map = self._derive_att_map(x, num_type1, num_type2) + + # directional edge for master node + master = self._update_master(x, master) + + # projection + x = self._project(x, att_map) + + # apply batch norm + x = self._apply_BN(x) + x = self.act(x) + + x1 = x.narrow(1, 0, num_type1) + + x2 = x.narrow(1, num_type1, num_type2) + + return x1, x2, master + + def _update_master(self, x, master): + + att_map = self._derive_att_map_master(x, master) + master = self._project_master(x, master, att_map) + + return master + + def _pairwise_mul_nodes(self, x): + """ + Calculates pairwise multiplication of nodes. + - for attention map + x :(#bs, #node, #dim) + out_shape :(#bs, #node, #node, #dim) + """ + + nb_nodes = x.size(1) + x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1) + x_mirror = x.transpose(1, 2) + + return x * x_mirror + + def _derive_att_map_master(self, x, master): + """ + x :(#bs, #node, #dim) + out_shape :(#bs, #node, #node, 1) + """ + att_map = x * master + att_map = torch.tanh(self.att_projM(att_map)) + + att_map = torch.matmul(att_map, self.att_weightM) + + # apply temperature + att_map = att_map / self.temp + + att_map = F.softmax(att_map, dim=-2) + + return att_map + + def _derive_att_map(self, x, num_type1, num_type2): + """ + x :(#bs, #node, #dim) + out_shape :(#bs, #node, #node, 1) + """ + att_map = self._pairwise_mul_nodes(x) + # size: (#bs, #node, #node, #dim_out) + att_map = torch.tanh(self.att_proj(att_map)) + # size: (#bs, #node, #node, 1) + + att_board = torch.zeros_like(att_map[:, :, :, 0]).unsqueeze(-1) + + att_board[:, :num_type1, :num_type1, :] = torch.matmul( + att_map[:, :num_type1, :num_type1, :], self.att_weight11 + ) + att_board[:, num_type1:, num_type1:, :] = torch.matmul( + att_map[:, num_type1:, num_type1:, :], self.att_weight22 + ) + att_board[:, :num_type1, num_type1:, :] = torch.matmul( + att_map[:, :num_type1, num_type1:, :], self.att_weight12 + ) + att_board[:, num_type1:, :num_type1, :] = torch.matmul( + att_map[:, num_type1:, :num_type1, :], self.att_weight12 + ) + + att_map = att_board + + # apply temperature + att_map = att_map / self.temp + + att_map = F.softmax(att_map, dim=-2) + + return att_map + + def _project(self, x, att_map): + x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x)) + x2 = self.proj_without_att(x) + + return x1 + x2 + + def _project_master(self, x, master, att_map): + + x1 = self.proj_with_attM(torch.matmul(att_map.squeeze(-1).unsqueeze(1), x)) + x2 = self.proj_without_attM(master) + + return x1 + x2 + + def _apply_BN(self, x): + org_size = x.size() + x = x.view(-1, org_size[-1]) + x = self.bn(x) + x = x.view(org_size) + + return x + + def _init_new_params(self, *size): + out = nn.Parameter(torch.FloatTensor(*size)) + nn.init.xavier_normal_(out) + return out + + +class GraphPool(nn.Module): + def __init__(self, k: float, in_dim: int, p: Union[float, int]): + super().__init__() + self.k = k + self.sigmoid = nn.Sigmoid() + self.proj = nn.Linear(in_dim, 1) + self.drop = nn.Dropout(p=p) if p > 0 else nn.Identity() + self.in_dim = in_dim + + def forward(self, h): + Z = self.drop(h) + weights = self.proj(Z) + scores = self.sigmoid(weights) + new_h = self.top_k_graph(scores, h, self.k) + + return new_h + + def top_k_graph(self, scores, h, k): + """ + args + ===== + scores: attention-based weights (#bs, #node, 1) + h: graph data (#bs, #node, #dim) + k: ratio of remaining nodes, (float) + returns + ===== + h: graph pool applied data (#bs, #node', #dim) + """ + _, n_nodes, n_feat = h.size() + n_nodes = max(int(n_nodes * k), 1) + _, idx = torch.topk(scores, n_nodes, dim=1) + idx = idx.expand(-1, -1, n_feat) + + h = h * scores + h = torch.gather(h, 1, idx) + + return h + + +class Residual_block(nn.Module): + def __init__(self, nb_filts, first=False): + super().__init__() + self.first = first + + if not self.first: + self.bn1 = nn.BatchNorm2d(num_features=nb_filts[0]) + self.conv1 = nn.Conv2d( + in_channels=nb_filts[0], + out_channels=nb_filts[1], + kernel_size=(2, 3), + padding=(1, 1), + stride=1, + ) + self.selu = nn.SELU(inplace=True) + + self.bn2 = nn.BatchNorm2d(num_features=nb_filts[1]) + self.conv2 = nn.Conv2d( + in_channels=nb_filts[1], + out_channels=nb_filts[1], + kernel_size=(2, 3), + padding=(0, 1), + stride=1, + ) + + if nb_filts[0] != nb_filts[1]: + self.downsample = True + self.conv_downsample = nn.Conv2d( + in_channels=nb_filts[0], + out_channels=nb_filts[1], + padding=(0, 1), + kernel_size=(1, 3), + stride=1, + ) + + else: + self.downsample = False + + def forward(self, x): + identity = x + if not self.first: + out = self.bn1(x) + out = self.selu(out) + else: + out = x + + out = self.conv1(x) + + out = self.bn2(out) + out = self.selu(out) + + out = self.conv2(out) + + if self.downsample: + identity = self.conv_downsample(identity) + + out += identity + + return out + + +class Model(nn.Module): + def __init__(self, args, device): + super().__init__() + self.device = device + + # AASIST parameters + filts = [128, [1, 32], [32, 32], [32, 64], [64, 64]] + gat_dims = [64, 32] + pool_ratios = [0.5, 0.5, 0.5, 0.5] + temperatures = [2.0, 2.0, 100.0, 100.0] + + #### + # create network wav2vec 2.0 + #### + self.ssl_model = SSLModel(self.device) + self.LL = nn.Linear(self.ssl_model.out_dim, 128) + + self.first_bn = nn.BatchNorm2d(num_features=1) + self.first_bn1 = nn.BatchNorm2d(num_features=64) + self.drop = nn.Dropout(0.5, inplace=True) + self.drop_way = nn.Dropout(0.2, inplace=True) + self.selu = nn.SELU(inplace=True) + + # RawNet2 encoder + self.encoder = nn.Sequential( + nn.Sequential(Residual_block(nb_filts=filts[1], first=True)), + nn.Sequential(Residual_block(nb_filts=filts[2])), + nn.Sequential(Residual_block(nb_filts=filts[3])), + nn.Sequential(Residual_block(nb_filts=filts[4])), + nn.Sequential(Residual_block(nb_filts=filts[4])), + nn.Sequential(Residual_block(nb_filts=filts[4])), + ) + + self.attention = nn.Sequential( + nn.Conv2d(64, 128, kernel_size=(1, 1)), + nn.SELU(inplace=True), + nn.BatchNorm2d(128), + nn.Conv2d(128, 64, kernel_size=(1, 1)), + ) + # position encoding + self.pos_S = nn.Parameter(torch.randn(1, 42, filts[-1][-1])) + + self.master1 = nn.Parameter(torch.randn(1, 1, gat_dims[0])) + self.master2 = nn.Parameter(torch.randn(1, 1, gat_dims[0])) + + # Graph module + self.GAT_layer_S = GraphAttentionLayer( + filts[-1][-1], gat_dims[0], temperature=temperatures[0] + ) + self.GAT_layer_T = GraphAttentionLayer( + filts[-1][-1], gat_dims[0], temperature=temperatures[1] + ) + # HS-GAL layer + self.HtrgGAT_layer_ST11 = HtrgGraphAttentionLayer( + gat_dims[0], gat_dims[1], temperature=temperatures[2] + ) + self.HtrgGAT_layer_ST12 = HtrgGraphAttentionLayer( + gat_dims[1], gat_dims[1], temperature=temperatures[2] + ) + self.HtrgGAT_layer_ST21 = HtrgGraphAttentionLayer( + gat_dims[0], gat_dims[1], temperature=temperatures[2] + ) + self.HtrgGAT_layer_ST22 = HtrgGraphAttentionLayer( + gat_dims[1], gat_dims[1], temperature=temperatures[2] + ) + + # Graph pooling layers + self.pool_S = GraphPool(pool_ratios[0], gat_dims[0], 0.3) + self.pool_T = GraphPool(pool_ratios[1], gat_dims[0], 0.3) + self.pool_hS1 = GraphPool(pool_ratios[2], gat_dims[1], 0.3) + self.pool_hT1 = GraphPool(pool_ratios[2], gat_dims[1], 0.3) + + self.pool_hS2 = GraphPool(pool_ratios[2], gat_dims[1], 0.3) + self.pool_hT2 = GraphPool(pool_ratios[2], gat_dims[1], 0.3) + + self.out_layer = nn.Linear(5 * gat_dims[1], 2) + + def forward(self, x): + # -------pre-trained Wav2vec model fine tunning ------------------------## + x_ssl_feat = self.ssl_model.extract_feat(x.squeeze(-1)) + x = self.LL(x_ssl_feat) # (bs,frame_number,feat_out_dim) + + # post-processing on front-end features + x = x.transpose(1, 2) # (bs,feat_out_dim,frame_number) + x = x.unsqueeze(dim=1) # add channel + x = F.max_pool2d(x, (3, 3)) + x = self.first_bn(x) + x = self.selu(x) + + # RawNet2-based encoder + x = self.encoder(x) + x = self.first_bn1(x) + x = self.selu(x) + + w = self.attention(x) + + # ------------SA for spectral feature-------------# + w1 = F.softmax(w, dim=-1) + m = torch.sum(x * w1, dim=-1) + e_S = m.transpose(1, 2) + self.pos_S + + # graph module layer + gat_S = self.GAT_layer_S(e_S) + out_S = self.pool_S(gat_S) # (#bs, #node, #dim) + + # ------------SA for temporal feature-------------# + w2 = F.softmax(w, dim=-2) + m1 = torch.sum(x * w2, dim=-2) + + e_T = m1.transpose(1, 2) + + # graph module layer + gat_T = self.GAT_layer_T(e_T) + out_T = self.pool_T(gat_T) + + # learnable master node + master1 = self.master1.expand(x.size(0), -1, -1) + master2 = self.master2.expand(x.size(0), -1, -1) + + # inference 1 + out_T1, out_S1, master1 = self.HtrgGAT_layer_ST11( + out_T, out_S, master=self.master1 + ) + + out_S1 = self.pool_hS1(out_S1) + out_T1 = self.pool_hT1(out_T1) + + out_T_aug, out_S_aug, master_aug = self.HtrgGAT_layer_ST12( + out_T1, out_S1, master=master1 + ) + out_T1 = out_T1 + out_T_aug + out_S1 = out_S1 + out_S_aug + master1 = master1 + master_aug + + # inference 2 + out_T2, out_S2, master2 = self.HtrgGAT_layer_ST21( + out_T, out_S, master=self.master2 + ) + out_S2 = self.pool_hS2(out_S2) + out_T2 = self.pool_hT2(out_T2) + + out_T_aug, out_S_aug, master_aug = self.HtrgGAT_layer_ST22( + out_T2, out_S2, master=master2 + ) + out_T2 = out_T2 + out_T_aug + out_S2 = out_S2 + out_S_aug + master2 = master2 + master_aug + + out_T1 = self.drop_way(out_T1) + out_T2 = self.drop_way(out_T2) + out_S1 = self.drop_way(out_S1) + out_S2 = self.drop_way(out_S2) + master1 = self.drop_way(master1) + master2 = self.drop_way(master2) + + out_T = torch.max(out_T1, out_T2) + out_S = torch.max(out_S1, out_S2) + master = torch.max(master1, master2) + + # Readout operation + T_max, _ = torch.max(torch.abs(out_T), dim=1) + T_avg = torch.mean(out_T, dim=1) + + S_max, _ = torch.max(torch.abs(out_S), dim=1) + S_avg = torch.mean(out_S, dim=1) + + last_hidden = torch.cat([T_max, T_avg, S_max, S_avg, master.squeeze(1)], dim=1) + + last_hidden = self.drop(last_hidden) + output = self.out_layer(last_hidden) + + return output diff --git a/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/requirements.txt b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..8e66e4ec7072a3fe2a0592db49b9c3f352b37f73 --- /dev/null +++ b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/requirements.txt @@ -0,0 +1,4 @@ +librosa==0.9.1 +python-dotenv==1.0.1 +tensorboardX +wandb \ No newline at end of file diff --git a/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/startup_config.py b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/startup_config.py new file mode 100644 index 0000000000000000000000000000000000000000..b482bd43721e1dcb496577a5a39bcc229d2a7af2 --- /dev/null +++ b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/startup_config.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python +""" +startup_config + +Startup configuration utilities + +""" + +from __future__ import absolute_import + +import importlib +import os +import random +import sys + +import numpy as np +import torch + +__author__ = "Xin Wang" +__email__ = "wangxin@nii.ac.jp" +__copyright__ = "Copyright 2020, Xin Wang" + + +def set_random_seed(random_seed, args=None): + """set_random_seed(random_seed, args=None) + + Set the random_seed for numpy, python, and cudnn + + input + ----- + random_seed: integer random seed + args: argue parser + """ + + # initialization + torch.manual_seed(random_seed) + random.seed(random_seed) + np.random.seed(random_seed) + os.environ["PYTHONHASHSEED"] = str(random_seed) + + # For torch.backends.cudnn.deterministic + # Note: this default configuration may result in RuntimeError + # see https://pytorch.org/docs/stable/notes/randomness.html + if args is None: + cudnn_deterministic = True + cudnn_benchmark = False + else: + cudnn_deterministic = args.cudnn_deterministic_toggle + cudnn_benchmark = args.cudnn_benchmark_toggle + + if not cudnn_deterministic: + print("cudnn_deterministic set to False") + if cudnn_benchmark: + print("cudnn_benchmark set to True") + + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(random_seed) + torch.backends.cudnn.deterministic = cudnn_deterministic + torch.backends.cudnn.benchmark = cudnn_benchmark + return diff --git a/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/train.py b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/train.py new file mode 100644 index 0000000000000000000000000000000000000000..86eeb35bb700309939f1f26952b4943003f32261 --- /dev/null +++ b/audio/shiftyspeech/synthetic_speech_detection/SSL_Anti-spoofing/train.py @@ -0,0 +1,446 @@ +import argparse +import os +import sys + +import librosa +import numpy as np +import torch +import wandb +import yaml +from data_utils import ( + Dataset_ASVspoof2019_train, + Dataset_ASVspoof2021_eval, + genSpoof_list, + pad, + process_Rawboost_feature, +) +from dotenv import load_dotenv +from model import Model +from sklearn.metrics import roc_auc_score +from startup_config import set_random_seed +from tensorboardX import SummaryWriter +from torch import Tensor, nn +from torch.utils.data import DataLoader +from tqdm import tqdm + +__author__ = "Hemlata Tak" +__email__ = "tak@eurecom.fr" + + +def compute_det_curve(target_scores, nontarget_scores): + + n_scores = target_scores.size + nontarget_scores.size + all_scores = np.concatenate((target_scores, nontarget_scores)) + labels = np.concatenate( + (np.ones(target_scores.size), np.zeros(nontarget_scores.size)) + ) + + indices = np.argsort(all_scores, kind="mergesort") + labels = labels[indices] + tar_trial_sums = np.cumsum(labels) + nontarget_trial_sums = nontarget_scores.size - ( + np.arange(1, n_scores + 1) - tar_trial_sums + ) + + frr = np.concatenate((np.atleast_1d(0), tar_trial_sums / target_scores.size)) + far = np.concatenate( + (np.atleast_1d(1), nontarget_trial_sums / nontarget_scores.size) + ) + # Thresholds are the sorted scores + thresholds = np.concatenate( + (np.atleast_1d(all_scores[indices[0]] - 0.001), all_scores[indices]) + ) + + return frr, far, thresholds + + +def compute_eer(target_scores, nontarget_scores): + """Returns equal error rate (EER) and the corresponding threshold.""" + frr, far, thresholds = compute_det_curve(target_scores, nontarget_scores) + abs_diffs = np.abs(frr - far) + min_index = np.argmin(abs_diffs) + eer = np.mean((frr[min_index], far[min_index])) + return eer, thresholds[min_index], frr, far + + +def calculate_tDCF_EER(cm_scores_file, output_file, printout=True): + # Load CM scores + cm_data = np.genfromtxt(cm_scores_file, dtype=str) + cm_utt_id = cm_data[:, 0] + cm_keys = cm_data[:, 1] + cm_scores = cm_data[:, 2].astype(float) + # Extract bona fide (real human) and spoof scores from the CM scores + bona_cm = cm_scores[cm_keys == "bonafide"] + spoof_cm = cm_scores[cm_keys == "spoof"] + all_scores = np.concatenate([bona_cm, spoof_cm]) + all_true_labels = np.concatenate([np.ones_like(bona_cm), np.zeros_like(spoof_cm)]) + + auc = roc_auc_score(all_true_labels, all_scores, max_fpr=0.05) + eer_cm, eer_threshold, frr, far = compute_eer(bona_cm, spoof_cm) + + if printout: + with open(output_file, "w") as f_res: + f_res.write("\nCM SYSTEM\n") + f_res.write( + "\tEER\t\t= {:8.9f} % " + "(Equal error rate for countermeasure)\n".format(eer_cm * 100) + ) + f_res.write("\t pAUC with max fpr - 0.05 is :{}".format(auc)) + + +def evaluate_accuracy(dev_loader, model, device, args): + val_loss = 0.0 + num_total = 0.0 + algo = args.algo + cut = 64600 + model.eval() + + weight = torch.FloatTensor([0.1, 0.9]).to(device) + criterion = nn.CrossEntropyLoss(weight=weight) + progress_bar = tqdm(dev_loader, desc=f"Epoch {epoch+1}/{args.num_epochs}") + for current_step, (batch_pths, batch_y) in enumerate(progress_bar): + batch_x = batch_pths + batch_size = batch_x.size(0) + num_total += batch_size + batch_x = batch_x.to(device) + batch_y = batch_y.view(-1).type(torch.int64).to(device) + batch_out = model(batch_x) + + batch_loss = criterion(batch_out, batch_y) + val_loss += batch_loss.item() * batch_size + + val_loss /= num_total + + return val_loss + + +def produce_evaluation_file(dataset, model, device, save_path, trial_path): + data_loader = DataLoader(dataset, batch_size=10, shuffle=False, drop_last=False) + num_correct = 0.0 + num_total = 0.0 + model.eval() + with open(trial_path, "r") as f_trl: + trial_lines = f_trl.readlines() + + fname_list = [] + score_list = [] + + for batch_x, utt_id in data_loader: + + batch_size = batch_x.size(0) + batch_x = batch_x.to(device) + + batch_out = model(batch_x) + + batch_score = (batch_out[:, 1]).data.cpu().numpy().ravel() + # add outputs + fname_list.extend(utt_id) + score_list.extend(batch_score.tolist()) + assert len(trial_lines) == len(fname_list) == len(score_list) + + with open(save_path, "a+") as fh: + for fname, cm, trl in zip(fname_list, score_list, trial_lines): + utt_id, key = trl.strip().split(" ") + assert fname == utt_id + fh.write("{} {} {}\n".format(fname, key, cm)) + fh.close() + print("Scores saved to {}".format(save_path)) + + +def train_epoch(train_loader, model, lr, optim, device, args): + running_loss = 0 + + num_total = 0.0 + algo = args.algo + model.train() + cut = 64600 + # set objective (Loss) functions + weight = torch.FloatTensor([0.1, 0.9]).to(device) + criterion = nn.CrossEntropyLoss(weight=weight) + progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.num_epochs}") + for current_step, (batch_pths, batch_y) in enumerate(progress_bar): + batch_x = batch_pths + batch_size = batch_x.size(0) + num_total += batch_size + + batch_x = batch_x.to(device) + batch_y = batch_y.view(-1).type(torch.int64).to(device) + batch_out = model(batch_x) + + batch_loss = criterion(batch_out, batch_y) + + running_loss += batch_loss.item() * batch_size + + optimizer.zero_grad() + batch_loss.backward() + optimizer.step() + + running_loss /= num_total + + return running_loss + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="SSL-AASIST baseline system") + + # Hyperparameters + parser.add_argument("--batch_size", type=int, default=64) + parser.add_argument("--num_epochs", type=int, default=100) + parser.add_argument("--lr", type=float, default=0.000001) + parser.add_argument("--weight_decay", type=float, default=0.0001) + parser.add_argument("--model_name", type=str, default="SSL-AASIST") + parser.add_argument("--loss", type=str, default="weighted_CCE") + parser.add_argument("--trn_list_path", default=None, help="path to train file") + parser.add_argument("--dev_list_path", default=None, help="path to validation file") + parser.add_argument("--test_list_path", default=None, help="path to test file") + parser.add_argument( + "--test_score_dir", default=None, help="path to save test scores" + ) + # model + parser.add_argument( + "--seed", type=int, default=1234, help="random seed (default: 1234)" + ) + parser.add_argument("--save_path", type=str, default=".", help="Model save path") + parser.add_argument("--model_path", type=str, default=None, help="Model checkpoint") + parser.add_argument( + "--comment", type=str, default=None, help="Comment to describe the saved model" + ) + # Auxiliary arguments + + parser.add_argument("--eval", action="store_true", default=False, help="eval mode") + parser.add_argument("--eval_part", type=int, default=0) + # backend options + parser.add_argument( + "--cudnn-deterministic-toggle", + action="store_false", + default=True, + help="use cudnn-deterministic? (default true)", + ) + + parser.add_argument( + "--cudnn-benchmark-toggle", + action="store_true", + default=False, + help="use cudnn-benchmark? (default false)", + ) + + ##===================================================Rawboost data augmentation ======================================================================# + + parser.add_argument( + "--algo", + type=int, + default=5, + help="Rawboost algos discriptions. 0: No augmentation 1: LnL_convolutive_noise, 2: ISD_additive_noise, 3: SSI_additive_noise, 4: series algo (1+2+3), \ + 5: series algo (1+2), 6: series algo (1+3), 7: series algo(2+3), 8: parallel algo(1,2) .[default=0]", + ) + + # LnL_convolutive_noise parameters + parser.add_argument( + "--nBands", + type=int, + default=5, + help="number of notch filters.The higher the number of bands, the more aggresive the distortions is.[default=5]", + ) + parser.add_argument( + "--minF", + type=int, + default=20, + help="minimum centre frequency [Hz] of notch filter.[default=20] ", + ) + parser.add_argument( + "--maxF", + type=int, + default=8000, + help="maximum centre frequency [Hz] ( bytes: + """Generate a simple sine wave WAV file as bytes.""" + import soundfile as sf + + t = np.linspace(0, duration_s, int(sr * duration_s), endpoint=False) + audio = 0.5 * np.sin(2 * np.pi * 440 * t).astype(np.float32) + buf = io.BytesIO() + sf.write(buf, audio, sr, format="WAV") + buf.seek(0) + return buf.read() + + +def _encode_file(path: str) -> str: + """Read a file and return base64 encoded string.""" + with open(path, "rb") as f: + return base64.b64encode(f.read()).decode("utf-8") + + +class TestHealthEndpoint: + """Tests for the /health endpoint.""" + + def test_health_returns_200(self): + response = client.get("/health") + assert response.status_code == 200 + + def test_health_contains_model_name(self): + response = client.get("/health") + data = response.json() + assert data["model"] == "shiftyspeech" + + def test_health_contains_device(self): + response = client.get("/health") + data = response.json() + assert data["device"] == "cpu" + + def test_health_contains_status(self): + response = client.get("/health") + data = response.json() + assert data["status"] in ("healthy", "degraded") + + +class TestPreprocessAudio: + """Tests for audio preprocessing logic.""" + + def test_preprocess_short_audio_tiles(self): + """Short audio should be tiled to TARGET_SAMPLES.""" + wav_bytes = _make_wav_bytes(duration_s=0.5, sr=16000) + tensor = api.preprocess_audio(wav_bytes) + assert tensor.shape == (1, api.TARGET_SAMPLES) + + def test_preprocess_long_audio_trims(self): + """Long audio should be trimmed to TARGET_SAMPLES.""" + wav_bytes = _make_wav_bytes(duration_s=10.0, sr=16000) + tensor = api.preprocess_audio(wav_bytes) + assert tensor.shape == (1, api.TARGET_SAMPLES) + + def test_preprocess_exact_length(self): + """Audio at exact TARGET_SAMPLES should pass through.""" + duration = api.TARGET_SAMPLES / api.SAMPLE_RATE + wav_bytes = _make_wav_bytes(duration_s=duration, sr=16000) + tensor = api.preprocess_audio(wav_bytes) + assert tensor.shape == (1, api.TARGET_SAMPLES) + + def test_preprocess_resamples_from_8khz(self): + """Audio at 8kHz should be resampled to 16kHz.""" + wav_bytes = _make_wav_bytes(duration_s=1.0, sr=8000) + tensor = api.preprocess_audio(wav_bytes) + assert tensor.shape == (1, api.TARGET_SAMPLES) + + def test_preprocess_invalid_input_raises(self): + """Invalid audio bytes should raise ValueError.""" + with pytest.raises(ValueError): + api.preprocess_audio(b"not audio data") + + +@pytest.mark.skipif( + not WEIGHTS_AVAILABLE, + reason="Model weights not available locally", +) +class TestPredictEndpoint: + """Integration tests for the /predict endpoint (requires weights).""" + + def test_predict_returns_200(self): + wav_bytes = _make_wav_bytes(duration_s=2.0) + b64 = base64.b64encode(wav_bytes).decode("utf-8") + response = client.post("/predict", json={"audio_data": b64}) + assert response.status_code == 200 + + def test_predict_response_schema(self): + wav_bytes = _make_wav_bytes(duration_s=2.0) + b64 = base64.b64encode(wav_bytes).decode("utf-8") + response = client.post("/predict", json={"audio_data": b64}) + data = response.json() + assert "model" in data + assert "probability" in data + assert "prediction" in data + assert "class" in data + assert "inference_time" in data + assert data["model"] == "shiftyspeech" + + def test_predict_probability_in_range(self): + wav_bytes = _make_wav_bytes(duration_s=2.0) + b64 = base64.b64encode(wav_bytes).decode("utf-8") + response = client.post("/predict", json={"audio_data": b64}) + data = response.json() + assert 0.0 <= data["probability"] <= 1.0 + + def test_predict_class_matches_prediction(self): + wav_bytes = _make_wav_bytes(duration_s=2.0) + b64 = base64.b64encode(wav_bytes).decode("utf-8") + response = client.post("/predict", json={"audio_data": b64}) + data = response.json() + if data["prediction"] == 1: + assert data["class"] == "fake" + else: + assert data["class"] == "real" + + def test_predict_custom_threshold(self): + wav_bytes = _make_wav_bytes(duration_s=2.0) + b64 = base64.b64encode(wav_bytes).decode("utf-8") + response = client.post( + "/predict", + json={"audio_data": b64, "threshold": 0.99}, + ) + data = response.json() + assert response.status_code == 200 + # With threshold=0.99, only very high prob_fake => fake + if data["probability"] < 0.99: + assert data["prediction"] == 0 + assert data["class"] == "real" + + def test_predict_inference_time_positive(self): + wav_bytes = _make_wav_bytes(duration_s=2.0) + b64 = base64.b64encode(wav_bytes).decode("utf-8") + response = client.post("/predict", json={"audio_data": b64}) + data = response.json() + assert data["inference_time"] > 0 + + +@pytest.mark.skipif( + not WEIGHTS_AVAILABLE or not DATASET_AVAILABLE, + reason="Model weights or dataset not available", +) +class TestRealDataset: + """Integration tests using actual dataset files.""" + + def test_predict_real_audio(self): + """Test prediction on a real audio file.""" + path = os.path.join(REAL_DIR, "real_0.wav") + if not os.path.exists(path): + pytest.skip("real_0.wav not found") + b64 = _encode_file(path) + response = client.post("/predict", json={"audio_data": b64}) + assert response.status_code == 200 + data = response.json() + assert 0.0 <= data["probability"] <= 1.0 + + def test_predict_fake_audio(self): + """Test prediction on a fake audio file.""" + path = os.path.join(FAKE_DIR, "fake_1.wav") + if not os.path.exists(path): + pytest.skip("fake_1.wav not found") + b64 = _encode_file(path) + response = client.post("/predict", json={"audio_data": b64}) + assert response.status_code == 200 + data = response.json() + assert 0.0 <= data["probability"] <= 1.0 + + +class TestPredictValidation: + """Tests for input validation on /predict.""" + + def test_predict_missing_audio_data(self): + response = client.post("/predict", json={}) + assert response.status_code == 422 + + def test_predict_invalid_base64(self): + response = client.post("/predict", json={"audio_data": "not-valid-base64!!!"}) + # Should return 500 (decode error) or 422 + assert response.status_code in (400, 422, 500) + + def test_predict_threshold_out_of_range(self): + response = client.post( + "/predict", + json={"audio_data": "dGVzdA==", "threshold": 1.5}, + ) + assert response.status_code == 422 diff --git a/audio/sonics/Dockerfile b/audio/sonics/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..a4fd9ed2b29ea33ce737efeb6385218c1e59e3c1 --- /dev/null +++ b/audio/sonics/Dockerfile @@ -0,0 +1,36 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3 python3-pip python3-dev \ + git ffmpeg libsndfile1 \ + build-essential g++ \ + && rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +WORKDIR /app + +# Install PyTorch with CUDA 12.1 +RUN pip install --no-cache-dir \ + torch==2.5.1 torchaudio==2.5.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY app.py . + +# Pre-download model weights +RUN python -c "from sonics import HFAudioClassifier; HFAudioClassifier.from_pretrained('awsaf49/sonics-spectttra-alpha-120s')" + +EXPOSE 8003 + +RUN adduser --disabled-password --gecos '' appuser +USER appuser + +ENV PRELOAD_MODEL=true + +CMD ["python", "app.py"] diff --git a/audio/sonics/app.py b/audio/sonics/app.py new file mode 100644 index 0000000000000000000000000000000000000000..d2295c37b3c334530888a28182df41a61081db43 --- /dev/null +++ b/audio/sonics/app.py @@ -0,0 +1,285 @@ +"""SONICS (SpecTTTra) Synthetic Music Detection API. + +Detects AI-generated music (Suno, Udio, etc.) using the SpecTTTra +architecture from the SONICS project (ICLR 2025). + +The model performs binary classification on raw audio waveforms via +internal MelSpectrogram features. It outputs a single logit; we apply +sigmoid to obtain the fake probability. + +Reference: https://github.com/awsaf49/sonics +""" + +import base64 +import io +import logging +import os +import platform +import time +from typing import Optional + +import librosa +import numpy as np +import torch +import uvicorn +from fastapi import FastAPI, HTTPException +from pydantic import BaseModel, Field + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", +) +logger = logging.getLogger("sonics_api") + +# Constants +MODEL_NAME = "sonics_detection" +HF_MODEL_ID = "awsaf49/sonics-spectttra-alpha-120s" +SAMPLE_RATE = 16000 +MAX_TIME = 120 # seconds (matches alpha-120s config) +MAX_LEN = MAX_TIME * SAMPLE_RATE # 1_920_000 samples +PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "true").lower() == "true" +MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", "300")) + + +def _get_device(): + """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if override: + pass # Invalid override, fall through to auto-detect + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + return torch.device("cpu") + + +DEVICE = _get_device() + +if DEVICE.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + +if DEVICE.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) +else: + logger.warning( + "Device: %s (no CUDA available -- check nvidia-container-toolkit)", + DEVICE, + ) + +# Global model instance +model = None + + +class AudioInput(BaseModel): + """Request schema for audio deepfake detection.""" + + audio_data: str = Field( + ..., description="Base64 encoded audio string (WAV/MP3/FLAC/etc)" + ) + threshold: Optional[float] = Field( + 0.5, ge=0.0, le=1.0, description="Classification threshold" + ) + + +app = FastAPI( + title="SONICS Synthetic Music Detection API", + description=( + "Service for detecting AI-generated music using the " + "SpecTTTra model from the SONICS project (ICLR 2025)." + ), + version="1.0.0", +) + + +def load_model(): + """Load the SONICS HFAudioClassifier from HuggingFace Hub. + + Returns: + The loaded model, or None if loading fails. + """ + global model + if model is not None: + return model + + logger.info("Loading SONICS model '%s' onto %s...", HF_MODEL_ID, DEVICE) + + try: + from sonics import HFAudioClassifier + + model = HFAudioClassifier.from_pretrained( + HF_MODEL_ID, + map_location=str(DEVICE), + ) + model.to(DEVICE) + model.eval() + + logger.info("SONICS model loaded successfully.") + return model + except Exception: + logger.exception("Failed to load SONICS model") + model = None + return None + + +@app.on_event("startup") +async def startup_event(): + """Optionally preload model on service startup.""" + if PRELOAD_MODEL: + load_model() + + +@app.get("/") +async def root(): + """Root info endpoint.""" + return { + "service": "SONICS Synthetic Music Detection", + "model": MODEL_NAME, + "version": "1.0.0", + } + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and DEVICE.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +async def health(): + """Health check endpoint.""" + return { + "status": "healthy" if model is not None else "degraded", + "model": MODEL_NAME, + "device": str(DEVICE), + **_gpu_health_info(), + } + + +def preprocess_audio(audio_bytes: bytes) -> torch.Tensor: + """Preprocess audio for SONICS inference. + + Loads audio from raw bytes, resamples to 16 kHz mono, + crops or zero-pads to MAX_LEN samples, and normalises by + standard deviation (matching the training pipeline). + + Args: + audio_bytes: Raw audio file bytes (WAV, MP3, FLAC, etc.). + + Returns: + Audio tensor of shape (1, MAX_LEN) on DEVICE. + + Raises: + ValueError: If audio preprocessing fails. + """ + try: + logger.info("Starting audio preprocessing...") + audio, sr = librosa.load(io.BytesIO(audio_bytes), sr=SAMPLE_RATE, mono=True) + logger.info("Audio loaded. Length: %d samples at %dHz", len(audio), sr) + + # Crop or pad to fixed length (matching SONICS dataset.py) + if len(audio) > MAX_LEN: + # Crop from 3/4 position (matching eval-mode logic) + idx = int((len(audio) - MAX_LEN) / 4 * 3) + audio = audio[idx : idx + MAX_LEN] + elif len(audio) < MAX_LEN: + audio = np.pad(audio, (0, MAX_LEN - len(audio)), mode="constant") + + # Normalise by standard deviation (matching training pipeline) + audio /= np.maximum(np.std(audio), 1e-6) + + logger.info("Audio preprocessed to %d samples", len(audio)) + + audio_tensor = torch.from_numpy(audio).float().unsqueeze(0) + audio_tensor = audio_tensor.to(DEVICE) + return audio_tensor + except Exception as e: + logger.error("Error preprocessing audio: %s", e) + raise ValueError(f"Audio preprocessing failed: {str(e)}") + + +@app.post("/predict") +async def predict(input_data: AudioInput): + """Run synthetic music detection on base64-encoded audio. + + The model uses BCEWithLogitsLoss with num_classes=1, so it + outputs a single logit. We apply sigmoid to obtain the fake + probability. + """ + if model is None: + if load_model() is None: + raise HTTPException(status_code=503, detail="Model not loaded") + + try: + start_time = time.time() + logger.info( + "Prediction request. Data size: %d chars", + len(input_data.audio_data), + ) + + # Decode base64 audio + audio_bytes = base64.b64decode(input_data.audio_data) + + # Preprocess + audio_tensor = preprocess_audio(audio_bytes) + + # Inference + logger.info("Starting model inference...") + with torch.no_grad(): + logits = model(audio_tensor) + # logits shape: (1, 1) -- single logit for binary classification + prob_fake = torch.sigmoid(logits).squeeze().item() + + prediction = 1 if prob_fake >= input_data.threshold else 0 + verdict = "fake" if prediction == 1 else "real" + inference_time = time.time() - start_time + + logger.info( + "Prediction: %s (prob_fake=%.4f, time=%.3fs)", + verdict, + prob_fake, + inference_time, + ) + + return { + "model": MODEL_NAME, + "probability": float(prob_fake), + "prediction": int(prediction), + "class": verdict, + "inference_time": float(inference_time), + } + + except Exception as e: + logger.exception("Error during prediction: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +if __name__ == "__main__": + port = int(os.environ.get("MODEL_PORT", 8003)) + uvicorn.run(app, host="0.0.0.0", port=port) diff --git a/audio/sonics/requirements.txt b/audio/sonics/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..f64e0b63bd09ac47a50599c94226356ea557fb25 --- /dev/null +++ b/audio/sonics/requirements.txt @@ -0,0 +1,22 @@ +# SONICS model (pip-installable from GitHub) +sonics @ git+https://github.com/awsaf49/sonics.git + +# Core inference +torch>=2.4.0 +torchaudio>=2.4.0 + +# Audio processing +librosa>=0.9.0 + +# ML utilities (SONICS dependency) +timm>=1.0.7 +fvcore + +# HuggingFace Hub (for model download) +huggingface-hub + +# API +fastapi +uvicorn[standard] +pydantic>=2.0 +numpy diff --git a/ensemble-core/Dockerfile b/ensemble-core/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..2753a64cfe60f3afe9d66c70e58edbb84d0fe7a4 --- /dev/null +++ b/ensemble-core/Dockerfile @@ -0,0 +1,10 @@ +FROM python:3.9-slim +WORKDIR /app +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt +COPY main.py . +# Drop root privileges +RUN adduser --disabled-password --gecos '' appuser +USER appuser + +CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8003"] diff --git a/ensemble-core/main.py b/ensemble-core/main.py new file mode 100644 index 0000000000000000000000000000000000000000..b30f7d6c016e883676bef8008e1b8985d0d666ee --- /dev/null +++ b/ensemble-core/main.py @@ -0,0 +1,23 @@ +from typing import Any, Dict + +from fastapi import FastAPI +from pydantic import BaseModel + +app = FastAPI(title="Ensemble Core Service") + + +class EnsembleRequest(BaseModel): + media_type: str + model_results: Dict[str, Any] + method: str = "stacking" + + +@app.post("/calculate") +async def calculate_ensemble(request: EnsembleRequest): + # TODO: Migrate proprietary ensemble logic from gateway + return {"verdict": "fake", "confidence": 0.95, "method_used": request.method} + + +@app.get("/health") +async def health(): + return {"status": "healthy"} diff --git a/ensemble-core/requirements.txt b/ensemble-core/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..cd0459c9ef13d8d085d062d22bfdd21b888f4097 --- /dev/null +++ b/ensemble-core/requirements.txt @@ -0,0 +1 @@ +fastapi\nuvicorn\npydantic diff --git a/ensemble-core/scripts/create_dataset.py b/ensemble-core/scripts/create_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..d5822c58638529a427a514e2a66331e28bf0d47a --- /dev/null +++ b/ensemble-core/scripts/create_dataset.py @@ -0,0 +1,169 @@ +#!/usr/bin/env python3 +""" +Build a balanced 10k-real / 10k-fake face-image folder with maximum +deepfake-tech variety. + +Usage: + python build_face_dataset.py --out_dir ./faces20k --seed 42 +""" + +import argparse +import os +import pathlib +import random +import shutil +import subprocess +import sys +import zipfile +from collections import defaultdict + +import pandas as pd +from tqdm import tqdm + +# ---------------------------------------------------------------------- +# 1. Edit here to add / remove sources +# ---------------------------------------------------------------------- +DATASETS = [ + { + "name": "140k", + "slug": "xhlulu/140k-real-and-fake-faces", + "subdirs": {"real": "real", "fake": "fake"}, + "fake_label": "stylegan2", + }, + { + "name": "deepfake_real", + "slug": "manjilkarki/deepfake-and-real-images", + "subdirs": {"real": "real", "fake": "fake"}, + "fake_label": "pggan_stylegan_mix", + }, + { + "name": "dfdc_f150", + "slug": "sciarrilli/dfdc-f150", + "subdirs": {"real": "real", "fake": "fake"}, + "fake_label": "dfdc_swaps", + }, + { + "name": "faceforensics_imgs", + "slug": "greatgamedota/faceforensics", + "subdirs": {"real": "real", "fake": "fake"}, + "fake_label": "ffpp_swaps", + }, +] + +TARGET_PER_CLASS = 10_000 +# ---------------------------------------------------------------------- + + +def kaggle_download(slug: str, dest: pathlib.Path) -> pathlib.Path: + """Download to dest/. Returns path of the zip.""" + dest.mkdir(parents=True, exist_ok=True) + zip_path = dest / f"{slug.split('/')[-1]}.zip" + if zip_path.exists(): + return zip_path + print(f"Downloading {slug} …") + subprocess.run( + ["kaggle", "datasets", "download", "-d", slug, "-p", str(dest), "--quiet"], + check=True, + ) + return zip_path + + +def extract(zip_path: pathlib.Path, dest: pathlib.Path) -> pathlib.Path: + """Unzip if needed. Returns extraction dir.""" + extract_dir = dest / zip_path.stem + if extract_dir.exists(): + return extract_dir + print(f"Extracting {zip_path.name} …") + with zipfile.ZipFile(zip_path) as zf: + zf.extractall(path=extract_dir) + return extract_dir + + +def glob_images(root: pathlib.Path, pattern: str): + return list(root.glob(pattern)) + list(root.glob(pattern.replace("jpg", "png"))) + + +def main(out_dir: pathlib.Path, seed: int): + random.seed(seed) + temp_root = out_dir / "_raw" + real_pool, fake_pool = [], [] + fake_source_tag = {} # path -> dataset tag + + # ------------------------------------------------------------------ + # 2. Pull sources + # ------------------------------------------------------------------ + for ds in DATASETS: + zip_path = kaggle_download(ds["slug"], temp_root) + extract_dir = extract(zip_path, temp_root) + real_dir = extract_dir / ds["subdirs"]["real"] + fake_dir = extract_dir / ds["subdirs"]["fake"] + real_pool += glob_images(real_dir, "**/*.jpg") + fakes = glob_images(fake_dir, "**/*.jpg") + fake_pool += fakes + for fp in fakes: + fake_source_tag[str(fp)] = ds["fake_label"] + + # sanity check + if len(real_pool) < TARGET_PER_CLASS or len(fake_pool) < TARGET_PER_CLASS: + print("Not enough images – add another dataset.", file=sys.stderr) + sys.exit(1) + + # ------------------------------------------------------------------ + # 3. Sample + # ------------------------------------------------------------------ + random.shuffle(real_pool) + random.shuffle(fake_pool) + + # try to spread fake quota equally over sources + per_source_quota = TARGET_PER_CLASS // len(DATASETS) + selected_fake = [] + taken = defaultdict(int) + for fp in fake_pool: + tag = fake_source_tag[str(fp)] + if taken[tag] < per_source_quota: + selected_fake.append(fp) + taken[tag] += 1 + if len(selected_fake) == TARGET_PER_CLASS: + break + # top-up if we’re short (some sets too small) + if len(selected_fake) < TARGET_PER_CLASS: + needed = TARGET_PER_CLASS - len(selected_fake) + selected_fake += fake_pool[len(selected_fake) : len(selected_fake) + needed] + + selected_real = real_pool[:TARGET_PER_CLASS] + + # ------------------------------------------------------------------ + # 4. Copy to final tree + manifest + # ------------------------------------------------------------------ + for cls in ("real", "fake"): + (out_dir / cls).mkdir(parents=True, exist_ok=True) + + manifest_rows = [] + + def copy_files(file_list, cls): + for src in tqdm(file_list, desc=f"Copying {cls}"): + dst = out_dir / cls / src.name + shutil.copy(src, dst) + manifest_rows.append( + { + "filename": dst.name, + "label": cls, + "source": fake_source_tag.get(str(src), "n/a"), + } + ) + + copy_files(selected_real, "real") + copy_files(selected_fake, "fake") + + pd.DataFrame(manifest_rows).to_csv(out_dir / "manifest.csv", index=False) + print("Done →", out_dir) + + +if __name__ == "__main__": + p = argparse.ArgumentParser() + p.add_argument( + "--out_dir", default="faces20k", type=pathlib.Path, help="destination folder" + ) + p.add_argument("--seed", default=42, type=int) + args = p.parse_args() + main(args.out_dir, args.seed) diff --git a/ensemble-core/scripts/meta_feature_generator.py b/ensemble-core/scripts/meta_feature_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..5735fad298d174f72241f67b72bfd7a27c2572db --- /dev/null +++ b/ensemble-core/scripts/meta_feature_generator.py @@ -0,0 +1,366 @@ +#!/usr/bin/env python3 +""" +DeepSafe Meta-Feature Generator +=============================== + +Orchestrates the generation of meta-feature datasets for training stacking ensembles. +This component acts as a data ingestion pipeline that: +1. Scans a target directory for labeled media (Real/Fake). +2. Queries the distributed model microservices to obtain base probability scores. +3. Aggregates these scores into a structured feature matrix (CSV) for the meta-learner. + +Architectural Note: +This script is designed to be fault-tolerant. If a specific model microservice is unreachable +or fails for a subset of files, the pipeline continues, recording NaNs for those features. +This ensures that a single model failure does not halt the entire training data generation process, +though downstream imputers must handle these missing values. +""" + +import argparse +import gc +import json +import os +import sys +import time +from typing import Any, Dict, List, Optional + +import numpy as np +import pandas as pd +from rich.console import Console +from rich.panel import Panel +from rich.progress import ( + BarColumn, + MofNCompleteColumn, + Progress, + SpinnerColumn, + TextColumn, + TimeElapsedColumn, +) +from rich.table import Table + +# Ensure utils is importable regardless of execution context. +# This fallback is necessary when running the script directly from the project root +# without an installed package structure. +try: + from utils.api_client import APIClient + from utils.config_manager import ConfigManager + from utils.media_handler import MediaHandler +except ImportError: + project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")) + if project_root not in sys.path: + sys.path.insert(0, project_root) + try: + from utils.api_client import APIClient + from utils.config_manager import ConfigManager + from utils.media_handler import MediaHandler + except ImportError as e: + print(f"Critical Error: Failed to resolve utils dependency. {e}") + sys.exit(1) + + +console = Console(width=120) + + +class MetaFeatureGenerator: + """ + Manages the ETL process for meta-learning datasets. + + Attributes: + media_type (str): The domain of operation (image, video, audio). + config_manager (ConfigManager): Centralized configuration handler. + api_client (APIClient): Interface for communicating with model microservices. + """ + + def __init__(self, media_type: str, config_manager: ConfigManager): + self.media_type = media_type + self.config_manager = config_manager + # run_from_host=True implies we are running outside the docker network (e.g., local dev), + # so we use localhost ports mapped in docker-compose. + self.api_client = APIClient(config_manager, media_type, run_from_host=True) + self.media_handler = MediaHandler(config_manager) + self.base_model_names = list( + config_manager.get_model_endpoints(media_type).keys() + ) + + if not self.base_model_names: + console.print( + f"[bold red]Configuration Error: No base models defined for '{media_type}'.[/bold red]" + ) + sys.exit(1) + + def generate( + self, + input_dir: str, + output_csv_path: str, + default_threshold: float, + specific_models: Optional[List[str]] = None, + ): + """ + Executes the generation pipeline. + + Args: + input_dir: Root directory containing 'Real' and 'Fake' subdirectories. + output_csv_path: Destination for the resulting feature matrix. + default_threshold: Decision threshold passed to models (mostly for logging/reference). + specific_models: Optional filter to run only a subset of available models. + """ + + console.print( + Panel( + f"[bold cyan]Meta-Feature Generation Protocol ({self.media_type.capitalize()})[/bold cyan]\n" + f"Source: {input_dir}\n" + f"Target: {output_csv_path}\n" + f"Active Models: {specific_models or 'All configured'}", + title="Pipeline Configuration", + border_style="blue", + expand=False, + ) + ) + + # Discovery phase: Scan filesystem for valid media files and infer ground truth from directory structure. + media_files_with_gt = self.media_handler.find_media_files( + input_dir, self.media_type + ) + if not media_files_with_gt: + console.print( + f"[bold red]Abort: No valid {self.media_type} files found in '{input_dir}'.[/bold red]" + ) + return + + # Determine the execution scope (subset of models vs all). + models_to_query = self.base_model_names + if specific_models: + models_to_query = [m for m in specific_models if m in self.base_model_names] + if not models_to_query: + console.print( + f"[bold red]Configuration Mismatch: Requested models {specific_models} are not configured for '{self.media_type}'.[/bold red]" + ) + return + console.print(f"Scope restricted to: {models_to_query}") + + all_feature_data = [] + + # Execution phase: Iterate through files and query models. + # We use a rich progress bar for observability during long-running batch processes. + with Progress( + SpinnerColumn(), + TextColumn("[progress.description]{task.description}"), + BarColumn(), + MofNCompleteColumn(), + TimeElapsedColumn(), + ) as progress: + total_files = len(media_files_with_gt) + outer_task = progress.add_task( + f"Processing {self.media_type} corpus...", total=total_files + ) + + for file_idx, (media_path, ground_truth_label) in enumerate( + media_files_with_gt + ): + media_file_name = os.path.basename(media_path) + progress.update( + outer_task, + description=f"Processing: [cyan]{media_file_name}[/cyan]", + ) + + # Pre-encode media to base64 once to avoid redundant I/O operations per model. + encoded_media = self.media_handler.encode_media_to_base64(media_path) + if not encoded_media: + console.print( + f"[yellow]Skip: Encoding failed for {media_file_name}.[/yellow]" + ) + progress.advance(outer_task) + continue + + # Feature vector initialization + current_media_features: Dict[str, Any] = { + "media_path": media_path, + "media_name": media_file_name, + # Map string labels to numeric binary targets: Fake=1, Real=0. + "ground_truth": ( + 1 + if ground_truth_label == "Fake" + else (0 if ground_truth_label == "Real" else -1) + ), + } + + # Initialize feature columns with NaN. This ensures structural consistency in the DataFrame + # even if specific model queries fail. + for model_name_cfg in self.base_model_names: + current_media_features[f"{model_name_cfg}_prob"] = np.nan + + # Query loop + for model_name_query in models_to_query: + model_result = self.api_client.test_with_individual_model( + model_name_query, media_path, encoded_media, default_threshold + ) + + if ( + "error" not in model_result + and model_result.get("probability") is not None + ): + current_media_features[f"{model_name_query}_prob"] = ( + model_result["probability"] + ) + else: + # Log failure but do not interrupt the pipeline. Robustness is key here. + error_msg = model_result.get( + "error", "Invalid response payload" + ) + console.print( + f"[yellow]Model Failure: {model_name_query} on {media_file_name}. Reason: {error_msg}.[/yellow]", + highlight=False, + ) + + all_feature_data.append(current_media_features) + progress.advance(outer_task) + + # Explicit garbage collection to prevent memory bloat during large dataset processing. + gc.collect() + + if not all_feature_data: + console.print( + "[bold red]Pipeline Failure: No features generated.[/bold red]" + ) + return + + # Data serialization and validation + meta_features_df = pd.DataFrame(all_feature_data) + + # Filter invalid ground truth (should be handled by discovery, but defensive programming is good). + meta_features_df = meta_features_df[meta_features_df["ground_truth"] != -1] + + if meta_features_df.empty: + console.print( + "[bold red]Data Error: No valid labeled data remaining after processing.[/bold red]" + ) + return + + # Schema enforcement: Ensure all expected columns exist. + expected_prob_cols = [f"{mn}_prob" for mn in self.base_model_names] + for col in expected_prob_cols: + if col not in meta_features_df.columns: + meta_features_df[col] = np.nan + + # Column ordering for readability and consistency. + ordered_prob_cols = sorted( + [col for col in meta_features_df.columns if col.endswith("_prob")] + ) + final_cols_order = ( + ["media_path", "media_name"] + ordered_prob_cols + ["ground_truth"] + ) + meta_features_df = meta_features_df[final_cols_order] + + try: + os.makedirs( + os.path.dirname(os.path.abspath(output_csv_path)), exist_ok=True + ) + meta_features_df.to_csv(output_csv_path, index=False, float_format="%.6f") + + console.print( + f"\n[bold green]Success: Dataset persisted to {os.path.abspath(output_csv_path)}[/bold green]" + ) + console.print(f"Dimensions: {meta_features_df.shape}") + + # Quality Assurance: Report missing values to inform downstream handling strategies. + nan_summary_table = Table( + title="Data Quality Report (Missing Values)", show_lines=True + ) + nan_summary_table.add_column("Feature", style="cyan") + nan_summary_table.add_column( + "Missing Count", style="magenta", justify="right" + ) + nan_summary_table.add_column("Missing %", style="yellow", justify="right") + + for col in ordered_prob_cols: + nan_count = meta_features_df[col].isnull().sum() + nan_percent = ( + (nan_count / len(meta_features_df)) * 100 + if len(meta_features_df) > 0 + else 0 + ) + nan_summary_table.add_row(col, str(nan_count), f"{nan_percent:.2f}%") + console.print(nan_summary_table) + + except Exception as e: + console.print( + f"[bold red]I/O Error: Failed to write output CSV. {e}[/bold red]" + ) + + +def main(): + parser = argparse.ArgumentParser( + description="DeepSafe Meta-Feature Generator: ETL for Stacking Ensemble Training Data.", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--media-type", + type=str, + choices=["image", "video", "audio"], + required=True, + help="Target domain. Defines the model registry subset.", + ) + parser.add_argument( + "--input-dir", + type=str, + required=True, + help="Source directory. Must contain 'Real' and 'Fake' subdirectories for label inference.", + ) + parser.add_argument( + "--output-csv", + type=str, + required=True, + help="Destination path for the generated feature matrix.", + ) + parser.add_argument( + "--threshold", + type=float, + help="Decision threshold override (0.0-1.0). Defaults to system config.", + ) + parser.add_argument( + "--specific-models", + type=str, + help="Optional filter: Comma-separated list of model identifiers to query.", + ) + parser.add_argument( + "--config-path", type=str, default=None, help=f"Configuration override path." + ) + + args = parser.parse_args() + + # Initialize configuration subsystem + cfg_manager = ConfigManager(config_path=args.config_path) + if not cfg_manager.is_config_loaded_successfully(): + sys.exit(1) + + default_thresh_from_config = cfg_manager.get_default("default_threshold", 0.5) + query_threshold = ( + args.threshold if args.threshold is not None else default_thresh_from_config + ) + + specific_models_list = ( + [m.strip() for m in args.specific_models.split(",")] + if args.specific_models + else None + ) + + generator = MetaFeatureGenerator(args.media_type, cfg_manager) + generator.generate( + args.input_dir, args.output_csv, query_threshold, specific_models_list + ) + + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + console.print("\n[bold yellow]Process Interrupted by User.[/bold yellow]") + sys.exit(0) + except Exception as e: + console.print(f"\n[bold red]Fatal Error: {e}[/bold red]") + import traceback + + console.print( + Panel(traceback.format_exc(), title="Stack Trace", border_style="red") + ) + sys.exit(1) diff --git a/ensemble-core/scripts/train_meta_learner_advanced.py b/ensemble-core/scripts/train_meta_learner_advanced.py new file mode 100644 index 0000000000000000000000000000000000000000..5239df112dfc65dbacde1b1f47563057717416b4 --- /dev/null +++ b/ensemble-core/scripts/train_meta_learner_advanced.py @@ -0,0 +1,1228 @@ +#!/usr/bin/env python3 +""" +DeepSafe Advanced Meta-Learner Training Suite (train_meta_learner_advanced.py) +============================================================================== + +This script trains and evaluates various meta-learners (stacking ensembles) +for deepfake detection. It takes a CSV file of meta-features (outputs from +base deepfake detection models) and ground truth labels as input. + +Key Features: +------------- +1. Modality-Specific Training: Supports training separate meta-learners for + different media types (image, video, audio) using the `--media-type` argument. + This ensures that the meta-learner is optimized for the characteristics of + the base models relevant to that modality. +2. Data Preprocessing: Includes imputation for missing values (e.g., if a base + model failed) and feature scaling. +3. Multiple Meta-Learner Models: Trains and evaluates several standard classifiers + (Logistic Regression, Random Forest, Gradient Boosting, SVC, KNN, Naive Bayes) + and, if available, advanced models like XGBoost and LightGBM. +4. Hyperparameter Optimization: + - Supports Optuna for efficient hyperparameter search. + - Falls back to GridSearchCV if Optuna is not installed or if specified. +5. Comprehensive Evaluation: + - Calculates Accuracy, F1-Score, Precision, Recall, and ROC AUC for each model. + - Generates classification reports and confusion matrices. + - Plots ROC curves for visual comparison of all trained meta-learners and + simple ensemble baselines. +6. Simple Ensemble Baselines: Also evaluates simple averaging and majority vote + ensembles for comparison against more complex stacking models. Includes an + option for optimized weighted averaging. +7. Artifact Generation: + - Saves all trained meta-learner models (e.g., .joblib files). + - Saves the data preprocessor (imputer + scaler). + - Saves the list of feature columns used during training. + - Saves a summary of all experiment metrics in JSON format. + - The final, best-performing trainable meta-learner and its associated + preprocessors are saved with generic names inside media-type specific + subfolders (e.g., api_artifacts_dir/image/deepsafe_meta_learner.joblib). +8. Configurable Output: Allows specifying separate directories for general + experiment outputs and for API-ready deployment artifacts. + +CLI Usage: +---------- +python train_meta_learner_advanced.py \\ + --media-type [image|video|audio] \\ + --meta-file /path/to/meta_features_[media_type].csv \\ + --output-dir ./meta_learning_experiment_runs/ \\ + --api-artifacts-dir ./api/meta_model_artifacts/ \\ + [--optimizer optuna|gridsearch] \\ + [--optuna-trials 50] \\ + [--weights /path/to/custom_weights.json] + +Arguments: +---------- + --media-type {image,video,audio} + (Required) The type of media for which the meta-learner + is being trained. This affects output artifact naming. + --meta-file META_FILE + (Required) Path to the CSV file containing meta-features + (base model outputs) and a 'ground_truth' column. + --output-dir OUTPUT_DIR + Base directory for saving all experiment-related outputs + (logs, plots, individual model files from this run). + A timestamped, media-type-specific subdirectory will be created. + (Default: ./meta_learning_experiment_runs/) + --api-artifacts-dir API_ARTIFACTS_DIR + Directory to save the final, API-ready deployment artifacts + (e.g., ./api/meta_model_artifacts/image/deepsafe_meta_learner.joblib). + (Default: ./api/meta_model_artifacts/) + --optimizer {optuna,gridsearch} + Hyperparameter optimization strategy (Default: optuna). + --optuna-trials N + Number of trials for Optuna optimization (Default: 50). + --weights WEIGHTS_PATH_OR_JSON + Optional. Path to a JSON file or a JSON string defining + custom weights for the 'Provided_Weighted_Average' ensemble. + Keys should be base model names (without '_prob' suffix). + +Example (Image Meta-Learner): +----------------------------- +python train_meta_learner_advanced.py \\ + --media-type image \\ + --meta-file ./meta_learning_data/meta_features_image.csv \\ + --output-dir ./ml_experiments_images \\ + --api-artifacts-dir ./deepsafe_private/api/meta_model_artifacts \\ + --optimizer optuna \\ + --optuna-trials 100 + +This will train image-specific meta-learners, save experiment details in +`./ml_experiments_images/experiments_image_YYYYMMDD_HHMMSS/`, and place +API-ready artifacts like `deepsafe_meta_learner.joblib` into +`./deepsafe_private/api/meta_model_artifacts/image/`. +""" + +import argparse +import itertools +import json +import os +import time +from typing import Any, Dict, List, Optional, Tuple + +import joblib +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import seaborn as sns +from rich.console import Console +from rich.panel import Panel +from rich.progress import ( + BarColumn, + MofNCompleteColumn, + Progress, + SpinnerColumn, + TextColumn, + TimeElapsedColumn, +) +from rich.table import Table +from sklearn.ensemble import GradientBoostingClassifier, RandomForestClassifier +from sklearn.impute import SimpleImputer +from sklearn.linear_model import LogisticRegression +from sklearn.metrics import ( + accuracy_score, + auc, + classification_report, + confusion_matrix, + f1_score, + precision_score, + recall_score, + roc_auc_score, + roc_curve, +) +from sklearn.model_selection import StratifiedKFold, train_test_split +from sklearn.naive_bayes import GaussianNB +from sklearn.neighbors import KNeighborsClassifier +from sklearn.pipeline import Pipeline +from sklearn.preprocessing import StandardScaler +from sklearn.svm import SVC + +# --- Optional Advanced Hyperparameter Optimization & Models --- +OPTIMIZER_CHOICE_DEFAULT = "optuna" + +try: + import optuna + + OPTIMIZER_AVAILABLE_OPTUNA = True +except ImportError: + optuna = None + OPTIMIZER_AVAILABLE_OPTUNA = False + +from sklearn.model_selection import GridSearchCV + +try: + from xgboost import XGBClassifier + + XGBOOST_AVAILABLE = True +except ImportError: + XGBClassifier = None + XGBOOST_AVAILABLE = False + +try: + from lightgbm import LGBMClassifier + + LIGHTGBM_AVAILABLE = True +except ImportError: + LGBMClassifier = None + LIGHTGBM_AVAILABLE = False + +console = Console(width=120) + +# --- Configuration --- +DEFAULT_EXPERIMENT_OUTPUT_DIR_BASE = "./meta_learning_experiment_runs" +DEFAULT_API_ARTIFACTS_DIR = "./api/meta_model_artifacts" +DEFAULT_THRESHOLD_FOR_SIMPLE_ENSEMBLES = 0.5 +N_OPTUNA_TRIALS_DEFAULT = 50 +CV_FOLDS_DEFAULT = 5 + + +# --- Helper Functions --- +class NpEncoder(json.JSONEncoder): + def default(self, o: Any) -> Any: + if isinstance(o, np.integer): + return int(o) + if isinstance(o, np.floating): + return float(o) + if isinstance(o, np.ndarray): + return o.tolist() + return super(NpEncoder, self).default(o) + + +def evaluate_model_predictions( + y_true: np.ndarray, + y_pred_class: np.ndarray, + y_pred_proba: Optional[np.ndarray], + model_name: str = "Model", +) -> Dict[str, Any]: + metrics: Dict[str, Any] = {"name": model_name} + try: + metrics["accuracy"] = accuracy_score(y_true, y_pred_class) + metrics["f1_score"] = f1_score(y_true, y_pred_class, zero_division=0) + metrics["precision"] = precision_score(y_true, y_pred_class, zero_division=0) + metrics["recall"] = recall_score(y_true, y_pred_class, zero_division=0) + + roc_auc_val = np.nan + if y_pred_proba is not None and len(np.unique(y_true)) > 1: + if not ( + len(np.unique(y_pred_proba)) < 2 and len(y_pred_proba) == len(y_true) + ): + try: + roc_auc_val = roc_auc_score(y_true, y_pred_proba) + except ValueError: + pass + metrics["roc_auc"] = roc_auc_val + + metrics["classification_report_dict"] = classification_report( + y_true, y_pred_class, digits=4, zero_division=0, output_dict=True + ) + metrics["confusion_matrix_list"] = confusion_matrix( + y_true, y_pred_class + ).tolist() + metrics["y_pred_test_classes_list"] = ( + y_pred_class.tolist() + if isinstance(y_pred_class, np.ndarray) + else y_pred_class + ) + metrics["y_prob_test_scores_list"] = ( + y_pred_proba.tolist() + if y_pred_proba is not None and isinstance(y_pred_proba, np.ndarray) + else y_pred_proba + ) + except Exception as e: + console.print( + f"[bold red]Error during evaluation for {model_name}: {e}[/bold red]" + ) + for m_key in ["accuracy", "f1_score", "precision", "recall", "roc_auc"]: + metrics[m_key] = np.nan + metrics["classification_report_dict"] = {} + metrics["confusion_matrix_list"] = [] + return metrics + + +def plot_roc_curves_all( + experiment_results_dict: Dict[str, Dict[str, Any]], + y_true_labels: np.ndarray, + output_dir_path: str, + media_type: str, +): + plt.figure(figsize=(12, 10)) + plot_count = 0 + for model_key, result_data in experiment_results_dict.items(): + if ( + "y_prob_test_scores_list" in result_data + and result_data["y_prob_test_scores_list"] is not None + ): + proba_scores = np.array(result_data["y_prob_test_scores_list"]) + if len(np.unique(y_true_labels)) < 2 or ( + proba_scores.ndim > 0 + and len(np.unique(proba_scores)) < 2 + and len(proba_scores) == len(y_true_labels) + ): + continue + try: + fpr, tpr, _ = roc_curve(y_true_labels, proba_scores) + roc_auc_value = result_data.get("roc_auc", auc(fpr, tpr)) + if pd.notna(roc_auc_value): + plt.plot( + fpr, + tpr, + lw=1.8, + label=f"{model_key} (AUC = {roc_auc_value:.4f})", + ) + plot_count += 1 + except ValueError as e: + console.print( + f"[yellow]Could not plot ROC for {model_key} ({media_type}): {e}[/yellow]" + ) + + if plot_count > 0: + plt.plot([0, 1], [0, 1], color="grey", lw=1.5, linestyle="--") + plt.xlim([-0.01, 1.0]) + plt.ylim([0.0, 1.01]) + plt.xlabel("False Positive Rate", fontsize=13) + plt.ylabel("True Positive Rate", fontsize=13) + plt.title( + f"Meta-Learner & Ensemble ROC Curves ({media_type.capitalize()})", + fontsize=15, + ) + plt.legend(loc="lower right", fontsize="small", frameon=True) + plt.grid(alpha=0.35, linestyle=":") + plt.tight_layout() + plot_path = os.path.join( + output_dir_path, f"all_meta_learners_roc_curves_{media_type}.png" + ) + plt.savefig(plot_path, dpi=150) + console.print( + f"Combined ROC curves plot for {media_type} saved to [green]{plot_path}[/green]" + ) + else: + console.print(f"[yellow]No valid ROC curves to plot for {media_type}.[/yellow]") + plt.close() + + +def optimize_average_weights_simple_grid( + X_val_probs: np.ndarray, + y_val_true: np.ndarray, + num_base_models: int, + weight_options: Optional[List[float]] = None, +) -> np.ndarray: + if weight_options is None: + weight_options = [0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0] + best_auc_val = -1.0 + best_weights_val = np.ones(num_base_models) + + max_combinations_exhaustive = 5**4 + num_random_samples_if_large = 2000 + + if num_base_models <= 0: + console.print( + "[yellow]No base models to optimize weights for. Returning default weights.[/yellow]" + ) + return best_weights_val + + if ( + num_base_models <= 4 + and len(weight_options) ** num_base_models <= max_combinations_exhaustive + ): + weight_candidates = list( + itertools.product(weight_options, repeat=num_base_models) + ) + console.print( + f"Optimizing average weights with exhaustive grid search ({len(weight_candidates)} trials)." + ) + else: + console.print( + f"[yellow]Optimizing average weights with random sampling ({num_random_samples_if_large} trials due to {num_base_models} models).[/yellow]" + ) + weight_candidates = [ + np.array(np.random.choice(weight_options, num_base_models)) + for _ in range(num_random_samples_if_large) + ] + + with Progress( + SpinnerColumn(), + TextColumn("[progress.description]{task.description}"), + BarColumn(), + TextColumn("{task.percentage:>3.1f}%"), + TimeElapsedColumn(), + MofNCompleteColumn(), + ) as progress: + task = progress.add_task("Weight Grid Search", total=len(weight_candidates)) + for current_weights_tuple in weight_candidates: + current_weights = np.array(current_weights_tuple) + if np.sum(current_weights) == 0: + progress.update(task, advance=1) + continue + + if X_val_probs.shape[0] == 0: + progress.update(task, advance=1) + continue + weighted_avg_probs_val_set = np.average( + X_val_probs, axis=1, weights=current_weights + ) + + current_auc_val = 0.0 + if len(np.unique(y_val_true)) > 1 and not ( + len(np.unique(weighted_avg_probs_val_set)) < 2 + and len(weighted_avg_probs_val_set) == len(y_val_true) + ): + try: + current_auc_val = roc_auc_score( + y_val_true, weighted_avg_probs_val_set + ) + except ValueError: + pass + if current_auc_val > best_auc_val: + best_auc_val, best_weights_val = current_auc_val, current_weights + progress.update(task, advance=1) + + console.print( + f"Best weights from validation grid search: {best_weights_val.tolist()} with Val AUC: {best_auc_val:.4f}" + ) + return best_weights_val + + +# --- Main Experimentation Function --- +def run_meta_learning_experiments( + meta_features_file: str, + output_dir_base: str, + api_artifacts_dir: str, + media_type: str, + optimizer_type: str, + n_optuna_trials_config: int, + provided_custom_weights: Optional[Dict[str, float]] = None, +): + global OPTIMIZER_CHOICE, N_OPTUNA_TRIALS + OPTIMIZER_CHOICE = optimizer_type + N_OPTUNA_TRIALS = n_optuna_trials_config + + if OPTIMIZER_CHOICE == "optuna" and not OPTIMIZER_AVAILABLE_OPTUNA: + console.print( + "[yellow]Optuna chosen but not installed. Falling back to GridSearchCV.[/yellow]" + ) + OPTIMIZER_CHOICE = "gridsearch" + + experiment_run_output_dir = os.path.join( + output_dir_base, f"experiments_{media_type}_{time.strftime('%Y%m%d_%H%M%S')}" + ) + os.makedirs(experiment_run_output_dir, exist_ok=True) + + # Main API artifacts directory (parent for media-specific subfolders) + os.makedirs(api_artifacts_dir, exist_ok=True) + # Media-type specific subdirectory within the main api_artifacts_dir + media_type_api_artifacts_subdir = os.path.join(api_artifacts_dir, media_type) + os.makedirs(media_type_api_artifacts_subdir, exist_ok=True) + + console.rule( + f"[bold cyan]DeepSafe Meta-Learning: {media_type.upper()} (Optimizer: {OPTIMIZER_CHOICE})[/bold cyan]" + ) + console.print( + Panel( + f"Meta-features: {meta_features_file}\n" + f"Experiment outputs: {os.path.abspath(experiment_run_output_dir)}\n" + f"API artifacts subfolder: {os.path.abspath(media_type_api_artifacts_subdir)}", + title="Paths", + border_style="dim blue", + expand=False, + ) + ) + all_experiment_results: Dict[str, Dict[str, Any]] = {} + + console.rule("[bold]1. Data Loading and Preprocessing[/bold]") + try: + df_meta = pd.read_csv(meta_features_file) + console.print( + f"Loaded {media_type} meta-features from: [cyan]{meta_features_file}[/cyan], shape: {df_meta.shape}" + ) + except Exception as e: + console.print( + f"[bold red]Fatal Error: Could not load meta-features file: {e}[/bold red]" + ) + return + + base_model_prob_features = sorted( + [col for col in df_meta.columns if col.endswith("_prob")] + ) + if not base_model_prob_features: + console.print( + "[bold red]Fatal Error: No base model probability columns (ending with '_prob') found in CSV.[/bold red]" + ) + return + + console.print( + f"Identified [magenta]{len(base_model_prob_features)}[/magenta] base model probability features: {base_model_prob_features}" + ) + + temp_exp_feature_cols_path = os.path.join( + experiment_run_output_dir, f"experiment_feature_columns_{media_type}.json" + ) + with open(temp_exp_feature_cols_path, "w") as f: + json.dump(base_model_prob_features, f, indent=2) + + X_meta_all = df_meta[base_model_prob_features].copy() + y_meta_all = df_meta["ground_truth"] + + cols_to_drop_all_nan = X_meta_all.columns[X_meta_all.isnull().all()].tolist() + if cols_to_drop_all_nan: + console.print( + f"[yellow]Warning: Dropping fully NaN columns: {cols_to_drop_all_nan}[/yellow]" + ) + X_meta_all = X_meta_all.drop(columns=cols_to_drop_all_nan) + base_model_prob_features = [ + col for col in base_model_prob_features if col not in cols_to_drop_all_nan + ] + if not base_model_prob_features: + console.print( + "[bold red]Fatal Error: All features became NaN after dropping some columns.[/bold red]" + ) + return + with open(temp_exp_feature_cols_path, "w") as f: + json.dump(base_model_prob_features, f, indent=2) + + X_meta_train_val, X_meta_test, y_meta_train_val, y_meta_test = train_test_split( + X_meta_all, + y_meta_all, + test_size=0.25, + random_state=42, + stratify=y_meta_all if len(np.unique(y_meta_all)) > 1 else None, + ) + console.print( + f"Data split: Meta-Train/Val shape {X_meta_train_val.shape}, Meta-Test shape {X_meta_test.shape}" + ) + + ml_preprocessor = Pipeline( + [("imputer", SimpleImputer(strategy="median")), ("scaler", StandardScaler())] + ) + X_meta_train_val_processed = ml_preprocessor.fit_transform(X_meta_train_val) + X_meta_test_processed = ml_preprocessor.transform(X_meta_test) + + joblib.dump( + ml_preprocessor, + os.path.join( + experiment_run_output_dir, f"experiment_ml_preprocessor_{media_type}.joblib" + ), + ) + console.print( + f"ML preprocessor for {media_type} (imputer + scaler) fitted and saved for this run." + ) + + imputer_for_simple_ensembles = ml_preprocessor.named_steps["imputer"] + X_meta_test_imputed_only_df = pd.DataFrame( + imputer_for_simple_ensembles.transform(X_meta_test), columns=X_meta_test.columns + ) + + console.rule("[bold]2. Defining ML Meta-Learners and Hyperparameter Spaces[/bold]") + models_and_param_spaces: Dict[str, Tuple[Any, Dict[str, Any]]] = { + "LogisticRegression": ( + LogisticRegression( + solver="liblinear", + random_state=42, + class_weight="balanced", + max_iter=3000, + ), + { + "C": ( + (0.01, 1000.0, "loguniform") + if OPTIMIZER_CHOICE == "optuna" + else [0.01, 0.1, 1, 10, 100, 500] + ) + }, + ), + "RandomForest": ( + RandomForestClassifier(random_state=42, class_weight="balanced"), + { + "n_estimators": ( + (100, 500, "int") + if OPTIMIZER_CHOICE == "optuna" + else [100, 200, 300, 400] + ), + "max_depth": ( + (5, 25, "int", True) + if OPTIMIZER_CHOICE == "optuna" + else [5, 10, 15, 20, None] + ), + "min_samples_split": ( + (2, 20, "int") if OPTIMIZER_CHOICE == "optuna" else [2, 5, 10, 15] + ), + "min_samples_leaf": ( + (1, 15, "int") if OPTIMIZER_CHOICE == "optuna" else [1, 5, 10, 15] + ), + }, + ), + "GradientBoosting": ( + GradientBoostingClassifier(random_state=42), + { + "n_estimators": ( + (100, 500, "int") + if OPTIMIZER_CHOICE == "optuna" + else [100, 200, 300, 400] + ), + "learning_rate": ( + (0.005, 0.2, "loguniform") + if OPTIMIZER_CHOICE == "optuna" + else [0.01, 0.05, 0.1, 0.15] + ), + "max_depth": ( + (3, 10, "int") if OPTIMIZER_CHOICE == "optuna" else [3, 5, 7, 9] + ), + }, + ), + "SVC_Linear": ( + SVC( + kernel="linear", + probability=True, + random_state=42, + class_weight="balanced", + max_iter=10000, + ), + { + "C": ( + (0.01, 100.0, "loguniform") + if OPTIMIZER_CHOICE == "optuna" + else [0.1, 1, 10, 100] + ) + }, + ), + "KNeighbors": ( + KNeighborsClassifier(), + { + "n_neighbors": ( + (3, 25, "int", False, 2) + if OPTIMIZER_CHOICE == "optuna" + else [3, 5, 7, 11, 15, 19, 23] + ), + "weights": ( + (["uniform", "distance"], "categorical") + if OPTIMIZER_CHOICE == "optuna" + else ["uniform", "distance"] + ), + }, + ), + "GaussianNB": (GaussianNB(), {}), + } + if XGBOOST_AVAILABLE and XGBClassifier: + models_and_param_spaces["XGBoost"] = ( + XGBClassifier(random_state=42, eval_metric="auc"), + { + "n_estimators": ( + (100, 600, "int") + if OPTIMIZER_CHOICE == "optuna" + else [100, 200, 300, 400, 500] + ), + "learning_rate": ( + (0.005, 0.2, "loguniform") + if OPTIMIZER_CHOICE == "optuna" + else [0.01, 0.05, 0.1] + ), + "max_depth": ( + (3, 12, "int") if OPTIMIZER_CHOICE == "optuna" else [3, 5, 7, 9, 11] + ), + "scale_pos_weight": ( + ( + (np.sum(y_meta_train_val == 0) / np.sum(y_meta_train_val == 1)) + if np.sum(y_meta_train_val == 1) > 0 + else 1.0 + ), + ), + }, + ) + if LIGHTGBM_AVAILABLE and LGBMClassifier: + models_and_param_spaces["LightGBM"] = ( + LGBMClassifier( + random_state=42, class_weight="balanced", metric="auc", verbosity=-1 + ), + { + "n_estimators": ( + (100, 600, "int") + if OPTIMIZER_CHOICE == "optuna" + else [100, 200, 300, 400, 500] + ), + "learning_rate": ( + (0.005, 0.2, "loguniform") + if OPTIMIZER_CHOICE == "optuna" + else [0.01, 0.05, 0.1] + ), + "num_leaves": ( + (20, 150, "int") + if OPTIMIZER_CHOICE == "optuna" + else [31, 50, 70, 100, 130] + ), + }, + ) + + console.rule( + f"[bold]3. Training and Evaluating ML-based Meta-Learners ({media_type.capitalize()} Stacking)[/bold]" + ) + cv_strategy = StratifiedKFold( + n_splits=CV_FOLDS_DEFAULT, shuffle=True, random_state=42 + ) + trained_ml_model_objects: Dict[str, Any] = {} + + for model_name_key, ( + model_instance_template, + param_def, + ) in models_and_param_spaces.items(): + console.rule( + f"[bold blue]Optimizing & Training {media_type.capitalize()} Meta-Learner: {model_name_key}[/bold blue]", + style="blue", + ) + start_train_time = time.time() + best_estimator_for_model = None + + if not param_def: + model_instance_template.fit(X_meta_train_val_processed, y_meta_train_val) + best_estimator_for_model = model_instance_template + console.print( + f"{model_name_key} fitted directly (no hyperparameters tuned)." + ) + elif OPTIMIZER_CHOICE == "optuna" and optuna: + + def optuna_objective(trial: optuna.Trial): + current_params = {} + for p_name, p_opts in param_def.items(): + if isinstance(p_opts, tuple) and len(p_opts) >= 2: + suggestion_type_or_values = ( + p_opts[1] + if p_name == "weights" and p_opts[1] == "categorical" + else p_opts[2] + ) + if suggestion_type_or_values == "loguniform": + current_params[p_name] = trial.suggest_float( + p_name, p_opts[0], p_opts[1], log=True + ) + elif suggestion_type_or_values == "uniform": + current_params[p_name] = trial.suggest_float( + p_name, p_opts[0], p_opts[1] + ) + elif suggestion_type_or_values == "int": + low, high = p_opts[0], p_opts[1] + can_be_none = p_opts[3] if len(p_opts) > 3 else False + step = p_opts[4] if len(p_opts) > 4 else 1 + val = trial.suggest_int(p_name, low, high, step=step) + if can_be_none and trial.suggest_categorical( + f"{p_name}_use_none", [True, False] + ): + val = None + current_params[p_name] = val + elif suggestion_type_or_values == "categorical": + current_params[p_name] = trial.suggest_categorical( + p_name, p_opts[0] + ) + elif len(p_opts) == 1 and not isinstance(p_opts[0], list): + current_params[p_name] = p_opts[0] + else: + console.print( + f"[red]Warning: Unknown Optuna parameter definition for {p_name}: {p_opts}[/red]" + ) + else: + if ( + p_name in model_instance_template.get_params() + and not isinstance(p_opts, tuple) + ): + current_params[p_name] = p_opts + + model_trial = model_instance_template.__class__( + **model_instance_template.get_params() + ) + valid_model_params = model_trial.get_params().keys() + filtered_current_params = { + k: v for k, v in current_params.items() if k in valid_model_params + } + model_trial.set_params(**filtered_current_params) + + scores = [] + for train_idx, val_idx in cv_strategy.split( + X_meta_train_val_processed, y_meta_train_val + ): + X_fold_train, X_fold_val = ( + X_meta_train_val_processed[train_idx], + X_meta_train_val_processed[val_idx], + ) + y_fold_train, y_fold_val = ( + y_meta_train_val.iloc[train_idx], + y_meta_train_val.iloc[val_idx], + ) + model_trial.fit(X_fold_train, y_fold_train) + if hasattr(model_trial, "predict_proba"): + try: + y_val_pred_proba = model_trial.predict_proba(X_fold_val)[ + :, 1 + ] + if len(np.unique(y_fold_val)) < 2 or ( + len(np.unique(y_val_pred_proba)) < 2 + and len(y_val_pred_proba) == len(y_fold_val) + ): + scores.append(0.5) + else: + scores.append( + roc_auc_score(y_fold_val, y_val_pred_proba) + ) + except Exception: + scores.append(0.0) + else: + scores.append( + f1_score( + y_fold_val, + model_trial.predict(X_fold_val), + zero_division=0, + ) + ) + return np.mean(scores) + + study = optuna.create_study( + direction="maximize", pruner=optuna.pruners.MedianPruner() + ) + study.optimize( + optuna_objective, + n_trials=N_OPTUNA_TRIALS, + show_progress_bar=True, + gc_after_trial=True, + ) + + sklearn_best_params = {} + for p_name_orig_def, p_opts_def in param_def.items(): + if p_name_orig_def in study.best_params: + sklearn_best_params[p_name_orig_def] = study.best_params[ + p_name_orig_def + ] + if len(p_opts_def) > 3 and p_opts_def[3] is True: + if ( + study.best_params.get(f"{p_name_orig_def}_use_none", False) + is True + ): + sklearn_best_params[p_name_orig_def] = None + console.print( + f"Best Optuna params for {model_name_key} ({media_type}): {sklearn_best_params}" + ) + best_estimator_for_model = model_instance_template.__class__( + **model_instance_template.get_params() + ) + best_estimator_for_model.set_params(**sklearn_best_params) + best_estimator_for_model.fit(X_meta_train_val_processed, y_meta_train_val) + else: + grid_search = GridSearchCV( + model_instance_template, + param_def, + cv=cv_strategy, + scoring="roc_auc", + n_jobs=-1, + verbose=0, + ) + grid_search.fit(X_meta_train_val_processed, y_meta_train_val) + best_estimator_for_model = grid_search.best_estimator_ + console.print( + f"Best GridSearchCV params for {model_name_key} ({media_type}): {grid_search.best_params_}" + ) + + joblib.dump( + best_estimator_for_model, + os.path.join( + experiment_run_output_dir, + f"{model_name_key}_meta_learner_{media_type}.joblib", + ), + ) + trained_ml_model_objects[model_name_key] = best_estimator_for_model + + y_test_pred_classes = best_estimator_for_model.predict(X_meta_test_processed) + y_test_pred_probas = ( + best_estimator_for_model.predict_proba(X_meta_test_processed)[:, 1] + if hasattr(best_estimator_for_model, "predict_proba") + else None + ) + metrics_results = evaluate_model_predictions( + y_meta_test.values, y_test_pred_classes, y_test_pred_probas, model_name_key + ) + all_experiment_results[model_name_key] = metrics_results + train_time = time.time() - start_train_time + console.print( + f"[bold]{model_name_key} Test Set Perf. ({media_type}):[/bold] AUC: {metrics_results.get('roc_auc', np.nan):.4f}, F1: {metrics_results.get('f1_score', np.nan):.4f}, Acc: {metrics_results.get('accuracy', np.nan):.4f} (Train time: {train_time:.2f}s)" + ) + + console.rule( + f"[bold]4. Evaluating Simple Ensemble Baselines ({media_type.capitalize()} Meta-Test Set)[/bold]" + ) + avg_probs_meta_test = X_meta_test_imputed_only_df.mean(axis=1).values + avg_preds_meta_test_classes = ( + avg_probs_meta_test >= DEFAULT_THRESHOLD_FOR_SIMPLE_ENSEMBLES + ).astype(int) + all_experiment_results["Simple_Average_Prob"] = evaluate_model_predictions( + y_meta_test.values, + avg_preds_meta_test_classes, + avg_probs_meta_test, + "Simple_Average_Prob", + ) + console.print( + f"[bold]Simple Average Prob Test ({media_type}):[/bold] AUC: {all_experiment_results['Simple_Average_Prob'].get('roc_auc', np.nan):.4f}, F1: {all_experiment_results['Simple_Average_Prob'].get('f1_score', np.nan):.4f}" + ) + + binarized_X_meta_test = ( + X_meta_test_imputed_only_df.values >= DEFAULT_THRESHOLD_FOR_SIMPLE_ENSEMBLES + ).astype(int) + num_models_for_vote = X_meta_test_imputed_only_df.shape[1] + fake_votes_per_item_meta_test = binarized_X_meta_test.sum(axis=1) + maj_vote_preds_meta_test_classes = ( + fake_votes_per_item_meta_test >= (num_models_for_vote / 2.0) + ).astype(int) + maj_vote_prob_scores_meta_test = ( + fake_votes_per_item_meta_test / num_models_for_vote + if num_models_for_vote > 0 + else np.full_like(fake_votes_per_item_meta_test, 0.5, dtype=float) + ) + all_experiment_results["Simple_Majority_Vote"] = evaluate_model_predictions( + y_meta_test.values, + maj_vote_preds_meta_test_classes, + maj_vote_prob_scores_meta_test, + "Simple_Majority_Vote", + ) + console.print( + f"[bold]Simple Majority Vote Test ({media_type}):[/bold] AUC: {all_experiment_results['Simple_Majority_Vote'].get('roc_auc', np.nan):.4f}, F1: {all_experiment_results['Simple_Majority_Vote'].get('f1_score', np.nan):.4f}" + ) + + if provided_custom_weights: + current_weights_values = [ + provided_custom_weights.get(fc.replace("_prob", ""), 1.0) + for fc in base_model_prob_features + ] + current_weights_array = np.array(current_weights_values) + + if ( + len(current_weights_array) == X_meta_test_imputed_only_df.shape[1] + and np.sum(current_weights_array) > 0 + ): + prov_weighted_avg_probs_meta_test = np.average( + X_meta_test_imputed_only_df.values, + axis=1, + weights=current_weights_array, + ) + prov_weighted_avg_preds_meta_test_classes = ( + prov_weighted_avg_probs_meta_test + >= DEFAULT_THRESHOLD_FOR_SIMPLE_ENSEMBLES + ).astype(int) + all_experiment_results["Provided_Weighted_Average"] = ( + evaluate_model_predictions( + y_meta_test.values, + prov_weighted_avg_preds_meta_test_classes, + prov_weighted_avg_probs_meta_test, + "Provided_Weighted_Average", + ) + ) + console.print( + f"[bold]Provided Weighted Average Test ({media_type}):[/bold] AUC: {all_experiment_results['Provided_Weighted_Average'].get('roc_auc', np.nan):.4f}, F1: {all_experiment_results['Provided_Weighted_Average'].get('f1_score', np.nan):.4f}" + ) + else: + console.print( + f"[yellow]Warning: Mismatch in provided_custom_weights keys vs. features for {media_type}, or sum of weights is zero. Skipping.[/yellow]" + ) + + X_train_val_imputed_for_opt_df = pd.DataFrame( + ml_preprocessor.named_steps["imputer"].transform(X_meta_train_val), + columns=base_model_prob_features, + ) + stratify_opt_split = ( + y_meta_train_val if len(np.unique(y_meta_train_val)) > 1 else None + ) + X_opt_train_df, X_opt_val_df, y_opt_train_series, y_opt_val_series = ( + train_test_split( + X_train_val_imputed_for_opt_df, + y_meta_train_val, + test_size=0.33, + random_state=123, + stratify=stratify_opt_split, + ) + ) + if X_opt_val_df.shape[0] > 10 and X_opt_val_df.shape[1] > 0: + console.print( + f"Optimizing weights for averaging ({media_type}) using a validation split of meta-train data..." + ) + optimized_avg_weights = optimize_average_weights_simple_grid( + X_opt_val_df.values, y_opt_val_series.values, X_opt_val_df.shape[1] + ) + opt_w_avg_probs_meta_test = np.average( + X_meta_test_imputed_only_df.values, axis=1, weights=optimized_avg_weights + ) + opt_w_avg_preds_meta_test_classes = ( + opt_w_avg_probs_meta_test >= DEFAULT_THRESHOLD_FOR_SIMPLE_ENSEMBLES + ).astype(int) + all_experiment_results["Optimized_Grid_Weighted_Average"] = ( + evaluate_model_predictions( + y_meta_test.values, + opt_w_avg_preds_meta_test_classes, + opt_w_avg_probs_meta_test, + "Optimized_Grid_Weighted_Average", + ) + ) + console.print( + f"[bold]Optimized Grid Weighted Average Test ({media_type}):[/bold] AUC: {all_experiment_results['Optimized_Grid_Weighted_Average'].get('roc_auc', np.nan):.4f}, F1: {all_experiment_results['Optimized_Grid_Weighted_Average'].get('f1_score', np.nan):.4f}" + ) + + # Save optimized weights to media-type specific subdirectory with generic name + # (or keep media_type in name if preferred, but API loads generic name from subdir) + # opt_weights_api_path_generic = os.path.join(media_type_api_artifacts_subdir, "optimized_grid_average_weights.json") + # For now, keeping the original behavior of saving to main api_artifacts_dir with media_type in name + opt_weights_api_path_typed = os.path.join( + api_artifacts_dir, f"optimized_grid_average_weights_{media_type}.json" + ) + with open(opt_weights_api_path_typed, "w") as f: + json.dump( + { + feat: w + for feat, w in zip(base_model_prob_features, optimized_avg_weights) + }, + f, + indent=2, + ) + console.print( + f"Optimized weights for {media_type} saved to API artifacts: [green]{opt_weights_api_path_typed}[/green]" + ) + else: + console.print( + f"[yellow]Validation set for weight optimization ({media_type}) too small or no features. Skipping.[/yellow]" + ) + + console.rule( + f"[bold green]5. Overall Experiment Summary & Artifacts ({media_type.capitalize()})[/bold green]" + ) + summary_table = Table( + title=f"Meta-Learner & Simple Ensemble Experiment Summary ({media_type.capitalize()} Meta-Test Set)" + ) + summary_table.add_column( + "Method/Model", style="cyan", overflow="fold", max_width=35 + ) + summary_table.add_column("Test AUC", style="magenta") + summary_table.add_column("Test F1", style="green") + summary_table.add_column("Test Acc.", style="blue") + summary_table.add_column("Test Prec.", style="yellow") + summary_table.add_column("Test Recall", style="red") + + sorted_results_list = sorted( + all_experiment_results.items(), + key=lambda item: ( + item[1].get("roc_auc", -1) if pd.notna(item[1].get("roc_auc")) else -1 + ), + reverse=True, + ) + best_method_overall_name = "None" + best_method_overall_auc = -1.0 + best_trainable_ml_model_for_api = None + + for method_name_result, metrics_result in sorted_results_list: + summary_table.add_row( + method_name_result, + ( + f"{metrics_result.get('roc_auc', 'N/A'):.4f}" + if pd.notna(metrics_result.get("roc_auc")) + else "N/A" + ), + f"{metrics_result.get('f1_score', 'N/A'):.4f}", + f"{metrics_result.get('accuracy', 'N/A'):.4f}", + f"{metrics_result.get('precision', 'N/A'):.4f}", + f"{metrics_result.get('recall', 'N/A'):.4f}", + ) + current_auc_val_result = metrics_result.get("roc_auc", -1) + if ( + pd.notna(current_auc_val_result) + and current_auc_val_result > best_method_overall_auc + ): + best_method_overall_auc = current_auc_val_result + best_method_overall_name = method_name_result + if method_name_result in trained_ml_model_objects: + best_trainable_ml_model_for_api = trained_ml_model_objects[ + method_name_result + ] + + console.print(summary_table) + console.print( + f"\n[bold gold1]Best performing method overall for {media_type.upper()} (Test AUC): [white]{best_method_overall_name}[/white] (AUC: {best_method_overall_auc:.4f})[/bold gold1]" + ) + + results_json_path = os.path.join( + experiment_run_output_dir, f"all_experiments_metrics_summary_{media_type}.json" + ) + with open(results_json_path, "w") as f: + json.dump(all_experiment_results, f, indent=2, cls=NpEncoder) + console.print( + f"All experiment metrics summaries for {media_type} saved to [green]{results_json_path}[/green]" + ) + + plot_roc_curves_all( + all_experiment_results, + y_meta_test.values, + experiment_run_output_dir, + media_type, + ) + + console.print( + f"\n[bold]Deployment Artifacts Preparation for {media_type.upper()} (in '{media_type_api_artifacts_subdir}'):[/bold]" + ) + + joblib.dump( + ml_preprocessor.named_steps["imputer"], + os.path.join(media_type_api_artifacts_subdir, "deepsafe_meta_imputer.joblib"), + ) + joblib.dump( + ml_preprocessor.named_steps["scaler"], + os.path.join(media_type_api_artifacts_subdir, "deepsafe_meta_scaler.joblib"), + ) + + api_feature_cols_path = os.path.join( + media_type_api_artifacts_subdir, "deepsafe_meta_feature_columns.json" + ) + if os.path.exists(temp_exp_feature_cols_path): + try: + with ( + open(temp_exp_feature_cols_path, "r") as src_f, + open(api_feature_cols_path, "w") as dst_f, + ): + json.dump(json.load(src_f), dst_f, indent=2) + console.print( + f"Feature columns for {media_type} API saved to [green]{api_feature_cols_path}[/green]" + ) + except Exception as e: + console.print( + f"[red]Error copying/saving feature columns file: {e}. Manual copy might be needed from {temp_exp_feature_cols_path} to {api_feature_cols_path}[/red]" + ) + else: + console.print( + f"[yellow]Temporary feature columns file {temp_exp_feature_cols_path} not found. API artifact for feature columns may be missing for {media_type}.[/yellow]" + ) + + console.print( + f"Common imputer, scaler, and feature columns for {media_type} saved for API in '{media_type_api_artifacts_subdir}'." + ) + + if best_trainable_ml_model_for_api: + api_model_joblib_path = os.path.join( + media_type_api_artifacts_subdir, "deepsafe_meta_learner.joblib" + ) + joblib.dump(best_trainable_ml_model_for_api, api_model_joblib_path) + console.print( + f"Best trainable ML meta-learner ([white]{best_method_overall_name}[/white]) for {media_type} saved as '{os.path.basename(api_model_joblib_path)}' in '{media_type_api_artifacts_subdir}'." + ) + console.print( + f"The 4 artifacts in '{media_type_api_artifacts_subdir}' are ready for the API." + ) + elif best_method_overall_name.startswith(("Simple", "Provided", "Optimized")): + console.print( + f"[yellow]The overall best method for {media_type} ([white]{best_method_overall_name}[/white]) is rule-based.[/yellow]" + ) + console.print( + f"[yellow]To deploy a trainable ML model, choose the best one from this run and ensure its '.joblib' is saved as 'deepsafe_meta_learner.joblib' inside '{media_type_api_artifacts_subdir}'.[/yellow]" + ) + + opt_weights_main_dir_path = os.path.join( + api_artifacts_dir, f"optimized_grid_average_weights_{media_type}.json" + ) + opt_weights_subdir_path_generic = os.path.join( + media_type_api_artifacts_subdir, "optimized_grid_average_weights.json" + ) + + if "Optimized_Grid_Weighted_Average" in best_method_overall_name: + if os.path.exists(opt_weights_main_dir_path): + console.print( + f" Optimized weights for this method are currently in '{opt_weights_main_dir_path}'. Consider standardizing its location if desired (e.g., to '{opt_weights_subdir_path_generic}')." + ) + elif os.path.exists( + opt_weights_subdir_path_generic + ): # If you adjust saving logic for weights too + console.print( + f" Optimized weights for this method are in '{opt_weights_subdir_path_generic}'." + ) + else: + console.print( + f"[bold red]Error: Could not determine a best trainable model to save for {media_type}. Please review results.[/bold red]" + ) + + console.rule( + f"[bold green]Experimentation Suite for {media_type.upper()} Completed[/bold green]" + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Run Meta-Learning Experiments for DeepSafe Ensemble." + ) + parser.add_argument( + "--media-type", + type=str, + choices=["image", "video", "audio"], + required=True, + help="Type of media for which the meta-learner is being trained (image, video, or audio).", + ) + parser.add_argument( + "--meta-file", + type=str, + required=True, + help="Path to the media-specific meta-features CSV (e.g., ./meta_data/meta_features_image.csv)", + ) + parser.add_argument( + "--output-dir", + type=str, + default=DEFAULT_EXPERIMENT_OUTPUT_DIR_BASE, + help=f"Base directory for saving all experiment-related outputs (default: {DEFAULT_EXPERIMENT_OUTPUT_DIR_BASE}).", + ) + parser.add_argument( + "--api-artifacts-dir", + type=str, + default=DEFAULT_API_ARTIFACTS_DIR, + help=f"Directory to save final API-ready artifacts (default: {DEFAULT_API_ARTIFACTS_DIR})", + ) + parser.add_argument( + "--optimizer", + type=str, + choices=["optuna", "gridsearch"], + default=OPTIMIZER_CHOICE_DEFAULT, + help=f"Hyperparameter optimizer (default: {OPTIMIZER_CHOICE_DEFAULT})", + ) + parser.add_argument( + "--optuna-trials", + type=int, + default=N_OPTUNA_TRIALS_DEFAULT, + help=f"Number of Optuna trials (default: {N_OPTUNA_TRIALS_DEFAULT})", + ) + parser.add_argument( + "--weights", + type=str, + default=None, + help='JSON string or path to JSON file for custom base model weights (for "Provided_Weighted_Average"). Keys should be base model names (e.g., "npr_deepfakedetection").', + ) + + args = parser.parse_args() + + if OPTIMIZER_CHOICE_DEFAULT == "optuna" and not OPTIMIZER_AVAILABLE_OPTUNA: + console.print( + "[yellow]Default optimizer is Optuna, but it's not installed. GridSearchCV will be used if Optuna is chosen via CLI and not available.[/yellow]" + ) + if not XGBOOST_AVAILABLE: + console.print( + "[yellow]XGBoost not installed. XGBoost experiments will be skipped if its block is reached.[/yellow]" + ) + if not LIGHTGBM_AVAILABLE: + console.print( + "[yellow]LightGBM not installed. LightGBM experiments will be skipped if its block is reached.[/yellow]" + ) + + custom_weights_dict_main = None + if args.weights: + try: + if os.path.exists(args.weights): + with open(args.weights, "r") as f: + custom_weights_dict_main = json.load(f) + else: + custom_weights_dict_main = json.loads(args.weights) + console.print( + f"Using provided custom base model weights: {custom_weights_dict_main}" + ) + except Exception as e_weights: + console.print( + f"[bold red]Error parsing --weights argument: {e_weights}. Proceeding without them.[/bold red]" + ) + + run_meta_learning_experiments( + meta_features_file=args.meta_file, + output_dir_base=args.output_dir, + api_artifacts_dir=args.api_artifacts_dir, + media_type=args.media_type, + optimizer_type=args.optimizer, + n_optuna_trials_config=args.optuna_trials, + provided_custom_weights=custom_weights_dict_main, + ) diff --git a/image/aide/.gitignore b/image/aide/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..18e583caf7ee16a2331ad502b9bf18445fe20bbf --- /dev/null +++ b/image/aide/.gitignore @@ -0,0 +1,5 @@ +weights/*.pth +weights/*.pt +weights/*/ +hf_cache/ +model_code/ diff --git a/image/aide/Dockerfile b/image/aide/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..7447b1cbec22249801843f3f38b9997910453ce1 --- /dev/null +++ b/image/aide/Dockerfile @@ -0,0 +1,47 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +WORKDIR /app + +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + python3 python3-pip python3-dev \ + git wget build-essential && \ + rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +RUN pip install --no-cache-dir --upgrade pip "setuptools>=68" wheel + +# Install PyTorch with CUDA 12.1 (replaces CPU-only arch-conditional install) +RUN pip install --no-cache-dir \ + torch==2.5.1 torchvision==0.20.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +RUN git clone https://github.com/shilinyan99/AIDE.git model_code && \ + touch model_code/__init__.py && \ + touch model_code/models/__init__.py && \ + touch model_code/data/__init__.py + +RUN mkdir -p /app/weights /app/hf_cache +COPY weights/ /app/weights/ + +COPY app.py . + +ENV MODEL_PORT=5004 +ENV PRELOAD_MODEL=false +ENV MODEL_TIMEOUT=600 +ENV AIDE_CHECKPOINT=GenImage_train.pth +ENV PYTHONPATH=/app/model_code:$PYTHONPATH + +EXPOSE ${MODEL_PORT} + +RUN adduser --disabled-password --gecos '' appuser +USER appuser + +CMD ["python", "app.py"] diff --git a/image/aide/app.py b/image/aide/app.py new file mode 100644 index 0000000000000000000000000000000000000000..4c6c990cad5094a8fa72a528009ca935a52d3c4c --- /dev/null +++ b/image/aide/app.py @@ -0,0 +1,433 @@ +import base64 +import gc +import io +import logging +import os +import platform +import sys +import threading +import time +from typing import Any, Dict, Optional + +import torch +import uvicorn +from fastapi import FastAPI, HTTPException +from fastapi.middleware.cors import CORSMiddleware +from PIL import Image, ImageFile +from pydantic import BaseModel + +ImageFile.LOAD_TRUNCATED_IMAGES = True + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + handlers=[logging.StreamHandler(sys.stdout)], +) +logger = logging.getLogger(__name__) + +# ── Path setup ────────────────────────────────────────────────────────────── +current_dir = os.path.dirname(os.path.abspath(__file__)) +model_code_dir = os.path.join(current_dir, "model_code") +sys.path.insert(0, model_code_dir) +sys.path.insert(0, os.path.join(model_code_dir, "models")) +sys.path.insert(0, os.path.join(model_code_dir, "data")) + +# ── Compatibility shim ─────────────────────────────────────────────────────── +# AIDE's models/AIDE.py imports `clip` (openai-clip) at module level, but the +# package is not used during inference — only open_clip is. The openai-clip +# package relies on pkg_resources which was removed in Python 3.13. We inject +# a lightweight stub so the import succeeds without installing the full package. +import types as _types + +if "clip" not in sys.modules: + _clip_stub = _types.ModuleType("clip") + sys.modules["clip"] = _clip_stub + +# ── Config ────────────────────────────────────────────────────────────────── +MODEL_NAME = "aide_detection" +WEIGHTS_DIR = os.environ.get("WEIGHTS_DIR", os.path.join(current_dir, "weights")) +HF_HOME = os.environ.get("HF_HOME", os.path.join(current_dir, "hf_cache")) +os.environ["HF_HOME"] = HF_HOME +# The checkpoint is self-contained (includes ConvNeXt weights), so we initialise +# the architecture with no pretrained weights and load everything from the checkpoint. +# Set CONVNEXT_PRETRAINED to a HuggingFace tag only if running without a checkpoint. +CONVNEXT_PRETRAINED = os.environ.get("CONVNEXT_PRETRAINED", None) + +# Preferred checkpoint filename (GenImage trains on the most diverse generators) +PREFERRED_CHECKPOINT = os.environ.get("AIDE_CHECKPOINT", "GenImage_train.pth") + + +def _get_device(): + """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if override: + pass # Invalid override, fall through to auto-detect + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + return torch.device("cpu") + + +DEVICE = _get_device() +if DEVICE.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") +if DEVICE.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) +else: + logger.warning( + "Device: %s (no CUDA available -- check nvidia-container-toolkit)", + DEVICE, + ) +PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true" +MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", "600")) + +# ── Globals ───────────────────────────────────────────────────────────────── +model = None +model_lock = threading.Lock() +last_used_time = 0 + + +class ImageInput(BaseModel): + """Request body for /predict.""" + + image_data: str + threshold: Optional[float] = 0.5 + + +# ── Weight discovery ───────────────────────────────────────────────────────── + + +def find_aide_checkpoint() -> Optional[str]: + """ + Return path to the best AIDE checkpoint in WEIGHTS_DIR. + + Priority order: + 1. PREFERRED_CHECKPOINT filename (GenImage_train.pth by default) + 2. Any other .pth file (largest wins) + """ + if not os.path.exists(WEIGHTS_DIR): + logger.warning(f"Weights directory not found: {WEIGHTS_DIR}") + return None + + # Try the preferred checkpoint first + preferred = os.path.join(WEIGHTS_DIR, PREFERRED_CHECKPOINT) + if os.path.exists(preferred): + logger.info(f"Using preferred checkpoint: {preferred}") + return preferred + + # Fall back to the largest available checkpoint + candidates = [ + os.path.join(WEIGHTS_DIR, f) + for f in os.listdir(WEIGHTS_DIR) + if f.endswith(".pth") or f.endswith(".pt") + ] + if not candidates: + logger.warning("No .pth checkpoint found in weights directory.") + return None + best = max(candidates, key=os.path.getsize) + logger.info(f"Using checkpoint: {best}") + return best + + +# ── Preprocessing ──────────────────────────────────────────────────────────── + + +def preprocess_image(image_bytes: bytes) -> torch.Tensor: + """ + Preprocess raw image bytes into AIDE's 5-view tensor. + + Args: + image_bytes: Raw bytes of a JPEG/PNG/etc. image. + + Returns: + Tensor of shape [1, 5, 3, 256, 256] on CPU. + + Raises: + Exception: If bytes cannot be decoded or processed. + """ + from data.dct import DCT_base_Rec_Module + from torchvision import transforms + + pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB") + + # Ensure minimum 256x256 so DCT unfold has enough patches + w, h = pil_image.size + if w < 256 or h < 256: + pil_image = pil_image.resize((256, 256), Image.BICUBIC) + + to_tensor = transforms.ToTensor() + image_tensor = to_tensor(pil_image) # [3, H, W] + + # DCT frequency decomposition → 4 patches [3, 32, 32] each + dct_module = DCT_base_Rec_Module() + x_minmin, x_maxmax, x_minmin1, x_maxmax1 = dct_module(image_tensor) + + # Resize all views to 256×256 and normalise with ImageNet stats + transform = transforms.Compose( + [ + transforms.Resize([256, 256]), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + ] + ) + + x_0 = transform(image_tensor) + x_minmin = transform(x_minmin) + x_maxmax = transform(x_maxmax) + x_minmin1 = transform(x_minmin1) + x_maxmax1 = transform(x_maxmax1) + + # Stack → [5, 3, 256, 256], unsqueeze batch → [1, 5, 3, 256, 256] + stacked = torch.stack([x_minmin, x_maxmax, x_minmin1, x_maxmax1, x_0], dim=0) + return stacked.unsqueeze(0).to(DEVICE) + + +# ── Model loading ───────────────────────────────────────────────────────────── + + +def load_model_internal(): + """Load AIDE_Model onto CPU with the best available checkpoint.""" + global model, last_used_time + + with model_lock: + if model is not None: + last_used_time = time.time() + return + + logger.info("Loading AIDE model...") + try: + import models.AIDE as AIDE_module + + aide_model = AIDE_module.AIDE( + resnet_path=None, + convnext_path=CONVNEXT_PRETRAINED, + ) + aide_model.to(DEVICE) + + checkpoint_path = find_aide_checkpoint() + if checkpoint_path: + logger.info(f"Loading checkpoint: {checkpoint_path}") + ckpt = torch.load(checkpoint_path, map_location=DEVICE) + if isinstance(ckpt, dict): + state_dict = ckpt.get("model") or ckpt.get("state_dict") or ckpt + else: + state_dict = ckpt + # Strip DataParallel "module." prefix if present + cleaned = { + k[7:] if k.startswith("module.") else k: v + for k, v in state_dict.items() + } + missing, unexpected = aide_model.load_state_dict(cleaned, strict=False) + logger.info( + f"Checkpoint loaded. Missing keys: {len(missing)}, " + f"Unexpected keys: {len(unexpected)}" + ) + else: + logger.warning( + "No checkpoint found — model uses pretrained-only weights. " + "Run download_weights.sh to fetch the fine-tuned checkpoint." + ) + + # Switch to inference mode (no gradient tracking, batch-norm uses running stats) + aide_model.train(mode=False) + model = aide_model + last_used_time = time.time() + logger.info("AIDE model ready.") + + except Exception as exc: + logger.exception(f"Failed to load AIDE model: {exc}") + model = None + raise + finally: + gc.collect() + + +def ensure_model_loaded(): + """Load model on first request (lazy loading).""" + global last_used_time + if model is None: + load_model_internal() + else: + last_used_time = time.time() + + +def unload_model_if_idle(): + """Evict model from RAM after MODEL_TIMEOUT seconds of inactivity.""" + global model + if model is None or PRELOAD_MODEL: + return + with model_lock: + if model is not None and (time.time() - last_used_time > MODEL_TIMEOUT): + logger.info("Unloading idle AIDE model to free RAM.") + del model + model = None + gc.collect() + + +# ── FastAPI app ─────────────────────────────────────────────────────────────── + +app = FastAPI( + title="AIDE Deepfake Detection Service", + description=( + "AI-generated image detection using AIDE (ICLR 2025) — " + "hybrid DCT frequency analysis + ConvNeXt-xxlarge semantic features." + ), + version="1.0.0", +) + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + + +@app.get("/") +async def root(): + """Root endpoint with service information.""" + return { + "model_name": MODEL_NAME, + "description": "AIDE (ICLR 2025) AI-generated image detector", + "device": str(DEVICE), + "model_loaded": model is not None, + } + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and DEVICE.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +async def health(): + """Health check endpoint.""" + return { + "status": "healthy", + "model_name": MODEL_NAME, + "device": str(DEVICE), + "model_loaded": model is not None, + **_gpu_health_info(), + } + + +@app.post("/unload") +async def unload_model_endpoint(): + """Manually unload the model to free RAM.""" + global model + if model is None: + return {"status": "not_loaded"} + del model + model = None + gc.collect() + return {"status": "success", "message": "Model unloaded."} + + +@app.post("/predict") +async def predict(image_input: ImageInput) -> Dict[str, Any]: + """ + Predict whether the submitted image is AI-generated. + + Args: + image_input: Base64-encoded image and optional classification threshold. + + Returns: + Dict with model name, fake probability, binary prediction, class label, and inference time. + """ + try: + ensure_model_loaded() + if model is None: + raise HTTPException(status_code=503, detail="Model not loaded.") + + start = time.time() + + try: + image_bytes = base64.b64decode(image_input.image_data) + input_tensor = preprocess_image(image_bytes) + except Exception as exc: + raise HTTPException(status_code=400, detail=f"Invalid image data: {exc}") + + with torch.no_grad(): + logits = model(input_tensor) # [1, 2] + probs = torch.softmax(logits, dim=-1) # [1, 2] + probability_fake = probs[0, 1].item() + + prediction = 1 if probability_fake >= image_input.threshold else 0 + class_label = "fake" if prediction == 1 else "real" + inference_time = time.time() - start + + logger.info( + f"Prediction: {class_label} (prob={probability_fake:.4f}, {inference_time:.3f}s)" + ) + + if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: + threading.Timer(MODEL_TIMEOUT + 5.0, unload_model_if_idle).start() + + return { + "model": MODEL_NAME, + "probability": float(probability_fake), + "prediction": int(prediction), + "class": class_label, + "inference_time": float(inference_time), + } + + except HTTPException: + raise + except Exception as exc: + logger.exception(f"Prediction error: {exc}") + raise HTTPException(status_code=500, detail=str(exc)) + + +@app.on_event("startup") +async def startup_event(): + """Startup handler — preloads model if PRELOAD_MODEL=true, else lazy-loads.""" + if PRELOAD_MODEL: + logger.info("Preloading AIDE model at startup.") + try: + load_model_internal() + except Exception as exc: + logger.error(f"Preload failed: {exc}") + else: + logger.info("AIDE service ready — model loads on first request.") + + if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: + + def _periodic_check(): + unload_model_if_idle() + threading.Timer(MODEL_TIMEOUT / 2.0, _periodic_check).start() + + threading.Timer(MODEL_TIMEOUT / 2.0, _periodic_check).start() + + +if __name__ == "__main__": + port = int(os.environ.get("MODEL_PORT", 5004)) + logger.info(f"Starting AIDE service on port {port}") + uvicorn.run(app, host="0.0.0.0", port=port) diff --git a/image/aide/download_weights.sh b/image/aide/download_weights.sh new file mode 100644 index 0000000000000000000000000000000000000000..653964a49305f74cc7d5ba55aaaaa400e7e1f776 --- /dev/null +++ b/image/aide/download_weights.sh @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +# download_weights.sh — Download all weights required by the AIDE service. +# +# Run ONCE from the repo root before building/running the service: +# bash services/image/aide/download_weights.sh +# +# Downloads: +# 1. AIDE fine-tuned checkpoint from Google Drive -> services/image/aide/weights/ +# 2. ConvNeXt-xxlarge OpenCLIP model (~4.8 GB) -> services/image/aide/hf_cache/ + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +WEIGHTS_DIR="$SCRIPT_DIR/weights" +HF_CACHE_DIR="$SCRIPT_DIR/hf_cache" + +mkdir -p "$WEIGHTS_DIR" "$HF_CACHE_DIR" + +echo "============================================================" +echo " DeepSafe - AIDE Weight Downloader" +echo "============================================================" + +# ── 1. AIDE checkpoint from Google Drive ────────────────────────────────── +echo "" +echo "[1/2] Downloading AIDE fine-tuned checkpoint from Google Drive..." + +pip install gdown -q + +FOLDER_ID="1qx76UFvDpgCxaPLBCmsA2WY-SSzeJrd4" + +gdown --folder "https://drive.google.com/drive/folders/${FOLDER_ID}" \ + --output "$WEIGHTS_DIR" \ + --remaining-ok || { + echo "" + echo " [WARN] Automatic download failed." + echo " Please manually download the checkpoint from:" + echo " https://drive.google.com/drive/folders/${FOLDER_ID}" + echo " and place the .pth file(s) in: $WEIGHTS_DIR" + echo "" +} + +echo "" +echo " Contents of weights dir:" +ls -lh "$WEIGHTS_DIR" 2>/dev/null || echo " (empty — manual download required)" + +# ── 2. ConvNeXt-xxlarge via open_clip (HuggingFace Hub, ~4.8 GB) ────────── +echo "" +echo "[2/2] Pre-caching ConvNeXt-xxlarge OpenCLIP model (~4.8 GB)..." +echo " This downloads once and is reused across container restarts." +echo "" + +export HF_HOME="$HF_CACHE_DIR" + +python3 -c " +import os, sys +os.environ['HF_HOME'] = os.environ.get('HF_HOME', '.') +print(' Importing open_clip...') +import open_clip +print(' Downloading convnext_xxlarge (laion2b_s34b_b82k_augreg)...') +m, _, _ = open_clip.create_model_and_transforms( + 'convnext_xxlarge', + pretrained='laion2b_s34b_b82k_augreg', +) +del m +print(' ConvNeXt-xxlarge cached successfully.') +" || { + echo " [WARN] open_clip pre-cache failed." + echo " The container will download ConvNeXt-xxlarge (~4.8 GB) on first startup." + echo " Ensure hf_cache/ is volume-mounted so it only downloads once." +} + +echo "" +echo "============================================================" +echo " Done." +echo " AIDE checkpoint : $WEIGHTS_DIR" +echo " ConvNeXt cache : $HF_CACHE_DIR" +echo "============================================================" diff --git a/image/aide/requirements.txt b/image/aide/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..4e67ff6d89453c35d1570bf9ce2e36affa236e39 --- /dev/null +++ b/image/aide/requirements.txt @@ -0,0 +1,13 @@ +fastapi==0.103.1 +uvicorn[standard]==0.23.2 +python-multipart +pydantic<2.0.0 +pillow==9.5.0 +numpy==1.24.3 +scipy==1.11.4 +scikit-image==0.20.0 +einops==0.6.1 +kornia==0.7.2 +open-clip-torch==2.24.0 +gdown +timm==0.9.6 diff --git a/image/cospy/Dockerfile b/image/cospy/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..18d01b5ed68744299dc0d7e71274a77a4f498bdd --- /dev/null +++ b/image/cospy/Dockerfile @@ -0,0 +1,49 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +WORKDIR /app + +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + python3 python3-pip python3-dev \ + git build-essential && \ + rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +RUN pip install --no-cache-dir --upgrade pip "setuptools>=68" wheel + +# Install PyTorch with CUDA 12.1 +RUN pip install --no-cache-dir \ + torch==2.5.1 torchvision==0.20.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +RUN git clone https://github.com/Megum1/CO-SPY.git model_code && \ + cd model_code && \ + git checkout 9fb781d && \ + cd .. + +RUN mkdir -p /app/weights /app/hf_cache +COPY weights/ /app/weights/ + +COPY app.py . + +ENV MODEL_PORT=5007 +ENV PRELOAD_MODEL=true +ENV MODEL_TIMEOUT=600 +ENV WEIGHTS_DIR=/app/weights +ENV HF_HOME=/app/hf_cache +ENV PYTHONPATH=/app/model_code:$PYTHONPATH + +EXPOSE ${MODEL_PORT} + +RUN adduser --disabled-password --gecos '' appuser && \ + chown -R appuser:appuser /app/hf_cache +USER appuser + +CMD ["python", "app.py"] diff --git a/image/cospy/app.py b/image/cospy/app.py new file mode 100644 index 0000000000000000000000000000000000000000..09c40e24b9a22479c8d9436049ce60ff1c9dece7 --- /dev/null +++ b/image/cospy/app.py @@ -0,0 +1,447 @@ +"""CO-SPY (CVPR 2025, Sony Research) AI-generated image detection service. + +Wraps the CO-SPY dual-branch fusion detector with a FastAPI endpoint +matching the DeepSafe contract. CO-SPY combines: + - Semantic branch: frozen SigLIP ViT-SO400M-14-SigLIP-384 (open_clip) + producing 1152-dim features -> Linear(1152, 1) + - Artifact branch: SD v1-4 VAE encode+decode reconstruction error fed + through a ResNet encoder producing 512-dim features -> Linear(512, 1) + - Fusion layer: Linear(2, 1) calibrating the two branch logits + +Key design decisions: +- Model architecture code is loaded from the cloned CO-SPY repo via + sys.path at startup. The API layer in this file is fully self-contained. +- Preprocessing mirrors the fusion detector's test_transform: + Resize(384) -> CenterCrop(384) -> ToTensor(), with per-branch + normalisation applied in the forward pass. +- Weights are downloaded from HuggingFace at model load time, not + at build time. Three weight files are needed: semantic_weights.pth, + artifact_weights.pth, fusion_weights.pth. +- The SD v1-4 VAE and SigLIP model are loaded lazily by the detector + constructors (heavy; ~2-3 GB combined). +""" + +import base64 +import gc +import io +import logging +import os +import platform +import sys +import threading +import time +from typing import Any, Dict, Optional + +import numpy as np +import torch +import uvicorn +from fastapi import FastAPI, HTTPException +from fastapi.middleware.cors import CORSMiddleware +from PIL import Image, ImageFile +from pydantic import BaseModel +from torchvision import transforms + +ImageFile.LOAD_TRUNCATED_IMAGES = True + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + handlers=[logging.StreamHandler(sys.stdout)], +) +logger = logging.getLogger(__name__) + +# -- Path setup --------------------------------------------------------------- +current_dir = os.path.dirname(os.path.abspath(__file__)) +model_code_dir = os.path.join(current_dir, "model_code") + +# Add the cloned CO-SPY repo root so that ``detectors``, ``utils``, and +# ``dataSets`` resolve during import. +if model_code_dir not in sys.path: + sys.path.insert(0, model_code_dir) + +# -- Config ------------------------------------------------------------------- +MODEL_NAME = "cospy_detection" +WEIGHTS_DIR = os.environ.get("WEIGHTS_DIR", os.path.join(current_dir, "weights")) +HF_HOME = os.environ.get("HF_HOME", os.path.join(current_dir, "hf_cache")) +os.environ["HF_HOME"] = HF_HOME + +# HuggingFace repo containing the three weight files. +HF_WEIGHTS_REPO = "ruojiruoli/Co-Spy-Pretrained-Weights" +HF_WEIGHTS_SUBDIR = "sd-v1_4" + +# Expected checkpoint filenames inside WEIGHTS_DIR// +SEMANTIC_CKPT = "semantic_weights.pth" +ARTIFACT_CKPT = "artifact_weights.pth" +FUSION_CKPT = "fusion_weights.pth" + +# Fusion preprocessing constants (matches fusion test_transform) +LOAD_SIZE = 384 +CROP_SIZE = 384 + + +def _get_device(): + """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if override: + pass # Invalid override, fall through to auto-detect + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + return torch.device("cpu") + + +DEVICE = _get_device() +if DEVICE.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + +if DEVICE.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) +else: + logger.warning( + "Device: %s (no CUDA available -- check nvidia-container-toolkit)", + DEVICE, + ) + +PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true" +MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", "600")) + +# -- Globals ------------------------------------------------------------------ +model = None +model_lock = threading.Lock() +last_used_time = 0 + + +class ImageInput(BaseModel): + """Request body for /predict.""" + + image_data: str + threshold: Optional[float] = 0.5 + + +# -- Weight download ---------------------------------------------------------- + + +def _ensure_weights_downloaded() -> str: + """Download CO-SPY pretrained weights from HuggingFace if missing. + + Returns: + Path to the directory containing the three weight files. + """ + weights_subdir = os.path.join(WEIGHTS_DIR, HF_WEIGHTS_SUBDIR) + semantic_path = os.path.join(weights_subdir, SEMANTIC_CKPT) + artifact_path = os.path.join(weights_subdir, ARTIFACT_CKPT) + fusion_path = os.path.join(weights_subdir, FUSION_CKPT) + + if ( + os.path.exists(semantic_path) + and os.path.exists(artifact_path) + and os.path.exists(fusion_path) + ): + logger.info("All CO-SPY weight files found in %s", weights_subdir) + return weights_subdir + + logger.info( + "Downloading CO-SPY weights from HuggingFace (%s)...", + HF_WEIGHTS_REPO, + ) + os.makedirs(weights_subdir, exist_ok=True) + + try: + from huggingface_hub import hf_hub_download + except ImportError: + raise RuntimeError( + "huggingface_hub is required to download weights. " + "Install with: pip install huggingface_hub" + ) + + for fname in (SEMANTIC_CKPT, ARTIFACT_CKPT, FUSION_CKPT): + remote_path = f"{HF_WEIGHTS_SUBDIR}/{fname}" + dest = os.path.join(weights_subdir, fname) + if not os.path.exists(dest): + logger.info(" Downloading %s ...", remote_path) + downloaded = hf_hub_download( + repo_id=HF_WEIGHTS_REPO, + filename=remote_path, + local_dir=WEIGHTS_DIR, + ) + logger.info(" Saved to %s", downloaded) + + return weights_subdir + + +# -- Preprocessing ------------------------------------------------------------ + + +def preprocess_image(image_bytes: bytes) -> torch.Tensor: + """Preprocess raw image bytes into a CO-SPY fusion-compatible tensor. + + Mirrors the fusion detector's test_transform: + Resize(384) -> CenterCrop(384) -> ToTensor(). Per-branch normalisation + (semantic: [0.5,0.5,0.5]/[0.5,0.5,0.5]; artifact: resize to 224 + + identity normalisation) is applied inside the model's forward() method. + + Args: + image_bytes: Raw bytes of a JPEG/PNG/etc. image. + + Returns: + Tensor of shape [1, 3, 384, 384]. + + Raises: + Exception: If bytes cannot be decoded or processed. + """ + pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB") + + transform = transforms.Compose( + [ + transforms.Resize(LOAD_SIZE), + transforms.CenterCrop(CROP_SIZE), + transforms.ToTensor(), + ] + ) + + tensor = transform(pil_image) # [3, 384, 384] + return tensor.unsqueeze(0) # [1, 3, 384, 384] + + +# -- Model loading ------------------------------------------------------------ + + +def load_model_internal(): + """Load the CO-SPY fusion detector onto the selected device.""" + global model, last_used_time + + with model_lock: + if model is not None: + last_used_time = time.time() + return + + logger.info("Loading CO-SPY fusion model...") + try: + # Ensure weights are present (download from HF if needed) + weights_subdir = _ensure_weights_downloaded() + + semantic_path = os.path.join(weights_subdir, SEMANTIC_CKPT) + artifact_path = os.path.join(weights_subdir, ARTIFACT_CKPT) + fusion_path = os.path.join(weights_subdir, FUSION_CKPT) + + # Import from the cloned CO-SPY repo. The directory name + # ``sd-v1_4`` contains a hyphen, so we must use importlib. + import importlib + + detector_module = importlib.import_module("detectors.sd-v1_4") + CoSpyFusionDetector = detector_module.CoSpyFusionDetector + + cospy_model = CoSpyFusionDetector( + semantic_weights_path=semantic_path, + artifact_weights_path=artifact_path, + ) + cospy_model.load_weights(fusion_path) + cospy_model.to(DEVICE) + cospy_model.train(mode=False) + + model = cospy_model + last_used_time = time.time() + logger.info("CO-SPY fusion model ready on %s.", DEVICE) + + except Exception as exc: + logger.exception("Failed to load CO-SPY model: %s", exc) + model = None + raise + finally: + gc.collect() + + +def ensure_model_loaded(): + """Load model on first request (lazy loading).""" + global last_used_time + if model is None: + load_model_internal() + else: + last_used_time = time.time() + + +def unload_model_if_idle(): + """Evict model from RAM after MODEL_TIMEOUT seconds of inactivity.""" + global model + if model is None or PRELOAD_MODEL: + return + with model_lock: + if model is not None and (time.time() - last_used_time > MODEL_TIMEOUT): + logger.info("Unloading idle CO-SPY model to free RAM.") + del model + model = None + gc.collect() + + +# -- FastAPI app -------------------------------------------------------------- + +app = FastAPI( + title="CO-SPY Deepfake Detection Service", + description=( + "AI-generated image detection using CO-SPY (CVPR 2025) -- " + "dual-branch fusion of SigLIP semantic features and SD VAE " + "reconstruction artifact features." + ), + version="1.0.0", +) + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + + +@app.get("/") +async def root(): + """Root endpoint with service information.""" + return { + "model_name": MODEL_NAME, + "description": ("CO-SPY (CVPR 2025) AI-generated image detector"), + "device": str(DEVICE), + "model_loaded": model is not None, + } + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and DEVICE.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +async def health(): + """Health check endpoint.""" + return { + "status": "healthy", + "model": MODEL_NAME, + "device": str(DEVICE), + "model_loaded": model is not None, + **_gpu_health_info(), + } + + +@app.post("/unload") +async def unload_model_endpoint(): + """Manually unload the model to free RAM.""" + global model + if model is None: + return {"status": "not_loaded"} + del model + model = None + gc.collect() + return {"status": "success", "message": "Model unloaded."} + + +@app.post("/predict") +async def predict(image_input: ImageInput) -> Dict[str, Any]: + """Predict whether the submitted image is AI-generated. + + Args: + image_input: Base64-encoded image and optional threshold. + + Returns: + Dict with model name, fake probability, binary prediction, + class label, and inference time. + """ + try: + ensure_model_loaded() + if model is None: + raise HTTPException(status_code=503, detail="Model not loaded.") + + start = time.time() + + try: + image_bytes = base64.b64decode(image_input.image_data) + input_tensor = preprocess_image(image_bytes) + input_tensor = input_tensor.to(DEVICE) + except Exception as exc: + raise HTTPException(status_code=400, detail=f"Invalid image data: {exc}") + + with torch.no_grad(): + output = model(input_tensor) # raw logit [B, 1] + probability_fake = output.sigmoid().item() # scalar in [0, 1] + + prediction = 1 if probability_fake >= image_input.threshold else 0 + class_label = "fake" if prediction == 1 else "real" + inference_time = time.time() - start + + logger.info( + "Prediction: %s (prob=%.4f, %.3fs)", + class_label, + probability_fake, + inference_time, + ) + + if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: + threading.Timer(MODEL_TIMEOUT + 5.0, unload_model_if_idle).start() + + return { + "model": MODEL_NAME, + "probability": float(probability_fake), + "prediction": int(prediction), + "class": class_label, + "inference_time": float(inference_time), + } + + except HTTPException: + raise + except Exception as exc: + logger.exception("Prediction error: %s", exc) + raise HTTPException(status_code=500, detail=str(exc)) + + +@app.on_event("startup") +async def startup_event(): + """Preload model if PRELOAD_MODEL=true, else lazy-load on first request.""" + if PRELOAD_MODEL: + logger.info("Preloading CO-SPY model at startup.") + try: + load_model_internal() + except Exception as exc: + logger.error("Preload failed: %s", exc) + else: + logger.info("CO-SPY service ready -- model loads on first request.") + + if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: + + def _periodic_check(): + unload_model_if_idle() + threading.Timer(MODEL_TIMEOUT / 2.0, _periodic_check).start() + + threading.Timer(MODEL_TIMEOUT / 2.0, _periodic_check).start() + + +if __name__ == "__main__": + port = int(os.environ.get("MODEL_PORT", 5007)) + logger.info("Starting CO-SPY service on port %d", port) + uvicorn.run(app, host="0.0.0.0", port=port) diff --git a/image/cospy/requirements.txt b/image/cospy/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..e480d664fafeaaad82d3d703081a2752b5562e98 --- /dev/null +++ b/image/cospy/requirements.txt @@ -0,0 +1,16 @@ +torch>=2.0.0 +torchvision>=0.15.0 +fastapi==0.103.1 +uvicorn[standard]==0.23.2 +python-multipart +pydantic>=2.0.0 +pillow>=9.5.0 +numpy>=1.24.3 +opencv-python-headless>=4.8.0 +open-clip-torch>=2.24.0 +diffusers==0.27.2 +huggingface_hub==0.23.5 +transformers==4.38.2 +accelerate>=0.25.0 +scikit-learn>=1.3.0 +scipy>=1.10.0 diff --git a/image/effort/Dockerfile b/image/effort/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..11376bf8595c18433a1f35ce7e437dd484a586b5 --- /dev/null +++ b/image/effort/Dockerfile @@ -0,0 +1,50 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +WORKDIR /app + +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + python3 python3-pip python3-dev \ + git build-essential && \ + rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +RUN pip install --no-cache-dir --upgrade pip "setuptools>=68" wheel + +# Install PyTorch with CUDA 12.1 +RUN pip install --no-cache-dir \ + torch==2.5.1 torchvision==0.20.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +RUN git clone https://github.com/YZY-stack/Effort-AIGI-Detection.git model_code && \ + cd model_code && \ + git checkout 96f5dea2b534d400cfd7003f053c7e93c8e16461 && \ + cd .. + +RUN mkdir -p /app/weights /app/hf_cache +COPY weights/ /app/weights/ + +COPY app.py . + +ENV MODEL_PORT=5006 +ENV PRELOAD_MODEL=false +ENV MODEL_TIMEOUT=600 +ENV EFFORT_CHECKPOINT=genimage_effort.pth +ENV HF_HOME=/app/hf_cache +ENV TRANSFORMERS_CACHE=/app/hf_cache/hub +ENV PYTHONPATH=/app/model_code/DeepfakeBench/training:/app/model_code/DeepfakeBench/training/detectors:/app/model_code/DeepfakeBench/training/utils:$PYTHONPATH + +EXPOSE ${MODEL_PORT} + +RUN adduser --disabled-password --gecos '' appuser && \ + chown -R appuser:appuser /app/hf_cache +USER appuser + +CMD ["python", "app.py"] diff --git a/image/effort/app.py b/image/effort/app.py new file mode 100644 index 0000000000000000000000000000000000000000..256b24e13fcde9a5be3cb8e8643b9e82df2b13b2 --- /dev/null +++ b/image/effort/app.py @@ -0,0 +1,501 @@ +"""Effort (ICML 2025) AI-generated image detection service. + +Wraps the Effort detector — SVD orthogonal subspace decomposition on +CLIP ViT-L/14 — with a FastAPI endpoint matching the DeepSafe contract. + +Key design decisions: +- Model architecture code (EffortDetector, SVDResidualLinear, + apply_svd_residual_to_self_attn) is loaded from the cloned repo via + sys.path at startup. The API layer in this file is fully self-contained. +- Preprocessing mirrors demo.py:preprocess_face() — resize to 224x224, + CLIP normalisation ([0.4815, 0.4578, 0.4082], [0.2686, 0.2613, 0.2758]). +- No face-alignment (dlib) dependency: the image is resized directly, + matching the repo's non-landmark path (landmark_model=False). +- Weights are volume-mounted at /app/weights/ — not baked into the image. +""" + +import base64 +import gc +import io +import logging +import math +import os +import platform +import sys +import threading +import time +from typing import Any, Dict, Optional + +import cv2 +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import uvicorn +from fastapi import FastAPI, HTTPException +from fastapi.middleware.cors import CORSMiddleware +from PIL import Image, ImageFile +from pydantic import BaseModel +from torchvision import transforms + +ImageFile.LOAD_TRUNCATED_IMAGES = True + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + handlers=[logging.StreamHandler(sys.stdout)], +) +logger = logging.getLogger(__name__) + +# ── Path setup ────────────────────────────────────────────────────────────── +current_dir = os.path.dirname(os.path.abspath(__file__)) +model_code_dir = os.path.join(current_dir, "model_code") + +# Add model code paths so the repo's detector / network / loss modules resolve +for _p in [ + model_code_dir, + os.path.join(model_code_dir, "training"), + os.path.join(model_code_dir, "training", "detectors"), + os.path.join(model_code_dir, "training", "networks"), + os.path.join(model_code_dir, "training", "loss"), + os.path.join(model_code_dir, "training", "utils"), + os.path.join(model_code_dir, "training", "metrics"), +]: + if _p not in sys.path: + sys.path.insert(0, _p) + +# ── Config ────────────────────────────────────────────────────────────────── +MODEL_NAME = "effort_detection" +WEIGHTS_DIR = os.environ.get("WEIGHTS_DIR", os.path.join(current_dir, "weights")) +HF_HOME = os.environ.get("HF_HOME", os.path.join(current_dir, "hf_cache")) +os.environ["HF_HOME"] = HF_HOME +os.environ["TRANSFORMERS_CACHE"] = os.path.join(HF_HOME, "hub") + +# Preferred checkpoint filename +PREFERRED_CHECKPOINT = os.environ.get("EFFORT_CHECKPOINT", "genimage_effort.pth") + +# CLIP normalization constants (from the Effort config / OpenAI CLIP) +CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073] +CLIP_STD = [0.26862954, 0.26130258, 0.27577711] +INPUT_RESOLUTION = 224 + + +def _get_device(): + """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if override: + pass # Invalid override, fall through to auto-detect + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + return torch.device("cpu") + + +DEVICE = _get_device() +if DEVICE.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + +if DEVICE.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) +else: + logger.warning( + "Device: %s (no CUDA available -- check nvidia-container-toolkit)", + DEVICE, + ) + +PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true" +MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", "600")) + +# ── Globals ───────────────────────────────────────────────────────────────── +model = None +model_lock = threading.Lock() +last_used_time = 0 + + +class ImageInput(BaseModel): + """Request body for /predict.""" + + image_data: str + threshold: Optional[float] = 0.5 + + +# ── Monkey-patch CLIP loading ────────────────────────────────────────────── +# The Effort repo's effort_detector.py hardcodes a relative path +# CLIPModel.from_pretrained("../models--openai--clip-vit-large-patch14") +# which is an invalid HuggingFace repo ID. We intercept from_pretrained and +# redirect to the canonical hub identifier so transformers downloads/caches +# the model correctly. +_CLIP_PATH_FIXUPS = { + "../models--openai--clip-vit-large-patch14": "openai/clip-vit-large-patch14", + "models--openai--clip-vit-large-patch14": "openai/clip-vit-large-patch14", +} + + +def _patch_clip_from_pretrained(): + """Wrap CLIPModel.from_pretrained to fix hardcoded relative paths.""" + from transformers import CLIPModel + + _original_from_pretrained = CLIPModel.from_pretrained.__func__ + + @classmethod # type: ignore[misc] + def _patched_from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): + fixed = _CLIP_PATH_FIXUPS.get( + pretrained_model_name_or_path, pretrained_model_name_or_path + ) + if fixed != pretrained_model_name_or_path: + logger.info( + "Intercepted CLIP path '%s' -> '%s'", + pretrained_model_name_or_path, + fixed, + ) + return _original_from_pretrained(cls, fixed, *args, **kwargs) + + CLIPModel.from_pretrained = _patched_from_pretrained + + +_patch_clip_from_pretrained() + +# ── Weight discovery ──────────────────────────────────────────────────────── + + +def find_effort_checkpoint() -> Optional[str]: + """Return path to the best Effort checkpoint in WEIGHTS_DIR. + + Priority order: + 1. PREFERRED_CHECKPOINT filename + 2. Any other .pth / .pt file (largest wins) + """ + if not os.path.exists(WEIGHTS_DIR): + logger.warning("Weights directory not found: %s", WEIGHTS_DIR) + return None + + preferred = os.path.join(WEIGHTS_DIR, PREFERRED_CHECKPOINT) + if os.path.exists(preferred): + logger.info("Using preferred checkpoint: %s", preferred) + return preferred + + candidates = [ + os.path.join(WEIGHTS_DIR, f) + for f in os.listdir(WEIGHTS_DIR) + if f.endswith(".pth") or f.endswith(".pt") + ] + if not candidates: + logger.warning("No .pth checkpoint found in weights directory.") + return None + best = max(candidates, key=os.path.getsize) + logger.info("Using checkpoint: %s", best) + return best + + +# ── Preprocessing ─────────────────────────────────────────────────────────── + + +def preprocess_image(image_bytes: bytes) -> torch.Tensor: + """Preprocess raw image bytes into an Effort-compatible tensor. + + Mirrors demo.py:preprocess_face() — resize to 224x224, convert + to PIL, apply CLIP normalisation. + + Args: + image_bytes: Raw bytes of a JPEG/PNG/etc. image. + + Returns: + Tensor of shape [1, 3, 224, 224]. + + Raises: + Exception: If bytes cannot be decoded or processed. + """ + pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB") + pil_image = pil_image.resize((INPUT_RESOLUTION, INPUT_RESOLUTION), Image.BICUBIC) + + transform = transforms.Compose( + [ + transforms.ToTensor(), + transforms.Normalize(mean=CLIP_MEAN, std=CLIP_STD), + ] + ) + + tensor = transform(pil_image) # [3, 224, 224] + return tensor.unsqueeze(0) # [1, 3, 224, 224] + + +# ── Model loading ─────────────────────────────────────────────────────────── + + +def _build_effort_config() -> dict: + """Build the minimal config dict required by EffortDetector.__init__.""" + return { + "model_name": "effort", + "backbone_name": "vit", + "pretrained": None, + "backbone_config": { + "mode": "original", + "num_classes": 2, + "inc": 3, + "dropout": False, + }, + "resolution": INPUT_RESOLUTION, + "mean": CLIP_MEAN, + "std": CLIP_STD, + "loss_func": "cross_entropy", + } + + +def load_model_internal(): + """Load EffortDetector onto the selected device.""" + global model, last_used_time + + with model_lock: + if model is not None: + last_used_time = time.time() + return + + logger.info("Loading Effort model...") + try: + # Import detector from the cloned repo code + from detectors.effort_detector import ( + EffortDetector, + ) + + cfg = _build_effort_config() + effort_model = EffortDetector(config=cfg) + effort_model.to(DEVICE) + + checkpoint_path = find_effort_checkpoint() + if checkpoint_path: + logger.info("Loading checkpoint: %s", checkpoint_path) + ckpt = torch.load(checkpoint_path, map_location=DEVICE) + if isinstance(ckpt, dict): + state_dict = ckpt.get("model") or ckpt.get("state_dict") or ckpt + else: + state_dict = ckpt + # Strip DataParallel "module." prefix if present + cleaned = { + k[7:] if k.startswith("module.") else k: v + for k, v in state_dict.items() + } + missing, unexpected = effort_model.load_state_dict( + cleaned, strict=False + ) + logger.info( + "Checkpoint loaded. Missing keys: %d, " "Unexpected keys: %d", + len(missing), + len(unexpected), + ) + else: + logger.warning( + "No checkpoint found -- model uses pretrained-only " + "weights. Download the fine-tuned checkpoint for " + "accurate predictions." + ) + + effort_model.train(mode=False) + model = effort_model + last_used_time = time.time() + logger.info("Effort model ready on %s.", DEVICE) + + except Exception as exc: + logger.exception("Failed to load Effort model: %s", exc) + model = None + raise + finally: + gc.collect() + + +def ensure_model_loaded(): + """Load model on first request (lazy loading).""" + global last_used_time + if model is None: + load_model_internal() + else: + last_used_time = time.time() + + +def unload_model_if_idle(): + """Evict model from RAM after MODEL_TIMEOUT seconds of inactivity.""" + global model + if model is None or PRELOAD_MODEL: + return + with model_lock: + if model is not None and (time.time() - last_used_time > MODEL_TIMEOUT): + logger.info("Unloading idle Effort model to free RAM.") + del model + model = None + gc.collect() + + +# ── FastAPI app ───────────────────────────────────────────────────────────── + +app = FastAPI( + title="Effort Deepfake Detection Service", + description=( + "AI-generated image detection using Effort (ICML 2025 Oral) -- " + "SVD orthogonal subspace decomposition on CLIP ViT-L/14." + ), + version="1.0.0", +) + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + + +@app.get("/") +async def root(): + """Root endpoint with service information.""" + return { + "model_name": MODEL_NAME, + "description": ("Effort (ICML 2025 Oral) AI-generated image detector"), + "device": str(DEVICE), + "model_loaded": model is not None, + } + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and DEVICE.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +async def health(): + """Health check endpoint.""" + return { + "status": "healthy", + "model": MODEL_NAME, + "device": str(DEVICE), + "model_loaded": model is not None, + **_gpu_health_info(), + } + + +@app.post("/unload") +async def unload_model_endpoint(): + """Manually unload the model to free RAM.""" + global model + if model is None: + return {"status": "not_loaded"} + del model + model = None + gc.collect() + return {"status": "success", "message": "Model unloaded."} + + +@app.post("/predict") +async def predict(image_input: ImageInput) -> Dict[str, Any]: + """Predict whether the submitted image is AI-generated. + + Args: + image_input: Base64-encoded image and optional threshold. + + Returns: + Dict with model name, fake probability, binary prediction, + class label, and inference time. + """ + try: + ensure_model_loaded() + if model is None: + raise HTTPException(status_code=503, detail="Model not loaded.") + + start = time.time() + + try: + image_bytes = base64.b64decode(image_input.image_data) + input_tensor = preprocess_image(image_bytes) + input_tensor = input_tensor.to(DEVICE) + except Exception as exc: + raise HTTPException(status_code=400, detail=f"Invalid image data: {exc}") + + with torch.no_grad(): + data_dict = { + "image": input_tensor, + "label": torch.tensor([0]).to(DEVICE), + } + preds = model(data_dict, inference=True) + probability_fake = preds["prob"].squeeze().cpu().item() + + prediction = 1 if probability_fake >= image_input.threshold else 0 + class_label = "fake" if prediction == 1 else "real" + inference_time = time.time() - start + + logger.info( + "Prediction: %s (prob=%.4f, %.3fs)", + class_label, + probability_fake, + inference_time, + ) + + if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: + threading.Timer(MODEL_TIMEOUT + 5.0, unload_model_if_idle).start() + + return { + "model": MODEL_NAME, + "probability": float(probability_fake), + "prediction": int(prediction), + "class": class_label, + "inference_time": float(inference_time), + } + + except HTTPException: + raise + except Exception as exc: + logger.exception("Prediction error: %s", exc) + raise HTTPException(status_code=500, detail=str(exc)) + + +@app.on_event("startup") +async def startup_event(): + """Preload model if PRELOAD_MODEL=true, else lazy-load on first request.""" + if PRELOAD_MODEL: + logger.info("Preloading Effort model at startup.") + try: + load_model_internal() + except Exception as exc: + logger.error("Preload failed: %s", exc) + else: + logger.info("Effort service ready -- model loads on first request.") + + if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: + + def _periodic_check(): + unload_model_if_idle() + threading.Timer(MODEL_TIMEOUT / 2.0, _periodic_check).start() + + threading.Timer(MODEL_TIMEOUT / 2.0, _periodic_check).start() + + +if __name__ == "__main__": + port = int(os.environ.get("MODEL_PORT", 5006)) + logger.info("Starting Effort service on port %d", port) + uvicorn.run(app, host="0.0.0.0", port=port) diff --git a/image/effort/requirements.txt b/image/effort/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..be09dc0440fa393455c4e9f1c9317d0cb9766878 --- /dev/null +++ b/image/effort/requirements.txt @@ -0,0 +1,17 @@ +torch>=2.0.0 +torchvision>=0.15.0 +fastapi==0.103.1 +uvicorn[standard]==0.23.2 +python-multipart +pydantic>=2.0.0 +pillow>=9.5.0 +numpy>=1.24.3 +opencv-python-headless>=4.8.0 +transformers>=4.30.0 +open-clip-torch>=2.24.0 +timm>=0.9.6 +einops>=0.6.1 +loralib>=0.1.2 +scikit-learn>=1.3.0 +tensorboard>=2.14.0 +efficientnet-pytorch>=0.7.1 diff --git a/image/fsd/Dockerfile b/image/fsd/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..6de6be5ad5b1c03b4471f3be44c5398a7b93bae3 --- /dev/null +++ b/image/fsd/Dockerfile @@ -0,0 +1,61 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +WORKDIR /app + +# System deps (using Python 3.10 from Ubuntu -- bypassing model_code's >=3.12 constraint) +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3 python3-pip python3-dev \ + git wget \ + && rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +RUN pip install --no-cache-dir --upgrade pip setuptools wheel + +# Clone the FSD repo +RUN git clone https://github.com/ductai199x/Forensic-Self-Descriptions-CVPR25.git model_code + +# Install PyTorch with CUDA 12.1 (replaces CPU-only install) +RUN pip install --no-cache-dir \ + torch==2.5.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +# Install FSD inference dependencies manually +RUN pip install --no-cache-dir \ + "scipy>=1.10.0" \ + "pillow>=10.0.0" \ + pillow-heif \ + click \ + tqdm \ + "scikit-learn>=1.6.0" + +# Install FSD package: bypass both deps AND python-requires check +# The model_code requires Python >=3.12 but works fine on 3.10 +RUN pip install --no-deps --no-build-isolation -e ./model_code 2>/dev/null || \ + (cd model_code && pip install --no-deps -e . --config-settings="--build-option=--skip-requires-python") 2>/dev/null || \ + (cd model_code && python3 setup.py develop --no-deps 2>/dev/null) || \ + echo "FSD model_code install skipped -- will use direct imports" + +# Install service-layer deps +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY app.py . + +RUN mkdir -p /app/weights +COPY weights/ /app/weights/ + +ENV MODEL_PORT=5005 +ENV PRELOAD_MODEL=false +ENV MODEL_TIMEOUT=600 +ENV WEIGHTS_DIR=/app/weights + +EXPOSE 5005 + +RUN adduser --disabled-password --gecos '' appuser +USER appuser + +CMD ["python", "app.py"] diff --git a/image/fsd/app.py b/image/fsd/app.py new file mode 100644 index 0000000000000000000000000000000000000000..13ab7c2d29dd656bb8ea7ce86fa7156486b8d79b --- /dev/null +++ b/image/fsd/app.py @@ -0,0 +1,347 @@ +"""FSD (Forensic Self-Descriptions) deepfake detection service. + +Wraps FSDDetector (CVPR 2025) with a FastAPI endpoint that matches +the contract of all other DeepSafe image services. + +Key design decisions: +- FSDDetector imported lazily inside load_detector_internal() so that + the service can be imported in tests without the fsd package installed. +- z_score_to_probability() converts FSD's z-score to [0,1] using a sigmoid + centred at THRESHOLD_Z=-2.0 (the paper's default boundary). +- SCALE=1.5 is an empirical choice — validated during evaluation. +""" + +import base64 +import gc +import io +import logging +import math +import os +import platform +import sys +import threading +import time +from contextlib import asynccontextmanager +from typing import Any, Dict, Optional + +import torch +import uvicorn +from fastapi import FastAPI, HTTPException +from fastapi.middleware.cors import CORSMiddleware +from PIL import Image, ImageFile +from pydantic import BaseModel + +ImageFile.LOAD_TRUNCATED_IMAGES = True + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + handlers=[logging.StreamHandler(sys.stdout)], +) +logger = logging.getLogger(__name__) + +# ── Configuration ───────────────────────────────────────────────────────────── +MODEL_NAME = "fsd_detection" +WEIGHTS_DIR = os.environ.get( + "WEIGHTS_DIR", os.path.join(os.path.dirname(__file__), "weights") +) +PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true" +MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", "600")) + + +def _get_device(): + """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return "cpu" + if override == "cuda" and torch.cuda.is_available(): + return "cuda" + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return "mps" + if override: + pass # Invalid override, fall through to auto-detect + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return "mps" + if torch.cuda.is_available(): + return "cuda" + return "cpu" + + +DEVICE = _get_device() +if DEVICE == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + +if DEVICE == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) +else: + logger.warning( + "Device: %s (no CUDA available -- check nvidia-container-toolkit)", + DEVICE, + ) + +# Z-score conversion constants (see design spec §4) +THRESHOLD_Z: float = -2.0 # FSD paper's default fake/real boundary +SCALE: float = 1.5 # Empirical sigmoid sharpness — validated during evaluation + +# ── Global state ────────────────────────────────────────────────────────────── +detector = None +detector_lock = threading.Lock() +last_used_time: float = 0.0 + + +class ImageInput(BaseModel): + """Request body for /predict.""" + + image_data: str + threshold: Optional[float] = 0.5 + + +# ── Conversion ──────────────────────────────────────────────────────────────── + + +def z_score_to_probability(z_score: float) -> float: + """Map an FSD z-score to a [0, 1] fake probability via sigmoid. + + Args: + z_score: FSD detector output. More negative = more likely AI-generated. + + Returns: + Float in [0, 1]. At z=THRESHOLD_Z the output is 0.5. + """ + exponent = -(THRESHOLD_Z - z_score) * SCALE + # Clamp to prevent math.exp overflow on extreme inputs + exponent = max(-500.0, min(500.0, exponent)) + return 1.0 / (1.0 + math.exp(exponent)) + + +# ── Model loading ───────────────────────────────────────────────────────────── + + +def load_detector_internal(): + """Load FSDDetector onto CPU from the local weights directory.""" + global detector, last_used_time + + with detector_lock: + if detector is not None: + last_used_time = time.time() + return + + logger.info("Loading FSD detector...") + try: + from fsd import FSDDetector # deferred import — keeps tests import-safe + + if not os.path.exists(WEIGHTS_DIR): + raise RuntimeError( + f"Weights directory not found: {WEIGHTS_DIR}. " + "Run download_weights.sh first." + ) + + loaded = FSDDetector.load( + weights_dir=WEIGHTS_DIR, + device=DEVICE, + threshold=THRESHOLD_Z, + ) + detector = loaded + last_used_time = time.time() + logger.info("FSD detector ready.") + + except Exception as exc: + logger.exception(f"Failed to load FSD detector: {exc}") + detector = None + raise + finally: + gc.collect() + + +def ensure_detector_loaded(): + """Load detector on first request (lazy loading).""" + global last_used_time + if detector is None: + load_detector_internal() + else: + last_used_time = time.time() + + +def unload_detector_if_idle(): + """Evict detector from RAM after MODEL_TIMEOUT seconds of inactivity.""" + global detector + if detector is None or PRELOAD_MODEL: + return + with detector_lock: + if detector is not None and (time.time() - last_used_time > MODEL_TIMEOUT): + logger.info("Unloading idle FSD detector to free RAM.") + del detector + detector = None + gc.collect() + + +# ── FastAPI lifespan ────────────────────────────────────────────────────────── + + +@asynccontextmanager +async def lifespan(app: FastAPI): + """Handle startup and shutdown logic for the FSD service.""" + if PRELOAD_MODEL: + logger.info("Preloading FSD detector at startup.") + try: + load_detector_internal() + except Exception as exc: + logger.error(f"Preload failed: {exc}") + else: + logger.info("FSD service ready — detector loads on first request.") + + if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: + + def _periodic_check(): + unload_detector_if_idle() + threading.Timer(MODEL_TIMEOUT / 2.0, _periodic_check).start() + + threading.Timer(MODEL_TIMEOUT / 2.0, _periodic_check).start() + + yield + + logger.info("Shutting down FSD service.") + + +# ── FastAPI app ─────────────────────────────────────────────────────────────── + +app = FastAPI( + title="FSD Deepfake Detection Service", + description=( + "Zero-shot AI-generated image detection using Forensic Self-Descriptions " + "(CVPR 2025). Trained on real images only; generalises to any generator." + ), + version="1.0.0", + lifespan=lifespan, +) + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + + +@app.get("/") +async def root(): + """Service info.""" + return { + "model_name": MODEL_NAME, + "description": "FSD (CVPR 2025) zero-shot AI-generated image detector", + "device": DEVICE, + "model_loaded": detector is not None, + } + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and DEVICE == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +async def health(): + """Health check endpoint.""" + return { + "status": "healthy", + "model_name": MODEL_NAME, + "device": DEVICE, + "model_loaded": detector is not None, + **_gpu_health_info(), + } + + +@app.post("/unload") +async def unload_model_endpoint(): + """Manually evict detector from RAM.""" + global detector + with detector_lock: + if detector is None: + return {"status": "not_loaded"} + del detector + detector = None + gc.collect() + return {"status": "success", "message": "Detector unloaded."} + + +@app.post("/predict") +async def predict(image_input: ImageInput) -> Dict[str, Any]: + """Predict whether the submitted image is AI-generated. + + Args: + image_input: Base64-encoded image and optional probability threshold. + + Returns: + Dict with model name, fake probability, binary prediction, class label, + and inference time. + """ + try: + ensure_detector_loaded() + if detector is None: + raise HTTPException(status_code=503, detail="Detector not loaded.") + + start = time.time() + + try: + image_bytes = base64.b64decode(image_input.image_data) + image_pil = Image.open(io.BytesIO(image_bytes)).convert("RGB") + except Exception as exc: + raise HTTPException(status_code=400, detail=f"Invalid image data: {exc}") + + with torch.no_grad(): + result = detector.score(image_pil) + + probability = z_score_to_probability(result.z_score) + prediction = 1 if probability >= image_input.threshold else 0 + class_label = "fake" if prediction == 1 else "real" + inference_time = time.time() - start + + logger.info( + f"z={result.z_score:.4f} prob={probability:.4f} " + f"→ {class_label} ({inference_time:.3f}s)" + ) + + if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: + threading.Timer(MODEL_TIMEOUT + 5.0, unload_detector_if_idle).start() + + return { + "model": MODEL_NAME, + "probability": float(probability), + "prediction": int(prediction), + "class": class_label, + "inference_time": float(inference_time), + } + + except HTTPException: + raise + except Exception as exc: + logger.exception(f"Prediction error: {exc}") + raise HTTPException(status_code=500, detail=str(exc)) + + +if __name__ == "__main__": + port = int(os.environ.get("MODEL_PORT", 5005)) + logger.info(f"Starting FSD service on port {port}") + uvicorn.run(app, host="0.0.0.0", port=port) diff --git a/image/fsd/download_weights.sh b/image/fsd/download_weights.sh new file mode 100644 index 0000000000000000000000000000000000000000..b585cfcda4b1afc1a901ebcf2c992a17a3f30776 --- /dev/null +++ b/image/fsd/download_weights.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# Downloads FSD detection weights from GitHub releases v1.2.0 +# Idempotent: skips files that already exist. +set -euo pipefail + +RELEASE_TAG="v1.2.0" +BASE_URL="https://github.com/ductai199x/Forensic-Self-Descriptions-CVPR25/releases/download/${RELEASE_TAG}" +WEIGHTS_DIR="$(dirname "$0")/weights" + +mkdir -p "$WEIGHTS_DIR" + +FILES=("config.json" "fre.pt" "gmm.pt" "fsd_transforms.pt") + +for filename in "${FILES[@]}"; do + filepath="$WEIGHTS_DIR/$filename" + if [ -f "$filepath" ]; then + echo "Already exists: $filename — skipping" + continue + fi + echo "Downloading $filename..." + wget -q --show-progress -O "$filepath" "${BASE_URL}/${filename}" + echo "Done: $filename" +done + +echo "All FSD weights downloaded to $WEIGHTS_DIR" diff --git a/image/fsd/requirements.txt b/image/fsd/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..1363737a1a47a6620a0750efd9a3a448095ad893 --- /dev/null +++ b/image/fsd/requirements.txt @@ -0,0 +1,3 @@ +fastapi +uvicorn +python-multipart diff --git a/image/npr/Dockerfile b/image/npr/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..580456cbc84dc1963f731bc95026526901520bb5 --- /dev/null +++ b/image/npr/Dockerfile @@ -0,0 +1,40 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +WORKDIR /app + +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3 python3-pip \ + git wget \ + && rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +# Install PyTorch with CUDA 12.1 (replaces CPU-only install) +RUN pip install --no-cache-dir \ + torch==2.5.1 torchvision==0.20.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY app.py . + +RUN git clone https://github.com/chuangchuangtan/NPR-DeepfakeDetection.git npr_deepfakedetection + +RUN mkdir -p npr_deepfakedetection/weights && \ + wget -O npr_deepfakedetection/weights/NPR.pth \ + https://github.com/chuangchuangtan/NPR-DeepfakeDetection/raw/main/model_epoch_last_3090.pth + +ENV MODEL_PORT=5001 +ENV PRELOAD_MODEL=false +ENV MODEL_TIMEOUT=600 + +EXPOSE 5001 + +RUN adduser --disabled-password --gecos '' appuser +USER appuser + +CMD ["python", "app.py"] diff --git a/image/npr/app.py b/image/npr/app.py new file mode 100644 index 0000000000000000000000000000000000000000..7feaa182750dd393de9b12ffd1b09a8ae79a216d --- /dev/null +++ b/image/npr/app.py @@ -0,0 +1,444 @@ +""" +NPR-DeepfakeDetection Model Service +This service loads the NPR-DeepfakeDetection model and exposes an API endpoint to analyze images. +It uses the Neural Pattern Residual (NPR) mechanism as described in the original paper. +""" + +import base64 +import gc +import io +import logging +import os +import sys +import threading +import time +from typing import Any, Dict, Optional + +import torch +import torchvision.transforms as transforms +import uvicorn +from fastapi import FastAPI, HTTPException +from fastapi.middleware.cors import CORSMiddleware +from PIL import Image, ImageFile +from pydantic import BaseModel, Field + +ImageFile.LOAD_TRUNCATED_IMAGES = True + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s [%(filename)s:%(lineno)d] - %(message)s", + handlers=[logging.StreamHandler(sys.stdout)], +) +logger = logging.getLogger(__name__) + +MODEL_REPO_SUBDIR = "npr_deepfakedetection" +current_dir = os.path.dirname(os.path.abspath(__file__)) +model_code_path = os.path.join(current_dir, MODEL_REPO_SUBDIR) + +# --- CRITICAL IMPORT: Make this fail loudly if model code isn't found --- +if model_code_path not in sys.path: + sys.path.insert(0, model_code_path) + logger.info(f"Added {model_code_path} to sys.path") + +try: + from networks.resnet import resnet50 # This is now a global variable + + logger.info("Successfully imported resnet50 from npr_deepfakedetection.networks") +except ImportError as e: + logger.critical( + f"CRITICAL: Failed to import resnet50 from {model_code_path}/networks. " + f"Error: {e}. Service cannot start without the model definition. " + f"Ensure '{MODEL_REPO_SUBDIR}' is correctly cloned and accessible.", + exc_info=True, + ) + # Exit if the core model class cannot be imported, as the service is non-functional. + # Or, you could let FastAPI start but have health checks fail catastrophically. + # For a model service, exiting might be cleaner. + sys.exit(f"Fatal Error: Could not import resnet50: {e}") +# --- END CRITICAL IMPORT --- + + +app = FastAPI( + title="NPR-DeepfakeDetection Model Service", + description="Service for detecting deepfake images using NPR-DeepfakeDetection model.", + version="1.0.2", # Incremented version +) + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +MODEL_NAME = "npr_deepfakedetection" +WEIGHTS_FILENAME = "NPR.pth" +MODEL_FULL_PATH = os.path.join(model_code_path, "weights", WEIGHTS_FILENAME) + + +def _get_device(): + """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU.""" + import platform + + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if override: + pass # Invalid override, fall through to auto-detect + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + return torch.device("cpu") + + +DEVICE = _get_device() +if DEVICE.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") +if DEVICE.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) +else: + logger.warning( + "Device: %s (no CUDA available -- check nvidia-container-toolkit)", + DEVICE, + ) +logger.info(f"Using device: {DEVICE}") + +PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true" +MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", "600")) + +from pydantic import BaseModel, ConfigDict, Field # Add ConfigDict if needed + + +class ImageInput(BaseModel): + image_data: str = Field( + ..., description="Base64 encoded image string" + ) # Renamed field + threshold: Optional[float] = Field( + 0.5, ge=0.0, le=1.0, description="Classification threshold" + ) + model_config = ConfigDict(protected_namespaces=()) # Pydantic V2 style for config + + +model: Optional[torch.nn.Module] = None +model_lock = threading.Lock() +last_used_time: float = 0.0 + + +def load_model_internal(): + global model, last_used_time + with model_lock: + if model is not None: + last_used_time = time.time() + return + + logger.info( + f"Loading {MODEL_NAME} model from {MODEL_FULL_PATH} onto {DEVICE}..." + ) + if not os.path.exists(MODEL_FULL_PATH): + logger.error(f"Model weights not found at {MODEL_FULL_PATH}") + raise FileNotFoundError(f"Model weights not found: {MODEL_FULL_PATH}") + + try: + # 'resnet50' is now guaranteed to be in the global scope if the script reached this point + _model = resnet50(num_classes=1) + state_dict = torch.load(MODEL_FULL_PATH, map_location=DEVICE) + + if all(key.startswith("module.") for key in state_dict.keys()): + logger.info("Removing 'module.' prefix from state_dict keys.") + state_dict = {k[len("module.") :]: v for k, v in state_dict.items()} + + _model.load_state_dict(state_dict) + _model.to(DEVICE) + _model.eval() + + model = _model + last_used_time = time.time() + logger.info(f"{MODEL_NAME} model loaded successfully to {DEVICE}.") + + except ( + FileNotFoundError + ) as e_fnf: # Should be caught by earlier check, but good to have + logger.error(f"Model file not found during load: {e_fnf}", exc_info=True) + model = None + raise RuntimeError( + f"Model file not found for {MODEL_NAME}: {e_fnf}" + ) from e_fnf + except Exception as e_load: # Catch other torch.load or model init errors + logger.error( + f"An unexpected error occurred loading {MODEL_NAME} model: {e_load}", + exc_info=True, + ) + model = None + raise RuntimeError( + f"Model loading error for {MODEL_NAME}: {e_load}" + ) from e_load + finally: + if DEVICE.type == "cuda": + torch.cuda.empty_cache() + gc.collect() + + +def ensure_model_loaded(): + global last_used_time + if model is None: + load_model_internal() + else: + last_used_time = time.time() + + +def unload_model_if_idle(): + global model + if model is None or PRELOAD_MODEL: + return + with model_lock: + if model is not None and (time.time() - last_used_time > MODEL_TIMEOUT): + logger.info( + f"Unloading {MODEL_NAME} model due to inactivity (timeout: {MODEL_TIMEOUT}s)." + ) + del model + model = None + if DEVICE.type == "cuda": + torch.cuda.empty_cache() + gc.collect() + logger.info(f"{MODEL_NAME} model unloaded and memory cleared.") + + +def preprocess_image(image_bytes: bytes) -> torch.Tensor: + try: + image = Image.open(io.BytesIO(image_bytes)).convert("RGB") + preprocess = transforms.Compose( + [ + transforms.Resize((256, 256)), + transforms.CenterCrop(224), + transforms.ToTensor(), + transforms.Normalize( + mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] + ), + ] + ) + return preprocess(image).unsqueeze(0) + except Exception as e: + logger.error(f"Error preprocessing image: {e}", exc_info=True) + raise ValueError(f"Image preprocessing failed: {e}") from e + + +@app.on_event("startup") +async def startup_event_handler(): + if PRELOAD_MODEL: + logger.info(f"Preloading {MODEL_NAME} model at startup (PRELOAD_MODEL=true).") + try: + load_model_internal() + except Exception as e: # Catch RuntimeError from load_model_internal + logger.error( + f"Fatal error during {MODEL_NAME} model preloading: {e}. Service might not function correctly.", + exc_info=True, + ) + else: + logger.info( + f"{MODEL_NAME} model will be loaded on first request (PRELOAD_MODEL=false)." + ) + + if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: + active_timer_npr = None + + def periodic_unload_check_runner(): + nonlocal active_timer_npr + unload_model_if_idle() + if ( + model is not None or PRELOAD_MODEL + ): # Keep timer running if model still loaded or preloading + active_timer_npr = threading.Timer( + MODEL_TIMEOUT / 2.0, periodic_unload_check_runner + ) + active_timer_npr.daemon = ( + True # Allow main program to exit even if timer is active + ) + active_timer_npr.start() + else: + logger.info( + f"[{MODEL_NAME}] Model unloaded, stopping idle check timer." + ) + + active_timer_npr = threading.Timer( + MODEL_TIMEOUT / 2.0, periodic_unload_check_runner + ) + active_timer_npr.daemon = True + active_timer_npr.start() + logger.info( + f"{MODEL_NAME} model idle check timer initiated (interval: {MODEL_TIMEOUT / 2.0}s)." + ) + + +@app.get("/", tags=["Info"]) +async def root_endpoint(): + return { + "model_name": MODEL_NAME, + "description": "NPR-based deepfake image detection, using Neural Pattern Residuals.", + "weights_path": MODEL_FULL_PATH, + "device_used": str(DEVICE), + "cuda_available_on_host": torch.cuda.is_available(), + "model_currently_loaded": model is not None, + "lazy_loading_enabled": not PRELOAD_MODEL, + "model_idle_timeout_seconds": ( + MODEL_TIMEOUT if not PRELOAD_MODEL else "N/A (preloaded)" + ), + } + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and DEVICE.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health", tags=["System"]) +async def health_check_endpoint(): + model_file_exists = os.path.exists(MODEL_FULL_PATH) + status_msg = "healthy" + if not model_file_exists: + status_msg = "error_missing_weights" + # Check if the resnet50 symbol is available globally (means critical import succeeded) + elif "resnet50" not in globals() or not callable(globals()["resnet50"]): + status_msg = "error_missing_model_definition" + + return { + "status": status_msg, + "model_name": MODEL_NAME, + "device_configured": str(DEVICE), + "model_weights_found": model_file_exists, + "model_definition_imported": "resnet50" in globals() + and callable(globals()["resnet50"]), + "model_loaded": model is not None, + **_gpu_health_info(), + } + + +@app.post("/unload", tags=["System"], include_in_schema=True) +async def unload_model_endpoint(): + global model + if model is None: + return { + "status": "not_loaded", + "message": f"{MODEL_NAME} model is not currently loaded.", + } + with model_lock: + if model is not None: + logger.info(f"Manually unloading {MODEL_NAME} model via /unload endpoint.") + del model + model = None + if DEVICE.type == "cuda": + torch.cuda.empty_cache() + gc.collect() + logger.info(f"{MODEL_NAME} model unloaded and memory cleared.") + return { + "status": "unloaded", + "message": f"{MODEL_NAME} model unloaded successfully.", + } + else: + return { + "status": "already_unloaded", + "message": f"{MODEL_NAME} model was already unloaded by another request.", + } + + +@app.post("/predict", response_model=Dict[str, Any]) +async def predict_image_endpoint(input_data: ImageInput): + try: + ensure_model_loaded() + if model is None: + logger.error( + f"{MODEL_NAME} model is not available for prediction (ensure_model_loaded failed)." + ) + raise HTTPException( + status_code=503, detail="Model is not loaded or failed to load." + ) + + start_time_pred = time.time() + image_bytes = base64.b64decode(input_data.image_data) + image_tensor = preprocess_image(image_bytes) + image_tensor = image_tensor.to(DEVICE) + + with torch.no_grad(): + output_logit = model(image_tensor) + logger.debug(f"Raw logit from model: {output_logit.item()}") + probability = torch.sigmoid(output_logit).item() + + prediction = 1 if probability >= input_data.threshold else 0 + class_label = "fake" if prediction == 1 else "real" + inference_time = time.time() - start_time_pred + + logger.info( + f"Prediction for {MODEL_NAME} completed in {inference_time:.4f}s. Prob Fake: {probability:.4f}, Threshold: {input_data.threshold}, Class: {class_label}" + ) + + return { + "model": MODEL_NAME, + "probability": float(probability), + "prediction": int(prediction), + "class": class_label, + "inference_time": float(inference_time), + } + except ( + RuntimeError + ) as e_rt_pred: # Catch errors from load_model_internal or PyTorch + logger.error(f"Model runtime or loading error: {e_rt_pred}", exc_info=True) + # Check if the error message matches the one from the log + if ( + "local variable 'resnet50' referenced before assignment" in str(e_rt_pred) + or "Model loading error" in str(e_rt_pred) + or "Model file not found" in str(e_rt_pred) + ): + # This indicates a problem during model loading, likely due to import or file issues + raise HTTPException( + status_code=503, detail=f"Model failed to load: {e_rt_pred}" + ) + raise HTTPException(status_code=500, detail=f"Model runtime error: {e_rt_pred}") + except ValueError as e_val_pred: + logger.error(f"Image processing error: {e_val_pred}", exc_info=True) + raise HTTPException( + status_code=400, detail=f"Image processing error: {e_val_pred}" + ) + except HTTPException: + raise + except Exception as e_pred: + logger.error( + f"Unexpected error processing prediction for {MODEL_NAME}: {e_pred}", + exc_info=True, + ) + raise HTTPException( + status_code=500, detail=f"Unexpected prediction error: {e_pred}" + ) + + +if __name__ == "__main__": + port = int(os.environ.get("MODEL_PORT", 5001)) + # The critical import of resnet50 is now at the top. + # If it fails, the script exits before uvicorn.run is called. + logger.info(f"Starting {MODEL_NAME} server on port {port} with device: {DEVICE}") + uvicorn.run("app:app", host="0.0.0.0", port=port, reload=False) diff --git a/image/npr/requirements.txt b/image/npr/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..42ed185ce549e8f13f83d0c36025d39d98d9f36f --- /dev/null +++ b/image/npr/requirements.txt @@ -0,0 +1,11 @@ +fastapi==0.103.1 +uvicorn==0.23.2 +pydantic==2.3.0 +torch==2.5.1 +torchvision==0.20.1 +pillow==10.0.1 +numpy==1.23.5 +scipy +scikit-learn +opencv-python-headless +python-multipart \ No newline at end of file diff --git a/image/universal/Dockerfile b/image/universal/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..2908fc9618f89df30240e720edc771f8d157e86e --- /dev/null +++ b/image/universal/Dockerfile @@ -0,0 +1,49 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +WORKDIR /app + +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3 python3-pip git wget \ + && rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +# Install PyTorch with CUDA 12.1 (upgraded from 1.11/CUDA 11.3 for Blackwell support) +RUN pip install --no-cache-dir \ + torch==2.5.1 torchvision==0.20.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +# Copy application code +COPY app.py . + +# Clone the repository +RUN git clone https://github.com/WisconsinAIVision/UniversalFakeDetect.git universalfakedetect && \ + mkdir -p universalfakedetect/pretrained_weights + +RUN python -c "import torch; import os; os.makedirs('/root/.cache/clip', exist_ok=True)" && \ + wget -O /root/.cache/clip/ViT-L-14.pt https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt + +# Download the weights file separately to avoid timeouts during startup +RUN wget -O universalfakedetect/pretrained_weights/fc_weights.pth \ + https://github.com/WisconsinAIVision/UniversalFakeDetect/raw/main/pretrained_weights/fc_weights.pth + +# Set environment variables +ENV MODEL_PORT=5003 +ENV USE_GPU=false +ENV TIMEOUT=600 + +# Expose port +EXPOSE 5003 + +# Drop root privileges +RUN adduser --disabled-password --gecos '' appuser +USER appuser + +# Run the application with gunicorn for better timeout handling +CMD ["python", "-m", "uvicorn", "app:app", "--host", "0.0.0.0", "--port", "5003", "--timeout-keep-alive", "600"] \ No newline at end of file diff --git a/image/universal/app.py b/image/universal/app.py new file mode 100644 index 0000000000000000000000000000000000000000..b704f5e836e34e75108b151da463ed7e7b5123fd --- /dev/null +++ b/image/universal/app.py @@ -0,0 +1,466 @@ +import base64 +import gc +import io +import logging +import os +import platform +import sys +import threading +import time +import traceback +from typing import Any, Dict, Optional + +import torch +from fastapi import FastAPI, HTTPException +from fastapi.middleware.cors import CORSMiddleware +from PIL import Image +from pydantic import BaseModel + +# Setup more detailed logging +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" +) +logger = logging.getLogger(__name__) + +# Initialize FastAPI +app = FastAPI( + title="UniversalFakeDetect API", + description="API for Universal Fake Image Detector", + version="1.0.0", +) + +# Add CORS middleware +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + + +# Environment variables +def _get_device(): + """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if override: + pass # Invalid override, fall through to auto-detect + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + return torch.device("cpu") + + +DEVICE = _get_device() +USE_GPU = DEVICE.type != "cpu" +if DEVICE.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") +if DEVICE.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) +else: + logger.warning( + "Device: %s (no CUDA available -- check nvidia-container-toolkit)", + DEVICE, + ) +MODEL_PORT = int(os.environ.get("MODEL_PORT", 5003)) +PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true" +MODEL_TIMEOUT = int( + os.environ.get("MODEL_TIMEOUT", "600") +) # Seconds to keep model loaded +logger.info(f"Using device: {DEVICE}") +logger.info(f"CUDA available: {torch.cuda.is_available()}") +if torch.cuda.is_available(): + logger.info(f"CUDA device count: {torch.cuda.device_count()}") + logger.info(f"CUDA device name: {torch.cuda.get_device_name(0)}") + +from typing import Optional + +# Define input model +from pydantic import BaseModel + + +class ImageInput(BaseModel): + image_data: str # Renamed field + threshold: Optional[float] = 0.5 + + +# Global variables for model +model = None +model_lock = threading.Lock() +last_used_time = 0 +model_loading = False + + +# Download weights if not present +def download_weights(): + weights_path = "universalfakedetect/pretrained_weights/fc_weights.pth" + if not os.path.exists("universalfakedetect/pretrained_weights"): + os.makedirs("universalfakedetect/pretrained_weights", exist_ok=True) + logger.info("Created pretrained_weights directory") + + if not os.path.exists(weights_path): + logger.info("Weights file not found, downloading...") + import urllib.request + + url = "https://github.com/WisconsinAIVision/UniversalFakeDetect/raw/main/pretrained_weights/fc_weights.pth" + urllib.request.urlretrieve(url, weights_path) + logger.info(f"Downloaded weights to {weights_path}") + + +# Load the model +def load_model(): + """Load the UniversalFakeDetect model.""" + global model, last_used_time, model_loading + + # If model is already loaded, update timestamp and return + if model is not None: + last_used_time = time.time() + return model + + with model_lock: # Thread safety for concurrent requests + # Check again after acquiring the lock + if model is not None: + last_used_time = time.time() + return model + + # Set flag to indicate model is loading + model_loading = True + + try: + logger.info(f"Loading UniversalFakeDetect model on {DEVICE}...") + + # List directory contents for debugging + logger.info(f"Current directory: {os.getcwd()}") + logger.info(f"Directory contents: {os.listdir('.')}") + + if os.path.exists("universalfakedetect"): + logger.info( + f"universalfakedetect directory contents: {os.listdir('universalfakedetect')}" + ) + else: + logger.error("universalfakedetect directory not found!") + model_loading = False + return None + + # Ensure weights exist + download_weights() + weights_path = "universalfakedetect/pretrained_weights/fc_weights.pth" + + # Import model modules + logger.info("Adding universalfakedetect to sys.path") + sys.path.append(os.path.abspath("universalfakedetect")) + + logger.info("Importing get_model from models") + try: + # Try direct import first + from models import get_model + + logger.info("Successfully imported get_model") + except ImportError as e: + logger.warning( + f"Direct import failed: {str(e)}, trying alternate import" + ) + from universalfakedetect.models import get_model + + logger.info("Successfully imported get_model with alternate path") + except Exception as e: + logger.error(f"Error importing get_model: {str(e)}") + logger.error(traceback.format_exc()) + model_loading = False + return None + + # Initialize the model + logger.info("Initializing model with CLIP:ViT-L/14") + model = get_model("CLIP:ViT-L/14") + + # Load the weights + logger.info(f"Loading weights from {weights_path}") + state_dict = torch.load(weights_path, map_location="cpu") + model.fc.load_state_dict(state_dict) + + # Move model to device and set to evaluation mode + logger.info(f"Moving model to device: {DEVICE}") + model.to(DEVICE) + model.eval() + + # Update last used time + last_used_time = time.time() + + # Clear CUDA cache to free up memory + if USE_GPU: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + gc.collect() + + logger.info("Model loaded successfully!") + model_loading = False + return model + + except Exception as e: + logger.error(f"Error loading model: {str(e)}") + logger.error(traceback.format_exc()) + model_loading = False + return None + + +@app.post("/unload", include_in_schema=True) +async def unload_model_endpoint(): + """Endpoint to manually unload the model.""" + global model + + if model_loading: + return { + "status": "loading_busy", + "message": "Model is currently being loaded, cannot unload now.", + } + + if model is None: + return {"status": "not_loaded", "message": "Model is not currently loaded."} + + with model_lock: + if model is not None: # Check again inside lock + logger.info( + "Manually unloading UniversalFakeDetect model via /unload endpoint." + ) + del model + model = None + # Clear CUDA cache if it was used (though DEVICE is 'cpu' here, good practice) + if DEVICE.type == "cuda": + if torch.cuda.is_available(): + torch.cuda.empty_cache() + gc.collect() + logger.info("UniversalFakeDetect model unloaded and memory cleared.") + return {"status": "unloaded", "message": "Model unloaded successfully."} + else: # Should not happen if initial check was model is not None + return { + "status": "already_unloaded", + "message": "Model was already unloaded.", + } + + +def unload_model_if_idle(): + """Unload model if it's been idle for too long.""" + global model + + if model is None: + return + + if time.time() - last_used_time > MODEL_TIMEOUT: + with model_lock: + if model is not None and time.time() - last_used_time > MODEL_TIMEOUT: + logger.info( + f"Unloading model after {MODEL_TIMEOUT} seconds of inactivity" + ) + # Delete model and clear memory + del model + model = None + # Clear CUDA cache + if USE_GPU: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + gc.collect() + logger.info("Model unloaded and memory cleared") + + +# Preprocess image for inference +def preprocess_image(image_bytes): + try: + # Read the image + image = Image.open(io.BytesIO(image_bytes)).convert("RGB") + + # Preprocess image + from torchvision import transforms + + # These values are from the validate.py file in the repository + mean = [0.48145466, 0.4578275, 0.40821073] # CLIP values + std = [0.26862954, 0.26130258, 0.27577711] # CLIP values + + transform = transforms.Compose( + [ + transforms.CenterCrop(224), + transforms.ToTensor(), + transforms.Normalize(mean=mean, std=std), + ] + ) + + img_tensor = transform(image).unsqueeze(0) # Add batch dimension + return img_tensor.to(DEVICE) + except Exception as e: + logger.error(f"Error preprocessing image: {str(e)}") + logger.error(traceback.format_exc()) + raise HTTPException(status_code=400, detail=f"Invalid image format: {str(e)}") + + +# Predict function +def predict(image_tensor, threshold=0.5): + try: + with torch.no_grad(): + # Forward pass through the model + output = model(image_tensor).sigmoid().flatten().item() + # The model outputs a score between 0 and 1, where higher values indicate fake images + prediction = 1 if output >= threshold else 0 + + return { + "probability": float(output), + "prediction": int(prediction), + "class": "fake" if prediction == 1 else "real", + } + except Exception as e: + logger.error(f"Prediction error: {str(e)}") + logger.error(traceback.format_exc()) + raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}") + + +# Root endpoint +@app.get("/") +def read_root(): + return { + "model": "UniversalFakeDetect", + "description": "Universal Fake Image Detector that Generalizes Across Generative Models", + "authors": "Utkarsh Ojha, Yuheng Li, Yong Jae Lee", + "paper": "https://arxiv.org/abs/2302.10174", + "source": "https://github.com/WisconsinAIVision/UniversalFakeDetect", + "model_loaded": model is not None, + "lazy_loading": not PRELOAD_MODEL, + } + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and DEVICE.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +# Health check endpoint +@app.get("/health") +def health_check(): + weights_path = "universalfakedetect/pretrained_weights/fc_weights.pth" + model_file_exists = os.path.exists(weights_path) + + if model is not None: + return { + "status": "healthy", + "device": str(DEVICE), + "model_loaded": True, + **_gpu_health_info(), + } + elif model_loading: + return { + "status": "loading", + "message": "Model is being loaded", + "device": str(DEVICE), + } + elif not model_file_exists: + return { + "status": "missing_weights", + "message": "Model weights not found", + "device": str(DEVICE), + } + else: + return { + "status": "not_loaded", + "message": "Model not loaded yet", + "device": str(DEVICE), + } + + +# Prediction endpoint +@app.post("/predict") +async def predict_image(input_data: ImageInput) -> Dict[str, Any]: + # Check if model is loaded + global model, last_used_time + + try: + # Load model if not already loaded + if model is None: + logger.info("Model not loaded. Loading model now...") + model = load_model() + if model is None: + raise HTTPException(status_code=500, detail="Failed to load model") + else: + # Update timestamp if already loaded + last_used_time = time.time() + + # Decode base64 image + image_bytes = base64.b64decode(input_data.image_data) + + # Start timing + start_time = time.time() + + # Optimize memory during inference + if USE_GPU: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + # Preprocess the image + image_tensor = preprocess_image(image_bytes) + + # Get predictions + results = predict(image_tensor, input_data.threshold) + + # Calculate inference time + inference_time = time.time() - start_time + + # Schedule unloading after timeout - run in background + threading.Timer(5.0, unload_model_if_idle).start() + + # Return results + return { + "model": "UniversalFakeDetect", + "probability": results["probability"], + "prediction": results["prediction"], + "class": results["class"], + "inference_time": inference_time, + } + except Exception as e: + logger.error(f"Error during prediction: {str(e)}") + logger.error(traceback.format_exc()) + raise HTTPException(status_code=500, detail=str(e)) + + +# Load model on startup +@app.on_event("startup") +async def startup_event(): + """Load model on startup only if PRELOAD_MODEL is true.""" + if PRELOAD_MODEL: + logger.info("Preloading model at startup (PRELOAD_MODEL=true)") + try: + load_model() + except Exception as e: + logger.error(f"Preloading failed: {str(e)}") + else: + logger.info("Model will be loaded on first request (PRELOAD_MODEL=false)") + + +# Run the server +if __name__ == "__main__": + import uvicorn + + uvicorn.run("app:app", host="0.0.0.0", port=MODEL_PORT, reload=False) diff --git a/image/universal/requirements.txt b/image/universal/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..91bfffed779ec0bfce89e5e75cf94cea00b0f6c3 --- /dev/null +++ b/image/universal/requirements.txt @@ -0,0 +1,11 @@ +fastapi +uvicorn +pillow +pydantic +numpy +# torch and torchvision installed separately in Dockerfile with CUDA support +# DO NOT add torch or torchvision pins here -- they are installed with CUDA +# support in the Dockerfile. Adding torch==X here would downgrade the install. +torch>=2.5.1 +ftfy +regex \ No newline at end of file diff --git a/image/yermandy/Dockerfile b/image/yermandy/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..f1f5b74c486f268f46aee2446458aae04e5aa010 --- /dev/null +++ b/image/yermandy/Dockerfile @@ -0,0 +1,55 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +WORKDIR /app + +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3 python3-pip \ + git wget \ + && rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +# Install PyTorch with CUDA 12.1 (downgraded from 2.6.0 -- no cu121 wheels for 2.6+) +RUN pip install --no-cache-dir \ + torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +RUN git clone https://github.com/yermandy/deepfake-detection.git model_code + +RUN mkdir -p model_code/weights && \ + wget https://huggingface.co/yermandy/deepfake-detection/resolve/main/model.ckpt -O model_code/weights/model.ckpt + +ENV HF_HOME=/app/huggingface_cache +ENV TRANSFORMERS_CACHE=/app/huggingface_cache +ENV HF_HUB_CACHE=/app/huggingface_cache +RUN mkdir -p $HF_HOME + +RUN pip install --no-cache-dir --upgrade transformers + +RUN python -c "from transformers import CLIPProcessor, CLIPModel; \ + CLIPProcessor.from_pretrained('openai/clip-vit-base-patch16'); \ + CLIPProcessor.from_pretrained('openai/clip-vit-base-patch32'); \ + CLIPProcessor.from_pretrained('openai/clip-vit-large-patch14'); \ + CLIPProcessor.from_pretrained('openai/clip-vit-large-patch14-336'); \ + CLIPModel.from_pretrained('openai/clip-vit-large-patch14')" + +COPY app.py . + +ENV MODEL_PORT=5002 +ENV MODEL_PATH=model_code/weights/model.ckpt +ENV PRELOAD_MODEL="false" +ENV MODEL_TIMEOUT="600" +ENV PYTHONPATH=/app/model_code:/app/model_code/src:$PYTHONPATH + +EXPOSE ${MODEL_PORT} + +RUN adduser --disabled-password --gecos '' appuser +USER appuser + +CMD ["python", "app.py"] diff --git a/image/yermandy/app.py b/image/yermandy/app.py new file mode 100644 index 0000000000000000000000000000000000000000..e0da1be1cc1075bc7687a6bc03cac14f6695c6f0 --- /dev/null +++ b/image/yermandy/app.py @@ -0,0 +1,464 @@ +import base64 +import gc +import io +import logging +import os +import platform +import sys +import threading +import time +from contextlib import asynccontextmanager +from typing import Any, Dict, Optional + +import torch +import uvicorn +from fastapi import FastAPI, HTTPException +from fastapi.middleware.cors import CORSMiddleware +from PIL import Image, ImageFile +from pydantic import BaseModel, Field + +# Allow loading of truncated images +ImageFile.LOAD_TRUNCATED_IMAGES = True + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + handlers=[logging.StreamHandler(sys.stdout)], +) +logger = logging.getLogger(__name__) + +# Model-specific imports from the cloned repository +# PYTHONPATH is set in Dockerfile to find 'model_code.src' or 'src' directly +try: + from lightning.fabric import Fabric + from src.config import Config + from src.model.dfdet import DeepfakeDetectionModel +except ImportError as e: + logger.error(f"Error importing model-specific modules: {e}. Check PYTHONPATH.") + # Attempt to add path manually if Docker's PYTHONPATH isn't picked up as expected in some envs + current_dir_for_app = os.path.dirname(os.path.abspath(__file__)) + model_code_path = os.path.join(current_dir_for_app, "model_code") + if model_code_path not in sys.path: + sys.path.insert(0, model_code_path) + logger.info(f"Added {model_code_path} to sys.path") + + # Retry imports + from lightning.fabric import Fabric + from src.config import Config + from src.model.dfdet import DeepfakeDetectionModel + + +@asynccontextmanager +async def lifespan(app: FastAPI): + # Code to run on startup + if PRELOAD_MODEL: + logger.info("Preloading model at startup (PRELOAD_MODEL=true)") + try: + load_model_internal() + except Exception as e: + logger.error( + f"Fatal error during model preloading: {e}. Service might not function." + ) + else: + logger.info("Model will be loaded on first request (PRELOAD_MODEL=false).") + + # Start a background timer to check for model unloading if not preloading + timer_thread = None + if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: + + def periodic_unload_check(): + unload_model_if_idle() + # Reschedule the check only if the model is still loaded or the app is running + # A more robust way would be to manage the timer cancellation in the 'yield' part + if model is not None and not PRELOAD_MODEL: # Check if model still exists + timer_thread = threading.Timer(MODEL_TIMEOUT / 2, periodic_unload_check) + timer_thread.start() + + # Initial call after a short delay + timer_thread = threading.Timer(MODEL_TIMEOUT / 2, periodic_unload_check) + timer_thread.start() + logger.info(f"Model idle check timer started (interval: {MODEL_TIMEOUT / 2}s).") + + yield # Application runs here + + # Code to run on shutdown + logger.info("Shutting down Yermandy CLIP detection service.") + if timer_thread and timer_thread.is_alive(): + timer_thread.cancel() + logger.info("Cancelled model idle check timer.") + # Ensure model is unloaded on shutdown if it exists + global model, preprocessing_fn, fabric + if model is not None: + logger.info("Unloading model on shutdown.") + del model, preprocessing_fn, fabric + model, preprocessing_fn, fabric = None, None, None + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +# Initialize FastAPI app with lifespan +app = FastAPI( + title="Yermandy CLIP Deepfake Detection Model Service", + description="Service for detecting deepfake images using the Yermandy CLIP-based model.", + version="1.0.0", + lifespan=lifespan, # Add the lifespan manager here +) + +# CORS setup +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +# --- Model Configuration & Globals --- +MODEL_NAME = "yermandy_clip_detection" +MODEL_PATH = os.environ.get("MODEL_PATH", "model_code/weights/model.ckpt") + + +def _get_device(): + """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if override: + pass # Invalid override, fall through to auto-detect + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + return torch.device("cpu") + + +DEVICE = _get_device() +if DEVICE.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") +if DEVICE.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) +else: + logger.warning( + "Device: %s (no CUDA available -- check nvidia-container-toolkit)", + DEVICE, + ) +PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true" +MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", "600")) # Default 10 minutes + +# Global variables for the model and related components +model: Optional[DeepfakeDetectionModel] = None +preprocessing_fn: Optional[callable] = None +fabric: Optional[Fabric] = None +model_lock = threading.Lock() +last_used_time = 0 + +from pydantic import BaseModel, Field # ConfigDict not needed here if no extra config + + +class ImageInput(BaseModel): + image_data: str = Field( + ..., description="Base64 encoded image string" + ) # Renamed field + threshold: Optional[float] = Field( + 0.35, ge=0.0, le=1.0, description="Classification threshold" + ) + + +def load_model_internal(): + """Loads the deepfake detection model and its components.""" + global model, preprocessing_fn, fabric, last_used_time + + with model_lock: + if model is not None: # Check again after acquiring lock + last_used_time = time.time() + return + + logger.info(f"Loading model from {MODEL_PATH} onto {DEVICE}...") + if not os.path.exists(MODEL_PATH): + logger.error(f"Model weights not found at {MODEL_PATH}") + raise FileNotFoundError(f"Model weights not found at {MODEL_PATH}") + + try: + ckpt = torch.load(MODEL_PATH, map_location="cpu") # Load to CPU first + + # Ensure hyper_parameters are available + if "hyper_parameters" not in ckpt: + logger.error( + "Checkpoint does not contain 'hyper_parameters'. Cannot initialize model." + ) + raise ValueError("Invalid checkpoint: missing 'hyper_parameters'") + + model_config = Config(**ckpt["hyper_parameters"]) + + _model = DeepfakeDetectionModel(model_config) + _model.load_state_dict(ckpt["state_dict"]) + _model.eval() # Set to evaluation mode + _model.to(DEVICE) # Move to CPU + + _preprocessing_fn = _model.get_preprocessing() + + # For CPU, Fabric precision should be "32-true" or similar, not from checkpoint if it's like "bf16-mixed" + # The original inference.py loads precision from ckpt["hyper_parameters"]["precision"] + # For CPU, it's safer to override this. + _accelerator = "cuda" if DEVICE.type == "cuda" else "cpu" + _fabric = Fabric(accelerator=_accelerator, devices=1, precision="32-true") + # fabric.launch() is not needed here as we are not in a distributed script + _model = _fabric.setup_module(_model) # Prepare model with Fabric + + # Assign to global variables + model = _model + preprocessing_fn = _preprocessing_fn + fabric = _fabric + + last_used_time = time.time() + logger.info("Model loaded successfully.") + + except Exception as e: + logger.exception(f"Failed to load model: {e}") + # Reset globals if loading failed + model = None + preprocessing_fn = None + fabric = None + raise + finally: + gc.collect() + + +def ensure_model_loaded(): + """Ensures the model is loaded, loading it if necessary.""" + global last_used_time + if model is None: + load_model_internal() + else: + last_used_time = time.time() # Update last used time if already loaded + + +def unload_model_if_idle(): + """Unloads the model if it has been idle for longer than MODEL_TIMEOUT.""" + global model, preprocessing_fn, fabric + if model is None or PRELOAD_MODEL: # Don't unload if preloaded or already unloaded + return + + with model_lock: + if model is not None and (time.time() - last_used_time > MODEL_TIMEOUT): + logger.info( + f"Unloading model due to inactivity (timeout: {MODEL_TIMEOUT}s)." + ) + del model + del preprocessing_fn + del fabric + model = None + preprocessing_fn = None + fabric = None + gc.collect() + logger.info("Model unloaded and memory cleared.") + + +# --- FastAPI Endpoints --- +@app.on_event("startup") +async def startup_event(): + if PRELOAD_MODEL: + logger.info("Preloading model at startup (PRELOAD_MODEL=true)") + try: + load_model_internal() + except Exception as e: + logger.error( + f"Fatal error during model preloading: {e}. Service might not function." + ) + else: + logger.info("Model will be loaded on first request (PRELOAD_MODEL=false).") + + # Start a background timer to check for model unloading if not preloading + if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: + + def periodic_unload_check(): + unload_model_if_idle() + if ( + model is None and not PRELOAD_MODEL + ): # if model got unloaded and we are not preloading, stop timer + return + # Reschedule the check + threading.Timer(MODEL_TIMEOUT / 2, periodic_unload_check).start() + + # Initial call after a short delay + threading.Timer(MODEL_TIMEOUT / 2, periodic_unload_check).start() + logger.info(f"Model idle check timer started (interval: {MODEL_TIMEOUT / 2}s).") + + +@app.get("/") +async def root(): + return { + "model_name": MODEL_NAME, + "description": "Deepfake detection model based on Yermandy's CLIP work.", + "model_path": MODEL_PATH, + "device": str(DEVICE), + "model_loaded": model is not None, + "lazy_loading_enabled": not PRELOAD_MODEL, + "model_timeout_seconds": ( + MODEL_TIMEOUT if not PRELOAD_MODEL else "N/A (preloaded)" + ), + } + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and DEVICE.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +async def health(): + model_file_exists = os.path.exists(MODEL_PATH) + status_message = "healthy" + if not model_file_exists: + status_message = "error_missing_weights" + # Could add a quick inference test here if model is loaded for a more thorough check + + return { + "status": status_message, + "model_name": MODEL_NAME, + "device": str(DEVICE), + "model_weights_found": model_file_exists, + "model_loaded": model is not None, + **_gpu_health_info(), + } + + +@app.post("/unload") +async def unload_model(): + """Endpoint to manually unload the model""" + global model + + if model is None: + return {"status": "not_loaded", "message": "Model is not currently loaded"} + + logger.info("Manually unloading model") + # Delete model and clear memory + del model + model = None + # Check if cuda is available instead of using undefined USE_GPU + if torch.cuda.is_available(): # Modified line + torch.cuda.empty_cache() + gc.collect() + logger.info("Model unloaded and memory cleared") + + return {"status": "success", "message": "Model unloaded successfully"} + + +@app.post("/predict", response_model=Dict[str, Any]) +async def predict(image_input: ImageInput): + try: + ensure_model_loaded() # This will load the model if it's not already loaded + + if model is None or preprocessing_fn is None or fabric is None: + logger.error("Model components are not available for prediction.") + raise HTTPException( + status_code=503, detail="Model is not loaded or failed to load." + ) + + start_time = time.time() + + # Decode base64 image + try: + image_bytes = base64.b64decode(image_input.image_data) + pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB") + except Exception as e: + logger.error(f"Invalid image data: {e}") + raise HTTPException(status_code=400, detail=f"Invalid image data: {e}") + + # Preprocess the image + # The preprocessing_fn from the model expects a PIL image and returns a tensor. + # The model's forward pass expects a batch, so unsqueeze(0). + image_tensor = preprocessing_fn(pil_image).unsqueeze(0) + + # Move tensor to the device Fabric prepared the model for (CPU in this case) + # The dtype should be handled by Fabric setup_module or .to(DEVICE) + image_tensor = image_tensor.to( + DEVICE + ) # fabric.to_device(image_tensor) could also be used if fabric object is accessible + + # Perform inference + with torch.no_grad(): + output = model(image_tensor) # model is already setup by fabric + + # Process output + # output.logits_labels is what the original inference.py uses + # It's a tensor of shape [batch_size, num_classes], e.g., [1, 2] for [prob_real, prob_fake] + probabilities_tensor = output.logits_labels.softmax(dim=1) + + # Probability of being FAKE is the second element (index 1) + probability_fake = probabilities_tensor[0, 1].item() + + prediction = 1 if probability_fake >= image_input.threshold else 0 + class_label = "fake" if prediction == 1 else "real" + + inference_time_seconds = time.time() - start_time + logger.info( + f"Prediction for {MODEL_NAME} completed in {inference_time_seconds:.4f}s. Prob Fake: {probability_fake:.4f}" + ) + + # Schedule model unload if not preloaded + if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: + threading.Timer( + MODEL_TIMEOUT + 5.0, unload_model_if_idle + ).start() # Check slightly after timeout + + return { + "model": MODEL_NAME, + "probability": probability_fake, + "prediction": prediction, + "class": class_label, + "inference_time": inference_time_seconds, + } + + except FileNotFoundError as e: + logger.error(f"Model file not found: {e}") + raise HTTPException(status_code=503, detail=f"Model weights missing: {e}") + except HTTPException: + raise # Re-raise HTTPException directly + except Exception as e: + logger.exception(f"Error during prediction: {e}") # Log full traceback + raise HTTPException( + status_code=500, detail=f"Internal server error during prediction: {e}" + ) + + +if __name__ == "__main__": + port = int(os.environ.get("MODEL_PORT", 5002)) + logger.info(f"Starting {MODEL_NAME} server on port {port} with CPU: {DEVICE}") + # When running app.py directly, PYTHONPATH might need to be set if model_code is not in the same dir + # This is handled by Dockerfile's ENV PYTHONPATH for containerized execution. + + # For local dev, if model_code is sibling to this app.py's dir (e.g. in models/image/yermandy_clip_detection/) + # and model_code is cloned as models/image/yermandy_clip_detection/model_code/ + # local_mc_path = os.path.join(os.path.dirname(__file__), 'model_code') + # if os.path.isdir(local_mc_path) and local_mc_path not in sys.path: + # sys.path.insert(0, local_mc_path) + + uvicorn.run(app, host="0.0.0.0", port=port) diff --git a/image/yermandy/requirements.txt b/image/yermandy/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..28ca8cdbb444ea88c7ff9619fd2218107b07f622 --- /dev/null +++ b/image/yermandy/requirements.txt @@ -0,0 +1,24 @@ +fastapi +uvicorn[standard] +# torch and torchvision are expected from the base pytorch image +Pillow>=9.0.0 +python-multipart +numpy + +torch==2.5.1 +torchaudio==2.5.1 +torchvision==0.20.1 +lightning==2.5.0 +transformers==4.50.0 +tqdm==4.67.1 # progress bar +timm==1.0.14 # torch models +matplotlib +seaborn==0.13.2 # visualization +scikit-learn==1.6.1 # metrics +rich==13.9.4 # logging +wandb==0.19.4 # logging +pydantic==2.9.2 # config +ruff==0.9.3 # formatting +fire==0.7.0 # CLI +peft==0.14.0 # parameter-efficient fine-tuning +# ipykernel==6.29.5 # jupyter \ No newline at end of file diff --git a/provenance/audioseal/Dockerfile b/provenance/audioseal/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..437ae0e6fe5ff68d5e0392dda6f631e74087fa73 --- /dev/null +++ b/provenance/audioseal/Dockerfile @@ -0,0 +1,34 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + python3 python3-pip \ + libsndfile1 ffmpeg && \ + rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +WORKDIR /app + +# Install PyTorch with CUDA 12.1 +RUN pip install --no-cache-dir \ + torch==2.5.1 torchaudio==2.5.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY app.py . + +ENV MODEL_PORT=9003 +ENV PRELOAD_MODEL=false +ENV MODEL_TIMEOUT=600 +EXPOSE 9003 + +RUN adduser --disabled-password --gecos "" appuser +USER appuser + +CMD ["python", "app.py"] diff --git a/provenance/audioseal/app.py b/provenance/audioseal/app.py new file mode 100644 index 0000000000000000000000000000000000000000..23fb3ca69dda0805a822fd100413cfe73c655e4e --- /dev/null +++ b/provenance/audioseal/app.py @@ -0,0 +1,405 @@ +"""AudioSeal watermark detection service. + +Detects Meta AudioSeal watermarks embedded in audio files using the +audioseal_detector_16bits model. Returns a probability score indicating +likelihood that the audio contains a watermark (indicating AI generation). + +Unlike detection models that analyze acoustic artifacts, this service +inspects only for embedded AudioSeal watermarks. +""" + +import base64 +import gc +import logging +import os +import platform +import sys +import tempfile +import threading +import time +from typing import Any, Dict, Optional + +import torch +import torchaudio +import uvicorn +from fastapi import FastAPI, HTTPException +from fastapi.middleware.cors import CORSMiddleware +from pydantic import BaseModel + +# ── Logging ──────────────────────────────────────────────────────────────── + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + handlers=[logging.StreamHandler(sys.stdout)], +) +logger = logging.getLogger(__name__) + +# ── Config ───────────────────────────────────────────────────────────────── + +MODEL_NAME = "audioseal_detector" +MODEL_PORT = int(os.environ.get("MODEL_PORT", "9003")) + +_PRODUCTION = os.environ.get("PRODUCTION", "false").lower() == "true" + +DETECTION_THRESHOLD = 0.3 +TARGET_SAMPLE_RATE = 16000 + + +def _get_device(): + """Select optimal device: CUDA > MPS > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if override: + pass # Invalid override, fall through to auto-detect + if torch.cuda.is_available(): + return torch.device("cuda") + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + return torch.device("cpu") + + +DEVICE = _get_device() +if DEVICE.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + +if DEVICE.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) +else: + logger.warning( + "Device: %s (no CUDA available -- check nvidia-container-toolkit)", + DEVICE, + ) + +PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true" +MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", "600")) + +# ── Globals ──────────────────────────────────────────────────────────────── + +detector = None +model_lock = threading.Lock() +last_used_time = 0 + + +# ── Pydantic models ─────────────────────────────────────────────────────── + + +class AudioInput(BaseModel): + """Request body for /predict.""" + + audio_data: str + threshold: float = 0.5 + + +# ── Confidence mapping ──────────────────────────────────────────────────── + + +def _map_confidence_to_probability(confidence: float) -> float: + """Map AudioSeal detection confidence to a probability score. + + Confidence values below the detection threshold are treated as + noise and mapped to 0.5 (neutral). Values at or above the + threshold are linearly scaled into [0.5, 1.0]. + + Args: + confidence: Raw detection confidence from AudioSeal in [0, 1]. + + Returns: + Probability in [0.5, 1.0]. 0.5 means neutral (no watermark). + """ + if confidence < DETECTION_THRESHOLD: + return 0.5 + return 0.5 + confidence * 0.5 + + +# ── Model loading ───────────────────────────────────────────────────────── + + +def load_model_internal(): + """Load AudioSeal detector model onto the selected device.""" + global detector, last_used_time + + with model_lock: + if detector is not None: + last_used_time = time.time() + return + + logger.info("Loading AudioSeal detector model...") + try: + from audioseal import AudioSeal + + loaded = AudioSeal.load_detector("audioseal_detector_16bits") + loaded = loaded.to(DEVICE) + # Switch to inference mode (no gradient tracking) + loaded.train(mode=False) + detector = loaded + last_used_time = time.time() + logger.info("AudioSeal detector ready on %s.", str(DEVICE)) + + except Exception as exc: + logger.exception("Failed to load AudioSeal detector: %s", exc) + detector = None + raise + finally: + gc.collect() + + +def ensure_model_loaded(): + """Load model on first request (lazy loading).""" + global last_used_time + if detector is None: + load_model_internal() + else: + last_used_time = time.time() + + +def unload_model_if_idle(): + """Evict model from RAM after MODEL_TIMEOUT seconds of inactivity.""" + global detector + if detector is None or PRELOAD_MODEL: + return + with model_lock: + if detector is not None and (time.time() - last_used_time > MODEL_TIMEOUT): + logger.info("Unloading idle AudioSeal detector to free RAM.") + del detector + detector = None + gc.collect() + + +# ── FastAPI app ──────────────────────────────────────────────────────────── + +app = FastAPI( + title="AudioSeal Watermark Detection Service", + description=( + "Detects Meta AudioSeal watermarks embedded in audio files " + "using the audioseal_detector_16bits model." + ), + version="1.0.0", + docs_url=None if _PRODUCTION else "/docs", + redoc_url=None if _PRODUCTION else "/redoc", + openapi_url=None if _PRODUCTION else "/openapi.json", +) + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + + +@app.get("/") +async def root(): + """Root endpoint with service information.""" + if _PRODUCTION: + return {"status": "ok"} + return { + "model_name": MODEL_NAME, + "description": ("AudioSeal watermark detector for AI-generated audio"), + "device": str(DEVICE), + "model_loaded": detector is not None, + } + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and DEVICE.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +async def health(): + """Health check endpoint.""" + if _PRODUCTION: + return {"status": "healthy"} + return { + "status": "healthy", + "model_name": MODEL_NAME, + "device": str(DEVICE), + "model_loaded": detector is not None, + **_gpu_health_info(), + } + + +@app.post("/predict") +async def predict(payload: AudioInput) -> Dict[str, Any]: + """Detect AudioSeal watermark in base64-encoded audio. + + Decodes the audio, resamples to 16 kHz mono, and runs the + AudioSeal detector to check for embedded watermarks. + + Args: + payload: Base64-encoded audio data and optional threshold. + + Returns: + Dict with model name, probability, prediction, class, and + inference time. + """ + try: + ensure_model_loaded() + if detector is None: + raise HTTPException(status_code=503, detail="Model not loaded.") + + start = time.time() + + # Decode base64 audio + try: + audio_bytes = base64.b64decode(payload.audio_data) + except Exception as exc: + raise HTTPException( + status_code=400, + detail=f"Invalid base64 data: {exc}", + ) + + if not audio_bytes: + raise HTTPException( + status_code=400, + detail="Empty audio data after base64 decode.", + ) + + # Write to temp file and load with torchaudio + tmp_path = None + try: + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: + tmp.write(audio_bytes) + tmp_path = tmp.name + + # Try torchaudio first, fall back to soundfile if torchcodec + # or other backends are unavailable. + try: + waveform, sample_rate = torchaudio.load(tmp_path) + except Exception as ta_exc: + logger.info( + "torchaudio.load failed (%s), falling back to soundfile.", + ta_exc, + ) + import soundfile as sf + + data, sample_rate = sf.read(tmp_path, dtype="float32") + waveform = torch.from_numpy(data).T # [channels, samples] + if waveform.dim() == 1: + waveform = waveform.unsqueeze(0) + except Exception as exc: + logger.warning("Failed to load audio: %s", exc) + inference_time = time.time() - start + return { + "model": MODEL_NAME, + "probability": 0.5, + "prediction": 0, + "class": "real", + "inference_time": float(inference_time), + } + finally: + if tmp_path and os.path.exists(tmp_path): + try: + os.unlink(tmp_path) + except OSError: + pass + + # Stereo to mono: average channels + if waveform.shape[0] > 1: + waveform = waveform.mean(dim=0, keepdim=True) + + # Resample to target sample rate + if sample_rate != TARGET_SAMPLE_RATE: + resampler = torchaudio.transforms.Resample( + orig_freq=sample_rate, + new_freq=TARGET_SAMPLE_RATE, + ) + waveform = resampler(waveform) + + # Add batch dimension: [1, channels, samples] + if waveform.dim() == 2: + waveform = waveform.unsqueeze(0) + + waveform = waveform.to(DEVICE) + + # Run detection + with torch.no_grad(): + result, message = detector.detect_watermark( + waveform, sample_rate=TARGET_SAMPLE_RATE + ) + + confidence = float(result) + probability = _map_confidence_to_probability(confidence) + + prediction = 1 if probability > payload.threshold else 0 + class_label = "fake" if prediction == 1 else "real" + inference_time = time.time() - start + + logger.info( + "AudioSeal check: %s (conf=%.4f, prob=%.4f, %.3fs)", + class_label, + confidence, + probability, + inference_time, + ) + + if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: + threading.Timer(MODEL_TIMEOUT + 5.0, unload_model_if_idle).start() + + return { + "model": MODEL_NAME, + "probability": float(probability), + "prediction": int(prediction), + "class": class_label, + "inference_time": float(inference_time), + } + + except HTTPException: + raise + except Exception as exc: + logger.exception("Prediction error: %s", exc) + raise HTTPException(status_code=500, detail=str(exc)) + + +@app.on_event("startup") +async def startup_event(): + """Startup: preload model if configured, else lazy-load on first request.""" + if PRELOAD_MODEL: + logger.info("Preloading AudioSeal detector at startup.") + try: + load_model_internal() + except Exception as exc: + logger.error("Preload failed: %s", exc) + else: + logger.info("AudioSeal service ready — model loads on first request.") + + if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: + + def _periodic_check(): + unload_model_if_idle() + threading.Timer(MODEL_TIMEOUT / 2.0, _periodic_check).start() + + threading.Timer(MODEL_TIMEOUT / 2.0, _periodic_check).start() + + +if __name__ == "__main__": + logger.info("Starting AudioSeal detector service on port %d", MODEL_PORT) + uvicorn.run(app, host="0.0.0.0", port=MODEL_PORT) diff --git a/provenance/audioseal/requirements.txt b/provenance/audioseal/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..bac57441f2b8e024c64d30e95914b7fefe3f381e --- /dev/null +++ b/provenance/audioseal/requirements.txt @@ -0,0 +1,8 @@ +fastapi==0.115.12 +uvicorn==0.34.2 +pydantic==2.11.1 +audioseal==0.1.4 +torch>=2.0.0 +torchaudio>=2.0.0 +numpy>=1.24,<2.0 +soundfile>=0.12.0 diff --git a/provenance/audioseal/test_app.py b/provenance/audioseal/test_app.py new file mode 100644 index 0000000000000000000000000000000000000000..a16d4a08b0751f90652c2e83e5b947b893b62e9f --- /dev/null +++ b/provenance/audioseal/test_app.py @@ -0,0 +1,447 @@ +"""Tests for the AudioSeal watermark detection service.""" + +import base64 +import struct +from unittest.mock import MagicMock, patch + +import pytest +from fastapi.testclient import TestClient + +from app import _map_confidence_to_probability, app + +client = TestClient(app) + + +# ── Fixture helpers ──────────────────────────────────────────────────────── + + +def _make_wav( + sample_rate: int = 16000, + duration_s: float = 1.0, + channels: int = 1, +) -> bytes: + """Create a valid PCM WAV file with RIFF header. + + Generates silence (all zeros) as audio data. + + Args: + sample_rate: Sample rate in Hz. + duration_s: Duration in seconds. + channels: Number of audio channels (1=mono, 2=stereo). + + Returns: + Raw WAV file bytes. + """ + bits_per_sample = 16 + num_samples = int(sample_rate * duration_s) + byte_rate = sample_rate * channels * bits_per_sample // 8 + block_align = channels * bits_per_sample // 8 + data_size = num_samples * channels * bits_per_sample // 8 + + header = struct.pack( + "<4sI4s4sIHHIIHH4sI", + b"RIFF", + 36 + data_size, + b"WAVE", + b"fmt ", + 16, + 1, # PCM format + channels, + sample_rate, + byte_rate, + block_align, + bits_per_sample, + b"data", + data_size, + ) + return header + b"\x00" * data_size + + +def _make_noise_wav( + sample_rate: int = 16000, + duration_s: float = 1.0, + channels: int = 1, +) -> bytes: + """Create a valid PCM WAV with white noise. + + Args: + sample_rate: Sample rate in Hz. + duration_s: Duration in seconds. + channels: Number of audio channels. + + Returns: + Raw WAV file bytes. + """ + import numpy as np + + bits_per_sample = 16 + num_samples = int(sample_rate * duration_s) + byte_rate = sample_rate * channels * bits_per_sample // 8 + block_align = channels * bits_per_sample // 8 + data_size = num_samples * channels * bits_per_sample // 8 + + header = struct.pack( + "<4sI4s4sIHHIIHH4sI", + b"RIFF", + 36 + data_size, + b"WAVE", + b"fmt ", + 16, + 1, # PCM format + channels, + sample_rate, + byte_rate, + block_align, + bits_per_sample, + b"data", + data_size, + ) + + rng = np.random.RandomState(42) + samples = rng.randint(-32768, 32767, size=num_samples * channels, dtype=np.int16) + return header + samples.tobytes() + + +def _b64(raw: bytes) -> str: + """Encode raw bytes as base64 string.""" + return base64.b64encode(raw).decode("utf-8") + + +# ── Mock detector ───────────────────────────────────────────────────────── + +# All /predict tests mock the AudioSeal model so tests run without GPU +# or the audioseal package installed. The mock returns (confidence, message). + + +def _mock_detect_watermark(waveform, sample_rate=16000): + """Mock detector that returns low confidence (no watermark).""" + return (0.1, "no watermark") + + +def _patch_detector(): + """Create a mock detector with detect_watermark method.""" + mock = MagicMock() + mock.detect_watermark = MagicMock(side_effect=_mock_detect_watermark) + return mock + + +# ── Health endpoint ──────────────────────────────────────────────────────── + + +class TestHealthEndpoint: + """Tests for GET /health.""" + + def test_health_returns_healthy(self): + """Health endpoint returns healthy status.""" + resp = client.get("/health") + assert resp.status_code == 200 + data = resp.json() + assert data["status"] == "healthy" + + def test_health_returns_model_name(self): + """Health endpoint includes correct model_name.""" + resp = client.get("/health") + data = resp.json() + assert data["model_name"] == "audioseal_detector" + + def test_health_returns_device(self): + """Health endpoint includes device field.""" + resp = client.get("/health") + data = resp.json() + assert "device" in data + + def test_health_returns_model_loaded(self): + """Health endpoint includes model_loaded field.""" + resp = client.get("/health") + data = resp.json() + assert "model_loaded" in data + + +# ── Root endpoint ────────────────────────────────────────────────────────── + + +class TestRootEndpoint: + """Tests for GET /.""" + + def test_root_returns_service_info(self): + """Root endpoint returns model info.""" + resp = client.get("/") + assert resp.status_code == 200 + data = resp.json() + assert data["model_name"] == "audioseal_detector" + + def test_root_returns_description(self): + """Root endpoint includes a description.""" + resp = client.get("/") + data = resp.json() + assert "description" in data + assert len(data["description"]) > 0 + + +# ── Predict endpoint: silent audio ────────────────────────────────────── + + +class TestPredictSilentAudio: + """Tests for POST /predict with silent (no watermark) WAV audio.""" + + @patch("app.detector", new_callable=_patch_detector) + @patch("app.ensure_model_loaded") + def test_silent_audio_returns_neutral(self, mock_load, mock_det): + """Silent WAV with no watermark returns probability=0.5.""" + resp = client.post( + "/predict", + json={"audio_data": _b64(_make_wav())}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["probability"] == 0.5 + assert data["class"] == "real" + + +# ── Predict endpoint: white noise ──────────────────────────────────────── + + +class TestPredictWhiteNoise: + """Tests for POST /predict with white noise audio.""" + + @patch("app.detector", new_callable=_patch_detector) + @patch("app.ensure_model_loaded") + def test_white_noise_returns_neutral(self, mock_load, mock_det): + """White noise WAV returns probability=0.5 (no watermark).""" + resp = client.post( + "/predict", + json={"audio_data": _b64(_make_noise_wav())}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["probability"] == 0.5 + assert data["class"] == "real" + + +# ── Predict endpoint: sample rate handling ─────────────────────────────── + + +class TestPredictSampleRates: + """Tests for different audio sample rates.""" + + @patch("app.detector", new_callable=_patch_detector) + @patch("app.ensure_model_loaded") + def test_16khz_handled(self, mock_load, mock_det): + """16 kHz audio is processed without error.""" + resp = client.post( + "/predict", + json={"audio_data": _b64(_make_wav(sample_rate=16000))}, + ) + assert resp.status_code == 200 + data = resp.json() + assert 0.0 <= data["probability"] <= 1.0 + + @patch("app.detector", new_callable=_patch_detector) + @patch("app.ensure_model_loaded") + def test_44100hz_handled(self, mock_load, mock_det): + """44.1 kHz audio is resampled and processed without error.""" + resp = client.post( + "/predict", + json={"audio_data": _b64(_make_wav(sample_rate=44100))}, + ) + assert resp.status_code == 200 + data = resp.json() + assert 0.0 <= data["probability"] <= 1.0 + + @patch("app.detector", new_callable=_patch_detector) + @patch("app.ensure_model_loaded") + def test_48khz_handled(self, mock_load, mock_det): + """48 kHz audio is resampled and processed without error.""" + resp = client.post( + "/predict", + json={"audio_data": _b64(_make_wav(sample_rate=48000))}, + ) + assert resp.status_code == 200 + data = resp.json() + assert 0.0 <= data["probability"] <= 1.0 + + +# ── Predict endpoint: stereo audio ────────────────────────────────────── + + +class TestPredictStereoAudio: + """Tests for stereo audio handling.""" + + @patch("app.detector", new_callable=_patch_detector) + @patch("app.ensure_model_loaded") + def test_stereo_audio_handled(self, mock_load, mock_det): + """Stereo WAV is averaged to mono and processed without error.""" + resp = client.post( + "/predict", + json={ + "audio_data": _b64(_make_wav(channels=2, duration_s=1.0)), + }, + ) + assert resp.status_code == 200 + data = resp.json() + assert 0.0 <= data["probability"] <= 1.0 + + +# ── Predict endpoint: short audio ─────────────────────────────────────── + + +class TestPredictShortAudio: + """Tests for short duration audio.""" + + @patch("app.detector", new_callable=_patch_detector) + @patch("app.ensure_model_loaded") + def test_short_audio_handled(self, mock_load, mock_det): + """Short 0.5s WAV is processed without error.""" + resp = client.post( + "/predict", + json={ + "audio_data": _b64(_make_wav(duration_s=0.5)), + }, + ) + assert resp.status_code == 200 + data = resp.json() + assert 0.0 <= data["probability"] <= 1.0 + + +# ── Predict endpoint: error handling ───────────────────────────────────── + + +class TestPredictErrorHandling: + """Tests for error cases.""" + + @patch("app.detector", new_callable=_patch_detector) + @patch("app.ensure_model_loaded") + def test_invalid_base64_returns_error(self, mock_load, mock_det): + """Invalid base64 string returns 400.""" + resp = client.post( + "/predict", + json={"audio_data": "!!!not-valid-base64!!!"}, + ) + assert resp.status_code == 400 + + @patch("app.detector", new_callable=_patch_detector) + @patch("app.ensure_model_loaded") + def test_non_audio_bytes_returns_neutral(self, mock_load, mock_det): + """Non-audio binary data returns probability=0.5, not crash.""" + random_bytes = base64.b64encode(b"\x00\x01\x02\x03" * 20).decode() + resp = client.post( + "/predict", + json={"audio_data": random_bytes}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["probability"] == 0.5 + + +# ── Predict endpoint: response format ──────────────────────────────────── + + +class TestPredictResponseSchema: + """Tests for the standard response format.""" + + @patch("app.detector", new_callable=_patch_detector) + @patch("app.ensure_model_loaded") + def test_response_has_all_required_fields(self, mock_load, mock_det): + """Response contains model, probability, prediction, class, + inference_time.""" + resp = client.post( + "/predict", + json={"audio_data": _b64(_make_wav())}, + ) + data = resp.json() + required = { + "model", + "probability", + "prediction", + "class", + "inference_time", + } + assert required.issubset(data.keys()) + + @patch("app.detector", new_callable=_patch_detector) + @patch("app.ensure_model_loaded") + def test_model_name_in_response(self, mock_load, mock_det): + """Response model field matches audioseal_detector.""" + resp = client.post( + "/predict", + json={"audio_data": _b64(_make_wav())}, + ) + data = resp.json() + assert data["model"] == "audioseal_detector" + + @patch("app.detector", new_callable=_patch_detector) + @patch("app.ensure_model_loaded") + def test_probability_is_float_in_range(self, mock_load, mock_det): + """Probability is a float in [0, 1].""" + resp = client.post( + "/predict", + json={"audio_data": _b64(_make_wav())}, + ) + data = resp.json() + assert isinstance(data["probability"], float) + assert 0.0 <= data["probability"] <= 1.0 + + @patch("app.detector", new_callable=_patch_detector) + @patch("app.ensure_model_loaded") + def test_inference_time_is_positive_float(self, mock_load, mock_det): + """Inference time is a positive float.""" + resp = client.post( + "/predict", + json={"audio_data": _b64(_make_wav())}, + ) + data = resp.json() + assert isinstance(data["inference_time"], float) + assert data["inference_time"] > 0 + + @patch("app.detector", new_callable=_patch_detector) + @patch("app.ensure_model_loaded") + def test_threshold_controls_prediction(self, mock_load, mock_det): + """Prediction uses > threshold (0.5 is neutral, not fake).""" + resp = client.post( + "/predict", + json={"audio_data": _b64(_make_wav()), "threshold": 0.5}, + ) + data = resp.json() + # probability=0.5 (from mock), threshold=0.5 -> not > -> real + assert data["prediction"] == 0 + assert data["class"] == "real" + + +# ── _map_confidence_to_probability helper ──────────────────────────────── + + +class TestMapConfidenceToProbability: + """Tests for the _map_confidence_to_probability helper function.""" + + def test_high_confidence_returns_high(self): + """High confidence (0.9) maps to high probability.""" + result = _map_confidence_to_probability(0.9) + assert result == 0.5 + 0.9 * 0.5 # 0.95 + assert result > 0.9 + + def test_medium_confidence_returns_medium(self): + """Medium confidence (0.5) maps to medium-high probability.""" + result = _map_confidence_to_probability(0.5) + assert result == 0.5 + 0.5 * 0.5 # 0.75 + + def test_below_threshold_returns_neutral(self): + """Confidence below 0.3 threshold returns 0.5 (neutral).""" + assert _map_confidence_to_probability(0.1) == 0.5 + assert _map_confidence_to_probability(0.0) == 0.5 + assert _map_confidence_to_probability(0.29) == 0.5 + + def test_boundary_at_threshold(self): + """Confidence at exactly 0.3 maps above neutral.""" + result = _map_confidence_to_probability(0.3) + assert result == 0.5 + 0.3 * 0.5 # 0.65 + assert result > 0.5 + + def test_just_below_threshold(self): + """Confidence just below 0.3 returns neutral.""" + result = _map_confidence_to_probability(0.299) + assert result == 0.5 + + def test_confidence_one_returns_max(self): + """Confidence of 1.0 returns maximum probability (1.0).""" + result = _map_confidence_to_probability(1.0) + assert result == 1.0 diff --git a/provenance/c2pa-checker/Dockerfile b/provenance/c2pa-checker/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..9439f075f8da17b428792a5f9208b71089f876c8 --- /dev/null +++ b/provenance/c2pa-checker/Dockerfile @@ -0,0 +1,22 @@ +FROM python:3.10-slim + +# System deps: exiftool for EXIF/IPTC metadata extraction +RUN apt-get update && \ + apt-get install -y --no-install-recommends libimage-exiftool-perl && \ + rm -rf /var/lib/apt/lists/* + +WORKDIR /app + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY app.py . + +ENV MODEL_PORT=9001 +EXPOSE 9001 + +# Drop root privileges +RUN adduser --disabled-password --gecos "" appuser +USER appuser + +CMD ["python", "app.py"] diff --git a/provenance/c2pa-checker/app.py b/provenance/c2pa-checker/app.py new file mode 100644 index 0000000000000000000000000000000000000000..40be0556420ee43ae163d1e9f3cd8f34e738bd6c --- /dev/null +++ b/provenance/c2pa-checker/app.py @@ -0,0 +1,417 @@ +"""C2PA + IPTC + EXIF provenance detection service. + +Checks media files for AI-generation provenance signals embedded in C2PA +manifests, EXIF metadata, and IPTC fields. Returns a probability score +indicating likelihood of AI generation based on metadata evidence alone. + +Unlike detection models that analyze visual/audio artifacts, this service +inspects only embedded metadata and watermarks. +""" + +import base64 +import logging +import os +import sys +import tempfile +import time +from typing import Any, Dict, Optional, Set + +import uvicorn +from fastapi import FastAPI, HTTPException +from fastapi.middleware.cors import CORSMiddleware +from pydantic import BaseModel + +# ── Logging ──────────────────────────────────────────────────────────────── + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + handlers=[logging.StreamHandler(sys.stdout)], +) +logger = logging.getLogger(__name__) + +# ── Config ───────────────────────────────────────────────────────────────── + +MODEL_NAME = "c2pa_checker" +MODEL_PORT = int(os.environ.get("MODEL_PORT", "9001")) + +_PRODUCTION = os.environ.get("PRODUCTION", "false").lower() == "true" + +# Known AI generator strings to match in C2PA manifests and EXIF metadata. +_AI_GENERATORS: Set[str] = { + "adobe firefly", + "dall-e", + "dall\u00b7e", + "openai", + "midjourney", + "stable diffusion", + "stability ai", + "google gemini", + "imagen", + "microsoft designer", + "copilot", + "meta ai", + "leonardo ai", + "runway", + "sora", + "kling", + "flux", + "ideogram", + "veo", + "google generative ai", + "trainedalgorithmicmedia", +} + +# ── Pydantic models ─────────────────────────────────────────────────────── + + +class PredictInput(BaseModel): + """Request body for /predict.""" + + image_data: Optional[str] = None + audio_data: Optional[str] = None + video_data: Optional[str] = None + threshold: float = 0.5 + + +# ── Detection helpers ────────────────────────────────────────────────────── + + +def _check_c2pa(file_path: str) -> tuple: + """Check file for C2PA manifest with AI generator assertions. + + Verifies the cryptographic signature chain. Unsigned or + invalid manifests are downweighted because they can be forged. + + Args: + file_path: Path to the media file on disk. + + Returns: + Tuple of (score_or_none, has_valid_c2pa). + Score is 0.95 for signed AI generator, 0.60 for unsigned, + None if no C2PA manifest. has_valid_c2pa indicates whether + a cryptographically valid manifest exists (used to gate + EXIF/IPTC trust). + """ + try: + import c2pa + + reader = c2pa.Reader.try_create(file_path) + if reader is None: + return None, False + manifest_json = reader.json() + + # Check signature validation status via manifest JSON + # reader.is_valid is a bool property (not a method) + sig_valid = True + try: + if hasattr(reader, "is_valid"): + sig_valid = reader.is_valid + # Also check validation_status in the JSON for untrusted certs + import json as _json + + mdata = _json.loads(manifest_json) + vstatus = mdata.get("validation_status", []) + if vstatus: + # untrusted signing credential is common for valid AI certs + # only mark invalid for actual signature failures + sig_codes = [v.get("code", "") for v in vstatus] + has_sig_failure = any( + "signature" in c and "untrusted" not in c for c in sig_codes + ) + if has_sig_failure: + sig_valid = False + logger.info("C2PA manifest has signature issues: %s", sig_codes) + else: + logger.info( + "C2PA manifest has untrusted cert (normal for AI): %s", + sig_codes, + ) + except Exception: + pass + + # Search for AI generator strings + manifest_lower = manifest_json.lower() + for generator in _AI_GENERATORS: + if generator in manifest_lower: + if sig_valid: + logger.info( + "C2PA (verified): AI generator '%s' detected.", + generator, + ) + return 0.95, True + else: + logger.info( + "C2PA (UNVERIFIED): AI generator '%s' — " + "signature invalid, downweighting.", + generator, + ) + return 0.60, False + + # C2PA manifest exists but no AI generator + logger.info("C2PA manifest found but no AI generator detected.") + return None, sig_valid + + except Exception: + return None, False + + +def _check_exif_iptc( + file_path: str, + has_valid_c2pa: bool = False, +) -> Optional[float]: + """Check file EXIF/IPTC metadata for AI generation indicators. + + IPTC/EXIF tags are trivially forgeable (one exiftool command). + They are only trusted when backed by a valid C2PA manifest. + Without C2PA, they receive near-neutral scores (0.55) that + contribute minimal ensemble impact. + + Args: + file_path: Path to the media file on disk. + has_valid_c2pa: Whether a cryptographically valid C2PA + manifest was found. Raises trust in EXIF/IPTC tags. + + Returns: + 0.90 if backed by C2PA + DigitalSourceType AI tag, + 0.85 if backed by C2PA + EXIF AI tool signature, + None if IPTC/EXIF AI indicators found WITHOUT C2PA (no trust), + None if no AI indicators found. + """ + try: + import exiftool + + with exiftool.ExifToolHelper() as et: + metadata_list = et.get_metadata(file_path) + if not metadata_list: + return None + + metadata = metadata_list[0] + + # Check IPTC DigitalSourceType + for key, value in metadata.items(): + key_lower = key.lower() + if "digitalsourcetype" in key_lower: + value_str = str(value).lower() + if "trainedalgorithmicmedia" in value_str: + if has_valid_c2pa: + logger.info("IPTC AI tag (C2PA-backed): %s", value) + return 0.90 + else: + logger.info( + "IPTC AI tag (NO C2PA — low trust): %s", + value, + ) + return None # No trust without C2PA + + # Check EXIF Software, Description, UserComment + text_fields = [] + for key, value in metadata.items(): + key_lower = key.lower() + if any( + field in key_lower + for field in ("software", "description", "usercomment") + ): + text_fields.append(str(value).lower()) + + combined_text = " ".join(text_fields) + for generator in _AI_GENERATORS: + if generator in combined_text: + if has_valid_c2pa: + logger.info( + "EXIF AI signature (C2PA-backed): %s", + generator, + ) + return 0.85 + else: + logger.info( + "EXIF AI signature (NO C2PA — ignored): %s", + generator, + ) + return None # No trust without C2PA + + return None + + except Exception: + return None + + +def _detect_provenance(file_bytes: bytes, data_key: str = "") -> float: + """Run all provenance checks on raw file bytes. + + Writes bytes to a temporary file, runs C2PA and EXIF/IPTC checks, + and returns the maximum signal found. Returns 0.5 (neutral) when + no AI provenance is detected. + + Args: + file_bytes: Raw bytes of the media file. + data_key: The payload key used (image_data, audio_data, + video_data) to determine file extension for C2PA. + + Returns: + Float probability in [0, 1]. 0.5 means neutral (no signal). + """ + if not file_bytes: + return 0.5 + + # C2PA needs a file extension to determine MIME type + _EXT_MAP = { + "video_data": ".mp4", + "image_data": ".jpg", + "audio_data": ".wav", + } + suffix = _EXT_MAP.get(data_key, ".bin") + + tmp_path = None + try: + with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: + tmp.write(file_bytes) + tmp_path = tmp.name + + signals = [] + + c2pa_score, has_valid_c2pa = _check_c2pa(tmp_path) + if c2pa_score is not None: + signals.append(c2pa_score) + + exif_result = _check_exif_iptc(tmp_path, has_valid_c2pa) + if exif_result is not None: + signals.append(exif_result) + + if signals: + return max(signals) + return 0.5 + + except Exception as exc: + logger.warning("Provenance detection error: %s", exc) + return 0.5 + + finally: + if tmp_path and os.path.exists(tmp_path): + try: + os.unlink(tmp_path) + except OSError: + pass + + +# ── FastAPI app ──────────────────────────────────────────────────────────── + +app = FastAPI( + title="C2PA Provenance Checker Service", + description=( + "Checks media files for AI-generation provenance signals in " + "C2PA manifests, EXIF metadata, and IPTC fields." + ), + version="1.0.0", + docs_url=None if _PRODUCTION else "/docs", + redoc_url=None if _PRODUCTION else "/redoc", + openapi_url=None if _PRODUCTION else "/openapi.json", +) + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + + +@app.get("/") +async def root(): + """Root endpoint with service information.""" + if _PRODUCTION: + return {"status": "ok"} + return { + "model_name": MODEL_NAME, + "description": ("C2PA + EXIF/IPTC provenance checker for AI-generated media"), + "modalities": ["image", "audio", "video"], + } + + +@app.get("/health") +async def health(): + """Health check endpoint.""" + if _PRODUCTION: + return {"status": "healthy"} + return { + "status": "healthy", + "model_name": MODEL_NAME, + "ready": True, + } + + +@app.post("/predict") +async def predict(payload: PredictInput) -> Dict[str, Any]: + """Check media for AI-generation provenance signals. + + Accepts image_data, audio_data, or video_data as base64-encoded + file bytes. Returns a standard detection response with probability + based on metadata evidence. + + Args: + payload: Base64-encoded media and optional threshold. + + Returns: + Dict with model name, probability, prediction, class, and + inference time. + """ + # Find the first non-None payload key + raw_b64 = None + data_key = "" + for key in ("image_data", "audio_data", "video_data"): + value = getattr(payload, key) + if value is not None: + raw_b64 = value + data_key = key + break + + if raw_b64 is None: + raise HTTPException( + status_code=400, + detail=( + "No media data provided. Include one of: " + "image_data, audio_data, video_data." + ), + ) + + start = time.time() + + try: + file_bytes = base64.b64decode(raw_b64) + except Exception as exc: + raise HTTPException( + status_code=400, + detail=f"Invalid base64 data: {exc}", + ) + + if not file_bytes: + raise HTTPException( + status_code=400, + detail="Empty file data after base64 decode.", + ) + + probability = _detect_provenance(file_bytes, data_key) + prediction = 1 if probability > payload.threshold else 0 + class_label = "fake" if prediction == 1 else "real" + inference_time = time.time() - start + + logger.info( + "Provenance check: %s (prob=%.4f, %.3fs)", + class_label, + probability, + inference_time, + ) + + return { + "model": MODEL_NAME, + "probability": float(probability), + "prediction": int(prediction), + "class": class_label, + "inference_time": float(inference_time), + } + + +if __name__ == "__main__": + logger.info("Starting C2PA checker service on port %d", MODEL_PORT) + uvicorn.run(app, host="0.0.0.0", port=MODEL_PORT) diff --git a/provenance/c2pa-checker/requirements.txt b/provenance/c2pa-checker/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..53a49028e3b59e0152ffa0fc6b92d26521000d68 --- /dev/null +++ b/provenance/c2pa-checker/requirements.txt @@ -0,0 +1,5 @@ +fastapi==0.115.12 +uvicorn==0.34.2 +pydantic==2.11.1 +c2pa-python==0.30.0 +PyExifTool==0.5.6 diff --git a/provenance/c2pa-checker/test_app.py b/provenance/c2pa-checker/test_app.py new file mode 100644 index 0000000000000000000000000000000000000000..1c903ed15da37b4796958fb4d7c11167761b385a --- /dev/null +++ b/provenance/c2pa-checker/test_app.py @@ -0,0 +1,275 @@ +"""Tests for the c2pa-checker provenance detection service.""" + +import base64 +import io +import struct + +import pytest +from fastapi.testclient import TestClient + +from app import app + +client = TestClient(app) + + +# ── Fixture helpers ──────────────────────────────────────────────────────── + + +def _make_minimal_jpeg() -> bytes: + """Create a minimal valid JPEG from an 8x8 RGB image.""" + from PIL import Image + + img = Image.new("RGB", (8, 8)) + buf = io.BytesIO() + img.save(buf, format="JPEG") + return buf.getvalue() + + +def _make_minimal_wav() -> bytes: + """Create a minimal valid WAV: RIFF/WAVE header + 1s silence 16kHz mono 16-bit.""" + sample_rate = 16000 + num_samples = sample_rate # 1 second + bits_per_sample = 16 + num_channels = 1 + byte_rate = sample_rate * num_channels * bits_per_sample // 8 + block_align = num_channels * bits_per_sample // 8 + data_size = num_samples * block_align + + header = struct.pack( + "<4sI4s4sIHHIIHH4sI", + b"RIFF", + 36 + data_size, + b"WAVE", + b"fmt ", + 16, + 1, # PCM + num_channels, + sample_rate, + byte_rate, + block_align, + bits_per_sample, + b"data", + data_size, + ) + return header + b"\x00" * data_size + + +def _make_minimal_mp4() -> bytes: + """Create a minimal valid MP4: ftyp box only.""" + return b"\x00\x00\x00\x14ftypmp42\x00\x00\x00\x00mp42" + + +def _b64(raw: bytes) -> str: + """Encode raw bytes as base64 string.""" + return base64.b64encode(raw).decode("utf-8") + + +# ── Health endpoint ──────────────────────────────────────────────────────── + + +class TestHealthEndpoint: + """Tests for GET /health.""" + + def test_health_returns_healthy(self): + """Health endpoint returns healthy status.""" + resp = client.get("/health") + assert resp.status_code == 200 + data = resp.json() + assert data["status"] == "healthy" + + def test_health_returns_model_name(self): + """Health endpoint includes correct model_name.""" + resp = client.get("/health") + data = resp.json() + assert data["model_name"] == "c2pa_checker" + + +# ── Root endpoint ────────────────────────────────────────────────────────── + + +class TestRootEndpoint: + """Tests for GET /.""" + + def test_root_returns_service_info(self): + """Root endpoint returns model info.""" + resp = client.get("/") + assert resp.status_code == 200 + data = resp.json() + assert "model_name" in data + assert data["model_name"] == "c2pa_checker" + + +# ── Predict endpoint: clean files ────────────────────────────────────────── + + +class TestPredictCleanFiles: + """Tests for POST /predict with clean (no provenance) media.""" + + def test_clean_jpeg_returns_neutral(self): + """Clean JPEG with no C2PA/EXIF markers returns probability=0.5.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_minimal_jpeg())}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["probability"] == 0.5 + assert data["class"] == "real" + + def test_clean_wav_returns_neutral(self): + """Clean WAV with no provenance returns probability=0.5.""" + resp = client.post( + "/predict", + json={"audio_data": _b64(_make_minimal_wav())}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["probability"] == 0.5 + + def test_clean_mp4_returns_neutral(self): + """Clean MP4 with no provenance returns probability=0.5.""" + resp = client.post( + "/predict", + json={"video_data": _b64(_make_minimal_mp4())}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["probability"] == 0.5 + + +# ── Predict endpoint: payload key acceptance ─────────────────────────────── + + +class TestPredictPayloadKeys: + """Tests for accepted payload keys.""" + + def test_accepts_image_data(self): + """image_data key is accepted.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_minimal_jpeg())}, + ) + assert resp.status_code == 200 + + def test_accepts_audio_data(self): + """audio_data key is accepted.""" + resp = client.post( + "/predict", + json={"audio_data": _b64(_make_minimal_wav())}, + ) + assert resp.status_code == 200 + + def test_accepts_video_data(self): + """video_data key is accepted.""" + resp = client.post( + "/predict", + json={"video_data": _b64(_make_minimal_mp4())}, + ) + assert resp.status_code == 200 + + def test_no_payload_key_returns_400(self): + """Request with no recognized payload key returns 400.""" + resp = client.post("/predict", json={"threshold": 0.5}) + assert resp.status_code == 400 + + +# ── Predict endpoint: error handling ─────────────────────────────────────── + + +class TestPredictErrorHandling: + """Tests for error cases.""" + + def test_invalid_base64_returns_error(self): + """Invalid base64 string returns 400.""" + resp = client.post( + "/predict", + json={"image_data": "!!!not-valid-base64!!!"}, + ) + assert resp.status_code == 400 + + def test_empty_base64_returns_error(self): + """Empty base64 string returns 400.""" + resp = client.post( + "/predict", + json={"image_data": ""}, + ) + assert resp.status_code == 400 + + def test_corrupt_file_bytes_returns_neutral(self): + """Corrupt file bytes should return probability=0.5, not crash.""" + corrupt = base64.b64encode(b"\x00\x01\x02\x03" * 20).decode() + resp = client.post( + "/predict", + json={"image_data": corrupt}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["probability"] == 0.5 + + +# ── Predict endpoint: response schema ───────────────────────────────────── + + +class TestPredictResponseSchema: + """Tests for the standard response format.""" + + def test_response_has_all_required_fields(self): + """Response contains model, probability, prediction, class, inference_time.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_minimal_jpeg())}, + ) + data = resp.json() + required = {"model", "probability", "prediction", "class", "inference_time"} + assert required.issubset(data.keys()) + + def test_probability_is_float_in_range(self): + """Probability is a float in [0, 1].""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_minimal_jpeg())}, + ) + data = resp.json() + assert isinstance(data["probability"], float) + assert 0.0 <= data["probability"] <= 1.0 + + def test_prediction_matches_threshold_logic(self): + """Prediction=1 when probability > threshold, else 0.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_minimal_jpeg()), "threshold": 0.6}, + ) + data = resp.json() + # probability=0.5, threshold=0.6 -> 0.5 < 0.6 -> prediction=0, class=real + assert data["prediction"] == 0 + assert data["class"] == "real" + + def test_prediction_fake_when_above_threshold(self): + """When probability >= threshold, prediction=1 and class=fake.""" + # With threshold=0.4 and probability=0.5: 0.5 >= 0.4 -> prediction=1 + resp = client.post( + "/predict", + json={"image_data": _b64(_make_minimal_jpeg()), "threshold": 0.4}, + ) + data = resp.json() + assert data["prediction"] == 1 + assert data["class"] == "fake" + + def test_model_name_in_response(self): + """Response model field matches c2pa_checker.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_minimal_jpeg())}, + ) + data = resp.json() + assert data["model"] == "c2pa_checker" + + def test_inference_time_is_positive_float(self): + """Inference time is a positive float.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_minimal_jpeg())}, + ) + data = resp.json() + assert isinstance(data["inference_time"], float) + assert data["inference_time"] > 0 diff --git a/provenance/invisible-watermark/Dockerfile b/provenance/invisible-watermark/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..c118471c913651bea84735b433190b8f31d0faa7 --- /dev/null +++ b/provenance/invisible-watermark/Dockerfile @@ -0,0 +1,22 @@ +FROM python:3.10-slim + +# System deps: OpenCV requires libgl1 and libglib2.0 +RUN apt-get update && \ + apt-get install -y --no-install-recommends libgl1 libglib2.0-0 && \ + rm -rf /var/lib/apt/lists/* + +WORKDIR /app + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY app.py . + +ENV MODEL_PORT=9002 +EXPOSE 9002 + +# Drop root privileges +RUN adduser --disabled-password --gecos "" appuser +USER appuser + +CMD ["python", "app.py"] diff --git a/provenance/invisible-watermark/app.py b/provenance/invisible-watermark/app.py new file mode 100644 index 0000000000000000000000000000000000000000..e366de8748403c66402fe62d7682a01efb496668 --- /dev/null +++ b/provenance/invisible-watermark/app.py @@ -0,0 +1,311 @@ +"""SDXL invisible watermark detection service. + +Detects Stable Diffusion XL invisible watermarks embedded in images using +the invisible-watermark library. SDXL embeds a specific 136-bit pattern +(starting with "SDV2") via DWT-DCT encoding during generation. + +Returns a probability score based on whether the decoded bytes match +known SDXL watermark patterns or exhibit non-random structure. +""" + +import base64 +import io +import logging +import math +import os +import sys +import time +from collections import Counter +from typing import Any, Dict + +import cv2 +import numpy as np +import uvicorn +from fastapi import FastAPI, HTTPException +from fastapi.middleware.cors import CORSMiddleware +from PIL import Image +from pydantic import BaseModel + +# ── Logging ──────────────────────────────────────────────────────────────── + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + handlers=[logging.StreamHandler(sys.stdout)], +) +logger = logging.getLogger(__name__) + +# ── Config ───────────────────────────────────────────────────────────────── + +MODEL_NAME = "sdxl_watermark_detector" +MODEL_PORT = int(os.environ.get("MODEL_PORT", "9002")) +MAX_IMAGE_DIMENSION = 4096 + +_PRODUCTION = os.environ.get("PRODUCTION", "false").lower() == "true" + +# SDXL embeds 136 bits (17 bytes). The first 4 bytes are "SDV2". +_WATERMARK_BITS = 136 +_SDV2_PREFIX = b"SDV2" + +# ── Pydantic models ─────────────────────────────────────────────────────── + + +class ImageInput(BaseModel): + """Request body for /predict.""" + + image_data: str + threshold: float = 0.5 + + +# ── Detection helpers ────────────────────────────────────────────────────── + + +def _compute_byte_entropy(data: bytes) -> float: + """Compute Shannon entropy of a byte sequence. + + Args: + data: Raw bytes to analyze. + + Returns: + Entropy value in bits per byte (0.0 to 8.0). + Returns 0.0 for empty input. + """ + if not data: + return 0.0 + + length = len(data) + counts = Counter(data) + entropy = 0.0 + for count in counts.values(): + probability = count / length + if probability > 0: + entropy -= probability * math.log2(probability) + return entropy + + +def _hamming_distance(a: bytes, b: bytes) -> int: + """Count the number of differing bits between two byte sequences. + + Args: + a: First byte sequence. + b: Second byte sequence (same length as a). + + Returns: + Number of differing bits. + """ + dist = 0 + for x, y in zip(a, b): + dist += bin(x ^ y).count("1") + return dist + + +def _decode_watermark(bgr_array: np.ndarray) -> float: + """Attempt to decode SDXL watermark from a BGR numpy array. + + Tries two decoding methods (dwtDct, dwtDctSvd) and scores the + result based on pattern matching and entropy analysis. + + The SDV2 prefix check uses Hamming distance to tolerate minor + bit errors introduced by lossy compression or PNG round-trips. + The entropy threshold is strict (< 2.0) with a minimum unique + byte count to avoid false positives on JPEG DCT noise. + + Args: + bgr_array: OpenCV BGR image as numpy array. + + Returns: + Probability score: 0.70 for exact SDV2 match, 0.65 for + fuzzy SDV2 match, 0.5 for random/no signal. Scores are + capped low to prevent flipping borderline model verdicts + via PROVENANCE_WEIGHT boost. + """ + from imwatermark import WatermarkDecoder + + methods = ["dwtDct", "dwtDctSvd"] + + for method in methods: + try: + decoder = WatermarkDecoder("bytes", _WATERMARK_BITS) + watermark_bytes = decoder.decode(bgr_array, method) + + if watermark_bytes is None or len(watermark_bytes) == 0: + continue + + logger.debug( + "Decoded %d bytes via %s: %s", + len(watermark_bytes), + method, + watermark_bytes[:8].hex(), + ) + + # Exact SDV2 prefix match + if watermark_bytes[:4] == _SDV2_PREFIX: + logger.info( + "SDXL watermark detected via %s: SDV2 prefix matched", + method, + ) + return 0.70 + + # Fuzzy SDV2 match: tolerate up to 2 bit errors in the + # first 4 bytes (32 bits) to handle minor round-trip noise. + # Higher tolerance causes false positives on real images. + prefix_hamming = _hamming_distance(watermark_bytes[:4], _SDV2_PREFIX) + if prefix_hamming <= 2: + logger.info( + "SDXL watermark detected via %s: fuzzy SDV2 match " "(hamming=%d)", + method, + prefix_hamming, + ) + return 0.65 + + except Exception as exc: + logger.debug("Watermark decode failed with %s: %s", method, exc) + continue + + return 0.5 + + +def _detect_watermark(image_bytes: bytes) -> float: + """Run watermark detection on raw image bytes. + + Decodes image bytes with PIL, converts to RGB numpy array, + and attempts watermark decoding. + + Args: + image_bytes: Raw bytes of the image file. + + Returns: + Float probability in [0, 1]. 0.5 means neutral (no signal). + """ + try: + image = Image.open(io.BytesIO(image_bytes)) + if image.width > MAX_IMAGE_DIMENSION or image.height > MAX_IMAGE_DIMENSION: + logger.warning( + "Image too large: %dx%d, max %d", + image.width, + image.height, + MAX_IMAGE_DIMENSION, + ) + return 0.5 + image = image.convert("RGB") + # invisible-watermark library operates on RGB arrays (not BGR). + # The encoder/decoder use channel-independent DWT transforms. + rgb_array = np.array(image) + return _decode_watermark(rgb_array) + + except Exception as exc: + logger.warning("Watermark detection error: %s", exc) + return 0.5 + + +# ── FastAPI app ──────────────────────────────────────────────────────────── + +app = FastAPI( + title="SDXL Invisible Watermark Detector Service", + description=( + "Detects Stable Diffusion XL invisible watermarks embedded " + "in images using DWT-DCT steganographic decoding." + ), + version="1.0.0", + docs_url=None if _PRODUCTION else "/docs", + redoc_url=None if _PRODUCTION else "/redoc", + openapi_url=None if _PRODUCTION else "/openapi.json", +) + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + + +@app.get("/") +async def root(): + """Root endpoint with service information.""" + if _PRODUCTION: + return {"status": "ok"} + return { + "model_name": MODEL_NAME, + "description": ( + "SDXL invisible watermark detector using " "DWT-DCT steganographic decoding" + ), + "device": "cpu", + "model_loaded": True, + } + + +@app.get("/health") +async def health(): + """Health check endpoint.""" + if _PRODUCTION: + return {"status": "healthy"} + return { + "status": "healthy", + "model_name": MODEL_NAME, + "device": "cpu", + "model_loaded": True, + } + + +@app.post("/predict") +async def predict(payload: ImageInput) -> Dict[str, Any]: + """Detect SDXL invisible watermark in a base64-encoded image. + + Args: + payload: Base64-encoded image data and optional threshold. + + Returns: + Dict with model name, probability, prediction, class, and + inference time. + """ + if not payload.image_data: + raise HTTPException( + status_code=400, + detail="No image data provided.", + ) + + start = time.time() + + try: + image_bytes = base64.b64decode(payload.image_data) + except Exception as exc: + raise HTTPException( + status_code=400, + detail=f"Invalid base64 data: {exc}", + ) + + if not image_bytes: + raise HTTPException( + status_code=400, + detail="Empty image data after base64 decode.", + ) + + probability = _detect_watermark(image_bytes) + prediction = 1 if probability > payload.threshold else 0 + class_label = "fake" if prediction == 1 else "real" + inference_time = time.time() - start + + logger.info( + "Watermark check: %s (prob=%.4f, %.3fs)", + class_label, + probability, + inference_time, + ) + + return { + "model": MODEL_NAME, + "probability": float(probability), + "prediction": int(prediction), + "class": class_label, + "inference_time": float(inference_time), + } + + +if __name__ == "__main__": + logger.info( + "Starting SDXL watermark detector service on port %d", + MODEL_PORT, + ) + uvicorn.run(app, host="0.0.0.0", port=MODEL_PORT) diff --git a/provenance/invisible-watermark/requirements.txt b/provenance/invisible-watermark/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..483dfb1aae0cbaf80d96de515c73d9b429899427 --- /dev/null +++ b/provenance/invisible-watermark/requirements.txt @@ -0,0 +1,7 @@ +fastapi==0.115.12 +uvicorn==0.34.2 +pydantic==2.11.1 +invisible-watermark==0.2.0 +opencv-python-headless==4.11.0.86 +numpy>=1.24,<2.0 +Pillow>=10.0 diff --git a/provenance/invisible-watermark/test_app.py b/provenance/invisible-watermark/test_app.py new file mode 100644 index 0000000000000000000000000000000000000000..c8dac5c4619ef5b32bca7e91f8549dadc411579c --- /dev/null +++ b/provenance/invisible-watermark/test_app.py @@ -0,0 +1,351 @@ +"""Tests for the SDXL invisible watermark detection service.""" + +import base64 +import io + +import numpy as np +import pytest +from fastapi.testclient import TestClient +from PIL import Image + +from app import _compute_byte_entropy, app + +client = TestClient(app) + + +# ── Fixture helpers ──────────────────────────────────────────────────────── + + +def _make_clean_image( + width: int = 256, + height: int = 256, + mode: str = "RGB", +) -> bytes: + """Create a clean noisy image with no watermark. + + Uses random pixel noise so that watermark decoder output has high + entropy (appears random), ensuring no false-positive detection. + + Args: + width: Image width in pixels. + height: Image height in pixels. + mode: PIL image mode (RGB, RGBA, L). + + Returns: + Raw PNG bytes. + """ + rng = np.random.RandomState(42) + if mode == "L": + arr = rng.randint(0, 256, (height, width), dtype=np.uint8) + elif mode == "RGBA": + arr = rng.randint(0, 256, (height, width, 4), dtype=np.uint8) + else: + arr = rng.randint(0, 256, (height, width, 3), dtype=np.uint8) + img = Image.fromarray(arr, mode=mode) + buf = io.BytesIO() + img.save(buf, format="PNG") + return buf.getvalue() + + +def _make_watermarked_image() -> bytes: + """Create a 512x512 image with an SDXL-style watermark embedded. + + Uses WatermarkEncoder to embed the SDV2 pattern (b"SDV2" + 13 null + bytes = 17 bytes = 136 bits) via the dwtDct method. + + Returns: + Raw PNG bytes of the watermarked image. + """ + from imwatermark import WatermarkEncoder + + # Create a 512x512 test image (larger for reliable watermark encoding) + img = Image.new("RGB", (512, 512), color=(100, 150, 200)) + rgb_array = np.array(img) + + # Convert RGB to BGR for OpenCV/imwatermark + import cv2 + + bgr_array = cv2.cvtColor(rgb_array, cv2.COLOR_RGB2BGR) + + # Encode SDV2 watermark: 4 bytes "SDV2" + 13 null bytes = 17 bytes + watermark_payload = b"SDV2" + b"\x00" * 13 + encoder = WatermarkEncoder() + encoder.set_watermark("bytes", watermark_payload) + watermarked_bgr = encoder.encode(bgr_array, "dwtDct") + + # Convert back to RGB PIL image + watermarked_rgb = cv2.cvtColor(watermarked_bgr, cv2.COLOR_BGR2RGB) + watermarked_img = Image.fromarray(watermarked_rgb) + + buf = io.BytesIO() + watermarked_img.save(buf, format="PNG") + return buf.getvalue() + + +def _b64(raw: bytes) -> str: + """Encode raw bytes as base64 string.""" + return base64.b64encode(raw).decode("utf-8") + + +# ── Health endpoint ──────────────────────────────────────────────────────── + + +class TestHealthEndpoint: + """Tests for GET /health.""" + + def test_health_returns_healthy(self): + """Health endpoint returns healthy status.""" + resp = client.get("/health") + assert resp.status_code == 200 + data = resp.json() + assert data["status"] == "healthy" + + def test_health_returns_model_name(self): + """Health endpoint includes correct model_name.""" + resp = client.get("/health") + data = resp.json() + assert data["model_name"] == "sdxl_watermark_detector" + + def test_health_returns_device(self): + """Health endpoint includes device=cpu.""" + resp = client.get("/health") + data = resp.json() + assert data["device"] == "cpu" + + def test_health_returns_model_loaded(self): + """Health endpoint includes model_loaded=True.""" + resp = client.get("/health") + data = resp.json() + assert data["model_loaded"] is True + + +# ── Root endpoint ────────────────────────────────────────────────────────── + + +class TestRootEndpoint: + """Tests for GET /.""" + + def test_root_returns_service_info(self): + """Root endpoint returns model info.""" + resp = client.get("/") + assert resp.status_code == 200 + data = resp.json() + assert data["model_name"] == "sdxl_watermark_detector" + assert data["device"] == "cpu" + assert data["model_loaded"] is True + + def test_root_returns_description(self): + """Root endpoint includes a description.""" + resp = client.get("/") + data = resp.json() + assert "description" in data + assert len(data["description"]) > 0 + + +# ── Predict endpoint: clean image ───────────────────────────────────────── + + +class TestPredictCleanImage: + """Tests for POST /predict with a clean (no watermark) image.""" + + def test_clean_image_returns_neutral(self): + """Clean image with no watermark returns probability=0.5.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_clean_image())}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["probability"] == 0.5 + assert data["class"] == "real" + + +# ── Predict endpoint: watermarked image ─────────────────────────────────── + + +class TestPredictWatermarkedImage: + """Tests for POST /predict with an SDXL-watermarked image.""" + + def test_sdxl_watermark_detected(self): + """Image with SDXL watermark returns high probability.""" + watermarked = _make_watermarked_image() + resp = client.post( + "/predict", + json={"image_data": _b64(watermarked)}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["probability"] > 0.5 + assert data["class"] == "fake" + + +# ── Predict endpoint: image format handling ─────────────────────────────── + + +class TestPredictImageFormats: + """Tests for various image formats and sizes.""" + + def test_grayscale_image_handled(self): + """Grayscale image is converted and processed without error.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_clean_image(mode="L"))}, + ) + assert resp.status_code == 200 + data = resp.json() + assert 0.0 <= data["probability"] <= 1.0 + + def test_rgba_image_handled(self): + """RGBA image is converted and processed without error.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_clean_image(mode="RGBA"))}, + ) + assert resp.status_code == 200 + data = resp.json() + assert 0.0 <= data["probability"] <= 1.0 + + def test_small_image_handled(self): + """Small 32x32 image is processed without error.""" + resp = client.post( + "/predict", + json={ + "image_data": _b64(_make_clean_image(width=32, height=32)), + }, + ) + assert resp.status_code == 200 + data = resp.json() + assert 0.0 <= data["probability"] <= 1.0 + + +# ── Predict endpoint: error handling ────────────────────────────────────── + + +class TestPredictErrorHandling: + """Tests for error cases.""" + + def test_invalid_base64_returns_error(self): + """Invalid base64 string returns 400.""" + resp = client.post( + "/predict", + json={"image_data": "!!!not-valid-base64!!!"}, + ) + assert resp.status_code == 400 + + def test_non_image_bytes_returns_neutral(self): + """Non-image binary data returns probability=0.5, not crash.""" + random_bytes = base64.b64encode(b"\x00\x01\x02\x03" * 20).decode() + resp = client.post( + "/predict", + json={"image_data": random_bytes}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["probability"] == 0.5 + + +# ── Predict endpoint: response schema ───────────────────────────────────── + + +class TestPredictResponseSchema: + """Tests for the standard response format.""" + + def test_response_has_all_required_fields(self): + """Response contains model, probability, prediction, class, + inference_time.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_clean_image())}, + ) + data = resp.json() + required = { + "model", + "probability", + "prediction", + "class", + "inference_time", + } + assert required.issubset(data.keys()) + + def test_model_name_in_response(self): + """Response model field matches sdxl_watermark_detector.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_clean_image())}, + ) + data = resp.json() + assert data["model"] == "sdxl_watermark_detector" + + def test_probability_is_float_in_range(self): + """Probability is a float in [0, 1].""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_clean_image())}, + ) + data = resp.json() + assert isinstance(data["probability"], float) + assert 0.0 <= data["probability"] <= 1.0 + + def test_inference_time_is_positive_float(self): + """Inference time is a positive float.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_clean_image())}, + ) + data = resp.json() + assert isinstance(data["inference_time"], float) + assert data["inference_time"] > 0 + + def test_threshold_controls_prediction(self): + """Prediction uses > threshold (0.5 is neutral, not fake).""" + resp = client.post( + "/predict", + json={ + "image_data": _b64(_make_clean_image()), + "threshold": 0.5, + }, + ) + data = resp.json() + # probability=0.5, threshold=0.5 -> 0.5 > 0.5 is False -> real + assert data["prediction"] == 0 + assert data["class"] == "real" + + +# ── Entropy helper ───────────────────────────────────────────────────────── + + +class TestComputeByteEntropy: + """Tests for the _compute_byte_entropy helper function.""" + + def test_empty_bytes_returns_zero(self): + """Empty input returns 0.0 entropy.""" + assert _compute_byte_entropy(b"") == 0.0 + + def test_single_byte_returns_zero(self): + """All identical bytes have zero entropy.""" + assert _compute_byte_entropy(b"\x00" * 100) == 0.0 + + def test_two_equally_distributed_returns_one(self): + """Two equally distributed byte values have entropy=1.0.""" + data = b"\x00\x01" * 50 + entropy = _compute_byte_entropy(data) + assert abs(entropy - 1.0) < 0.01 + + def test_high_entropy_random_bytes(self): + """Random-like bytes have high entropy (close to 8.0).""" + # 256 distinct byte values each appearing once + data = bytes(range(256)) + entropy = _compute_byte_entropy(data) + assert entropy == 8.0 + + def test_low_entropy_structured_bytes(self): + """Structured bytes have low entropy.""" + # "SDV2" repeated = limited unique bytes + data = b"SDV2" * 10 + entropy = _compute_byte_entropy(data) + assert entropy < 3.0 + + def test_entropy_is_non_negative(self): + """Entropy is always non-negative.""" + assert _compute_byte_entropy(b"\xff") >= 0.0 + assert _compute_byte_entropy(b"hello world") >= 0.0 diff --git a/provenance/videoseal/Dockerfile b/provenance/videoseal/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..f6a3d26891b7fc15f2f8a1fd3342e2a01f219990 --- /dev/null +++ b/provenance/videoseal/Dockerfile @@ -0,0 +1,45 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +# System deps: OpenCV headless needs libgl, ffmpeg for video decoding, +# build-essential for pycocotools compilation +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + python3 python3-pip python3-dev \ + libgl1 libglib2.0-0 ffmpeg build-essential && \ + rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +WORKDIR /app + +# Install PyTorch with CUDA 12.1 +RUN pip install --no-cache-dir \ + torch==2.5.1 torchvision==0.20.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +COPY requirements.txt . +# Install videoseal without deps to avoid decord (no ARM64 wheels). +RUN pip install --no-cache-dir -r requirements.txt && \ + pip install --no-cache-dir --no-deps videoseal==1.0.1 && \ + apt-get purge -y --auto-remove build-essential && \ + rm -rf /var/lib/apt/lists/* + +COPY app.py . +COPY configs/ configs/ + +ENV MODEL_PORT=9004 +ENV PRELOAD_MODEL=false +ENV MODEL_TIMEOUT=600 +EXPOSE 9004 + +# Pre-create cache dirs that videoseal needs (before dropping root) +RUN mkdir -p /app/ckpts /app/.cache +COPY ckpts/ /app/ckpts/ +RUN adduser --disabled-password --gecos "" appuser && \ + chown -R appuser:appuser /app/ckpts /app/.cache +USER appuser + +CMD ["python", "app.py"] diff --git a/provenance/videoseal/app.py b/provenance/videoseal/app.py new file mode 100644 index 0000000000000000000000000000000000000000..d7d560d850cbe4f2a584e5ade88b2daf1e35fc58 --- /dev/null +++ b/provenance/videoseal/app.py @@ -0,0 +1,500 @@ +"""VideoSeal/PixelSeal watermark detection service. + +Detects Meta VideoSeal watermarks embedded in images and videos using the +videoseal_base model. Returns a probability score indicating likelihood +that the media contains a watermark (indicating AI generation). + +Unlike detection models that analyze visual artifacts, this service +inspects only for embedded VideoSeal watermarks. +""" + +import base64 +import gc +import io +import logging +import os +import platform +import sys +import tempfile +import threading +import time +from typing import Any, Dict, Optional + +import cv2 +import numpy as np +import torch +import torchvision.transforms.functional as TF +import uvicorn +from fastapi import FastAPI, HTTPException +from fastapi.middleware.cors import CORSMiddleware +from PIL import Image +from pydantic import BaseModel + +# ── Logging ──────────────────────────────────────────────────────────────── + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + handlers=[logging.StreamHandler(sys.stdout)], +) +logger = logging.getLogger(__name__) + +# ── Config ───────────────────────────────────────────────────────────────── + +MODEL_NAME = "videoseal_detector" +MODEL_PORT = int(os.environ.get("MODEL_PORT", "9004")) +MAX_IMAGE_DIMENSION = 4096 + +_PRODUCTION = os.environ.get("PRODUCTION", "false").lower() == "true" + +DETECTION_THRESHOLD = 0.6 # Bit confidence ratio: clean ~0.01, watermarked ~0.76 +MAX_FRAMES = 16 + + +def _get_device(): + """Select optimal device: CUDA > MPS > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if override: + pass # Invalid override, fall through to auto-detect + if torch.cuda.is_available(): + return torch.device("cuda") + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + return torch.device("cpu") + + +DEVICE = _get_device() +if DEVICE.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + +if DEVICE.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) +else: + logger.warning( + "Device: %s (no CUDA available -- check nvidia-container-toolkit)", + DEVICE, + ) + +PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true" +MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", "600")) + +# ── Globals ──────────────────────────────────────────────────────────────── + +model = None +model_lock = threading.Lock() +last_used_time = 0 + + +# ── Pydantic models ─────────────────────────────────────────────────────── + + +class MediaInput(BaseModel): + """Request body for /predict.""" + + image_data: Optional[str] = None + video_data: Optional[str] = None + threshold: float = 0.5 + + +# ── Confidence mapping ──────────────────────────────────────────────────── + + +def _map_confidence_to_probability(confidence: float) -> float: + """Map VideoSeal detection confidence to a probability score. + + Confidence values below the detection threshold are treated as + noise and mapped to 0.5 (neutral). Values at or above the + threshold are linearly scaled into [0.5, 1.0]. + + Args: + confidence: Raw detection confidence from VideoSeal in [0, 1]. + + Returns: + Probability in [0.5, 1.0]. 0.5 means neutral (no watermark). + """ + if confidence < DETECTION_THRESHOLD: + return 0.5 + return 0.5 + confidence * 0.5 + + +# ── Model loading ───────────────────────────────────────────────────────── + + +def load_model_internal(): + """Load VideoSeal detector model onto the selected device.""" + global model, last_used_time + + with model_lock: + if model is not None: + last_used_time = time.time() + return + + logger.info("Loading VideoSeal detector model...") + try: + # Use absolute path to model card (videoseal uses relative + # Path("videoseal/cards") which fails when CWD != package root) + import pathlib + + import videoseal + + cards_dir = pathlib.Path(videoseal.__file__).parent / "cards" + card_path = cards_dir / "videoseal_1.0.yaml" + loaded = videoseal.load(card_path) + loaded = loaded.to(DEVICE) + loaded.eval() + model = loaded + last_used_time = time.time() + logger.info("VideoSeal detector ready on %s.", str(DEVICE)) + + except Exception as exc: + logger.exception("Failed to load VideoSeal detector: %s", exc) + model = None + raise + finally: + gc.collect() + + +def ensure_model_loaded(): + """Load model on first request (lazy loading).""" + global last_used_time + if model is None: + load_model_internal() + else: + last_used_time = time.time() + + +def unload_model_if_idle(): + """Evict model from RAM after MODEL_TIMEOUT seconds of inactivity.""" + global model + if model is None or PRELOAD_MODEL: + return + with model_lock: + if model is not None and (time.time() - last_used_time > MODEL_TIMEOUT): + logger.info("Unloading idle VideoSeal detector to free RAM.") + del model + model = None + gc.collect() + + +# ── Detection helpers ───────────────────────────────────────────────────── + + +def _detect_image(image_bytes: bytes) -> float: + """Detect VideoSeal watermark in image bytes. + + Decodes the image, converts to a [1, 3, H, W] tensor in [0, 1], + and runs the VideoSeal detector. Returns the mean p-value as the + raw confidence score. + + Args: + image_bytes: Raw bytes of a JPEG/PNG/etc. image. + + Returns: + Confidence float in [0, 1]. + """ + try: + pil_image = Image.open(io.BytesIO(image_bytes)) + if ( + pil_image.width > MAX_IMAGE_DIMENSION + or pil_image.height > MAX_IMAGE_DIMENSION + ): + logger.warning( + "Image too large: %dx%d, max %d", + pil_image.width, + pil_image.height, + MAX_IMAGE_DIMENSION, + ) + return 0.0 + pil_image = pil_image.convert("RGB") + except Exception as exc: + logger.warning("Failed to decode image: %s", exc) + return 0.0 + + tensor = TF.to_tensor(pil_image).unsqueeze(0).to(DEVICE) # [1, 3, H, W] + + with torch.no_grad(): + outputs = model.detect(tensor) + + # VideoSeal outputs 'preds': [1, nbits+1]. The first value is a + # detection logit (unreliable). The real signal is the message bits: + # watermarked images have many confident bits (far from 0.5), + # clean images have bits clustered near 0.5. + preds = outputs["preds"] + bit_confs = torch.sigmoid(preds[0, 1:]) # message bit confidences + confident_bits = ((bit_confs > 0.8) | (bit_confs < 0.2)).float().mean() + confidence = float(confident_bits) + return confidence + + +def _detect_video(video_bytes: bytes) -> float: + """Detect VideoSeal watermark in video bytes. + + Writes bytes to a temporary file, extracts up to MAX_FRAMES + evenly sampled frames via OpenCV, runs detection on each frame, + and returns the maximum confidence across all frames. + + Args: + video_bytes: Raw bytes of a video file. + + Returns: + Maximum confidence float across frames in [0, 1]. + """ + tmp_path = None + try: + with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp: + tmp.write(video_bytes) + tmp_path = tmp.name + + cap = cv2.VideoCapture(tmp_path) + if not cap.isOpened(): + logger.warning("Failed to open video file.") + return 0.0 + + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + if total_frames <= 0: + cap.release() + return 0.0 + + # Sample frame indices evenly + if total_frames <= MAX_FRAMES: + indices = list(range(total_frames)) + else: + indices = np.linspace(0, total_frames - 1, MAX_FRAMES, dtype=int).tolist() + + max_confidence = 0.0 + + for idx in indices: + cap.set(cv2.CAP_PROP_POS_FRAMES, idx) + ret, frame = cap.read() + if not ret or frame is None: + continue + + # OpenCV BGR -> RGB -> PIL -> tensor + rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + pil_frame = Image.fromarray(rgb_frame) + tensor = TF.to_tensor(pil_frame).unsqueeze(0).to(DEVICE) + + with torch.no_grad(): + outputs = model.detect(tensor) + + preds = outputs["preds"] + bit_confs = torch.sigmoid(preds[0, 1:]) + confidence = float(((bit_confs > 0.8) | (bit_confs < 0.2)).float().mean()) + if confidence > max_confidence: + max_confidence = confidence + + cap.release() + return max_confidence + + except Exception as exc: + logger.warning("Video detection error: %s", exc) + return 0.0 + + finally: + if tmp_path and os.path.exists(tmp_path): + try: + os.unlink(tmp_path) + except OSError: + pass + + +# ── FastAPI app ──────────────────────────────────────────────────────────── + +app = FastAPI( + title="VideoSeal Watermark Detection Service", + description=( + "Detects Meta VideoSeal/PixelSeal watermarks embedded in images " + "and videos using the videoseal_base model." + ), + version="1.0.0", + docs_url=None if _PRODUCTION else "/docs", + redoc_url=None if _PRODUCTION else "/redoc", + openapi_url=None if _PRODUCTION else "/openapi.json", +) + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + + +@app.get("/") +async def root(): + """Root endpoint with service information.""" + if _PRODUCTION: + return {"status": "ok"} + return { + "model_name": MODEL_NAME, + "description": ( + "VideoSeal/PixelSeal watermark detector for " + "AI-generated images and videos" + ), + "device": str(DEVICE), + "model_loaded": model is not None, + } + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and DEVICE.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +async def health(): + """Health check endpoint.""" + if _PRODUCTION: + return {"status": "healthy"} + return { + "status": "healthy", + "model_name": MODEL_NAME, + "device": str(DEVICE), + "model_loaded": model is not None, + **_gpu_health_info(), + } + + +@app.post("/predict") +async def predict(payload: MediaInput) -> Dict[str, Any]: + """Detect VideoSeal watermark in base64-encoded image or video. + + Accepts image_data or video_data as base64-encoded file bytes. + Determines modality from which key is non-None. Returns a + standard detection response with probability based on watermark + evidence. + + Args: + payload: Base64-encoded media and optional threshold. + + Returns: + Dict with model name, probability, prediction, class, and + inference time. + """ + try: + ensure_model_loaded() + if model is None: + raise HTTPException(status_code=503, detail="Model not loaded.") + + start = time.time() + + # Determine modality from input keys + if payload.image_data is not None: + raw_b64 = payload.image_data + modality = "image" + elif payload.video_data is not None: + raw_b64 = payload.video_data + modality = "video" + else: + raise HTTPException( + status_code=400, + detail=( + "No media data provided. Include one of: " "image_data, video_data." + ), + ) + + # Decode base64 payload + try: + media_bytes = base64.b64decode(raw_b64) + except Exception as exc: + raise HTTPException( + status_code=400, + detail=f"Invalid base64 data: {exc}", + ) + + if not media_bytes: + raise HTTPException( + status_code=400, + detail="Empty media data after base64 decode.", + ) + + # Run detection based on modality + if modality == "image": + confidence = _detect_image(media_bytes) + else: + confidence = _detect_video(media_bytes) + + probability = _map_confidence_to_probability(confidence) + prediction = 1 if probability > payload.threshold else 0 + class_label = "fake" if prediction == 1 else "real" + inference_time = time.time() - start + + logger.info( + "VideoSeal check (%s): %s (conf=%.4f, prob=%.4f, %.3fs)", + modality, + class_label, + confidence, + probability, + inference_time, + ) + + if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: + threading.Timer(MODEL_TIMEOUT + 5.0, unload_model_if_idle).start() + + return { + "model": MODEL_NAME, + "probability": float(probability), + "prediction": int(prediction), + "class": class_label, + "inference_time": float(inference_time), + } + + except HTTPException: + raise + except Exception as exc: + logger.exception("Prediction error: %s", exc) + raise HTTPException(status_code=500, detail=str(exc)) + + +@app.on_event("startup") +async def startup_event(): + """Startup: preload model if configured, else lazy-load on first request.""" + if PRELOAD_MODEL: + logger.info("Preloading VideoSeal detector at startup.") + try: + load_model_internal() + except Exception as exc: + logger.error("Preload failed: %s", exc) + else: + logger.info("VideoSeal service ready — model loads on first request.") + + if not PRELOAD_MODEL and MODEL_TIMEOUT > 0: + + def _periodic_check(): + unload_model_if_idle() + threading.Timer(MODEL_TIMEOUT / 2.0, _periodic_check).start() + + threading.Timer(MODEL_TIMEOUT / 2.0, _periodic_check).start() + + +if __name__ == "__main__": + logger.info("Starting VideoSeal detector service on port %d", MODEL_PORT) + uvicorn.run(app, host="0.0.0.0", port=MODEL_PORT) diff --git a/provenance/videoseal/configs/attenuation.yaml b/provenance/videoseal/configs/attenuation.yaml new file mode 100644 index 0000000000000000000000000000000000000000..76d25e2c609b1dfc577abfd63990aa5e8efc9f0f --- /dev/null +++ b/provenance/videoseal/configs/attenuation.yaml @@ -0,0 +1,15 @@ +jnd_1_1: + in_channels: 1 + out_channels: 1 + +jnd_3_3: + in_channels: 3 + out_channels: 3 + +jnd_1_3: + in_channels: 1 + out_channels: 3 + +jnd_3_1: + in_channels: 3 + out_channels: 1 diff --git a/provenance/videoseal/requirements.txt b/provenance/videoseal/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..58b923950f1f26e50a5cdacc8ecff60d803551e7 --- /dev/null +++ b/provenance/videoseal/requirements.txt @@ -0,0 +1,17 @@ +fastapi==0.115.12 +uvicorn==0.34.2 +pydantic==2.11.1 +torch>=2.0.0 +torchvision>=0.15.0 +numpy>=1.24,<2.0 +Pillow>=10.0 +opencv-python-headless==4.11.0.86 +# videoseal runtime deps (installed with --no-deps to skip decord) +requests>=2.28.0 +omegaconf>=2.3.0 +timm==0.9.16 +einops>=0.6.0 +av>=10.0.0 +pycocotools>=2.0.0 +tqdm>=4.64.0 +scipy>=1.10.0 diff --git a/provenance/videoseal/test_app.py b/provenance/videoseal/test_app.py new file mode 100644 index 0000000000000000000000000000000000000000..f37f30600eb42dff475fabb90527c526235c44e5 --- /dev/null +++ b/provenance/videoseal/test_app.py @@ -0,0 +1,371 @@ +"""Tests for the VideoSeal watermark detection service.""" + +import base64 +import io +from unittest.mock import MagicMock, patch + +import pytest +import torch +from fastapi.testclient import TestClient +from PIL import Image + +from app import _map_confidence_to_probability, app + +client = TestClient(app) + + +# ── Fixture helpers ──────────────────────────────────────────────────────── + + +def _make_image(width: int = 256, height: int = 256) -> bytes: + """Create a valid PNG image of the given size. + + Generates a solid red image. + + Args: + width: Image width in pixels. + height: Image height in pixels. + + Returns: + Raw PNG file bytes. + """ + img = Image.new("RGB", (width, height), color=(255, 0, 0)) + buf = io.BytesIO() + img.save(buf, format="PNG") + return buf.getvalue() + + +def _make_mp4() -> bytes: + """Create a minimal MP4 file with a valid ftyp box header. + + Returns: + Raw bytes of a minimal MP4 file. + """ + return b"\x00\x00\x00\x14ftypmp42\x00\x00\x00\x00mp42" + + +def _b64(raw: bytes) -> str: + """Encode raw bytes as base64 string.""" + return base64.b64encode(raw).decode("utf-8") + + +# ── Mock detector ───────────────────────────────────────────────────────── + +# All /predict tests mock the VideoSeal model so tests run without GPU +# or the videoseal package installed. The mock returns a dict with pvalue. + + +def _mock_detect(tensor): + """Mock detector that returns low confidence (no watermark). + + Returns a dict with pvalue tensor of 0.1 (below threshold). + """ + return {"pvalue": torch.tensor(0.1)} + + +def _patch_model(): + """Create a mock model with detect method.""" + mock = MagicMock() + mock.detect = MagicMock(side_effect=_mock_detect) + return mock + + +# ── Health endpoint ──────────────────────────────────────────────────────── + + +class TestHealthEndpoint: + """Tests for GET /health.""" + + def test_health_returns_healthy(self): + """Health endpoint returns healthy status.""" + resp = client.get("/health") + assert resp.status_code == 200 + data = resp.json() + assert data["status"] == "healthy" + + def test_health_returns_model_name(self): + """Health endpoint includes correct model_name.""" + resp = client.get("/health") + data = resp.json() + assert data["model_name"] == "videoseal_detector" + + def test_health_returns_device(self): + """Health endpoint includes device field.""" + resp = client.get("/health") + data = resp.json() + assert "device" in data + + def test_health_returns_model_loaded(self): + """Health endpoint includes model_loaded field.""" + resp = client.get("/health") + data = resp.json() + assert "model_loaded" in data + + +# ── Root endpoint ────────────────────────────────────────────────────────── + + +class TestRootEndpoint: + """Tests for GET /.""" + + def test_root_returns_service_info(self): + """Root endpoint returns model info.""" + resp = client.get("/") + assert resp.status_code == 200 + data = resp.json() + assert data["model_name"] == "videoseal_detector" + + def test_root_returns_description(self): + """Root endpoint includes a description.""" + resp = client.get("/") + data = resp.json() + assert "description" in data + assert len(data["description"]) > 0 + + +# ── Predict endpoint: clean image returns neutral ───────────────────────── + + +class TestPredictCleanImage: + """Tests for POST /predict with a clean (no watermark) image.""" + + @patch("app.model", new_callable=_patch_model) + @patch("app.ensure_model_loaded") + def test_clean_image_returns_neutral(self, mock_load, mock_mdl): + """Clean image with no watermark returns probability=0.5.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_image())}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["probability"] == 0.5 + assert data["class"] == "real" + + +# ── Predict endpoint: different image sizes ─────────────────────────────── + + +class TestPredictImageSizes: + """Tests for POST /predict with various image dimensions.""" + + @patch("app.model", new_callable=_patch_model) + @patch("app.ensure_model_loaded") + def test_128x128_image(self, mock_load, mock_mdl): + """128x128 image is processed without error.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_image(128, 128))}, + ) + assert resp.status_code == 200 + data = resp.json() + assert 0.0 <= data["probability"] <= 1.0 + + @patch("app.model", new_callable=_patch_model) + @patch("app.ensure_model_loaded") + def test_512x512_image(self, mock_load, mock_mdl): + """512x512 image is processed without error.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_image(512, 512))}, + ) + assert resp.status_code == 200 + data = resp.json() + assert 0.0 <= data["probability"] <= 1.0 + + @patch("app.model", new_callable=_patch_model) + @patch("app.ensure_model_loaded") + def test_1920x1080_image(self, mock_load, mock_mdl): + """1920x1080 image is processed without error.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_image(1920, 1080))}, + ) + assert resp.status_code == 200 + data = resp.json() + assert 0.0 <= data["probability"] <= 1.0 + + +# ── Predict endpoint: accepts both modalities ───────────────────────────── + + +class TestPredictModalities: + """Tests that /predict accepts both image_data and video_data keys.""" + + @patch("app.model", new_callable=_patch_model) + @patch("app.ensure_model_loaded") + def test_accepts_image_data(self, mock_load, mock_mdl): + """Endpoint accepts image_data key.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_image())}, + ) + assert resp.status_code == 200 + + @patch("app.model", new_callable=_patch_model) + @patch("app.ensure_model_loaded") + def test_accepts_video_data(self, mock_load, mock_mdl): + """Endpoint accepts video_data key.""" + resp = client.post( + "/predict", + json={"video_data": _b64(_make_mp4())}, + ) + assert resp.status_code == 200 + data = resp.json() + # Minimal MP4 cannot be decoded by OpenCV, so confidence=0.0 + # which maps to probability=0.5 (neutral) + assert data["probability"] == 0.5 + + +# ── Predict endpoint: error handling ─────────────────────────────────────── + + +class TestPredictErrorHandling: + """Tests for error cases.""" + + @patch("app.model", new_callable=_patch_model) + @patch("app.ensure_model_loaded") + def test_invalid_base64_returns_error(self, mock_load, mock_mdl): + """Invalid base64 string returns 400.""" + resp = client.post( + "/predict", + json={"image_data": "!!!not-valid-base64!!!"}, + ) + assert resp.status_code == 400 + + @patch("app.model", new_callable=_patch_model) + @patch("app.ensure_model_loaded") + def test_non_media_bytes_returns_neutral(self, mock_load, mock_mdl): + """Non-media binary data returns probability=0.5, not crash.""" + random_bytes = base64.b64encode(b"\x00\x01\x02\x03" * 20).decode() + resp = client.post( + "/predict", + json={"image_data": random_bytes}, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["probability"] == 0.5 + + @patch("app.model", new_callable=_patch_model) + @patch("app.ensure_model_loaded") + def test_no_media_returns_error(self, mock_load, mock_mdl): + """Request with neither image_data nor video_data returns 400.""" + resp = client.post( + "/predict", + json={"threshold": 0.5}, + ) + assert resp.status_code == 400 + + +# ── Predict endpoint: response format ────────────────────────────────────── + + +class TestPredictResponseSchema: + """Tests for the standard response format.""" + + @patch("app.model", new_callable=_patch_model) + @patch("app.ensure_model_loaded") + def test_response_has_all_required_fields(self, mock_load, mock_mdl): + """Response contains model, probability, prediction, class, + inference_time.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_image())}, + ) + data = resp.json() + required = { + "model", + "probability", + "prediction", + "class", + "inference_time", + } + assert required.issubset(data.keys()) + + @patch("app.model", new_callable=_patch_model) + @patch("app.ensure_model_loaded") + def test_model_name_in_response(self, mock_load, mock_mdl): + """Response model field matches videoseal_detector.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_image())}, + ) + data = resp.json() + assert data["model"] == "videoseal_detector" + + @patch("app.model", new_callable=_patch_model) + @patch("app.ensure_model_loaded") + def test_probability_is_float_in_range(self, mock_load, mock_mdl): + """Probability is a float in [0, 1].""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_image())}, + ) + data = resp.json() + assert isinstance(data["probability"], float) + assert 0.0 <= data["probability"] <= 1.0 + + @patch("app.model", new_callable=_patch_model) + @patch("app.ensure_model_loaded") + def test_inference_time_is_positive_float(self, mock_load, mock_mdl): + """Inference time is a positive float.""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_image())}, + ) + data = resp.json() + assert isinstance(data["inference_time"], float) + assert data["inference_time"] > 0 + + @patch("app.model", new_callable=_patch_model) + @patch("app.ensure_model_loaded") + def test_threshold_controls_prediction(self, mock_load, mock_mdl): + """Prediction uses > threshold (0.5 is neutral, not fake).""" + resp = client.post( + "/predict", + json={"image_data": _b64(_make_image()), "threshold": 0.5}, + ) + data = resp.json() + # probability=0.5 (from mock), threshold=0.5 -> not > -> real + assert data["prediction"] == 0 + assert data["class"] == "real" + + +# ── _map_confidence_to_probability helper ────────────────────────────────── + + +class TestMapConfidenceToProbability: + """Tests for the _map_confidence_to_probability helper function.""" + + def test_high_confidence_returns_high(self): + """High confidence (0.9) maps to high probability.""" + result = _map_confidence_to_probability(0.9) + assert result == 0.5 + 0.9 * 0.5 # 0.95 + assert result > 0.9 + + def test_medium_confidence_returns_medium(self): + """Medium confidence (0.5) maps to medium-high probability.""" + result = _map_confidence_to_probability(0.5) + assert result == 0.5 + 0.5 * 0.5 # 0.75 + + def test_below_threshold_returns_neutral(self): + """Confidence below 0.3 threshold returns 0.5 (neutral).""" + assert _map_confidence_to_probability(0.1) == 0.5 + assert _map_confidence_to_probability(0.0) == 0.5 + assert _map_confidence_to_probability(0.29) == 0.5 + + def test_boundary_at_threshold(self): + """Confidence at exactly 0.3 maps above neutral.""" + result = _map_confidence_to_probability(0.3) + assert result == 0.5 + 0.3 * 0.5 # 0.65 + assert result > 0.5 + + def test_just_below_threshold(self): + """Confidence just below 0.3 returns neutral.""" + result = _map_confidence_to_probability(0.299) + assert result == 0.5 + + def test_confidence_one_returns_max(self): + """Confidence of 1.0 returns maximum probability (1.0).""" + result = _map_confidence_to_probability(1.0) + assert result == 1.0 diff --git a/video/dfd-fcg/Dockerfile b/video/dfd-fcg/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..5855e6ff9b6b1880be7836b5f58dd52341b19e10 --- /dev/null +++ b/video/dfd-fcg/Dockerfile @@ -0,0 +1,59 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +WORKDIR /app + +# Install Python 3.10 and system dependencies for OpenCV +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3 python3-pip \ + libgl1 \ + libglib2.0-0 \ + libsm6 \ + libxext6 \ + libxrender-dev \ + && rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +# Install PyTorch with CUDA 12.1 +RUN pip install --no-cache-dir \ + torch==2.5.1 torchvision==0.20.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +# Copy requirements and install +# facenet-pytorch needs --no-deps due to torch<2.3 pin +COPY requirements.txt . +RUN pip install --no-cache-dir --no-deps facenet-pytorch && \ + pip install --no-cache-dir -r requirements.txt + +# Create logs and weights directories +RUN mkdir -p logs weights model_code + +# Copy model code (read-only reference, never modified) +COPY model_code/ /app/model_code/ + +# Copy weights +COPY weights/ /app/weights/ + +# Copy application code +COPY app.py . + +# Environment variables +ENV MODEL_PORT=7003 +ENV PRELOAD_MODEL=false +ENV MODEL_TIMEOUT=1800 +ENV WEIGHTS_PATH=/app/weights/dfd_fcg_checkpoint.pth +ENV MODEL_CODE_DIR=/app/model_code + +# Expose port +EXPOSE 7003 + +# Drop root privileges +RUN adduser --disabled-password --gecos '' appuser && \ + chown -R appuser:appuser /app/logs /app/weights +USER appuser + +# Run the service +CMD ["python", "app.py"] diff --git a/video/dfd-fcg/app.py b/video/dfd-fcg/app.py new file mode 100644 index 0000000000000000000000000000000000000000..397a4f30d24c77e674ce40c0b2bb0d221af7cd48 --- /dev/null +++ b/video/dfd-fcg/app.py @@ -0,0 +1,607 @@ +"""DFD-FCG (Deepfake Detection via Facial Component Guidance) service. + +Wraps the DFD-FCG (CVPR 2025) video deepfake detection model with a +FastAPI endpoint. Uses CLIP ViT-L/14 with a Synoptic Video Learner +and Facial Component Guidance for robust face forgery detection. + +The model analyses temporal and spatial inconsistencies across facial +components (lips, skin, eyes, nose) using learned synoptic attention +over multi-frame CLIP embeddings. + +Reference: "Deepfake Detection via Facial Component Guidance", +CVPR 2025. +""" + +import base64 +import gc +import logging +import math +import os +import platform +import sys +import tempfile +import threading +import time +from typing import Any, Dict, List, Optional, Tuple + +import cv2 +import numpy as np +import torch +import torch.nn.functional as F +import uvicorn +from fastapi import FastAPI, HTTPException +from PIL import Image +from pydantic import BaseModel, ConfigDict, Field + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +MODEL_PORT = int(os.environ.get("MODEL_PORT", 7003)) +PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true" +MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", 1800)) +WEIGHTS_PATH = os.environ.get("WEIGHTS_PATH", "/app/weights/dfd_fcg_checkpoint.pth") +# Path to serialized face semantic features used by FFG module +FACE_FEATURES_PATH = os.environ.get( + "FACE_FEATURES_PATH", + "/app/model_code/misc/L14_real_semantic_patches_v4_2000.pickle", +) + +# DFD-FCG uses CLIP ViT-L/14 at 224x224 (standard CLIP resolution) +IMAGE_SIZE = 224 +# Number of frames to sample per clip (matches config num_frames=10) +NUM_FRAMES = 10 +# Frame sampling stride in seconds (matches demo.py stride=0.333) +FRAME_STRIDE = 0.333 +# Face crop margin factor for MTCNN bounding boxes +MARGIN_FACTOR = 0.5 + + +def _get_device() -> torch.device: + """Select optimal device: CUDA (NVIDIA) > MPS (Apple) > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + return torch.device("cpu") + + +def _build_model_from_config(): + """Build FFGSynoVideoLearner from config parameters. + + Instantiates the model directly using the same parameters + defined in configs/clip/L14/ffg.yaml and configs/base.yaml, + avoiding the heavy ODLightningCLI/trainer machinery that is + unnecessary for inference. + + Returns: + An FFGSynoVideoLearner instance (unloaded weights). + """ + # Add model_code to sys.path so src.* imports resolve + model_code_dir = os.environ.get("MODEL_CODE_DIR", "/app/model_code") + if model_code_dir not in sys.path: + sys.path.insert(0, model_code_dir) + + # Suppress wandb import in svl.py (it imports wandb at top) + if "wandb" not in sys.modules: + import types + + wandb_stub = types.ModuleType("wandb") + sys.modules["wandb"] = wandb_stub + + from src.model.clip.svl import FFGSynoVideoLearner + + model = FFGSynoVideoLearner( + # FFG-specific params (from configs/clip/L14/ffg.yaml) + face_feature_path=FACE_FEATURES_PATH, + face_parts=["lips", "skin", "eyes", "nose"], + architecture="ViT-L/14", + num_frames=NUM_FRAMES, + ksize_s=5, + ksize_t=5, + s_k_attr="k", + s_v_attr="emb", + t_attrs=["q", "k", "v"], + # Defaults from base model + text_embed=False, + op_mode=["S", "T"], + ) + return model + + +# ── Global state ─────────────────────────────────────────────────── + +_model = None +_face_detector = None +_transform = None +_device: Optional[torch.device] = None +_load_lock = threading.Lock() + + +def _load_models() -> None: + """Load DFD-FCG model and MTCNN face detector (thread-safe).""" + global _model, _face_detector, _transform, _device + + if _model is not None: + return + + with _load_lock: + if _model is not None: + return + + _device = _get_device() + if _device.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + if _device.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) + else: + logger.warning( + "Device: %s (no CUDA -- check " "nvidia-container-toolkit)", + _device, + ) + logger.info("Loading DFD-FCG model on %s ...", _device) + + # ── Face detector (MTCNN) ────────────────────────────── + from facenet_pytorch import MTCNN + + _face_detector = MTCNN( + keep_all=True, + device=_device, + post_process=False, + ) + + # ── DFD-FCG model ────────────────────────────────────── + if not os.path.exists(WEIGHTS_PATH): + raise FileNotFoundError(f"DFD-FCG weights not found at {WEIGHTS_PATH}") + + model = _build_model_from_config() + model_cls = model.__class__ + + # The checkpoint stores face_feature_path as a + # relative path (e.g. "misc/..."). We must set CWD + # to model_code/ so load_from_checkpoint can resolve + # it when replaying the constructor. + model_code_dir = os.environ.get("MODEL_CODE_DIR", "/app/model_code") + original_cwd = os.getcwd() + os.chdir(model_code_dir) + + try: + model = model_cls.load_from_checkpoint(WEIGHTS_PATH) + except Exception: + logger.info("Strict checkpoint load failed, " "retrying non-strict.") + model = model_cls.load_from_checkpoint(WEIGHTS_PATH, strict=False) + finally: + os.chdir(original_cwd) + + model = model.to(_device) + model.requires_grad_(False) + # Use torch eval -- demo.py uses the same pattern + model.train(False) + + _transform = model.transform + _model = model + logger.info("DFD-FCG model loaded successfully.") + + +def _is_model_loaded() -> bool: + """Return True if model and face detector are loaded.""" + return _model is not None and _face_detector is not None + + +# ── FastAPI app ──────────────────────────────────────────────────── + +app = FastAPI( + title="DFD-FCG Detection Service", + description=( + "Deepfake Detection via Facial Component Guidance " "(CLIP ViT-L/14, CVPR 2025)" + ), + version="1.0.0", +) + + +class PredictRequest(BaseModel): + """Incoming prediction request.""" + + video_data: str # Base64-encoded video bytes + threshold: float = 0.5 + + +class PredictResponse(BaseModel): + """Outgoing prediction result.""" + + model_config = ConfigDict(populate_by_name=True) + + model: str = "dfd_fcg_detection" + probability: float + prediction: int + class_name: str = Field(..., alias="class") + inference_time: float + metadata: Dict[str, Any] + + +@app.on_event("startup") +async def startup_event(): + """Optionally preload model at startup.""" + if PRELOAD_MODEL: + _load_models() + + +@app.get("/") +def root(): + """Service info endpoint.""" + return { + "service": "dfd_fcg_detection", + "port": MODEL_PORT, + "model_loaded": _is_model_loaded(), + "device": str(_device) if _device else "unknown", + } + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and _device is not None and _device.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +def health(): + """Health check endpoint.""" + return { + "status": "healthy", + "model": "dfd_fcg_detection", + "device": str(_device) if _device else "cpu", + "model_loaded": _is_model_loaded(), + "weights_exist": os.path.exists(WEIGHTS_PATH), + **_gpu_health_info(), + } + + +# ── Video / face utilities ───────────────────────────────────────── + + +def _extract_dense_frames( + video_path: str, +) -> Tuple[List[np.ndarray], float]: + """Extract all frames from a video for clip-based sampling. + + Args: + video_path: Path to the video on disk. + + Returns: + Tuple of (list of RGB uint8 arrays, fps). + """ + cap = cv2.VideoCapture(video_path) + fps = cap.get(cv2.CAP_PROP_FPS) or 30.0 + frames: List[np.ndarray] = [] + + while True: + ret, frame = cap.read() + if not ret: + break + frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + + cap.release() + return frames, fps + + +def _crop_face( + img: np.ndarray, + bbox: Tuple[float, float, float, float], + margin: float = MARGIN_FACTOR, +) -> np.ndarray: + """Crop a face region from an image with a relative margin. + + Args: + img: RGB image array (H, W, 3). + bbox: (x0, y0, x1, y1) face bounding box. + margin: Fraction of bbox dimension to add as padding. + + Returns: + Cropped face region as a numpy array. + """ + h_img, w_img = img.shape[:2] + x0, y0, x1, y1 = bbox + w = x1 - x0 + h = y1 - y0 + + x0_new = max(0, int(x0 - w * margin / 2)) + x1_new = min(w_img, int(x1 + w * margin / 2) + 1) + y0_new = max(0, int(y0 - h * margin / 2)) + y1_new = min(h_img, int(y1 + h * margin / 2) + 1) + + return img[y0_new:y1_new, x0_new:x1_new] + + +def _detect_and_crop_face( + frame: np.ndarray, +) -> Optional[np.ndarray]: + """Detect the largest face in a frame and return the crop. + + Uses MTCNN for detection, selects the largest bounding box, + crops with margin, and returns the face region. + + Args: + frame: RGB image array (H, W, 3). + + Returns: + Face crop as uint8 array, or None if no face found. + """ + assert _face_detector is not None + + pil_img = Image.fromarray(frame) + boxes, _ = _face_detector.detect(pil_img) + + if boxes is None or len(boxes) == 0: + return None + + # Select the largest face by bounding box area + areas = [(b[2] - b[0]) * (b[3] - b[1]) for b in boxes] + best_idx = int(np.argmax(areas)) + box = boxes[best_idx] + + x0, y0, x1, y1 = box.tolist() + face = _crop_face(frame, (x0, y0, x1, y1)) + if face.size == 0: + return None + + return face + + +def _prepare_clip_tensor( + face_crops: List[np.ndarray], +) -> Optional[torch.Tensor]: + """Apply the model's CLIP transform to face crops. + + Converts a list of face crop arrays into a single tensor + of shape (1, num_frames, 3, 224, 224) suitable for the + DFD-FCG model. + + Args: + face_crops: List of RGB uint8 face crop arrays. + + Returns: + Tensor of shape (1, T, 3, 224, 224) or None. + """ + if not face_crops or _transform is None: + return None + + # The model transform expects either PIL or tensor input. + # Convert each crop to a torch tensor (C, H, W) uint8, + # which the custom _to_tensor in clip.py handles (divides + # by 255 if max > 1). + transformed = [] + for crop in face_crops: + # Convert HWC uint8 -> CHW tensor + t = torch.from_numpy(crop).permute(2, 0, 1) + t = _transform(t) + transformed.append(t) + + # Stack into (T, C, H, W), then add batch dim + clip_tensor = torch.stack(transformed, dim=0) + clip_tensor = clip_tensor.unsqueeze(0) # (1, T, C, H, W) + return clip_tensor + + +def _find_nearest_crop( + crops: List[Optional[np.ndarray]], + idx: int, +) -> Optional[np.ndarray]: + """Find the nearest non-None face crop to a given index. + + Searches outward from idx in both directions to find the + closest frame that had a detected face. + + Args: + crops: List of face crops (may contain None entries). + idx: Target index. + + Returns: + The nearest non-None crop, or None if all are None. + """ + n = len(crops) + for offset in range(n): + for candidate in (idx - offset, idx + offset): + if 0 <= candidate < n and crops[candidate] is not None: + return crops[candidate] + return None + + +# ── Prediction endpoint ──────────────────────────────────────────── + + +@app.post("/predict", response_model=PredictResponse) +async def predict(request: PredictRequest): + """Run DFD-FCG deepfake detection on a base64-encoded video. + + Pipeline: + 1. Decode video and write to temp file. + 2. Extract all frames from the video. + 3. Detect and crop the main face per frame (MTCNN). + 4. Build overlapping clips of NUM_FRAMES face crops. + 5. Run each clip through the DFD-FCG model. + 6. Average per-clip fake probabilities. + + If no faces are detected in any frame, the service returns + probability=0.5 (undetermined) rather than raising an error. + """ + if not _is_model_loaded(): + _load_models() + + start_time = time.time() + + # ── Decode video ─────────────────────────────────────────── + with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp: + try: + video_bytes = base64.b64decode(request.video_data) + tmp.write(video_bytes) + tmp_path = tmp.name + except Exception as e: + raise HTTPException( + status_code=400, + detail=f"Failed to decode video: {e}", + ) + + try: + # ── Extract frames ───────────────────────────────────── + frames, fps = _extract_dense_frames(tmp_path) + if not frames: + raise HTTPException( + status_code=400, + detail="Could not extract frames from video.", + ) + + # ── Detect and crop faces ────────────────────────────── + face_crops: List[Optional[np.ndarray]] = [] + faces_detected = 0 + for frame in frames: + crop = _detect_and_crop_face(frame) + face_crops.append(crop) + if crop is not None: + faces_detected += 1 + + if faces_detected == 0: + # No faces in any frame -- undetermined + return PredictResponse( + probability=0.5, + prediction=0, + class_name="real", + inference_time=time.time() - start_time, + metadata={ + "frames_total": len(frames), + "frames_with_faces": 0, + "clips_evaluated": 0, + "device": str(_device), + }, + ) + + # ── Build clips and run inference ────────────────────── + # Sample clip indices matching the demo.py pattern: + # indices spaced by stride*fps, NUM_FRAMES per clip + stride_frames = max(1, int(math.floor(FRAME_STRIDE * fps))) + clip_indices = [i * stride_frames for i in range(NUM_FRAMES)] + max_start = len(frames) - clip_indices[-1] - 1 + + probs: List[float] = [] + batch_size = 8 + + if max_start <= 0: + # Video too short for stride-based sampling; + # uniformly sample NUM_FRAMES frames instead. + sample_idx = np.linspace( + 0, + len(frames) - 1, + NUM_FRAMES, + endpoint=True, + dtype=int, + ) + clip_crops = [] + for idx in sample_idx: + c = face_crops[idx] + if c is None: + c = _find_nearest_crop(face_crops, idx) + if c is not None: + clip_crops.append(c) + + if len(clip_crops) == NUM_FRAMES: + tensor = _prepare_clip_tensor(clip_crops) + if tensor is not None: + tensor = tensor.to(_device) + with torch.no_grad(): + result = _model.evaluate(tensor) + p = result["logits"].softmax(dim=-1)[:, 1].cpu().item() + probs.append(p) + else: + # Process clips in batches + clip_starts = list(range(0, max_start + 1)) + for batch_start in range(0, len(clip_starts), batch_size): + batch_clips = clip_starts[batch_start : batch_start + batch_size] + tensors = [] + for start in batch_clips: + clip_crops = [] + for offset in clip_indices: + idx = start + offset + c = face_crops[idx] + if c is None: + c = _find_nearest_crop(face_crops, idx) + if c is not None: + clip_crops.append(c) + + if len(clip_crops) == NUM_FRAMES: + t = _prepare_clip_tensor(clip_crops) + if t is not None: + tensors.append(t) + + if tensors: + batch_tensor = torch.cat(tensors, dim=0).to(_device) + with torch.no_grad(): + result = _model.evaluate(batch_tensor) + batch_probs = ( + result["logits"] + .softmax(dim=-1)[:, 1] + .flatten() + .cpu() + .tolist() + ) + probs.extend(batch_probs) + + # ── Aggregate ────────────────────────────────────────── + if probs: + probability = float(np.mean(probs)) + else: + probability = 0.5 + + prediction = 1 if probability >= request.threshold else 0 + class_name = "fake" if prediction == 1 else "real" + + return PredictResponse( + probability=probability, + prediction=prediction, + class_name=class_name, + inference_time=time.time() - start_time, + metadata={ + "frames_total": len(frames), + "frames_with_faces": faces_detected, + "clips_evaluated": len(probs), + "device": str(_device), + }, + ) + + except HTTPException: + raise + except Exception as e: + logger.exception("Error during DFD-FCG prediction") + raise HTTPException(status_code=500, detail=str(e)) + finally: + if os.path.exists(tmp_path): + os.remove(tmp_path) + gc.collect() + + +if __name__ == "__main__": + uvicorn.run(app, host="0.0.0.0", port=MODEL_PORT) diff --git a/video/dfd-fcg/requirements.txt b/video/dfd-fcg/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..5cfd49618f01a07e4bc411fb2668a4308d67c823 --- /dev/null +++ b/video/dfd-fcg/requirements.txt @@ -0,0 +1,17 @@ +fastapi +uvicorn +pydantic +python-multipart +torch>=2.0.0 +torchvision>=0.15.0 +facenet-pytorch +opencv-python-headless +numpy<2.0.0 +Pillow +open-clip-torch +lightning>=2.0.0 +torchmetrics +jsonargparse[omegaconf] +omegaconf +ftfy +regex diff --git a/video/fake-stormer/Dockerfile b/video/fake-stormer/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..133b87580fb3281d9160246849fdd9967edd73ca --- /dev/null +++ b/video/fake-stormer/Dockerfile @@ -0,0 +1,44 @@ +FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-devel + +ENV DEBIAN_FRONTEND=noninteractive + +WORKDIR /app + +# Fix potential GPG key issues on older CUDA base images +RUN (apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub 2>/dev/null || true) && \ + apt-get update && apt-get install -y \ + libgl1 \ + libglib2.0-0 \ + libsm6 \ + libxext6 \ + libxrender-dev \ + libopencv-dev \ + build-essential \ + git \ + && rm -rf /var/lib/apt/lists/* + +# Copy requirements and install +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +# Create logs directory +RUN mkdir -p logs + +# Copy model code and app +COPY model_code/ ./model_code/ +COPY app.py . + +# Environment variables +ENV MODEL_PORT=7001 +ENV PRELOAD_MODEL=false + +# Expose port +EXPOSE 7001 + +# Drop root privileges +RUN adduser --disabled-password --gecos '' appuser && \ + chown -R appuser:appuser /app/logs +USER appuser + +# Run the service +CMD ["python", "app.py"] diff --git a/video/fake-stormer/app.py b/video/fake-stormer/app.py new file mode 100644 index 0000000000000000000000000000000000000000..b762e5c01e4bbf916d81288e15092e7240d9a1ea --- /dev/null +++ b/video/fake-stormer/app.py @@ -0,0 +1,311 @@ +"""FakeSTormer deepfake video detection service. + +Wraps the FakeSTormer (ICCV 2025) video deepfake detection model +with a FastAPI endpoint. +""" + +import base64 +import gc +import logging +import os +import platform +import sys +import tempfile +import time +from typing import Any, Dict, List + +import cv2 +import numpy as np +import torch +import uvicorn +from fastapi import FastAPI, HTTPException +from PIL import Image +from pydantic import BaseModel, ConfigDict, Field + +# Add model_code to path to allow imports +sys.path.insert(0, os.path.join(os.getcwd(), "model_code")) + +from configs.get_config import load_config +from models import MODELS, build_model, load_pretrained +from package_utils.transform import final_transform + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +MODEL_PORT = int(os.environ.get("MODEL_PORT", 7001)) +PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true" +CONFIG_PATH = os.environ.get( + "CONFIG_PATH", "model_code/configs/temporal/FakeSFormer_base_c23.yaml" +) +WEIGHTS_PATH = os.environ.get( + "WEIGHTS_PATH", + "model_code/weights/TopDownDetector_C23_ViTBase224_ST_hm100_tempLOC0.2_4SBI_SAM_mp0.01_temp2_0vidAug_0.35_temp_normlms_normREAL_model_best.pth", +) + + +def _get_device(): + """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if override: + pass # Invalid override, fall through to auto-detect + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + return torch.device("cpu") + + +# Global model references +_model = None +_cfg = None +_transforms = None +_device = None + + +def _load_models(): + """Load FakeSTormer model.""" + global _model, _cfg, _transforms, _device # noqa: F824 + + if _model is not None: + return + + _device = _get_device() + + if _device.type == "cuda": + torch.backends.cudnn.benchmark = True + if hasattr(torch, "set_float32_matmul_precision"): + torch.set_float32_matmul_precision("high") + + if _device.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) + else: + logger.warning( + "Device: %s (no CUDA available -- check nvidia-container-toolkit)", + _device, + ) + + logger.info(f"Loading FakeSTormer model on {_device}...") + + try: + # Load config + _cfg = load_config(CONFIG_PATH) + + # Build model + _model = build_model(_cfg.MODEL, MODELS).to(torch.float) + + # Load weights + if not os.path.exists(WEIGHTS_PATH): + logger.error(f"Weights file not found at {WEIGHTS_PATH}") + # Use placeholder if weights missing? User asked for as true as possible results, + # so we should fail if weights are missing, but let's see. + raise FileNotFoundError(f"Weights missing: {WEIGHTS_PATH}") + + logger.info(f"Loading weight ... {WEIGHTS_PATH}") + _model = load_pretrained(_model, WEIGHTS_PATH) + _model = _model.to(_device) + _model.eval() + + # Setup transforms + _transforms = final_transform(_cfg.DATASET) + + logger.info("FakeSTormer model loaded successfully.") + + except Exception as e: + logger.error(f"Failed to load FakeSTormer model: {e}") + raise e + + +def _is_model_loaded(): + return _model is not None + + +app = FastAPI( + title="FakeSTormer Detection Service", + description="Vulnerability-Aware Spatio-Temporal Learning for Generalizable Deepfake Video Detection", + version="1.0.0", +) + + +class PredictRequest(BaseModel): + video_data: str # Base64 encoded video + threshold: float = 0.5 + + +class PredictResponse(BaseModel): + model_config = ConfigDict(populate_by_name=True) + model: str = "fakestormer" + probability: float + prediction: int + class_name: str = Field(..., alias="class") + inference_time: float + metadata: Dict[str, Any] + + +@app.on_event("startup") +async def startup_event(): + if PRELOAD_MODEL: + _load_models() + + +@app.get("/") +def root(): + return { + "service": "fakestormer", + "port": MODEL_PORT, + "model_loaded": _is_model_loaded(), + "device": str(_device) if _device else "unknown", + } + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and _device is not None and _device.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +def health(): + return { + "status": "ok", + "model_loaded": _is_model_loaded(), + "weights_exist": os.path.exists(WEIGHTS_PATH), + "device": str(_device) if _device else "cpu", + **_gpu_health_info(), + } + + +def extract_frames(video_path: str, num_frames: int = 4) -> List[Image.Image]: + """Extract frames from video file uniformly.""" + cap = cv2.VideoCapture(video_path) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + + if total_frames <= 0: + cap.release() + return [] + + # Get indices for uniform sampling + indices = np.linspace(0, total_frames - 1, num_frames, dtype=int) + + frames = [] + for idx in indices: + cap.set(cv2.CAP_PROP_POS_FRAMES, idx) + ret, frame = cap.read() + if ret: + # Convert BGR to RGB + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + # Crop 15 pixels from each side as in test.py + H, W, _ = frame.shape + if H > 30 and W > 30: + frame = frame[15 : H - 15, 15 : W - 15] + frames.append(Image.fromarray(frame)) + + cap.release() + + # Pad if not enough frames + while len(frames) < num_frames and len(frames) > 0: + frames.append(frames[-1]) + + return frames + + +@app.post("/predict", response_model=PredictResponse) +async def predict(request: PredictRequest): + global _model, _cfg, _transforms, _device # noqa: F824 + + if not _is_model_loaded(): + _load_models() + + start_time = time.time() + + # Create a temporary file to save the video + with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp_video: + try: + video_bytes = base64.b64decode(request.video_data) + tmp_video.write(video_bytes) + tmp_video_path = tmp_video.name + except Exception as e: + raise HTTPException(status_code=400, detail=f"Failed to decode video: {e}") + + try: + num_frames = _cfg.DATASET.DATA.SAMPLES_PER_VIDEO.NUM_FRAMES or 4 + frames = extract_frames(tmp_video_path, num_frames) + + if not frames: + raise HTTPException( + status_code=500, detail="Failed to extract frames from video." + ) + + # Preprocess frames + transformed_imgs = [] + image_size = _cfg.DATASET.IMAGE_SIZE # [224, 224] + + for frame in frames: + img_resize = frame.resize((int(image_size[0]), int(image_size[1]))) + img_resize = np.array(img_resize) / 255.0 + img_tensor = _transforms(img_resize).to(torch.float) + transformed_imgs.append(img_tensor.unsqueeze(0)) + + # Stack and prepare for model [B, T, C, H, W] + input_tensor = torch.cat(transformed_imgs, 0) # [T, C, H, W] + input_tensor = input_tensor.to(_device) + input_tensor = input_tensor.unsqueeze(0) # [1, T, C, H, W] + + # FakeSTormer model expects [1, C, T, H, W] + input_tensor = input_tensor.transpose(1, 2) # [1, C, T, H, W] + + with torch.no_grad(): + outputs = _model(input_tensor) + + if isinstance(outputs, list): + outputs = outputs[0] + + prob = outputs["cls"].sigmoid().cpu().item() + + prediction = 1 if prob >= request.threshold else 0 + class_name = "fake" if prediction == 1 else "real" + + return PredictResponse( + probability=float(prob), + prediction=prediction, + class_name=class_name, + inference_time=time.time() - start_time, + metadata={"frames_extracted": len(frames), "device": str(_device)}, + ) + + except Exception as e: + logger.exception("Error during FakeSTormer prediction") + raise HTTPException(status_code=500, detail=str(e)) + finally: + # Cleanup temporary file + if os.path.exists(tmp_video_path): + os.remove(tmp_video_path) + gc.collect() + + +if __name__ == "__main__": + uvicorn.run(app, host="0.0.0.0", port=MODEL_PORT) diff --git a/video/fake-stormer/model_code/LICENSE b/video/fake-stormer/model_code/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..a251c2be404e71aadcf03faaf2310ae89d7f7f9c --- /dev/null +++ b/video/fake-stormer/model_code/LICENSE @@ -0,0 +1,14 @@ +SnT academic license +(based on the MIT license with academic limitations) + +Copyright 2025 University of Luxembourg + +Permission is hereby granted, free of charge, to any academic and research institution and researcher obtaining a copy of this software, all derivative works and associated documentation files (the “Software”), to deal in the Software for academic research and development, testing, validation and academic or scientific research purposes only, including without limitation the rights to use, copy, modify, merge, and/or publish copies of the Software, and to permit persons to whom the Software are furnished to do so for such purposes only, subject to the following conditions: + +All copies or substantial portions of the Software shall include the above copyright notice and this permission notice. + +All copies of the Software shall be distributed under the terms of this license only. + +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. + +Any other use of the Software requires a separate license agreement. Should you be interested in making other uses of the Software, please send an email to snt-tto@uni.lu. diff --git a/video/fake-stormer/model_code/NOTICE b/video/fake-stormer/model_code/NOTICE new file mode 100644 index 0000000000000000000000000000000000000000..0565e669ae2d8594c7fa898d23e24070755d466c --- /dev/null +++ b/video/fake-stormer/model_code/NOTICE @@ -0,0 +1,8 @@ +FakeSTormer is © 2024 - 2025 University of Luxembourg + +Authors: +- Dat NGUYEN +- Marcella ASTRID +- Anis KACEM +- Enjie GHORBEL +- Djamila AOUADA diff --git a/video/fake-stormer/model_code/README.md b/video/fake-stormer/model_code/README.md new file mode 100644 index 0000000000000000000000000000000000000000..060383b0db8c6e938a49a5d518ba4acf019b2548 --- /dev/null +++ b/video/fake-stormer/model_code/README.md @@ -0,0 +1,187 @@ +# [ICCV2025] [FakeSTormer] Vulnerability-Aware Spatio-Temporal Learning for Generalizable Deepfake Video Detection + +![alt text](./demo/method.png?raw=true) +This is an official implementation of FakeSTormer! [[📜Paper](https://openaccess.thecvf.com/content/ICCV2025/papers/Nguyen_Vulnerability-Aware_Spatio-Temporal_Learning_for_Generalizable_Deepfake_Video_Detection_ICCV_2025_paper.pdf)] + + +## Updates +- [x] 26/11/2025:*Official release of code (v1) and pretrained weights 🌈.* +- [x] 08/07/2025: *First version pre-released for this open source code 🌱.* +- [x] 26/06/2025: *FakeSTormer has been accepted to ICCV2025 🎉.* + + +## Abstract +Detecting deepfake videos is highly challenging given the complexity of characterizing spatio-temporal artifacts. Most existing methods rely on binary classifiers trained using real and fake image sequences, therefore hindering their generalization capabilities to unseen generation methods. Moreover, with the constant progress in generative Artificial Intelligence (AI), deepfake artifacts are becoming imperceptible at both the spatial and the temporal levels, making them extremely difficult to capture. To address these issues, we propose a fine-grained deepfake video detection approach called FakeSTormer that enforces the modeling of subtle spatio-temporal inconsistencies while avoiding overfitting. Specifically, we introduce a multi-task learning framework that incorporates two auxiliary branches for explicitly attending artifact-prone spatial and temporal regions. Additionally, we propose a video-level data synthesis strategy that generates pseudo-fake videos with subtle spatio-temporal artifacts, providing high-quality samples and hand-free annotations for our additional branches. Extensive experiments on several challenging benchmarks demonstrate the superiority of our approach compared to recent state-of-the-art methods. + + +## Main Results +Results on 6 datasets ([CDF2](https://github.com/yuezunli/celeb-deepfakeforensics), [DFW](https://github.com/deepfakeinthewild/deepfake-in-the-wild), [DFD](https://blog.research.google/2019/09/contributing-data-to-deepfake-detection.html), [DFDC, DFDCP](https://ai.meta.com/datasets/dfdc/), and [DiffSwap](https://openaccess.thecvf.com/content/CVPR2023/papers/Zhao_DiffSwap_High-Fidelity_and_Controllable_Face_Swapping_via_3D-Aware_Masked_Diffusion_CVPR_2023_paper.pdf)) under cross-dataset evaluation setting reported by AUC (%) at video-level. + +| | CDF2 | DFW | DFD | DFDC | DFDCP | DiffSwap | +|--|--------|------------|------------|------------|---------|-----------| +|
Compression
c23
c0
|
AUC
92.4
96.5
|
AUC
74.2
76.3
|
AUC
98.5
98.9
|
AUC
74.6
77.6
|
AUC
90.0
94.1
|
AUC
96.9
97.7
+ + +## Recommended Environment +*For experimental purposes, we encourage the installation of the following libraries. Both Conda or Python virtual env should work.* + +* CUDA: 11.4 +* [Python](https://www.python.org/): >= 3.8.x +* [PyTorch](https://pytorch.org/get-started/previous-versions/): 1.8.0 +* [TensorboardX](https://github.com/lanpa/tensorboardX): 2.5.1 +* [ImgAug](https://github.com/aleju/imgaug): 0.4.0 +* [Scikit-image](https://scikit-image.org/): 0.17.2 +* [Torchvision](https://pytorch.org/vision/stable/index.html): 0.9.0 +* [Albumentations](https://albumentations.ai/): 1.1.0 +* [mmcv](https://github.com/open-mmlab/mmcv): 1.6.1 +* [natsort](https://pypi.org/project/natsort/): 8.4.0 + + + +## Pre-trained Models +* 📌 *The pre-trained weights of FakeSTormer can be found [here](https://www.dropbox.com/scl/fo/elk2szqf0du4l6zm5job9/AAdVmNH--6ywHBZGNQJlR5o?rlkey=j8xesf2fu4ahxdw99w5ndrkb2&st=fe6drzpx&dl=0)* + + +## Docker Build (Optional) +*We further provide an optional Docker file that can be used to build a working env with Docker. More detailed steps can be found [here](dockerfiles/README.md).* + +1. Install docker to the system (skip the step if docker has already been installed): + ```shell + sudo apt install docker + ``` +2. To start your docker environment, please go to the folder **dockerfiles**: + ```shell + cd dockerfiles + ``` +3. Create a docker image (you can put any name you want): + ```shell + docker build --tag 'fakestormer' . + ``` + + +## Quickstart +1. **Preparation** + + 1. ***Prepare environment*** + + Installing main packages as the recommended environment. *Note that we recommend building mmcv from source as below.* + > git clone https://github.com/open-mmlab/mmcv.git \ + cd mmcv \ + git checkout v1.6.1 \ + MMCV_WITH_OPS=1 pip install -e . + + 2. ***Prepare dataset*** + + 1. Downloading [FF++](https://github.com/ondyari/FaceForensics) *Original* dataset for training data preparation. Following the original split convention, it is firstly used to randomly extract frames and facial crops: + ``` + python package_utils/images_crop.py -d {dataset} \ + -c {compression} \ + -n {num_frames} \ + -t {task} + ``` + (*This script can also be utilized for cropping faces in other datasets such as [CDF2](https://github.com/yuezunli/celeb-deepfakeforensics), [DFD](https://blog.research.google/2019/09/contributing-data-to-deepfake-detection.html), [DFDCP, DFDC](https://ai.meta.com/datasets/dfdc/) for cross-evaluation test. You do not need to run crop for [DFW](https://github.com/deepfakeinthewild/deepfake-in-the-wild) as the data is already preprocessed*). + + | Parameter | Value | Definition | + | --- | --- | --- | + | -d | Subfolder in each dataset. For example: *['Face2Face','Deepfakes','FaceSwap','NeuralTextures', ...]*| You can use one of those datasets.| + | -c | *['raw','c23','c40']*| You can use one of those compressions| + | -n | *256* | Number of frames (*default* 32 for val/test and 256 for train) | + | -t | *['train', 'val', 'test']* | Default train| + + These faces cropped are saved for online pseudo-fake generation in the training process, following the data structure below: + + ``` + ROOT = '/data/deepfake_cluster/datasets_df' + └── Celeb-DFv2 + └──... + └── FF++ + └── c0 + └── c23 + ├── test + │   └── videos + │   └── Deepfakes + | ├── 000_003 + | ├── 044_945 + | ├── 138_142 + | ├── ... + │   ├── Face2Face + │   ├── FaceSwap + │   ├── NeuralTextures + │   └── original + | └── frames + ├── train + │   └── videos + │   └── aligned + | ├── 001 + | ├── 002 + | ├── ... + │   └── original + | ├── 001 + | ├── 002 + | ├── ... + | └── frames + └── val + └── videos + ├── aligned + └── original + └── frames + └── c40 + ``` + + 2. Downloading **Dlib** [[81]](https://github.com/codeniko/shape_predictor_81_face_landmarks) facial landmarks detector pretrained and place into ```/pretrained/``` for *SBI* synthesis. + + 3. Landmarks detection. After completing the following script running, a file that stores metadata information of the data is saved at ```processed_data/c23/{SPLIT}_FaceForensics_videos_.json```. + ``` + python package_utils/geo_landmarks_extraction.py \ + --config configs/data_preprocessing_c23.yaml \ + --extract_landmarks + ``` + +2. **Training script** + + We offer a number of config files for different compression levels of training data. For *c23*, opening ```configs/temporal/FakeSTormer_base_c23.yaml```, please make sure you set ```TRAIN: True``` and ```FROM_FILE: True``` and run: + ``` + .scripts/fakestormer_sbi.sh + ``` + + Otherwise, with *[c0, c40]*, the config file is ```configs/temporal/FakeSTormer_base_[c0, c40].yaml```. You can also find other configs for other network architectures in the ```configs/``` folder. + + +3. **Testing script** + + Opening ```configs/temporal/FakeSTormer_base_c23.yaml```, with ```subtask: eval``` in the *test* section, we support evaluation mode, please turn off ```TRAIN: False``` and ```FROM_FILE: False``` and run: + ``` + .scripts/test_fakestormer.sh + ``` + For others (.e.g., data compression levels, network architectures), please change the path of the corresponding config file. + + > ⚠️ *Please make sure you set the correct path to your downloaded pre-trained weights in the config files.* + + > ℹ️ *Flip test can be used by setting ```flip_test: True```* + + > ℹ️ *The mode for single video inference is also provided, please set ```sub_task: test_vid``` and pass a video path as an argument in test.py* + + +## Contact +Please contact dat.nguyen@uni.lu. Any questions or discussions are welcomed! + + +## License +This software is © University of Luxembourg and is licensed under the snt academic license. See [LICENSE](LICENSE) + + +## Acknowledge +We acknowledge the excellent implementation from [OpenMMLab](https://github.com/open-mmlab) ([mmengine](https://github.com/open-mmlab/mmengine), [mmcv](https://github.com/open-mmlab/mmcv)), [SBI](https://github.com/mapooon/SelfBlendedImages), and [LAA-Net](https://github.com/10Ring/LAA-Net). + + +## Citation +Please kindly consider citing our papers in your publications. +``` +@inproceedings{nguyen2025vulnerability, + title={Vulnerability-Aware Spatio-Temporal Learning for Generalizable Deepfake Video Detection}, + author={Nguyen, Dat and Astrid, Marcella and Kacem, Anis and Ghorbel, Enjie and Aouada, Djamila}, + booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision}, + pages={10786--10796}, + year={2025} +} +``` diff --git a/video/fake-stormer/model_code/Third_Party_License_Notice b/video/fake-stormer/model_code/Third_Party_License_Notice new file mode 100644 index 0000000000000000000000000000000000000000..728074becf0d25831d1868a35adc309af6415ae4 --- /dev/null +++ b/video/fake-stormer/model_code/Third_Party_License_Notice @@ -0,0 +1,1606 @@ +FakeSTormer is © 2024 - 2025 University of Luxembourg +Developed by: Dat Nguyen at CVI2/SnT + +FakeSTormer is licensed under the SnT academic license (see #LICENSE) + +FakeSTormer includes the following components: + +From PyTorch: + +Copyright (c) 2016- Facebook, Inc (Adam Paszke) +Copyright (c) 2014- Facebook, Inc (Soumith Chintala) +Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert) +Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu) +Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu) +Copyright (c) 2011-2013 NYU (Clement Farabet) +Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston) +Copyright (c) 2006 Idiap Research Institute (Samy Bengio) +Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz) + + +From Caffe2: + +Copyright (c) 2016-present, Facebook Inc. All rights reserved. + +All contributions by Facebook: +Copyright (c) 2016 Facebook Inc. + +All contributions by Google: +Copyright (c) 2015 Google Inc. +All rights reserved. + +All contributions by Yangqing Jia: +Copyright (c) 2015 Yangqing Jia +All rights reserved. + +All contributions by Kakao Brain: +Copyright 2019-2020 Kakao Brain + +All contributions by Cruise LLC: +Copyright (c) 2022 Cruise LLC. +All rights reserved. + +All contributions from Caffe: +Copyright(c) 2013, 2014, 2015, the respective contributors +All rights reserved. + +All other contributions: +Copyright(c) 2015, 2016 the respective contributors +All rights reserved. + +Caffe2 uses a copyright model similar to Caffe: each contributor holds +copyright over their contributions to Caffe2. The project versioning records +all such contribution and copyright details. If a contributor wants to further +mark their specific copyright on a particular contribution, they should +indicate their copyright solely in the commit message of the change when it is +committed. + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. 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. + +3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America + and IDIAP Research Institute 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 OWNER 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 OF THIS SOFTWARE, EVEN IF ADVISED OF THE +POSSIBILITY OF SUCH DAMAGE. + + +From TorchVision: + +BSD 3-Clause License + +Copyright (c) Soumith Chintala 2016, +All rights reserved. + +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 +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +From Tensorboard: + +Copyright 2017 The TensorFlow Authors. All rights reserved. + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2017, The TensorFlow Authors. + + 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. + + +From TensorboardX: + +MIT License + +Copyright (c) 2017 Tzu-Wei Huang + +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. + + +From ImgAug: + +The MIT License (MIT) + +Copyright (c) 2015 aleju + +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. + + +From Torchcam: + +Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. + + +From Dlib: + +Boost Software License - Version 1.0 - August 17th, 2003 + +Permission is hereby granted, free of charge, to any person or organization +obtaining a copy of the software and accompanying documentation covered by +this license (the "Software") to use, reproduce, display, distribute, +execute, and transmit the Software, and to prepare derivative works of the +Software, and to permit third-parties to whom the Software is furnished to +do so, all subject to the following: + +The copyright notices in the Software and this entire statement, including +the above license grant, this restriction and the following disclaimer, +must be included in all copies of the Software, in whole or in part, and +all derivative works of the Software, unless such copies or derivative +works are solely in the form of machine-executable object code generated by +a source language processor. + +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, TITLE AND NON-INFRINGEMENT. IN NO EVENT +SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE +FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE, +ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +DEALINGS IN THE SOFTWARE. + + +From Matplotlib: + +License agreement for matplotlib versions 1.3.0 and later +========================================================= + +1. This LICENSE AGREEMENT is between the Matplotlib Development Team +("MDT"), and the Individual or Organization ("Licensee") accessing and +otherwise using matplotlib software in source or binary form and its +associated documentation. + +2. Subject to the terms and conditions of this License Agreement, MDT +hereby grants Licensee a nonexclusive, royalty-free, world-wide license +to reproduce, analyze, test, perform and/or display publicly, prepare +derivative works, distribute, and otherwise use matplotlib +alone or in any derivative version, provided, however, that MDT's +License Agreement and MDT's notice of copyright, i.e., "Copyright (c) +2012- Matplotlib Development Team; All Rights Reserved" are retained in +matplotlib alone or in any derivative version prepared by +Licensee. + +3. In the event Licensee prepares a derivative work that is based on or +incorporates matplotlib or any part thereof, and wants to +make the derivative work available to others as provided herein, then +Licensee hereby agrees to include in any such work a brief summary of +the changes made to matplotlib . + +4. MDT is making matplotlib available to Licensee on an "AS +IS" basis. MDT MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, MDT MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF MATPLOTLIB +WILL NOT INFRINGE ANY THIRD PARTY RIGHTS. + +5. MDT SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF MATPLOTLIB + FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR +LOSS AS A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING +MATPLOTLIB , OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF +THE POSSIBILITY THEREOF. + +6. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +7. Nothing in this License Agreement shall be deemed to create any +relationship of agency, partnership, or joint venture between MDT and +Licensee. This License Agreement does not grant permission to use MDT +trademarks or trade name in a trademark sense to endorse or promote +products or services of Licensee, or any third party. + +8. By copying, installing or otherwise using matplotlib , +Licensee agrees to be bound by the terms and conditions of this License +Agreement. + +License agreement for matplotlib versions prior to 1.3.0 +======================================================== + +1. This LICENSE AGREEMENT is between John D. Hunter ("JDH"), and the +Individual or Organization ("Licensee") accessing and otherwise using +matplotlib software in source or binary form and its associated +documentation. + +2. Subject to the terms and conditions of this License Agreement, JDH +hereby grants Licensee a nonexclusive, royalty-free, world-wide license +to reproduce, analyze, test, perform and/or display publicly, prepare +derivative works, distribute, and otherwise use matplotlib +alone or in any derivative version, provided, however, that JDH's +License Agreement and JDH's notice of copyright, i.e., "Copyright (c) +2002-2011 John D. Hunter; All Rights Reserved" are retained in +matplotlib alone or in any derivative version prepared by +Licensee. + +3. In the event Licensee prepares a derivative work that is based on or +incorporates matplotlib or any part thereof, and wants to +make the derivative work available to others as provided herein, then +Licensee hereby agrees to include in any such work a brief summary of +the changes made to matplotlib. + +4. JDH is making matplotlib available to Licensee on an "AS +IS" basis. JDH MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, JDH MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF MATPLOTLIB +WILL NOT INFRINGE ANY THIRD PARTY RIGHTS. + +5. JDH SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF MATPLOTLIB + FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR +LOSS AS A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING +MATPLOTLIB , OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF +THE POSSIBILITY THEREOF. + +6. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +7. Nothing in this License Agreement shall be deemed to create any +relationship of agency, partnership, or joint venture between JDH and +Licensee. This License Agreement does not grant permission to use JDH +trademarks or trade name in a trademark sense to endorse or promote +products or services of Licensee, or any third party. + +8. By copying, installing or otherwise using matplotlib, +Licensee agrees to be bound by the terms and conditions of this License +Agreement. + + +From tqdm: + +`tqdm` is a product of collaborative work. +Unless otherwise stated, all authors (see commit logs) retain copyright +for their respective work, and release the work under the MIT licence +(text below). + +Exceptions or notable authors are listed below +in reverse chronological order: + +* files: * + MPL-2.0 2015-2024 (c) Casper da Costa-Luis + [casperdcl](https://github.com/casperdcl). +* files: tqdm/_tqdm.py + MIT 2016 (c) [PR #96] on behalf of Google Inc. +* files: tqdm/_tqdm.py README.rst .gitignore + MIT 2013 (c) Noam Yorav-Raphael, original author. + +[PR #96]: https://github.com/tqdm/tqdm/pull/96 + + +Mozilla Public Licence (MPL) v. 2.0 - Exhibit A +----------------------------------------------- + +This Source Code Form is subject to the terms of the +Mozilla Public License, v. 2.0. +If a copy of the MPL was not distributed with this project, +You can obtain one at https://mozilla.org/MPL/2.0/. + + +MIT License (MIT) +----------------- + +Copyright (c) 2013 noamraph + +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. + + +From scikit-image: + +Copyright: 2009-2022 the scikit-image team +License: BSD-3-Clause + +Files: doc/source/themes/scikit-image/layout.html +Copyright: 2007-2010 the Sphinx team +License: BSD-3-Clause + +Files: skimage/feature/_canny.py + skimage/filters/edges.py + skimage/filters/_rank_order.py + skimage/morphology/_skeletonize.py + skimage/morphology/tests/test_watershed.py + skimage/morphology/watershed.py + skimage/segmentation/heap_general.pxi + skimage/segmentation/heap_watershed.pxi + skimage/segmentation/_watershed.py + skimage/segmentation/_watershed_cy.pyx +Copyright: 2003-2009 Massachusetts Institute of Technology + 2009-2011 Broad Institute + 2003 Lee Kamentsky + 2003-2005 Peter J. Verveer +License: BSD-3-Clause + +Files: skimage/filters/thresholding.py + skimage/graph/_mcp.pyx + skimage/graph/heap.pyx +Copyright: 2009-2015 Board of Regents of the University of + Wisconsin-Madison, Broad Institute of MIT and Harvard, + and Max Planck Institute of Molecular Cell Biology and + Genetics + 2009 Zachary Pincus + 2009 Almar Klein +License: BSD-2-Clause + +File: skimage/morphology/grayreconstruct.py + skimage/morphology/tests/test_reconstruction.py +Copyright: 2003-2009 Massachusetts Institute of Technology + 2009-2011 Broad Institute + 2003 Lee Kamentsky +License: BSD-3-Clause + +File: skimage/morphology/_grayreconstruct.pyx +Copyright: 2003-2009 Massachusetts Institute of Technology + 2009-2011 Broad Institute + 2003 Lee Kamentsky + 2022 Gregory Lee (added a 64-bit integer variant for large images) +License: BSD-3-Clause + +File: skimage/segmentation/_expand_labels.py +Copyright: 2020 Broad Institute + 2020 CellProfiler team +License: BSD-3-Clause + +File: skimage/exposure/_adapthist.py +Copyright: 1994 Karel Zuiderveld +License: BSD-3-Clause + +Function: skimage/morphology/_skeletonize_cy.pyx:_skeletonize_loop +Copyright: 2003-2009 Massachusetts Institute of Technology + 2009-2011 Broad Institute + 2003 Lee Kamentsky +License: BSD-3-Clause + +Function: skimage/_shared/version_requirements.py:_check_version +Copyright: 2013 The IPython Development Team +License: BSD-3-Clause + +Function: skimage/_shared/version_requirements.py:is_installed +Copyright: 2009-2011 Pierre Raybaut +License: MIT + +File: skimage/feature/_fisher_vector.py +Copyright: 2014 2014 Dan Oneata +License: MIT + +File: skimage/_vendored/numpy_lookfor.py +Copyright: 2005-2023, NumPy Developers +License: BSD-3-Clause + + +License: BSD-2-Clause + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions +are met: +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. +2. 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. +. +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 HOLDERS 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 OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +License: BSD-3-Clause + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions +are met: +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. +2. 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. +3. Neither the name of the University 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 HOLDERS 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 OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +License: MIT + +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. + + +From opencv-python: + +MIT License + +Copyright (c) Olli-Pekka Heinisuo + +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. + + +From numpy: + +Copyright (c) 2005-2024, NumPy Developers. +All rights reserved. + +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 NumPy Developers nor the names of any + 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 +OWNER 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 +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +From python-box: + +MIT License + +Copyright (c) 2017-2023 Chris Griffith + +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. + + +From Imutils: + +The MIT License (MIT) + +Copyright (c) 2015-2016 Adrian Rosebrock, http://www.pyimagesearch.com + +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. + + +From Pillow: + +The Python Imaging Library (PIL) is + + Copyright © 1997-2011 by Secret Labs AB + Copyright © 1995-2011 by Fredrik Lundh + +Pillow is the friendly PIL fork. It is + + Copyright © 2010-2024 by Jeffrey A. Clark (Alex) and contributors. + +Like PIL, Pillow is licensed under the open source HPND License: + +By obtaining, using, and/or copying this software and/or its associated +documentation, you agree that you have read, understood, and will comply +with the following terms and conditions: + +Permission to use, copy, modify and distribute this software and its +documentation for any purpose and without fee is hereby granted, +provided that the above copyright notice appears in all copies, and that +both that copyright notice and this permission notice appear in supporting +documentation, and that the name of Secret Labs AB or the author not be +used in advertising or publicity pertaining to distribution of the software +without specific, written prior permission. + +SECRET LABS AB AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS +SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. +IN NO EVENT SHALL SECRET LABS AB OR THE AUTHOR BE LIABLE FOR ANY SPECIAL, +INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE +OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +PERFORMANCE OF THIS SOFTWARE. + + +From PyYAML: + +Copyright (c) 2017-2021 Ingy döt Net +Copyright (c) 2006-2016 Kirill Simonov + +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. + + +From scipy: + +Copyright (c) 2001-2002 Enthought, Inc. 2003-2024, SciPy Developers. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions +are met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. 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. + +3. 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 +OWNER 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 +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +From scikit-learn: + +BSD 3-Clause License + +Copyright (c) 2007-2023 The scikit-learn developers. +All rights reserved. + +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 +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +From simplejson: + +simplejson is dual-licensed software. It is available under the terms +of the MIT license, or the Academic Free License version 2.1. The full +text of each license agreement is included below. This code is also +licensed to the Python Software Foundation (PSF) under a Contributor +Agreement. + +MIT License +=========== + +Copyright (c) 2006 Bob Ippolito + +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. + +Academic Free License v. 2.1 +============================ + +Copyright (c) 2006 Bob Ippolito. All rights reserved. + +This Academic Free License (the "License") applies to any original work of authorship (the "Original Work") whose owner (the "Licensor") has placed the following notice immediately following the copyright notice for the Original Work: + +Licensed under the Academic Free License version 2.1 + +1) Grant of Copyright License. Licensor hereby grants You a world-wide, royalty-free, non-exclusive, perpetual, sublicenseable license to do the following: + +a) to reproduce the Original Work in copies; + +b) to prepare derivative works ("Derivative Works") based upon the Original Work; + +c) to distribute copies of the Original Work and Derivative Works to the public; + +d) to perform the Original Work publicly; and + +e) to display the Original Work publicly. + +2) Grant of Patent License. Licensor hereby grants You a world-wide, royalty-free, non-exclusive, perpetual, sublicenseable license, under patent claims owned or controlled by the Licensor that are embodied in the Original Work as furnished by the Licensor, to make, use, sell and offer for sale the Original Work and Derivative Works. + +3) Grant of Source Code License. The term "Source Code" means the preferred form of the Original Work for making modifications to it and all available documentation describing how to modify the Original Work. Licensor hereby agrees to provide a machine-readable copy of the Source Code of the Original Work along with each copy of the Original Work that Licensor distributes. Licensor reserves the right to satisfy this obligation by placing a machine-readable copy of the Source Code in an information repository reasonably calculated to permit inexpensive and convenient access by You for as long as Licensor continues to distribute the Original Work, and by publishing the address of that information repository in a notice immediately following the copyright notice that applies to the Original Work. + +4) Exclusions From License Grant. Neither the names of Licensor, nor the names of any contributors to the Original Work, nor any of their trademarks or service marks, may be used to endorse or promote products derived from this Original Work without express prior written permission of the Licensor. Nothing in this License shall be deemed to grant any rights to trademarks, copyrights, patents, trade secrets or any other intellectual property of Licensor except as expressly stated herein. No patent license is granted to make, use, sell or offer to sell embodiments of any patent claims other than the licensed claims defined in Section 2. No right is granted to the trademarks of Licensor even if such marks are included in the Original Work. Nothing in this License shall be interpreted to prohibit Licensor from licensing under different terms from this License any Original Work that Licensor otherwise would have a right to license. + +5) This section intentionally omitted. + +6) Attribution Rights. You must retain, in the Source Code of any Derivative Works that You create, all copyright, patent or trademark notices from the Source Code of the Original Work, as well as any notices of licensing and any descriptive text identified therein as an "Attribution Notice." You must cause the Source Code for any Derivative Works that You create to carry a prominent Attribution Notice reasonably calculated to inform recipients that You have modified the Original Work. + +7) Warranty of Provenance and Disclaimer of Warranty. Licensor warrants that the copyright in and to the Original Work and the patent rights granted herein by Licensor are owned by the Licensor or are sublicensed to You under the terms of this License with the permission of the contributor(s) of those copyrights and patent rights. Except as expressly stated in the immediately proceeding sentence, the Original Work is provided under this License on an "AS IS" BASIS and WITHOUT WARRANTY, either express or implied, including, without limitation, the warranties of NON-INFRINGEMENT, MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. THE ENTIRE RISK AS TO THE QUALITY OF THE ORIGINAL WORK IS WITH YOU. This DISCLAIMER OF WARRANTY constitutes an essential part of this License. No license to Original Work is granted hereunder except under this disclaimer. + +8) Limitation of Liability. Under no circumstances and under no legal theory, whether in tort (including negligence), contract, or otherwise, shall the Licensor be liable to any person for any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or the use of the Original Work including, without limitation, damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses. This limitation of liability shall not apply to liability for death or personal injury resulting from Licensor's negligence to the extent applicable law prohibits such limitation. Some jurisdictions do not allow the exclusion or limitation of incidental or consequential damages, so this exclusion and limitation may not apply to You. + +9) Acceptance and Termination. If You distribute copies of the Original Work or a Derivative Work, You must make a reasonable effort under the circumstances to obtain the express assent of recipients to the terms of this License. Nothing else but this License (or another written agreement between Licensor and You) grants You permission to create Derivative Works based upon the Original Work or to exercise any of the rights granted in Section 1 herein, and any attempt to do so except under the terms of this License (or another written agreement between Licensor and You) is expressly prohibited by U.S. copyright law, the equivalent laws of other countries, and by international treaty. Therefore, by exercising any of the rights granted to You in Section 1 herein, You indicate Your acceptance of this License and all of its terms and conditions. + +10) Termination for Patent Action. This License shall terminate automatically and You may no longer exercise any of the rights granted to You by this License as of the date You commence an action, including a cross-claim or counterclaim, against Licensor or any licensee alleging that the Original Work infringes a patent. This termination provision shall not apply for an action alleging patent infringement by combinations of the Original Work with other software or hardware. + +11) Jurisdiction, Venue and Governing Law. Any action or suit relating to this License may be brought only in the courts of a jurisdiction wherein the Licensor resides or in which Licensor conducts its primary business, and under the laws of that jurisdiction excluding its conflict-of-law provisions. The application of the United Nations Convention on Contracts for the International Sale of Goods is expressly excluded. Any use of the Original Work outside the scope of this License or after its termination shall be subject to the requirements and penalties of the U.S. Copyright Act, 17 U.S.C. § 101 et seq., the equivalent laws of other countries, and international treaty. This section shall survive the termination of this License. + +12) Attorneys Fees. In any action to enforce the terms of this License or seeking damages relating thereto, the prevailing party shall be entitled to recover its costs and expenses, including, without limitation, reasonable attorneys' fees and costs incurred in connection with such action, including any appeal of such action. This section shall survive the termination of this License. + +13) Miscellaneous. This License represents the complete agreement concerning the subject matter hereof. If any provision of this License is held to be unenforceable, such provision shall be reformed only to the extent necessary to make it enforceable. + +14) Definition of "You" in This License. "You" throughout this License, whether in upper or lower case, means an individual or a legal entity exercising rights under, and complying with all of the terms of, this License. For legal entities, "You" includes any entity that controls, is controlled by, or is under common control with you. For purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. + +15) Right to Use. You may use the Original Work in all ways not otherwise restricted or conditioned by this License or by law, and Licensor promises not to interfere with or be responsible for such uses by You. + +This license is Copyright (C) 2003-2004 Lawrence E. Rosen. All rights reserved. Permission is hereby granted to copy and distribute this license without modification. This license may not be modified without the express written permission of its copyright owner. + + +From pandas: + +SD 3-Clause License + +Copyright (c) 2008-2011, AQR Capital Management, LLC, Lambda Foundry, Inc. and PyData Development Team +All rights reserved. + +Copyright (c) 2011-2024, Open source contributors. + +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 +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +From shape_predictor_81_face_landmarks: + +BSD 3-Clause License + +Copyright (c) 2019, Nikolay Feldman +All rights reserved. + +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 +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +From retinaface: + +MIT License + +Copyright (c) 2020 Vladimir Iglovikov + +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. + + +From EfficientNet-PyTorch: + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. + + +From albumentations: + +MIT License + +Copyright (c) 2017 Buslaev Alexander, Alexander Parinov, Vladimir Iglovikov + +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. + + +From SAM: + +MIT License + +Copyright (c) 2021 David Samuel + +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. + + +From natsort: + +Copyright (c) 2012-2023 Seth M. Morton + +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. + + +************END OF THIRD-PARTY SOFTWARE NOTICES AND INFORMATION********** diff --git a/video/fake-stormer/model_code/configs/base.yaml b/video/fake-stormer/model_code/configs/base.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1faca4d31ec18af2bc8a153a05f024c8f59e2cbb --- /dev/null +++ b/video/fake-stormer/model_code/configs/base.yaml @@ -0,0 +1,79 @@ +TASK: heatmap +PRECISION: float32 +DATASET: + type: HeatmapFaceForensic + TRAIN: True + ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c23/ + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + FAKETYPE: [FaceXRay] + IMAGE_SUFFIX: jpg + FROM_FILE: True + NUM_WORKERS: 8 + PIN_MEMORY: True + IMAGE_SIZE: [256, 256] + HEATMAP_SIZE: [64, 64] + SIGMA: 2 + HEATMAP_TYPE: gaussian + DATA: + TYPE: images + TRAIN: + ANNO_FILE: FaceXRay/train/train_FF_Xray.json + VAL: + ANNO_FILE: FaceXRay/val/val_FF_Xray.json + TEST: + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + TRANSFORM: + geometry: + type: GeometryTransform + resize: [256, 256, 0] #h, w, p=probability. If no affine transform, set p=1 + color: + type: ColorJitterTransform + clahe: 0.5 + colorjitter: 0.5 + gaussianblur: 0.5 + jpegcompression: 0.5 + rgbshift: 0.5 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False +MODEL: + type: SimpleClassificationDF + backbone: + type: ResNet + num_layers: 50 + drop_ratio: 0.5 + mode: ir_se + head: + type: SimpleClassificationHead + drop_ratio: 0.5 + in_planes: 512 +TRAIN: + gpus: [0,1,2,3] + batch_size: 32 + lr: 0.001 + epochs: 100 + begin_epoch: 0 + warm_up: 5 + every_val_epochs: 3 + loss: + type: CombinedLoss + use_target_weight: False + optimizer: Adam + distributed: False + pretrained: pretrained/model_ir_se50.pth + tensorboard: True + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [50, 80, 90] + gamma: 0.1 +PREPROCESSING: + DATASET: FaceForensics + SPLIT: train + ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c23/ + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + IMAGE_SUFFIX: jpg + DATA_TYPE: images + LABEL: real + facial_lm_pretrained: pretrained/shape_predictor_68_face_landmarks.dat diff --git a/video/fake-stormer/model_code/configs/get_config.py b/video/fake-stormer/model_code/configs/get_config.py new file mode 100644 index 0000000000000000000000000000000000000000..e47c3393fdbabb52d5063bda7b5604fe68a27ae1 --- /dev/null +++ b/video/fake-stormer/model_code/configs/get_config.py @@ -0,0 +1,19 @@ +# -*- coding: utf-8 -*- +import os + +from yaml import dump, load + +try: + from yaml import CDumper as Dumper + from yaml import CLoader as Loader +except ImportError: + from yaml import Loader, Dumper + +from box import Box as edict + + +def load_config(cfg): + with open(cfg) as f: + config = load(f, Loader=Loader) + + return edict(config) diff --git a/video/fake-stormer/model_code/configs/preprocessing/data_preprocessing_c0.yaml b/video/fake-stormer/model_code/configs/preprocessing/data_preprocessing_c0.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9f499a523d4212dfd294786f0d1275c70ddc729a --- /dev/null +++ b/video/fake-stormer/model_code/configs/preprocessing/data_preprocessing_c0.yaml @@ -0,0 +1,16 @@ +PREPROCESSING: + DATASET: FF++ + COMPRESSION: c0 + SPLIT: val + ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + FAKETYPE: [original] + IMAGE_SUFFIX: png + DATA_TYPE: videos + LABEL: [real] + facial_lm_pretrained: pretrained/shape_predictor_81_face_landmarks.dat + N_LANDMARKS: 81 + SAMPLING: + ACTIVE: False + NUMBERS: 8 + DEBUG: False diff --git a/video/fake-stormer/model_code/configs/preprocessing/data_preprocessing_c23.yaml b/video/fake-stormer/model_code/configs/preprocessing/data_preprocessing_c23.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d071da68974c657dc13b640e72de68c4176a089b --- /dev/null +++ b/video/fake-stormer/model_code/configs/preprocessing/data_preprocessing_c23.yaml @@ -0,0 +1,16 @@ +PREPROCESSING: + DATASET: FaceForensics + COMPRESSION: c23 + SPLIT: train + ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c23/ + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + FAKETYPE: [original] + IMAGE_SUFFIX: png + DATA_TYPE: videos + LABEL: [real] + facial_lm_pretrained: pretrained/shape_predictor_81_face_landmarks.dat + N_LANDMARKS: 81 + SAMPLING: + ACTIVE: False + NUMBERS: 8 + DEBUG: False diff --git a/video/fake-stormer/model_code/configs/spatial/binary_cls/efns/efn_4.yaml b/video/fake-stormer/model_code/configs/spatial/binary_cls/efns/efn_4.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0dc7d27a39fbd9b8a5fd20c8dfc6be9313d578ce --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/binary_cls/efns/efn_4.yaml @@ -0,0 +1,160 @@ +TASK: EFNB4_BCE_Adam_5e4_Batch32_50epochs_abl_FS +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: False + +DATASET: + type: BinaryFaceForensic + DATA_TYPE: image + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 7 + PIN_MEMORY: True + IMAGE_SIZE: [256, 256] + HEATMAP_SIZE: [28, 28] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: False + TRAIN: 8 # Dynamically random number of frames in each epoch + VAL: 8 + TEST: 128 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/DFW/ + FROM_FILE: False + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # FAKETYPE: [new_fake_train, new_real_train] + # FAKETYPE: [original, NeuralTextures] + # FAKETYPE: [original, Face2Face] + # FAKETYPE: [original, Deepfakes] + FAKETYPE: [original, FaceSwap] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/train_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/DFW/ + FROM_FILE: False + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [original, NeuralTextures] + # FAKETYPE: [original, Face2Face] + # FAKETYPE: [original, Deepfakes] + FAKETYPE: [original, FaceSwap] + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/val_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: Celeb-DFv1 + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [256, 256, 1] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.5, 0.5, 0.5] + std: [0.5, 0.5, 0.5] + DEBUG: False + +MODEL: + type: PoseEfficientNet + model_name: efficientnet-b4 + num_layers: B4 + include_top: True + num_classes: 1 + include_hm_decoder: False + INIT_WEIGHTS: + pretrained: True + advprop: True + +TRAIN: + gpus: [0] + batch_size: 32 + lr: 0.0005 + epochs: 50 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + loss: + type: BinaryCrossEntropy + reduction: mean + optimizer: Adam + distributed: False + pretrained: '' + tensorboard: True + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: False + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 4 + +TEST: + gpus: [0,1,2,3] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: True + video_level: True + pretrained: 'logs/10-01-2024/TopDownDetector_ViTSmall112_hm10_8SBI_NonOverlap100_MSE_AdamW_IN_5e5_FZ5_Batch32_200epochs_Drop0.2_AVG_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/spatial/binary_cls/vits/vit_small_112p16.yaml b/video/fake-stormer/model_code/configs/spatial/binary_cls/vits/vit_small_112p16.yaml new file mode 100644 index 0000000000000000000000000000000000000000..19d65db341d1dcc7b02a5bb36472412692b1dea7 --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/binary_cls/vits/vit_small_112p16.yaml @@ -0,0 +1,180 @@ +TASK: ViTSmall112_BCE_AdamW_IN_5e5_Batch32_50epochs_Drop0.2_abl_p16_FS +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: False + +DATASET: + type: BinaryFaceForensic + DATA_TYPE: image + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 7 + PIN_MEMORY: True + IMAGE_SIZE: [112, 112] + HEATMAP_SIZE: [7, 7] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: False + TRAIN: 8 # Dynamically random number of frames in each epoch + VAL: 8 + TEST: 128 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/DFW/ + FROM_FILE: False + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # FAKETYPE: [new_fake_train, new_real_train] + # FAKETYPE: [original, NeuralTextures] + # FAKETYPE: [original, Face2Face] + # FAKETYPE: [original, Deepfakes] + FAKETYPE: [original, FaceSwap] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/train_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/DFW/ + FROM_FILE: False + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [original, NeuralTextures] + # FAKETYPE: [original, Face2Face] + # FAKETYPE: [original, Deepfakes] + FAKETYPE: [original, FaceSwap] + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/val_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: Celeb-DFv1 + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [112, 112, 1] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + +MODEL: + type: TopDownDetector + backbone: + type: ViT + img_size: [112, 112] + patch_size: 16 + embed_dim: 384 + depth: 12 + num_heads: 6 + mlp_ratio: 4 + drop_path_rate: 0.2 + qkv_bias: True + class_token: True + pretrained: pretrained/dino_deitsmall16_pretrain.pth + # pretrained: null + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 384 + heads: + # hm: 1 + cls: 1 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256, 256] + num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + +TRAIN: + gpus: [0] + batch_size: 32 + lr: 0.00005 + epochs: 50 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + loss: + type: BinaryCrossEntropy + reduction: mean + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: True + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: False + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 4 + +TEST: + gpus: [0,1,2,3] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: True + video_level: True + pretrained: 'logs/10-01-2024/TopDownDetector_ViTSmall112_hm10_8SBI_NonOverlap100_MSE_AdamW_IN_5e5_FZ5_Batch32_200epochs_Drop0.2_AVG_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/spatial/binary_cls/vits/vit_small_112p8.yaml b/video/fake-stormer/model_code/configs/spatial/binary_cls/vits/vit_small_112p8.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7e716fe1c0991f206d66075c6137d590d3a00292 --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/binary_cls/vits/vit_small_112p8.yaml @@ -0,0 +1,180 @@ +TASK: ViTSmall112_BCE_AdamW_IN_5e5_Batch32_50epochs_Drop0.2_abl_p8_FS +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: False + +DATASET: + type: BinaryFaceForensic + DATA_TYPE: image + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 7 + PIN_MEMORY: True + IMAGE_SIZE: [112, 112] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: False + TRAIN: 8 # Dynamically random number of frames in each epoch + VAL: 8 + TEST: 128 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/DFW/ + FROM_FILE: False + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # FAKETYPE: [new_fake_train, new_real_train] + # FAKETYPE: [original, NeuralTextures] + # FAKETYPE: [original, Face2Face] + # FAKETYPE: [original, Deepfakes] + FAKETYPE: [original, FaceSwap] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/train_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/DFW/ + FROM_FILE: False + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [original, NeuralTextures] + # FAKETYPE: [original, Face2Face] + # FAKETYPE: [original, Deepfakes] + FAKETYPE: [original, FaceSwap] + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/val_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: Celeb-DFv1 + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [112, 112, 1] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + +MODEL: + type: TopDownDetector + backbone: + type: ViT + img_size: [112, 112] + patch_size: 8 + embed_dim: 384 + depth: 12 + num_heads: 6 + mlp_ratio: 4 + drop_path_rate: 0.2 + qkv_bias: True + class_token: True + pretrained: pretrained/dino_deitsmall8_pretrain.pth + # pretrained: null + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 384 + heads: + # hm: 1 + cls: 1 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256, 256] + num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + +TRAIN: + gpus: [0] + batch_size: 32 + lr: 0.00005 + epochs: 50 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + loss: + type: BinaryCrossEntropy + reduction: mean + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: True + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: False + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 4 + +TEST: + gpus: [0,1,2,3] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: True + video_level: True + pretrained: 'logs/10-01-2024/TopDownDetector_ViTSmall112_hm10_8SBI_NonOverlap100_MSE_AdamW_IN_5e5_FZ5_Batch32_200epochs_Drop0.2_AVG_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/spatial/binary_cls/vits/vit_small_224p8.yaml b/video/fake-stormer/model_code/configs/spatial/binary_cls/vits/vit_small_224p8.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7e716fe1c0991f206d66075c6137d590d3a00292 --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/binary_cls/vits/vit_small_224p8.yaml @@ -0,0 +1,180 @@ +TASK: ViTSmall112_BCE_AdamW_IN_5e5_Batch32_50epochs_Drop0.2_abl_p8_FS +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: False + +DATASET: + type: BinaryFaceForensic + DATA_TYPE: image + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 7 + PIN_MEMORY: True + IMAGE_SIZE: [112, 112] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: False + TRAIN: 8 # Dynamically random number of frames in each epoch + VAL: 8 + TEST: 128 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/DFW/ + FROM_FILE: False + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # FAKETYPE: [new_fake_train, new_real_train] + # FAKETYPE: [original, NeuralTextures] + # FAKETYPE: [original, Face2Face] + # FAKETYPE: [original, Deepfakes] + FAKETYPE: [original, FaceSwap] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/train_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/DFW/ + FROM_FILE: False + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [original, NeuralTextures] + # FAKETYPE: [original, Face2Face] + # FAKETYPE: [original, Deepfakes] + FAKETYPE: [original, FaceSwap] + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/val_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: Celeb-DFv1 + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [112, 112, 1] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + +MODEL: + type: TopDownDetector + backbone: + type: ViT + img_size: [112, 112] + patch_size: 8 + embed_dim: 384 + depth: 12 + num_heads: 6 + mlp_ratio: 4 + drop_path_rate: 0.2 + qkv_bias: True + class_token: True + pretrained: pretrained/dino_deitsmall8_pretrain.pth + # pretrained: null + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 384 + heads: + # hm: 1 + cls: 1 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256, 256] + num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + +TRAIN: + gpus: [0] + batch_size: 32 + lr: 0.00005 + epochs: 50 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + loss: + type: BinaryCrossEntropy + reduction: mean + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: True + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: False + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 4 + +TEST: + gpus: [0,1,2,3] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: True + video_level: True + pretrained: 'logs/10-01-2024/TopDownDetector_ViTSmall112_hm10_8SBI_NonOverlap100_MSE_AdamW_IN_5e5_FZ5_Batch32_200epochs_Drop0.2_AVG_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/spatial/binary_cls/xcepts/xcept.yaml b/video/fake-stormer/model_code/configs/spatial/binary_cls/xcepts/xcept.yaml new file mode 100644 index 0000000000000000000000000000000000000000..011398f572d211804271b290757d5d63552dd4a5 --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/binary_cls/xcepts/xcept.yaml @@ -0,0 +1,156 @@ +TASK: Exeption_BCE_Adam_5e4_Batch32_50epochs_abl_F2F +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: False + +DATASET: + type: BinaryFaceForensic + DATA_TYPE: image + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 14 + PIN_MEMORY: True + IMAGE_SIZE: [256, 256] + HEATMAP_SIZE: [28, 28] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: False + TRAIN: 8 # Dynamically random number of frames in each epoch + VAL: 8 + TEST: 128 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/DFW/ + FROM_FILE: False + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # FAKETYPE: [new_fake_train, new_real_train] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [original, NeuralTextures] + FAKETYPE: [original, Face2Face] + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, FaceSwap] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/train_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/DFW/ + FROM_FILE: False + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [original, NeuralTextures] + FAKETYPE: [original, Face2Face] + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, FaceSwap] + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/val_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: Celeb-DFv1 + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [256, 256, 1] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.5, 0.5, 0.5] + std: [0.5, 0.5, 0.5] + DEBUG: False + +MODEL: + type: Xception + num_classes: 1 + INIT_WEIGHTS: + pretrained: True + +TRAIN: + gpus: [0,1] + batch_size: 16 + lr: 0.0005 + epochs: 50 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + loss: + type: BinaryCrossEntropy + reduction: mean + optimizer: Adam + distributed: False + pretrained: '' + tensorboard: True + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: False + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 4 + +TEST: + gpus: [0,1,2,3] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: True + video_level: True + pretrained: 'logs/10-01-2024/TopDownDetector_ViTSmall112_hm10_8SBI_NonOverlap100_MSE_AdamW_IN_5e5_FZ5_Batch32_200epochs_Drop0.2_AVG_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/spatial/efn4_fpn_hm.yaml b/video/fake-stormer/model_code/configs/spatial/efn4_fpn_hm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..675d428cee3a016963a136954fb1496a3ebd13e5 --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/efn4_fpn_hm.yaml @@ -0,0 +1,147 @@ +TASK: EFN_hm100_FPN_NoBasedCLS_Focal_C3_256Cstency10_32FXRayv1_SAM(Adam) +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True +DATASET: + type: HeatmapFaceForensic + NAME: Celeb-DFv1 # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + TRAIN: False #Switch to True for training mode, False for testing mode + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + COMPRESSION: c0 + IMAGE_SUFFIX: png + FROM_FILE: False + NUM_WORKERS: 7 + PIN_MEMORY: True + IMAGE_SIZE: [384, 384] + HEATMAP_SIZE: [96, 96] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 2 + ADAPTIVE_SIGMA: True + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 32 + VAL: 32 + TEST: 32 + TRAIN: + FAKETYPE: [FaceXRay] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/dynamic_trainBI_FFv2.json + LABEL_FOLDER: [real, fake] + VAL: + FAKETYPE: [FaceXRay] + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/dynamic_valBI_FFv2.json + LABEL_FOLDER: [real, fake] + TEST: + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [384, 384, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.5, 3] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True +MODEL: + type: PoseEfficientNet + model_name: efficientnet-b4 + num_layers: B4 + include_top: False + include_hm_decoder: True + head_conv: 64 + use_c2: False + use_c3: True + use_c4: True + use_c51: True + fpn: True + heads: + hm: 1 + cls: 1 + cstency: 256 + INIT_WEIGHTS: + pretrained: True +TRAIN: + gpus: [0] + batch_size: 32 + lr: 0.00025 + epochs: 100 + begin_epoch: 0 + warm_up: 6 + every_val_epochs: 1 + loss: + type: CombinedFocalLoss + use_target_weight: False + cls_lmda: 1 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 100 + cstency_lmda: 10 + mse_reduction: sum + ce_reduction: mean + optimizer: SAM + distributed: False + pretrained: '' + tensorboard: True + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: False + pretrained: 'logs/20-02-2023/PoseEfficientNet_EFN_hm100_FPN_Based_CLS_Focal_NoC2_256Cstency10_32FXRayv2_SAM(Adam)_model_best.pth' +PREPROCESSING: + DATASET: FaceForensics + SPLIT: train + ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c23/ + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + IMAGE_SUFFIX: jpg + DATA_TYPE: images + LABEL: [real, fake] + facial_lm_pretrained: pretrained/shape_predictor_68_face_landmarks.dat + DEBUG: False diff --git a/video/fake-stormer/model_code/configs/spatial/efn4_fpn_hm_adv.yaml b/video/fake-stormer/model_code/configs/spatial/efn4_fpn_hm_adv.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c97167c6892b2d69b6d5e13c7ac501bef4f958ac --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/efn4_fpn_hm_adv.yaml @@ -0,0 +1,165 @@ +TASK: EFN_hm10_EFPN_NoBasedCLS_Focal_C3_256Cst100_32FXRayv2_SAM(Adam)_ADV_Era1_OutSigmoid_1e7_boost500_UnFZ +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True +DATASET: + type: HeatmapFaceForensic + DATA_TYPE: image + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 28 + PIN_MEMORY: True + IMAGE_SIZE: [384, 384] + HEATMAP_SIZE: [96, 96] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 2 + ADAPTIVE_SIGMA: True + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 32 # Dynamically random number of frames in each epoch + VAL: 32 + TEST: 32 + DIST: [1., 1.] # Distribution of Real, Fake - 1.,1. by default + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/dynamic_trainBI_FFv2.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/dynamic_valBI_FFv2.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: FF++ + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # FAKETYPE: [frames_BW_5/original, frames_BW_5/Deepfakes, frames_BW_5/Face2Face, frames_BW_5/FaceSwap, frames_BW_5/NeuralTextures] + FAKETYPE: [frames_JPEG_1/original, frames_JPEG_1/Deepfakes, frames_JPEG_1/Face2Face, frames_JPEG_1/FaceSwap, frames_JPEG_1/NeuralTextures, + frames_JPEG_2/original, frames_JPEG_2/Deepfakes, frames_JPEG_2/Face2Face, frames_JPEG_2/FaceSwap, frames_JPEG_2/NeuralTextures, + frames_JPEG_3/original, frames_JPEG_3/Deepfakes, frames_JPEG_3/Face2Face, frames_JPEG_3/FaceSwap, frames_JPEG_3/NeuralTextures, + frames_JPEG_4/original, frames_JPEG_4/Deepfakes, frames_JPEG_4/Face2Face, frames_JPEG_4/FaceSwap, frames_JPEG_4/NeuralTextures, + frames_JPEG_5/original, frames_JPEG_5/Deepfakes, frames_JPEG_5/Face2Face, frames_JPEG_5/FaceSwap, frames_JPEG_5/NeuralTextures] + # FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [real_test, fake_test] + # ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [384, 384, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.5, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.5, 0.5, 0.5] + std: [0.5, 0.5, 0.5] + DEBUG: False + DYNAMIC_FXRAY: True +MODEL: + type: PoseEfficientNet + model_name: efficientnet-b4 + num_layers: B4 + include_top: False + include_hm_decoder: True + head_conv: 64 + use_c2: False + use_c3: True + use_c4: True + use_c51: True + efpn: True + tfpn: False + se_layer: False + norm_c2: False + heads: + hm: 1 + cls: 1 + cstency: 256 + INIT_WEIGHTS: + pretrained: True + advprop: True +TRAIN: + gpus: [0] + batch_size: 16 + lr: 0.0000001 + epochs: 100 + begin_epoch: -1 + warm_up: 6 + every_val_epochs: 1 + loss: + type: CombinedFocalLoss + use_target_weight: False + cls_lmda: 1 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 10 + cstency_lmda: 100 + mse_reduction: sum + ce_reduction: mean + optimizer: SAM + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: False + debug: + active: False + save_hm_gt: True + save_hm_pred: True +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: True + video_level: True + apr: True # Average precision/recall or normal precision/recall + no_shot_preds: 1 + pred_file: /project/home/p200249/XXX/saved_predictions/laanet_wBI_DFDC_100_100.json # File to save predictions + save_preds: False + pretrained: 'pretrained/PoseEfficientNet_EFN_hm100_EFPN_NoBasedCLS_Focal_C3_256Cst100_8FXRayv2_SAM(Adam)_ADV_Era1_OutSigmoid_5e5_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/spatial/efn4_fpn_sbi_adv.yaml b/video/fake-stormer/model_code/configs/spatial/efn4_fpn_sbi_adv.yaml new file mode 100644 index 0000000000000000000000000000000000000000..06446a290a2f6c9832faf7dde32eb31cde2afc8e --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/efn4_fpn_sbi_adv.yaml @@ -0,0 +1,180 @@ +TASK: EFN_hm10_EFPN_NoBasedCLS_Focal_C3_256Cst100_8SBI_SAM(Adam)_ADV_Era1_OutSigmoid_1e7_boost500_UnFZ +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True +DATASET: + type: SBIFaceForensic + DATA_TYPE: image + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 7 + PIN_MEMORY: True + IMAGE_SIZE: [384, 384] + HEATMAP_SIZE: [96, 96] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 2 + ADAPTIVE_SIGMA: True + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: False + TRAIN: 8 # Dynamically random number of frames in each epoch + VAL: 8 + TEST: 32 + DIST: [1., 1.] # Distribution of Real, Fake - 1.,1. by default + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/train_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/val_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: FF++ + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + # ROOT: /project/home/p200249/XXX/DiffSwap/ + # ROOT: /project/home/p200249/XXX/DF40_test/ + FROM_FILE: False + # FAKETYPE: [original, Face2Face] + # FAKETYPE: [frames_JPEG_1/original, frames_JPEG_1/Deepfakes, frames_JPEG_1/Face2Face, frames_JPEG_1/FaceSwap, frames_JPEG_1/NeuralTextures] + FAKETYPE: [frames_JPEG_1/original, frames_JPEG_1/Deepfakes, frames_JPEG_1/Face2Face, frames_JPEG_1/FaceSwap, frames_JPEG_1/NeuralTextures, + frames_JPEG_2/original, frames_JPEG_2/Deepfakes, frames_JPEG_2/Face2Face, frames_JPEG_2/FaceSwap, frames_JPEG_2/NeuralTextures, + frames_JPEG_3/original, frames_JPEG_3/Deepfakes, frames_JPEG_3/Face2Face, frames_JPEG_3/FaceSwap, frames_JPEG_3/NeuralTextures, + frames_JPEG_4/original, frames_JPEG_4/Deepfakes, frames_JPEG_4/Face2Face, frames_JPEG_4/FaceSwap, frames_JPEG_4/NeuralTextures, + frames_JPEG_5/original, frames_JPEG_5/Deepfakes, frames_JPEG_5/Face2Face, frames_JPEG_5/FaceSwap, frames_JPEG_5/NeuralTextures] + # FAKETYPE: [Celeb-real-0.6-0.8-v2, Celeb-synthesis-0.6-0.8-v2, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [Real_DiffSwap1, DiffSwap_1] + # FAKETYPE: [blendface, danet, deepfacelab, e4e, e4s, facedancer, faceswap, facevid2vid, fomm, fsgan, heygen, + # hyperreenact, inswap, lia, mcnet, mobileswap, MRAA, one_shot_free, pirender, sadtalker, simswap, tpsm, uniface, + # wav2lip, real_videos] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [384, 384, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.5, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.5, 0.5, 0.5] + std: [0.5, 0.5, 0.5] + DEBUG: False + DYNAMIC_FXRAY: True +MODEL: + type: PoseEfficientNet + model_name: efficientnet-b4 + num_layers: B4 + include_top: False + include_hm_decoder: True + head_conv: 64 + use_c2: False + use_c3: True + use_c4: True + use_c51: True + efpn: True + tfpn: False + se_layer: False + norm_c2: False + heads: + hm: 1 + cls: 1 + cstency: 256 + INIT_WEIGHTS: + pretrained: True + advprop: True +TRAIN: + gpus: [0] + batch_size: 8 + lr: 0.0000001 + epochs: 100 + begin_epoch: -1 + warm_up: 6 + every_val_epochs: 1 + loss: + type: CombinedFocalLoss + use_target_weight: False + cls_lmda: 1 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 10 + cstency_lmda: 100 + mse_reduction: sum + ce_reduction: mean + optimizer: SAM + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: False + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 8 +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: True + video_level: True + apr: True # Average precision/recall or normal precision/recall + no_shot_preds: 1 + pred_file: /project/home/p200249/XXX/saved_predictions/laanet_DiffSwap_100_100.json # File to save predictions + save_preds: True + pretrained: pretrained/PoseEfficientNet_EFN_hm100_EFPN_NoBasedCLS_Focal_C3_256Cst100_8FXRayv2_SAM(Adam)_ADV_Era1_OutSigmoid_5e5_model_best.pth diff --git a/video/fake-stormer/model_code/configs/spatial/hrnet_sbi.yaml b/video/fake-stormer/model_code/configs/spatial/hrnet_sbi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a7ab3ece0262d23a7ee1e7612d3654e13c8bdc68 --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/hrnet_sbi.yaml @@ -0,0 +1,170 @@ +TASK: heatmap_sbi_separated_CLS_Focal_C2 +PRECISION: float64 +METRICS_BASE: combine +DATASET: + type: SBIFaceForensic + TRAIN: True #Switch to True for training mode, False for testing mode + ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/FaceForensics++/c23/ + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + IMAGE_SUFFIX: jpg + FROM_FILE: True + NUM_WORKERS: 8 + PIN_MEMORY: True + IMAGE_SIZE: [384, 384] + HEATMAP_SIZE: [96, 96] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 3 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: gaussian + DATA: + TYPE: images + SAMPLES_PER_VIDEO: 32 + TRAIN: + FAKETYPE: [FaceXRay] + # ANNO_FILE: FaceXRay/train/train_FF_Xray.json + # ANNO_FILE: FaceXRay/train/new_trainBI_FF.json + ANNO_FILE: processed_data/new_valBI_FF.json + LABEL_FOLDER: [real, fake] + VAL: + FAKETYPE: [FaceXRay] + # ANNO_FILE: FaceXRay/val/val_FF_Xray.json + ANNO_FILE: processed_data/new_valBI_FF.json + LABEL_FOLDER: [real, fake] + TEST: + FAKETYPE: [FaceXRay] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [256, 256, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + jpegcompression: 0.5 + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True +MODEL: + type: PoseHighResolutionNet + INIT_WEIGHTS: + pretrained: 'pretrained/hrnet_w48-8ef0771d.pth' + MODEL: + NAME: pose_hrnet + NUM_JOINTS: 1 + HEATMAP_SIZE: [96, 96] + cls_based_hm: True + heads: + hm: 1 + cls: 1 + EXTRA: + PRETRAINED_LAYERS: + - 'conv1' + - 'bn1' + - 'conv2' + - 'bn2' + - 'layer1' + - 'transition1' + - 'stage2' + - 'transition2' + - 'stage3' + - 'transition3' + - 'stage4' + FINAL_CONV_KERNEL: 1 + STAGE2: + NUM_MODULES: 1 + NUM_BRANCHES: 2 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + NUM_CHANNELS: + - 48 + - 96 + FUSE_METHOD: SUM + STAGE3: + NUM_MODULES: 4 + NUM_BRANCHES: 3 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + - 4 + NUM_CHANNELS: + - 48 + - 96 + - 192 + FUSE_METHOD: SUM + STAGE4: + NUM_MODULES: 3 + NUM_BRANCHES: 4 + BLOCK: BASIC + NUM_BLOCKS: + - 4 + - 4 + - 4 + - 4 + NUM_CHANNELS: + - 48 + - 96 + - 192 + - 384 + FUSE_METHOD: SUM +TRAIN: + gpus: [0,1,2] + batch_size: 16 + lr: 0.001 + epochs: 30 + begin_epoch: 0 + warm_up: 5 + every_val_epochs: 1 + loss: + type: CombinedFocalLoss + use_target_weight: False + cls_lmda: 0.08 + # dst_lmda: 0.05 + reduction: 'mean' + # dist_cal: False + cls_cal: True + combine_compute: False + optimizer: SAM + distributed: False + # pretrained: 'logs/05-12-2022/PoseResNet_heatmap_FPN_Separated_CLS_Focal_NoFrZ_model_best.pth' + tensorboard: True + resume: False + lr_scheduler: + type: LinearDecayLR + milestones: [5, 12] + gamma: 0.5 + freeze_backbone: False + debug: + save_hm_gt: True + save_hm_pred: True +TEST: + gpus: [0,1,2] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + # pretrained: 'logs/05-12-2022/PoseResNet_heatmap_FPN_Separated_CLS_Focal_NoFrZ_model_best.pth' +PREPROCESSING: + DATASET: FaceForensics + SPLIT: train + ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c23/ + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + IMAGE_SUFFIX: jpg + DATA_TYPE: images + LABEL: [real, fake] + facial_lm_pretrained: pretrained/shape_predictor_68_face_landmarks.dat + DEBUG: False diff --git a/video/fake-stormer/model_code/configs/spatial/resnet_fpn_hm.yaml b/video/fake-stormer/model_code/configs/spatial/resnet_fpn_hm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f9e67ef529105897440753bf399a5f150b693a08 --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/resnet_fpn_hm.yaml @@ -0,0 +1,153 @@ +TASK: PoseRes50_100hm_EFPN_NoBased_CLS_Focal_C2_256Cst100_32FXRayv2_SAM(Adam)_Era1_OutSigmoid_5e5_div4_FZ +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True +DATASET: + type: HeatmapFaceForensic + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 7 + PIN_MEMORY: True + IMAGE_SIZE: [256, 256] + HEATMAP_SIZE: [64, 64] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 2 + ADAPTIVE_SIGMA: True + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 32 # Dynamically random number of frames in each epoch + VAL: 32 + TEST: 32 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/dynamic_trainBI_FFv2.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/dynamic_valBI_FFv2.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: Celeb-DFv1 + ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /data/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [256, 256, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.5, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True +MODEL: + type: PoseResNet + num_layers: 50 + block: Bottleneck + layers: [3, 4, 6, 3] + heads: + hm: 1 + cls: 1 + cstency: 256 + head_conv: 64 + dropout_prob: 0.5 + fpn: True + cls_based_hm: False + use_c2: True + INIT_WEIGHTS: + pretrained: True + num_layers: 50 +TRAIN: + gpus: [0] + batch_size: 32 + lr: 0.00005 + epochs: 100 + begin_epoch: -1 + warm_up: 6 + every_val_epochs: 1 + loss: + type: CombinedFocalLoss + use_target_weight: False + cls_lmda: 1 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 100 + cstency_lmda: 100 + mse_reduction: sum + ce_reduction: mean + optimizer: SAM + distributed: False + # pretrained: '' + tensorboard: True + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: False + pretrained: 'logs/03-03-2023/PoseResNet_PoseRes50_100hm_FPN_NoBased_CLS_Focal_NoC2_256Cstency10_32FXRayv2_SAM(Adam)_NoErasing_model_best.pth' +PREPROCESSING: + DATASET: FaceForensics + SPLIT: train + ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c23/ + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + IMAGE_SUFFIX: jpg + DATA_TYPE: images + LABEL: [real, fake] + facial_lm_pretrained: pretrained/shape_predictor_68_face_landmarks.dat + DEBUG: False diff --git a/video/fake-stormer/model_code/configs/spatial/resnet_fpn_hm_deepfakes.yaml b/video/fake-stormer/model_code/configs/spatial/resnet_fpn_hm_deepfakes.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d65e890f6234dde1d303ddf501d80b42893aa3b2 --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/resnet_fpn_hm_deepfakes.yaml @@ -0,0 +1,115 @@ +TASK: heatmap_FPN_Separated_CLS_Focal_Deepfakes +PRECISION: float64 +DATASET: + type: HeatmapFaceForensic + TRAIN: False #Switch to True for training mode, False for testing mode + ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/FaceForensics++/c23/ + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + IMAGE_SUFFIX: jpg + FROM_FILE: False + NUM_WORKERS: 14 + PIN_MEMORY: True + IMAGE_SIZE: [256, 256] + HEATMAP_SIZE: [64, 64] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 2 + ADAPTIVE_SIGMA: True + HEATMAP_TYPE: gaussian + DATA: + TYPE: images + TRAIN: + FAKETYPE: [Deepfakes_FaceXRay] + ANNO_FILE: Deepfakes_FaceXRay/val/val_FF_Deepfakes_Xray.json + # ANNO_FILE: FaceXRay/train/new_trainBI_FF.json + # ANNO_FILE: processed_data/new_valBI_FF.json + LABEL_FOLDER: [real, fake] + VAL: + FAKETYPE: [Deepfakes_FaceXRay] + ANNO_FILE: Deepfakes_FaceXRay/val/val_FF_Deepfakes_Xray.json + # ANNO_FILE: processed_data/new_valBI_FF.json + LABEL_FOLDER: [real, fake] + TEST: + FAKETYPE: [Deepfakes] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [256, 256, 0] # h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.3 + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + jpegcompression: 0.5 + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: False +MODEL: + type: PoseResNet + num_layers: 50 + # block: Bottleneck + layers: [3, 4, 6, 3] + heads: + hm: 1 + cls: 1 + head_conv: 64 + dropout_prob: 0.5 + fpn: True + cls_based_hm: False + use_c2: False +TRAIN: + gpus: [0,1,2,3] + batch_size: 32 + lr: 0.001 + epochs: 30 + begin_epoch: 0 + warm_up: 6 + every_val_epochs: 1 + loss: + type: CombinedFocalLoss + use_target_weight: False + cls_lmda: 1 + # dst_lmda: 0.05 + reduction: 'mean' + # dist_cal: False + cls_cal: True + optimizer: Adam + distributed: False + # pretrained: 'logs/05-12-2022/PoseResNet_heatmap_FPN_Separated_CLS_Focal_NoFrZ_model_best.pth' + tensorboard: True + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [6, 15, 20, 25] + gamma: 0.5 + freeze_backbone: False + debug: + save_hm_gt: True + save_hm_pred: True +TEST: + gpus: [0,1,2,3] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + pretrained: 'logs/08-12-2022/PoseResNet_heatmap_FPN_Separated_CLS_Focal_Deepfakes_model_best.pth' +PREPROCESSING: + DATASET: FaceForensics + SPLIT: train + ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c23/ + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + IMAGE_SUFFIX: jpg + DATA_TYPE: images + LABEL: [real, fake] + facial_lm_pretrained: pretrained/shape_predictor_68_face_landmarks.dat + DEBUG: False diff --git a/video/fake-stormer/model_code/configs/spatial/resnet_fpn_hmbin.yaml b/video/fake-stormer/model_code/configs/spatial/resnet_fpn_hmbin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..65d272ed59ff458c502b37b5645101b844f7932a --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/resnet_fpn_hmbin.yaml @@ -0,0 +1,102 @@ +TASK: heatmapbin +PRECISION: float64 +DATASET: + type: HeatmapFaceForensic + TRAIN: True #Switch to True for training mode, False for testing mode + ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c23/ + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + IMAGE_SUFFIX: jpg + FROM_FILE: True + NUM_WORKERS: 8 + PIN_MEMORY: True + IMAGE_SIZE: [256, 256] + HEATMAP_SIZE: [64, 64] + SIGMA: 2 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: gaussian + DATA: + TYPE: images + TRAIN: + FAKETYPE: [FaceXRay] + ANNO_FILE: FaceXRay/train/train_FF_Xray.json + LABEL_FOLDER: [real, fake] + VAL: + FAKETYPE: [FaceXRay] + ANNO_FILE: FaceXRay/val/val_FF_Xray.json + LABEL_FOLDER: [real, fake] + TEST: + FAKETYPE: [NeuralTextures] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [256, 256, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + color: + type: ColorJitterTransform + clahe: 0.5 + colorjitter: 0.5 + gaussianblur: 0.5 + jpegcompression: 0.0 + rgbshift: 0.5 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: False +MODEL: + type: PoseResNet + num_layers: 50 + # block: Bottleneck + layers: [3, 4, 6, 3] + heads: + hm: 1 + cls: 1 + head_conv: 64 + is_fpn: False + dropout_prob: 0.4 +TRAIN: + gpus: [0,1,2,3] + batch_size: 32 + lr: 0.001 + epochs: 30 + begin_epoch: 0 + warm_up: 6 + every_val_epochs: 1 + loss: + type: HeatmapBinaryLoss + use_target_weight: False + cls_lmda: 0.05 + reduction: 'mean' + cls_cal: True + optimizer: Adam + distributed: False + pretrained: 'logs/23-11-2022/PoseResNet_heatmapbin_model_best.pth' + tensorboard: True + resume: True + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + save_hm_gt: True + save_hm_pred: True +TEST: + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + pretrained: 'logs/16-11-2022/PoseResNet_model_best.pth' +PREPROCESSING: + DATASET: FaceForensics + SPLIT: train + ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c23/ + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + IMAGE_SUFFIX: jpg + DATA_TYPE: images + LABEL: [real, fake] + facial_lm_pretrained: pretrained/shape_predictor_68_face_landmarks.dat + DEBUG: False diff --git a/video/fake-stormer/model_code/configs/spatial/resnet_fpn_sbi.yaml b/video/fake-stormer/model_code/configs/spatial/resnet_fpn_sbi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..26793c43cbe38e4883de1e7b40df15dc53978e7b --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/resnet_fpn_sbi.yaml @@ -0,0 +1,133 @@ +TASK: heatmap_sbi_Separated_CLS_Focal_NoC2_50_SoftDISCRE +PRECISION: float64 +METRICS_BASE: combine +SEED: 5 +DATASET: + type: SBIFaceForensic + TRAIN: True #Switch to True for training mode, False for testing mode + ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/FaceForensics++/c23/ + # FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + IMAGE_SUFFIX: jpg + FROM_FILE: True + NUM_WORKERS: 4 + PIN_MEMORY: True + IMAGE_SIZE: [256, 256] + HEATMAP_SIZE: [64, 64] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 2 + ADAPTIVE_SIGMA: True + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: True + DATA: + TYPE: images + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 32 + VAL: 32 + TEST: 32 + TRAIN: + FAKETYPE: [FaceSBI] + # ANNO_FILE: FaceXRay/train/train_FF_Xray.json + ANNO_FILE: FaceSBI/train/train_FF_SBI.json + ANNO_FILE_R: processed_data/train_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + VAL: + FAKETYPE: [FaceSBI] + ANNO_FILE: FaceSBI/val/val_FF_SBI.json + ANNO_FILE_R: processed_data/val_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + TEST: + FAKETYPE: [FaceSBI] + ANNO_FILE_R: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [256, 256, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + #Either Scaling or Cropping, not do at the same time + cropping: [0.75, 1.25, 1] #Format: [low, high, prob] + scale: [0.1, 0.25, 1] #Format: [shift, scale, prob] + erasing: 0.5 + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + jpegcompression: 0.5 + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True +MODEL: + type: PoseResNet + num_layers: 50 + # block: Bottleneck + layers: [3, 4, 6, 3] + heads: + hm: 1 + cls: 1 + # offset: 1 + head_conv: 64 + dropout_prob: 0.5 + fpn: True + cls_based_hm: False + use_c2: False + INIT_WEIGHTS: + pretrained: True + num_layers: 50 +TRAIN: + gpus: [0] + batch_size: 16 + lr: 0.0001 + epochs: 50 + begin_epoch: 0 + warm_up: 3 + every_val_epochs: 1 + loss: + type: CombinedFocalLoss + use_target_weight: False + cls_lmda: 1 + dst_hm_cls_lmda: 0.01 + offset_lmda: 0 + hm_lmda: 100 + mse_reduction: sum + ce_reduction: mean + optimizer: SAM + distributed: False + # pretrained: 'logs/09-01-2023/PoseResNet_heatmap_sbi_Separated_CLS_Focal_NoC2_50_SoftDISCRE_model_best.pth' + tensorboard: True + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [25, 35, 45] + gamma: 0.5 + freeze_backbone: True + debug: + active: True + save_hm_gt: True + save_hm_pred: True +TEST: + gpus: [0,1,2,3] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + pretrained: 'logs/28-12-2022/PoseResNet_heatmap_sbi_based_CLS_Focal_C2_101_model_best.pth' +PREPROCESSING: + DATASET: FaceForensics + SPLIT: train + ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c23/ + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + IMAGE_SUFFIX: jpg + DATA_TYPE: images + LABEL: [real, fake] + facial_lm_pretrained: pretrained/shape_predictor_68_face_landmarks.dat + DEBUG: False diff --git a/video/fake-stormer/model_code/configs/spatial/swin_bi_small.yaml b/video/fake-stormer/model_code/configs/spatial/swin_bi_small.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cb9e4009422357c5ab7e9cc3b39a92bdb050ebd3 --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/swin_bi_small.yaml @@ -0,0 +1,180 @@ +TASK: SwinSmall224_hm10_16BI_Overlap100_Focal_AdamW_IN_1e7_FZ5_200epochs_Drop0.2_Boost250 +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True + +DATASET: + type: FakeSFormerBI + DATA_TYPE: image + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 7 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 16 # Dynamically random number of frames in each epoch + VAL: 16 + TEST: 32 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/dynamic_trainBI_FFv2.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/dynamic_valBI_FFv2.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: DFDC + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: False + TARGET_OVERLAP: True + +MODEL: + type: TopDownDetector + backbone: + type: SwinTransformer + embed_dims: 96 + depths: [2, 2, 18, 2] + num_heads: [3, 6, 12, 24] + window_size: 7 + mlp_ratio: 4 + qkv_bias: True + drop_rate: 0. + attn_drop_rate: 0. + drop_path_rate: 0.2 + patch_norm: True + with_cp: False + convert_weights: True + pretrained: pretrained/swin_small_patch4_window7_224_22k.pth + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 768 + heads: + hm: 1 + cls: 1 + num_deconv_layers: 1 #Config n deconv layers to build the decoder + num_deconv_filters: [768] + num_deconv_kernels: [4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + +TRAIN: + gpus: [0] + batch_size: 32 + lr: 0.0000001 + epochs: 200 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + loss: + type: CombinedFocalLoss + use_target_weight: False + cls_lmda: 1 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 10 + cstency_lmda: 0 + mse_reduction: mean + ce_reduction: mean + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 250 + start_decay: 4 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: True + video_level: True + pretrained: 'logs/12-03-2024/TopDownDetector_SwinSmall224_hm10_16BI_Overlap100_Focal_AdamW_IN_1e7_FZ5_Batch32_200epochs_Drop0.2_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/spatial/swin_sbi_base.yaml b/video/fake-stormer/model_code/configs/spatial/swin_sbi_base.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8e902c1e0cd1116b7400a047540401680f58d0a2 --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/swin_sbi_base.yaml @@ -0,0 +1,197 @@ +TASK: SwinBase224_hm10_8SBI_NonOverlap100_MSE_AdamW_IN_5e5_FZ5_Batch32_200epochs_Drop0.2 +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: image + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 7 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 8 # Dynamically random number of frames in each epoch + VAL: 8 + TEST: 32 + DIST: [1., 1.] # Distribution of Real, Fake - 1.,1. by default + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/train_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/val_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: FF++ + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + # ROOT: /project/home/p200249/XXX/DF40_test/ + # ROOT: /project/home/p200249/XXX/DiffSwap/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + FAKETYPE: [original_pixel_random, Deepfakes_pixel_random, Face2Face_pixel_random, FaceSwap_pixel_random, NeuralTextures_pixel_random] + # FAKETYPE: [frames_CS_1/original, frames_CS_1/Deepfakes, frames_CS_1/Face2Face, frames_CS_1/FaceSwap, frames_CS_1/NeuralTextures, + # frames_CS_2/original, frames_CS_2/Deepfakes, frames_CS_2/Face2Face, frames_CS_2/FaceSwap, frames_CS_2/NeuralTextures, + # frames_CS_3/original, frames_CS_3/Deepfakes, frames_CS_3/Face2Face, frames_CS_3/FaceSwap, frames_CS_3/NeuralTextures, + # frames_CS_4/original, frames_CS_4/Deepfakes, frames_CS_4/Face2Face, frames_CS_4/FaceSwap, frames_CS_4/NeuralTextures, + # frames_CS_5/original, frames_CS_5/Deepfakes, frames_CS_5/Face2Face, frames_CS_5/FaceSwap, frames_CS_5/NeuralTextures] + # FAKETYPE: [Celeb-real-0.97-1.0-v2, Celeb-synthesis-0.97-1.0-v2, YouTube-real] + # FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [mobileswap, real_videos] + # FAKETYPE: [Real_DiffSwap1, DiffSwap_1] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: False + TARGET_OVERLAP: False + +MODEL: + type: TopDownDetector + backbone: + type: SwinTransformer + embed_dims: 128 + depths: [2, 2, 18, 2] + num_heads: [4, 8, 16, 32] + window_size: 7 + mlp_ratio: 4 + qkv_bias: True + drop_rate: 0. + attn_drop_rate: 0. + drop_path_rate: 0.2 + patch_norm: True + with_cp: False + convert_weights: True + # pretrained: pretrained/swin_base_patch4_window7_224_22k.pth + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 512 + heads: + hm: 1 + cls: 1 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [512] + num_deconv_kernels: [4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + INIT_WEIGHTS: + pretrained: pretrained/swin_base_patch4_window7_224_22k.pth + +TRAIN: + gpus: [0,1,2,3] + batch_size: 32 + lr: 0.00005 + epochs: 200 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + loss: + type: CombinedMSELoss + use_target_weight: False + cls_lmda: 1 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 10 + cstency_lmda: 0 + mse_reduction: mean + ce_reduction: mean + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: True + video_level: True + apr: True # Average precision/recall or normal precision/recall + no_shot_preds: 1 + pred_file: /project/home/p200249/XXX/saved_predictions/FakeSwin-B_DiffSwap_100_100.json # File to save predictions + save_preds: False + pretrained: pretrained/fakeformer/TopDownDetector_SwinBase224_hm10_8SBI_NonOverlap100_MSE_AdamW_MAE_5e5_FZ5_Batch32_200epochs_Drop0.2_model_best.pth diff --git a/video/fake-stormer/model_code/configs/spatial/swin_sbi_small.yaml b/video/fake-stormer/model_code/configs/spatial/swin_sbi_small.yaml new file mode 100644 index 0000000000000000000000000000000000000000..790ea28721e73bf191bc566add146f436bb9b5d5 --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/swin_sbi_small.yaml @@ -0,0 +1,201 @@ +TASK: SwinSmall224_hm10_8SBI_Overlap100_Focal_AdamW_IN22k_1e3_FZ5_Cutout_abl +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: image + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 7 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 8 # Dynamically random number of frames in each epoch + VAL: 8 + TEST: 32 + DIST: [1., 1.] # Distribution of Real, Fake - 1.,1. by default + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/train_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/val_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: FF++ + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + # ROOT: /project/home/p200249/XXX/DF40_test/ + # ROOT: /project/home/p200249/XXX/DiffSwap/ + # ROOT: /home/users/XXX/data/Combine/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # FAKETYPE: [original_pixel_random, Deepfakes_pixel_random, Face2Face_pixel_random, FaceSwap_pixel_random, NeuralTextures_pixel_random] + # FAKETYPE: [Celeb-real-0.6-0.8-v2, Celeb-synthesis-0.6-0.8-v2, YouTube-real] + # FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [Celeb_DFv1-real, Celeb_DFv2-real, DeepFakeDetection, Deepfakes, FaceSwap, method_A, NeuralTextures, + # original_videos, YouTube_DFv1-real, Celeb_DFv1-synthesis, Celeb_DFv2-synthesis, DeepFakeDetection_original, Face2Face, + # fake_test, method_B, original, real_test, YouTube_DFv2-real] + # FAKETYPE: [e4s, real_videos] + # FAKETYPE: [Real_DiffSwap1, DiffSwap_1] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.5, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: False + TARGET_OVERLAP: True + +MODEL: + type: TopDownDetector + backbone: + type: SwinTransformer + embed_dims: 96 + depths: [2, 2, 18, 2] + num_heads: [3, 6, 12, 24] + window_size: 7 + mlp_ratio: 4 + qkv_bias: True + drop_rate: 0. + attn_drop_rate: 0. + drop_path_rate: 0.2 + patch_norm: True + with_cp: False + convert_weights: True + # pretrained: pretrained/swin_small_patch4_window7_224.pth + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 768 + heads: + # hm: 1 + cls: 1 + num_deconv_layers: 1 #Config n deconv layers to build the decoder + num_deconv_filters: [384] + num_deconv_kernels: [4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + INIT_WEIGHTS: + pretrained: pretrained/swin_small_patch4_window7_224_22k.pth + +TRAIN: + gpus: [0,1,2,3] + batch_size: 32 + lr: 0.001 + epochs: 100 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + accumulation_steps: 1 + use_amp: False + loss: + type: CombinedFocalLoss # For L2-Att + use_target_weight: False + cls_lmda: 1 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 10 + cstency_lmda: 0 + mse_reduction: mean + ce_reduction: mean + # type: BinaryCrossEntropy # For binary only + # reduction: mean + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 10 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: True + video_level: True + apr: True # Average precision/recall or normal precision/recall + no_shot_preds: 1 + pred_file: /project/home/p200249/XXX/saved_predictions/FakeSwin_CDF2_100_100.json # File to save predictions + save_preds: False + pretrained: logs/29-10-2025/TopDownDetector_SwinSmall224_hm0_8SBI_Overlap100_BCE_AdamW_IN22k_5e4_FZ5_0Cutout_abl_model_best.pth diff --git a/video/fake-stormer/model_code/configs/spatial/swin_sbi_tiny.yaml b/video/fake-stormer/model_code/configs/spatial/swin_sbi_tiny.yaml new file mode 100644 index 0000000000000000000000000000000000000000..276608e93755f9942473eacf0a8502d791865358 --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/swin_sbi_tiny.yaml @@ -0,0 +1,178 @@ +TASK: SwinTiny224_hm10_16SBI_Overlap100_MSE_AdamW_IN_5e5_FZ5_Batch32_200epochs_Drop0.2 +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: image + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 32 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 16 # Dynamically random number of frames in each epoch + VAL: 16 + TEST: 32 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/train_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/val_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: Celeb-DFv2 + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: False + TARGET_OVERLAP: True + +MODEL: + type: TopDownDetector + backbone: + type: SwinTransformer + embed_dims: 96 + depths: [2, 2, 6, 2] + num_heads: [3, 6, 12, 24] + window_size: 7 + mlp_ratio: 4 + qkv_bias: True + drop_rate: 0. + attn_drop_rate: 0. + drop_path_rate: 0.2 + patch_norm: True + with_cp: False + convert_weights: True + pretrained: pretrained/swin_tiny_patch4_window7_224.pth + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 768 + heads: + hm: 1 + cls: 1 + num_deconv_layers: 1 #Config n deconv layers to build the decoder + num_deconv_filters: [384] + num_deconv_kernels: [4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + +TRAIN: + gpus: [0,1,2,3] + batch_size: 32 + lr: 0.00005 + epochs: 200 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + loss: + type: CombinedMSELoss + use_target_weight: False + cls_lmda: 1 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 10 + cstency_lmda: 0 + mse_reduction: mean + ce_reduction: mean + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + +TEST: + gpus: [0,1,2,3] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: True + video_level: True + pretrained: 'logs/12-01-2024/TopDownDetector_SwinTiny224_hm10_32SBI_Overlap100_MSE_AdamW_IN_5e5_FZ5_Batch32_200epochs_Drop0.2_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/spatial/vit_bi_small.yaml b/video/fake-stormer/model_code/configs/spatial/vit_bi_small.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d226cff7e98ca824575e916f422dc6bc46a2825c --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/vit_bi_small.yaml @@ -0,0 +1,184 @@ +TASK: ViTSmall112_hm100_32BI_Overlap100_MSE_AdamW_5e5_FZ5_Batch32_200epochs_Drop0.2 +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True + +DATASET: + type: FakeSFormerBI + DATA_TYPE: image + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 7 + PIN_MEMORY: True + IMAGE_SIZE: [112, 112] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: False + TRAIN: 32 # Dynamically random number of frames in each epoch + VAL: 32 + TEST: 32 + DIST: [1., 1.] # Distribution of Real, Fake - 1.,1. by default + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/dynamic_trainBI_FFv2.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/dynamic_valBI_FFv2.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: DFW + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + ROOT: /project/home/p200249/XXX/DFW/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + FAKETYPE: [fake_test, real_test] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [112, 112, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.5, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: False + TARGET_OVERLAP: True + +MODEL: + type: TopDownDetector + backbone: + type: ViT + img_size: [112, 112] + patch_size: 8 + embed_dim: 384 + depth: 12 + num_heads: 6 + mlp_ratio: 4 + drop_path_rate: 0.2 + qkv_bias: True + class_token: True + # pretrained: pretrained/dino_deitsmall8_pretrain.pth + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 384 + heads: + hm: 1 + cls: 1 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256, 256] + num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + INIT_WEIGHTS: + pretrained: pretrained/dino_deitsmall8_pretrain.pth + +TRAIN: + gpus: [0,1,2,3] + batch_size: 32 + lr: 0.00005 + epochs: 200 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + accumulation_steps: 1 + use_amp: False + loss: + type: CombinedFocalLoss + use_target_weight: False + cls_lmda: 1 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 0 + cstency_lmda: 0 + mse_reduction: mean + ce_reduction: mean + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: True + video_level: True + apr: True # Average precision/recall or normal precision/recall + no_shot_preds: 1 + pred_file: /project/home/p200249/XXX/saved_predictions/FakeFormer_wBI_DFW_100_100.json # File to save predictions + save_preds: True + pretrained: pretrained/fakeformer/TopDownDetector_ViTSmall112_hm10_8SBI_Overlap100_MSE_AdamW_IN_5e5_FZ5_Batch32_200epochs_Drop0.2_Decay1e4_model_best.pth diff --git a/video/fake-stormer/model_code/configs/spatial/vit_sbi_base.yaml b/video/fake-stormer/model_code/configs/spatial/vit_sbi_base.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f0225675a083abde4da6dfe9779e5a0c98fa09b4 --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/vit_sbi_base.yaml @@ -0,0 +1,187 @@ +TASK: ViTBase224_hm10_8SBI_NonOverlap100_MSE_AdamW_IN_5e5_FZ5_Batch32_200epochs_Drop0.3 +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: image + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 7 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 8 # Dynamically random number of frames in each epoch + VAL: 8 + TEST: 32 + DIST: [1., 1.] # Distribution of Real, Fake - 1.,1. by default + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/train_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/val_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: DiffSwap + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + ROOT: /project/home/p200249/XXX/DiffSwap/ + # ROOT: /project/home/p200249/XXX/DF40_test/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # FAKETYPE: [Celeb-real-0.6-0.8-v2, Celeb-synthesis-0.6-0.8-v2, YouTube-real] + # FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + FAKETYPE: [Real_DiffSwap1, DiffSwap_1] + # FAKETYPE: [mobileswap, real_videos] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: False + TARGET_OVERLAP: False + +MODEL: + type: TopDownDetector + backbone: + type: ViT + img_size: [224, 224] + patch_size: 16 + embed_dim: 768 + depth: 12 + num_heads: 12 + mlp_ratio: 4 + drop_path_rate: 0.3 + qkv_bias: True + class_token: True + # pretrained: pretrained/dino_vitbase16_pretrain.pth + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 768 + heads: + hm: 1 + cls: 1 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256, 256] + num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + INIT_WEIGHTS: + pretrained: pretrained/dino_vitbase16_pretrain.pth + +TRAIN: + gpus: [0,1,2,3] + batch_size: 32 + lr: 0.00005 + epochs: 200 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + loss: + type: CombinedMSELoss + use_target_weight: False + cls_lmda: 1 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 10 + cstency_lmda: 0 + mse_reduction: mean + ce_reduction: mean + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: True + video_level: True + apr: True # Average precision/recall or normal precision/recall + no_shot_preds: 1 + pred_file: /project/home/p200249/XXX/saved_predictions/FakeFormer-B_DiffSwap_100_100.json # File to save predictions + save_preds: True + pretrained: pretrained/fakeformer/TopDownDetector_ViTBase224_hm10_8SBI_NonOverlap100_MSE_AdamW_IN_5e5_FZ5_Batch32_200epochs_Drop0.3_model_best.pth diff --git a/video/fake-stormer/model_code/configs/spatial/vit_sbi_large.yaml b/video/fake-stormer/model_code/configs/spatial/vit_sbi_large.yaml new file mode 100644 index 0000000000000000000000000000000000000000..59767190e46178ebd8df4ffd2b32a8260fc4e7c4 --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/vit_sbi_large.yaml @@ -0,0 +1,171 @@ +TASK: ViTLarge224_hm1_8SBI_NonOverlap100_MSE_AdamW_MAE_1e4_FZ5_Batch16_200epochs +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: image + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 36 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 8 # Dynamically random number of frames in each epoch + VAL: 8 + TEST: 32 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/train_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/val_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: Celeb-DFv1 + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: False + TARGET_OVERLAP: False + +MODEL: + type: TopDownDetector + backbone: + type: ViT + img_size: [224, 224] + patch_size: 16 + embed_dim: 1024 + depth: 24 + num_heads: 16 + mlp_ratio: 4 + qkv_bias: True + class_token: True + pretrained: pretrained/mae_pretrain_vit_large.pth + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 1024 + heads: + hm: 1 + cls: 1 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256, 256] + num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + +TRAIN: + gpus: [0,1,2,3] + batch_size: 16 + lr: 0.0001 + epochs: 200 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + loss: + type: CombinedMSELoss + use_target_weight: False + cls_lmda: 1 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 1 + cstency_lmda: 0 + mse_reduction: mean + ce_reduction: mean + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: True + video_level: True + pretrained: 'logs/14-12-2023/TopDownDetector_ViTLarge224_hm1_8SBI_NonOverlap100_MSE_AdamW_MAE_1e4_FZ5_Batch16_200epochs_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/spatial/vit_sbi_small.yaml b/video/fake-stormer/model_code/configs/spatial/vit_sbi_small.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ff45c683c4c794646a289dd86494d6a6c9b87bac --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/vit_sbi_small.yaml @@ -0,0 +1,208 @@ +TASK: ViTSmall112_hm0_8SBI_Overlap100_BCE_AdamW_IN22k_5e4_FZ5_0Cutout_abl +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: image + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 7 + PIN_MEMORY: True + IMAGE_SIZE: [112, 112] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 8 # Dynamically random number of frames in each epoch + VAL: 8 + TEST: 32 + DIST: [1., 1.] # Distribution of Real, Fake - 1.,1. by default + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/train_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/val_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: DFD + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + # ROOT: /project/home/p200249/XXX/DF40_test/ + # ROOT: /project/home/p200249/XXX/DiffSwap/ + # ROOT: /home/users/XXX/data/Combine/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # FAKETYPE: [original_pixel_random, Deepfakes_pixel_random, Face2Face_pixel_random, FaceSwap_pixel_random, NeuralTextures_pixel_random] + # FAKETYPE: [frames_JPEG_1/original, frames_JPEG_1/Deepfakes, frames_JPEG_1/Face2Face, frames_JPEG_1/FaceSwap, frames_JPEG_1/NeuralTextures, + # frames_JPEG_2/original, frames_JPEG_2/Deepfakes, frames_JPEG_2/Face2Face, frames_JPEG_2/FaceSwap, frames_JPEG_2/NeuralTextures, + # frames_JPEG_3/original, frames_JPEG_3/Deepfakes, frames_JPEG_3/Face2Face, frames_JPEG_3/FaceSwap, frames_JPEG_3/NeuralTextures, + # frames_JPEG_4/original, frames_JPEG_4/Deepfakes, frames_JPEG_4/Face2Face, frames_JPEG_4/FaceSwap, frames_JPEG_4/NeuralTextures, + # frames_JPEG_5/original, frames_JPEG_5/Deepfakes, frames_JPEG_5/Face2Face, frames_JPEG_5/FaceSwap, frames_JPEG_5/NeuralTextures] + # FAKETYPE: [Celeb-real-0.97-1.0-v2, Celeb-synthesis-0.97-1.0-v2, YouTube-real] + # FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [blendface, danet, deepfacelab, e4e, e4s, facedancer, faceswap, facevid2vid, fomm, fsgan, heygen, + # hyperreenact, inswap, lia, mcnet, mobileswap, MRAA, one_shot_free, pirender, sadtalker, simswap, tpsm, uniface, + # wav2lip, real_videos] + # FAKETYPE: [Real_DiffSwap1, DiffSwap_1] + # FAKETYPE: [Celeb_DFv1-real, Celeb_DFv2-real, DeepFakeDetection, Deepfakes, FaceSwap, method_A, NeuralTextures, + # original_videos, YouTube_DFv1-real, Celeb_DFv1-synthesis, Celeb_DFv2-synthesis, DeepFakeDetection_original, Face2Face, + # fake_test, method_B, original, real_test, YouTube_DFv2-real] + # FAKETYPE: [fomm, real_videos] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [112, 112, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: False + TARGET_OVERLAP: True + +MODEL: + type: TopDownDetector + backbone: + type: ViT + img_size: [112, 112] + patch_size: 8 + embed_dim: 384 + depth: 12 + num_heads: 6 + mlp_ratio: 4 + drop_path_rate: 0.2 + qkv_bias: True + class_token: True + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 384 + heads: + hm: 1 + cls: 1 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256] + num_deconv_kernels: [4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + # conv_2direction: False + # features: 2D + INIT_WEIGHTS: + pretrained: pretrained/dino_deitsmall8_pretrain.pth + +TRAIN: + gpus: [0,1,2,3] + batch_size: 32 + lr: 0.0005 + epochs: 100 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + accumulation_steps: 1 + use_amp: False + loss: + # type: CombinedFocalLoss # For L2-Att + # use_target_weight: False + # cls_lmda: 1 + # dst_hm_cls_lmda: 0 + # offset_lmda: 0 + # hm_lmda: 10 + # cstency_lmda: 0 + # mse_reduction: mean + # ce_reduction: mean + type: BinaryCrossEntropy # For binary only + reduction: mean + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 10 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: True + video_level: True + apr: True # Average precision/recall or normal precision/recall + no_shot_preds: 1 + pred_file: /project/home/p200249/XXX/saved_predictions/FakeFormer_DF40_FOMM_100_100.json # File to save predictions + save_preds: False + pretrained: logs/30-10-2025/TopDownDetector_ViTSmall112_hm10_8SBI_Overlap100_Focal_AdamW_IN22k_1e3_FZ5_Cutout_abl_model_best.pth diff --git a/video/fake-stormer/model_code/configs/spatial/xception_sbi.yaml b/video/fake-stormer/model_code/configs/spatial/xception_sbi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c1f14d7c86d8e7d7d682571ff845c07f6856bb27 --- /dev/null +++ b/video/fake-stormer/model_code/configs/spatial/xception_sbi.yaml @@ -0,0 +1,169 @@ +TASK: Xception_bin_sbi_abl +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: image + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 28 + PIN_MEMORY: True + IMAGE_SIZE: [299, 299] + HEATMAP_SIZE: [96, 96] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 3 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: gaussian + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 8 # Dynamically random number of frames in each epoch + VAL: 8 + TEST: 32 + DIST: [1., 1.] # Distribution of Real, Fake - 1.,1. by default + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/train_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/val_81_FF++_processed.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: Celeb-DFv2 + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + # ROOT: /home/users/XXX/data/Combine/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + FAKETYPE: [Celeb-real-0.95-0.97-v2, Celeb-synthesis-0.95-0.97-v2, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [Celeb_DFv1-real, Celeb_DFv2-real, DeepFakeDetection, Deepfakes, FaceSwap, method_A, NeuralTextures, + # original_videos, YouTube_DFv1-real, Celeb_DFv1-synthesis, Celeb_DFv2-synthesis, DeepFakeDetection_original, Face2Face, + # fake_test, method_B, original, real_test, YouTube_DFv2-real] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [380, 380, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + #Either Scaling or Cropping, not do at the same time + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.5, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.5, 0.5, 0.5] + std: [0.5, 0.5, 0.5] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: False + TARGET_OVERLAP: True + +MODEL: + type: TopDownDetector + backbone: + type: Xception + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 2048 + heads: + cls: 1 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256] + num_deconv_kernels: [4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + INIT_WEIGHTS: + pretrained: True + +TRAIN: + gpus: [0,1,2,3] + batch_size: 32 + lr: 0.0001 + epochs: 100 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + accumulation_steps: 1 + use_amp: False + loss: + type: BinaryCrossEntropy + reduction: mean + optimizer: SAM + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: LinearDecayLR + milestones: [30] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 4 +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: True + video_level: True + apr: True # Average precision/recall or normal precision/recall + no_shot_preds: 1 + pred_file: /project/home/p200249/XXX/saved_predictions/FakeFormer_DF40_90_10.json # File to save predictions + save_preds: False # VERY CAREFUL + pretrained: logs/15-09-2025/TopDownDetector_Xception_bin_sbi_abl_model_best.pth diff --git a/video/fake-stormer/model_code/configs/temporal/FakeSFormer_base_c0.yaml b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_base_c0.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d8e84ea97c9a09a078b6ecc5706c1a874853f24b --- /dev/null +++ b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_base_c0.yaml @@ -0,0 +1,213 @@ +TASK: C0_ViTBase224_ST_hm100_tempLOC0.2_4SBI_SAM_mp0.01_temp2_0vidAug_0.25_temp_normlmsFT_IN_1e7 +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: video + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 28 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: m_std_normalized + SPLIT_IMAGE: False + DATA: + TYPE: videos # We employ "sequential sampling" as "videos" type for Training while "uniform sampling" as "frames" type for Testing + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 4 # Dynamically random number of frames in each epoch + VAL: 4 + TEST: 4 + NUM_FRAMES: 4 + DIST: [1., 1.] # Distribution of Real, Fake - 1.,1. by default + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/train_FF++_videos_81.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/val_FF++_videos_81.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: Celeb-DFv2 + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + # ROOT: /project/home/p200249/XXX/DeeperForensics/ + # ROOT: /project/home/p200249/XXX/DiffSwap/ + # ROOT: /project/home/p200249/XXX/DF40_test/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [fake_val, real_val] + # FAKETYPE: [DFo_source_videos, DFo_manipulated_videos] + # FAKETYPE: [Real_DiffSwap1, DiffSwap_1] + # FAKETYPE: [blendface, danet, deepfacelab, e4e, e4s, facedancer, faceswap, facevid2vid, fomm, fsgan, heygen, + # hyperreenact, inswap, lia, mcnet, mobileswap, MRAA, one_shot_free, pirender, sadtalker, simswap, tpsm, uniface, + # wav2lip, real_videos] + # FAKETYPE: [mobileswap, real_videos] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.] #Format: [crop_limit, prob] + scale: [0.15, 0.] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: True # Normally only use for temporal modeling + TARGET_OVERLAP: True + MASK_PROB: 0.01 + TEMP_MASKOUT: False + +MODEL: + type: TopDownDetector + backbone: + type: TimeViT + img_size: [224, 224] + patch_size: 16 + embed_dim: 768 + depth: 12 + num_heads: 12 + mlp_ratio: 4 + drop_path_rate: 0.3 + qkv_bias: True + class_token: True + register_token: False + temp_token: True + attention_type: divided_space_time #joint_space_time, space_only + num_frames: 4 + low_level_enhanced: False + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 768 + hm_size: [14, 14] #img_size // patch_size + heads: + hm: 1 + cls: 1 + temp_loc: 4 + # num_deconv_layers: 0 #Config n deconv layers to build the decoder + # num_deconv_filters: [256, 256] + # num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + conv_2direction: False + use_temp_token: True + features: 3D + avg_pool: False + INIT_WEIGHTS: + pretrained: pretrained/mae_pretrain_vit_base.pth + +TRAIN: + gpus: [0,1,2,3] + batch_size: 16 + lr: 0.0000001 + epochs: 120 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + accumulation_steps: 1 + use_amp: False + loss: + type: CombinedMSELoss + use_target_weight: False + cls_lmda: 0.8 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 100 + cstency_lmda: 0 + tmp_loc_lmda: 0.2 + mse_reduction: mean + ce_reduction: mean + use_ce: False + temperature: 2 + optimizer: SAM + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 500 + start_decay: 4 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: False + video_level: True + apr: True # Average precision/recall or normal precision/recall + no_shot_preds: 1 + pred_file: /project/home/p200249/XXX/saved_predictions/fakestormer_DF40_90_10.json # File to save predictions + save_preds: False + pretrained: logs/03-12-2024/TopDownDetector_C0_ViTBase224_ST_hm100_tempLOC0.2_4SBI_SAM_mp0.01_temp2_0vidAug_0.35_temp_normlms_IN_model_best.pth diff --git a/video/fake-stormer/model_code/configs/temporal/FakeSFormer_base_c23.yaml b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_base_c23.yaml new file mode 100644 index 0000000000000000000000000000000000000000..54659169bc7ab430d5a8fc175a839912419f6d98 --- /dev/null +++ b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_base_c23.yaml @@ -0,0 +1,206 @@ +TASK: C23_ViTBase224_ST_hm100_tempLOC0.2_4SBI_SAM_mp0.01_temp2_0vidAug_0.25_temp_0normlms_0CSP_spatialHead +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: video + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c23 + IMAGE_SUFFIX: png + NUM_WORKERS: 32 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: m_std_normalized + SPLIT_IMAGE: False + DATA: + TYPE: videos # We employ "sequential sampling" as "videos" type for Training while "uniform sampling" as "frames" type for Testing (better results) + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 4 # Dynamically random number of frames in each epoch + VAL: 4 + TEST: 4 + NUM_FRAMES: 4 + DIST: [1., 1.] # Distribution of Real, Fake - 1.,1. by default + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/train_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/val_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: Celeb-DFv2 + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + # ROOT: /project/home/p200249/XXX/DeeperForensics/ + # ROOT: /project/home/p200249/XXX/DiffSwap/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [frames_BW_5/original, frames_BW_5/Deepfakes, frames_BW_5/Face2Face, frames_BW_5/FaceSwap, frames_BW_5/NeuralTextures] + FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [DFo_source_videos, DFo_manipulated_videos] + # FAKETYPE: [Real_DiffSwap1, DiffSwap_1] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.] #Format: [crop_limit, prob] + scale: [0.15, 0.] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: True # Normally only use for temporal modeling + TARGET_OVERLAP: True + MASK_PROB: 0.01 + TEMP_MASKOUT: False + +MODEL: + type: TopDownDetector + backbone: + type: TimeViT + img_size: [224, 224] + patch_size: 16 + embed_dim: 768 + depth: 12 + num_heads: 12 + mlp_ratio: 4 + drop_path_rate: 0.3 + qkv_bias: True + class_token: True + register_token: False + temp_token: True + attention_type: divided_space_time #joint_space_time, space_only + num_frames: 4 + low_level_enhanced: False + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 768 + hm_size: [14, 14] #img_size // patch_size + heads: + hm: 1 + cls: 1 + temp_loc: 4 + # num_deconv_layers: 0 #Config n deconv layers to build the decoder + # num_deconv_filters: [256, 256] + # num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + conv_2direction: False + use_temp_token: True + features: 3D + avg_pool: False + INIT_WEIGHTS: + pretrained: pretrained/mae_pretrain_vit_base.pth + +TRAIN: + gpus: [0,1,2,3] + batch_size: 16 + lr: 0.0000001 + epochs: 200 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + accumulation_steps: 1 + use_amp: False + loss: + type: CombinedMSELoss + use_target_weight: False + cls_lmda: 0.8 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 100 + cstency_lmda: 0 + tmp_loc_lmda: 0.2 + mse_reduction: mean + ce_reduction: mean + use_ce: False + temperature: 2 + optimizer: SAM + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 500 + start_decay: 4 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: False + video_level: True + apr: True # Average precision/recall or normal precision/recall + no_shot_preds: 1 + pred_file: /project/home/p200249/XXX/saved_predictions/fakestormer_DF40_90_10.json # File to save predictions + save_preds: False + pretrained: 'logs/20-10-2024_FakeSTomer_gassian_abl/TopDownDetector_C23_ViTBase224_ST_hm100_tempLOC0.2_4SBI_SAM_mp0.01_temp2_0vidAug_0.35_temp_normlms_gaussian_abl_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/temporal/FakeSFormer_base_c23_224p8.yaml b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_base_c23_224p8.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3679e8e20865552398673c876e26debc34a5cf96 --- /dev/null +++ b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_base_c23_224p8.yaml @@ -0,0 +1,200 @@ +TASK: C23_ViTBase224p8_ST_hm10_tempLOC0.2_4SBI_SAM_mp0.01_temp2_0vidAug_0.35_temp_normlms_normREAL +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: video + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c23 + IMAGE_SUFFIX: png + NUM_WORKERS: 32 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [28, 28] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: m_std_normalized + SPLIT_IMAGE: False + DATA: + TYPE: videos # We employ "sequential sampling" as "videos" type for Training while "uniform sampling" as "frames" type for Testing + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 4 # Dynamically random number of frames in each epoch + VAL: 4 + TEST: 4 + NUM_FRAMES: 4 + DIST: [1., 1.] # Distribution of Real, Fake - 1.,1. by default + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/train_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/val_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: Celeb-DFv2 + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + # ROOT: /project/home/p200249/XXX/DeeperForensics/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [frames_BW_5/original, frames_BW_5/Deepfakes, frames_BW_5/Face2Face, frames_BW_5/FaceSwap, frames_BW_5/NeuralTextures] + FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [DFo_source_videos, DFo_manipulated_videos] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.] #Format: [crop_limit, prob] + scale: [0.15, 0.] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: True # Normally only use for temporal modeling + TARGET_OVERLAP: True + MASK_PROB: 0.01 + TEMP_MASKOUT: False + +MODEL: + type: TopDownDetector + backbone: + type: TimeViT + img_size: [224, 224] + patch_size: 8 + embed_dim: 768 + depth: 12 + num_heads: 12 + mlp_ratio: 4 + drop_path_rate: 0.3 + qkv_bias: True + class_token: True + register_token: False + temp_token: True + attention_type: divided_space_time #joint_space_time, space_only + num_frames: 4 + low_level_enhanced: False + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 768 + hm_size: [14, 14] #img_size // patch_size + heads: + hm: 1 + cls: 1 + temp_loc: 4 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256, 256] + num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + conv_2direction: False + use_temp_token: True + features: 3D + INIT_WEIGHTS: + pretrained: pretrained/dino_vitbase8_pretrain.pth + +TRAIN: + gpus: [0,1,2,3] + batch_size: 2 + lr: 0.00005 + epochs: 100 + begin_epoch: -1 + warm_up: 6 + every_val_epochs: 1 + accumulation_steps: 8 + use_amp: True + loss: + type: CombinedMSELoss + use_target_weight: False + cls_lmda: 0.8 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 10 + cstency_lmda: 0 + tmp_loc_lmda: 0.2 + mse_reduction: mean + ce_reduction: mean + use_ce: False + temperature: 2 + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 2.5 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: False + video_level: True + no_shot_preds: 1 + pretrained: 'logs/13-09-2024_BestOVA_FakeSTormer/TopDownDetector_C23_ViTBase224_ST_hm100_tempLOC0.2_4SBI_SAM_mp0.01_temp2_0vidAug_0.35_temp_normlms_normREAL_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/temporal/FakeSFormer_base_c40.yaml b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_base_c40.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2013d6fd2673904a1227e86530ddb2425c0ac83f --- /dev/null +++ b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_base_c40.yaml @@ -0,0 +1,203 @@ +TASK: C40_ViTBase224_ST_hm100_tempLOC0.1_4SBI_SAM_mp0.01_temp2_0vidAug_0.35_temp_normlms_normREALtarget +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: video + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c40 + IMAGE_SUFFIX: png + NUM_WORKERS: 32 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: m_std_normalized + SPLIT_IMAGE: False + DATA: + TYPE: videos # We employ "sequential sampling" as "videos" type for Training while "uniform sampling" as "frames" type for Testing (better results) + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 4 # Dynamically random number of frames in each epoch + VAL: 4 + TEST: 4 + NUM_FRAMES: 4 + DIST: [1., 1.] # Distribution of Real, Fake - 1.,1. by default + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c40/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c40/train_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c40/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c40/val_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: FF++ + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c40/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + # ROOT: /project/home/p200249/XXX/DeeperForensics/ + FROM_FILE: False + FAKETYPE: [original, Deepfakes] + # FAKETYPE: [videos_BW_5/original, videos_BW_5/Deepfakes, videos_BW_5/Face2Face, videos_BW_5/FaceSwap, videos_BW_5/NeuralTextures] + # FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [DFo_source_videos, DFo_manipulated_videos] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.] #Format: [crop_limit, prob] + scale: [0.15, 0.] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: True # Normally only use for temporal modeling + TARGET_OVERLAP: True + MASK_PROB: 0.01 + TEMP_MASKOUT: False + +MODEL: + type: TopDownDetector + backbone: + type: TimeViT + img_size: [224, 224] + patch_size: 16 + embed_dim: 768 + depth: 12 + num_heads: 12 + mlp_ratio: 4 + drop_path_rate: 0.3 + qkv_bias: True + class_token: True + register_token: False + temp_token: True + attention_type: divided_space_time #joint_space_time, space_only + num_frames: 4 + low_level_enhanced: False + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 768 + hm_size: [14, 14] #img_size // patch_size + heads: + hm: 1 + cls: 1 + temp_loc: 4 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256, 256] + num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + conv_2direction: False + use_temp_token: True + features: 3D + INIT_WEIGHTS: + pretrained: pretrained/mae_pretrain_vit_base.pth + +TRAIN: + gpus: [0,1,2,3] + batch_size: 16 + lr: 0.00005 + epochs: 100 + begin_epoch: -1 + warm_up: 3 + every_val_epochs: 1 + accumulation_steps: 1 + use_amp: False + loss: + type: CombinedMSELoss + use_target_weight: False + cls_lmda: 0.9 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 100 + cstency_lmda: 0 + tmp_loc_lmda: 0.1 + mse_reduction: mean + ce_reduction: mean + use_ce: False + temperature: 2 + optimizer: SAM + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 2.5 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: False + video_level: True + no_shot_preds: 1 + pretrained: 'logs/13-09-2024/TopDownDetector_C40_ViTBase224_ST_hm100_tempLOC0.1_4SBI_SAM_mp0.01_temp2_0vidAug_0.35_temp_normlms_normREALtarget_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/temporal/FakeSFormer_large_c23.yaml b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_large_c23.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bbc6cbd4e406e56f67175357ac9767cbd510b666 --- /dev/null +++ b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_large_c23.yaml @@ -0,0 +1,198 @@ +TASK: ViTLarge224_ST_hm100_tempLOC0.2_4SBI_AdamW_mp0.01_temp2_reload_0vidAug +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: video + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c23 + IMAGE_SUFFIX: png + NUM_WORKERS: 32 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: m_std_normalized + SPLIT_IMAGE: False + DATA: + TYPE: videos # We employ "sequential sampling" as "videos" type for Training while "uniform sampling" as "frames" type for Testing (better results) + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 4 # Dynamically random number of frames in each epoch + VAL: 4 + TEST: 4 + NUM_FRAMES: 4 + DIST: [1., 1.] # Distribution of Real, Fake - 1.,1. by default + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/train_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/val_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: Celeb-DFv2 + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DeeperForensics/ + FROM_FILE: False + # FAKETYPE: [original, Face2Face, FaceSwap, NeuralTextures] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [DFo_source_videos, DFo_manipulated_videos] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.] #Format: [crop_limit, prob] + scale: [0.15, 0.] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: True # Normally only use for temporal modeling + TARGET_OVERLAP: True + MASK_PROB: 0.01 + TEMP_MASKOUT: False + +MODEL: + type: TopDownDetector + backbone: + type: TimeViT + img_size: [224, 224] + patch_size: 16 + embed_dim: 1024 + depth: 24 + num_heads: 16 + mlp_ratio: 4 + drop_path_rate: 0.3 + qkv_bias: True + class_token: True + register_token: False + temp_token: True + attention_type: divided_space_time #joint_space_time, space_only + num_frames: 4 + low_level_enhanced: False + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 1024 + hm_size: [14, 14] #img_size // patch_size + heads: + hm: 1 + cls: 1 + temp_loc: 4 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256, 256] + num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + conv_2direction: False + use_temp_token: True + features: 3D + INIT_WEIGHTS: + pretrained: pretrained/mae_pretrain_vit_large.pth + +TRAIN: + gpus: [0,1,2,3] + batch_size: 4 + lr: 0.00005 + epochs: 100 + begin_epoch: -1 + warm_up: 6 + every_val_epochs: 1 + accumulation_steps: 4 + use_amp: True + loss: + type: CombinedMSELoss + use_target_weight: False + cls_lmda: 0.8 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 100 + cstency_lmda: 0 + tmp_loc_lmda: 0.2 + mse_reduction: mean + ce_reduction: mean + use_ce: False + temperature: 2 + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 4 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: False + video_level: True + no_shot_preds: 1 + pretrained: '' diff --git a/video/fake-stormer/model_code/configs/temporal/FakeSFormer_small.yaml b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_small.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ffd117d57f7a18cb55dc0da3f66b000142c04ba9 --- /dev/null +++ b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_small.yaml @@ -0,0 +1,187 @@ +TASK: ViTSmall112_ST_hm10_tempLOC0.1_4SBI_DymIntens_MSE_AdamW_5e5_boost1_mp0.01_discrete_fsdyn_fullSA_m_std +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: False + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: video + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c0 + IMAGE_SUFFIX: png + NUM_WORKERS: 32 + PIN_MEMORY: True + IMAGE_SIZE: [112, 112] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: m_std_normalized + SPLIT_IMAGE: False + DATA: + TYPE: videos # frames or videos + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 4 # Dynamically random number of frames in each epoch + VAL: 4 + TEST: 4 + NUM_FRAMES: 4 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/train_FF++_videos_81.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c0/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/val_FF++_videos_81.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: FF++ + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200328/XXX/data/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200328/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200328/XXX/data/DFDCP/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + FROM_FILE: False + FAKETYPE: [original, NeuralTextures] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [112, 112, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0 + cropping: [0.15, 0] #Format: [crop_limit, prob] + scale: [0.15, 0] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: True # Normally only use for temporal modeling + TARGET_OVERLAP: True + MASK_PROB: 0.01 + +MODEL: + type: TopDownDetector + backbone: + type: TimeViT + img_size: [112, 112] + patch_size: 8 + embed_dim: 384 + depth: 12 + num_heads: 6 + mlp_ratio: 4 + drop_path_rate: 0.3 + qkv_bias: True + class_token: True + register_token: True + attention_type: divided_space_time #joint_space_time, space_only + num_frames: 8 + pretrained: pretrained/dino_deitsmall8_pretrain.pth + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 384 + hm_size: [14, 14] #img_size // patch_size + heads: + hm: 1 + cls: 1 + temp_loc: 4 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256, 256] + num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + conv_2direction: True + use_temp_token: True + features: 3D + +TRAIN: + gpus: [0,1,2,3] + batch_size: 4 + lr: 0.00005 + epochs: 200 + begin_epoch: -1 + warm_up: 3 + every_val_epochs: 1 + loss: + type: CombinedMSELoss + use_target_weight: False + cls_lmda: 1 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 10 + cstency_lmda: 0 + tmp_loc_lmda: 0.1 + mse_reduction: mean + ce_reduction: mean + use_ce: False + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 4 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: False + video_level: True + no_shot_preds: 3 + pretrained: 'logs/18-07-2024/TopDownDetector_ViTSmall112_ST_hm10_tempLOC0.1_4SBI_DymIntens_MSE_SAM(AdamW)_5e5_boost1_mp0.01_discrete_fsdyn_fullSA_m_std_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/temporal/FakeSFormer_small_c23.yaml b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_small_c23.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ae2e2a4fe3176b72bce90bab691015235ced9065 --- /dev/null +++ b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_small_c23.yaml @@ -0,0 +1,195 @@ +TASK: ViTSmall112_ST_hm100_tempLOC0.3_8SBI_DymIntens_MSE_SAM(AdamW)_5e5_boost1_mp0.01_fsdyn_fullSA_m_std_c23_no_compress +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: False + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: video + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c23 + IMAGE_SUFFIX: png + NUM_WORKERS: 32 + PIN_MEMORY: True + IMAGE_SIZE: [112, 112] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: m_std_normalized + SPLIT_IMAGE: False + DATA: + TYPE: videos # frames or videos + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 4 # Dynamically random number of frames in each epoch + VAL: 4 + TEST: 4 + NUM_FRAMES: 4 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/FaceForensics++/c23/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/train_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/FaceForensics++/c23/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/val_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: FF++ + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + # ROOT: /project/home/p200328/XXX/data/DeeperForensics/ + FROM_FILE: False + # FAKETYPE: [original, FaceSwap] + FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [DFo_source_videos, DFo_manipulated_videos] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [112, 112, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0 + cropping: [0.15, 0] #Format: [crop_limit, prob] + scale: [0.15, 0] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: True # Normally only use for temporal modeling + TARGET_OVERLAP: True + MASK_PROB: 0.01 + +MODEL: + type: TopDownDetector + backbone: + type: TimeViT + img_size: [112, 112] + patch_size: 8 + embed_dim: 384 + depth: 12 + num_heads: 6 + mlp_ratio: 4 + drop_path_rate: 0.3 + qkv_bias: True + class_token: True + register_token: True + attention_type: divided_space_time #joint_space_time, space_only + num_frames: 8 + pretrained: pretrained/dino_deitsmall8_pretrain.pth + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 384 + hm_size: [14, 14] #img_size // patch_size + heads: + hm: 1 + cls: 1 + temp_loc: 4 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256, 256] + num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + conv_2direction: True + use_temp_token: True + features: 3D + +TRAIN: + gpus: [0,1,2,3] + batch_size: 4 + lr: 0.00005 + epochs: 100 + begin_epoch: -1 + warm_up: 3 + every_val_epochs: 1 + accumulation_steps: 1 + loss: + type: CombinedMSELoss + use_target_weight: False + cls_lmda: 0.7 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 100 + cstency_lmda: 0 + tmp_loc_lmda: 0.3 + mse_reduction: mean + ce_reduction: mean + use_ce: False + optimizer: SAM + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 2 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: True + video_level: True + no_shot_preds: 1 + pretrained: 'logs/18-07-2024/TopDownDetector_ViTSmall112_ST_hm10_tempLOC0.1_8SBI_DymIntens_MSE_SAM(AdamW)_5e5_boost1_mp0.01_discrete_fsdyn_fullSA_m_std_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/temporal/FakeSFormer_small_c23_224p16.yaml b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_small_c23_224p16.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d86bd1908f53d9a5c63a9ebc8c6c2a7b93993d37 --- /dev/null +++ b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_small_c23_224p16.yaml @@ -0,0 +1,194 @@ +TASK: ViTSmall224_hm100_tempLOC0.5_4SBI_DymIntens_MSE_SAM(AdamW)_mp0.01_fsdyn_fullSA_m_std_c23_0aug_newmaskdeform_test +PRECISION: float +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: False + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: video + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c23 + IMAGE_SUFFIX: png + NUM_WORKERS: 32 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: m_std_normalized + SPLIT_IMAGE: False + DATA: + TYPE: videos # frames or videos + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 4 # Dynamically random number of frames in each epoch + VAL: 4 + TEST: 4 + NUM_FRAMES: 4 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/FaceForensics++/c23/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/train_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/FaceForensics++/c23/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/val_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: FF++ + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /project/home/p200328/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200328/XXX/data/DeeperForensics/ + FROM_FILE: False + FAKETYPE: [original, NeuralTextures] + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [DFo_source_videos, DFo_manipulated_videos] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0 + cropping: [0.15, 0] #Format: [crop_limit, prob] + scale: [0.15, 0] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: True # Normally only use for temporal modeling + TARGET_OVERLAP: True + MASK_PROB: 0.01 + +MODEL: + type: TopDownDetector + backbone: + type: TimeViT + img_size: [224, 224] + patch_size: 16 + embed_dim: 384 + depth: 12 + num_heads: 6 + mlp_ratio: 4 + drop_path_rate: 0.3 + qkv_bias: True + class_token: True + register_token: True + attention_type: divided_space_time #joint_space_time, space_only + num_frames: 4 + pretrained: pretrained/dino_deitsmall16_pretrain.pth + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 384 + hm_size: [14, 14] #img_size // patch_size + heads: + hm: 1 + cls: 1 + temp_loc: 4 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256, 256] + num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + conv_2direction: True + use_temp_token: True + features: 3D + +TRAIN: + gpus: [0,1,2,3] + batch_size: 16 + lr: 0.00005 + epochs: 100 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + accumulation_steps: 1 + loss: + type: CombinedMSELoss + use_target_weight: False + cls_lmda: 0.5 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 100 + cstency_lmda: 0 + tmp_loc_lmda: 0.5 + mse_reduction: mean + ce_reduction: mean + use_ce: False + optimizer: SAM + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 4 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: False + video_level: True + no_shot_preds: 1 + pretrained: 'logs/04-08-2024/TopDownDetector_ViTSmall224_hm100_tempLOC0.3_4SBI_DymIntens_MSE_SAM(AdamW)_mp0.01_fsdyn_fullSA_m_std_c23_0comp_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/temporal/FakeSFormer_small_c23_224p8.yaml b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_small_c23_224p8.yaml new file mode 100644 index 0000000000000000000000000000000000000000..02b42eb7b7bb408ce144f9e3a8aa4ba854f27c13 --- /dev/null +++ b/video/fake-stormer/model_code/configs/temporal/FakeSFormer_small_c23_224p8.yaml @@ -0,0 +1,197 @@ +TASK: ViTSmall224p8_ST_hm100_tempLOC0.5_4SBI_SAM_mp0.01_temp2_reload_0vidAug_0.5_harderBI +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: False + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: video + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c23 + IMAGE_SUFFIX: png + NUM_WORKERS: 32 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [28, 28] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: m_std_normalized + SPLIT_IMAGE: False + DATA: + TYPE: videos # frames or videos + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 4 # Dynamically random number of frames in each epoch + VAL: 4 + TEST: 4 + NUM_FRAMES: 4 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/FaceForensics++/c23/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/train_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /home/users/XXX/data/FaceForensics++/c23/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/val_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: FF++ + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /project/home/p200328/XXX/data/DFDCP/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200328/XXX/data/DeeperForensics/ + FROM_FILE: False + # FAKETYPE: [original, NeuralTextures] + FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [DFo_source_videos, DFo_manipulated_videos] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [112, 112, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0 + cropping: [0.15, 0] #Format: [crop_limit, prob] + scale: [0.15, 0] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: True # Normally only use for temporal modeling + TARGET_OVERLAP: True + MASK_PROB: 0.01 + TEMP_MASKOUT: False + +MODEL: + type: TopDownDetector + backbone: + type: TimeViT + img_size: [224, 224] + patch_size: 8 + embed_dim: 384 + depth: 12 + num_heads: 6 + mlp_ratio: 4 + drop_path_rate: 0.3 + qkv_bias: True + class_token: True + register_token: True + attention_type: divided_space_time #joint_space_time, space_only + num_frames: 4 + low_level_enhanced: False + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 384 + hm_size: [28, 28] #img_size // patch_size + heads: + hm: 1 + cls: 1 + temp_loc: 4 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256, 256] + num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + conv_2direction: True + use_temp_token: True + features: 3D + INIT_WEIGHTS: + pretrained: pretrained/dino_deitsmall8_pretrain.pth + +TRAIN: + gpus: [0,1,2,3] + batch_size: 8 + lr: 0.00005 + epochs: 100 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + accumulation_steps: 1 + use_amp: False + loss: + type: CombinedMSELoss + use_target_weight: False + cls_lmda: 0.5 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 100 + cstency_lmda: 0 + tmp_loc_lmda: 0.5 + mse_reduction: mean + ce_reduction: mean + use_ce: False + temperature: 2 + optimizer: SAM + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: False + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 4 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: False + video_level: True + no_shot_preds: 1 + pretrained: 'logs/28-08-2024/TopDownDetector_ViTSmall224p8_ST_hm10_tempLOC0.5_4SBI_SAM_mp0.01_temp2_reload_0vidAug_0.5_harderBI_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/temporal/FakeSwin3D_base_c0.yaml b/video/fake-stormer/model_code/configs/temporal/FakeSwin3D_base_c0.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ff16e2be5a54eb7de1f39ff4d1fa40678e06ded1 --- /dev/null +++ b/video/fake-stormer/model_code/configs/temporal/FakeSwin3D_base_c0.yaml @@ -0,0 +1,206 @@ +TASK: C0_Swin3DBase224_hm10_128CST100_EFPN_C3_32SBI_AdamW_temp2_0.25_normlms_1e8 +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: video + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c23 + IMAGE_SUFFIX: png + NUM_WORKERS: 28 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [56, 56] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: m_std_normalized + SPLIT_IMAGE: False + DATA: + TYPE: videos + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 32 # Dynamically random number of frames in each epoch + VAL: 32 + TEST: 32 + NUM_FRAMES: 32 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/train_FF++_videos_81.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c0/val_FF++_videos_81.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: Celeb-DFv2 + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + # ROOT: /project/home/p200249/XXX/DeeperForensics/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [frames_BW_5/original, frames_BW_5/Deepfakes, frames_BW_5/Face2Face, frames_BW_5/FaceSwap, frames_BW_5/NeuralTextures] + FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [DFo_source_videos, DFo_manipulated_videos] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.] #Format: [crop_limit, prob] + scale: [0.15, 0.] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: True # Normally only use for temporal modeling + TARGET_OVERLAP: True + MASK_PROB: 0. + TEMP_MASKOUT: False + +MODEL: + type: TopDownDetector + backbone: + type: SwinTransformer3D + embed_dim: 128 + patch_size: [2,4,4] + depths: [2,2,18,2] + num_heads: [4,8,16,32] + window_size: [8,7,7] + mlp_ratio: 4 + qkv_bias: True + drop_rate: 0. + attn_drop_rate: 0. + drop_path_rate: 0.2 + patch_norm: True + pretrained2d: True + pretrained: pretrained/swin_base_patch4_window7_224_22k.pth + neck: + type: EFPN3D + in_channels: 1024 + num_deconv_layers: 4 #Config n deconv layers to build the decoder + num_deconv_filters: [1024, 512, 256, 128] + num_deconv_kernels: [[1,1,1], [1,4,4], [1,4,4], [4,4,4]] + num_deconv_strides: [[1,1,1], [1,2,2], [1,2,2], [2,2,2]] + efpn: True + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 128 + hm_size: [14, 14] #img_size // patch_size + heads: + hm: 1 + cls: 1 + cstency: 128 + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + conv_2direction: False + use_temp_token: True + features: 3D + avg_pool: True + act: GELU + INIT_WEIGHTS: + pretrained: pretrained/swin_base_patch4_window7_224_22k.pth + +TRAIN: + gpus: [0,1,2,3] + batch_size: 2 + lr: 0.0000001 + epochs: 200 + begin_epoch: -1 + warm_up: 10 + every_val_epochs: 1 + accumulation_steps: 8 + use_amp: True + loss: + type: CombinedMSELoss + use_target_weight: False + cls_lmda: 1 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 10 + cstency_lmda: 100 + tmp_loc_lmda: 0 + mse_reduction: mean + ce_reduction: mean + use_ce: False + feature: 3D + temperature: 2 + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 500 + start_decay: 4 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: False + video_level: True + no_shot_preds: 1 + pretrained: 'logs/22-01-2025/TopDownDetector_C23_Swin3DBase224_hm10_256CST100_EFPN_C3_32SBI_AdamW_temp2_0.25_normlms_normREAL_1e8_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/temporal/FakeSwin3D_base_c23.yaml b/video/fake-stormer/model_code/configs/temporal/FakeSwin3D_base_c23.yaml new file mode 100644 index 0000000000000000000000000000000000000000..27b483d65a8e64ce0b1aa0b7cba76d6750b2d9eb --- /dev/null +++ b/video/fake-stormer/model_code/configs/temporal/FakeSwin3D_base_c23.yaml @@ -0,0 +1,206 @@ +TASK: C23_Swin3DBase224_hm10_128CST100_EFPN_C3_4SBI_AdamW_temp2_0.25_normlms_1e8 +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: video + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c23 + IMAGE_SUFFIX: png + NUM_WORKERS: 28 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [56, 56] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: m_std_normalized + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 32 # Dynamically random number of frames in each epoch + VAL: 32 + TEST: 32 + NUM_FRAMES: 32 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/train_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/val_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: DFDC + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + # ROOT: /project/home/p200249/XXX/DeeperForensics/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes] + # FAKETYPE: [frames_BW_5/original, frames_BW_5/Deepfakes, frames_BW_5/Face2Face, frames_BW_5/FaceSwap, frames_BW_5/NeuralTextures] + # FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [DFo_source_videos, DFo_manipulated_videos] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.] #Format: [crop_limit, prob] + scale: [0.15, 0.] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: True # Normally only use for temporal modeling + TARGET_OVERLAP: True + MASK_PROB: 0. + TEMP_MASKOUT: False + +MODEL: + type: TopDownDetector + backbone: + type: SwinTransformer3D + embed_dim: 128 + patch_size: [2,4,4] + depths: [2,2,18,2] + num_heads: [4,8,16,32] + window_size: [8,7,7] + mlp_ratio: 4 + qkv_bias: True + drop_rate: 0. + attn_drop_rate: 0. + drop_path_rate: 0.2 + patch_norm: True + pretrained2d: True + pretrained: pretrained/swin_base_patch4_window7_224_22k.pth + neck: + type: EFPN3D + in_channels: 1024 + num_deconv_layers: 4 #Config n deconv layers to build the decoder + num_deconv_filters: [1024, 512, 256, 128] + num_deconv_kernels: [[1,1,1], [1,4,4], [1,4,4], [4,4,4]] + num_deconv_strides: [[1,1,1], [1,2,2], [1,2,2], [2,2,2]] + efpn: True + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 128 + hm_size: [14, 14] #img_size // patch_size + heads: + hm: 1 + cls: 1 + cstency: 128 + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + conv_2direction: False + use_temp_token: True + features: 3D + avg_pool: True + act: GELU + INIT_WEIGHTS: + pretrained: pretrained/swin_base_patch4_window7_224_22k.pth + +TRAIN: + gpus: [0,1,2,3] + batch_size: 8 + lr: 0.00000001 + epochs: 100 + begin_epoch: -1 + warm_up: 5 + every_val_epochs: 1 + accumulation_steps: 1 + use_amp: False + loss: + type: CombinedMSELoss + use_target_weight: False + cls_lmda: 1 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 10 + cstency_lmda: 100 + tmp_loc_lmda: 0 + mse_reduction: mean + ce_reduction: mean + use_ce: False + feature: 3D + temperature: 2 + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: False + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 5000 + start_decay: 4 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: False + video_level: True + no_shot_preds: 1 + pretrained: 'logs/18-03-2025/TopDownDetector_C23_Swin3DBase224_hm10_128CST100_EFPN_C3_32SBI_AdamW_temp2_0.25_normlms_1e7_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/temporal/ResNet3D_EFPN3D_hm3D_c23.yaml b/video/fake-stormer/model_code/configs/temporal/ResNet3D_EFPN3D_hm3D_c23.yaml new file mode 100644 index 0000000000000000000000000000000000000000..052f30491258b142cb7252a70c34588ea7661c6a --- /dev/null +++ b/video/fake-stormer/model_code/configs/temporal/ResNet3D_EFPN3D_hm3D_c23.yaml @@ -0,0 +1,204 @@ +TASK: C23_ResNet3D224_hm10_256CST100_EFPN_C2_32SBI_AdamW_temp2_0.25_normlms_1e7_mstd +PRECISION: float64 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: True + +DATASET: + type: FakeSFormerSBI + DATA_TYPE: video + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c23 + IMAGE_SUFFIX: png + NUM_WORKERS: 32 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [56, 56] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: m_std_normalized + SPLIT_IMAGE: False + DATA: + TYPE: videos + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 32 # Dynamically random number of frames in each epoch + VAL: 32 + TEST: 32 + NUM_FRAMES: 32 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/train_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: True + FAKETYPE: [Deepfakes, Face2Face, FaceSwap, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/val_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: Celeb-DFv2 + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /data/deepfake_cluster/datasets_df/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /data/deepfake_cluster/datasets_df/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /data/deepfake_cluster/datasets_df/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + # ROOT: /data/deepfake_cluster/datasets_df/DFW/ + # ROOT: /project/home/p200249/XXX/DeeperForensics/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes, Face2Face, FaceSwap, NeuralTextures] + # FAKETYPE: [original, FaceShifter] + # FAKETYPE: [videos_BW_5/original, videos_BW_5/Deepfakes, videos_BW_5/Face2Face, videos_BW_5/FaceSwap, videos_BW_5/NeuralTextures] + FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake_method, real_method] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [DFo_source_videos, DFo_manipulated_videos] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 0] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.] #Format: [crop_limit, prob] + scale: [0.15, 0.] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + DEBUG: False + DYNAMIC_FXRAY: True + DYNAMIC_BLENDING_PROB: True # Normally only use for temporal modeling + TARGET_OVERLAP: True + MASK_PROB: 0. + TEMP_MASKOUT: False + +MODEL: + type: TopDownDetector + backbone: + type: ResNet3D + block: Bottleneck + layers: [3, 4, 6, 3] + block_inplanes: [64, 128, 256, 512] + neck: + type: EFPN3D + in_channels: 2048 + num_deconv_layers: 4 #Config n deconv layers to build the decoder + num_deconv_filters: [1024, 512, 256, 128] + num_deconv_kernels: [[4,4,4], [4,4,4], [4,4,4], [4,1,1]] + num_deconv_strides: [[2,2,2], [2,2,2], [2,2,2], [2,1,1]] + efpn: True + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 128 + heads: + cls: 1 + hm: 1 + cstency: 128 + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + conv_2direction: False + use_temp_token: True + features: 3D + avg_pool: True + act: RELU + INIT_WEIGHTS: + pretrained: pretrained/r3d50_KS_200ep.pth + +TRAIN: + gpus: [0,1,2,3] + batch_size: 4 + lr: 0.0000001 + epochs: 200 + begin_epoch: -1 + warm_up: 10 + every_val_epochs: 1 + accumulation_steps: 4 + use_amp: True + loss: + type: CombinedMSELoss + use_target_weight: False + cls_lmda: 1 + dst_hm_cls_lmda: 0 + offset_lmda: 0 + hm_lmda: 10 + cstency_lmda: 100 + tmp_loc_lmda: 0 + mse_reduction: mean + ce_reduction: mean + use_ce: False + feature: 3D + temperature: 2 + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 500 + start_decay: 4 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: False + video_level: True + no_shot_preds: 1 + pretrained: 'logs/30-09-2024/TopDownDetector_C23_Res3D224_hm0_tempLOC0_8SBI_Adam_mp0.01_temp1_0vidAug_0.35_normlms_SBV_ablation_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/temporal/bin_cls/ResNet3D_c23.yaml b/video/fake-stormer/model_code/configs/temporal/bin_cls/ResNet3D_c23.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ef2fceea9d28e80ea3f68cd244a120683f174dc8 --- /dev/null +++ b/video/fake-stormer/model_code/configs/temporal/bin_cls/ResNet3D_c23.yaml @@ -0,0 +1,166 @@ +TASK: C23_Res3D224_binary_8SBI_Adam_0vidAug_Deepfakes +PRECISION: float32 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: False + +DATASET: + type: BinaryFaceForensic + DATA_TYPE: video + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c23 + IMAGE_SUFFIX: png + NUM_WORKERS: 32 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: m_std_normalized + SPLIT_IMAGE: False + DATA: + TYPE: frames + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 8 # Dynamically random number of frames in each epoch + VAL: 8 + TEST: 8 + NUM_FRAMES: 8 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: False + FAKETYPE: [original, Deepfakes] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/train_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: False + FAKETYPE: [original, Deepfakes] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/val_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: DFW + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + ROOT: /project/home/p200249/XXX/DFW/ + # ROOT: /project/home/p200249/XXX/DeeperForensics/ + FROM_FILE: False + # FAKETYPE: [original, Deepfakes, FaceSwap, Face2Face, NeuralTextures] + # FAKETYPE: [videos_BW_5/original, videos_BW_5/Deepfakes, videos_BW_5/Face2Face, videos_BW_5/FaceSwap, videos_BW_5/NeuralTextures] + # FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + FAKETYPE: [fake_test, real_test] + # FAKETYPE: [DFo_source_videos, DFo_manipulated_videos] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 1] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + +MODEL: + type: TopDownDetector + backbone: + type: ResNet3D + block: Bottleneck + layers: [3, 4, 6, 3] + block_inplanes: [64, 128, 256, 512] + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 2048 + heads: + cls: 1 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256, 256] + num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + INIT_WEIGHTS: + pretrained: null + +TRAIN: + gpus: [0,1,2,3] + batch_size: 16 + lr: 0.00005 + epochs: 100 + begin_epoch: -1 + warm_up: 3 + every_val_epochs: 1 + accumulation_steps: 1 + use_amp: False + loss: + type: BinaryCrossEntropy + reduction: mean + optimizer: Adam + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 2.5 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: False + video_level: True + no_shot_preds: 1 + pretrained: 'logs/17-09-2024/TopDownDetector_C23_Res3D224_binary_8SBI_Adam_0vidAug_Deepfakes_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/temporal/bin_cls/ResNet3D_c40.yaml b/video/fake-stormer/model_code/configs/temporal/bin_cls/ResNet3D_c40.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1233deee5e9cbbd51b94e3f1e189a224ccd5009c --- /dev/null +++ b/video/fake-stormer/model_code/configs/temporal/bin_cls/ResNet3D_c40.yaml @@ -0,0 +1,166 @@ +TASK: C40_Res3D224_binary_8SBI_Adam_0vidAug_NeuralTextures +PRECISION: float32 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: False + +DATASET: + type: BinaryFaceForensic + DATA_TYPE: video + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c40 + IMAGE_SUFFIX: png + NUM_WORKERS: 32 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: m_std_normalized + SPLIT_IMAGE: False + DATA: + TYPE: videos + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 8 # Dynamically random number of frames in each epoch + VAL: 8 + TEST: 8 + NUM_FRAMES: 8 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c40/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: False + FAKETYPE: [original, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/train_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c40/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: False + FAKETYPE: [original, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/val_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: FF++ + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c40/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + # ROOT: /project/home/p200249/XXX/DeeperForensics/ + FROM_FILE: False + FAKETYPE: [original, Deepfakes] + # FAKETYPE: [videos_BW_5/original, videos_BW_5/Deepfakes, videos_BW_5/Face2Face, videos_BW_5/FaceSwap, videos_BW_5/NeuralTextures] + # FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [DFo_source_videos, DFo_manipulated_videos] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 1] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + +MODEL: + type: TopDownDetector + backbone: + type: ResNet3D + block: Bottleneck + layers: [3, 4, 6, 3] + block_inplanes: [64, 128, 256, 512] + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 2048 + heads: + cls: 1 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256, 256] + num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + INIT_WEIGHTS: + pretrained: null + +TRAIN: + gpus: [0,1,2,3] + batch_size: 16 + lr: 0.00005 + epochs: 100 + begin_epoch: -1 + warm_up: 3 + every_val_epochs: 1 + accumulation_steps: 1 + use_amp: False + loss: + type: BinaryCrossEntropy + reduction: mean + optimizer: Adam + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 2.5 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: False + video_level: True + no_shot_preds: 1 + pretrained: 'logs/02-10-2024/TopDownDetector_C40_Res3D224_binary_8SBI_Adam_0vidAug_NeuralTextures_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/temporal/bin_cls/TimeSFormer_base_c23.yaml b/video/fake-stormer/model_code/configs/temporal/bin_cls/TimeSFormer_base_c23.yaml new file mode 100644 index 0000000000000000000000000000000000000000..94df7b21216391e0beb5dadc75a9785179cf5994 --- /dev/null +++ b/video/fake-stormer/model_code/configs/temporal/bin_cls/TimeSFormer_base_c23.yaml @@ -0,0 +1,181 @@ +TASK: C23_TimeBase224_ST_binary_8SBI_AdamW_0vidAug_temp_FaceSwap +PRECISION: float32 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: False + +DATASET: + type: BinaryFaceForensic + DATA_TYPE: video + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c23 + IMAGE_SUFFIX: png + NUM_WORKERS: 32 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: m_std_normalized + SPLIT_IMAGE: False + DATA: + TYPE: videos + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 8 # Dynamically random number of frames in each epoch + VAL: 8 + TEST: 8 + NUM_FRAMES: 8 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: False + FAKETYPE: [original, FaceSwap] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/train_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: False + FAKETYPE: [original, FaceSwap] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/val_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: Celeb-DFv2 + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + # ROOT: /project/home/p200249/XXX/FaceForensics++/c23/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + # ROOT: /project/home/p200249/XXX/DeeperForensics/ + FROM_FILE: False + # FAKETYPE: [original, FaceSwap] + # FAKETYPE: [videos_BW_5/original, videos_BW_5/Deepfakes, videos_BW_5/Face2Face, videos_BW_5/FaceSwap, videos_BW_5/NeuralTextures] + FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [DFo_source_videos, DFo_manipulated_videos] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 1] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + +MODEL: + type: TopDownDetector + backbone: + type: TimeViT + img_size: [224, 224] + patch_size: 16 + embed_dim: 768 + depth: 12 + num_heads: 12 + mlp_ratio: 4 + drop_path_rate: 0.3 + qkv_bias: True + class_token: True + register_token: False + temp_token: True + attention_type: divided_space_time #joint_space_time, space_only + num_frames: 8 + low_level_enhanced: False + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 768 + hm_size: [14, 14] #img_size // patch_size + heads: + cls: 1 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256, 256] + num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + conv_2direction: False + use_temp_token: True + features: 3D + INIT_WEIGHTS: + pretrained: null + +TRAIN: + gpus: [0,1,2,3] + batch_size: 16 + lr: 0.00005 + epochs: 100 + begin_epoch: -1 + warm_up: 3 + every_val_epochs: 1 + accumulation_steps: 1 + use_amp: False + loss: + type: BinaryCrossEntropy + reduction: mean + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 2.5 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: False + video_level: True + no_shot_preds: 1 + pretrained: 'logs/03-10-2024/TopDownDetector_C23_TimeBase224_ST_binary_8SBI_AdamW_0vidAug_temp_FaceSwap_model_best.pth' diff --git a/video/fake-stormer/model_code/configs/temporal/bin_cls/TimeSFormer_base_c40.yaml b/video/fake-stormer/model_code/configs/temporal/bin_cls/TimeSFormer_base_c40.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b62fdf3390fbe87fc00f146ae81cf248112bfd76 --- /dev/null +++ b/video/fake-stormer/model_code/configs/temporal/bin_cls/TimeSFormer_base_c40.yaml @@ -0,0 +1,181 @@ +TASK: C40_TimeBase224_ST_binary_8SBI_AdamW_0vidAug_temp_NeuralTextures +PRECISION: float32 +METRICS_BASE: binary +SEED: 317 +DATA_RELOAD: False + +DATASET: + type: BinaryFaceForensic + DATA_TYPE: video + TRAIN: True #Switch to True for training mode, False for testing mode + COMPRESSION: c40 + IMAGE_SUFFIX: png + NUM_WORKERS: 32 + PIN_MEMORY: True + IMAGE_SIZE: [224, 224] + HEATMAP_SIZE: [14, 14] #[IMAGE_SIZE//4, IMAGE_SIZE//4] + SIGMA: 1 + ADAPTIVE_SIGMA: False + HEATMAP_TYPE: m_std_normalized + SPLIT_IMAGE: False + DATA: + TYPE: videos + SAMPLES_PER_VIDEO: + ACTIVE: True + TRAIN: 8 # Dynamically random number of frames in each epoch + VAL: 8 + TEST: 8 + NUM_FRAMES: 8 + TRAIN: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c40/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: False + FAKETYPE: [original, NeuralTextures] + # ANNO_FILE: train/frames/FaceXRay/train_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/train_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + VAL: + NAME: FF++ # This field to define datasets that can be used to train/in-dataset/cross-dataset evaluation + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c40/ + # ROOT: /data/deepfake_cluster/datasets_df/FaceForensics++/c0/ + FROM_FILE: False + FAKETYPE: [original, NeuralTextures] # Choosing Deepfake techniques to be loaded for dataloader + # ANNO_FILE: val/frames/FaceXRay/val_FF_FaceXRay.json + ANNO_FILE: processed_data/c23/val_FaceForensics_videos_81.json + LABEL_FOLDER: [real, fake] + TEST: + NAME: FF++ + # ROOT: /home/users/XXX/data/FaceForensics++/c0/ + ROOT: /project/home/p200249/XXX/FaceForensics++/c40/ + # ROOT: /home/users/XXX/data/Celeb-DFv1/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv1/ + # ROOT: /data/deepfake_cluster/datasets_df/Celeb-DFv1/ + # ROOT: /home/users/XXX/data/Celeb-DFv2/ + # ROOT: /project/home/p200249/XXX/Celeb-DFv2/ + # ROOT: /home/users/XXX/data/DFDCP/ + # ROOT: /project/home/p200249/XXX/DFDCP/ + # ROOT: /home/users/XXX/data/DFDC/ + # ROOT: /project/home/p200249/XXX/DFDC/ + # ROOT: /home/users/XXX/data/DFD/ + # ROOT: /project/home/p200249/XXX/DFD/ + # ROOT: /home/users/XXX/data/DFW/ + # ROOT: /project/home/p200249/XXX/DFW/ + # ROOT: /project/home/p200249/XXX/DeeperForensics/ + FROM_FILE: False + FAKETYPE: [original, FaceSwap] + # FAKETYPE: [videos_BW_5/original, videos_BW_5/Deepfakes, videos_BW_5/Face2Face, videos_BW_5/FaceSwap, videos_BW_5/NeuralTextures] + # FAKETYPE: [Celeb-real, Celeb-synthesis, YouTube-real] + # FAKETYPE: [method_A, method_B, original_videos] + # FAKETYPE: [fake, real] + # FAKETYPE: [DeepFakeDetection_original, DeepFakeDetection] + # FAKETYPE: [fake_test, real_test] + # FAKETYPE: [DFo_source_videos, DFo_manipulated_videos] + ANNO_FILE: FaceXRay/test/test_FF_Xray.json + LABEL_FOLDER: [real, fake] + TRANSFORM: + geometry: + type: GeometryTransform + resize: [224, 224, 1] #h, w, p=probability. If no affine transform, set p=1 + normalize: 0 + horizontal_flip: 0.5 + cropping: [0.15, 0.5] #Format: [crop_limit, prob] + scale: [0.15, 0.5] #Format: [scale_limit, prob] + rand_erasing: [0.0, 1] #Format: [p, max_count] + color: + type: ColorJitterTransform + clahe: 0.0 + colorjitter: 0.3 + gaussianblur: 0.3 + gaussnoise: 0.3 + jpegcompression: [0.5, 40, 100] # prob, lower and upper quality respectively + rgbshift: 0.3 + randomcontrast: 0.0 + randomgamma: 0.5 + randombrightness: 1 + huesat: 1 + normalize: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + +MODEL: + type: TopDownDetector + backbone: + type: TimeViT + img_size: [224, 224] + patch_size: 16 + embed_dim: 768 + depth: 12 + num_heads: 12 + mlp_ratio: 4 + drop_path_rate: 0.3 + qkv_bias: True + class_token: True + register_token: False + temp_token: True + attention_type: divided_space_time #joint_space_time, space_only + num_frames: 8 + low_level_enhanced: False + keypoint_head: + type: TopdownHeatmapSimpleHead + in_channels: 768 + hm_size: [14, 14] #img_size // patch_size + heads: + cls: 1 + num_deconv_layers: 0 #Config n deconv layers to build the decoder + num_deconv_filters: [256, 256] + num_deconv_kernels: [4, 4] + loss_keypoint: + type: JointsMSELoss + use_target_weight: False + extra: + final_conv_kernel: 3 + num_conv_layers: 1 + conv_2direction: False + use_temp_token: True + features: 3D + INIT_WEIGHTS: + pretrained: null + +TRAIN: + gpus: [0,1,2,3] + batch_size: 16 + lr: 0.00005 + epochs: 100 + begin_epoch: -1 + warm_up: 3 + every_val_epochs: 1 + accumulation_steps: 1 + use_amp: False + loss: + type: BinaryCrossEntropy + reduction: mean + optimizer: AdamW + distributed: False + pretrained: '' + tensorboard: False + resume: False + lr_scheduler: + # type: MultiStepLR + milestones: [5, 15, 20, 25] + gamma: 0.5 + freeze_backbone: True + debug: + active: False + save_hm_gt: True + save_hm_pred: True + booster: 1 + start_decay: 2.5 + +TEST: + gpus: [0] + subtask: 'eval' + test_file: '' + vis_hm: True + threshold: 0.5 + flip_test: False + video_level: True + no_shot_preds: 1 + pretrained: 'logs/03-10-2024/TopDownDetector_C40_TimeBase224_ST_binary_8SBI_AdamW_0vidAug_temp_NeuralTextures_model_best.pth' diff --git a/video/fake-stormer/model_code/datasets/__init__.py b/video/fake-stormer/model_code/datasets/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2fa76c040a21d56ea365398bcff98e9866f2eacc --- /dev/null +++ b/video/fake-stormer/model_code/datasets/__init__.py @@ -0,0 +1,21 @@ +# -*- coding: utf-8 -*- +from .builder import DATASETS, PIPELINES, build_dataset +from .face_forensic_binary import BinaryFaceForensic +from .face_forensic_hm import HeatmapFaceForensic +from .face_forensic_sbi import SBIFaceForensic +from .fakesformer_bi import FakeSFormerBI +from .fakesformer_sbi import FakeSFormerSBI +from .pipelines import * + +__all__ = [ + "GeometryTransform", + "BinaryFaceForensic", + "ColorJitterTransform", + "PIPELINES", + "DATASETS", + "build_dataset", + "HeatmapFaceForensic", + "SBIFaceForensic", + "FakeSFormerSBI", + "FakeSFormerBI", +] diff --git a/video/fake-stormer/model_code/datasets/builder.py b/video/fake-stormer/model_code/datasets/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..35f09797618a41e24f0933fad3408696987a36f8 --- /dev/null +++ b/video/fake-stormer/model_code/datasets/builder.py @@ -0,0 +1,30 @@ +# -*- coding: utf-8 -*- +import os +import sys +from typing import Any, Dict, Optional + +if not os.getcwd() in sys.path: + sys.path.append(os.getcwd()) + +from register.register import Registry, build_from_cfg + +PIPELINES = Registry("Pipeline", build_func=build_from_cfg) +DATASETS = Registry("Dataset", build_func=build_from_cfg) + + +def build_pipeline( + cfg, + pipeline: Registry, + build_func=build_from_cfg, + default_args: Optional[Dict] = None, +) -> Any: + return build_func(cfg, pipeline, default_args) + + +def build_dataset( + cfg, + dataset: Registry, + build_func=build_from_cfg, + default_args: Optional[Dict] = None, +) -> Any: + return build_func(cfg, dataset, default_args) diff --git a/video/fake-stormer/model_code/datasets/celebDF_v1.py b/video/fake-stormer/model_code/datasets/celebDF_v1.py new file mode 100644 index 0000000000000000000000000000000000000000..dfa23f74ff8cd4c121f356541048249cb3666346 --- /dev/null +++ b/video/fake-stormer/model_code/datasets/celebDF_v1.py @@ -0,0 +1,41 @@ +# -*- coding: utf-8 -*- +import os +from glob import glob + +import numpy as np + +from .builder import DATASETS +from .common import CommonDataset + + +@DATASETS.register_module() +class CDFV1(CommonDataset): + def __init__(self, cfg, **kwargs): + super().__init__(cfg, **kwargs) + + def _load_from_path(self, split): + assert os.path.exists( + self._cfg.DATA[self.split.upper()].ROOT + ), "Root path to dataset can not be None!" + data = self._cfg["DATA"] + data_type = data.TYPE + fake_types = self._cfg.DATA[split.upper()]["FAKETYPE"] + img_paths, labels, mask_paths, ot_props = [], [], [], [] + + # Load image data for each type of fake techniques + for idx, ft in enumerate(fake_types): + data_dir = os.path.join( + self._cfg.DATA[self.split.upper()].ROOT, self.split, data_type, ft + ) + if not os.path.exists(data_dir): + raise ValueError("Data Directory can not be invalid!") + + for sub_dir in os.listdir(data_dir): + sub_dir_path = os.path.join(data_dir, sub_dir) + img_paths_ = glob(f"{sub_dir_path}/*.{self._cfg.IMAGE_SUFFIX}") + + img_paths.extend(img_paths_) + labels.extend(np.full(len(img_paths_), int("Celeb-synthesis" in ft))) + + print("{} image paths have been loaded from CDFv1!".format(len(img_paths))) + return img_paths, labels, mask_paths, ot_props diff --git a/video/fake-stormer/model_code/datasets/celebDF_v2.py b/video/fake-stormer/model_code/datasets/celebDF_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..61c7de7d9e982f9d82add4afba40c5b8e27db569 --- /dev/null +++ b/video/fake-stormer/model_code/datasets/celebDF_v2.py @@ -0,0 +1,58 @@ +# -*- coding: utf-8 -*- +import os +from glob import glob + +import numpy as np + +from .builder import DATASETS +from .common import CommonDataset + + +@DATASETS.register_module() +class CDFV2(CommonDataset): + def __init__(self, cfg, **kwargs): + super().__init__(cfg, **kwargs) + + def _load_from_path(self, split): + assert os.path.exists( + self._cfg.DATA[self.split.upper()].ROOT + ), "Root path to dataset can not be None!" + data = self._cfg["DATA"] + data_type = data.TYPE + fake_types = self._cfg.DATA[split.upper()]["FAKETYPE"] + img_paths, labels, mask_paths, ot_props = [], [], [], [] + + count = 0 + n_samples = 100000 + + # Load image data for each type of fake techniques + for idx, ft in enumerate(fake_types): + data_dir = os.path.join( + self._cfg.DATA[self.split.upper()].ROOT, self.split, data_type, ft + ) + if not os.path.exists(data_dir): + raise ValueError("Data Directory can not be invalid!") + + for sub_dir in os.listdir(data_dir): + sub_dir_path = os.path.join(data_dir, sub_dir) + img_paths_ = glob(f"{sub_dir_path}/*.{self._cfg.IMAGE_SUFFIX}") + + if "Celeb-synthesis" in ft: + if count < n_samples: + n_add = ( + len(img_paths_) + if ((n_samples - count) > len(img_paths_)) + else (n_samples - count) + ) + count += n_add + print(f"n fake samples added --- {count}") + else: + continue + else: + n_add = len(img_paths_) + + img_paths.extend(img_paths_[:n_add]) + labels.extend(np.full(n_add, int("Celeb-synthesis" in ft))) + + print("{} image paths have been loaded from CDFv2!".format(len(img_paths))) + return img_paths, labels, mask_paths, ot_props diff --git a/video/fake-stormer/model_code/datasets/combine.py b/video/fake-stormer/model_code/datasets/combine.py new file mode 100644 index 0000000000000000000000000000000000000000..2f69a304f4fac0411b498b1fa89f97f82c08d25d --- /dev/null +++ b/video/fake-stormer/model_code/datasets/combine.py @@ -0,0 +1,47 @@ +# -*- coding: utf-8 -*- +import os +from glob import glob + +import numpy as np + +from .builder import DATASETS +from .common import CommonDataset + + +@DATASETS.register_module() +class Combine(CommonDataset): + def __init__(self, cfg, **kwargs): + super().__init__(cfg, **kwargs) + + def _load_from_path(self, split): + assert os.path.exists( + self._cfg.DATA[self.split.upper()].ROOT + ), "Root path to dataset can not be None!" + data = self._cfg["DATA"] + data_type = data.TYPE + fake_types = self._cfg.DATA[split.upper()]["FAKETYPE"] + img_paths, labels, mask_paths, ot_props = [], [], [], [] + + # Load image data for each type of fake techniques + for idx, ft in enumerate(fake_types): + data_dir = os.path.join( + self._cfg.DATA[self.split.upper()].ROOT, self.split, data_type, ft + ) + if not os.path.exists(data_dir): + raise ValueError("Data Directory can not be invalid!") + + for sub_dir in os.listdir(data_dir): + sub_dir_path = os.path.join(data_dir, sub_dir) + # sub_dir_path = data_dir + img_paths_ = glob(f"{sub_dir_path}/*.{self._cfg.IMAGE_SUFFIX}") + + img_paths.extend(img_paths_) + labels.extend( + np.full( + len(img_paths_), + int(("real" not in ft) and ("original" not in ft)), + ) + ) + + print("{} image paths have been loaded from Combine!".format(len(img_paths))) + return img_paths, labels, mask_paths, ot_props diff --git a/video/fake-stormer/model_code/datasets/common.py b/video/fake-stormer/model_code/datasets/common.py new file mode 100644 index 0000000000000000000000000000000000000000..c716a3ec3fdfdc5eaaa3b5268d8304692a02ac9b --- /dev/null +++ b/video/fake-stormer/model_code/datasets/common.py @@ -0,0 +1,717 @@ +# -*- coding: utf-8 -*- +import math +import os +import random +import sys +from abc import ABC, abstractmethod +from glob import glob + +import numpy as np +import simplejson as json +from box import Box as edict +from natsort import natsorted +from package_utils.image_utils import cal_mask_wh, gaussian_radius +from package_utils.transform import final_transform +from package_utils.utils import file_extention +from PIL import Image +from torch.utils.data import Dataset + +from .builder import DATASETS +from .utils import _extract_data_based_dist + +PREFIX_PATH = "/data/deepfake_cluster/datasets_df/FaceForensics++/c0/" + + +class ParameterStore: + _instance = None + _parameters = {} + + @classmethod + def get_instance(cls): + if cls._instance is None: + cls._instance = cls() + return cls._instance + + @classmethod + def add_parameters(cls, param_name, param_value): + cls._parameters[param_name] = param_value + + @classmethod + def get_parameters(cls, param_name): + return cls._parameters.get(param_name) + + @classmethod + def del_parameters(cls): + for k in cls._parameters.keys(): + del cls._parameters[k] + + @classmethod + def has_key(cls, key): + return key in cls._parameters + + @classmethod + def reset(cls): + cls._parameters.clear() + cls._instance = None + + +@DATASETS.register_module() +class CommonDataset(Dataset, ABC): + def __init__(self, cfg, **kwargs): + super().__init__() + self._cfg = edict(cfg) if not isinstance(cfg, edict) else cfg + self.dataset = self._cfg.DATA[self.split.upper()].NAME + # self.train = self._cfg["TRAIN"] + self.train = self.split != "test" + self.final_transforms = final_transform(self._cfg) + self.sigma_adaptive = self._cfg.ADAPTIVE_SIGMA + self.sampler_active = self._cfg.DATA.SAMPLES_PER_VIDEO.ACTIVE + self.samples_per_video = self._cfg.DATA.SAMPLES_PER_VIDEO[self.split.upper()] + self.sampler_dist = ( + self._cfg.DATA.SAMPLES_PER_VIDEO.DIST + ) # Distribution of [Real, Fake] + self.heatmap_w = self._cfg.HEATMAP_SIZE[1] + self.heatmap_h = self._cfg.HEATMAP_SIZE[0] + self.split_image = self._cfg.SPLIT_IMAGE + self.compression = self._cfg.COMPRESSION + self.data_type = self._cfg.DATA_TYPE + + if kwargs is not None: + for k, v in kwargs.items(): + if v is None: + raise ValueError(f"{k}:{v} retrieve a None value!") + self.__setattr__(k, v) + + @abstractmethod + def _load_from_path(self, split): + return NotImplemented + + def _load_from_file(self, split, anno_file=None): + """ + @split: train/val + This function for loading data from file for 4 types of manipulated images FF++ and FaceXray generation data + """ + assert os.path.exists( + self._cfg.DATA[self.split.upper()].ROOT + ), "Root path to dataset can not be invalid!" + data_cfg = self._cfg.DATA + + if anno_file is None: + anno_file = data_cfg[split.upper()].ANNO_FILE + if not os.access(anno_file, os.R_OK): + anno_file = os.path.join(self._cfg.DATA[self.split.upper()].ROOT, anno_file) + assert os.access(anno_file, os.R_OK), "Annotation file can not be invalid!!" + + f_name, f_extention = file_extention(anno_file) + data = None + image_paths, labels, mask_paths, ot_props = [], [], [], [] + f = open(anno_file) + if f_extention == ".json": + data = json.load(f) + data = edict(data)[ + "data" + ] # A list of proprocessed data objects containing image properties + + for item in data: + assert ( + "image_path" in item.keys() + ), "Image path must be available in item dict!" + image_path = item.image_path + ot_prop = {} + + # Custom base on the specific data structure + if not "label" in item.keys(): + lb = ("fake" in image_path) or ( + ("original" not in image_path) and ("aligned" not in image_path) + ) + else: + lb = item.label == "fake" + lb_encoded = int(lb) + labels.append(lb_encoded) + + if PREFIX_PATH in item.image_path: + image_path = item.image_path.replace( + PREFIX_PATH, self._cfg.DATA[self.split.upper()].ROOT + ) + else: + image_path = os.path.join( + self._cfg.DATA[self.split.upper()].ROOT, item.image_path + ) + image_paths.append(image_path) + + # Appending more data properties for data loader + if "mask_path" in item.keys(): + mask_path = item.mask_path + if PREFIX_PATH in item.mask_path: + mask_path = item.mask_path.replace( + PREFIX_PATH, self._cfg.DATA[self.split.upper()].ROOT + ) + else: + mask_path = os.path.join( + self._cfg.DATA[self.split.upper()].ROOT, item.mask_path + ) + mask_paths.append(mask_path) + if "best_match" in item.keys(): + best_match = item.best_match + best_match = [ + os.path.join(self._cfg.DATA[self.split.upper()].ROOT, bm) + for bm in best_match + if self._cfg.DATA[self.split.upper()].ROOT not in bm + ] + ot_prop["best_match"] = best_match + for lms_key in ["aligned_lms", "orig_lms"]: + if lms_key in item.keys(): + f_lms = np.array(item[lms_key]) + ot_prop[lms_key] = f_lms + + ot_props.append(ot_prop) + else: + raise Exception( + f"{f_extention} has not been supported yet! Please change to Json file!" + ) + + print("{} image paths have been loaded!".format(len(image_paths))) + return image_paths, labels, mask_paths, ot_props + + def _gen_vul_parts(self, blending_mask): + H, W, C = blending_mask.shape + Hp, Wp = self.heatmap_h, self.heatmap_w + py, px = int(H // Hp), int(W // Wp) + + assert (H // Hp) == (W // Wp) + vul_parts = np.zeros((Hp, Wp)) + + for i in range(0, Hp): + for j in range(0, Wp): + blending_part = blending_mask[ + (py * i) : (py * (i + 1)), (px * j) : (px * (j + 1)), 0 + ] + part_intensity = np.mean(blending_part) + vul_parts[i, j] = part_intensity + vul_parts_out = np.tile(vul_parts[:, :, np.newaxis], (1, 1, 3)).astype(np.uint8) + + return vul_parts_out + + def _mask_out_vulnerability(self, input, mask, fake_intensity, mask_prob=0.9): + if self.dynamic_blending_prob: + p_h = self._cfg.IMAGE_SIZE[0] // self.heatmap_h + p_w = self._cfg.IMAGE_SIZE[1] // self.heatmap_w + + max_value = max(0.1, mask[..., 0].max()) + upper_bound_intensity = min(1.0, fake_intensity) + upper_bound_value = max( + 0.1, mask[mask[..., 0] < max_value * upper_bound_intensity].max() + ) + target_mask_ = (mask[..., 0] > upper_bound_value).astype(int) + + # Randomly mask out mask if the input is real + if np.count_nonzero(target_mask_) == 0: + pos_matrix = (self.heatmap_h, self.heatmap_w) + all_indices = [ + (i, j) for i in range(self.heatmap_h) for j in range(self.heatmap_w) + ] + selected_indices = np.random.choice( + len(all_indices), + size=math.floor(mask_prob * np.prod((pos_matrix))), + replace=False, + ) + selected_indices_2d = [all_indices[i] for i in selected_indices] + i_indices, j_indices = zip(*selected_indices_2d) + else: + all_indices = [ + (i, j) + for i in range(self.heatmap_h) + for j in range(self.heatmap_w) + if ( + (target_mask_[i, j] == 0) + and (mask[..., 0][i, j] < upper_bound_value) + ) + ] + idxes = np.where(target_mask_ == 1) + n_mask_pos_h = len(idxes[0]) + pos_matrix = (self.heatmap_h, self.heatmap_w) + + if len(all_indices) < math.floor( + mask_prob * np.prod((pos_matrix)) - n_mask_pos_h + ): + size = math.ceil(len(all_indices) * mask_prob) + else: + size = max( + 0, math.floor(mask_prob * np.prod((pos_matrix)) - n_mask_pos_h) + ) + + selected_indices = np.random.choice( + len(all_indices), size=size, replace=False + ) + selected_indices_2d = [all_indices[i] for i in selected_indices] + i_indices_, j_indices_ = zip(*selected_indices_2d) + i_indices = np.hstack((idxes[0], np.array(i_indices_))) + j_indices = np.hstack((idxes[1], np.array(j_indices_))) + + target_mask_[i_indices, j_indices] = 1 + idxes = np.where(target_mask_ == 1) + masked_matrix = 1 - target_mask_ + + for i, j in zip(idxes[0], idxes[1]): + input[ + int(i * p_h) : int((i + 1) * p_h), + int(j * p_w) : int((j + 1) * p_w), + :, + ] = np.zeros((1, 1, input.shape[2]), dtype=input.dtype) + # mask[i, j] = np.zeros((mask.shape[2]), dtype=mask.dtype) + + return input, mask, masked_matrix, upper_bound_value + + def _mask_out_vulnerability2(self, input, mask, fake_intensity, **kwargs): + mask_prob = kwargs.get("mask_prob") + param_store_ins = ParameterStore.get_instance() + masked_matrix = param_store_ins.get_parameters("masked_matrix") + + p_h = self._cfg.IMAGE_SIZE[0] // self.heatmap_h + p_w = self._cfg.IMAGE_SIZE[1] // self.heatmap_w + upper_bound_value = None + + if self.dynamic_blending_prob: + if masked_matrix is not None: + upper_bound_value = max(1, np.max(mask[..., 0] * masked_matrix)) + fake_intensity = upper_bound_value / 255 + target_mask_ = 1 - masked_matrix + idxes = np.where(target_mask_ == 1) + else: + max_value = max(1, mask[..., 0].max()) + max_f_intensity = max_value / 255 + fake_intensity = min(fake_intensity, max_f_intensity) + + upper_bound_value = max( + 1, mask[mask[..., 0] < 255 * fake_intensity].max() + ) + target_mask_ = (mask[..., 0] > upper_bound_value).astype(int) + + # Randomly mask out mask if the input is real + if np.count_nonzero(target_mask_) == 0: + pos_matrix = (self.heatmap_h, self.heatmap_w) + all_indices = [ + (i, j) + for i in range(self.heatmap_h) + for j in range(self.heatmap_w) + ] + selected_indices = np.random.choice( + len(all_indices), + size=math.ceil(mask_prob * np.prod((pos_matrix))), + replace=False, + ) + selected_indices_2d = [all_indices[i] for i in selected_indices] + i_indices, j_indices = zip(*selected_indices_2d) + else: + all_indices = [ + (i, j) + for i in range(self.heatmap_h) + for j in range(self.heatmap_w) + if ( + (target_mask_[i, j] == 0) + and (mask[..., 0][i, j] < upper_bound_value) + ) + ] + idxes = np.where(target_mask_ == 1) + n_mask_pos_h = len(idxes[0]) + pos_matrix = (self.heatmap_h, self.heatmap_w) + + if len(all_indices) < math.floor( + mask_prob * np.prod((pos_matrix)) - n_mask_pos_h + ): + size = math.ceil(len(all_indices) * mask_prob) + else: + size = max( + 1, + math.ceil(mask_prob * np.prod((pos_matrix)) - n_mask_pos_h), + ) + + selected_indices = np.random.choice( + len(all_indices), size=size, replace=False + ) + selected_indices_2d = [all_indices[i] for i in selected_indices] + i_indices_, j_indices_ = zip(*selected_indices_2d) + i_indices = np.hstack((idxes[0], np.array(i_indices_))) + j_indices = np.hstack((idxes[1], np.array(j_indices_))) + + target_mask_[i_indices, j_indices] = 1 + idxes = np.where(target_mask_ == 1) + masked_matrix = 1 - target_mask_ + param_store_ins.add_parameters("masked_matrix", masked_matrix) + + for i, j in zip(idxes[0], idxes[1]): + rand_val = np.random.randint(0, 255) + input[ + int(i * p_h) : int((i + 1) * p_h), + int(j * p_w) : int((j + 1) * p_w), + :, + ] = np.full((1, 1, input.shape[2]), 0, dtype=input.dtype) + # mask[i, j] = np.zeros((mask.shape[2]), dtype=mask.dtype) + + return input, mask, masked_matrix, upper_bound_value, fake_intensity + + def _encode_temporal_target(self, target_mask, **kwargs): + """ + Adaptively encode targets based on the vulnerability levels for spatial-temporal outputs (3D) + """ + assert self.heatmap_type in ["gaussian", "m_std_normalized", "max_normalized"] + + if isinstance(target_mask, list): + target_mask = np.array(target_mask) + + saved_params = {} + hm_w = self._cfg.HEATMAP_SIZE[1] + hm_h = self._cfg.HEATMAP_SIZE[0] + ndim = len(target_mask) + heatmap = np.zeros((ndim, hm_h, hm_w), dtype=np.float32) # dimension d, h, w + # cstency_hm = np.zeros((ndim, hm_h, hm_w), dtype=np.float32) + + derivative = np.diff(target_mask[:, :, :, 0], axis=0) + derivative = np.absolute(derivative) + d_max = max(1, np.max(derivative)) + + if bool(derivative.max()) and kwargs.get("vis_derivative"): + idx = kwargs.get("idx") + for i in range(len(derivative)): + di = derivative[i].astype(np.uint8) + di = np.repeat(di[:, :, np.newaxis], 3, axis=2) + Image.fromarray(di).save(f"samples/debugs/derivative_f_{idx}_{i}.png") + + if self.heatmap_type == "gaussian": + x = np.arange(0, hm_w, 1, float) + y = np.arange(0, hm_h, 1, float) + y = np.expand_dims(y, -1) + z = np.arange(0, ndim, 1, float) + z = np.expand_dims((np.expand_dims(z, -1)), -1) + derivative = np.concatenate( + (np.zeros_like(derivative[0][np.newaxis, ...]), derivative), axis=0 + ) + centers = np.where(derivative == max(0.1, derivative.max())) + + for i, j, k in zip(centers[0], centers[1], centers[2]): + heatmap_ijk = np.exp( + -( + ((z - i) ** 2) / (2.0 * (self.sigma / 2) ** 2) + + ((y - j) ** 2) / (2.0 * (self.sigma / 2) ** 2) + + ((x - k) ** 2) / (2.0 * (self.sigma / 2) ** 2) + ) + ) + heatmap = np.maximum(heatmap_ijk, heatmap) + elif self.heatmap_type == "m_std_normalized": + d_m = kwargs.get("d_mean") or np.mean(derivative) + d_std = kwargs.get("d_std") or np.std(derivative) + derivative = np.concatenate( + (np.zeros_like(derivative[0][np.newaxis, ...]), derivative), axis=0 + ) + # Calculating the 3D self-consistency map + # cstency_hm = 255 - np.absolute(d_max - derivative) + + if d_std != 0: + heatmap = (derivative - d_m) / d_std + + saved_params = {"d_mean": d_m, "d_std": d_std} + elif self.heatmap_type == "max_normalized": + derivative = np.concatenate( + (np.zeros_like(derivative[0][np.newaxis, ...]), derivative), axis=0 + ) + # Calculating the 3D self-consistency map + # cstency_hm = 255 - np.absolute(d_max - derivative) + if d_max != 0: + heatmap[1:] = derivative / d_max + else: + raise ValueError("Now only support gaussian or mean std normalization") + + return heatmap, derivative / d_max, saved_params + + def _encode_target(self, target_mask, fake_intensity=0.5): + """ + Adaptively encode targets based on the vulnerability levels + """ + assert ( + self.heatmap_type == "gaussian" + ), "Only Gaussian Heatmap is supported now!" + hm_w = self._cfg.HEATMAP_SIZE[1] + hm_h = self._cfg.HEATMAP_SIZE[0] + heatmap = np.zeros((1, hm_h, hm_w), dtype=np.float32) + + # Draw heatmap for blending region + max_val_all = target_mask[..., 0].max() + max_val = max_val_all if max_val_all > 0 else 255 + + # Select value to draw attention masks + if self.data_type == "video": + lower_bound_intensity = max(0.0, (fake_intensity - 0.1)) + upper_bound_intensity = min(1.0, (fake_intensity + 0.1)) + target_mask_ = ( + (target_mask[..., 0] >= 255 * lower_bound_intensity) + & (target_mask[..., 0] < 255 * upper_bound_intensity) + ).astype(np.int8) + else: + target_mask_ = (target_mask[..., 0] >= max_val * fake_intensity).astype( + np.int8 + ) + + points = np.where(target_mask_ == 1) + + for j, i in zip(points[0], points[1]): + if self.sigma_adaptive: + w_sbi, h_sbi = cal_mask_wh((j, i), target_mask[..., 0]) + radius = gaussian_radius((h_sbi, w_sbi)) + self.sigma = radius / 3 + 1e-4 + tmp = self.sigma * 3 + size = tmp * 2 + 1 + ul = [int(i - tmp), int(j - tmp)] + br = [int(i + tmp + 1), int(j + tmp + 1)] + x = np.arange(0, size, 1, np.float32) + y = x[:, np.newaxis] + + x0 = y0 = size // 2 + g = np.exp(-((x - x0) ** 2 + (y - y0) ** 2) / (2 * (self.sigma**2))) + + g_x = max(0, -ul[0]), min(br[0], hm_w) - ul[0] + g_y = max(0, -ul[1]), min(br[1], hm_h) - ul[1] + + img_x = max(0, ul[0]), min(br[0], hm_w) + img_y = max(0, ul[1]), min(br[1], hm_h) + + heatmap[0][img_y[0] : img_y[1], img_x[0] : img_x[1]] = np.maximum( + g[g_y[0] : g_y[1], g_x[0] : g_x[1]], + heatmap[0][img_y[0] : img_y[1], img_x[0] : img_x[1]], + ) + + return heatmap, None + + def _encode_target_v1(self, target_mask, fake_intensity=0.5): + assert ( + self.heatmap_type == "gaussian" + ), "Only Gaussian Heatmap is supported now!" + # fake_ratio = [0.5, 0.6, 0.7, 0.8, 0.9, 1.0] + n_outputs = 1 + hm_w = self._cfg.HEATMAP_SIZE[1] + hm_h = self._cfg.HEATMAP_SIZE[0] + patches = [[0, 0], [0, 1 / 2], [1 / 2, 0], [1 / 2, 1 / 2]] + target_H, target_W = target_mask[..., 0].shape[:2] + heatmap = np.zeros((n_outputs, target_H, target_W), dtype=np.float32) + cstency_hm = np.zeros((n_outputs, target_H, target_W), dtype=np.float32) + max_val_all = target_mask[..., 0].max() + + # Draw heatmap for blending region + for fr in range(len(patches)): + # target_mask_ = np.where(((target_mask[..., 0] > 255*fake_ratio[fr]) & (target_mask[..., 0] <= 255*fake_ratio[fr+1])), 1, 0) + p_x1, p_y1 = int(target_W * patches[fr][0]), int(target_H * patches[fr][1]) + p_x2, p_y2 = int(target_W * (patches[fr][0] + 1 / 2)), int( + target_H * (patches[fr][1] + 1 / 2) + ) + + max_value = target_mask[p_y1:p_y2, p_x1:p_x2, 0].max() + max_value = max_value if max_value > 0 else 1 + target_mask_ = (target_mask[p_y1:p_y2, p_x1:p_x2, 0] == (max_value)).astype( + np.uint8 + ) + points = np.where(target_mask_ == 1) + + if len(points[0]): + p = (points[0] + p_y1, points[1] + p_x1) + + for j, i in zip(p[0], p[1]): + if self.sigma_adaptive: + w_sbi, h_sbi = cal_mask_wh((j, i), target_mask[..., 0]) + radius = gaussian_radius((h_sbi, w_sbi)) + self.sigma = radius / 3 + 1e-4 + tmp = self.sigma * 3 + size = tmp * 2 + 1 + ul = [int(i - tmp), int(j - tmp)] + br = [int(i + tmp + 1), int(j + tmp + 1)] + x = np.arange(0, size, 1, np.float32) + y = x[:, np.newaxis] + + x0 = y0 = size // 2 + g = np.exp(-((x - x0) ** 2 + (y - y0) ** 2) / (2 * (self.sigma**2))) + + g_x = max(0, -ul[0]), min(br[0], hm_w) - ul[0] + g_y = max(0, -ul[1]), min(br[1], hm_h) - ul[1] + + img_x = max(0, ul[0]), min(br[0], hm_w) + img_y = max(0, ul[1]), min(br[1], hm_h) + + heatmap[0][img_y[0] : img_y[1], img_x[0] : img_x[1]] = np.maximum( + g[g_y[0] : g_y[1], g_x[0] : g_x[1]], + heatmap[0][img_y[0] : img_y[1], img_x[0] : img_x[1]], + ) + + if n_outputs > 1: + cstency_hm[fr][p_y1:p_y2, p_x1:p_x2] = 255 - np.absolute( + max_value - target_mask[p_y1:p_y2, p_x1:p_x2, 0] + ) + else: + cstency_hm[0][p_y1:p_y2, p_x1:p_x2] = 255 - np.absolute( + max_val_all - target_mask[p_y1:p_y2, p_x1:p_x2, 0] + ) + + return heatmap, cstency_hm + + def _encode_target_v2(self, target_mask, fake_intensity=0.5): + assert ( + self.heatmap_type == "gaussian" + ), "Only Gaussian Heatmap is supported now!" + n_outputs = 1 + hm_w = self._cfg.HEATMAP_SIZE[1] + hm_h = self._cfg.HEATMAP_SIZE[0] + target_H, target_W = target_mask[..., 0].shape[:2] + heatmap = np.zeros((n_outputs, target_H, target_W), dtype=np.float32) + cstency_hm = np.zeros((n_outputs, target_H, target_W), dtype=np.float32) + + # Draw heatmap for blending region + target_mask_ = (target_mask[..., 0] > 128).astype(np.uint8) + points = np.where(target_mask_ == 1) + + if len(points[0]): + p = (int(points[0].mean()), int(points[1].mean())) + j, i = p + + if self.sigma_adaptive: + w_sbi, h_sbi = cal_mask_wh((j, i), target_mask[..., 0]) + radius = gaussian_radius((h_sbi, w_sbi)) + self.sigma = radius / 3 + 1e-4 + tmp = self.sigma * 3 + size = tmp * 2 + 1 + ul = [int(i - tmp), int(j - tmp)] + br = [int(i + tmp + 1), int(j + tmp + 1)] + x = np.arange(0, size, 1, np.float32) + y = x[:, np.newaxis] + + x0 = y0 = size // 2 + g = np.exp(-((x - x0) ** 2 + (y - y0) ** 2) / (2 * (self.sigma**2))) + + g_x = max(0, -ul[0]), min(br[0], hm_w) - ul[0] + g_y = max(0, -ul[1]), min(br[1], hm_h) - ul[1] + + img_x = max(0, ul[0]), min(br[0], hm_w) + img_y = max(0, ul[1]), min(br[1], hm_h) + + heatmap[0][img_y[0] : img_y[1], img_x[0] : img_x[1]] = np.maximum( + g[g_y[0] : g_y[1], g_x[0] : g_x[1]], + heatmap[0][img_y[0] : img_y[1], img_x[0] : img_x[1]], + ) + + cstency_hm[0] = 255 - np.absolute( + target_mask[j, i, 0] - target_mask[..., 0] + ) + + return heatmap, cstency_hm + + def _sampler(self, image_paths, labels, epoch=0, **params): + if self.sampler_dist[0] != 1.0 or self.sampler_dist[1] != 1.0: + image_paths, labels, params = _extract_data_based_dist( + self.data_type, image_paths, labels, self.sampler_dist, **params + ) + + vid_dict = {} + data = {"image_paths": [], "labels": []} + + for k, v in params.items(): + if v is not None and len(v): + data[k] = [] + + for idx, ip in enumerate(image_paths): + f_name = ip.split("/")[-1] + + if self.compression in ["c0", "c23", "c40"]: + vid_id = os.path.dirname(ip) + if self.dataset == "FF++" and self.train: + f_type = ip.split("/")[-3] + vid_id = "_".join([f_type, vid_id]) + else: + raise NotImplementedError( + "Only c23, c40, and c0 compression mode is supported now! Please check again!" + ) + lb = labels[idx] + + data_per_vid = dict(image=ip, label=lb) + for k, v in params.items(): + if k in data.keys(): + data_per_vid[k] = v[idx] + + if vid_id in vid_dict.keys(): + vid_dict[vid_id].append(data_per_vid) + else: + vid_dict[vid_id] = [data_per_vid] + + if self.data_type == "image": + """ + Samples data for the mode of working with single images + """ + for vid_id in vid_dict.keys(): + if self.train: + samples_per_vid = random.choices( + vid_dict[vid_id], k=self.samples_per_video + ) + else: + samples_per_vid = random.sample( + vid_dict[vid_id], k=len(vid_dict[vid_id]) + ) + + for spl in samples_per_vid: + data["image_paths"].append(spl["image"]) + data["labels"].append(spl["label"]) + for k in params.keys(): + if k in data.keys(): + data[k].append(spl[k]) + return data + elif self.data_type == "video": + # Sorting to obtain successive frames for videos, important for temporal modeling + for vid_id in vid_dict.keys(): + vid_dict[vid_id] = natsorted(vid_dict[vid_id], key=lambda x: x["image"]) + + # if self.train: + """ + Generate new video data for training + """ + assert "NUM_FRAMES" in self._cfg.DATA.SAMPLES_PER_VIDEO + new_vid_dict = {} + n_fs = self._cfg.DATA.SAMPLES_PER_VIDEO.NUM_FRAMES + + for vid_id in vid_dict.keys(): + vid_len = len(vid_dict[vid_id]) + start_idx = 0 if epoch == 0 else np.random.randint(0, n_fs - 1) + + for k in range(start_idx, vid_len, n_fs): + try: + if (k + n_fs) <= vid_len: + new_vid_id = "+++".join( + [vid_id, str(k)] + ) # Adding index segment to original video to create sub videos + new_vid_dict[new_vid_id] = vid_dict[vid_id][k : (k + n_fs)] + except: + break + + return new_vid_dict + else: + raise ValueError( + f'{self.data_type} has not been supported! Only "image" or "video" data can be extracted!' + ) + + def select_encode_method(self, version=0, dimension="spatial"): + if dimension == "spatial": + if version == 2: + return self._encode_target_v2 + elif version == 1: + return self._encode_target_v1 + else: + return self._encode_target + elif dimension == "temporal": + return self._encode_temporal_target + else: + raise ValueError(f"The input {dimension} has not been supported yet!") + + @abstractmethod + def __len__(self): + return NotImplemented + + @abstractmethod + def __getitem__(self, idx): + return NotImplemented + + @property + def __repr__(self): + return self.__class__.__name__ diff --git a/video/fake-stormer/model_code/datasets/df40.py b/video/fake-stormer/model_code/datasets/df40.py new file mode 100644 index 0000000000000000000000000000000000000000..b98cfebc66d3649cf292e34baadf97c2e1902fc6 --- /dev/null +++ b/video/fake-stormer/model_code/datasets/df40.py @@ -0,0 +1,46 @@ +# -*- coding: utf-8 -*- +import os +from glob import glob + +import numpy as np + +from .builder import DATASETS +from .common import CommonDataset + + +@DATASETS.register_module() +class DF40(CommonDataset): + def __init__(self, cfg, **kwargs): + super().__init__(cfg, **kwargs) + + def _load_from_path(self, split): + # Check if the root directory exists. + root_dir = self._cfg.DATA[self.split.upper()].ROOT + assert os.path.exists(root_dir), "Root path to dataset cannot be None!" + + data = self._cfg["DATA"] + data_type = data.TYPE + fake_types = self._cfg.DATA[split.upper()]["FAKETYPE"] + img_paths, labels, mask_paths, ot_props = [], [], [], [] + + # Load image data for each type of fake technique. + for ft in fake_types: + # Construct path: ROOT/split/fake_type (skipping data_type) + data_dir = os.path.join(root_dir, self.split, data_type, ft) + if not os.path.exists(data_dir): + raise ValueError("Data Directory is invalid!") + + # Define common image extensions. + extensions = ["jpg", "jpeg", "png", "tif", "webp"] + img_paths_ = [] + # Recursively search for images in the fake type directory. + for ext in extensions: + pattern = os.path.join(data_dir, "**", f"*.{ext}") + img_paths_.extend(glob(pattern, recursive=True)) + + # Extend the main lists with the images and corresponding labels. + img_paths.extend(img_paths_) + labels.extend(np.full(len(img_paths_), int(ft != "real_videos"))) + + print("{} image paths have been loaded from DF40!".format(len(img_paths))) + return img_paths, labels, mask_paths, ot_props diff --git a/video/fake-stormer/model_code/datasets/dfd.py b/video/fake-stormer/model_code/datasets/dfd.py new file mode 100644 index 0000000000000000000000000000000000000000..62649bdb74b309d8807f19f9527c0227a52ad4ec --- /dev/null +++ b/video/fake-stormer/model_code/datasets/dfd.py @@ -0,0 +1,41 @@ +# -*- coding: utf-8 -*- +import os +from glob import glob + +import numpy as np + +from .builder import DATASETS +from .common import CommonDataset + + +@DATASETS.register_module() +class DFD(CommonDataset): + def __init__(self, cfg, **kwargs): + super().__init__(cfg, **kwargs) + + def _load_from_path(self, split): + assert os.path.exists( + self._cfg.DATA[self.split.upper()].ROOT + ), "Root path to dataset can not be None!" + data = self._cfg["DATA"] + data_type = data.TYPE + fake_types = self._cfg.DATA[split.upper()]["FAKETYPE"] + img_paths, labels, mask_paths, ot_props = [], [], [], [] + + # Load image data for each type of fake techniques + for idx, ft in enumerate(fake_types): + data_dir = os.path.join( + self._cfg.DATA[self.split.upper()].ROOT, self.split, data_type, ft + ) + if not os.path.exists(data_dir): + raise ValueError("Data Directory can not be invalid!") + + for sub_dir in os.listdir(data_dir): + sub_dir_path = os.path.join(data_dir, sub_dir) + img_paths_ = glob(f"{sub_dir_path}/*.{self._cfg.IMAGE_SUFFIX}") + + img_paths.extend(img_paths_) + labels.extend(np.full(len(img_paths_), int(ft == "DeepFakeDetection"))) + + print("{} image paths have been loaded from DFD!".format(len(img_paths))) + return img_paths, labels, mask_paths, ot_props diff --git a/video/fake-stormer/model_code/datasets/dfdc.py b/video/fake-stormer/model_code/datasets/dfdc.py new file mode 100644 index 0000000000000000000000000000000000000000..997d34befbeab2668bae022308c436c84e6372b5 --- /dev/null +++ b/video/fake-stormer/model_code/datasets/dfdc.py @@ -0,0 +1,41 @@ +# -*- coding: utf-8 -*- +import os +from glob import glob + +import numpy as np + +from .builder import DATASETS +from .common import CommonDataset + + +@DATASETS.register_module() +class DFDC(CommonDataset): + def __init__(self, cfg, **kwargs): + super().__init__(cfg, **kwargs) + + def _load_from_path(self, split): + assert os.path.exists( + self._cfg.DATA[self.split.upper()].ROOT + ), "Root path to dataset can not be None!" + data = self._cfg["DATA"] + data_type = data.TYPE + fake_types = self._cfg.DATA[split.upper()]["FAKETYPE"] + img_paths, labels, mask_paths, ot_props = [], [], [], [] + + # Load image data for each type of fake techniques + for idx, ft in enumerate(fake_types): + data_dir = os.path.join( + self._cfg.DATA[self.split.upper()].ROOT, self.split, data_type, ft + ) + if not os.path.exists(data_dir): + raise ValueError("Data Directory can not be invalid!") + + for sub_dir in os.listdir(data_dir): + sub_dir_path = os.path.join(data_dir, sub_dir) + img_paths_ = glob(f"{sub_dir_path}/*.{self._cfg.IMAGE_SUFFIX}") + + img_paths.extend(img_paths_) + labels.extend(np.full(len(img_paths_), int(ft == "fake"))) + + print("{} image paths have been loaded from DFDC!".format(len(img_paths))) + return img_paths, labels, mask_paths, ot_props diff --git a/video/fake-stormer/model_code/datasets/dfdcp.py b/video/fake-stormer/model_code/datasets/dfdcp.py new file mode 100644 index 0000000000000000000000000000000000000000..a0198bd0fc30fdee74b73c9040491ff0fd0b1423 --- /dev/null +++ b/video/fake-stormer/model_code/datasets/dfdcp.py @@ -0,0 +1,45 @@ +# -*- coding: utf-8 -*- +import os +from glob import glob + +import numpy as np + +from .builder import DATASETS +from .common import CommonDataset + + +@DATASETS.register_module() +class DFDCP(CommonDataset): + def __init__(self, cfg, **kwargs): + super().__init__(cfg, **kwargs) + + def _load_from_path(self, split): + assert os.path.exists( + self._cfg.DATA[self.split.upper()].ROOT + ), "Root path to dataset can not be None!" + data = self._cfg["DATA"] + data_type = data.TYPE + fake_types = self._cfg.DATA[split.upper()]["FAKETYPE"] + img_paths, labels, mask_paths, ot_props = [], [], [], [] + + # Load image data for each type of fake techniques + for idx, ft in enumerate(fake_types): + data_dir = os.path.join( + self._cfg.DATA[self.split.upper()].ROOT, self.split, data_type, ft + ) + if not os.path.exists(data_dir): + raise ValueError("Data Directory can not be invalid!") + + for sub_dir in os.listdir(data_dir): + sub_dir_path = os.path.join(data_dir, sub_dir) + img_paths_ = glob(f"{sub_dir_path}/*.{self._cfg.IMAGE_SUFFIX}") + + img_paths.extend(img_paths_) + labels.extend( + np.full( + len(img_paths_), int((ft == "method_A") or (ft == "method_B")) + ) + ) + + print("{} image paths have been loaded from DFDCP!".format(len(img_paths))) + return img_paths, labels, mask_paths, ot_props diff --git a/video/fake-stormer/model_code/datasets/dfo.py b/video/fake-stormer/model_code/datasets/dfo.py new file mode 100644 index 0000000000000000000000000000000000000000..b360108c2ea5e867fe87c473c8707979c1df88e3 --- /dev/null +++ b/video/fake-stormer/model_code/datasets/dfo.py @@ -0,0 +1,46 @@ +# -*- coding: utf-8 -*- +import os +from glob import glob + +import numpy as np + +from .builder import DATASETS +from .common import CommonDataset + + +@DATASETS.register_module() +class DFo(CommonDataset): + def __init__(self, cfg, **kwargs): + super().__init__(cfg, **kwargs) + + def _load_from_path(self, split): + assert os.path.exists( + self._cfg.DATA[self.split.upper()].ROOT + ), "Root path to dataset can not be None!" + data = self._cfg["DATA"] + data_type = data.TYPE + fake_types = self._cfg.DATA[split.upper()]["FAKETYPE"] + img_paths, labels, mask_paths, ot_props = [], [], [], [] + + # Load image data for each type of fake techniques + for idx, ft in enumerate(fake_types): + data_dir = os.path.join( + self._cfg.DATA[self.split.upper()].ROOT, self.split, data_type, ft + ) + if not os.path.exists(data_dir): + raise ValueError("Data Directory can not be invalid!") + + for sub_dir in os.listdir(data_dir): + sub_dir_path = os.path.join(data_dir, sub_dir) + # sub_dir_path = data_dir + img_paths_ = glob(f"{sub_dir_path}/*.{self._cfg.IMAGE_SUFFIX}") + + img_paths.extend(img_paths_) + labels.extend(np.full(len(img_paths_), int("manipulated" in ft))) + + print( + "{} image paths have been loaded from DeeperForensics!".format( + len(img_paths) + ) + ) + return img_paths, labels, mask_paths, ot_props diff --git a/video/fake-stormer/model_code/datasets/dfw.py b/video/fake-stormer/model_code/datasets/dfw.py new file mode 100644 index 0000000000000000000000000000000000000000..4e374150446bfecc72539c9e0f1c732df2e1fd36 --- /dev/null +++ b/video/fake-stormer/model_code/datasets/dfw.py @@ -0,0 +1,42 @@ +# -*- coding: utf-8 -*- +import os +from glob import glob + +import numpy as np + +from .builder import DATASETS +from .common import CommonDataset + + +@DATASETS.register_module() +class DFW(CommonDataset): + def __init__(self, cfg, **kwargs): + super().__init__(cfg, **kwargs) + + def _load_from_path(self, split): + assert os.path.exists( + self._cfg.DATA[self.split.upper()].ROOT + ), "Root path to dataset can not be None!" + data = self._cfg["DATA"] + data_type = data.TYPE + fake_types = self._cfg.DATA[split.upper()]["FAKETYPE"] + img_paths, labels, mask_paths, ot_props = [], [], [], [] + + # Load image data for each type of fake techniques + for idx, ft in enumerate(fake_types): + data_dir = os.path.join( + self._cfg.DATA[self.split.upper()].ROOT, self.split, data_type, ft + ) + if not os.path.exists(data_dir): + raise ValueError("Data Directory can not be invalid!") + + for sub_dir in os.listdir(data_dir): + sub_dir_path = os.path.join(data_dir, sub_dir) + # sub_dir_path = data_dir + img_paths_ = glob(f"{sub_dir_path}/*.{self._cfg.IMAGE_SUFFIX}") + + img_paths.extend(img_paths_) + labels.extend(np.full(len(img_paths_), int("fake" in ft))) + + print("{} image paths have been loaded from DFW!".format(len(img_paths))) + return img_paths, labels, mask_paths, ot_props diff --git a/video/fake-stormer/model_code/datasets/diffswap.py b/video/fake-stormer/model_code/datasets/diffswap.py new file mode 100644 index 0000000000000000000000000000000000000000..ef26e46129b2a92b259a63ee03fbc7cf544512e7 --- /dev/null +++ b/video/fake-stormer/model_code/datasets/diffswap.py @@ -0,0 +1,45 @@ +# -*- coding: utf-8 -*- +import os +from glob import glob + +import numpy as np + +from .builder import DATASETS +from .common import CommonDataset + + +@DATASETS.register_module() +class DiffSwap(CommonDataset): + def __init__(self, cfg, **kwargs): + super().__init__(cfg, **kwargs) + + def _load_from_path(self, split): + assert os.path.exists( + self._cfg.DATA[self.split.upper()].ROOT + ), "Root path to dataset can not be None!" + data = self._cfg["DATA"] + data_type = data.TYPE + fake_types = self._cfg.DATA[split.upper()]["FAKETYPE"] + img_paths, labels, mask_paths, ot_props = [], [], [], [] + + # Load image data for each type of fake techniques + for idx, ft in enumerate(fake_types): + data_dir = os.path.join( + self._cfg.DATA[self.split.upper()].ROOT, self.split, data_type, ft + ) + if not os.path.exists(data_dir): + raise ValueError("Data Directory can not be invalid!") + + if not os.path.isdir(data_dir): + continue + + for sub_dir in os.listdir(data_dir): + sub_dir_path = os.path.join(data_dir, sub_dir) + # sub_dir_path = data_dir + img_paths_ = glob(f"{sub_dir_path}/*.{self._cfg.IMAGE_SUFFIX}") + + img_paths.extend(img_paths_) + labels.extend(np.full(len(img_paths_), int("Real" not in ft))) + + print("{} image paths have been loaded from DiffSwap!".format(len(img_paths))) + return img_paths, labels, mask_paths, ot_props diff --git a/video/fake-stormer/model_code/datasets/face_forensic_binary.py b/video/fake-stormer/model_code/datasets/face_forensic_binary.py new file mode 100644 index 0000000000000000000000000000000000000000..e323f886c842b5692aa4dff27efd1eb63bc47cd0 --- /dev/null +++ b/video/fake-stormer/model_code/datasets/face_forensic_binary.py @@ -0,0 +1,180 @@ +# -*- coding: utf-8 -*- +import os +import sys +from abc import abstractmethod +from random import shuffle + +import numpy as np +import torch +from package_utils.image_utils import load_image +from PIL import Image +from torch.utils.data import default_collate + +from .builder import DATASETS, PIPELINES, build_pipeline +from .master import MasterDataset + + +@DATASETS.register_module() +class BinaryFaceForensic(MasterDataset): + def __init__(self, config, split, **kwargs): + """ + @params: + config: Dataset config + split: train/val/test which directs to the split folders + """ + self.split = split + if kwargs is not None: + for k, v in kwargs.items(): + if v is None: + raise ValueError(f"{k}:{v} retrieve a None value!") + self.__setattr__(k, v) + super().__init__(config, **kwargs) + + # Load data + self.data_sampler = self._load_data(split) + + # Parse data + self._parsing_data() + + # Calling transform methods for inputs + self.geo_transform = build_pipeline(config.TRANSFORM.geometry, PIPELINES) + self.colorjitter_transform = build_pipeline(config.TRANSFORM.color, PIPELINES) + + def _load_data(self, split, anno_file=None, epoch=0): + from_file = self._cfg.DATA[self.split.upper()].FROM_FILE + + if epoch == 0: + if not from_file: + self.image_paths, self.labels, self.mask_paths, self.ot_props = ( + self._load_from_path(split) + ) + else: + self.image_paths, self.labels, self.mask_paths, self.ot_props = ( + self._load_from_file(split, anno_file=anno_file) + ) + + assert ( + len(self.image_paths) != 0 + ), "Image paths have not been loaded! Please check image directory!" + assert ( + len(self.labels) != 0 + ), "Labels have not been loaded! Please check annotation file!" + + if self.sampler_active: + print("Running sampler...") + params = dict( + mask_paths=self.mask_paths, ot_props=self.ot_props, epoch=epoch + ) + data_sampler = self._sampler(self.image_paths, self.labels, **params) + return data_sampler + + def _parsing_data(self): + assert self.data_type in ["image", "video"] + # Parsing data for training + if self.data_type == "video": + return + self.image_paths, self.labels = ( + self.data_sampler["image_paths"], + self.data_sampler["labels"], + ) + + def _load_img(self, img_path): + return load_image(img_path) + + def __len__(self): + if self.data_type == "image": + assert "image_paths" in self.data_sampler.keys() + return len(self.labels) + elif self.data_type == "video": + return len(self.data_sampler.keys()) + else: + raise ValueError( + f'{self.data_type} has not been supported. Please use "image" or "video" instead!' + ) + + def __getitem__(self, idx): + if self.data_type == "image": + return self.__getitem_path__(idx=idx) + elif self.data_type == "video": + return self.__getitem_video__(idx=idx) + else: + raise ValueError( + f"{self.data_type} has not been supported. Only image or video are used for training!" + ) + + def __getitem_path__(self, idx): + img_path = self.image_paths[idx] + label = np.expand_dims(self.labels[idx], axis=-1) + img = self._load_img(img_path) + + # Applying geo transform to inputs + geo_transfomed = self.geo_transform(img) + img_trans = geo_transfomed["image"] + + # Applying color transform to inputs + color_transfomed = self.colorjitter_transform(img_trans) + img_trans = color_transfomed["image"] + + # Normalise + Convert numpy array to tensor + img_trans = img_trans / 255 + img_trans = self.final_transforms(img_trans) + return img_trans, label + + def __getitem_video__(self, idx): + inputs = [] + vid_id = [*self.data_sampler.keys()][idx] + vid_data = self.data_sampler[vid_id] + + label = np.expand_dims(vid_data[0]["label"], axis=-1) + + f_idxes = range(0, self.samples_per_video) + for ix, f_idx in enumerate(f_idxes): + it = vid_data[f_idx] + img_path = it["image"] + img = self._load_img(img_path) + + if self.train: + # Applying geo transform to inputs + geo_transfomed = self.geo_transform(img) + img_trans = geo_transfomed["image"] + + # Applying color transform to inputs + color_transfomed = self.colorjitter_transform(img_trans) + img_trans = color_transfomed["image"] + + # Normalise + Convert numpy array to tensor + img_trans = img_trans / 255 + else: + img_trans = img / 255 + img_trans = self.final_transforms(img_trans) + inputs.append(img_trans) + inputs = torch.tensor(np.array([ip.numpy() for ip in inputs])) + inputs = inputs.transpose(0, 1) + + if self.train: + return inputs, label + else: + return inputs, label, vid_id.split("-")[0] + + def train_collate_fn(self, batch): + return default_collate(batch) + + +if __name__ == "__main__": + from configs.get_config import load_config + from datasets import * + from pipelines.color_transform import ColorJitterTransform + from pipelines.geo_transform import GeometryTransform + from torch.utils.data import DataLoader + + PIPELINES.register_module(module=GeometryTransform) + PIPELINES.register_module(module=ColorJitterTransform) + + config = load_config("configs/temporal/bin_cls/TimeSFormer_base_c23.yaml") + bin_ff = DATASETS.build( + cfg=config.DATASET, default_args=dict(split="val", config=config.DATASET) + ) + bin_ff_loader = DataLoader(bin_ff, batch_size=10, shuffle=True) + for b, (X, y) in enumerate(bin_ff_loader): + print(f"X.shape - {X.shape}, y shape - {y.shape}") + break diff --git a/video/fake-stormer/model_code/datasets/face_forensic_hm.py b/video/fake-stormer/model_code/datasets/face_forensic_hm.py new file mode 100644 index 0000000000000000000000000000000000000000..4c4653b1d3a1dc0201d7a73c1d2ed1a404e0eeae --- /dev/null +++ b/video/fake-stormer/model_code/datasets/face_forensic_hm.py @@ -0,0 +1,447 @@ +# -*- coding: utf-8 -*- +import random + +import cv2 +import numpy as np +import torch +from imgaug import augmenters as iaa +from package_utils.bi_online_generation import ( + blendImages, + colorTransfer, + random_erode_dilate, + random_get_hull, +) +from package_utils.image_utils import load_image +from package_utils.transform import ( + get_affine_transform, + get_center_scale, +) +from package_utils.utils import draw_landmarks, vis_heatmap +from PIL import Image +from skimage import transform as sktransform + +from .builder import DATASETS, PIPELINES, build_pipeline +from .master import MasterDataset + + +@DATASETS.register_module() +class HeatmapFaceForensic(MasterDataset): + def __init__(self, config, split, **kwargs): + """ + @params: + config: Dataset config + split: train/val/test which directs to the split folders + """ + self.split = split + super().__init__(config, **kwargs) + if kwargs is not None: + for k, v in kwargs.items(): + if v is None: + raise ValueError(f"{k}:{v} retrieve a None value!") + self.__setattr__(k, v) + self.rot = 0 + self.pixel_std = 200 + self.target_w = self._cfg.IMAGE_SIZE[1] + self.target_h = self._cfg.IMAGE_SIZE[0] + self.aspect_ratio = self.target_w * 1.0 / self.target_h + self.sigma = self._cfg.SIGMA + self.heatmap_type = self._cfg.HEATMAP_TYPE + self.debug = self._cfg.DEBUG + # self.train = self._cfg.TRAIN + self.dynamic_fxray = self._cfg.DYNAMIC_FXRAY + + # Load data + self.image_paths, self.labels, self.mask_paths, self.ot_props = self._load_data( + split + ) + + # predefine mask distortion + self.distortion = iaa.Sequential([iaa.PiecewiseAffine(scale=(0.01, 0.15))]) + + # Calling transform methods for inputs + self.geo_transform = build_pipeline(config.TRANSFORM.geometry, PIPELINES) + self.colorjitter_transform = build_pipeline(config.TRANSFORM.color, PIPELINES) + + def _load_data(self, split): + from_file = self._cfg.DATA[self.split.upper()].FROM_FILE + + if not from_file: + image_paths, labels, mask_paths, ot_props = self._load_from_path(split) + else: + image_paths, labels, mask_paths, ot_props = self._load_from_file(split) + + if self.sampler_active and self.train: + print("Running sampler...") + params = dict(mask_paths=mask_paths, ot_props=ot_props) + data_sampler = self._sampler(image_paths, labels, **params) + image_paths, labels = data_sampler["image_paths"], data_sampler["labels"] + if len(mask_paths): + mask_paths = data_sampler["mask_paths"] + if len(ot_props): + ot_props = data_sampler["ot_props"] + print(f"n samples after running sampling --- {len(image_paths)}") + + assert ( + len(image_paths) != 0 + ), "Image paths have not been loaded! Please check image directory!" + assert ( + len(labels) != 0 + ), "Labels have not been loaded! Please check annotation file!" + # if not self.dynamic_fxray or self.split == 'val': + if from_file and not (self.dynamic_fxray): + assert ( + len(mask_paths) != 0 + ), "Mask paths have not been loaded! Please check mask directory!" + return image_paths, labels, mask_paths, ot_props + + def _reload_data(self): + self.image_paths, self.labels, self.mask_paths, self.ot_props = self._load_data( + self.split + ) + + def _gen_BI( + self, background_face, background_landmark, foreground_face_path, idx=None + ): + foreground_face = load_image(foreground_face_path) + + # down sample before blending + aug_size = random.randint(128, 317) + background_landmark = background_landmark * (aug_size / 317) + foreground_face = sktransform.resize( + foreground_face, (aug_size, aug_size), preserve_range=True + ).astype(np.uint8) + background_face = sktransform.resize( + background_face, (aug_size, aug_size), preserve_range=True + ).astype(np.uint8) + + # get random type of initial blending mask + mask = random_get_hull(background_landmark, background_face) + + if self.debug: + Image.fromarray((mask * 255).astype(np.uint8)).save( + f"samples/debugs/orig_CH_{idx}.jpg" + ) + + # random deform mask + mask = self.distortion.augment_image(mask) + mask = random_erode_dilate(mask) + + if self.debug: + Image.fromarray((mask * 255).astype(np.uint8)).save( + f"samples/debugs/deformed_CH_{idx}.jpg" + ) + + # filte empty mask after deformation + if np.sum(mask) == 0: + raise NotImplementedError + + # apply color transfer + foreground_face = colorTransfer(background_face, foreground_face, mask * 255) + + # blend two face + blended_face, mask = blendImages(foreground_face, background_face, mask * 255) + blended_face = blended_face.astype(np.uint8) + + # resize back to default resolution + blended_face = sktransform.resize( + blended_face, (317, 317), preserve_range=True + ).astype(np.uint8) + mask = sktransform.resize(mask, (317, 317), preserve_range=True) + mask = mask[:, :, 0:1] + return blended_face, mask + + def _gen_target( + self, background_face, background_landmark, foreground_face_path, idx=None + ): + data_type = "real" if random.randint(0, 1) else "fake" + + if not background_landmark.any(): + data_type = "real" + + if data_type == "fake": + face_img, mask = self._gen_BI( + background_face, background_landmark, foreground_face_path, idx=idx + ) + mask = (1 - mask) * mask * 4 + else: + face_img = background_face + mask = np.zeros((317, 317, 1)) + + face_img = Image.fromarray(face_img) + # randomly downsample after BI pipeline + if random.randint(0, 1): + aug_size = random.randint(64, 317) + if random.randint(0, 1): + face_img = face_img.resize((aug_size, aug_size), Image.BILINEAR) + else: + face_img = face_img.resize((aug_size, aug_size), Image.NEAREST) + face_img = face_img.resize((317, 317), Image.BILINEAR) + face_img = np.array(face_img) + + face_img = face_img[60:(317), 30:(287), :] + mask = mask[60:(317), 30:(287), :] + mask = np.repeat(mask, 3, 2) + mask = (mask * 255).astype(np.uint8) + return face_img, mask, int(data_type == "fake") + + def __len__(self): + return len(self.labels) + + def __getitem__(self, idx): + flag = True + while flag: + try: + img_path = self.image_paths[idx] + label = self.labels[idx] + vid_id = img_path.split("/")[-2] + img = load_image(img_path) + mask = None + + if self.debug: + Image.fromarray(img).save(f"samples/debugs/orig_{idx}.jpg") + + # Applying color transform to inputs + if self.split == "train": + color_transfomed = self.colorjitter_transform(img) + img = color_transfomed["image"] + + # if not self.dynamic_fxray or self.split == 'val': + if not self.dynamic_fxray: + if bool(self.mask_paths): + mask_path = self.mask_paths[idx] + mask = load_image(mask_path) + else: + if self.train: + best_match = ( + self.ot_props[idx]["best_match"] + if len(self.ot_props[idx]["best_match"]) + else [] + ) + if len(self.ot_props[idx]["aligned_lms"]): + f_lms = self.ot_props[idx]["aligned_lms"] + elif len(self.ot_props[idx]["orig_lms"]): + f_lms = self.ot_props[idx]["orig_lms"] + else: + f_lms = [] + + if self.debug: + img_lms_draw = draw_landmarks(img, f_lms) + Image.fromarray(img_lms_draw).save( + f"samples/debugs/orig_{idx}_lms.jpg" + ) + + if len(best_match): + best_match_idx = random.randint(0, len(best_match) - 1) + best_match_path = best_match[best_match_idx] + img, mask, label = self._gen_target( + img, f_lms, best_match_path, idx=idx + ) + else: + img, mask, label = self._gen_target(img, np.array([]), "") + else: + img = cv2.resize(img, (317, 317)) + # Best croppings from 18-299 for testing + img = img[18:(299), 18:(299), :] + target = None + + if mask is not None: + assert ( + mask.shape[:2] == img.shape[:2] + ), "Color Image and Mask must have the same shape!" + + # Applying geo transform to inputs and masks + if self.split == "train": + geo_transfomed = self.geo_transform(img, mask=mask) + img = geo_transfomed["image"] + mask = geo_transfomed["mask"] + + # Applying affine transform + c, s = get_center_scale( + img.shape[:2], self.aspect_ratio, pixel_std=self.pixel_std + ) + trans = get_affine_transform(c, s, self.rot, self._cfg.IMAGE_SIZE) + trans_heatmap = get_affine_transform( + c, s, self.rot, self._cfg.HEATMAP_SIZE + ) + + input = cv2.warpAffine( + img, + trans, + (int(self._cfg.IMAGE_SIZE[0]), int(self._cfg.IMAGE_SIZE[1])), + flags=cv2.INTER_LINEAR, + ) + + if mask is not None: + target = cv2.warpAffine( + mask, + trans_heatmap, + ( + int(self._cfg.HEATMAP_SIZE[0]), + int(self._cfg.HEATMAP_SIZE[1]), + ), + flags=cv2.INTER_LINEAR, + ) + + # Target encoding + # 0 for original, 1 for FXRay, 2 for NoFXRay. If 2, comment: mask = (1 - mask) * mask * 4 + heatmap, cstency_hm = ( + self.select_encode_method(version=1)(target) + if ( + target is not None + and self.heatmap_type == "gaussian" + and self.train + ) + else (None, None) + ) + + if self.debug: + Image.fromarray(input).save(f"samples/debugs/affine_{idx}.jpg") + Image.fromarray(target).save( + f"samples/debugs/mask_affine_{idx}.jpg" + ) + vis_heatmap( + input, + cstency_hm / 255, + f"samples/debugs/cstency_mask_{idx}.jpg", + ) + vis_heatmap(input, heatmap, f"samples/debugs/hm_{idx}.jpg") + + if self.train: + if self.split_image: + patch_img_trans = [] + patch_heatmap = [] + patch_cstency_hm = [] + patch_target = [] + patch_label = np.expand_dims(np.tile(label, len(heatmap)), -1) + + for i, (k, l) in enumerate( + [[0, 0], [1 / 2, 0], [0, 1 / 2], [1 / 2, 1 / 2]] + ): + input_ = input[ + int(self.target_h * k) : int( + self.target_h * (k + 1 / 2) + ), + int(self.target_w * l) : int( + self.target_w * (l + 1 / 2) + ), + :, + ] + heatmap_ = heatmap[i][ + int(self.heatmap_h * k) : int( + self.heatmap_h * (k + 1 / 2) + ), + int(self.heatmap_w * l) : int( + self.heatmap_w * (l + 1 / 2) + ), + ] + cstency_ = cstency_hm[i][ + int(self.heatmap_h * k) : int( + self.heatmap_h * (k + 1 / 2) + ), + int(self.heatmap_w * l) : int( + self.heatmap_w * (l + 1 / 2) + ), + ] + target_ = target[..., 0][ + int(self.heatmap_h * k) : int( + self.heatmap_h * (k + 1 / 2) + ), + int(self.heatmap_w * l) : int( + self.heatmap_w * (l + 1 / 2) + ), + ] + + # Normalise + Convert numpy array to tensor + input_ = input_ / 255 + patch_img_trans.append(self.final_transforms(input_)) + + patch_heatmap.append(heatmap_) + patch_cstency_hm.append(cstency_ / 255) + patch_target.append(target_ / 255) + else: + patch_img_trans = self.final_transforms(input / 255) + patch_heatmap = heatmap + patch_cstency_hm = cstency_hm / 255 + patch_target = target / 255 + patch_label = np.expand_dims(label, axis=-1) + # patch_label = label + else: + # Normalise + Convert numpy array to tensor + img_trans = input / 255 + img_trans = self.final_transforms(img_trans) + label = np.expand_dims(label, axis=-1) + flag = False + except Exception as e: + print("There is an exception during loading data, please check --- ", e) + idx = torch.randint(low=0, high=self.__len__(), size=(1,)).item() + + if self.train: + return ( + patch_img_trans, + patch_label, + patch_target, + patch_heatmap, + patch_cstency_hm, + ) + else: + meta = {"vid_id": vid_id, "img_path": img_path} + return img_trans, label, meta + + def train_collate_fn(self, batch): + batch_data = {} + img, label, target, hm, cstency_hm = zip(*batch) + + # Collating data in case of using spliting images into patches + if self.split_image: + hm_H, hm_W = hm[0][0].shape + + img = np.reshape(img, (-1)) + hm = np.reshape(hm, (-1, 1, hm_H, hm_W)) + cstency_hm = np.reshape(cstency_hm, (-1, 1, hm_H, hm_W)) + target = np.reshape(target, (-1, 1, hm_H, hm_W)) + label = np.reshape(label, (-1, 1)) + + img = torch.tensor([it.numpy() for it in img]) + heatmap = torch.tensor(hm).float() + cstency_heatmap = torch.tensor(cstency_hm).float() + target = torch.tensor(target).float() + label = torch.tensor(label) + + batch_data["img"] = img + batch_data["label"] = label + batch_data["target"] = target + batch_data["heatmap"] = heatmap + batch_data["cstency"] = cstency_heatmap + + return batch_data + + +if __name__ == "__main__": + # from datasets import * + from configs.get_config import load_config + from pipelines.color_transform import ColorJitterTransform + from pipelines.geo_transform import GeometryTransform + from torch.utils.data import DataLoader + + PIPELINES.register_module(module=GeometryTransform) + PIPELINES.register_module(module=ColorJitterTransform) + + config = load_config("configs/efn4_fpn_hm_adv.yaml") + hm_ff = DATASETS.build( + cfg=config.DATASET, default_args=dict(split="train", config=config.DATASET) + ) + hm_ff_loader = DataLoader( + hm_ff, batch_size=10, shuffle=False, collate_fn=hm_ff.train_collate_fn + ) + for b, batch_data in enumerate(hm_ff_loader): + inputs, labels, targets, heatmaps, cstency_heatmap = ( + batch_data["img"], + batch_data["label"], + batch_data["target"], + batch_data["heatmap"], + batch_data["cstency_heatmap"], + ) + print( + f"X.shape - {inputs.shape}, y shape - {labels.shape}, heatmaps - {heatmaps.shape}, consistency -- {cstency_heatmap.shape}" + ) + break diff --git a/video/fake-stormer/model_code/datasets/face_forensic_sbi.py b/video/fake-stormer/model_code/datasets/face_forensic_sbi.py new file mode 100644 index 0000000000000000000000000000000000000000..c92e169efd99e5fc1c9b2168c74fb73ad1837b1c --- /dev/null +++ b/video/fake-stormer/model_code/datasets/face_forensic_sbi.py @@ -0,0 +1,463 @@ +# -*- coding: utf-8 -*- +import os +import random +import sys + +import cv2 +import numpy as np +import torch +import torch.nn.functional as F +from package_utils.image_utils import crop_by_margin, load_image +from package_utils.transform import get_affine_transform, get_center_scale +from package_utils.utils import draw_landmarks, draw_most_vul_points, vis_heatmap +from PIL import Image + +from .builder import DATASETS, PIPELINES, build_pipeline +from .master import MasterDataset +from .pipelines.geo_transform import get_transforms +from .sbi.utils import * + + +@DATASETS.register_module() +class SBIFaceForensic(MasterDataset): + def __init__(self, config, split, **kwargs): + """ + @params: + config: Dataset config + split: train/val/test which directs to the split folders + """ + self.split = split + super(SBIFaceForensic, self).__init__(config, **kwargs) + if kwargs is not None: + for k, v in kwargs.items(): + if v is None: + raise ValueError(f"{k}:{v} retrieve a None value!") + self.__setattr__(k, v) + self.rot = 0 + self.pixel_std = 200 + self.target_w = self._cfg.IMAGE_SIZE[1] + self.target_h = self._cfg.IMAGE_SIZE[0] + self.aspect_ratio = self.target_w * 1.0 / self.target_h + self.sigma = self._cfg.SIGMA + self.heatmap_type = self._cfg.HEATMAP_TYPE + self.debug = self._cfg.DEBUG + # self.train = self._cfg.TRAIN + self.dynamic_fxray = self._cfg.DYNAMIC_FXRAY + + # Load data + self.image_paths_r, self.labels_r, self.mask_paths_r, self.ot_props_r = ( + self._load_data(split) + ) + + # Calling transform methods for inputs + self.geo_transform = build_pipeline( + config.TRANSFORM.geometry, + PIPELINES, + default_args={"additional_targets": {"image_f": "image", "mask_f": "mask"}}, + ) + # self.colorjitter_transform = build_pipeline(config.TRANSFORM.color, PIPELINES) + + self.transforms = get_transforms() + + def __len__(self): + return len(self.labels_r) + + def _load_img(self, img_path): + return load_image(img_path) + + def _reload_data(self): + self.image_paths, self.labels, self.mask_paths, self.ot_props = self._load_data( + self.split + ) + + def _load_data(self, split, anno_file=None): + from_file = self._cfg.DATA[self.split.upper()].FROM_FILE + + if not from_file: + image_paths, labels, mask_paths, ot_props = self._load_from_path(split) + else: + image_paths, labels, mask_paths, ot_props = self._load_from_file( + split, anno_file=anno_file + ) + + assert ( + len(image_paths) != 0 + ), "Image paths have not been loaded! Please check image directory!" + assert ( + len(labels) != 0 + ), "Labels have not been loaded! Please check annotation file!" + if not self.dynamic_fxray: + assert ( + len(mask_paths) != 0 + ), "Mask paths have not been loaded! Please check mask directory!" + + if self.sampler_active: + print("Running sampler...") + params = dict(mask_paths=mask_paths, ot_props=ot_props) + data_sampler = self._sampler(image_paths, labels, **params) + image_paths, labels = data_sampler["image_paths"], data_sampler["labels"] + if len(mask_paths): + mask_paths = data_sampler["mask_paths"] + if len(ot_props): + ot_props = data_sampler["ot_props"] + print(f"n samples after running sampling --- {len(image_paths)}") + return image_paths, labels, mask_paths, ot_props + + def __getitem__(self, idx): + flag = True + while flag: + try: + # Selecting data from data list + img_path = self.image_paths_r[idx] + label = self.labels_r[idx] + vid_id = img_path.split("/")[-2] + img = self._load_img(img_path) + if self.split == "test": + # Best is 17,17 and 0.0 and 5,5 + img = crop_by_margin(img, margin=[5, 5]) + + img_f = None + mask = None + mask_f = None + + # if not self.dynamic_fxray or self.split == 'val': + if not self.dynamic_fxray: + if bool(self.mask_paths_r): + mask_path = self.mask_paths_r[idx] + mask = self._load_img(mask_path) + else: + mask = np.zeros((img.shape[0], img.shape[1], 3)) + else: + if self.train: + if len(self.ot_props_r[idx]["aligned_lms"]): + f_lms = self.ot_props_r[idx]["aligned_lms"] + elif len(self.ot_props_r[idx]["orig_lms"]): + f_lms = self.ot_props_r[idx]["orig_lms"] + else: + f_lms = [] + f_lms = np.array(f_lms) + if not f_lms.any(): + raise ValueError( + "Can not find fake copy image of empty landmarks!" + ) + + if len(f_lms) > 68: + f_lms = reorder_landmark(f_lms) + + # if self.debug: + # img_lms_draw = draw_landmarks(img, f_lms) + # Image.fromarray(img_lms_draw).save(f'samples/debugs/orig_{idx}_lms.jpg') + + if self.split == "train": + if np.random.rand() < 0.5: + img, ___, f_lms, __ = sbi_hflip(img, None, f_lms, None) + + margin = np.random.randint(5, 25) + img_f, mask_f, img, mask, fake_intensity = gen_target( + img, f_lms, margin=[margin, margin], index=idx, debug=False + ) + target = None + target_f = None + + if mask is not None: + assert ( + mask.shape[:2] == img.shape[:2] + ), "Color Image and Mask must have the same shape!" + + # Applying affine transform + c, s = get_center_scale( + img.shape[:2], self.aspect_ratio, pixel_std=self.pixel_std + ) + + # Applying geo transform to images and masks + if self.split == "train": + geo_transfomed = self.geo_transform( + img, mask=mask, image_f=img_f, mask_f=mask_f + ) + img = geo_transfomed["image"] + mask = geo_transfomed["mask"] + img_f = geo_transfomed["image_f"] + mask_f = geo_transfomed["mask_f"] + + trans = get_affine_transform( + c, s, self.rot, self._cfg.IMAGE_SIZE, pixel_std=self.pixel_std + ) + trans_heatmap = get_affine_transform( + c, s, self.rot, self._cfg.HEATMAP_SIZE, pixel_std=self.pixel_std + ) + + input = cv2.warpAffine( + img, + trans, + (int(self._cfg.IMAGE_SIZE[0]), int(self._cfg.IMAGE_SIZE[1])), + flags=cv2.INTER_LINEAR, + ) + + if img_f is not None: + input_f = cv2.warpAffine( + img_f, + trans, + (int(self._cfg.IMAGE_SIZE[0]), int(self._cfg.IMAGE_SIZE[1])), + flags=cv2.INTER_LINEAR, + ) + + if mask is not None: + target = cv2.warpAffine( + mask, + trans_heatmap, + ( + int(self._cfg.HEATMAP_SIZE[0]), + int(self._cfg.HEATMAP_SIZE[1]), + ), + flags=cv2.INTER_LINEAR, + ) + + if mask_f is not None: + target_f = cv2.warpAffine( + mask_f, + trans_heatmap, + ( + int(self._cfg.HEATMAP_SIZE[0]), + int(self._cfg.HEATMAP_SIZE[1]), + ), + flags=cv2.INTER_LINEAR, + ) + + # Drawing the most vulnerable points (MVPs) + if self.debug: + mvp_f_drawed = draw_most_vul_points(target_f) + mvp_drawed = draw_most_vul_points(target) + Image.fromarray(mvp_f_drawed).save( + f"samples/debugs/mvp_f_{idx}.jpg" + ) + Image.fromarray(mvp_drawed).save(f"samples/debugs/mvp_{idx}.jpg") + + # Target encoding + # 0 for original, 1 for FXRay, 2 for NoFXRay. If 2, comment: mask = (1 - mask) * mask * 4 + heatmap, cstency_hm = ( + self.select_encode_method(version=1)(target) + if ( + target is not None + and self.heatmap_type == "gaussian" + and self.train + ) + else (None, None) + ) + heatmap_f, cstency_hm_f = ( + self.select_encode_method(version=1)(target_f) + if ( + target_f is not None + and self.heatmap_type == "gaussian" + and self.train + ) + else (None, None) + ) + + # Applying transform for blending images + if self.train: + if target_f is None: + transformed = self.transforms(image=input.astype("uint8")) + input = transformed["image"] + else: + transformed = self.transforms( + image=input.astype("uint8"), image_f=input_f.astype("uint8") + ) + input = transformed["image"] + input_f = transformed["image_f"] + + if self.debug: + Image.fromarray(input).save(f"samples/debugs/affine_{idx}.jpg") + Image.fromarray(input_f).save(f"samples/debugs/affine_f_{idx}.jpg") + Image.fromarray(target).save( + f"samples/debugs/mask_affine_{idx}.jpg" + ) + Image.fromarray(target_f).save( + f"samples/debugs/mask_affine_f_{idx}.jpg" + ) + Image.fromarray(mask).save(f"samples/debugs/mask_{idx}.jpg") + Image.fromarray(mask_f).save(f"samples/debugs/mask_f_{idx}.jpg") + vis_heatmap( + input, + cstency_hm_f / 255, + f"samples/debugs/cstency_mask_f_{idx}.jpg", + ) + vis_heatmap( + input, + cstency_hm / 255, + f"samples/debugs/cstency_mask_{idx}.jpg", + ) + vis_heatmap(input, heatmap, f"samples/debugs/hm_{idx}.jpg") + vis_heatmap(input_f, heatmap_f, f"samples/debugs/hm_f_{idx}.jpg") + + if self.train: + # if self.split_image: + # patch_img_trans = [] + # patch_heatmap_r = [] + # patch_img_trans_f = [] + # patch_heatmap_f = [] + # patch_target_r = [] + # patch_target_f = [] + + # for i, (k, l) in enumerate([[0,0], [1/2,0], [0,1/2], [1/2,1/2]]): + # input_ = input[int(self.target_h*k): int(self.target_h*(k+1/2)), int(self.target_w*l): int(self.target_w*(l+1/2)), :] + # input_f_ = input_f[int(self.target_h*k): int(self.target_h*(k+1/2)), int(self.target_w*l): int(self.target_w*(l+1/2)), :] + # heatmap_ = heatmap[i][int(self.heatmap_h*k): int(self.heatmap_h*(k+1/2)), int(self.heatmap_w*l): int(self.heatmap_w*(l+1/2))] + # heatmap_f_ = heatmap_f[i][int(self.heatmap_h*k): int(self.heatmap_h*(k+1/2)), int(self.heatmap_w*l): int(self.heatmap_w*(l+1/2))] + # target_ = target[..., 0][int(self.heatmap_h*k): int(self.heatmap_h*(k+1/2)), int(self.heatmap_w*l): int(self.heatmap_w*(l+1/2))] + # target_f_ = target_f[..., 0][int(self.heatmap_h*k): int(self.heatmap_h*(k+1/2)), int(self.heatmap_w*l): int(self.heatmap_w*(l+1/2))] + + # # Flipping + # if np.random.random() < 0.5: + # input_ = input_[:, ::-1, :] + # input_f_ = input_f_[:, ::-1, :] + # heatmap_f_ = heatmap_f_[:, ::-1] + + # #Normalise + Convert numpy array to tensor + # input_f_ = input_f_/255 + # patch_img_trans_f.append(self.final_transforms(input_f_)) + + # input_ = input_/255 + # patch_img_trans.append(self.final_transforms(input_)) + + # patch_heatmap_r.append(heatmap_) + # patch_heatmap_f.append(heatmap_f_) + # patch_target_r.append(target_/255) + # patch_target_f.append(target_f_/255) + # else: + patch_img_trans = self.final_transforms(input / 255) + patch_img_trans_f = self.final_transforms(input_f / 255) + patch_heatmap_f = heatmap_f + patch_heatmap_r = heatmap + patch_target_f = target_f / 255 + patch_target_r = target / 255 + patch_cstency_r = cstency_hm / 255 + patch_cstency_f = cstency_hm_f / 255 + else: + # Normalise + Convert numpy array to tensor + img_trans = input / 255 + img_trans = self.final_transforms(img_trans) + + label = np.expand_dims(label, axis=-1) + flag = False + except Exception as e: + # print(f'There is something wrong! Please check the DataLoader!, {e}') + flag = True + idx = torch.randint(low=0, high=self.__len__(), size=(1,)).item() + + if self.train: + return ( + patch_img_trans_f, + patch_heatmap_f, + patch_target_f, + patch_cstency_f, + patch_img_trans, + patch_heatmap_r, + patch_target_r, + patch_cstency_r, + ) + else: + meta = {"vid_id": vid_id, "img_path": img_path} + return img_trans, label, meta + + def train_collate_fn(self, batch): + batch_data = {} + + img_f, hm_f, target_f, cst_f, img_r, hm_r, target_r, cst_r = zip(*batch) + + # Collating data in case of using spliting images into patches + # if self.split_image: + # hm_H, hm_W = hm_r[0][0].shape + + # img_f = np.reshape(img_f, (-1)) + # hm_f = np.reshape(hm_f, (-1, 1, hm_H, hm_W)) + # target_f = np.reshape(target_f, (-1, 1, hm_H, hm_W)) + # img_r = np.reshape(img_r, (-1)) + # hm_r = np.reshape(hm_r, (-1, 1, hm_H, hm_W)) + # target_r = np.reshape(target_r, (-1, 1, hm_H, hm_W)) + + img = torch.cat( + [ + torch.tensor(np.array([it.numpy() for it in img_r])), + torch.tensor(np.array([it.numpy() for it in img_f])), + ], + 0, + ) + heatmap = torch.cat( + [ + torch.tensor(np.array(hm_r)).float(), + torch.tensor(np.array(hm_f)).float(), + ], + 0, + ) + target = torch.cat( + [ + torch.tensor(np.array(target_r)).float(), + torch.tensor(np.array(target_f)).float(), + ], + 0, + ) + label = torch.tensor([[0]] * len(img_r) + [[1]] * len(img_f)) + # label = torch.tensor([0] * len(img_r) + [1]*len(img_f)) + cst = torch.cat( + [ + torch.tensor(np.array(cst_r)).float(), + torch.tensor(np.array(cst_f)).float(), + ], + 0, + ) + + b_size = label.size(0) + + # Permute idxes + idxes = torch.randperm(b_size) + img, label, target, heatmap, cst = ( + img[idxes], + label[idxes], + target[idxes], + heatmap[idxes], + cst[idxes], + ) + + batch_data["img"] = img + batch_data["label"] = label + batch_data["target"] = target + batch_data["heatmap"] = heatmap + batch_data["cstency"] = cst + + return batch_data + + def train_worker_init_fn(self, worker_id): + # print('Current state {} --- worker id {}'.format(np.random.get_state()[1][0], worker_id)) + np.random.seed(np.random.get_state()[1][0] + worker_id) + + +if __name__ == "__main__": + from configs.get_config import load_config + from pipelines.geo_transform import GeometryTransform + from torch.utils.data import DataLoader + + PIPELINES.register_module(module=GeometryTransform) + + config = load_config("configs/efn4_fpn_sbi_adv.yaml") + hm_ff = DATASETS.build( + cfg=config.DATASET, default_args=dict(split="train", config=config.DATASET) + ) + hm_ff_loader = DataLoader( + hm_ff, + batch_size=10, + shuffle=True, + collate_fn=hm_ff.train_collate_fn, + worker_init_fn=hm_ff.train_worker_init_fn, + ) + + for b, batch_data in enumerate(hm_ff_loader): + inputs, labels, heatmaps, consistencies = ( + batch_data["img"], + batch_data["label"], + batch_data["heatmap"], + batch_data["cstency"], + ) + print( + f"X.shape - {inputs.shape}, y shape - {labels.shape}, heatmap shape - {heatmaps.shape}, {heatmaps.max()}, cst shape - {consistencies.shape}" + ) + break diff --git a/video/fake-stormer/model_code/datasets/fakesformer_bi.py b/video/fake-stormer/model_code/datasets/fakesformer_bi.py new file mode 100644 index 0000000000000000000000000000000000000000..6cdee4afac2dafc3c6f794ececfb8c3d7281b64d --- /dev/null +++ b/video/fake-stormer/model_code/datasets/fakesformer_bi.py @@ -0,0 +1,394 @@ +# -*- coding: utf-8 -*- +import random + +import cv2 +import numpy as np +import torch +from imgaug import augmenters as iaa +from package_utils.bi_online_generation import ( + blendImages, + colorTransfer, + random_erode_dilate, + random_get_hull, +) +from package_utils.image_utils import load_image +from package_utils.transform import ( + get_affine_transform, + get_center_scale, +) +from package_utils.utils import draw_landmarks, vis_heatmap +from PIL import Image +from skimage import transform as sktransform + +from .builder import DATASETS, PIPELINES, build_pipeline +from .master import MasterDataset + + +@DATASETS.register_module() +class FakeSFormerBI(MasterDataset): + def __init__(self, config, split, **kwargs): + """ + @params: + config: Dataset config + split: train/val/test which directs to the split folders + """ + self.split = split + super().__init__(config, **kwargs) + if kwargs is not None: + for k, v in kwargs.items(): + if v is None: + raise ValueError(f"{k}:{v} retrieve a None value!") + self.__setattr__(k, v) + self.rot = 0 + self.pixel_std = 200 + self.target_w = self._cfg.IMAGE_SIZE[1] + self.target_h = self._cfg.IMAGE_SIZE[0] + self.aspect_ratio = self.target_w * 1.0 / self.target_h + self.sigma = self._cfg.SIGMA + self.heatmap_type = self._cfg.HEATMAP_TYPE + self.debug = self._cfg.DEBUG + self.dynamic_fxray = self._cfg.DYNAMIC_FXRAY + self.target_overlap = self._cfg.TARGET_OVERLAP + + # Load data + self.image_paths, self.labels, self.mask_paths, self.ot_props = self._load_data( + split + ) + + # predefine mask distortion + self.distortion = iaa.Sequential([iaa.PiecewiseAffine(scale=(0.01, 0.15))]) + + # Calling transform methods for inputs + self.geo_transform = build_pipeline(config.TRANSFORM.geometry, PIPELINES) + self.colorjitter_transform = build_pipeline(config.TRANSFORM.color, PIPELINES) + + def _load_data(self, split): + from_file = self._cfg.DATA[self.split.upper()].FROM_FILE + + if not from_file: + image_paths, labels, mask_paths, ot_props = self._load_from_path(split) + else: + image_paths, labels, mask_paths, ot_props = self._load_from_file(split) + + if self.sampler_active and self.train: + print("Running sampler...") + params = dict(mask_paths=mask_paths, ot_props=ot_props) + data_sampler = self._sampler(image_paths, labels, **params) + image_paths, labels = data_sampler["image_paths"], data_sampler["labels"] + if len(mask_paths): + mask_paths = data_sampler["mask_paths"] + if len(ot_props): + ot_props = data_sampler["ot_props"] + print(f"n samples after running sampling --- {len(image_paths)}") + + assert ( + len(image_paths) != 0 + ), "Image paths have not been loaded! Please check image directory!" + assert ( + len(labels) != 0 + ), "Labels have not been loaded! Please check annotation file!" + # if not self.dynamic_fxray or self.split == 'val': + if from_file and not (self.dynamic_fxray): + assert ( + len(mask_paths) != 0 + ), "Mask paths have not been loaded! Please check mask directory!" + return image_paths, labels, mask_paths, ot_props + + def _reload_data(self): + self.image_paths, self.labels, self.mask_paths, self.ot_props = self._load_data( + self.split + ) + + def _gen_BI( + self, background_face, background_landmark, foreground_face_path, idx=None + ): + foreground_face = load_image(foreground_face_path) + + # down sample before blending + aug_size = random.randint(128, 317) + background_landmark = background_landmark * (aug_size / 317) + foreground_face = sktransform.resize( + foreground_face, (aug_size, aug_size), preserve_range=True + ).astype(np.uint8) + background_face = sktransform.resize( + background_face, (aug_size, aug_size), preserve_range=True + ).astype(np.uint8) + + # get random type of initial blending mask + mask = random_get_hull(background_landmark, background_face) + + # if self.debug: + # Image.fromarray((mask*255).astype(np.uint8)).save(f'samples/debugs/orig_CH_{idx}.jpg') + + # random deform mask + mask = self.distortion.augment_image(mask) + mask = random_erode_dilate(mask) + + # if self.debug: + # Image.fromarray((mask*255).astype(np.uint8)).save(f'samples/debugs/deformed_CH_{idx}.jpg') + + # filte empty mask after deformation + if np.sum(mask) == 0: + raise NotImplementedError + + # apply color transfer + foreground_face = colorTransfer(background_face, foreground_face, mask * 255) + + # blend two face + blended_face, mask = blendImages(foreground_face, background_face, mask * 255) + blended_face = blended_face.astype(np.uint8) + + # resize back to default resolution + blended_face = sktransform.resize( + blended_face, (317, 317), preserve_range=True + ).astype(np.uint8) + mask = sktransform.resize(mask, (317, 317), preserve_range=True) + mask = mask[:, :, 0:1] + return blended_face, mask + + def _gen_target( + self, background_face, background_landmark, foreground_face_path, idx=None + ): + data_label = "real" if random.randint(0, 1) else "fake" + + if not background_landmark.any(): + data_label = "real" + + if data_label == "fake": + face_img, mask = self._gen_BI( + background_face, background_landmark, foreground_face_path, idx=idx + ) + mask = (1 - mask) * mask * 4 + else: + face_img = background_face + mask = np.zeros((317, 317, 1)) + + face_img = Image.fromarray(face_img) + # randomly downsample after BI pipeline + if random.randint(0, 1): + aug_size = random.randint(64, 317) + if random.randint(0, 1): + face_img = face_img.resize((aug_size, aug_size), Image.BILINEAR) + else: + face_img = face_img.resize((aug_size, aug_size), Image.NEAREST) + face_img = face_img.resize((317, 317), Image.BILINEAR) + face_img = np.array(face_img) + + soft_margin = np.random.randint(-30, 30) + + face_img = face_img[ + 30 + soft_margin : (287 + soft_margin), + 30 + soft_margin : (287 + soft_margin), + :, + ] + mask = mask[ + 30 + soft_margin : (287 + soft_margin), + 30 + soft_margin : (287 + soft_margin), + :, + ] + mask = np.repeat(mask, 3, 2) + mask = (mask * 255).astype(np.uint8) + return face_img, mask, int(data_label == "fake") + + def __len__(self): + return len(self.labels) + + def __getitem__(self, idx): + flag = True + while flag: + try: + img_path = self.image_paths[idx] + label = self.labels[idx] + vid_id = img_path.split("/")[-2] + img = load_image(img_path) + mask = None + + # if self.debug: + # Image.fromarray(img).save(f'samples/debugs/orig_{idx}.jpg') + + # Applying color transform to inputs + if self.split == "train": + color_transfomed = self.colorjitter_transform(img) + img = color_transfomed["image"] + + # if not self.dynamic_fxray or self.split == 'val': + if not self.dynamic_fxray: + if bool(self.mask_paths): + mask_path = self.mask_paths[idx] + mask = load_image(mask_path) + else: + if self.train: + best_match = ( + self.ot_props[idx]["best_match"] + if len(self.ot_props[idx]["best_match"]) + else [] + ) + if len(self.ot_props[idx]["aligned_lms"]): + f_lms = self.ot_props[idx]["aligned_lms"] + elif len(self.ot_props[idx]["orig_lms"]): + f_lms = self.ot_props[idx]["orig_lms"] + else: + f_lms = [] + + # if self.debug: + # img_lms_draw = draw_landmarks(img, f_lms) + # Image.fromarray(img_lms_draw).save(f'samples/debugs/orig_{idx}_lms.jpg') + + if len(best_match): + best_match_idx = random.randint(0, len(best_match) - 10) + best_match_path = best_match[best_match_idx] + img, mask, label = self._gen_target( + img, f_lms, best_match_path, idx=idx + ) + else: + img, mask, label = self._gen_target(img, np.array([]), "") + else: + img = cv2.resize(img, (317, 317)) + # Best croppings from 18-299 for testing + img = img[18:(299), 18:(299), :] + target = None + + if mask is not None: + assert ( + mask.shape[:2] == img.shape[:2] + ), "Color Image and Mask must have the same shape!" + + # Applying geo transform to inputs and masks + if self.split == "train": + geo_transfomed = self.geo_transform(img, mask=mask) + img = geo_transfomed["image"] + mask = geo_transfomed["mask"] + + # Applying affine transform + c, s = get_center_scale( + img.shape[:2], self.aspect_ratio, pixel_std=self.pixel_std + ) + trans = get_affine_transform(c, s, self.rot, self._cfg.IMAGE_SIZE) + trans_heatmap = get_affine_transform( + c, s, self.rot, self._cfg.HEATMAP_SIZE + ) + + input = cv2.warpAffine( + img, + trans, + (int(self._cfg.IMAGE_SIZE[0]), int(self._cfg.IMAGE_SIZE[1])), + flags=cv2.INTER_LINEAR, + ) + + if mask is not None: + if self.target_overlap: + target = cv2.warpAffine( + mask, + trans_heatmap, + ( + int(self._cfg.HEATMAP_SIZE[0]), + int(self._cfg.HEATMAP_SIZE[1]), + ), + flags=cv2.INTER_LINEAR, + ) + else: + mask = cv2.warpAffine( + mask, + trans, + ( + int(self._cfg.IMAGE_SIZE[0]), + int(self._cfg.IMAGE_SIZE[1]), + ), + flags=cv2.INTER_LINEAR, + ) + target = self._gen_vul_parts(blending_mask=mask) + + # Target encoding + # 0 for original, 1 for FXRay, 2 for NoFXRay. If 2, comment: mask = (1 - mask) * mask * 4 + heatmap, cstency_hm = ( + self.select_encode_method(version=0)(target, fake_intensity=1.0) + if ( + target is not None + and self.heatmap_type == "gaussian" + and self.train + ) + else (None, None) + ) + + if self.debug: + Image.fromarray(input).save(f"samples/debugs/affine_{idx}.jpg") + # Image.fromarray(target).save(f'samples/debugs/mask_affine_{idx}.jpg') + vis_heatmap(input, heatmap, f"samples/debugs/hm_{idx}.jpg") + + if self.train: + patch_img_trans = self.final_transforms(input / 255) + patch_heatmap = heatmap + patch_label = np.expand_dims(label, axis=-1) + # patch_label = label + else: + # Normalise + Convert numpy array to tensor + img_trans = input / 255 + img_trans = self.final_transforms(img_trans) + label = np.expand_dims(label, axis=-1) + flag = False + except Exception as e: + print("There is an exception during loading data, please check --- ", e) + idx = torch.randint(low=0, high=self.__len__(), size=(1,)).item() + + if self.train: + return patch_img_trans, patch_label, patch_heatmap + else: + meta = {"vid_id": vid_id, "img_path": img_path} + return img_trans, label, meta + + def train_collate_fn(self, batch): + batch_data = {} + img, label, hm = zip(*batch) + + # Collating data in case of using spliting images into patches + if self.split_image: + hm_H, hm_W = hm[0][0].shape + + img = np.reshape(img, (-1)) + hm = np.reshape(hm, (-1, 1, hm_H, hm_W)) + # cstency_hm = np.reshape(cstency_hm, (-1, 1, hm_H, hm_W)) + # target = np.reshape(target, (-1, 1, hm_H, hm_W)) + label = np.reshape(label, (-1, 1)) + + img = torch.tensor(np.array([it.numpy() for it in img])) + heatmap = torch.tensor(np.array(hm)).float() + # cstency_heatmap = torch.tensor(cstency_hm).float() + # target = torch.tensor(target).float() + label = torch.tensor(np.array(label)) + + batch_data["img"] = img + batch_data["label"] = label + # batch_data["target"] = target + batch_data["heatmap"] = heatmap + # batch_data["cstency"] = cstency_heatmap + + return batch_data + + +if __name__ == "__main__": + # from datasets import * + from configs.get_config import load_config + from pipelines.color_transform import ColorJitterTransform + from pipelines.geo_transform import GeometryTransform + from torch.utils.data import DataLoader + + PIPELINES.register_module(module=GeometryTransform) + PIPELINES.register_module(module=ColorJitterTransform) + + config = load_config("configs/spatial/vit_bi_base.yaml") + hm_ff = DATASETS.build( + cfg=config.DATASET, default_args=dict(split="val", config=config.DATASET) + ) + hm_ff_loader = DataLoader( + hm_ff, batch_size=10, shuffle=False, collate_fn=hm_ff.train_collate_fn + ) + for b, batch_data in enumerate(hm_ff_loader): + inputs, labels, heatmaps = ( + batch_data["img"], + batch_data["label"], + batch_data["heatmap"], + ) + print( + f"X.shape - {inputs.shape}, y shape - {labels.shape}, heatmaps - {heatmaps.shape}" + ) + break diff --git a/video/fake-stormer/model_code/datasets/fakesformer_sbi.py b/video/fake-stormer/model_code/datasets/fakesformer_sbi.py new file mode 100644 index 0000000000000000000000000000000000000000..fd97cc8e34cdaca517200d6d959eb865f90d9e45 --- /dev/null +++ b/video/fake-stormer/model_code/datasets/fakesformer_sbi.py @@ -0,0 +1,1074 @@ +# -*- coding: utf-8 -*- +import os +import random +import sys + +import cv2 +import numpy as np +import torch +import torch.nn.functional as F +from numpy.random import randint +from package_utils.image_utils import crop_by_margin, load_image +from package_utils.transform import get_affine_transform, get_center_scale +from package_utils.utils import ( + draw_landmarks, + draw_most_vul_points, + vis_3d_heatmap, + vis_heatmap, +) +from PIL import Image + +from .builder import DATASETS, PIPELINES, build_pipeline +from .common import ParameterStore +from .master import MasterDataset +from .pipelines.geo_transform import get_transforms +from .sbi.utils import * + + +@DATASETS.register_module() +class FakeSFormerSBI(MasterDataset): + def __init__(self, config, split, **kwargs): + """ + @params: + config: Dataset config + split: train/val/test which directs to the split folders + """ + self.split = split + super(FakeSFormerSBI, self).__init__(config, **kwargs) + if kwargs is not None: + for k, v in kwargs.items(): + if v is None: + raise ValueError(f"{k}:{v} retrieve a None value!") + self.__setattr__(k, v) + self.rot = 0 + self.pixel_std = 200 + self.target_w = self._cfg.IMAGE_SIZE[1] + self.target_h = self._cfg.IMAGE_SIZE[0] + self.aspect_ratio = self.target_w * 1.0 / self.target_h + self.sigma = self._cfg.SIGMA + self.heatmap_type = self._cfg.HEATMAP_TYPE + self.debug = self._cfg.DEBUG + self.dynamic_blending_prob = self._cfg.DYNAMIC_BLENDING_PROB + self.dynamic_fxray = self._cfg.DYNAMIC_FXRAY + self.target_overlap = self._cfg.TARGET_OVERLAP + if self.data_type == "video": + self.mask_prob = self._cfg.MASK_PROB + self.temp_maskout = self._cfg.TEMP_MASKOUT + + # Load data + self.data_sampler = self._load_data(split) + + # Parse data + self._parsing_data() + + # Calling transform methods for inputs + self.geo_transform = build_pipeline( + config.TRANSFORM.geometry, + PIPELINES, + default_args={"additional_targets": {"image_f": "image", "mask_f": "mask"}}, + ) + + # predefine mask distortion + self.distortion = iaa.Sequential([iaa.PiecewiseAffine(scale=(0.01, 0.15))]) + + self.transforms = get_transforms(data_type=self.data_type) + + def __len__(self): + if self.data_type == "image": + assert "image_paths" in self.data_sampler.keys() + return len(self.labels_r) + elif self.data_type == "video": + return len(self.data_sampler.keys()) + else: + raise ValueError( + f'{self.data_type} has not been supported. Please use "image" or "video" instead!' + ) + + def _load_img(self, img_path): + return load_image(img_path) + + def _reload_data(self, epoch=0): + self.data_sampler = self._load_data(self.split, epoch=epoch) + + def _load_data(self, split, anno_file=None, epoch=0): + from_file = self._cfg.DATA[self.split.upper()].FROM_FILE + + if epoch == 0: + if not from_file: + self.image_paths, self.labels, self.mask_paths, self.ot_props = ( + self._load_from_path(split) + ) + else: + self.image_paths, self.labels, self.mask_paths, self.ot_props = ( + self._load_from_file(split, anno_file=anno_file) + ) + + assert ( + len(self.image_paths) != 0 + ), "Image paths have not been loaded! Please check image directory!" + assert ( + len(self.labels) != 0 + ), "Labels have not been loaded! Please check annotation file!" + if not self.dynamic_fxray: + assert ( + len(self.mask_paths) != 0 + ), "Mask paths have not been loaded! Please check mask directory!" + + if self.sampler_active: + print("Running sampler...") + params = dict( + mask_paths=self.mask_paths, ot_props=self.ot_props, epoch=epoch + ) + data_sampler = self._sampler(self.image_paths, self.labels, **params) + return data_sampler + + def _parsing_data(self): + assert self.data_type in ["image", "video"] + # Parsing data for training + if self.data_type == "video": + return + + self.image_paths_r, self.labels_r = ( + self.data_sampler["image_paths"], + self.data_sampler["labels"], + ) + if "mask_paths" in self.data_sampler.keys() and len( + self.data_sampler["mask_paths"] + ): + self.mask_paths_r = self.data_sampler["mask_paths"] + if "ot_props" in self.data_sampler.keys() and len( + self.data_sampler["ot_props"] + ): + self.ot_props_r = self.data_sampler["ot_props"] + + def __getitem_path__(self, idx): + param_store_ins = ParameterStore.get_instance() + # Use to store func parameters that can be reused to generate multiple blending, e.g. SBI synthesis frames + param_store_ins.add_parameters("data_type", self.data_type) + flag = True + + while flag: + try: + # Selecting data from data list + img_path = self.image_paths_r[idx] + label = self.labels_r[idx] + vid_id = img_path.split("/")[-2] + img = self._load_img(img_path) + if self.split == "test": + # Best is 9,9 and 0.0 and 11,11 + img = crop_by_margin(img, margin=[9, 9]) + + img_f = None + mask = None + mask_f = None + + # if not self.dynamic_fxray or self.split == 'val': + if not self.dynamic_fxray: + if bool(self.mask_paths_r): + mask_path = self.mask_paths_r[idx] + mask = self._load_img(mask_path) + else: + mask = np.zeros((img.shape[0], img.shape[1], 3)) + else: + if self.train: + if len(self.ot_props_r[idx]["aligned_lms"]): + f_lms = self.ot_props_r[idx]["aligned_lms"] + elif len(self.ot_props_r[idx]["orig_lms"]): + f_lms = self.ot_props_r[idx]["orig_lms"] + else: + f_lms = [] + f_lms = np.array(f_lms) + if not f_lms.any(): + raise ValueError( + "Can not find fake copy image of empty landmarks!" + ) + + if len(f_lms) > 68: + f_lms = reorder_landmark(f_lms) + + # if self.debug: + # img_lms_draw = draw_landmarks(img, f_lms) + # Image.fromarray(img_lms_draw).save(f'samples/debugs/orig_{idx}_lms.jpg') + + if self.split == "train": + if np.random.rand() < 0.5: + img, ___, f_lms, __ = sbi_hflip(img, None, f_lms, None) + + margin = np.random.randint(5, 25) + img_f, mask_f, img, mask, fake_intensity = gen_target( + img, + f_lms, + margin=[margin, margin], + index=idx, + debug=False, + dynamic_blending_prob=self.dynamic_blending_prob, + ) + target = None + target_f = None + + if mask is not None: + assert ( + mask.shape[:2] == img.shape[:2] + ), "Color Image and Mask must have the same shape!" + + # Applying affine transform + c, s = get_center_scale( + img.shape[:2], self.aspect_ratio, pixel_std=self.pixel_std + ) + + # Applying geo transform to images and masks + if self.split == "train": + geo_transfomed = self.geo_transform( + img, mask=mask, image_f=img_f, mask_f=mask_f + ) + img = geo_transfomed["image"] + mask = geo_transfomed["mask"] + img_f = geo_transfomed["image_f"] + mask_f = geo_transfomed["mask_f"] + + trans = get_affine_transform( + c, s, self.rot, self._cfg.IMAGE_SIZE, pixel_std=self.pixel_std + ) + trans_heatmap = get_affine_transform( + c, s, self.rot, self._cfg.HEATMAP_SIZE, pixel_std=self.pixel_std + ) + + input = cv2.warpAffine( + img, + trans, + (int(self._cfg.IMAGE_SIZE[0]), int(self._cfg.IMAGE_SIZE[1])), + flags=cv2.INTER_LINEAR, + ) + + if img_f is not None: + input_f = cv2.warpAffine( + img_f, + trans, + (int(self._cfg.IMAGE_SIZE[0]), int(self._cfg.IMAGE_SIZE[1])), + flags=cv2.INTER_LINEAR, + ) + + if mask is not None: + if self.target_overlap: + target = cv2.warpAffine( + mask, + trans_heatmap, + ( + int(self._cfg.HEATMAP_SIZE[0]), + int(self._cfg.HEATMAP_SIZE[1]), + ), + flags=cv2.INTER_LINEAR, + ) + else: + mask = cv2.warpAffine( + mask, + trans, + ( + int(self._cfg.IMAGE_SIZE[0]), + int(self._cfg.IMAGE_SIZE[1]), + ), + flags=cv2.INTER_LINEAR, + ) + target = self._gen_vul_parts(blending_mask=mask) + + if mask_f is not None: + if self.target_overlap: + target_f = cv2.warpAffine( + mask_f, + trans_heatmap, + ( + int(self._cfg.HEATMAP_SIZE[0]), + int(self._cfg.HEATMAP_SIZE[1]), + ), + flags=cv2.INTER_LINEAR, + ) + else: + mask_f = cv2.warpAffine( + mask_f, + trans, + ( + int(self._cfg.IMAGE_SIZE[0]), + int(self._cfg.IMAGE_SIZE[1]), + ), + flags=cv2.INTER_LINEAR, + ) + target_f = self._gen_vul_parts(blending_mask=mask_f) + + # Drawing the most vulnerable parts (MVPs) + if self.debug: + mvp_f_drawed = draw_most_vul_points(target_f) + mvp_drawed = draw_most_vul_points(target) + Image.fromarray(mvp_f_drawed).save( + f"samples/fakeformer_debugs/mvp_f_{idx}.jpg" + ) + Image.fromarray(mvp_drawed).save( + f"samples/fakeformer_debugs/mvp_{idx}.jpg" + ) + + # Target encoding + # 0 for original, 1 for FXRay, 2 for NoFXRay. If 2, comment: mask = (1 - mask) * mask * 4 + heatmap_f, cstency_hm_f = ( + self.select_encode_method(version=0)(target_f, fake_intensity=1.0) + if ( + target_f is not None + and self.heatmap_type == "gaussian" + and self.train + ) + else (None, None) + ) + heatmap, cstency_hm = ( + self.select_encode_method(version=0)(target, fake_intensity=1.0) + if ( + target is not None + and self.heatmap_type == "gaussian" + and self.train + ) + else (None, None) + ) + + # Applying transform for blending images + if self.train: + if target_f is None: + transformed = self.transforms(image=input.astype("uint8")) + input = transformed["image"] + else: + transformed = self.transforms( + image=input.astype("uint8"), image_f=input_f.astype("uint8") + ) + input = transformed["image"] + input_f = transformed["image_f"] + + if self.debug: + Image.fromarray(input).save( + f"samples/fakeformer_debugs/affine_{idx}.jpg" + ) + Image.fromarray(input_f).save( + f"samples/fakeformer_debugs/affine_f_{idx}.jpg" + ) + Image.fromarray(np.tile(target, 3)).save( + f"samples/fakeformer_debugs/mask_affine_{idx}.jpg" + ) + Image.fromarray(np.tile(target_f, 3)).save( + f"samples/fakeformer_debugs/mask_affine_f_{idx}.jpg" + ) + Image.fromarray(mask).save( + f"samples/fakeformer_debugs/mask_{idx}.jpg" + ) + Image.fromarray(mask_f).save( + f"samples/fakeformer_debugs/mask_f_{idx}.jpg" + ) + if cstency_hm is not None: + vis_heatmap( + input, + cstency_hm_f / 255, + f"samples/fakeformer_debugs/cstency_mask_f_{idx}.jpg", + ) + vis_heatmap( + input, + cstency_hm / 255, + f"samples/fakeformer_debugs/cstency_mask_{idx}.jpg", + ) + vis_heatmap( + input, heatmap, f"samples/fakeformer_debugs/hm_{idx}.jpg" + ) + vis_heatmap( + input_f, heatmap_f, f"samples/fakeformer_debugs/hm_f_{idx}.jpg" + ) + + if self.train: + patch_img_trans = self.final_transforms(input / 255) + patch_img_trans_f = self.final_transforms(input_f / 255) + patch_heatmap_f = heatmap_f + patch_heatmap_r = heatmap + patch_target_f = target_f / 255 + patch_target_r = target / 255 + patch_cstency_r = ( + cstency_hm / 255 if cstency_hm is not None else None + ) + patch_cstency_f = ( + cstency_hm_f / 255 if cstency_hm_f is not None else None + ) + else: + # Normalise + Convert numpy array to tensor + img_trans = input / 255 + img_trans = self.final_transforms(img_trans) + + label = np.expand_dims(label, axis=-1) + flag = False + except Exception as e: + print(f"There is something wrong! Please check the DataLoader!, {e}") + flag = True + idx = torch.randint(low=0, high=self.__len__(), size=(1,)).item() + + if self.train: + return ( + patch_img_trans_f, + patch_heatmap_f, + patch_target_f, + patch_cstency_f, + patch_img_trans, + patch_heatmap_r, + patch_target_r, + patch_cstency_r, + ) + else: + meta = {"vid_id": vid_id, "img_path": img_path} + return img_trans, label, meta + + def __getitem_video__(self, idx): + param_store_ins = ParameterStore.get_instance() + # Use to store func parameters that can be reused to generate multiple blending, e.g. SBI synthesis frames + param_store_ins.add_parameters("data_type", self.data_type) + flag = True + + while flag: + try: + # Real data section + inputs = [] + targets = [] + temp_loc = np.zeros(self.samples_per_video) + masked_matrixes = [] + + # Fake data section + inputs_f = [] + targets_f = [] + labels = [] + temp_loc_f = np.ones(self.samples_per_video) + masked_matrixes_f = [] + + vid_id = [*self.data_sampler.keys()][idx] + vid_data = self.data_sampler[vid_id] + + # f_idxes = randint(0, len(vid_data), self.samples_per_video) #randint might generate duplicate values, be careful! + f_idxes = range(0, self.samples_per_video) + pre_lms = None + vid_path = None + + if self.train: + seq_det = self.distortion.to_deterministic() + margin_ = np.random.randint(5, 25) + else: + # Optimal hyper-param for testing + margin_ = ( + 15 # 0 for DFW, 5 for DFDCP, 13 for DFD, and 15 for the others + ) + + for ix, f_idx in enumerate(f_idxes): + it = vid_data[f_idx] + img_path = it["image"] + if ix == 0: + vid_path = os.path.dirname(img_path) + label = it["label"] + img = self._load_img(img_path) + if self.split == "test": + img = crop_by_margin( + img, margin=[margin_, margin_] + ) # Best is 17,17 and 0.0 and 5,5 + + img_f = None + mask = None + mask_f = None + + # if not self.dynamic_fxray or self.split == 'val': + if self.train: + assert ( + self.dynamic_fxray + ), "Online blending (dynamic_fxray) is always TRUE when working with SBI!" + if "ot_props" in it.keys(): + ot_props = it["ot_props"] + + if ( + len(ot_props["aligned_lms"]) and "aligned" in img_path + ): # only take aligned lms when input already aligned + f_lms = ot_props["aligned_lms"] + elif len(ot_props["orig_lms"]): + f_lms = ot_props["orig_lms"] + else: + f_lms = [] + f_lms = np.array(f_lms) + if not f_lms.any(): + raise ValueError( + "Can not find fake copy image of empty landmarks!" + ) + + if len(f_lms) > 68: + f_lms = reorder_landmark(f_lms) + # Compute the variation of lms between each frame + if f_idx != 0: + l2_lms_dis = np.linalg.norm(f_lms - pre_lms) / len( + f_lms + ) + # print(f'Change of norm of landmark distance --- {l2_lms_dis}') + if l2_lms_dis > 0.35: + f_lms = pre_lms + (f_lms - pre_lms) / ( + round(l2_lms_dis / 0.2) + ) + pre_lms = f_lms + + if self.debug: + img_lms_draw = draw_landmarks(img, f_lms) + Image.fromarray(img_lms_draw).save( + f"samples/debugs/orig_{idx}_{f_idx}_lms.jpg" + ) + + if self.split == "train": + rand_flip = param_store_ins.get_parameters("rand_flip") + if rand_flip is None: + rand_flip = np.random.rand() < 0.5 + param_store_ins.add_parameters( + "rand_flip", rand_flip + ) + + if rand_flip < 0.5: + img, ___, f_lms, __ = sbi_hflip( + img, None, f_lms, None + ) + + img_f, mask_f, img, mask, fake_intensity = gen_target( + img, + f_lms, + margin=[margin_, margin_], + index=f_idx, + debug=False, + dynamic_blending_prob=self.dynamic_blending_prob, + distortion=seq_det, + ) + + target = None + target_f = None + + if mask is not None: + assert ( + mask.shape[:2] == img.shape[:2] + ), "Color Image and Mask must have the same shape!" + + # Applying affine transform + c, s = get_center_scale( + img.shape[:2], self.aspect_ratio, pixel_std=self.pixel_std + ) + + # Applying geo transform to images and masks + # if self.split == 'train': + # geo_transfomed = self.geo_transform(img, mask=mask, image_f=img_f, mask_f=mask_f) + # img = geo_transfomed['image'] + # mask = geo_transfomed['mask'] + # img_f = geo_transfomed['image_f'] + # mask_f = geo_transfomed['mask_f'] + + trans = get_affine_transform( + c, s, self.rot, self._cfg.IMAGE_SIZE, pixel_std=self.pixel_std + ) + trans_heatmap = get_affine_transform( + c, s, self.rot, self._cfg.HEATMAP_SIZE, pixel_std=self.pixel_std + ) + + input = cv2.warpAffine( + img, + trans, + (int(self._cfg.IMAGE_SIZE[0]), int(self._cfg.IMAGE_SIZE[1])), + flags=cv2.INTER_LINEAR, + ) + + if img_f is not None: + input_f = cv2.warpAffine( + img_f, + trans, + ( + int(self._cfg.IMAGE_SIZE[0]), + int(self._cfg.IMAGE_SIZE[1]), + ), + flags=cv2.INTER_LINEAR, + ) + + if mask is not None: + if self.target_overlap: + target = cv2.warpAffine( + mask, + trans_heatmap, + ( + int(self._cfg.HEATMAP_SIZE[0]), + int(self._cfg.HEATMAP_SIZE[1]), + ), + flags=cv2.INTER_LINEAR, + ) + else: + mask = cv2.warpAffine( + mask, + trans, + ( + int(self._cfg.IMAGE_SIZE[0]), + int(self._cfg.IMAGE_SIZE[1]), + ), + flags=cv2.INTER_LINEAR, + ) + target = self._gen_vul_parts(blending_mask=mask) + + if mask_f is not None: + if self.target_overlap: + target_f = cv2.warpAffine( + mask_f, + trans_heatmap, + ( + int(self._cfg.HEATMAP_SIZE[0]), + int(self._cfg.HEATMAP_SIZE[1]), + ), + flags=cv2.INTER_LINEAR, + ) + else: + mask_f = cv2.warpAffine( + mask_f, + trans, + ( + int(self._cfg.IMAGE_SIZE[0]), + int(self._cfg.IMAGE_SIZE[1]), + ), + flags=cv2.INTER_LINEAR, + ) + target_f = self._gen_vul_parts(blending_mask=mask_f) + + # Drawing the most vulnerable parts (MVPs) + if self.debug: + mvp_drawed = draw_most_vul_points(target_f) + Image.fromarray(mvp_drawed).save( + f"samples/debugs/mvp_{idx}_{f_idx}_{fake_intensity}.png" + ) + + # Applying transform for blending images + if self.train: + if target_f is None: + if f_idx == 0: + transformed = self.transforms( + image=input.astype(np.uint8) + ) + input = transformed["image"] + replay_params = transformed["replay"] + param_store_ins.add_parameters( + "replay_transform", replay_params + ) + else: + replay_params = param_store_ins.get_parameters( + "replay_transform" + ) + data = alb.ReplayCompose.replay( + replay_params, image=input.astype(np.uint8) + ) + input = data["image"] + else: + if f_idx == 0: + transformed = self.transforms( + image=input.astype(np.uint8), + image_f=input_f.astype(np.uint8), + ) + input = transformed["image"] + input_f = transformed["image_f"] + replay_params = transformed["replay"] + param_store_ins.add_parameters( + "replay_transform", replay_params + ) + else: + replay_params = param_store_ins.get_parameters( + "replay_transform" + ) + data = alb.ReplayCompose.replay( + replay_params, + image=input.astype(np.uint8), + image_f=input_f.astype(np.uint8), + ) + input = data["image"] + input_f = data["image_f"] + + mask_out_rand = ( + param_store_ins.get_parameters("mask_out_rand") + or np.random.rand() + ) + label = 1 + if f_idx == 0: + param_store_ins.add_parameters( + "mask_out_rand", mask_out_rand + ) + if mask_out_rand > 0.5 and self.mask_prob > 0: + # Mask out vulnerabilities + ( + input_f, + target_f, + masked_matrix_f, + upper_bound_value, + fake_intensity, + ) = self._mask_out_vulnerability2( + input_f, + target_f, + fake_intensity=fake_intensity, + mask_prob=self.mask_prob, + ) + if upper_bound_value == 1: + label = 0 + + temp_loc_f[f_idx] = ( + fake_intensity # Updating temporal location value, default 0 + ) + input, target, masked_matrix, upper_bound_value, _ = ( + self._mask_out_vulnerability2( + input, + target, + fake_intensity=fake_intensity, + mask_prob=self.mask_prob, + ) + ) + else: + masked_matrix = np.ones_like(target[..., 0]) + masked_matrix_f = np.ones_like(target_f[..., 0]) + + targets.append(target) + targets_f.append(target_f) + masked_matrixes.append(masked_matrix) + masked_matrixes_f.append(masked_matrix_f) + + if self.debug: + # Image.fromarray(mask).save(f'samples/debugs/mask_{fake_intensity}.jpg') + # Image.fromarray(mask_f).save(f'samples/debugs/mask_f_{fake_intensity}.jpg') + Image.fromarray(input).save( + f"samples/debugs/affine_{idx}_{f_idx}_{fake_intensity}.jpg" + ) + Image.fromarray(input_f).save( + f"samples/debugs/affine_f_{idx}_{f_idx}_{fake_intensity}.png" + ) + # Image.fromarray(target).save(f'samples/debugs/mask_affine_{idx}_{f_idx}_{fake_intensity}.jpg') + Image.fromarray(target_f).save( + f"samples/debugs/mask_affine_f_{idx}_{f_idx}_{fake_intensity}.png" + ) + Image.fromarray(mask_f).save( + f"samples/debugs/mask_f_{idx}_{f_idx}_{fake_intensity}.png" + ) + + if self.train: + img_trans = self.final_transforms(input / 255) + img_trans_f = self.final_transforms(input_f / 255) + inputs_f.append(img_trans_f) + else: + # Normalise + Convert numpy array to tensor + img_trans = input / 255 + img_trans = self.final_transforms(img_trans) + + inputs.append(img_trans) + labels.append(label) + + if self.train and self.temp_maskout: + if np.random.rand() < 0.5: + start_idx = np.random.randint(0, len(inputs) - 1) + end_idx = min( + len(inputs), + start_idx + np.random.randint(1, int(len(inputs) / 2 + 1)), + ) + targets_f = np.array(targets_f) + targets = np.array(targets) + + if np.random.rand() > 0.5: # Repeat + # inputs_f[start_idx:end_idx] = inputs[start_idx:end_idx] + inputs_f[start_idx:end_idx] = ( + inputs_f[start_idx] + .unsqueeze(0) + .repeat(end_idx - start_idx, 1, 1, 1) + ) + targets_f[start_idx:end_idx] = targets_f[ + np.newaxis, start_idx + ].repeat(end_idx - start_idx, 0) + temp_loc_f[start_idx:end_idx] = temp_loc_f[ + np.newaxis, start_idx + ].repeat(end_idx - start_idx) + else: # Temporal cutout + zero_transform = self.final_transforms( + np.zeros( + ( + int(self._cfg.IMAGE_SIZE[0]), + int(self._cfg.IMAGE_SIZE[1]), + 3, + ) + ) + ) + + # Assigning values for the range from start_idx to end_idx + inputs_f.extend(inputs_f[start_idx:end_idx]) + targets_f = np.append( + targets_f, targets_f[start_idx:end_idx], 0 + ) + temp_loc_f = np.append( + temp_loc_f, temp_loc_f[start_idx:end_idx], 0 + ) + + inputs_f[start_idx:end_idx] = zero_transform.unsqueeze( + 0 + ).repeat(end_idx - start_idx, 1, 1, 1) + targets_f[start_idx:end_idx] = targets[start_idx:end_idx] + temp_loc_f[start_idx:end_idx] = temp_loc[start_idx:end_idx] + + inputs_f = inputs_f[: self.samples_per_video] + targets_f = targets_f[: self.samples_per_video] + temp_loc_f = temp_loc_f[: self.samples_per_video] + + labels[start_idx:end_idx] = [ + 0 for i in range(end_idx - start_idx) + ] + + # Target encoding + # 0 for original, 1 for FXRay, 2 for NoFXRay. If 2, find and comment the line: mask = (1 - mask) * mask * 4 + heatmap_f, cstency_hm_f, normalized_params = ( + self.select_encode_method(version=0, dimension="temporal")( + targets_f, fake_intensity=fake_intensity + ) + if (len(targets_f) and self.train) + else (None, None, None) + ) + heatmap, cstency_hm, _ = ( + self.select_encode_method(version=0, dimension="temporal")( + targets, fake_intensity=fake_intensity, **normalized_params + ) + if (len(targets) and self.train) + else (None, None, None) + ) + + # Debugging 3D heatmap + if self.debug: + vis_3d_heatmap(heatmap, f"samples/debugs/hm_{idx}_{f_idx}.png") + vis_3d_heatmap(heatmap_f, f"samples/debugs/hm_f_{idx}_{f_idx}.png") + vis_3d_heatmap( + cstency_hm, f"samples/debugs/cstency_{idx}_{f_idx}.png" + ) + vis_3d_heatmap( + cstency_hm_f, f"samples/debugs/cstency_f_{idx}_{f_idx}.png" + ) + + # End for loop + if not self.train: + inputs = torch.tensor( + np.array([[j.numpy() for j in i] for i in inputs]) + ).transpose(0, 1) + labels = torch.tensor(np.array([np.array(it) for it in labels])) + label = torch.max(labels).unsqueeze(0) + else: + label = np.max(labels) + if label == 0: + raise ValueError( + "There is at least one frame containing artifacts!" + ) + + flag = False + ParameterStore.reset() + except Exception as e: + print(f"There is something wrong! Please check the DataLoader!, {e}") + flag = True + idx = torch.randint(low=0, high=self.__len__(), size=(1,)).item() + + if self.train: + return ( + inputs, + inputs_f, + heatmap, + heatmap_f, + 0, + 1, + temp_loc, + temp_loc_f, + masked_matrixes, + masked_matrixes_f, + cstency_hm, + cstency_hm_f, + ) + else: + meta = {"vid_id": vid_id.split("+++")[0], "vid_path": vid_path} + return inputs, label, meta + + def __getitem__(self, idx): + if self.data_type == "image": + return self.__getitem_path__(idx=idx) + elif self.data_type == "video": + return self.__getitem_video__(idx=idx) + else: + raise ValueError( + f"{self.data_type} has not been supported. Only image or video are used for training!" + ) + + def train_collate_fn(self, batch): + batch_data = {} + + if self.data_type == "image": + img_f, hm_f, target_f, cst_f, img_r, hm_r, target_r, cst_r = zip(*batch) + + img = torch.cat( + [ + torch.tensor(np.array([it.numpy() for it in img_r])), + torch.tensor(np.array([it.numpy() for it in img_f])), + ], + 0, + ) + heatmap = torch.cat( + [ + torch.tensor(np.array(hm_r)).float(), + torch.tensor(np.array(hm_f)).float(), + ], + 0, + ) + target = torch.cat( + [ + torch.tensor(np.array(target_r)).float(), + torch.tensor(np.array(target_f)).float(), + ], + 0, + ) + label = torch.tensor([[0]] * len(img_r) + [[1]] * len(img_f)) + cst = ( + torch.cat( + [ + torch.tensor(np.array(cst_r)).float(), + torch.tensor(np.array(cst_f)).float(), + ], + 0, + ) + if None not in cst_r + else None + ) + + b_size = label.size(0) + + # Permute idxes + idxes = torch.randperm(b_size) + img, label, target, heatmap = ( + img[idxes], + label[idxes], + target[idxes], + heatmap[idxes], + ) + if cst is not None: + cst = cst[idxes] + + batch_data["img"] = img + batch_data["label"] = label + batch_data["target"] = target + batch_data["heatmap"] = heatmap + batch_data["cstency"] = cst + else: + ( + img_r, + img_f, + hm_r, + hm_f, + label_r, + label_f, + temp_loc_r, + temp_loc_f, + mask_idx_r, + mask_idx_f, + cst_r, + cst_f, + ) = zip(*batch) + + img = torch.cat( + [ + torch.tensor( + np.array([[j.numpy() for j in i] for i in img_r]) + ).transpose(1, 2), + torch.tensor( + np.array([[j.numpy() for j in i] for i in img_f]) + ).transpose(1, 2), + ], + 0, + ) + label = torch.cat( + [ + torch.tensor([i for i in label_r]), + torch.tensor([i for i in label_f]), + ], + 0, + ).unsqueeze(1) + hm = torch.cat( + [ + torch.tensor(np.array(hm_r)).float(), + torch.tensor(np.array(hm_f)).float(), + ], + 0, + ).unsqueeze(1) + temp_loc = torch.cat( + [ + torch.tensor(np.array([i for i in temp_loc_r])), + torch.tensor(np.array([i for i in temp_loc_f])), + ], + 0, + ) + mask_idx = torch.cat( + [ + torch.tensor(np.array(mask_idx_r)), + torch.tensor(np.array(mask_idx_f)), + ], + 0, + ).unsqueeze(1) + cst = torch.cat( + [ + torch.tensor(np.array(cst_r)).float(), + torch.tensor(np.array(cst_f)).float(), + ], + 0, + ).unsqueeze(1) + + b_size = label.size(0) + + # Permute idxes + idxes = torch.randperm(b_size) + img, label, hm, temp_loc, mask_idx, cst = ( + img[idxes], + label[idxes], + hm[idxes], + temp_loc[idxes], + mask_idx[idxes], + cst[idxes], + ) + + batch_data["img"] = img + batch_data["label"] = label + batch_data["heatmap"] = hm + batch_data["temp_loc"] = temp_loc + batch_data["mask_out_p"] = mask_idx + batch_data["cstency"] = cst + + return batch_data + + def train_worker_init_fn(self, worker_id): + # print('Current state {} --- worker id {}'.format(np.random.get_state()[1][0], worker_id)) + np.random.seed(np.random.get_state()[1][0] + worker_id) + + +if __name__ == "__main__": + from configs.get_config import load_config + from pipelines.geo_transform import GeometryTransform + from torch.utils.data import DataLoader + + PIPELINES.register_module(module=GeometryTransform) + + config = load_config("configs/temporal/ResNet3D_EFPN3D_hm3D_c23.yaml") + + # Seed + seed = 529 + random.seed(seed) + torch.manual_seed(seed) + np.random.seed(seed) + torch.cuda.manual_seed(seed) + + hm_ff = DATASETS.build( + cfg=config.DATASET, default_args=dict(split="val", config=config.DATASET) + ) + hm_ff_loader = DataLoader( + hm_ff, + batch_size=8, + shuffle=True, + collate_fn=hm_ff.train_collate_fn, + worker_init_fn=hm_ff.train_worker_init_fn, + ) + + for b, batch_data in enumerate(hm_ff_loader): + inputs, labels, heatmaps = ( + batch_data["img"], + batch_data["label"], + batch_data["heatmap"], + ) + print( + f"X.shape - {inputs.shape}, y shape - {labels.shape}, heatmap shape - {heatmaps.shape}" + ) + break diff --git a/video/fake-stormer/model_code/datasets/ff.py b/video/fake-stormer/model_code/datasets/ff.py new file mode 100644 index 0000000000000000000000000000000000000000..9844a63f9094dc81607806e8320437411f161495 --- /dev/null +++ b/video/fake-stormer/model_code/datasets/ff.py @@ -0,0 +1,42 @@ +# -*- coding: utf-8 -*- +import os +from glob import glob + +import numpy as np + +from .builder import DATASETS +from .common import CommonDataset + + +@DATASETS.register_module() +class FF(CommonDataset): + def __init__(self, cfg, **kwargs): + super().__init__(cfg, **kwargs) + + def _load_from_path(self, split): + assert os.path.exists( + self._cfg.DATA[self.split.upper()].ROOT + ), "Root path to dataset can not be None!" + data = self._cfg["DATA"] + data_type = data.TYPE + fake_types = self._cfg.DATA[split.upper()]["FAKETYPE"] + label_folders = self._cfg.DATA[split.upper()]["LABEL_FOLDER"] + img_paths, labels, mask_paths, ot_props = [], [], [], [] + + # Load image data for each type of fake techniques + for idx, ft in enumerate(fake_types): + data_dir = os.path.join( + self._cfg.DATA[self.split.upper()].ROOT, self.split, data_type, ft + ) + if not os.path.exists(data_dir): + raise ValueError("Data Directory can not be invalid!") + + for sub_dir in os.listdir(data_dir): + sub_dir_path = os.path.join(data_dir, sub_dir) + img_paths_ = glob(f"{sub_dir_path}/*.{self._cfg.IMAGE_SUFFIX}") + + img_paths.extend(img_paths_) + labels.extend(np.full(len(img_paths_), int("original" not in ft))) + + print("{} image paths have been loaded from FF++!".format(len(img_paths))) + return img_paths, labels, mask_paths, ot_props diff --git a/video/fake-stormer/model_code/datasets/master.py b/video/fake-stormer/model_code/datasets/master.py new file mode 100644 index 0000000000000000000000000000000000000000..18d546bd481bf5709c06cc36f348d462a2ffd104 --- /dev/null +++ b/video/fake-stormer/model_code/datasets/master.py @@ -0,0 +1,48 @@ +# -*- coding: utf-8 -*- +from .builder import DATASETS +from .celebDF_v1 import CDFV1 +from .celebDF_v2 import CDFV2 +from .combine import Combine +from .df40 import DF40 +from .dfd import DFD +from .dfdc import DFDC +from .dfdcp import DFDCP +from .dfo import DFo +from .dfw import DFW +from .diffswap import DiffSwap +from .ff import FF + + +@DATASETS.register_module() +class MasterDataset( + CDFV1, FF, DFDCP, CDFV2, DFDC, DFD, DFW, DFo, DiffSwap, DF40, Combine +): + def __init__(self, cfg, **kwargs): + super().__init__(cfg, **kwargs) + + def _load_from_path(self, split): + # Explicitly overide some main methods from the dataset config + if self.dataset == "FF++": + return MasterDataset.__mro__[2]._load_from_path(self, split=split) + elif self.dataset == "Celeb-DFv1": + return MasterDataset.__mro__[1]._load_from_path(self, split=split) + elif self.dataset == "DFDCP": + return MasterDataset.__mro__[3]._load_from_path(self, split=split) + elif self.dataset == "Celeb-DFv2": + return MasterDataset.__mro__[4]._load_from_path(self, split=split) + elif self.dataset == "DFDC": + return MasterDataset.__mro__[5]._load_from_path(self, split=split) + elif self.dataset == "DFD": + return MasterDataset.__mro__[6]._load_from_path(self, split=split) + elif self.dataset == "DFW": + return MasterDataset.__mro__[7]._load_from_path(self, split=split) + elif self.dataset == "DFo": + return MasterDataset.__mro__[8]._load_from_path(self, split=split) + elif self.dataset == "DiffSwap": + return MasterDataset.__mro__[9]._load_from_path(self, split=split) + elif self.dataset == "DF40": + return MasterDataset.__mro__[10]._load_from_path(self, split=split) + elif self.dataset == "Combine": + return MasterDataset.__mro__[11]._load_from_path(self, split=split) + else: + return NotImplementedError(f"{self.dataset} has not been supported yet!") diff --git a/video/fake-stormer/model_code/datasets/pipelines/__init__.py b/video/fake-stormer/model_code/datasets/pipelines/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3be7f6cb584fa2e32c5a9f17f7511fc0ee65e738 --- /dev/null +++ b/video/fake-stormer/model_code/datasets/pipelines/__init__.py @@ -0,0 +1,5 @@ +# -*- coding: utf-8 -*- +from .color_transform import ColorJitterTransform +from .geo_transform import GeometryTransform + +__all__ = ["GeometryTransform", "ColorJitterTransform"] diff --git a/video/fake-stormer/model_code/datasets/pipelines/color_transform.py b/video/fake-stormer/model_code/datasets/pipelines/color_transform.py new file mode 100644 index 0000000000000000000000000000000000000000..d2175d72e0a423f8ec822bdc6cd38502565378dd --- /dev/null +++ b/video/fake-stormer/model_code/datasets/pipelines/color_transform.py @@ -0,0 +1,193 @@ +import os +import sys + +if not os.getcwd() in sys.path: + sys.path.append(os.getcwd()) + +import albumentations as A +from datasets.builder import PIPELINES + + +@PIPELINES.register_module() +class ColorJitterTransform(object): + def __init__( + self, + clahe: float, + colorjitter: float, + gaussianblur: float, + jpegcompression: list, + rgbshift: float, + randomcontrast: float, + randomgamma: float, + randombrightness: float, + huesat: float, + gaussnoise: float, + *args, + **kwargs, + ): + super().__init__() + self.clahe = clahe + self.colorjitter = colorjitter + self.gaussianblur = gaussianblur + self.jpegcompression = jpegcompression + self.rgbshift = rgbshift + self.randomcontrast = randomcontrast + self.randomgamma = randomgamma + self.randombrightness = randombrightness + self.huesat = huesat + self.gaussnoise = gaussnoise + + if kwargs is not None: + for k, v in kwargs.items(): + if v is None: + raise ValueError(f"{k}:{v} retrieve a None value!") + self.__setattr__(k, v) + + def _CLAHE(self, clip_limit=4.0, tile_grid_size=(8, 8), always_apply=False, p=0.5): + return A.CLAHE( + clip_limit=clip_limit, + tile_grid_size=tile_grid_size, + always_apply=always_apply, + p=p, + ) + + def _colorjitter( + self, + brightness=0.2, + contrast=0.2, + saturation=0.2, + hue=0.2, + always_apply=False, + p=0.5, + ): + return A.ColorJitter( + brightness=brightness, + contrast=contrast, + saturation=saturation, + hue=hue, + always_apply=always_apply, + p=p, + ) + + def _gaussianblur( + self, blur_limit=(3, 7), sigma_limit=0, always_apply=False, p=0.5 + ): + return A.GaussianBlur( + blur_limit=blur_limit, + sigma_limit=sigma_limit, + always_apply=always_apply, + p=p, + ) + + def _gauss_noise( + self, + var_limit=(10.0, 50.0), + mean=0, + per_channel=True, + always_apply=False, + p=0.5, + ): + return A.GaussNoise( + var_limit=var_limit, + mean=mean, + per_channel=per_channel, + always_apply=always_apply, + p=p, + ) + + def _jpegcompression( + self, quality_lower=70, quality_upper=100, always_apply=False, p=0.5 + ): + return A.ImageCompression( + quality_lower=quality_lower, + quality_upper=quality_upper, + always_apply=always_apply, + p=p, + ) + + def _rgbshift( + self, + r_shift_limit=20, + g_shift_limit=20, + b_shift_limit=20, + always_apply=False, + p=0.5, + ): + return A.RGBShift( + r_shift_limit=r_shift_limit, + g_shift_limit=g_shift_limit, + b_shift_limit=b_shift_limit, + always_apply=always_apply, + p=p, + ) + + def _randomcontrast(self, limit=0.2, always_apply=False, p=0.5): + return A.RandomContrast(limit=limit, always_apply=always_apply, p=p) + + def _randombrightness( + self, + brightness_limit=0.1, + contrast_limit=0.1, + brightness_by_max=True, + always_apply=False, + p=0.5, + ): + return A.RandomBrightnessContrast( + brightness_limit=brightness_limit, + contrast_limit=contrast_limit, + brightness_by_max=brightness_by_max, + always_apply=always_apply, + p=p, + ) + + def _randomgamma(self, gamma_limit=(80, 120), eps=None, always_apply=False, p=0.5): + return A.RandomGamma( + gamma_limit=gamma_limit, eps=eps, always_apply=always_apply, p=p + ) + + def _huesaturation( + self, + hue_shift_limit=20, + sat_shift_limit=20, + val_shift_limit=20, + always_apply=False, + p=0.5, + ): + return A.HueSaturationValue( + hue_shift_limit=hue_shift_limit, + sat_shift_limit=sat_shift_limit, + val_shift_limit=val_shift_limit, + always_apply=always_apply, + p=p, + ) + + def __call__(self, x): + transforms = [ + A.Compose( + [ + self._CLAHE(p=self.clahe), + self._randomcontrast(p=self.randomcontrast), + self._colorjitter(p=self.colorjitter), + self._jpegcompression( + p=self.jpegcompression[0], + quality_lower=self.jpegcompression[1], + quality_upper=self.jpegcompression[2], + ), + self._rgbshift(p=self.rgbshift), + self._randomgamma(p=self.randomgamma), + ] + ), + A.OneOf( + [ + self._gaussianblur(p=self.gaussianblur), + self._gauss_noise(p=self.gaussnoise), + ] + ), + A.OneOf( + [ + self._randombrightness(p=self.randombrightness), + self._huesaturation(p=self.huesat), + ] + ), + ] + return A.Compose(transforms)(image=x) diff --git a/video/fake-stormer/model_code/datasets/pipelines/functional.py b/video/fake-stormer/model_code/datasets/pipelines/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..2730494f2301b94bbc8debbe5c9394fb227544a8 --- /dev/null +++ b/video/fake-stormer/model_code/datasets/pipelines/functional.py @@ -0,0 +1,11 @@ +# -*- coding: utf-8 -*- +import numpy as np + + +def _get_pixels(per_pixel, rand_color, patch_size, dtype=np.float32): + if per_pixel: + return np.random.randint(0, 255, patch_size).astype(dtype=dtype) + elif rand_color: + return np.random.randint(0, 255, (1, 1, patch_size[2])).astype(dtype=dtype) + else: + return np.zeros((1, 1, patch_size[2]), dtype=dtype) diff --git a/video/fake-stormer/model_code/datasets/pipelines/geo_transform.py b/video/fake-stormer/model_code/datasets/pipelines/geo_transform.py new file mode 100644 index 0000000000000000000000000000000000000000..d7c39de0f531bd0cd77ff59606c0b4f7f7d27d71 --- /dev/null +++ b/video/fake-stormer/model_code/datasets/pipelines/geo_transform.py @@ -0,0 +1,372 @@ +# -*- coding: utf-8 -*- +import math +import os +import random +import sys +from typing import Dict + +if not os.getcwd() in sys.path: + sys.path.append(os.getcwd()) + +import albumentations as A +import cv2 +import numpy as np +from albumentations.augmentations.transforms import DualTransform +from albumentations.core.transforms_interface import ImageOnlyTransform +from datasets.builder import PIPELINES + +from .functional import _get_pixels + + +@PIPELINES.register_module() +class GeometryTransform(object): + def __init__( + self, + resize: list, + normalize: float, + horizontal_flip: float, + scale: list, + cropping: list, + rand_erasing: list, + *args, + **kwargs, + ): + super().__init__() + self.resize = resize # [H, W, p] + self.normalize = normalize # p + self.horizontal_flip = horizontal_flip # p + self.cropping = cropping # [crop_limit, p] + self.scale = scale # [scale_limit, p] + self.rand_erasing = rand_erasing # p + + if kwargs is not None: + for k, v in kwargs.items(): + if v is None: + raise ValueError(f"{k}:{v} retrieve a None value!") + self.__setattr__(k, v) + + def _resize(self): + hr, wr, p = self.resize + return A.Resize(hr, wr, interpolation=2, p=p) + + # We offen use normalize transform from torch, so set p=0.0 + def _normalize(self, p=0.0): + return A.Normalize(p=p) + + def _horizontal_flip(self, p=0.5, always_apply=False): + return A.HorizontalFlip(always_apply=always_apply, p=p) + + def _random_scale( + self, p=0.5, always_apply=False, scale_limit=0.1, interpolation=1 + ): + return A.RandomScale( + scale_limit=scale_limit, + interpolation=interpolation, + always_apply=always_apply, + p=p, + ) + + def _random_crop( + self, p=0.5, always_apply=False, crop_limit=0.1, img_h=256, img_w=256 + ): + crop_h = int((1 - np.random.choice(np.arange(0.0, crop_limit, 0.01))) * img_h) + crop_w = int((1 - np.random.choice(np.arange(0.0, crop_limit, 0.01))) * img_w) + return A.RandomCrop(height=crop_h, width=crop_w, always_apply=always_apply, p=p) + + def _random_erasing(self, p=0.5, always_apply=False, max_count=3, mode="const"): + return RandomErasing( + p=p, always_apply=always_apply, mode=mode, max_count=max_count + ) + + def __call__(self, x, mask=None, image_f=None, mask_f=None): + x_h, x_w = x.shape[:2] + + if hasattr(self, "additional_targets"): + additional_targets = self.__getattribute__("additional_targets") + else: + additional_targets = {} + + transform = A.Compose( + [ + A.OneOf( + [ + self._random_crop( + p=self.cropping[1], + crop_limit=self.cropping[0], + img_h=x_h, + img_w=x_w, + ), + self._random_scale(p=self.scale[1], scale_limit=self.scale[0]), + self._random_erasing( + p=self.rand_erasing[0], + mode="const", + max_count=self.rand_erasing[1], + ), + ] + ), + A.Compose( + [ + self._resize(), + self._normalize(p=self.normalize), + self._horizontal_flip(p=self.horizontal_flip), + ] + ), + ], + additional_targets=additional_targets, + ) + + if mask is not None: + if mask_f is not None: + assert ( + image_f is not None + ), "Image Fake sample can not be None in case of Mask sample!" + assert len( + additional_targets.keys() + ), "Additional targets for Albumentations can not be None!" + return transform(image=x, mask=mask, image_f=image_f, mask_f=mask_f) + else: + return transform(image=x, mask=mask) + else: + return transform(image=x) + + +class RandomErasing(DualTransform): + def __init__( + self, + always_apply: bool = False, + p: float = 0.5, + min_area=0.02, + max_area=1 / 3, + min_aspect=0.3, + max_aspect=None, + mode="const", + min_count=1, + max_count=None, + num_splits=0, + img_h=257, + img_w=257, + img_chan=3, + ): + super(RandomErasing, self).__init__(always_apply, p) + self.min_area = min_area + self.max_area = max_area + max_aspect = max_aspect or 1 / min_aspect + self.log_aspect_ratio = (math.log(min_aspect), math.log(max_aspect)) + self.min_count = min_count + self.max_count = max_count or min_count + self.num_splits = num_splits + mode = mode.lower() + self.rand_color = False + self.per_pixel = False + self.img_h = img_h + self.img_w = img_w + self.img_chan = img_chan + + if mode == "rand": + self.rand_color = True # per block random normal + elif mode == "pixel": + self.per_pixel = True # per pixel random normal + else: + assert not mode or mode == "const" + + def apply(self, img: np.array, **params): + return self._erase(img, **params) + + def get_params(self) -> Dict: + area = self.img_h * self.img_w + count = ( + self.min_count + if self.min_count == self.max_count + else random.randint(self.min_count, self.max_count) + ) + + tops, lefts, ws, hs = [], [], [], [] + for _ in range(count): + for attempt in range(10): + target_area = ( + random.uniform(self.min_area, self.max_area) * area / count + ) + aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio)) + h = int(round(math.sqrt(target_area * aspect_ratio))) + w = int(round(math.sqrt(target_area / aspect_ratio))) + + if w < self.img_w and h < self.img_h: + top = random.randint(0, self.img_h - h) + left = random.randint(0, self.img_w - w) + + tops.append(top) + lefts.append(left) + ws.append(w) + hs.append(h) + break + return { + "tops": tops, + "lefts": lefts, + "ws": ws, + "hs": hs, + "img_chan": self.img_chan, + } + + def _erase( + self, + img: np.array, + tops: list, + lefts: list, + hs: list, + ws: list, + img_chan: int, + **params, + ): + for i in range(len(tops)): + top = tops[i] + left = lefts[i] + w = ws[i] + h = hs[i] + + img[top : top + h, left : left + w, :] = _get_pixels( + self.per_pixel, self.rand_color, (h, w, img_chan), dtype=img.dtype + ) + return img + + +class RandomDownScale(ImageOnlyTransform): + def __init__( + self, always_apply: bool = False, p: float = 0.5, ratio_list: list = [2, 4] + ): + self.ratio_list = ratio_list + super().__init__(p=p, always_apply=always_apply) + + def apply(self, img: np.ndarray, ratio: int, **params): + return self.randomdownscale(img, ratio, **params) + + def get_params(self): + ratio = self.ratio_list[np.random.randint(len(self.ratio_list))] + return {"ratio": ratio} + + def randomdownscale(self, img, ratio, **kwargs): + keep_ratio = True + keep_input_shape = True + H, W, C = img.shape + + # r = np.random.uniform(2, 4) + img_ds = cv2.resize( + img, (int(W / ratio), int(H / ratio)), interpolation=cv2.INTER_NEAREST + ) + if keep_input_shape: + img_ds = cv2.resize(img_ds, (W, H), interpolation=cv2.INTER_LINEAR) + + return img_ds + + def get_transform_init_args_names(self): + return ("ratio_list",) + + +def get_source_transforms(data_type="image"): + """ + Transforms specially design for SBI synthesis + """ + assert data_type in ["image", "video"] + if data_type == "image": + return A.Compose( + [ + A.Compose( + [ + A.RGBShift((-20, 20), (-20, 20), (-20, 20), p=0.3), + A.HueSaturationValue( + hue_shift_limit=(-0.3, 0.3), + sat_shift_limit=(-0.3, 0.3), + val_shift_limit=(-0.3, 0.3), + p=1, + ), + A.RandomBrightnessContrast( + brightness_limit=(-0.1, 0.1), + contrast_limit=(-0.1, 0.1), + p=1, + ), + ], + p=1, + ), + A.OneOf( + [ + RandomDownScale(p=1), + A.Sharpen(alpha=(0.2, 0.5), lightness=(0.5, 1.0), p=1), + ], + p=1, + ), + ], + p=1.0, + ) + else: + return A.ReplayCompose( + [ + A.Compose( + [ + A.RGBShift((-20, 20), (-20, 20), (-20, 20), p=0.3), + A.HueSaturationValue( + hue_shift_limit=(-0.3, 0.3), + sat_shift_limit=(-0.3, 0.3), + val_shift_limit=(-0.3, 0.3), + p=1, + ), + A.RandomBrightnessContrast( + brightness_limit=(-0.1, 0.1), + contrast_limit=(-0.1, 0.1), + p=1, + ), + ], + p=1, + ), + A.OneOf( + [ + RandomDownScale(p=1), + A.Sharpen(alpha=(0.2, 0.5), lightness=(0.5, 1.0), p=1), + ], + p=1, + ), + ], + p=1.0, + ) + + +def get_transforms(data_type="image"): + """ + Transforms specially design for SBI synthesis + """ + assert data_type in ["image", "video"] + + if data_type == "image": + return A.Compose( + [ + A.RGBShift((-20, 20), (-20, 20), (-20, 20), p=0.3), + A.HueSaturationValue( + hue_shift_limit=(-0.3, 0.3), + sat_shift_limit=(-0.3, 0.3), + val_shift_limit=(-0.3, 0.3), + p=0.3, + ), + A.RandomBrightnessContrast( + brightness_limit=(-0.3, 0.3), contrast_limit=(-0.3, 0.3), p=0.3 + ), + A.ImageCompression(quality_lower=40, quality_upper=100, p=0.5), + ], + additional_targets={"image_f": "image"}, + p=1.0, + ) + else: + return A.ReplayCompose( + [ + A.RGBShift((-20, 20), (-20, 20), (-20, 20), p=0.3), + A.HueSaturationValue( + hue_shift_limit=(-0.3, 0.3), + sat_shift_limit=(-0.3, 0.3), + val_shift_limit=(-0.3, 0.3), + p=0.3, + ), + A.RandomBrightnessContrast( + brightness_limit=(-0.3, 0.3), contrast_limit=(-0.3, 0.3), p=0.3 + ), + A.ImageCompression(quality_lower=40, quality_upper=100, p=0.5), + ], + additional_targets={"image_f": "image"}, + p=1.0, + ) diff --git a/video/fake-stormer/model_code/datasets/sbi/utils.py b/video/fake-stormer/model_code/datasets/sbi/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e8fb51471d3a980eba15fda107e03a297cbfd1fb --- /dev/null +++ b/video/fake-stormer/model_code/datasets/sbi/utils.py @@ -0,0 +1,354 @@ +# -*- coding: utf-8 -*- +import os +import random +import sys + +if os.getcwd() not in sys.path: + sys.path.insert(0, os.getcwd()) + +import albumentations as alb +import cv2 +import numpy as np +from imgaug import augmenters as iaa +from package_utils.bi_online_generation import blendImages, random_erode_dilate +from package_utils.deepfake_mask import dynamic_blend, random_get_hull +from package_utils.image_utils import load_image +from package_utils.transform import randaffine +from PIL import Image +from skimage import transform as sktransform + +from ..common import ParameterStore +from ..pipelines.geo_transform import get_source_transforms + + +def gen_SBI(img, landmark, **kwargs): + """ + This function is adapted to process SBI generation for both image-level and video-level + """ + index = kwargs.get("index") + assert index is not None + debug = kwargs.get("debug") or False + + param_store_ins = ParameterStore.get_instance() + + use_lms68 = param_store_ins.get_parameters("use_lms68") or False + data_type = param_store_ins.get_parameters("data_type") + assert data_type is not None + + if ( + (data_type == "image" and np.random.rand() < 0.25) + or (data_type == "video" and index == 0 and np.random.rand() < 0.25) + or use_lms68 + ): + landmark = landmark[:68] + + if not param_store_ins.has_key("use_lms68") and data_type == "video": + param_store_ins.add_parameters("use_lms68", True) + + # Getting ConvexHull + mask, hull_type = random_get_hull( + landmark, img, hull_type=param_store_ins.get_parameters("hull_type") + ) + if index == 0 and data_type == "video": + param_store_ins.add_parameters("hull_type", hull_type) + + # For debugging + if index is not None and debug: + # Image.fromarray(img).save(f'samples/debugs/BG_{index}.jpg') + Image.fromarray((mask * 255).astype(np.uint8)).save( + f"samples/debugs/ConvexHull_{index}.jpg" + ) + + source = img.copy() + rand_value = param_store_ins.get_parameters("rand_value") or np.random.rand() + if index == 0 and data_type == "video": + param_store_ins.add_parameters("rand_value", rand_value) + + if rand_value < 0.5: + if data_type == "video": + if index == 0: + transform = get_source_transforms(data_type=data_type) + data = transform(image=source.astype(np.uint8)) + source = data["image"] + replay_params = data["replay"] + param_store_ins.add_parameters("s_replay_params", replay_params) + else: + replay_params = param_store_ins.get_parameters("s_replay_params") + data = alb.ReplayCompose.replay( + replay_params, image=source.astype(np.uint8) + ) + source = data["image"] + else: + source = get_source_transforms()(image=source.astype(np.uint8))["image"] + else: + if data_type == "video": + if index == 0: + transform = get_source_transforms(data_type=data_type) + data = transform(image=img.astype(np.uint8)) + img = data["image"] + replay_params = data["replay"] + param_store_ins.add_parameters("i_replay_params", replay_params) + else: + replay_params = param_store_ins.get_parameters("i_replay_params") + data = alb.ReplayCompose.replay( + replay_params, image=img.astype(np.uint8) + ) + img = data["image"] + else: + img = get_source_transforms()(image=img.astype(np.uint8))["image"] + + # if index is not None and debug: + # Image.fromarray(source.astype(np.uint8)).save(f'samples/debugs/FG_{index}.jpg') + + if data_type == "image": + source, mask, _ = randaffine(source, mask[:, :, 0]) + else: + # if use_lms68: + # seq_distortion = kwargs.get('distortion') + # img_h, img_w, img_c = mask.shape + # aug_size = param_store_ins.get_parameters('aug_size') or random.randint(int(img_h*0.8), int(img_h/0.8)) + # mask = sktransform.resize(mask,(aug_size,aug_size),preserve_range=True) # resize mask before deformation + # mask = seq_distortion.augment_image(mask) + # mask, ksize, rand_erode = random_erode_dilate(mask, + # ksize=param_store_ins.get_parameters('ksize'), + # rand_erode=param_store_ins.get_parameters('rand_erode')) # mask of shape (H,W,3) + # mask = sktransform.resize(mask,(img_h,img_w),preserve_range=True) # getting back mask + # mask = mask[:,:,0] + + # # filte empty mask after deformation + # if np.sum(mask) == 0 : + # raise ValueError('Deformed mask has no facial region for blending!!!') + + # if not param_store_ins.has_key('ksize'): + # param_store_ins.add_parameters('ksize', ksize) + # if not param_store_ins.has_key('rand_erode'): + # param_store_ins.add_parameters('rand_erode', rand_erode) + # if not param_store_ins.has_key('aug_size'): + # param_store_ins.add_parameters('aug_size', aug_size) + # else: + if index == 0: + source, mask, fg_replay_params = randaffine( + source, mask[:, :, 0], index=index, data_type=data_type + ) + param_store_ins.add_parameters( + "f_replay_params", fg_replay_params["f_replay_params"] + ) + param_store_ins.add_parameters( + "g_replay_params", fg_replay_params["g_replay_params"] + ) + else: + f_replay_params = param_store_ins.get_parameters("f_replay_params") + g_replay_params = param_store_ins.get_parameters("g_replay_params") + source, mask, _ = randaffine( + source, + mask[:, :, 0], + index=index, + data_type=data_type, + f_replay_params=f_replay_params, + g_replay_params=g_replay_params, + ) # mask of shape (H, W) + + # Getting Deformed ConvexHull + if index is not None and debug: + Image.fromarray((mask * 255).astype(np.uint8)).save( + f"samples/debugs/Deformed_ConvexHull_{index}.jpg" + ) + + if data_type == "image": + img_blended, mask, _ = dynamic_blend(source, img, mask) + else: + # use_BI = param_store_ins.get_parameters('use_BI') or np.random.rand() > 0.5 + # if not param_store_ins.has_key('use_BI'): + # param_store_ins.add_parameters('use_BI', use_BI) + # if use_lms68: + # if index == 0: + # img_blended, mask, blending_params = blendImages(source, + # img, + # mask*255) + # param_store_ins.add_parameters('blending_params', blending_params) + # else: + # img_blended, mask, _ = blendImages(source, + # img, + # mask*255, + # **param_store_ins.get_parameters('blending_params')) + # mask = mask[:,:,0:1] + # else: + if index == 0: + img_blended, mask, blending_params = dynamic_blend(source, img, mask) + param_store_ins.add_parameters("blending_params", blending_params) + else: + blending_params = param_store_ins.get_parameters("blending_params") + img_blended, mask, _ = dynamic_blend(source, img, mask, **blending_params) + img_blended = img_blended.astype(np.uint8) + img = img.astype(np.uint8) + + return img, img_blended, mask + + +def gen_target(background_face, background_landmark, margin=[20, 20], **kwargs): + index = kwargs.get("index") + assert index is not None + + if isinstance(background_face, str): + background_face = load_image(background_face) + + background_face, face_img, mask_f = gen_SBI( + background_face, background_landmark, **kwargs + ) + mask_f = (1 - mask_f) * mask_f * 4 + mask_r = np.zeros((mask_f.shape[0], mask_f.shape[1], 1)) + + margin_x, margin_y = margin + H, W = len(face_img), len(face_img[0]) + face_img = face_img[margin_y : (H - margin_y), margin_x : (W - margin_x), :] + background_face = background_face[ + margin_y : (H - margin_y), margin_x : (W - margin_x), : + ] + + mask_f = mask_f[margin_y : (H - margin_y), margin_x : (W - margin_x), :] + mask_r = mask_r[margin_y : (H - margin_y), margin_x : (W - margin_x), :] + + mask_f, mask_r = np.repeat(mask_f, 3, 2), np.repeat(mask_r, 3, 2) + mask_f, mask_r = (mask_f * 255).astype(np.uint8), (mask_r * 255).astype(np.uint8) + + # lower_bound = [0.5,0.75,1,1] + # fake_intensity = np.random.uniform(lower_bound[np.random.randint(len(lower_bound))], 1.) + fake_intensity = np.random.uniform(0.5, 1.0) + return face_img, mask_f, background_face, mask_r, fake_intensity + + +def reorder_landmark(landmark): + landmark_add = np.zeros((13, 2)) + for idx, idx_l in enumerate([77, 75, 76, 68, 69, 70, 71, 80, 72, 73, 79, 74, 78]): + landmark_add[idx] = landmark[idx_l] + landmark[68:] = landmark_add + return landmark + + +def sbi_hflip(img, mask=None, landmark=None, bbox=None): + H, W = img.shape[:2] + if landmark is not None: + landmark = landmark.copy() + + if bbox is not None: + bbox = bbox.copy() + + if landmark is not None: + landmark_new = np.zeros_like(landmark) + + landmark_new[:17] = landmark[:17][::-1] + landmark_new[17:27] = landmark[17:27][::-1] + + landmark_new[27:31] = landmark[27:31] + landmark_new[31:36] = landmark[31:36][::-1] + + landmark_new[36:40] = landmark[42:46][::-1] + landmark_new[40:42] = landmark[46:48][::-1] + + landmark_new[42:46] = landmark[36:40][::-1] + landmark_new[46:48] = landmark[40:42][::-1] + + landmark_new[48:55] = landmark[48:55][::-1] + landmark_new[55:60] = landmark[55:60][::-1] + + landmark_new[60:65] = landmark[60:65][::-1] + landmark_new[65:68] = landmark[65:68][::-1] + if len(landmark) == 68: + pass + elif len(landmark) == 81: + landmark_new[68:81] = landmark[68:81][::-1] + else: + raise NotImplementedError + landmark_new[:, 0] = W - landmark_new[:, 0] + else: + landmark_new = None + + if bbox is not None: + bbox_new = np.zeros_like(bbox) + bbox_new[0, 0] = bbox[1, 0] + bbox_new[1, 0] = bbox[0, 0] + bbox_new[:, 0] = W - bbox_new[:, 0] + bbox_new[:, 1] = bbox[:, 1].copy() + if len(bbox) > 2: + bbox_new[2, 0] = W - bbox[3, 0] + bbox_new[2, 1] = bbox[3, 1] + bbox_new[3, 0] = W - bbox[2, 0] + bbox_new[3, 1] = bbox[2, 1] + bbox_new[4, 0] = W - bbox[4, 0] + bbox_new[4, 1] = bbox[4, 1] + bbox_new[5, 0] = W - bbox[6, 0] + bbox_new[5, 1] = bbox[6, 1] + bbox_new[6, 0] = W - bbox[5, 0] + bbox_new[6, 1] = bbox[5, 1] + else: + bbox_new = None + + if mask is not None: + mask = mask[:, ::-1] + else: + mask = None + img = img[:, ::-1].copy() + return img, mask, landmark_new, bbox_new + + +def BI_postprocessing(img, face_img, mask): + param_store_ins = ParameterStore.get_instance() + face_img = Image.fromarray(face_img) + img = Image.fromarray(img) + + # randomly downsample after BI pipeline + rand_val_post = param_store_ins.get_parameters("rand_val_post") or random.randint( + 0, 1 + ) + if rand_val_post: + aug_size = param_store_ins.get_parameters("post_aug_size") or random.randint( + 64, 317 + ) + rand_resize = param_store_ins.get_parameters("rand_resize") or random.randint( + 0, 1 + ) + + if rand_resize: + face_img = face_img.resize((aug_size, aug_size), Image.BILINEAR) + img = img.resize((aug_size, aug_size), Image.BILINEAR) + else: + face_img = face_img.resize((aug_size, aug_size), Image.NEAREST) + img = img.resize((aug_size, aug_size), Image.NEAREST) + + if not param_store_ins.has_key("post_aug_size"): + param_store_ins.add_parameters("post_aug_size", aug_size) + + if not param_store_ins.has_key("rand_resize"): + param_store_ins.add_parameters("rand_resize", rand_resize) + + if not param_store_ins.has_key("rand_val_post"): + param_store_ins.add_parameters("rand_val_post", rand_val_post) + + face_img = face_img.resize((317, 317), Image.BILINEAR) + img = img.resize((317, 317), Image.BILINEAR) + face_img = np.array(face_img) + img = np.array(img) + + soft_margin = param_store_ins.get_parameters("soft_margin") or np.random.randint( + -30, 30 + ) + if not param_store_ins.has_key("soft_margin"): + param_store_ins.add_parameters("soft_margin", soft_margin) + + face_img = face_img[ + 30 + soft_margin : (287 + soft_margin), + 30 + soft_margin : (287 + soft_margin), + :, + ] + img = img[ + 30 + soft_margin : (287 + soft_margin), + 30 + soft_margin : (287 + soft_margin), + :, + ] + mask = mask[ + 30 + soft_margin : (287 + soft_margin), + 30 + soft_margin : (287 + soft_margin), + :, + ] + + return img, face_img, mask diff --git a/video/fake-stormer/model_code/datasets/utils.py b/video/fake-stormer/model_code/datasets/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a6b3c6fdb6ed9c080a893dc90147287108e5b1ae --- /dev/null +++ b/video/fake-stormer/model_code/datasets/utils.py @@ -0,0 +1,128 @@ +# -*- coding: utf-8 -*- +import os +import random + +import numpy as np + + +def _extract_data_based_dist( + data_type: str, image_paths: list, labels: list, dist: list, **params +): + def sampling_frames(image_paths: list, labels: list, dist: list, **params): + """ + Extracting image paths and labels based on given distribution + """ + total_f = sum(labels) + total_r = len(labels) - total_f + + assert ( + total_f > 0 + ), "Number of fake images must be greater than 0 for distribution sampling!" + assert ( + total_r > 0 + ), "Number of real images must be greater than 0 for distribution sampling!" + + r_dist, f_dist = dist[0], dist[1] + + idxes = sorted(range(0, len(labels)), key=lambda k: labels[k]) + r_idxes = idxes[:total_r] + f_idxes = idxes[total_r:] + + print(f"Original Number of Fake images --- {len(f_idxes)}") + print(f"Original Number of Real images --- {len(r_idxes)}") + + if int((total_f / f_dist) * r_dist) > total_r: + total_f = int((total_r / r_dist) * f_dist) + f_idxes = random.sample(f_idxes, total_f) + else: + total_r = int((total_f / f_dist) * r_dist) + r_idxes = random.sample(r_idxes, total_r) + + print( + f"Number of Fake images --- {len(f_idxes)} given Fake distribution --- {f_dist}" + ) + print( + f"Number of Real images --- {len(r_idxes)} given Real distribution --- {r_dist}" + ) + + new_idxes = r_idxes + f_idxes + + image_paths = [image_paths[i] for i in new_idxes] + labels = np.array(labels)[new_idxes] + + for k, v in params.items(): + if v is not None and len(v): + params[k] = [v[i] for i in new_idxes] + + return image_paths, labels, params + + def sampling_videos(image_paths: list, labels: list, dist: list, **params): + """ + Extracting image paths and labels based on given distribution for video data + """ + f_vid_ids = [] + r_vid_ids = [] + for ip in image_paths: + faketype = ip.split("/")[8] + vid_id = os.path.dirname(ip) + if ( + faketype == "real_videos" + or "real" in faketype + or "original" in faketype + ): + r_vid_ids.append(vid_id) + else: + f_vid_ids.append(vid_id) + + f_vid_ids = list(set(f_vid_ids)) + r_vid_ids = list(set(r_vid_ids)) + total_f = len(f_vid_ids) + total_r = len(r_vid_ids) + + assert ( + total_f > 0 + ), "Number of fake videos must be greater than 0 for distribution sampling!" + assert ( + total_r > 0 + ), "Number of real videos must be greater than 0 for distribution sampling!" + print(f"Original Number of Fake videos --- {total_f}") + print(f"Original Number of Real videos --- {total_r}") + + r_dist, f_dist = dist[0], dist[1] + + if int((total_f / f_dist) * r_dist) > total_r: + total_f = int((total_r / r_dist) * f_dist) + f_vid_ids = random.sample(f_vid_ids, total_f) + else: + total_r = int((total_f / f_dist) * r_dist) + r_vid_ids = random.sample(r_vid_ids, total_r) + + print( + f"Number of Fake videos --- {len(f_vid_ids)} given Fake distribution --- {f_dist}" + ) + print( + f"Number of Real videos --- {len(r_vid_ids)} given Real distribution --- {r_dist}" + ) + + vid_ids = r_vid_ids + f_vid_ids + new_idxes = [] + + for i in range(len(labels)): + ip = image_paths[i] + vid_id = "/".join([ip.split("/")[-3], ip.split("/")[-2]]) + if vid_id in vid_ids: + new_idxes.append(i) + + image_paths = [image_paths[i] for i in new_idxes] + labels = np.array(labels)[new_idxes] + + for k, v in params.items(): + if v is not None and len(v): + params[k] = [v[i] for i in new_idxes] + + return image_paths, labels, params + + if data_type == "image": + return sampling_frames(image_paths, labels, dist, **params) + else: + return sampling_videos(image_paths, labels, dist, **params) diff --git a/video/fake-stormer/model_code/demo/method.png b/video/fake-stormer/model_code/demo/method.png new file mode 100644 index 0000000000000000000000000000000000000000..ca956c5a695416cc55db07a44f7f0fb5caa7f48f --- /dev/null +++ b/video/fake-stormer/model_code/demo/method.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d7c6acdff1e9e7673b97b2ef5000e63b23c4e7be75cf9a8a3ba3989d90ab1131 +size 3612142 diff --git a/video/fake-stormer/model_code/dockerfiles/Dockerfile b/video/fake-stormer/model_code/dockerfiles/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..19c1c1ca1f7b98764f1a5e242e4b8056f612f3f8 --- /dev/null +++ b/video/fake-stormer/model_code/dockerfiles/Dockerfile @@ -0,0 +1,25 @@ +FROM pytorch/pytorch:1.8.0-cuda11.1-cudnn8-devel + +MAINTAINER DatNGUYEN + +ENV TZ=Europe/Luxembourg +RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone + +RUN apt-key del 7fa2af80 && \ + apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub && \ + apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu2004/x86_64/7fa2af80.pub + +RUN apt-get -y update && \ + apt-get upgrade -y && \ + apt-get install -y libprotobuf-dev protobuf-compiler && \ + apt-get install -y cmake && \ + apt-get install -y git && \ + apt-get install -y libgl1-mesa-dev && \ + apt-get -y update && apt-get install -y libopencv-dev + +RUN pip install albumentations==1.1.0 dlib==19.24.0 python-box==7.1.1 imgaug==0.4.0 && \ + pip install imutils==0.5.4 ipython numpy==1.23.3 opencv-python==4.5.1.48 && \ + pip install pandas==1.3.5 Pillow==9.3.0 scikit-image==0.19.3 scipy==1.9.3 simplejson && \ + pip install tensorboardX==2.5.1 natsort==8.4.0 tqdm PyYAML + +WORKDIR /workspace diff --git a/video/fake-stormer/model_code/dockerfiles/README.md b/video/fake-stormer/model_code/dockerfiles/README.md new file mode 100644 index 0000000000000000000000000000000000000000..fc135f5a640c1159aa76356a309c6079c65e4f83 --- /dev/null +++ b/video/fake-stormer/model_code/dockerfiles/README.md @@ -0,0 +1,59 @@ +## Docker Build (Optional) +*We further provide an optional Docker file which can be used to build working env with Docker.* + +1. Install docker to the system (skip the step if docker has been already installed): + ```shell + sudo apt install docker + ``` +2. To start your docker environment, please go to the folder **dockerfiles**: + ```shell + cd dockerfiles + ``` +3. Create a docker image (you can put any name you want): + ```shell + docker build --tag 'fakestormer' . + ``` +4. Check the status of the image created: + 1. Run command: + ``` shell + docker image ls + ``` + 2. You should see something similiar: + + |REPOSITORY| IMAGE ID|CREATED| SIZE |TAG| + |----------|---------|-------|------|---| + |fakestormer| efd422370750|12 minutes ago| 18.4GB |latest | +5. Run a container from the created image: + 1. Run command + ```shell + docker run -v : --gpus 'all,capabilities=utility' -it fakestormer /bin/bash + ``` + 2. Check the container created: + 1. Run command: + ```shell + docker ps + ``` + 2. Check the result: + CONTAINER ID|IMAGE|COMMAND|CREATED|STATUS|PORTS|NAMES| + |-----------|-----|-------|-------|------|-----|-----| + |0203c192febb|fakestormer|"/bin/bash"| 29 seconds ago |Up 28 seconds | |determined_cannon| + 3. To access the docker container: + ```shell + docker exec -it 0203c192febb /bin/bash + ``` + 4. To start the container: + ```shell + docker start 0203c192febb + ``` + 5. To stop the container: + ```shell + docker stop 0203c192febb + ``` + +6. Inside the docker container, you can clone or mount the repository from outside: + ```shell + cd /workspace/ + git clone https://github.com/10Ring/FakeSTormer.git + cd FakeSTormer/ + ``` +7. Now you are ready for [*QuickStart*](#quickstart) diff --git a/video/fake-stormer/model_code/lib/core_function.py b/video/fake-stormer/model_code/lib/core_function.py new file mode 100644 index 0000000000000000000000000000000000000000..3295477e3e0decdc92e6b63a9ce77752298886a3 --- /dev/null +++ b/video/fake-stormer/model_code/lib/core_function.py @@ -0,0 +1,552 @@ +# -*- coding: utf-8 -*- +import os +import time +from typing import Union + +import torch +from lib.metrics import bin_calculate_auc_ap_ar, get_acc_mesure_func +from logs.logger import board_writing +from numpy import arange +from package_utils.utils import debugging_panel +from tqdm import tqdm + + +class AverageMeter(object): + """Computes and stores the average and current value""" + + def __init__(self): + self.reset() + + 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 if self.count != 0 else 0 + + +def get_batch_data(batch_data: Union[dict]): + """ + Parsing data for model consumption (train/val/test) + """ + inputs = batch_data["img"] if isinstance(batch_data, dict) else batch_data[0] + labels = batch_data["label"] if isinstance(batch_data, dict) else batch_data[1] + heatmaps = None + cstency_heatmaps = None + offsets = None + targets = None + temp_locs = None + maskout_pes = None + + if "heatmap" in batch_data: + heatmaps = batch_data["heatmap"] + + if "target" in batch_data: + targets = batch_data["target"] + + if "cstency" in batch_data: + cstency_heatmaps = batch_data["cstency"] + + if "offset" in batch_data: + offsets = batch_data["offset"] + + if "temp_loc" in batch_data: + temp_locs = batch_data["temp_loc"] + + if "mask_out_p" in batch_data: + maskout_pes = batch_data["mask_out_p"] + + return ( + inputs, + labels, + targets, + heatmaps, + cstency_heatmaps, + offsets, + temp_locs, + maskout_pes, + ) + + +def train( + cfg, + model, + critetion, + optimizer, + epoch, + data_loader, + logger, + writer, + devices, + trainIters, + metrics_base="combine", + scaler=None, +): + calculate_acc = get_acc_mesure_func(metrics_base) + batch_time = AverageMeter() + data_time = AverageMeter() + losses = AverageMeter() + acc = AverageMeter() + + # Switch to train mode + model.train() + data_loader = tqdm(data_loader, dynamic_ncols=True) + accumulation_steps = cfg.TRAIN.accumulation_steps if cfg.TRAIN.use_amp else 1 + start = time.time() + optimizer.zero_grad() + for i, batch_data in enumerate(data_loader): + ( + inputs, + labels, + targets, + heatmaps, + cstency_heatmaps, + offsets, + temp_locs, + maskout_pes, + ) = get_batch_data(batch_data) + inputs = inputs.cuda().to(non_blocking=True, dtype=torch.float) + labels = labels.cuda().to(non_blocking=True, dtype=torch.float) + maskout_pes = ( + maskout_pes.cuda().to(non_blocking=True) + if maskout_pes is not None + else None + ) + # additional_targets = {"maskout_pes": maskout_pes} + + # Measuring data loading time + data_time.update(time.time() - start) + loop = arange(1) if cfg.TRAIN.optimizer != "SAM" else arange(2) + + for idx in loop: + with torch.cuda.amp.autocast(enabled=cfg.TRAIN.use_amp): + # outputs = model(inputs, **additional_targets) + outputs = model(inputs) + if isinstance(outputs, list): + outputs = outputs[0] + + # In case outputs contain a dict key + if isinstance(outputs, dict): + outputs_cls = outputs["cls"] + outputs_hm = outputs["hm"] if "hm" in outputs.keys() else None + outputs_offset = ( + outputs["offset"] if "offset" in outputs.keys() else None + ) + outputs_cstency = ( + outputs["cstency"] if "cstency" in outputs.keys() else None + ) + outputs_temp_loc = ( + outputs["temp_loc"] if "temp_loc" in outputs.keys() else None + ) + + if idx == 0: + first_outputs_hm = outputs_hm + first_outputs_cls = outputs_cls + + if "Combined" in cfg.TRAIN.loss.type: + # labels = labels.cuda().to(non_blocking=True).long() + + if offsets is not None: + offsets = offsets.cuda().to(non_blocking=True) + + if cstency_heatmaps is not None: + cstency_heatmaps = cstency_heatmaps.cuda().to(non_blocking=True) + + if temp_locs is not None: + temp_locs = temp_locs.cuda().to(non_blocking=True) + + if cfg.TRAIN.loss.type != "CombinedHeatmapBinaryLoss": + heatmaps = heatmaps.cuda().to(non_blocking=True) + else: + heatmaps = targets.cuda().to(non_blocking=True) + + loss_ = critetion( + outputs_hm, + heatmaps, + outputs_cls, + labels, + offset_preds=outputs_offset, + offset_gts=offsets, + cstency_preds=outputs_cstency, + cstency_gts=cstency_heatmaps, + temp_loc_preds=outputs_temp_loc, + temp_loc_gts=temp_locs, + hm_mask=maskout_pes, + ) + loss = loss_["hm"] + if "cls" in loss_.keys(): + loss += loss_["cls"] + if "dst_hm_cls" in loss_.keys(): + loss += loss_["dst_hm_cls"] + if "offset" in loss_.keys(): + loss += loss_["offset"] + if "cstency" in loss_.keys(): + loss += loss_["cstency"] + if "temp_loc" in loss_.keys(): + loss += loss_["temp_loc"] + else: + loss = critetion(outputs_cls, labels) + + loss /= accumulation_steps + + # gradients accumulation for larger batch + if cfg.TRAIN.use_amp: + scaler( + cfg, + loss, + optimizer, + parameters=model.parameters(), + step=idx, + update_grad=(i + 1) % accumulation_steps == 0, + ) + if (i + 1) % accumulation_steps == 0: + optimizer.zero_grad() + else: + loss.backward() + + if cfg.TRAIN.optimizer != "SAM": + optimizer.step() + else: + if idx == 0: + optimizer.first_step(zero_grad=True) + else: + optimizer.second_step(zero_grad=True) + optimizer.zero_grad() + + if cfg.TRAIN.use_amp: + torch.cuda.synchronize() + + if cfg.TRAIN.debug.active: + debugging_panel( + cfg.TRAIN.debug, + inputs, + heatmaps, + first_outputs_hm, + i, + batch_cls_pred=first_outputs_cls, + ) + + if metrics_base == "binary": + acc_ = calculate_acc(first_outputs_cls, targets=targets, labels=labels) + elif metrics_base == "heatmap": + acc_ = calculate_acc(first_outputs_hm, targets=targets, labels=labels) + else: + acc_ = calculate_acc( + first_outputs_hm, + first_outputs_cls, + targets=targets, + labels=labels, + cls_lamda=critetion.cls_lmda, + ) + + if isinstance(inputs, list): + batch_size = inputs[0].size(0) + else: + batch_size = inputs.size(0) + + # Measure accuracy and record loss + losses.update(loss.item() * accumulation_steps, n=batch_size) + acc.update(acc_, n=batch_size) + + batch_time.update(time.time() - start) + start = time.time() + + # Logging + if i % 5 == 0: + params = {} + if "Combined" in cfg.TRAIN.loss.type: + if ( + hasattr(critetion, "dst_hm_cls_lmda") + and critetion.dst_hm_cls_lmda > 0 + ): + params["loss_dst"] = loss_["dst_hm_cls"].item() + if hasattr(critetion, "offset_lmda") and critetion.offset_lmda > 0: + params["loss_offset"] = loss_["offset"].item() + if "cstency" in loss_.keys(): + params["loss_cstency"] = loss_["cstency"].item() + if "temp_loc" in loss_.keys(): + params["loss_temp_loc"] = loss_["temp_loc"].item() + logger.epochInfor( + epoch, + i, + len(data_loader), + batch_time=batch_time, + data_time=data_time, + losses=losses, + acc=acc, + speed=batch_size / batch_time.val, + loss_cls=loss_["cls"].item(), + **params, + ) + else: + logger.epochInfor( + epoch, + i, + len(data_loader), + batch_time=batch_time, + data_time=data_time, + losses=losses, + acc=acc, + speed=batch_size / batch_time.val, + ) + + trainIters += 1 + if cfg.TRAIN.tensorboard: + board_writing(writer, losses.avg, acc.avg, trainIters, "Train") + return losses, acc, trainIters + + +def validate( + cfg, + model, + critetion, + epoch, + data_loader, + logger, + writer, + devices, + valIters, + metrics_base="combine", +): + calculate_acc = get_acc_mesure_func(metrics_base) + batch_time = AverageMeter() + data_time = AverageMeter() + losses = AverageMeter() + acc = AverageMeter() + + # Switch to test mode + model.eval() + data_loader = tqdm(data_loader, dynamic_ncols=True) + start = time.time() + with torch.no_grad(): + for i, batch_data in enumerate(data_loader): + ( + inputs, + labels, + targets, + heatmaps, + cstency_heatmaps, + offsets, + temp_locs, + maskout_pes, + ) = get_batch_data(batch_data) + inputs = inputs.to(devices, non_blocking=True, dtype=torch.float).cuda() + labels = labels.cuda().to(non_blocking=True, dtype=torch.float) + maskout_pes = ( + maskout_pes.cuda().to(non_blocking=True) + if maskout_pes is not None + else None + ) + # additional_targets = {"maskout_pes": maskout_pes} + + # Measuring data loading time + data_time.update(time.time() - start) + + # outputs = model(inputs, **additional_targets) + outputs = model(inputs) + if isinstance(outputs, list): + outputs = outputs[0] + + # In case outputs contain a dict key + if isinstance(outputs, dict): + outputs_cls = outputs["cls"] + outputs_hm = outputs["hm"] if "hm" in outputs.keys() else None + outputs_offset = ( + outputs["offset"] if "offset" in outputs.keys() else None + ) + outputs_cstency = ( + outputs["cstency"] if "cstency" in outputs.keys() else None + ) + outputs_temp_loc = ( + outputs["temp_loc"] if "temp_loc" in outputs.keys() else None + ) + + if "Combined" in cfg.TRAIN.loss.type: + # labels = labels.cuda().to(non_blocking=True).long() + + if offsets is not None: + offsets = offsets.cuda().to(non_blocking=True) + + if cstency_heatmaps is not None: + cstency_heatmaps = cstency_heatmaps.cuda().to(non_blocking=True) + + if temp_locs is not None: + temp_locs = temp_locs.cuda().to(non_blocking=True) + + if cfg.TRAIN.loss.type != "CombinedHeatmapBinaryLoss": + heatmaps = heatmaps.cuda().to(non_blocking=True) + else: + heatmaps = targets.cuda().to(non_blocking=True) + + loss_ = critetion( + outputs_hm, + heatmaps, + outputs_cls, + labels, + offset_preds=outputs_offset, + offset_gts=offsets, + cstency_preds=outputs_cstency, + cstency_gts=cstency_heatmaps, + temp_loc_preds=outputs_temp_loc, + temp_loc_gts=temp_locs, + hm_mask=maskout_pes, + ) + loss = loss_["hm"] + if "cls" in loss_.keys(): + loss += loss_["cls"] + if "dst_hm_cls" in loss_.keys(): + loss += loss_["dst_hm_cls"] + if "offset" in loss_.keys(): + loss += loss_["offset"] + if "cstency" in loss_.keys(): + loss += loss_["cstency"] + if "temp_loc" in loss_.keys(): + loss += loss_["temp_loc"] + else: + loss = critetion(outputs_cls, labels) + + if cfg.TRAIN.debug.active: + debugging_panel( + cfg.TRAIN.debug, + inputs, + heatmaps, + outputs_hm, + i, + batch_cls_pred=outputs_cls, + split="val", + ) + + if metrics_base == "binary": + acc_ = calculate_acc(outputs_cls, targets=targets, labels=labels) + elif metrics_base == "heatmap": + acc_ = calculate_acc(outputs_hm, targets=targets, labels=labels) + else: + acc_ = calculate_acc( + outputs_hm, + outputs_cls, + targets=targets, + labels=labels, + cls_lamda=critetion.cls_lmda, + ) + + if isinstance(inputs, list): + batch_size = inputs[0].size(0) + else: + batch_size = inputs.size(0) + + # Measure accuracy and record loss + losses.update(loss.item(), n=batch_size) + acc.update(acc_, n=batch_size) + + batch_time.update(time.time() - start) + start = time.time() + + valIters += 1 + if cfg.TRAIN.tensorboard: + board_writing(writer, losses.avg, acc.avg, valIters, "Val") + + # Logging + params = {} + if "Combined" in cfg.TRAIN.loss.type: + if ( + hasattr(critetion, "dst_hm_cls_lmda") + and critetion.dst_hm_cls_lmda > 0 + ): + params["loss_dst"] = loss_["dst_hm_cls"].item() + if hasattr(critetion, "offset_lmda") and critetion.offset_lmda > 0: + params["loss_offset"] = loss_["offset"].item() + if "cstency" in loss_.keys(): + params["loss_cstency"] = loss_["cstency"].item() + if "temp_loc" in loss_.keys(): + params["loss_temp_loc"] = loss_["temp_loc"].item() + logger.epochInfor( + epoch, + i, + len(data_loader), + batch_time=batch_time, + data_time=data_time, + losses=losses, + acc=acc, + speed=batch_size / batch_time.val, + loss_cls=loss_["cls"].item(), + **params, + ) + else: + logger.epochInfor( + epoch, + i, + len(data_loader), + batch_time=batch_time, + data_time=data_time, + losses=losses, + acc=acc, + speed=batch_size / batch_time.val, + ) + return losses, acc, valIters + + +def test( + cfg, + model, + critetion, + epoch, + data_loader, + logger, + writer, + devices, + valIters, + metrics_base="combine", +): + calculate_acc = get_acc_mesure_func(metrics_base) + total_preds = torch.tensor([]).cuda().to(dtype=torch.float) + total_labels = torch.tensor([]).cuda().to(dtype=torch.float) + + # Switch to test mode + model.eval() + test_dataloader = tqdm(data_loader, dynamic_ncols=True) + with torch.no_grad(): + for b, (inputs, labels, vid_ids) in enumerate(test_dataloader): + inputs = inputs.to(dtype=torch.float).cuda() + labels = labels.to(dtype=torch.float).cuda() + + outputs = model(inputs) + # Applying Flip test + if isinstance(outputs, list): + outputs = outputs[0] + + # In case outputs contain a dict key + if isinstance(outputs, dict): + # hm_outputs = outputs['hm'] if 'hm' in outputs.keys() else None + cls_outputs = outputs["cls"] + # outputs_temp_loc = outputs['temp_loc'] if 'temp_loc' in outputs.keys() else None + + total_preds = torch.cat((total_preds, cls_outputs), 0) + total_labels = torch.cat((total_labels, labels), 0) + + acc_ = calculate_acc( + total_preds, targets=None, labels=total_labels, threshold=cfg.TEST.threshold + ) + metrics = bin_calculate_auc_ap_ar( + total_preds, + total_labels, + metrics_base=metrics_base, + threshold=cfg.TEST.threshold, + ) + auc_, ap_, ar_, mf1_ = ( + metrics["auc"], + metrics["ap"], + metrics["ar"], + metrics["mf1"], + ) + + logger.info( + f"Current ACC, AUC, AP, AR, mF1 for {cfg.DATASET.DATA.TEST.FAKETYPE} --- {cfg.DATASET.DATA.TEST.LABEL_FOLDER} -- \ + {acc_*100} -- {auc_*100} -- {ap_*100} -- {ar_*100} -- {mf1_*100}" + ) + + return acc_, auc_, ap_, ar_ diff --git a/video/fake-stormer/model_code/lib/metrics.py b/video/fake-stormer/model_code/lib/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..c25e662345713ca5b0dfa616130a988f9a8c5e88 --- /dev/null +++ b/video/fake-stormer/model_code/lib/metrics.py @@ -0,0 +1,201 @@ +# -*- coding: utf-8 -*- +import os + +import numpy as np +import torch +from losses.losses import _avg_sigmoid, _sigmoid +from scipy.stats import beta, gaussian_kde +from sklearn import metrics as cal_metrics +from sklearn.metrics import ( + average_precision_score, + balanced_accuracy_score, + f1_score, + precision_score, + recall_score, +) + + +def bin_calculate_acc(preds, labels, targets=None, threshold=0.5): + if torch.is_tensor(preds): + if preds.shape[-1] > 1: + preds = preds.softmax(dim=-1) + if preds.shape[-1] == 2: + preds = preds[:, -1] + else: + preds = preds.max(dim=-1, keepdim=True).values + labels = labels.max(dim=-1, keepdim=True).values + else: + preds = preds.sigmoid() + + preds = preds.detach().cpu().numpy() + labels = labels.detach().cpu().numpy() + + preds_ = (preds >= threshold).astype(int) + acc = np.mean((preds_ == labels).astype(int), axis=0) + + if acc.ndim >= 1: + return acc[0] + else: + return acc + + +def hm_calculate_acc(preds, targets=None, labels=None, threshold=0.5): + cls_ = _avg_sigmoid(preds) + acc = bin_calculate_acc(cls_, labels, threshold=threshold) + return acc + + +def hm_bin_calculate_acc( + hm_preds, cls_preds, targets=None, labels=None, cls_lamda=0.05 +): + # Select top hm_preds + hm_preds_ = _sigmoid(hm_preds.clone()) + hm_preds_ = torch.reshape(hm_preds_, (hm_preds_.shape[0], hm_preds_.shape[1], -1)) + top_k = torch.topk(hm_preds_, 10, -1).values + mean_hm_preds = torch.mean(top_k, -1) + + cls_preds_ = cls_lamda * cls_preds + (1 - cls_lamda) * mean_hm_preds + acc = bin_calculate_acc(cls_preds_, labels) + return acc + + +def bin_calculate_auc_ap_ar( + cls_preds, + labels, + metrics_base="binary", + hm_preds=None, + cls_lamda=0.1, + threshold=0.5, + apr=True, +): + assert metrics_base in [ + "binary", + "heatmap", + "combine", + ], "Metric base is only one of these values [binary, heatmap, combine]" + + if torch.is_tensor(cls_preds): + if cls_preds.shape[-1] > 1: + cls_preds = cls_preds.softmax(dim=-1) + if cls_preds.shape[-1] == 2: + cls_preds = cls_preds[:, -1] + else: + cls_preds = cls_preds.max(dim=-1, keepdim=True).values + labels = labels.max(dim=-1, keepdim=True).values + else: + cls_preds = cls_preds.sigmoid() + + if metrics_base == "combine": + assert ( + hm_preds is not None + ), "Heatmap predict can not be None if metrics-base is combine" + hm_preds = _sigmoid(hm_preds) + hm_preds = torch.reshape(hm_preds, (hm_preds.shape[0], 1, -1)) + top_k = torch.topk(hm_preds, 10, -1).values + mean_hm_preds = torch.mean(top_k, -1) + cls_preds = cls_lamda * cls_preds + (1 - cls_lamda) * mean_hm_preds + + labels = labels.cpu().numpy() + cls_preds = cls_preds.cpu().numpy() + fpr, tpr, thresholds = cal_metrics.roc_curve(labels, cls_preds, pos_label=1) + + metrics = {} + distances = np.sqrt((1 - tpr) ** 2 + fpr**2) + optimal_idx = np.argmin(distances) + optimal_threshold = thresholds[optimal_idx] + top_k_min_indices = np.argsort(distances)[:5] + top_k_thresholds = thresholds[top_k_min_indices] + metrics["best_thr"] = optimal_threshold + metrics["thr_var"] = np.var(top_k_thresholds, ddof=1) + + # AUC + metrics["auc"] = cal_metrics.auc(fpr, tpr) + + if apr: + # AP metric + ap = average_precision_score(labels, cls_preds) + metrics["ap"] = ap + + # AR metric + ar = recall_score(labels, (cls_preds >= threshold).astype(int), average="macro") + metrics["ar"] = ar + + # mF1 metric + metrics["mf1"] = (ap * ar * 2) / (ap + ar) + + return metrics + else: + # False Negative Rate + fnr = 1 - tpr + eer_threshold = thresholds[np.nanargmin(np.absolute(fpr - fnr))] + eer = fpr[np.nanargmin(np.absolute(fpr - fnr))] + metrics["eer"] = eer + + metrics["bacc"] = balanced_accuracy_score( + labels, (cls_preds >= threshold).astype(int) + ) + + metrics["p"] = precision_score(labels, (cls_preds >= threshold).astype(int)) + + metrics["r"] = recall_score(labels, (cls_preds >= threshold).astype(int)) + + metrics["s"] = recall_score( + labels, (cls_preds >= threshold).astype(int), pos_label=0 + ) + + metrics["f1"] = f1_score(labels, (cls_preds >= threshold).astype(int)) + + return metrics + + +def get_acc_mesure_func(task="binary"): + if task == "binary": + return bin_calculate_acc + elif task == "heatmap": + return hm_calculate_acc + else: + return hm_bin_calculate_acc + + +# Compute Empirical CDF +def empirical_cdf(data): + sorted_data = np.sort(data) + return sorted_data, np.arange(1, len(sorted_data) + 1) / len(sorted_data) + + +# Transformations +def apply_cdf_transform(neg_data, pos_data, cdf_type="empirical"): + if cdf_type == "empirical": + return ( + empirical_cdf(neg_data)[1], + empirical_cdf(pos_data)[1], + ) # Get the ECDF values + elif cdf_type == "kde": + kde_data1 = gaussian_kde(neg_data) + kde_data2 = gaussian_kde(pos_data) + + x_values = np.linspace(0, 1, 11) # 0.0; 0.1; 0.2; ... + + cdf1 = np.cumsum(kde_data1(x_values)) + cdf1 /= cdf1[-1] + + cdf2 = np.cumsum(kde_data2(x_values)) + cdf2 /= cdf2[-1] + return cdf1, cdf2 + elif cdf_type == "para": + alp1, bta1, _, _ = beta.fit(neg_data, floc=0, fscale=1.0001) + alp2, bta2, _, _ = beta.fit(pos_data, floc=0, fscale=1.0001) + + x_values = np.linspace(0, 1, 11) # 0.0; 0.1; 0.2; ... + + cdf1 = beta.cdf(x_values, alp1, bta1) + cdf2 = beta.cdf(x_values, alp2, bta2) + return cdf1, cdf2 + elif cdf_type == "quantile": + q = np.linspace(0, 1, min(len(neg_data), len(pos_data))) # Matching quantiles + F_inv_P0 = np.quantile(neg_data, q) + F_inv_P1 = np.quantile(pos_data, q) + + return F_inv_P0, F_inv_P1 + else: + raise ValueError("Unknown CDF type") diff --git a/video/fake-stormer/model_code/lib/optimizers/sam.py b/video/fake-stormer/model_code/lib/optimizers/sam.py new file mode 100644 index 0000000000000000000000000000000000000000..8247fbf4e0d1ed50169a897f238ac67a594e7bde --- /dev/null +++ b/video/fake-stormer/model_code/lib/optimizers/sam.py @@ -0,0 +1,96 @@ +# -*- coding: utf-8 -*- +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, scaler=None): + 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)" + + if scaler is None: + self.base_optimizer.step() # do the actual "sharpness-aware" update + else: + scaler.step(self.base_optimizer) + + if zero_grad: + self.zero_grad() + + if scaler is not None: + return scaler + + @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 diff --git a/video/fake-stormer/model_code/lib/scheduler/linear_decay.py b/video/fake-stormer/model_code/lib/scheduler/linear_decay.py new file mode 100644 index 0000000000000000000000000000000000000000..a218e60ee69a54f454f324fbef46a77071f07243 --- /dev/null +++ b/video/fake-stormer/model_code/lib/scheduler/linear_decay.py @@ -0,0 +1,38 @@ +# -*- coding: utf-8 -*- +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, booster=2): + self.start_decay = start_decay + self.n_epoch = n_epoch + self.booster = booster + 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[-1] + + if last_epoch > 0: + try: + cur_lr = self.get_last_lr() + except: + cur_lr = b_lr * self.booster + start_decay = self.start_decay + + if last_epoch >= start_decay: + lr = b_lr * self.booster - (b_lr * self.booster) / ( + n_epoch - start_decay + ) * (last_epoch - start_decay) + else: + if last_epoch < start_decay: + lr = b_lr + (b_lr * self.booster - b_lr) / start_decay * last_epoch + else: + lr = cur_lr + + self._last_lr = lr + print(f"Active Learning Rate --- {lr}") + return [lr] diff --git a/video/fake-stormer/model_code/logs/logger.py b/video/fake-stormer/model_code/logs/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..655ae2871372e8788a1c800918bf6571350f6920 --- /dev/null +++ b/video/fake-stormer/model_code/logs/logger.py @@ -0,0 +1,109 @@ +# -*- coding: utf-8 -*- +import logging +import os +from datetime import datetime +from types import MethodType + +import torch +import torch.nn.functional as F +from mmcv.utils import get_logger +from package_utils.utils import make_dir + +LOG_DIR = "logs/{}".format(datetime.today().strftime("%d-%m-%Y")) +make_dir(LOG_DIR) + + +class Logger: + def __init__(self, task="training", workdir=LOG_DIR): + super().__init__() + self.logger = logging.getLogger("") + self.logger.setLevel(logging.INFO) + + file_handler = logging.FileHandler("{}/{}.log".format(workdir, task)) + stream_handler = logging.StreamHandler() + self.logger.addHandler(file_handler) + self.logger.addHandler(stream_handler) + self.logger.epochInfor = MethodType(self.epochInfor, self.logger) + self.info = self.logger.info + + def epochInfor( + self, + epoch, + idx, + length, + batch_time, + speed, + data_time, + losses, + acc, + loss_cls=None, + **kwargs, + ): + msg = ( + "Epoch: [{0}][{1}/{2}]\t" + "Time {batch_time.val:.3f}s ({batch_time.avg:.3f}s)\t" + "Speed {speed:.1f} samples/s\t" + "Data {data_time.val:.3f}s ({data_time.avg:.3f}s)\t" + "Loss {loss.val:.5f} ({loss.avg:.5f})\t" + "Accuracy {acc.val:.3f} ({acc.avg:.3f})".format( + epoch, + idx, + length, + batch_time=batch_time, + speed=speed, + data_time=data_time, + loss=losses, + acc=acc, + ) + ) + if loss_cls is not None: + msg += "\t Cls Loss: {loss_cls:.5f}".format(loss_cls=loss_cls) + + for k, v in kwargs.items(): + if v is not None: + msg += f"\t {k}: {v:.5f}" + + self.logger.info(msg) + + +def board_writing(writer, loss, acc, iterations, dataset="Train"): + writer.add_scalar("{}/loss".format(dataset), loss, iterations) + writer.add_scalar("{}/acc".format(dataset), acc, iterations) + + +def debug_writing(writer, outputs, labels, inputs, iterations): + tmp_tar = torch.unsqueeze(labels.cpu().data[0], dim=1) + # tmp_out = torch.unsqueeze(outputs.cpu().data[0], dim=1) + + tmp_inp = inputs.cpu().data[0] + tmp_inp[0] += 0.406 + tmp_inp[1] += 0.457 + tmp_inp[2] += 0.480 + + tmp_inp[0] += torch.sum( + F.interpolate(tmp_tar, scale_factor=4, mode="bilinear"), dim=0 + )[0] + tmp_inp.clamp_(0, 1) + + writer.add_image("Data/input", tmp_inp, iterations) + + +def get_root_logger(log_file=None, log_level=logging.INFO): + """Use `get_logger` method in mmcv to get the root logger. + + The logger will be initialized if it has not been initialized. By default a + StreamHandler will be added. If `log_file` is specified, a FileHandler will + also be added. The name of the root logger is the top-level package name, + e.g., "mmpose". + + Args: + log_file (str | None): The log filename. If specified, a FileHandler + will be added to the root logger. + log_level (int): The root logger level. Note that only the process of + rank 0 is affected, while other processes will set the level to + "Error" and be silent most of the time. + + Returns: + logging.Logger: The root logger. + """ + return get_logger(LOG_DIR, log_file, log_level) diff --git a/video/fake-stormer/model_code/losses/__init__.py b/video/fake-stormer/model_code/losses/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6a92031f4fe8b37a637cb268e7d6f239c516e39c --- /dev/null +++ b/video/fake-stormer/model_code/losses/__init__.py @@ -0,0 +1,5 @@ +# -*- coding: utf-8 -*- +from .builder import LOSSES, build_losses +from .losses import BinaryCrossEntropy + +__all__ = ["LOSSES", "build_losses", "BinaryCrossEntropy"] diff --git a/video/fake-stormer/model_code/losses/builder.py b/video/fake-stormer/model_code/losses/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..3602ef6e21758e2de844ed9cad4f0294ab5aaa8a --- /dev/null +++ b/video/fake-stormer/model_code/losses/builder.py @@ -0,0 +1,15 @@ +# -*- coding: utf-8 -*- +from typing import Any, Dict, Optional + +from register.register import Registry, build_from_cfg + +LOSSES = Registry("Loss") + + +def build_losses( + cfg, + loss_func: Registry, + build_func=build_from_cfg, + default_args: Optional[Dict] = None, +) -> Any: + return build_func(cfg, loss_func, default_args) diff --git a/video/fake-stormer/model_code/losses/losses.py b/video/fake-stormer/model_code/losses/losses.py new file mode 100644 index 0000000000000000000000000000000000000000..62aa97883c9e368cf346bb0a5f3b5acffb51a9b2 --- /dev/null +++ b/video/fake-stormer/model_code/losses/losses.py @@ -0,0 +1,485 @@ +# -*- coding: utf-8 -*- +import math + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn import BCELoss, CrossEntropyLoss +from torch.nn.functional import binary_cross_entropy + +from .builder import LOSSES + + +def _sigmoid(hm): + x = hm + y = torch.clamp(x.sigmoid_(), min=1e-4, max=1 - 1e-4) + return y + + +def _avg_sigmoid(hm): + if hm.dim() == 4: + x = torch.mean(hm, [2, 3]) + else: + x = hm + y = torch.clamp(x.sigmoid_(), min=1e-4, max=1 - 1e-4) + return y + + +def f_cstency(cstency_hm_preds, cstency_hm_gt, feature="2D"): + # Heatmap here that is original is returned from model without any modification + cstency_matrix = torch.zeros_like(cstency_hm_gt).cuda() + b_size = cstency_hm_preds.size(0) + + indices_ = cstency_hm_gt.view(b_size, -1).argmax(dim=-1) + + cst_hm_dim = cstency_hm_preds.size(1) + + if feature == "2D": + # Handling 2D output features + cst_hm_h = cstency_hm_preds.size(2) + cst_hm_w = cstency_hm_preds.size(3) + + cstency_matrix_ = torch.matmul( + cstency_hm_preds.view(b_size, cst_hm_dim, -1)[ + np.arange(b_size), :, indices_ + ].view(b_size, 1, cst_hm_dim), + cstency_hm_preds.view(b_size, cst_hm_dim, -1), + ) + cstency_matrix_ = cstency_matrix_.view( + b_size, cstency_hm_gt.size(1), cst_hm_h, cst_hm_w + ) / math.sqrt(cst_hm_dim) + elif feature == "3D": + # Handling 3D output features + cst_hm_d = cstency_hm_preds.size(2) + cst_hm_h = cstency_hm_preds.size(3) + cst_hm_w = cstency_hm_preds.size(4) + + cstency_matrix_ = torch.matmul( + cstency_hm_preds.view(b_size, cst_hm_dim, -1)[ + np.arange(b_size), :, indices_ + ].view(b_size, 1, cst_hm_dim), + cstency_hm_preds.view(b_size, cst_hm_dim, -1), + ) + cstency_matrix_ = cstency_matrix_.view( + b_size, cstency_hm_gt.size(1), cst_hm_d, cst_hm_h, cst_hm_w + ) / math.sqrt(cst_hm_dim) + else: + raise ValueError(f"{feature} output shape has not been supported!") + + cstency_matrix = cstency_matrix_.sigmoid_() + + return cstency_matrix + + +def _neg_pos_loss(hm_pred, hm_gt): + pos_idxes = hm_gt > 0 + neg_idxes = ~pos_idxes + batch_size = hm_gt.size(0) + neg_pos_gt, neg_pos_pred = ( + torch.zeros(batch_size, 1, dtype=torch.float).cuda(), + torch.zeros(batch_size, 1, dtype=torch.float).cuda(), + ) + hm_pred_ = torch.squeeze(torch.clone(hm_pred)) + + for i in range(batch_size): + neg_pos_gt[i] = torch.sum(hm_gt[i][pos_idxes[i, :, :]]) - torch.sum( + hm_gt[i][neg_idxes[i, :, :]] + ) + neg_pos_pred[i] = torch.sum(hm_pred_[i][pos_idxes[i, :, :]]) - torch.sum( + hm_pred_[i][neg_idxes[i, :, :]] + ) + + return torch.abs(neg_pos_pred), torch.abs(neg_pos_gt) + + +def _neg_loss(pred, gt, epsilon=0.1, noise_distribution=0.2, alpha=0.25, **kwargs): + """Modified focal loss. Exactly the same as CornerNet. + Runs faster and costs a little bit more memory + Arguments: + pred (batch x c x h x w) + gt_regr (batch x c x h x w) + """ + hm_mask = kwargs.get( + "hm_mask" + ) # Removing non-computed self-attention positions in total loss + + loss = 0 + pos_inds = gt.eq(1.0).float() + neg_inds = gt.lt(1.0).float() + b_size = gt.shape[0] + + if hm_mask is not None: + pos_inds = hm_mask * pos_inds + neg_inds = hm_mask * neg_inds + + neg_weights = torch.pow(1 - gt, 4) + + # pos_loss = torch.log(pred) * torch.pow(1 - pred, 2) * pos_inds * alpha + pos_loss = (1 - epsilon) * torch.log(pred) * torch.pow(1 - pred, 2) * pos_inds + pos_loss_noise = ( + epsilon + * torch.log(pred) + * torch.pow(1 - pred, 2) + * noise_distribution + * pos_inds + ) + pos_loss = pos_loss + pos_loss_noise + neg_loss = torch.log(1 - pred) * torch.pow(pred, 2) * neg_inds * neg_weights + + num_pos = pos_inds.float().sum() + pos_loss = pos_loss.sum() + neg_loss = neg_loss.sum() + + if num_pos == 0: + loss = loss - neg_loss + else: + loss = loss - (pos_loss + neg_loss) / num_pos + loss *= alpha + return loss + + +def _distance_hm_cls_loss( + cos_sim_ins, hm_preds, hm_gts, label_preds, label_gts, alpha=0.25 +): + b_size = hm_preds.size(0) + hm_preds = hm_preds.view(b_size, -1) + hm_gts = hm_gts.view(b_size, -1) + pos_hm_loss = 0.0 + neg_hm_loss = 0.0 + + for i in range(0, b_size // 2): + for j in range(0, b_size // 2): + pos_hm_loss += (1 / 2) * (1 - cos_sim_ins(hm_preds[i], hm_preds[j])) + neg_hm_loss += (1 / 2) * ( + 1 - cos_sim_ins(hm_preds[i], hm_preds[j + b_size // 2]) + ) + + cos_loss = pos_hm_loss / ((b_size // 2) ** 2) - neg_hm_loss / ((b_size // 2) ** 2) + cos_loss = cos_loss * alpha + return cos_loss + + +@LOSSES.register_module() +class BaseLoss(nn.Module): + def __init__(self, cfg, **kwargs): + self.cfg = cfg + super().__init__() + + for k, v in kwargs.items(): + if v is not None: + self.__setattr__(k, v) + # Critetion ins + self.mse_critetion = nn.MSELoss(reduction=self.cfg.mse_reduction) + if hasattr(self, "use_ce") and getattr(self, "use_ce"): + self.bce_critetion = nn.CrossEntropyLoss(reduction=self.ce_reduction) + else: + self.bce_critetion = nn.BCEWithLogitsLoss( + reduction=self.cfg.ce_reduction + ) # For Binary Cross Entropy Loss + self.ce_critetion = CrossEntropyLoss( + reduction=self.cfg.ce_reduction + ) # For Cross Entropy Loss in general + + # Lambda coefs + self.offset_lmda = self.cfg.offset_lmda + self.cls_lmda = self.cfg.cls_lmda + self.dst_hm_cls_lmda = self.cfg.dst_hm_cls_lmda + self.hm_lmda = self.cfg.hm_lmda + self.cstency_lmda = self.cfg.cstency_lmda + + # Others + self.cos_sim_ins = nn.CosineSimilarity(dim=0, eps=1e-6) + + def _offset_loss(self, preds, gts, apply_filter=False): + loss = 0 + coefs = gts.gt(0).float() if apply_filter else 1 + n_coefs = coefs.float().sum() + + loss = 0.5 * self.mse_critetion(preds * coefs, gts * coefs) + loss /= n_coefs + 1e-6 + loss *= self.offset_lmda + return loss + + def _cls_loss(self, preds, gts): + loss = 0 + loss = self.bce_critetion(preds, gts) + loss *= self.cls_lmda + return loss + + def _consistency_loss(self, preds, gts, feature="2D"): + loss = torch.zeros(1).cuda() + encode_preds = f_cstency(preds, gts, feature=feature) + # loss = self.bce_critetion(encode_preds.view(-1, 1), gts.view(-1, 1)) + loss = self.mse_critetion(encode_preds, gts) + loss *= self.cstency_lmda + return loss.sum() + + def _temp_loc_loss(self, preds, gts, alpha=0.25, gamma=2): + """ + Calculating loss for temporal location + agrs: + preds: output prediction of temporal location + gts: gt of temporal location + """ + loss = self.bce_critetion( + preds.view(-1).unsqueeze(-1), gts.view(-1).unsqueeze(-1) + ) + loss *= self.cfg.tmp_loc_lmda + return loss + + +@LOSSES.register_module() +class BinaryCrossEntropy(nn.Module): + def __init__(self, cfg, reduction="mean"): + super(BinaryCrossEntropy, self).__init__() + self.reduction = reduction + self.bce = nn.BCEWithLogitsLoss(reduction=self.reduction) + + def __call__(self, pred, y): + return self.bce(pred, y) + + +@LOSSES.register_module() +class CombinedFocalLoss(BaseLoss): + """nn.Module warpper for focal loss""" + + def __init__(self, cfg, use_target_weight, **kwargs): + super(CombinedFocalLoss, self).__init__(cfg, **kwargs) + self.hm_loss = _neg_loss + self.feature = kwargs.get("feature") or "2D" + + def forward( + self, + hm_outputs, + hm_targets, + cls_preds, + cls_gts, + hm_mask=None, + offset_preds=None, + offset_gts=None, + cstency_preds=None, + cstency_gts=None, + target_weight=None, + temp_loc_preds=None, + temp_loc_gts=None, + ): + loss_return = {} + hm_outputs_ = torch.clone(hm_outputs) + hm_outputs_ = _sigmoid(hm_outputs_) + if hm_targets.dim() == 3: + hm_targets = torch.unsqueeze(hm_targets, 1) + + loss_hm = self.hm_loss( + hm_outputs_, hm_targets, alpha=self.hm_lmda, hm_mask=hm_mask + ) + loss_return["hm"] = loss_hm + loss_return["cls"] = self._cls_loss(cls_preds, cls_gts) + + if self.dst_hm_cls_lmda > 0: + loss_return["dst_hm_cls"] = _distance_hm_cls_loss( + self.cos_sim_ins, + hm_outputs, + hm_targets, + cls_preds, + cls_gts, + alpha=self.dst_hm_cls_lmda, + ) + + if self.offset_lmda > 0 and offset_preds is not None: + loss_return["offset"] = self._offset_loss( + offset_preds, offset_gts, apply_filter=True + ) + + if self.cstency_lmda > 0 and cstency_preds is not None: + loss_return["cstency"] = self._consistency_loss( + cstency_preds, cstency_gts, feature=self.feature + ) + + if temp_loc_preds is not None and self.cfg.tmp_loc_lmda is not None: + loss_temp_loc = self._temp_loc_loss(temp_loc_preds, temp_loc_gts) + loss_return["temp_loc"] = loss_temp_loc + + return loss_return + + +@LOSSES.register_module() +class JointsMSELoss(nn.Module): + def __init__(self, use_target_weight, reduction="mean", lmda=1): + super(JointsMSELoss, self).__init__() + self.reduction = reduction + self.criterion = nn.MSELoss(reduction=reduction) + self.use_target_weight = use_target_weight + self.lmda = lmda + + def forward(self, output, target, target_weight=None, **kwargs): + hm_mask = kwargs.get("hm_mask") + batch_size = output.size(0) + num_joints = output.size(1) + heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1) + heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1) + if hm_mask is not None: + hm_mask_ = hm_mask.reshape((batch_size, num_joints, -1)).split(1, 1) + loss = 0 + + for idx in range(num_joints): + heatmap_pred = heatmaps_pred[idx].squeeze() + heatmap_gt = heatmaps_gt[idx].squeeze() + if self.use_target_weight and target_weight is not None: + loss += 0.5 * self.criterion( + heatmap_pred.mul(target_weight[:, idx]), + heatmap_gt.mul(target_weight[:, idx]), + ) + else: + if hm_mask is not None: + heatmap_pred = heatmap_pred * hm_mask_[idx].squeeze() + heatmap_gt = heatmap_gt * hm_mask_[idx].squeeze() + loss += 0.5 * self.criterion(heatmap_pred, heatmap_gt) + + if self.reduction != "mean": + loss = loss * self.lmda / num_joints + else: + loss = loss * self.lmda + + return loss + + +@LOSSES.register_module() +class CombinedMSELoss(BaseLoss): + def __init__(self, cfg, use_target_weight=False, **kwargs): + super(CombinedMSELoss, self).__init__(cfg=cfg, **kwargs) + self.criterion_hm = JointsMSELoss( + use_target_weight=use_target_weight, + reduction=self.cfg.mse_reduction, + lmda=self.cfg.hm_lmda, + ) + self.use_target_weight = use_target_weight + self.feature = kwargs.get("feature") or "2D" + + def forward( + self, + hm_outputs, + hm_targets, + cls_preds, + cls_gts, + hm_mask=None, + target_weight=None, + cstency_preds=None, + cstency_gts=None, + temp_loc_preds=None, + temp_loc_gts=None, + **kwargs, + ): + loss_return = {} + loss_hm = self.criterion_hm( + hm_outputs, hm_targets, target_weight=target_weight, hm_mask=hm_mask + ) + loss_return["hm"] = loss_hm + + loss_cls = self._cls_loss(cls_preds / self.temperature, cls_gts) + loss_return["cls"] = loss_cls + + if self.cstency_lmda > 0 and cstency_preds is not None: + loss_return["cstency"] = self._consistency_loss( + cstency_preds, cstency_gts, feature=self.feature + ) + + if temp_loc_preds is not None and self.cfg.tmp_loc_lmda is not None: + loss_temp_loc = self._temp_loc_loss(temp_loc_preds, temp_loc_gts) + loss_return["temp_loc"] = loss_temp_loc + + return loss_return + + +@LOSSES.register_module() +class CombinedHeatmapBinaryLoss(nn.Module): + def __init__( + self, use_target_weight, cls_lmda=0.2, reduction="mean", cls_cal=True, **kwargs + ): + super(CombinedHeatmapBinaryLoss, self).__init__() + self.criterion_cls = BinaryCrossEntropy(reduction=reduction) + self.criterion_hm = BinaryCrossEntropy(reduction=reduction) + self.use_target_weight = use_target_weight + self.cls_lmda = cls_lmda if cls_cal else 0 + self.cls_cal = cls_cal + + def forward(self, hm_outputs, hm_targets, cls_preds, cls_gts, target_weight=None): + batch_size = hm_outputs.size(0) + hm_targets = hm_targets[:, :, :, 0] + hm_h = hm_outputs.size(2) + hm_w = hm_outputs.size(3) + total_pixels = hm_h * hm_w + loss_hm = torch.zeros(1).cuda() + hm_outputs_ = torch.clone(hm_outputs) + hm_outputs_ = _sigmoid(hm_outputs_) + + for i in range(hm_h): + for j in range(hm_w): + loss_hm_ = self.criterion_hm( + hm_outputs_[:, :, i, j], torch.unsqueeze(hm_targets[:, i, j], 1) + ) + loss_hm += loss_hm_ + + loss_hm = loss_hm / total_pixels + loss_return = {} + loss_return["hm"] = loss_hm + + loss_cls = self.criterion_cls(cls_preds, cls_gts) + loss_return["cls"] = loss_cls + return loss_return + + +@LOSSES.register_module() +class CombinedPolyLoss(nn.Module): + """ + PolyLoss: A Polynomial Expansion Perspective of Classification Loss Functions + """ + + def __init__( + self, + use_target_weight, + epsilon=2.0, + cls_lmda=0.05, + reduction="mean", + cls_cal=True, + **kwargs, + ): + super(CombinedPolyLoss, self).__init__() + self.cls_critetion = BinaryCrossEntropy(reduction=reduction) + self.use_target_weight = use_target_weight + self.epsilon = epsilon + self.cls_lmda = cls_lmda if cls_cal else 0 + self.reduction = reduction + self.cls_cal = cls_cal + + def forward(self, hm_outputs, hm_targets, cls_preds, cls_gts): + batch_size = hm_outputs.size(0) + n_classes = hm_outputs.size(1) + hm_h = hm_outputs.size(2) + hm_w = hm_outputs.size(3) + total_pixels = hm_h * hm_w + poly_loss = torch.zeros(batch_size, 1).cuda() + hm_outputs_ = _sigmoid(hm_outputs) + + for i in range(hm_h): + for j in range(hm_w): + ce = binary_cross_entropy( + hm_outputs_[:, :, i, j], + torch.unsqueeze(hm_targets[:, i, j], -1), + reduction="none", + ) + pt = hm_outputs_[:, :, i, j] + pt = torch.squeeze(pt) + pt = torch.where(hm_targets[:, i, j] > 0, pt, 1 - pt) + poly_loss += ce + self.epsilon * (1.0 - torch.unsqueeze(pt, -1)) + + if self.reduction == "mean": + poly_loss = poly_loss.sum() / total_pixels / batch_size + else: + poly_loss = poly_loss.sum() + loss_return = {} + loss_return["hm"] = poly_loss + + loss_cls = self.cls_critetion(cls_preds, cls_gts) + loss_return["cls"] = loss_cls * self.cls_lmda + return loss_return diff --git a/video/fake-stormer/model_code/models/__init__.py b/video/fake-stormer/model_code/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0e79a1504a6cc94e37644af9a91dc288e636974b --- /dev/null +++ b/video/fake-stormer/model_code/models/__init__.py @@ -0,0 +1,57 @@ +# -*- coding: utf-8 -*- +from .builder import MODELS, build_model +from .networks.backbones import ( + ResNet3D, + SwinTransformer, + SwinTransformer3D, + TimeViT, + ViT, + Xception, +) +from .networks.backbones.arcface import ( + SimpleClassificationDF, +) +from .networks.common import * +from .networks.detectors import TopDownDetector +from .networks.heads.hm_simple_head import TopdownHeatmapSimpleHead +from .networks.mrsa_resnet import Bottleneck, PoseResNet, resnet_spec +from .networks.necks import EFPN3D +from .networks.pose_efficientNet import PoseEfficientNet +from .networks.pose_hrnet import PoseHighResolutionNet +from .utils import ( + freeze_backbone, + load_model, + load_pretrained, + n_param_model, + preset_model, + save_model, + unfreeze_backbone, +) + +__all__ = [ + "SimpleClassificationDF", + "PoseResNet", + "MODELS", + "build_model", + "load_pretrained", + "freeze_backbone", + "resnet_spec", + "n_param_model", + "load_model", + "save_model", + "unfreeze_backbone", + "Bottleneck", + "preset_model", + "PoseHighResolutionNet", + "Xception", + "PoseEfficientNet", + "TopDownDetector", + "ViT", + "TopdownHeatmapSimpleHead", + "TimeViT", + "SwinTransformer", + "ResNet3D", + "Xception", + "SwinTransformer3D", + "EFPN3D", +] diff --git a/video/fake-stormer/model_code/models/builder.py b/video/fake-stormer/model_code/models/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..47f605ab1e84b61a59e375de95516fded85915f5 --- /dev/null +++ b/video/fake-stormer/model_code/models/builder.py @@ -0,0 +1,46 @@ +# -*- coding: utf-8 -*- +import os +import sys +from typing import Any, Dict, Optional + +if not os.getcwd() in sys.path: + sys.path.append(os.getcwd()) + +from register.register import Registry, build_from_cfg +from torch.nn import Sequential + + +def build_model_from_cfg(cfg, registry, default_args=None): + """Build a PyTorch model from config dict(s). Different from + ``build_from_cfg``, if cfg is a list, a ``nn.Sequential`` will be built. + Args: + cfg (dict, list[dict]): The config of modules, is is either a config + dict or a list of config dicts. If cfg is a list, a + the built modules will be wrapped with ``nn.Sequential``. + registry (:obj:`Registry`): A registry the module belongs to. + default_args (dict, optional): Default arguments to build the module. + Defaults to None. + Returns: + nn.Module: A built nn module. + """ + if isinstance(cfg, list): + modules = [build_from_cfg(cfg_, registry, default_args) for cfg_ in cfg] + return Sequential(*modules) + else: + return build_from_cfg(cfg, registry, default_args) + + +MODELS = Registry("model", build_func=build_model_from_cfg) +HEADS = MODELS +BACKBONES = MODELS +DETECTORS = MODELS +NECKS = MODELS + + +def build_model( + cfg: Dict, + model: Registry, + build_func=build_model_from_cfg, + default_args: Optional[Dict] = None, +) -> Any: + return build_func(cfg, model, default_args) diff --git a/video/fake-stormer/model_code/models/networks/backbones/__init__.py b/video/fake-stormer/model_code/models/networks/backbones/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b012fa9f3f3374d9c242e1a568a7d87d729bd5e0 --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/backbones/__init__.py @@ -0,0 +1,15 @@ +# -*- coding: utf-8 -*- +from .resnet3d import ResNet3D +from .swin import SwinTransformer +from .swin3d import SwinTransformer3D +from .vit import TimeViT, ViT +from .xception import Xception + +__all__ = [ + "ViT", + "TimeViT", + "SwinTransformer", + "ResNet3D", + "Xception", + "SwinTransformer3D", +] diff --git a/video/fake-stormer/model_code/models/networks/backbones/arcface.py b/video/fake-stormer/model_code/models/networks/backbones/arcface.py new file mode 100644 index 0000000000000000000000000000000000000000..33ee14e492c661bfa1abb38457bfc09874155ae2 --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/backbones/arcface.py @@ -0,0 +1,484 @@ +# -*- coding: utf-8 -*- +import math +import os +from collections import namedtuple + +import torch +import torch.nn.functional as F +from torch.nn import ( + AdaptiveAvgPool2d, + AvgPool2d, + BatchNorm1d, + BatchNorm2d, + Conv2d, + Dropout, + Dropout2d, + Linear, + MaxPool2d, + Module, + Parameter, + PReLU, + ReLU, + Sequential, + Sigmoid, + Softmax, +) + +from ...builder import ( + BACKBONES, + HEADS, + MODELS, + build_model, +) + +################################## Original Arcface Model ############################################################# + + +class Flatten(Module): + def forward(self, input): + return input.view(input.size(0), -1) + + +def l2_norm(input, axis=1): + norm = torch.norm(input, 2, axis, True) + output = torch.div(input, norm) + return output + + +class SEModule(Module): + def __init__(self, channels, reduction): + super(SEModule, self).__init__() + self.avg_pool = AdaptiveAvgPool2d(1) + self.fc1 = Conv2d( + channels, channels // reduction, kernel_size=1, padding=0, bias=False + ) + 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 bottleneck_IR(Module): + def __init__(self, in_channel, depth, stride): + super(bottleneck_IR, 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), + 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 bottleneck_IR_SE(Module): + def __init__(self, in_channel, depth, stride): + super(bottleneck_IR_SE, 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), + PReLU(depth), + Conv2d(depth, depth, (3, 3), stride, 1, bias=False), + BatchNorm2d(depth), + SEModule(depth, 16), + ) + + def forward(self, x): + shortcut = self.shortcut_layer(x) + res = self.res_layer(x) + return res + shortcut + + +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 == 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=64, num_units=3), + get_block(in_channel=64, depth=128, num_units=8), + get_block(in_channel=128, depth=256, num_units=36), + get_block(in_channel=256, depth=512, num_units=3), + ] + return blocks + + +@BACKBONES.register_module() +class ResNet(Module): + def __init__(self, num_layers=50, drop_ratio=0.6, mode="ir", **kwargs): + """ + Implementation for ResNet 50, 101, 152 with/out SE module + """ + super(ResNet, self).__init__() + assert num_layers in [50, 100, 152], "num_layers should be 50,100, or 152" + assert mode in ["ir", "ir_se"], "mode should be ir or ir_se" + blocks = get_blocks(num_layers) + if mode == "ir": + unit_module = bottleneck_IR + elif mode == "ir_se": + unit_module = bottleneck_IR_SE + self.input_layer = Sequential( + Conv2d(3, 64, (3, 3), 1, 1, bias=False), BatchNorm2d(64), PReLU(64) + ) + self.output_layer = Sequential( + BatchNorm2d(512), + Dropout(drop_ratio), + Flatten(), + Linear(512 * 7 * 7, 512), + BatchNorm1d(512), + ) + modules = [] + for block in blocks: + for bottleneck in block: + modules.append( + unit_module( + bottleneck.in_channel, bottleneck.depth, bottleneck.stride + ) + ) + self.body = Sequential(*modules) + + def forward(self, x): + x = self.input_layer(x) + x = self.body(x) + x = self.output_layer(x) + x = l2_norm(x) + return x + + +@HEADS.register_module() +class SimpleClassificationHead(Module): + def __init__(self, drop_ratio=0.6, in_planes=512, **kwargs): + super(SimpleClassificationHead, self).__init__() + self.classification_head = Sequential( + Dropout(drop_ratio), + Linear(in_planes, 256), + BatchNorm1d(256), + Dropout(drop_ratio), + Linear(256, 128), + BatchNorm1d(128), + Dropout(drop_ratio), + Linear(128, 64), + BatchNorm1d(64), + Dropout(drop_ratio), + Linear(64, 32), + BatchNorm1d(32), + # Dropout(drop_ratio), + Linear(32, 1), + Sigmoid(), + ) + + def forward(self, x): + x = self.classification_head(x) + return x + + +@MODELS.register_module() +class SimpleClassificationDF(Module): + def __init__(self, cfg: dict, **kwargs): + super(SimpleClassificationDF, self).__init__() + assert "backbone" in cfg, "Config for Backbones is mandatory!" + assert "head" in cfg, "Config for Heads is mandatory!" + + self.backbone = BACKBONES.get(cfg.backbone.type)(**cfg.backbone) + self.head = HEADS.get(cfg.head.type)(**cfg.head) + self.model = Sequential(*[self.backbone, self.head]) + + def forward(self, x): + x = self.model(x) + return x + + +################################## MobileFaceNet ############################################################# + + +class Conv_block(Module): + def __init__( + self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1 + ): + super(Conv_block, self).__init__() + self.conv = Conv2d( + in_c, + out_channels=out_c, + kernel_size=kernel, + groups=groups, + stride=stride, + padding=padding, + bias=False, + ) + self.bn = BatchNorm2d(out_c) + self.prelu = PReLU(out_c) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + x = self.prelu(x) + return x + + +class Linear_block(Module): + def __init__( + self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1 + ): + super(Linear_block, self).__init__() + self.conv = Conv2d( + in_c, + out_channels=out_c, + kernel_size=kernel, + groups=groups, + stride=stride, + padding=padding, + bias=False, + ) + self.bn = BatchNorm2d(out_c) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + return x + + +class Depth_Wise(Module): + def __init__( + self, + in_c, + out_c, + residual=False, + kernel=(3, 3), + stride=(2, 2), + padding=(1, 1), + groups=1, + ): + super(Depth_Wise, self).__init__() + self.conv = Conv_block( + in_c, out_c=groups, kernel=(1, 1), padding=(0, 0), stride=(1, 1) + ) + self.conv_dw = Conv_block( + groups, groups, groups=groups, kernel=kernel, padding=padding, stride=stride + ) + self.project = Linear_block( + groups, out_c, kernel=(1, 1), padding=(0, 0), stride=(1, 1) + ) + self.residual = residual + + def forward(self, x): + if self.residual: + short_cut = x + x = self.conv(x) + x = self.conv_dw(x) + x = self.project(x) + if self.residual: + output = short_cut + x + else: + output = x + return output + + +class Residual(Module): + def __init__( + self, c, num_block, groups, kernel=(3, 3), stride=(1, 1), padding=(1, 1) + ): + super(Residual, self).__init__() + modules = [] + for _ in range(num_block): + modules.append( + Depth_Wise( + c, + c, + residual=True, + kernel=kernel, + padding=padding, + stride=stride, + groups=groups, + ) + ) + self.model = Sequential(*modules) + + def forward(self, x): + return self.model(x) + + +class MobileFaceNet(Module): + def __init__(self, embedding_size): + super(MobileFaceNet, self).__init__() + self.conv1 = Conv_block(3, 64, kernel=(3, 3), stride=(2, 2), padding=(1, 1)) + self.conv2_dw = Conv_block( + 64, 64, kernel=(3, 3), stride=(1, 1), padding=(1, 1), groups=64 + ) + self.conv_23 = Depth_Wise( + 64, 64, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=128 + ) + self.conv_3 = Residual( + 64, num_block=4, groups=128, kernel=(3, 3), stride=(1, 1), padding=(1, 1) + ) + self.conv_34 = Depth_Wise( + 64, 128, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=256 + ) + self.conv_4 = Residual( + 128, num_block=6, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1) + ) + self.conv_45 = Depth_Wise( + 128, 128, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=512 + ) + self.conv_5 = Residual( + 128, num_block=2, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1) + ) + self.conv_6_sep = Conv_block( + 128, 512, kernel=(1, 1), stride=(1, 1), padding=(0, 0) + ) + self.conv_6_dw = Linear_block( + 512, 512, groups=512, kernel=(7, 7), stride=(1, 1), padding=(0, 0) + ) + self.conv_6_flatten = Flatten() + self.linear = Linear(512, embedding_size, bias=False) + self.bn = BatchNorm1d(embedding_size) + + def forward(self, x): + out = self.conv1(x) + out = self.conv2_dw(out) + out = self.conv_23(out) + out = self.conv_3(out) + out = self.conv_34(out) + out = self.conv_4(out) + out = self.conv_45(out) + out = self.conv_5(out) + out = self.conv_6_sep(out) + out = self.conv_6_dw(out) + out = self.conv_6_flatten(out) + out = self.linear(out) + out = self.bn(out) + + return l2_norm(out) + + +################################## Arcface head ############################################################# + + +class Arcface(Module): + # implementation of additive margin softmax loss in https://arxiv.org/abs/1801.05599 + def __init__(self, embedding_size=512, classnum=51332, s=64.0, m=0.5): + super(Arcface, self).__init__() + self.classnum = classnum + self.kernel = Parameter(torch.Tensor(embedding_size, classnum)) + # initial kernel + self.kernel.data.uniform_(-1, 1).renorm_(2, 1, 1e-5).mul_(1e5) + self.m = m # the margin value, default is 0.5 + self.s = s # scalar value default is 64, see normface https://arxiv.org/abs/1704.06369 + self.cos_m = math.cos(m) + self.sin_m = math.sin(m) + self.mm = self.sin_m * m # issue 1 + self.threshold = math.cos(math.pi - m) + + def forward(self, embbedings, label): + # weights norm + nB = len(embbedings) + kernel_norm = l2_norm(self.kernel, axis=0) + # cos(theta+m) + cos_theta = torch.mm(embbedings, kernel_norm) + # output = torch.mm(embbedings,kernel_norm) + cos_theta = cos_theta.clamp(-1, 1) # for numerical stability + cos_theta_2 = torch.pow(cos_theta, 2) + sin_theta_2 = 1 - cos_theta_2 + sin_theta = torch.sqrt(sin_theta_2) + cos_theta_m = cos_theta * self.cos_m - sin_theta * self.sin_m + # this condition controls the theta+m should in range [0, pi] + # 0<=theta+m<=pi + # -m<=theta<=pi-m + cond_v = cos_theta - self.threshold + cond_mask = cond_v <= 0 + keep_val = cos_theta - self.mm # when theta not in [0,pi], use cosface instead + cos_theta_m[cond_mask] = keep_val[cond_mask] + output = ( + cos_theta * 1.0 + ) # a little bit hacky way to prevent in_place operation on cos_theta + idx_ = torch.arange(0, nB, dtype=torch.long) + output[idx_, label] = cos_theta_m[idx_, label] + output *= ( + self.s + ) # scale up in order to make softmax work, first introduced in normface + return output + + +################################## Cosface head ############################################################# + + +class Am_softmax(Module): + # implementation of additive margin softmax loss in https://arxiv.org/abs/1801.05599 + def __init__(self, embedding_size=512, classnum=51332): + super(Am_softmax, self).__init__() + self.classnum = classnum + self.kernel = Parameter(torch.Tensor(embedding_size, classnum)) + # initial kernel + self.kernel.data.uniform_(-1, 1).renorm_(2, 1, 1e-5).mul_(1e5) + self.m = 0.35 # additive margin recommended by the paper + self.s = 30.0 # see normface https://arxiv.org/abs/1704.06369 + + def forward(self, embbedings, label): + kernel_norm = l2_norm(self.kernel, axis=0) + cos_theta = torch.mm(embbedings, kernel_norm) + cos_theta = cos_theta.clamp(-1, 1) # for numerical stability + phi = cos_theta - self.m + label = label.view(-1, 1) # size=(B,1) + index = cos_theta.data * 0.0 # size=(B,Classnum) + index.scatter_(1, label.data.view(-1, 1), 1) + index = index.byte() + output = cos_theta * 1.0 + output[index] = phi[index] # only change the correct predicted output + output *= ( + self.s + ) # scale up in order to make softmax work, first introduced in normface + return output + + +if __name__ == "__main__": + cfg = dict(num_layers=50, drop_ratio=0.6, mode="ir", type="Backbone") + backbone = MODELS.build(cfg) + print(backbone) diff --git a/video/fake-stormer/model_code/models/networks/backbones/base.py b/video/fake-stormer/model_code/models/networks/backbones/base.py new file mode 100644 index 0000000000000000000000000000000000000000..1154ff8b3d6d1e9a51eac290ad8186115d0be3b3 --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/backbones/base.py @@ -0,0 +1,50 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import logging +from abc import ABCMeta, abstractmethod + +import torch.nn as nn +from models.utils import load_checkpoint + + +class BaseBackbone(nn.Module, metaclass=ABCMeta): + """Base backbone. + + This class defines the basic functions of a backbone. Any backbone that + inherits this class should at least define its own `forward` function. + """ + + def init_weights(self, pretrained=None, patch_padding="pad", part_features=None): + """Init backbone weights. + + Args: + pretrained (str | None): If pretrained is a string, then it + initializes backbone weights by loading the pretrained + checkpoint. If pretrained is None, then it follows default + initializer or customized initializer in subclasses. + """ + if isinstance(pretrained, str): + logger = logging.getLogger() + load_checkpoint( + self, + pretrained, + strict=False, + logger=logger, + patch_padding=patch_padding, + part_features=part_features, + ) + elif pretrained is None: + # use default initializer or customized initializer in subclasses + pass + else: + raise TypeError( + "pretrained must be a str or None." f" But received {type(pretrained)}." + ) + + @abstractmethod + def forward(self, x): + """Forward function. + + Args: + x (Tensor | tuple[Tensor]): x could be a torch.Tensor or a tuple of + torch.Tensor, containing input data for forward computation. + """ diff --git a/video/fake-stormer/model_code/models/networks/backbones/efficientNet.py b/video/fake-stormer/model_code/models/networks/backbones/efficientNet.py new file mode 100644 index 0000000000000000000000000000000000000000..f38e265d058dfa8ca82998fcdf26451a591c55b8 --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/backbones/efficientNet.py @@ -0,0 +1,592 @@ +# -*- coding: utf-8 -*- +import collections +import math +import re +from functools import partial + +import torch +import torch.nn as nn +from torch.nn import functional as F +from torch.utils import model_zoo + +# Parameters for the entire model (stem, all blocks, and head) +GlobalParams = collections.namedtuple( + "GlobalParams", + [ + "width_coefficient", + "depth_coefficient", + "image_size", + "dropout_rate", + "num_classes", + "batch_norm_momentum", + "batch_norm_epsilon", + "drop_connect_rate", + "depth_divisor", + "min_depth", + "include_top", + "include_hm_decoder", + "head_conv", + "heads", + "num_layers", + "INIT_WEIGHTS", + "use_c2", + "use_c3", + "use_c4", + "use_c51", + "efpn", + "se_layer", + "tfpn", + "norm_c2", + ], +) + +# Parameters for an individual model block +BlockArgs = collections.namedtuple( + "BlockArgs", + [ + "num_repeat", + "kernel_size", + "stride", + "expand_ratio", + "input_filters", + "output_filters", + "se_ratio", + "id_skip", + ], +) + +# Set GlobalParams and BlockArgs's defaults +GlobalParams.__new__.__defaults__ = (None,) * len(GlobalParams._fields) +BlockArgs.__new__.__defaults__ = (None,) * len(BlockArgs._fields) + + +# Swish activation function +if hasattr(nn, "SiLU"): + Swish = nn.SiLU +else: + # For compatibility with old PyTorch versions + class Swish(nn.Module): + def forward(self, x): + return x * torch.sigmoid(x) + + +def round_filters(filters, global_params): + """Calculate and round number of filters based on width multiplier. + Use width_coefficient, depth_divisor and min_depth of global_params. + Args: + filters (int): Filters number to be calculated. + global_params (namedtuple): Global params of the model. + Returns: + new_filters: New filters number after calculating. + """ + multiplier = global_params.width_coefficient + if not multiplier: + return filters + # TODO: modify the params names. + # maybe the names (width_divisor,min_width) + # are more suitable than (depth_divisor,min_depth). + divisor = global_params.depth_divisor + min_depth = global_params.min_depth + filters *= multiplier + min_depth = min_depth or divisor # pay attention to this line when using min_depth + # follow the formula transferred from official TensorFlow implementation + new_filters = max(min_depth, int(filters + divisor / 2) // divisor * divisor) + if new_filters < 0.9 * filters: # prevent rounding by more than 10% + new_filters += divisor + return int(new_filters) + + +def round_repeats(repeats, global_params): + """Calculate module's repeat number of a block based on depth multiplier. + Use depth_coefficient of global_params. + Args: + repeats (int): num_repeat to be calculated. + global_params (namedtuple): Global params of the model. + Returns: + new repeat: New repeat number after calculating. + """ + multiplier = global_params.depth_coefficient + if not multiplier: + return repeats + # follow the formula transferred from official TensorFlow implementation + return int(math.ceil(multiplier * repeats)) + + +def drop_connect(inputs, p, training): + """Drop connect. + Args: + input (tensor: BCWH): Input of this structure. + p (float: 0.0~1.0): Probability of drop connection. + training (bool): The running mode. + Returns: + output: Output after drop connection. + """ + assert 0 <= p <= 1, "p must be in range of [0,1]" + + if not training: + return inputs + + batch_size = inputs.shape[0] + keep_prob = 1 - p + + # generate binary_tensor mask according to probability (p for 0, 1-p for 1) + random_tensor = keep_prob + random_tensor += torch.rand( + [batch_size, 1, 1, 1], dtype=inputs.dtype, device=inputs.device + ) + binary_tensor = torch.floor(random_tensor) + + output = inputs / keep_prob * binary_tensor + return output + + +def get_same_padding_conv2d(image_size=None): + """Chooses static padding if you have specified an image size, and dynamic padding otherwise. + Static padding is necessary for ONNX exporting of models. + Args: + image_size (int or tuple): Size of the image. + Returns: + Conv2dDynamicSamePadding or Conv2dStaticSamePadding. + """ + if image_size is None: + return Conv2dDynamicSamePadding + else: + return partial(Conv2dStaticSamePadding, image_size=image_size) + + +class Conv2dDynamicSamePadding(nn.Conv2d): + """2D Convolutions like TensorFlow, for a dynamic image size. + The padding is operated in forward function by calculating dynamically. + """ + + # Tips for 'SAME' mode padding. + # Given the following: + # i: width or height + # s: stride + # k: kernel size + # d: dilation + # p: padding + # Output after Conv2d: + # o = floor((i+p-((k-1)*d+1))/s+1) + # If o equals i, i = floor((i+p-((k-1)*d+1))/s+1), + # => p = (i-1)*s+((k-1)*d+1)-i + + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride=1, + dilation=1, + groups=1, + bias=True, + ): + super().__init__( + in_channels, out_channels, kernel_size, stride, 0, dilation, groups, bias + ) + self.stride = self.stride if len(self.stride) == 2 else [self.stride[0]] * 2 + + def forward(self, x): + ih, iw = x.size()[-2:] + kh, kw = self.weight.size()[-2:] + sh, sw = self.stride + oh, ow = math.ceil(ih / sh), math.ceil( + iw / sw + ) # change the output size according to stride ! ! ! + pad_h = max((oh - 1) * self.stride[0] + (kh - 1) * self.dilation[0] + 1 - ih, 0) + pad_w = max((ow - 1) * self.stride[1] + (kw - 1) * self.dilation[1] + 1 - iw, 0) + if pad_h > 0 or pad_w > 0: + x = F.pad( + x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2] + ) + return F.conv2d( + x, + self.weight, + self.bias, + self.stride, + self.padding, + self.dilation, + self.groups, + ) + + +class Conv2dStaticSamePadding(nn.Conv2d): + """2D Convolutions like TensorFlow's 'SAME' mode, with the given input image size. + The padding mudule is calculated in construction function, then used in forward. + """ + + # With the same calculation as Conv2dDynamicSamePadding + + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride=1, + image_size=None, + **kwargs, + ): + super().__init__(in_channels, out_channels, kernel_size, stride, **kwargs) + self.stride = self.stride if len(self.stride) == 2 else [self.stride[0]] * 2 + + # Calculate padding based on image size and save it + assert image_size is not None + ih, iw = (image_size, image_size) if isinstance(image_size, int) else image_size + kh, kw = self.weight.size()[-2:] + sh, sw = self.stride + oh, ow = math.ceil(ih / sh), math.ceil(iw / sw) + pad_h = max((oh - 1) * self.stride[0] + (kh - 1) * self.dilation[0] + 1 - ih, 0) + pad_w = max((ow - 1) * self.stride[1] + (kw - 1) * self.dilation[1] + 1 - iw, 0) + if pad_h > 0 or pad_w > 0: + self.static_padding = nn.ZeroPad2d( + (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2) + ) + else: + self.static_padding = nn.Identity() + + def forward(self, x): + x = self.static_padding(x) + x = F.conv2d( + x, + self.weight, + self.bias, + self.stride, + self.padding, + self.dilation, + self.groups, + ) + return x + + +def get_model_params(model_name, override_params): + """Get the block args and global params for a given model name. + Args: + model_name (str): Model's name. + override_params (dict): A dict to modify global_params. + Returns: + blocks_args, global_params + """ + if model_name.startswith("efficientnet"): + w, d, s, p = efficientnet_params(model_name) + # note: all models have drop connect rate = 0.2 + blocks_args, global_params = efficientnet( + width_coefficient=w, depth_coefficient=d, dropout_rate=p, image_size=s + ) + else: + raise NotImplementedError( + "model name is not pre-defined: {}".format(model_name) + ) + if override_params: + # ValueError will be raised here if override_params has fields not included in global_params. + global_params = global_params._replace(**override_params) + return blocks_args, global_params + + +def efficientnet_params(model_name): + """Map EfficientNet model name to parameter coefficients. + Args: + model_name (str): Model name to be queried. + Returns: + params_dict[model_name]: A (width,depth,res,dropout) tuple. + """ + params_dict = { + # Coefficients: width,depth,res,dropout + "efficientnet-b0": (1.0, 1.0, 224, 0.2), + "efficientnet-b1": (1.0, 1.1, 240, 0.2), + "efficientnet-b2": (1.1, 1.2, 260, 0.3), + "efficientnet-b3": (1.2, 1.4, 300, 0.3), + "efficientnet-b4": (1.4, 1.8, 380, 0.4), + "efficientnet-b5": (1.6, 2.2, 456, 0.4), + "efficientnet-b6": (1.8, 2.6, 528, 0.5), + "efficientnet-b7": (2.0, 3.1, 600, 0.5), + "efficientnet-b8": (2.2, 3.6, 672, 0.5), + "efficientnet-l2": (4.3, 5.3, 800, 0.5), + } + return params_dict[model_name] + + +def efficientnet( + width_coefficient=None, + depth_coefficient=None, + image_size=None, + dropout_rate=0.2, + drop_connect_rate=0.2, + num_classes=1000, + include_top=True, + include_hm_decoder=False, + head_conv=None, + heads=None, + use_c2=False, + use_c3=False, + use_c4=False, + use_c51=False, + num_layers=None, + INIT_WEIGHTS=None, + efpn=False, + se_layer=False, + tfpn=False, + norm_c2=False, +): + """Create BlockArgs and GlobalParams for efficientnet model. + Args: + width_coefficient (float) + depth_coefficient (float) + image_size (int) + dropout_rate (float) + drop_connect_rate (float) + num_classes (int) + Meaning as the name suggests. + Returns: + blocks_args, global_params. + """ + + # Blocks args for the whole model(efficientnet-b0 by default) + # It will be modified in the construction of EfficientNet Class according to model + blocks_args = [ + "r1_k3_s11_e1_i32_o16_se0.25", + "r2_k3_s22_e6_i16_o24_se0.25", + "r2_k5_s22_e6_i24_o40_se0.25", + "r3_k3_s22_e6_i40_o80_se0.25", + "r3_k5_s11_e6_i80_o112_se0.25", + "r4_k5_s22_e6_i112_o192_se0.25", + "r1_k3_s11_e6_i192_o320_se0.25", + ] + blocks_args = BlockDecoder.decode(blocks_args) + + global_params = GlobalParams( + width_coefficient=width_coefficient, + depth_coefficient=depth_coefficient, + image_size=image_size, + dropout_rate=dropout_rate, + num_classes=num_classes, + batch_norm_momentum=0.99, + batch_norm_epsilon=1e-3, + drop_connect_rate=drop_connect_rate, + depth_divisor=8, + min_depth=None, + include_top=include_top, + include_hm_decoder=include_hm_decoder, + head_conv=head_conv, + heads=heads, + use_c2=use_c2, + use_c3=use_c3, + use_c4=use_c4, + use_c51=use_c51, + efpn=efpn, + tfpn=tfpn, + se_layer=se_layer, + num_layers=num_layers, + norm_c2=norm_c2, + INIT_WEIGHTS=INIT_WEIGHTS, + ) + + return blocks_args, global_params + + +class BlockDecoder(object): + """Block Decoder for readability, + straight from the official TensorFlow repository. + """ + + @staticmethod + def _decode_block_string(block_string): + """Get a block through a string notation of arguments. + Args: + block_string (str): A string notation of arguments. + Examples: 'r1_k3_s11_e1_i32_o16_se0.25_noskip'. + Returns: + BlockArgs: The namedtuple defined at the top of this file. + """ + assert isinstance(block_string, str) + + ops = block_string.split("_") + options = {} + for op in ops: + splits = re.split(r"(\d.*)", op) + if len(splits) >= 2: + key, value = splits[:2] + options[key] = value + + # Check stride + assert ("s" in options and len(options["s"]) == 1) or ( + len(options["s"]) == 2 and options["s"][0] == options["s"][1] + ) + + return BlockArgs( + num_repeat=int(options["r"]), + kernel_size=int(options["k"]), + stride=[int(options["s"][0])], + expand_ratio=int(options["e"]), + input_filters=int(options["i"]), + output_filters=int(options["o"]), + se_ratio=float(options["se"]) if "se" in options else None, + id_skip=("noskip" not in block_string), + ) + + @staticmethod + def _encode_block_string(block): + """Encode a block to a string. + Args: + block (namedtuple): A BlockArgs type argument. + Returns: + block_string: A String form of BlockArgs. + """ + args = [ + "r%d" % block.num_repeat, + "k%d" % block.kernel_size, + "s%d%d" % (block.strides[0], block.strides[1]), + "e%s" % block.expand_ratio, + "i%d" % block.input_filters, + "o%d" % block.output_filters, + ] + if 0 < block.se_ratio <= 1: + args.append("se%s" % block.se_ratio) + if block.id_skip is False: + args.append("noskip") + return "_".join(args) + + @staticmethod + def decode(string_list): + """Decode a list of string notations to specify blocks inside the network. + Args: + string_list (list[str]): A list of strings, each string is a notation of block. + Returns: + blocks_args: A list of BlockArgs namedtuples of block args. + """ + assert isinstance(string_list, list) + blocks_args = [] + for block_string in string_list: + blocks_args.append(BlockDecoder._decode_block_string(block_string)) + return blocks_args + + @staticmethod + def encode(blocks_args): + """Encode a list of BlockArgs to a list of strings. + Args: + blocks_args (list[namedtuples]): A list of BlockArgs namedtuples of block args. + Returns: + block_strings: A list of strings, each string is a notation of block. + """ + block_strings = [] + for block in blocks_args: + block_strings.append(BlockDecoder._encode_block_string(block)) + return block_strings + + +class SwishImplementation(torch.autograd.Function): + @staticmethod + def forward(ctx, i): + result = i * torch.sigmoid(i) + ctx.save_for_backward(i) + return result + + @staticmethod + def backward(ctx, grad_output): + i = ctx.saved_tensors[0] + sigmoid_i = torch.sigmoid(i) + return grad_output * (sigmoid_i * (1 + i * (1 - sigmoid_i))) + + +def get_width_and_height_from_size(x): + """Obtain height and width from x. + Args: + x (int, tuple or list): Data size. + Returns: + size: A tuple or list (H,W). + """ + if isinstance(x, int): + return x, x + if isinstance(x, list) or isinstance(x, tuple): + return x + else: + raise TypeError() + + +def calculate_output_image_size(input_image_size, stride): + """Calculates the output image size when using Conv2dSamePadding with a stride. + Necessary for static padding. Thanks to mannatsingh for pointing this out. + Args: + input_image_size (int, tuple or list): Size of input image. + stride (int, tuple or list): Conv2d operation's stride. + Returns: + output_image_size: A list [H,W]. + """ + if input_image_size is None: + return None + image_height, image_width = get_width_and_height_from_size(input_image_size) + stride = stride if isinstance(stride, int) else stride[0] + image_height = int(math.ceil(image_height / stride)) + image_width = int(math.ceil(image_width / stride)) + return [image_height, image_width] + + +class MemoryEfficientSwish(nn.Module): + def forward(self, x): + return SwishImplementation.apply(x) + + +url_map_advprop = { + "efficientnet-b0": "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b0-b64d5a18.pth", + "efficientnet-b1": "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b1-0f3ce85a.pth", + "efficientnet-b2": "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b2-6e9d97e5.pth", + "efficientnet-b3": "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b3-cdd7c0f4.pth", + "efficientnet-b4": "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b4-44fb3a87.pth", + "efficientnet-b5": "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b5-86493f6b.pth", + "efficientnet-b6": "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b6-ac80338e.pth", + "efficientnet-b7": "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b7-4652b6dd.pth", + "efficientnet-b8": "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b8-22a8fe65.pth", +} + + +url_map = { + "efficientnet-b0": "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b0-355c32eb.pth", + "efficientnet-b1": "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b1-f1951068.pth", + "efficientnet-b2": "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b2-8bb594d6.pth", + "efficientnet-b3": "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b3-5fb5a3c3.pth", + "efficientnet-b4": "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b4-6ed6700e.pth", + "efficientnet-b5": "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b5-b6417697.pth", + "efficientnet-b6": "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b6-c76e70fd.pth", + "efficientnet-b7": "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b7-dcc49843.pth", +} + + +def load_pretrained_weights( + model, model_name, weights_path=None, load_fc=True, advprop=False, verbose=True +): + """Loads pretrained weights from weights path or download using url. + Args: + model (Module): The whole model of efficientnet. + model_name (str): Model name of efficientnet. + weights_path (None or str): + str: path to pretrained weights file on the local disk. + None: use pretrained weights downloaded from the Internet. + load_fc (bool): Whether to load pretrained weights for fc layer at the end of the model. + advprop (bool): Whether to load pretrained weights + trained with advprop (valid when weights_path is None). + """ + if isinstance(weights_path, str): + state_dict = torch.load(weights_path, map_location=torch.device("cpu")) + else: + # AutoAugment or Advprop (different preprocessing) + url_map_ = url_map_advprop if advprop else url_map + state_dict = model_zoo.load_url(url_map_[model_name]) + + if load_fc: + ret = model.load_state_dict(state_dict, strict=False) + assert ( + not ret.missing_keys + ), "Missing keys when loading pretrained weights: {}".format(ret.missing_keys) + else: + state_dict.pop("_fc.weight") + state_dict.pop("_fc.bias") + ret = model.load_state_dict(state_dict, strict=False) + + # if len(ret.missing_keys): + # assert set(ret.missing_keys) == set( + # ['_fc.weight', '_fc.bias']), 'Missing keys when loading pretrained weights: {}'.format(ret.missing_keys) + assert ( + not ret.unexpected_keys + ), "Missing keys when loading pretrained weights: {}".format(ret.unexpected_keys) + + if verbose: + print("Loaded pretrained weights for {}".format(model_name)) diff --git a/video/fake-stormer/model_code/models/networks/backbones/resnet3d.py b/video/fake-stormer/model_code/models/networks/backbones/resnet3d.py new file mode 100644 index 0000000000000000000000000000000000000000..7bb9055de832389eff045026729d1f0f152995e6 --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/backbones/resnet3d.py @@ -0,0 +1,296 @@ +# -*- coding:utf-8 -*- +import os +import sys + +if not (os.getcwd()) in sys.path: + sys.path.append(os.getcwd()) +import math +from functools import partial + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...builder import BACKBONES +from .base import BaseBackbone + + +def get_inplanes(): + return [64, 128, 256, 512] + + +def conv3x3x3(in_planes, out_planes, stride=1): + return nn.Conv3d( + in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False + ) + + +def conv1x1x1(in_planes, out_planes, stride=1): + return nn.Conv3d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) + + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, in_planes, planes, stride=1, downsample=None): + super().__init__() + + self.conv1 = conv3x3x3(in_planes, planes, stride) + self.bn1 = nn.BatchNorm3d(planes) + self.relu = nn.ReLU(inplace=True) + self.conv2 = conv3x3x3(planes, planes) + self.bn2 = nn.BatchNorm3d(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 + + @staticmethod + def __repr__(): + return "BasicBlock" + + +class Bottleneck(nn.Module): + expansion = 4 + + def __init__(self, in_planes, planes, stride=1, downsample=None): + super().__init__() + + self.conv1 = conv1x1x1(in_planes, planes) + self.bn1 = nn.BatchNorm3d(planes) + self.conv2 = conv3x3x3(planes, planes, stride) + self.bn2 = nn.BatchNorm3d(planes) + self.conv3 = conv1x1x1(planes, planes * self.expansion) + self.bn3 = nn.BatchNorm3d(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 + + @staticmethod + def __repr__(): + return "Bottleneck" + + +@BACKBONES.register_module() +class ResNet3D(BaseBackbone): + def __init__( + self, + block, + layers, + block_inplanes, + n_input_channels=3, + conv1_t_size=7, + conv1_t_stride=1, + no_max_pool=False, + shortcut_type="B", + widen_factor=1.0, + n_classes=400, + do_cls=False, + dropout_rate=0.25, + ): + super().__init__() + + # Convert Cls name into Cls Object + if isinstance(block, str): + for bl in [BasicBlock, Bottleneck]: + if block == bl.__repr__(): + block = bl + + block_inplanes = [int(x * widen_factor) for x in block_inplanes] + + self.in_planes = block_inplanes[0] + self.no_max_pool = no_max_pool + + self.conv1 = nn.Conv3d( + n_input_channels, + self.in_planes, + kernel_size=(conv1_t_size, 7, 7), + stride=(conv1_t_stride, 2, 2), + padding=(conv1_t_size // 2, 3, 3), + bias=False, + ) + self.bn1 = nn.BatchNorm3d(self.in_planes) + self.relu = nn.ReLU(inplace=True) + self.maxpool = nn.MaxPool3d(kernel_size=3, stride=2, padding=1) + self.layer1 = self._make_layer( + block, block_inplanes[0], layers[0], shortcut_type + ) + self.drop_layer1 = nn.Dropout(dropout_rate) + self.layer2 = self._make_layer( + block, block_inplanes[1], layers[1], shortcut_type, stride=2 + ) + self.drop_layer2 = nn.Dropout(dropout_rate) + self.layer3 = self._make_layer( + block, block_inplanes[2], layers[2], shortcut_type, stride=2 + ) + self.drop_layer3 = nn.Dropout(dropout_rate) + self.layer4 = self._make_layer( + block, block_inplanes[3], layers[3], shortcut_type, stride=2 + ) + self.drop_layer4 = nn.Dropout(dropout_rate) + + self.avgpool = nn.AdaptiveAvgPool3d((1, 1, 1)) + self.do_cls = do_cls + if self.do_cls: + self.fc = nn.Linear(block_inplanes[3] * block.expansion, n_classes) + + def _downsample_basic_block(self, x, planes, stride): + out = F.avg_pool3d(x, kernel_size=1, stride=stride) + zero_pads = torch.zeros( + out.size(0), planes - out.size(1), out.size(2), out.size(3), out.size(4) + ) + if isinstance(out.data, torch.cuda.FloatTensor): + zero_pads = zero_pads.cuda() + + out = torch.cat([out.data, zero_pads], dim=1) + + return out + + def _make_layer(self, block, planes, blocks, shortcut_type, stride=1): + downsample = None + if stride != 1 or self.in_planes != planes * block.expansion: + if shortcut_type == "A": + downsample = partial( + self._downsample_basic_block, + planes=planes * block.expansion, + stride=stride, + ) + else: + downsample = nn.Sequential( + conv1x1x1(self.in_planes, planes * block.expansion, stride), + nn.BatchNorm3d(planes * block.expansion), + ) + + layers = [] + layers.append( + block( + in_planes=self.in_planes, + planes=planes, + stride=stride, + downsample=downsample, + ) + ) + self.in_planes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.in_planes, planes)) + + return nn.Sequential(*layers) + + def forward(self, x): + outputs = [] + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + if not self.no_max_pool: + x = self.maxpool(x) + + x1 = self.layer1(x) # 256x16x56x56 + x2 = self.layer2(x1) # 512x8x28x28 + x3 = self.layer3(x2) # 1024x4x14x14 + x4 = self.layer4(x3) # 2048x2x7x7 + + if self.do_cls: + x_avg = self.avgpool(x4) + x = x_avg.view(x_avg.size(0), -1).unsqueeze(1) + x = self.fc(x) + + res = {} + res["embed"] = x4 + + x1 = self.drop_layer1(x1) + outputs.append(x1) + x2 = self.drop_layer2(x2) + outputs.append(x2) + x3 = self.drop_layer3(x3) + outputs.append(x3) + outputs.append(x4) + res["outputs"] = outputs + + return res + + def init_weights(self, pretrained=None): + if pretrained is not None: + super().init_weights(pretrained=pretrained) + else: + for m in self.modules(): + if isinstance(m, nn.Conv3d): + nn.init.kaiming_normal_( + m.weight, mode="fan_out", nonlinearity="relu" + ) + elif isinstance(m, nn.BatchNorm3d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + +def generate_model(model_depth, **kwargs): + assert model_depth in [10, 18, 34, 50, 101, 152, 200] + + if model_depth == 10: + model = ResNet3D(BasicBlock, [1, 1, 1, 1], get_inplanes(), **kwargs) + elif model_depth == 18: + model = ResNet3D(BasicBlock, [2, 2, 2, 2], get_inplanes(), **kwargs) + elif model_depth == 34: + model = ResNet3D(BasicBlock, [3, 4, 6, 3], get_inplanes(), **kwargs) + elif model_depth == 50: + model = ResNet3D(Bottleneck, [3, 4, 6, 3], get_inplanes(), **kwargs) + elif model_depth == 101: + model = ResNet3D(Bottleneck, [3, 4, 23, 3], get_inplanes(), **kwargs) + elif model_depth == 152: + model = ResNet3D(Bottleneck, [3, 8, 36, 3], get_inplanes(), **kwargs) + elif model_depth == 200: + model = ResNet3D(Bottleneck, [3, 24, 36, 3], get_inplanes(), **kwargs) + + return model + + +if __name__ == "__main__": + cfg = { + "type": "ResNet3D", + "block": Bottleneck, + "layers": [3, 4, 6, 3], + "block_inplanes": [64, 128, 256, 512], + } + net = BACKBONES.build(cfg=cfg, default_args=cfg) + input = torch.rand(1, 3, 32, 224, 224) + res = net(input) + print(res["embed"].shape) + for i in range(len(res["outputs"])): + print(f"Layer {i+1}", res["outputs"][i].shape) diff --git a/video/fake-stormer/model_code/models/networks/backbones/swin.py b/video/fake-stormer/model_code/models/networks/backbones/swin.py new file mode 100644 index 0000000000000000000000000000000000000000..9f0d18cd6b593e99e3d39f684d74c9d83fb46ac9 --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/backbones/swin.py @@ -0,0 +1,1151 @@ +# -*- coding: utf-8 -*- +# Copyright (c) OpenMMLab. All rights reserved. +import math +import os +import sys +from collections import OrderedDict +from copy import deepcopy +from typing import Sequence + +if not os.getcwd() in sys.path: + sys.path.append(os.getcwd()) + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as cp +from einops import rearrange, reduce, repeat +from logs.logger import get_root_logger +from mmcv.cnn import ( + build_conv_layer, + build_norm_layer, + constant_init, + trunc_normal_init, +) +from mmcv.cnn.bricks.transformer import FFN, build_dropout +from mmcv.cnn.utils.weight_init import trunc_normal_ +from mmcv.runner import _load_checkpoint +from mmcv.runner.base_module import BaseModule +from mmcv.utils import to_2tuple +from models.builder import BACKBONES +from models.utils import swin_converter + +from .base import BaseBackbone + + +class AdaptivePadding(nn.Module): + """Applies padding to input (if needed) so that input can get fully covered + by filter you specified. It support two modes "same" and "corner". The + "same" mode is same with "SAME" padding mode in TensorFlow, pad zero around + input. The "corner" mode would pad zero to bottom right. + + Args: + kernel_size (int | tuple): Size of the kernel: + stride (int | tuple): Stride of the filter. Default: 1: + dilation (int | tuple): Spacing between kernel elements. + Default: 1 + padding (str): Support "same" and "corner", "corner" mode + would pad zero to bottom right, and "same" mode would + pad zero around input. Default: "corner". + Example: + >>> kernel_size = 16 + >>> stride = 16 + >>> dilation = 1 + >>> input = torch.rand(1, 1, 15, 17) + >>> adap_pad = AdaptivePadding( + >>> kernel_size=kernel_size, + >>> stride=stride, + >>> dilation=dilation, + >>> padding="corner") + >>> out = adap_pad(input) + >>> assert (out.shape[2], out.shape[3]) == (16, 32) + >>> input = torch.rand(1, 1, 16, 17) + >>> out = adap_pad(input) + >>> assert (out.shape[2], out.shape[3]) == (16, 32) + """ + + def __init__(self, kernel_size=1, stride=1, dilation=1, padding="corner"): + + super(AdaptivePadding, self).__init__() + + assert padding in ("same", "corner") + + kernel_size = to_2tuple(kernel_size) + stride = to_2tuple(stride) + padding = to_2tuple(padding) + dilation = to_2tuple(dilation) + + self.padding = padding + self.kernel_size = kernel_size + self.stride = stride + self.dilation = dilation + + def get_pad_shape(self, input_shape): + input_h, input_w = input_shape + kernel_h, kernel_w = self.kernel_size + stride_h, stride_w = self.stride + output_h = math.ceil(input_h / stride_h) + output_w = math.ceil(input_w / stride_w) + pad_h = max( + (output_h - 1) * stride_h + (kernel_h - 1) * self.dilation[0] + 1 - input_h, + 0, + ) + pad_w = max( + (output_w - 1) * stride_w + (kernel_w - 1) * self.dilation[1] + 1 - input_w, + 0, + ) + return pad_h, pad_w + + def forward(self, x): + pad_h, pad_w = self.get_pad_shape(x.size()[-2:]) + if pad_h > 0 or pad_w > 0: + if self.padding == "corner": + x = F.pad(x, [0, pad_w, 0, pad_h]) + elif self.padding == "same": + x = F.pad( + x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2] + ) + return x + + +class PatchEmbed(BaseModule): + """Image to Patch Embedding. + + We use a conv layer to implement PatchEmbed. + + Args: + in_channels (int): The num of input channels. Default: 3 + embed_dims (int): The dimensions of embedding. Default: 768 + conv_type (str): The config dict for embedding + conv layer type selection. Default: "Conv2d. + kernel_size (int): The kernel_size of embedding conv. Default: 16. + stride (int): The slide stride of embedding conv. + Default: None (Would be set as `kernel_size`). + padding (int | tuple | string ): The padding length of + embedding conv. When it is a string, it means the mode + of adaptive padding, support "same" and "corner" now. + Default: "corner". + dilation (int): The dilation rate of embedding conv. Default: 1. + bias (bool): Bias of embed conv. Default: True. + norm_cfg (dict, optional): Config dict for normalization layer. + Default: None. + input_size (int | tuple | None): The size of input, which will be + used to calculate the out size. Only work when `dynamic_size` + is False. Default: None. + init_cfg (`mmcv.ConfigDict`, optional): The Config for initialization. + Default: None. + """ + + def __init__( + self, + in_channels=3, + embed_dims=768, + conv_type="Conv2d", + kernel_size=16, + stride=16, + padding="corner", + dilation=1, + bias=True, + norm_cfg=None, + input_size=None, + init_cfg=None, + ): + super(PatchEmbed, self).__init__(init_cfg=init_cfg) + + self.embed_dims = embed_dims + if stride is None: + stride = kernel_size + + kernel_size = to_2tuple(kernel_size) + stride = to_2tuple(stride) + dilation = to_2tuple(dilation) + + if isinstance(padding, str): + self.adap_padding = AdaptivePadding( + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + padding=padding, + ) + # disable the padding of conv + padding = 0 + else: + self.adap_padding = None + padding = to_2tuple(padding) + + self.projection = build_conv_layer( + dict(type=conv_type), + in_channels=in_channels, + out_channels=embed_dims, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + ) + + if norm_cfg is not None: + self.norm = build_norm_layer(norm_cfg, embed_dims)[1] + else: + self.norm = None + + if input_size: + input_size = to_2tuple(input_size) + # `init_out_size` would be used outside to + # calculate the num_patches + # when `use_abs_pos_embed` outside + self.init_input_size = input_size + if self.adap_padding: + pad_h, pad_w = self.adap_padding.get_pad_shape(input_size) + input_h, input_w = input_size + input_h = input_h + pad_h + input_w = input_w + pad_w + input_size = (input_h, input_w) + + # https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html + h_out = ( + input_size[0] + 2 * padding[0] - dilation[0] * (kernel_size[0] - 1) - 1 + ) // stride[0] + 1 + w_out = ( + input_size[1] + 2 * padding[1] - dilation[1] * (kernel_size[1] - 1) - 1 + ) // stride[1] + 1 + self.init_out_size = (h_out, w_out) + else: + self.init_input_size = None + self.init_out_size = None + + def forward(self, x): + """ + Args: + x (Tensor): Has shape (B, C, H, W). In most case, C is 3. + + Returns: + tuple: Contains merged results and its spatial shape. + + - x (Tensor): Has shape (B, out_h * out_w, embed_dims) + - out_size (tuple[int]): Spatial shape of x, arrange as + (out_h, out_w). + """ + + if self.adap_padding: + x = self.adap_padding(x) + + x = self.projection(x) + out_size = (x.shape[2], x.shape[3]) + x = x.flatten(2).transpose(1, 2) + if self.norm is not None: + x = self.norm(x) + return x, out_size + + +class PatchMerging(BaseModule): + """Merge patch feature map. + + This layer groups feature map by kernel_size, and applies norm and linear + layers to the grouped feature map. Our implementation uses `nn.Unfold` to + merge patch, which is about 25% faster than original implementation. + Instead, we need to modify pretrained models for compatibility. + + Args: + in_channels (int): The num of input channels. + to gets fully covered by filter and stride you specified.. + Default: True. + out_channels (int): The num of output channels. + kernel_size (int | tuple, optional): the kernel size in the unfold + layer. Defaults to 2. + stride (int | tuple, optional): the stride of the sliding blocks in the + unfold layer. Default: None. (Would be set as `kernel_size`) + padding (int | tuple | string ): The padding length of + embedding conv. When it is a string, it means the mode + of adaptive padding, support "same" and "corner" now. + Default: "corner". + dilation (int | tuple, optional): dilation parameter in the unfold + layer. Default: 1. + bias (bool, optional): Whether to add bias in linear layer or not. + Defaults: False. + norm_cfg (dict, optional): Config dict for normalization layer. + Default: dict(type='LN'). + init_cfg (dict, optional): The extra config for initialization. + Default: None. + """ + + def __init__( + self, + in_channels, + out_channels, + kernel_size=2, + stride=None, + padding="corner", + dilation=1, + bias=False, + norm_cfg=dict(type="LN"), + init_cfg=None, + ): + super().__init__(init_cfg=init_cfg) + self.in_channels = in_channels + self.out_channels = out_channels + if stride: + stride = stride + else: + stride = kernel_size + + kernel_size = to_2tuple(kernel_size) + stride = to_2tuple(stride) + dilation = to_2tuple(dilation) + + if isinstance(padding, str): + self.adap_padding = AdaptivePadding( + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + padding=padding, + ) + # disable the padding of unfold + padding = 0 + else: + self.adap_padding = None + + padding = to_2tuple(padding) + self.sampler = nn.Unfold( + kernel_size=kernel_size, dilation=dilation, padding=padding, stride=stride + ) + + sample_dim = kernel_size[0] * kernel_size[1] * in_channels + + if norm_cfg is not None: + self.norm = build_norm_layer(norm_cfg, sample_dim)[1] + else: + self.norm = None + + self.reduction = nn.Linear(sample_dim, out_channels, bias=bias) + + def forward(self, x, input_size): + """ + Args: + x (Tensor): Has shape (B, H*W, C_in). + input_size (tuple[int]): The spatial shape of x, arrange as (H, W). + Default: None. + + Returns: + tuple: Contains merged results and its spatial shape. + + - x (Tensor): Has shape (B, Merged_H * Merged_W, C_out) + - out_size (tuple[int]): Spatial shape of x, arrange as + (Merged_H, Merged_W). + """ + B, L, C = x.shape + assert isinstance(input_size, Sequence), ( + f"Expect " f"input_size is " f"`Sequence` " f"but get {input_size}" + ) + + H, W = input_size + assert L == H * W, "input feature has wrong size" + + x = x.view(B, H, W, C).permute([0, 3, 1, 2]) # B, C, H, W + # Use nn.Unfold to merge patch. About 25% faster than original method, + # but need to modify pretrained model for compatibility + + if self.adap_padding: + x = self.adap_padding(x) + H, W = x.shape[-2:] + + x = self.sampler(x) + # if kernel_size=2 and stride=2, x should has shape (B, 4*C, H/2*W/2) + + out_h = ( + H + + 2 * self.sampler.padding[0] + - self.sampler.dilation[0] * (self.sampler.kernel_size[0] - 1) + - 1 + ) // self.sampler.stride[0] + 1 + out_w = ( + W + + 2 * self.sampler.padding[1] + - self.sampler.dilation[1] * (self.sampler.kernel_size[1] - 1) + - 1 + ) // self.sampler.stride[1] + 1 + + output_size = (out_h, out_w) + x = x.transpose(1, 2) # B, H/2*W/2, 4*C + x = self.norm(x) if self.norm else x + x = self.reduction(x) + return x, output_size + + +class WindowMSA(nn.Module): + """Window based multi-head self-attention (W-MSA) module with relative + position bias. + + Args: + embed_dims (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (tuple[int]): The height and width of the window. + qkv_bias (bool, optional): If True, add a learnable bias to q, k, v. + Default: True. + qk_scale (float | None, optional): Override default qk scale of + head_dim ** -0.5 if set. Default: None. + attn_drop_rate (float, optional): Dropout ratio of attention weight. + Default: 0.0 + proj_drop_rate (float, optional): Dropout ratio of output. Default: 0. + """ + + def __init__( + self, + embed_dims, + num_heads, + window_size, + qkv_bias=True, + qk_scale=None, + attn_drop_rate=0.0, + proj_drop_rate=0.0, + ): + + super().__init__() + self.embed_dims = embed_dims + self.window_size = window_size # Wh, Ww + self.num_heads = num_heads + head_embed_dims = embed_dims // num_heads + self.scale = qk_scale or head_embed_dims**-0.5 + + # define a parameter table of relative position bias + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads) + ) # 2*Wh-1 * 2*Ww-1, nH + + # About 2x faster than original impl + Wh, Ww = self.window_size + rel_index_coords = self.double_step_seq(2 * Ww - 1, Wh, 1, Ww) + rel_position_index = rel_index_coords + rel_index_coords.T + rel_position_index = rel_position_index.flip(1).contiguous() + self.register_buffer("relative_position_index", rel_position_index) + + self.qkv = nn.Linear(embed_dims, embed_dims * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop_rate) + self.proj = nn.Linear(embed_dims, embed_dims) + self.proj_drop = nn.Dropout(proj_drop_rate) + + self.softmax = nn.Softmax(dim=-1) + + def init_weights(self): + trunc_normal_(self.relative_position_bias_table, std=0.02) + + def forward(self, x, mask=None): + """ + Args: + + x (tensor): input features with shape of (num_windows*B, N, C) + mask (tensor | None, Optional): mask with shape of (num_windows, + Wh*Ww, Wh*Ww), value should be between (-inf, 0]. + """ + B, N, C = x.shape + qkv = ( + self.qkv(x) + .reshape(B, N, 3, self.num_heads, C // self.num_heads) + .permute(2, 0, 3, 1, 4) + ) + # make torchscript happy (cannot use tensor as tuple) + q, k, v = qkv[0], qkv[1], qkv[2] + + q = q * self.scale + attn = q @ k.transpose(-2, -1) + + relative_position_bias = self.relative_position_bias_table[ + self.relative_position_index.view(-1) + ].view( + self.window_size[0] * self.window_size[1], + self.window_size[0] * self.window_size[1], + -1, + ) # Wh*Ww,Wh*Ww,nH + relative_position_bias = relative_position_bias.permute( + 2, 0, 1 + ).contiguous() # nH, Wh*Ww, Wh*Ww + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B // nW, nW, self.num_heads, N, N) + mask.unsqueeze( + 1 + ).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + attn = self.softmax(attn) + + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + @staticmethod + def double_step_seq(step1, len1, step2, len2): + seq1 = torch.arange(0, step1 * len1, step1) + seq2 = torch.arange(0, step2 * len2, step2) + return (seq1[:, None] + seq2[None, :]).reshape(1, -1) + + +class ShiftWindowMSA(nn.Module): + """Shifted Window Multihead Self-Attention Module. + + Args: + embed_dims (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (int): The height and width of the window. + shift_size (int, optional): The shift step of each window towards + right-bottom. If zero, act as regular window-msa. Defaults to 0. + qkv_bias (bool, optional): If True, add a learnable bias to q, k, v. + Default: True + qk_scale (float | None, optional): Override default qk scale of + head_dim ** -0.5 if set. Defaults: None. + attn_drop_rate (float, optional): Dropout ratio of attention weight. + Defaults: 0. + proj_drop_rate (float, optional): Dropout ratio of output. + Defaults: 0. + dropout_layer (dict, optional): The dropout_layer used before output. + Defaults: dict(type='DropPath', drop_prob=0.). + """ + + def __init__( + self, + embed_dims, + num_heads, + window_size, + shift_size=0, + qkv_bias=True, + qk_scale=None, + attn_drop_rate=0, + proj_drop_rate=0, + dropout_layer=dict(type="DropPath", drop_prob=0.0), + ): + super().__init__() + + self.window_size = window_size + self.shift_size = shift_size + assert 0 <= self.shift_size < self.window_size + + self.w_msa = WindowMSA( + embed_dims=embed_dims, + num_heads=num_heads, + window_size=to_2tuple(window_size), + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop_rate=attn_drop_rate, + proj_drop_rate=proj_drop_rate, + ) + + self.drop = build_dropout(dropout_layer) + + def forward(self, query, hw_shape): + B, L, C = query.shape + H, W = hw_shape + assert L == H * W, "input feature has wrong size" + query = query.view(B, H, W, C) + + # pad feature maps to multiples of window size + pad_r = (self.window_size - W % self.window_size) % self.window_size + pad_b = (self.window_size - H % self.window_size) % self.window_size + query = F.pad(query, (0, 0, 0, pad_r, 0, pad_b)) + H_pad, W_pad = query.shape[1], query.shape[2] + + # cyclic shift + if self.shift_size > 0: + shifted_query = torch.roll( + query, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2) + ) + + # calculate attention mask for SW-MSA + img_mask = torch.zeros((1, H_pad, W_pad, 1), device=query.device) + h_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) + w_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + # nW, window_size, window_size, 1 + mask_windows = self.window_partition(img_mask) + mask_windows = mask_windows.view(-1, self.window_size * self.window_size) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill( + attn_mask != 0, float(-100.0) + ).masked_fill(attn_mask == 0, float(0.0)) + else: + shifted_query = query + attn_mask = None + + # nW*B, window_size, window_size, C + query_windows = self.window_partition(shifted_query) + # nW*B, window_size*window_size, C + query_windows = query_windows.view(-1, self.window_size**2, C) + + # W-MSA/SW-MSA (nW*B, window_size*window_size, C) + attn_windows = self.w_msa(query_windows, mask=attn_mask) + + # merge windows + attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) + + # B H' W' C + shifted_x = self.window_reverse(attn_windows, H_pad, W_pad) + # reverse cyclic shift + if self.shift_size > 0: + x = torch.roll( + shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2) + ) + else: + x = shifted_x + + if pad_r > 0 or pad_b: + x = x[:, :H, :W, :].contiguous() + + x = x.view(B, H * W, C) + + x = self.drop(x) + return x + + def window_reverse(self, windows, H, W): + """ + Args: + windows: (num_windows*B, window_size, window_size, C) + H (int): Height of image + W (int): Width of image + Returns: + x: (B, H, W, C) + """ + window_size = self.window_size + B = int(windows.shape[0] / (H * W / window_size / window_size)) + x = windows.view( + B, H // window_size, W // window_size, window_size, window_size, -1 + ) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) + return x + + def window_partition(self, x): + """ + Args: + x: (B, H, W, C) + Returns: + windows: (num_windows*B, window_size, window_size, C) + """ + B, H, W, C = x.shape + window_size = self.window_size + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous() + windows = windows.view(-1, window_size, window_size, C) + return windows + + +class SwinBlock(nn.Module): + """ " + Args: + embed_dims (int): The feature dimension. + num_heads (int): Parallel attention heads. + feedforward_channels (int): The hidden dimension for FFNs. + window_size (int, optional): The local window scale. Default: 7. + shift (bool, optional): whether to shift window or not. Default False. + qkv_bias (bool, optional): enable bias for qkv if True. Default: True. + qk_scale (float | None, optional): Override default qk scale of + head_dim ** -0.5 if set. Default: None. + drop_rate (float, optional): Dropout rate. Default: 0. + attn_drop_rate (float, optional): Attention dropout rate. Default: 0. + drop_path_rate (float, optional): Stochastic depth rate. Default: 0. + act_cfg (dict, optional): The config dict of activation function. + Default: dict(type='GELU'). + norm_cfg (dict, optional): The config dict of normalization. + Default: dict(type='LN'). + with_cp (bool, optional): Use checkpoint or not. Using checkpoint + will save some memory while slowing down the training speed. + Default: False. + """ + + def __init__( + self, + embed_dims, + num_heads, + feedforward_channels, + window_size=7, + shift=False, + qkv_bias=True, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.0, + act_cfg=dict(type="GELU"), + norm_cfg=dict(type="LN"), + with_cp=False, + ): + + super(SwinBlock, self).__init__() + + self.with_cp = with_cp + + self.norm1 = build_norm_layer(norm_cfg, embed_dims)[1] + self.attn = ShiftWindowMSA( + embed_dims=embed_dims, + num_heads=num_heads, + window_size=window_size, + shift_size=window_size // 2 if shift else 0, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop_rate=attn_drop_rate, + proj_drop_rate=drop_rate, + dropout_layer=dict(type="DropPath", drop_prob=drop_path_rate), + ) + + self.norm2 = build_norm_layer(norm_cfg, embed_dims)[1] + self.ffn = FFN( + embed_dims=embed_dims, + feedforward_channels=feedforward_channels, + num_fcs=2, + ffn_drop=drop_rate, + dropout_layer=dict(type="DropPath", drop_prob=drop_path_rate), + act_cfg=act_cfg, + add_identity=True, + init_cfg=None, + ) + + def forward(self, x, hw_shape): + + def _inner_forward(x): + identity = x + x = self.norm1(x) + x = self.attn(x, hw_shape) + + x = x + identity + + identity = x + x = self.norm2(x) + x = self.ffn(x, identity=identity) + + return x + + if self.with_cp and x.requires_grad: + x = cp.checkpoint(_inner_forward, x) + else: + x = _inner_forward(x) + + return x + + +class SwinBlockSequence(nn.Module): + """Implements one stage in Swin Transformer. + + Args: + embed_dims (int): The feature dimension. + num_heads (int): Parallel attention heads. + feedforward_channels (int): The hidden dimension for FFNs. + depth (int): The number of blocks in this stage. + window_size (int, optional): The local window scale. Default: 7. + qkv_bias (bool, optional): enable bias for qkv if True. Default: True. + qk_scale (float | None, optional): Override default qk scale of + head_dim ** -0.5 if set. Default: None. + drop_rate (float, optional): Dropout rate. Default: 0. + attn_drop_rate (float, optional): Attention dropout rate. Default: 0. + drop_path_rate (float | list[float], optional): Stochastic depth + rate. Default: 0. + downsample (nn.Module | None, optional): The downsample operation + module. Default: None. + act_cfg (dict, optional): The config dict of activation function. + Default: dict(type='GELU'). + norm_cfg (dict, optional): The config dict of normalization. + Default: dict(type='LN'). + with_cp (bool, optional): Use checkpoint or not. Using checkpoint + will save some memory while slowing down the training speed. + Default: False. + """ + + def __init__( + self, + embed_dims, + num_heads, + feedforward_channels, + depth, + window_size=7, + qkv_bias=True, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.0, + downsample=None, + act_cfg=dict(type="GELU"), + norm_cfg=dict(type="LN"), + with_cp=False, + ): + super().__init__() + + if isinstance(drop_path_rate, list): + drop_path_rates = drop_path_rate + assert len(drop_path_rates) == depth + else: + drop_path_rates = [deepcopy(drop_path_rate) for _ in range(depth)] + + self.blocks = nn.ModuleList() + for i in range(depth): + block = SwinBlock( + embed_dims=embed_dims, + num_heads=num_heads, + feedforward_channels=feedforward_channels, + window_size=window_size, + shift=False if i % 2 == 0 else True, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop_rate=drop_rate, + attn_drop_rate=attn_drop_rate, + drop_path_rate=drop_path_rates[i], + act_cfg=act_cfg, + norm_cfg=norm_cfg, + with_cp=with_cp, + ) + self.blocks.append(block) + + self.downsample = downsample + + def forward(self, x, hw_shape): + for block in self.blocks: + x = block(x, hw_shape) + + if self.downsample: + x_down, down_hw_shape = self.downsample(x, hw_shape) + return x_down, down_hw_shape, x, hw_shape + else: + return x, hw_shape, x, hw_shape + + +@BACKBONES.register_module() +class SwinTransformer(BaseBackbone): + """Swin Transformer + A PyTorch implement of : `Swin Transformer: + Hierarchical Vision Transformer using Shifted Windows` - + https://arxiv.org/abs/2103.14030 + + Inspiration from + https://github.com/microsoft/Swin-Transformer + + Args: + pretrain_img_size (int | tuple[int]): The size of input image when + pretrain. Defaults: 224. + in_channels (int): The num of input channels. + Defaults: 3. + embed_dims (int): The feature dimension. Default: 96. + patch_size (int | tuple[int]): Patch size. Default: 4. + window_size (int): Window size. Default: 7. + mlp_ratio (int): Ratio of mlp hidden dim to embedding dim. + Default: 4. + depths (tuple[int]): Depths of each Swin Transformer stage. + Default: (2, 2, 6, 2). + num_heads (tuple[int]): Parallel attention heads of each Swin + Transformer stage. Default: (3, 6, 12, 24). + strides (tuple[int]): The patch merging or patch embedding stride of + each Swin Transformer stage. (In swin, we set kernel size equal to + stride.) Default: (4, 2, 2, 2). + out_indices (tuple[int]): Output from which stages. + Default: (0, 1, 2, 3). + qkv_bias (bool, optional): If True, add a learnable bias to query, key, + value. Default: True + qk_scale (float | None, optional): Override default qk scale of + head_dim ** -0.5 if set. Default: None. + patch_norm (bool): If add a norm layer for patch embed and patch + merging. Default: True. + drop_rate (float): Dropout rate. Defaults: 0. + attn_drop_rate (float): Attention dropout rate. Default: 0. + drop_path_rate (float): Stochastic depth rate. Defaults: 0.1. + use_abs_pos_embed (bool): If True, add absolute position embedding to + the patch embedding. Defaults: False. + act_cfg (dict): Config dict for activation layer. + Default: dict(type='LN'). + norm_cfg (dict): Config dict for normalization layer at + output of backone. Defaults: dict(type='LN'). + with_cp (bool, optional): Use checkpoint or not. Using checkpoint + will save some memory while slowing down the training speed. + Default: False. + pretrained (str, optional): model pretrained path. Default: None. + convert_weights (bool): The flag indicates whether the + pre-trained model is from the original repo. We may need + to convert some keys to make it compatible. + Default: False. + frozen_stages (int): Stages to be frozen (stop grad and set eval mode). + Default: -1 (-1 means not freezing any parameters). + """ + + def __init__( + self, + pretrain_img_size=224, + in_channels=3, + embed_dims=96, + patch_size=4, + window_size=7, + mlp_ratio=4, + depths=(2, 2, 6, 2), + num_heads=(3, 6, 12, 24), + strides=(4, 2, 2, 2), + out_indices=(0, 1, 2, 3), + qkv_bias=True, + qk_scale=None, + patch_norm=True, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.1, + use_abs_pos_embed=False, + act_cfg=dict(type="GELU"), + norm_cfg=dict(type="LN"), + with_cp=False, + convert_weights=False, + frozen_stages=-1, + pretrained=None, + ): + self.convert_weights = convert_weights + self.pretrained = pretrained + self.frozen_stages = frozen_stages + if isinstance(pretrain_img_size, int): + pretrain_img_size = to_2tuple(pretrain_img_size) + elif isinstance(pretrain_img_size, tuple): + if len(pretrain_img_size) == 1: + pretrain_img_size = to_2tuple(pretrain_img_size[0]) + assert len(pretrain_img_size) == 2, ( + f"The size of image should have length 1 or 2, " + f"but got {len(pretrain_img_size)}" + ) + + super(SwinTransformer, self).__init__() + + num_layers = len(depths) + self.out_indices = out_indices + self.use_abs_pos_embed = use_abs_pos_embed + + assert strides[0] == patch_size, "Use non-overlapping patch embed." + + self.patch_embed = PatchEmbed( + in_channels=in_channels, + embed_dims=embed_dims, + conv_type="Conv2d", + kernel_size=patch_size, + stride=strides[0], + norm_cfg=norm_cfg if patch_norm else None, + init_cfg=None, + ) + + if self.use_abs_pos_embed: + patch_row = pretrain_img_size[0] // patch_size + patch_col = pretrain_img_size[1] // patch_size + num_patches = patch_row * patch_col + self.absolute_pos_embed = nn.Parameter( + torch.zeros((1, num_patches, embed_dims)) + ) + + self.drop_after_pos = nn.Dropout(p=drop_rate) + + # set stochastic depth decay rule + total_depth = sum(depths) + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, total_depth)] + + self.stages = nn.ModuleList() + in_channels = embed_dims + for i in range(num_layers): + if i < num_layers - 1: + downsample = PatchMerging( + in_channels=in_channels, + out_channels=2 * in_channels, + stride=strides[i + 1], + norm_cfg=norm_cfg if patch_norm else None, + init_cfg=None, + ) + else: + downsample = None + + stage = SwinBlockSequence( + embed_dims=in_channels, + num_heads=num_heads[i], + feedforward_channels=mlp_ratio * in_channels, + depth=depths[i], + window_size=window_size, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop_rate=drop_rate, + attn_drop_rate=attn_drop_rate, + drop_path_rate=dpr[sum(depths[:i]) : sum(depths[: i + 1])], + downsample=downsample, + act_cfg=act_cfg, + norm_cfg=norm_cfg, + with_cp=with_cp, + ) + self.stages.append(stage) + if downsample: + in_channels = downsample.out_channels + + self.num_features = [int(embed_dims * 2**i) for i in range(num_layers)] + # Add a norm layer for each output + for i in out_indices: + layer = build_norm_layer(norm_cfg, self.num_features[i])[1] + layer_name = f"norm{i}" + self.add_module(layer_name, layer) + + def train(self, mode=True): + """Convert the model into training mode while keep layers freezed.""" + super(SwinTransformer, self).train(mode) + self._freeze_stages() + + def _freeze_stages(self): + if self.frozen_stages >= 0: + self.patch_embed.eval() + for param in self.patch_embed.parameters(): + param.requires_grad = False + if self.use_abs_pos_embed: + self.absolute_pos_embed.requires_grad = False + self.drop_after_pos.eval() + + for i in range(1, self.frozen_stages + 1): + + if (i - 1) in self.out_indices: + norm_layer = getattr(self, f"norm{i-1}") + norm_layer.eval() + for param in norm_layer.parameters(): + param.requires_grad = False + + m = self.stages[i - 1] + m.eval() + for param in m.parameters(): + param.requires_grad = False + + def init_weights(self, pretrained=None): + """Initialize the weights in backbone. + + Args: + pretrained (str, optional): Path to pre-trained weights. + Defaults to None. + """ + if isinstance(pretrained, str): + logger = get_root_logger() + ckpt = _load_checkpoint(pretrained, logger=None, map_location="cpu") + if "state_dict" in ckpt: + _state_dict = ckpt["state_dict"] + elif "model" in ckpt: + _state_dict = ckpt["model"] + else: + _state_dict = ckpt + if self.convert_weights: + # supported loading weight from original repo, + _state_dict = swin_converter(_state_dict) + + state_dict = OrderedDict() + for k, v in _state_dict.items(): + if k.startswith("backbone."): + state_dict[k[9:]] = v + + # strip prefix of state_dict + if list(state_dict.keys())[0].startswith("module."): + state_dict = {k[7:]: v for k, v in state_dict.items()} + + # reshape absolute position embedding + if state_dict.get("absolute_pos_embed") is not None: + absolute_pos_embed = state_dict["absolute_pos_embed"] + N1, L, C1 = absolute_pos_embed.size() + N2, C2, H, W = self.absolute_pos_embed.size() + if N1 != N2 or C1 != C2 or L != H * W: + logger.warning("Error in loading absolute_pos_embed, pass") + else: + state_dict["absolute_pos_embed"] = ( + absolute_pos_embed.view(N2, H, W, C2) + .permute(0, 3, 1, 2) + .contiguous() + ) + + # interpolate position bias table if needed + relative_position_bias_table_keys = [ + k for k in state_dict.keys() if "relative_position_bias_table" in k + ] + for table_key in relative_position_bias_table_keys: + table_pretrained = state_dict[table_key] + table_current = self.state_dict()[table_key] + L1, nH1 = table_pretrained.size() + L2, nH2 = table_current.size() + if nH1 != nH2: + logger.warning(f"Error in loading {table_key}, pass") + elif L1 != L2: + S1 = int(L1**0.5) + S2 = int(L2**0.5) + table_pretrained_resized = F.interpolate( + table_pretrained.permute(1, 0).reshape(1, nH1, S1, S1), + size=(S2, S2), + mode="bicubic", + ) + state_dict[table_key] = ( + table_pretrained_resized.view(nH2, L2) + .permute(1, 0) + .contiguous() + ) + + # load state_dict + self.load_state_dict(state_dict, False) + elif pretrained is None: + if self.use_abs_pos_embed: + trunc_normal_(self.absolute_pos_embed, std=0.02) + for m in self.modules(): + if isinstance(m, nn.Linear): + trunc_normal_init(m, std=0.02, bias=0.0) + elif isinstance(m, nn.LayerNorm): + constant_init(m, 1.0) + else: + raise TypeError("pretrained must be a str or None") + + def forward(self, x): + x, hw_shape = self.patch_embed(x) + + if self.use_abs_pos_embed: + x = x + self.absolute_pos_embed + x = self.drop_after_pos(x) + + outs = [] + for i, stage in enumerate(self.stages): + x, hw_shape, out, out_hw_shape = stage(x, hw_shape) + if i in self.out_indices: + norm_layer = getattr(self, f"norm{i}") + out = norm_layer(out) + out = ( + out.view(-1, *out_hw_shape, self.num_features[i]) + .permute(0, 3, 1, 2) + .contiguous() + ) + outs.append(out) + + res = {} + B, C, Hp, Wp = outs[3].shape + res["cls"] = torch.mean( + rearrange(outs[3], "b c h w -> b c (h w)", h=Hp, w=Wp), -1, False + ) + res["embed"] = outs[3] + + return res + + +if __name__ == "__main__": + model = SwinTransformer( + embed_dims=128, + depths=[2, 2, 18, 2], + num_heads=[4, 8, 16, 32], + window_size=7, + mlp_ratio=4, + qkv_bias=True, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.3, + patch_norm=True, + out_indices=(0, 1, 2, 3), + with_cp=False, + convert_weights=True, + ) + inputs = torch.randn(1, 3, 224, 224) + outputs = model(inputs) + pretrained_weights = "pretrained/swin_base_patch4_window7_224_22k.pth" + print(f"Loading pretrained from --- {pretrained_weights}") + model.init_weights(pretrained=pretrained_weights) + + for i in range(len(outputs)): + print(f"Output length --- {len(outputs)}, Output shape --- {outputs[i].shape}") diff --git a/video/fake-stormer/model_code/models/networks/backbones/swin3d.py b/video/fake-stormer/model_code/models/networks/backbones/swin3d.py new file mode 100644 index 0000000000000000000000000000000000000000..70b56e2171f0955c4869dc718184b41a569355bb --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/backbones/swin3d.py @@ -0,0 +1,877 @@ +# -*- coding:utf-8 -*- +import os +import sys + +if not os.getcwd() in sys.path: + sys.path.append(os.getcwd()) +from functools import lru_cache, reduce +from operator import mul + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +from einops import rearrange +from logs.logger import get_root_logger +from mmcv.cnn import build_norm_layer +from mmcv.runner import _load_checkpoint as load_checkpoint +from timm.models.layers import DropPath, trunc_normal_ + +from ...builder import BACKBONES +from .base import BaseBackbone + + +class Mlp(nn.Module): + """Multilayer perceptron.""" + + def __init__( + self, + in_features, + hidden_features=None, + out_features=None, + act_layer=nn.GELU, + drop=0.0, + ): + 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) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +def window_partition(x, window_size): + """ + Args: + x: (B, D, H, W, C) + window_size (tuple[int]): window size + + Returns: + windows: (B*num_windows, window_size*window_size, C) + """ + B, D, H, W, C = x.shape + x = x.view( + B, + D // window_size[0], + window_size[0], + H // window_size[1], + window_size[1], + W // window_size[2], + window_size[2], + C, + ) + windows = ( + x.permute(0, 1, 3, 5, 2, 4, 6, 7) + .contiguous() + .view(-1, reduce(mul, window_size), C) + ) + return windows + + +def window_reverse(windows, window_size, B, D, H, W): + """ + Args: + windows: (B*num_windows, window_size, window_size, C) + window_size (tuple[int]): Window size + H (int): Height of image + W (int): Width of image + + Returns: + x: (B, D, H, W, C) + """ + x = windows.view( + B, + D // window_size[0], + H // window_size[1], + W // window_size[2], + window_size[0], + window_size[1], + window_size[2], + -1, + ) + x = x.permute(0, 1, 4, 2, 5, 3, 6, 7).contiguous().view(B, D, H, W, -1) + return x + + +def get_window_size(x_size, window_size, shift_size=None): + use_window_size = list(window_size) + if shift_size is not None: + use_shift_size = list(shift_size) + for i in range(len(x_size)): + if x_size[i] <= window_size[i]: + use_window_size[i] = x_size[i] + if shift_size is not None: + use_shift_size[i] = 0 + + if shift_size is None: + return tuple(use_window_size) + else: + return tuple(use_window_size), tuple(use_shift_size) + + +class WindowAttention3D(nn.Module): + """Window based multi-head self attention (W-MSA) module with relative position bias. + It supports both of shifted and non-shifted window. + Args: + dim (int): Number of input channels. + window_size (tuple[int]): The temporal length, height and width of the window. + num_heads (int): Number of attention heads. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set + attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 + proj_drop (float, optional): Dropout ratio of output. Default: 0.0 + """ + + def __init__( + self, + dim, + window_size, + num_heads, + qkv_bias=False, + qk_scale=None, + attn_drop=0.0, + proj_drop=0.0, + ): + + super().__init__() + self.dim = dim + self.window_size = window_size # Wd, Wh, Ww + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim**-0.5 + + # define a parameter table of relative position bias + self.relative_position_bias_table = nn.Parameter( + torch.zeros( + (2 * window_size[0] - 1) + * (2 * window_size[1] - 1) + * (2 * window_size[2] - 1), + num_heads, + ) + ) # 2*Wd-1 * 2*Wh-1 * 2*Ww-1, nH + + # get pair-wise relative position index for each token inside the window + coords_d = torch.arange(self.window_size[0]) + coords_h = torch.arange(self.window_size[1]) + coords_w = torch.arange(self.window_size[2]) + coords = torch.stack( + torch.meshgrid(coords_d, coords_h, coords_w) + ) # 3, Wd, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 3, Wd*Wh*Ww + relative_coords = ( + coords_flatten[:, :, None] - coords_flatten[:, None, :] + ) # 3, Wd*Wh*Ww, Wd*Wh*Ww + relative_coords = relative_coords.permute( + 1, 2, 0 + ).contiguous() # Wd*Wh*Ww, Wd*Wh*Ww, 3 + relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 2] += self.window_size[2] - 1 + + relative_coords[:, :, 0] *= (2 * self.window_size[1] - 1) * ( + 2 * self.window_size[2] - 1 + ) + relative_coords[:, :, 1] *= 2 * self.window_size[2] - 1 + relative_position_index = relative_coords.sum(-1) # Wd*Wh*Ww, Wd*Wh*Ww + self.register_buffer("relative_position_index", relative_position_index) + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + trunc_normal_(self.relative_position_bias_table, std=0.02) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, x, mask=None): + """Forward function. + Args: + x: input features with shape of (num_windows*B, N, C) + mask: (0/-inf) mask with shape of (num_windows, N, N) or None + """ + B_, N, C = x.shape + qkv = ( + self.qkv(x) + .reshape(B_, N, 3, self.num_heads, C // self.num_heads) + .permute(2, 0, 3, 1, 4) + ) + q, k, v = qkv[0], qkv[1], qkv[2] # B_, nH, N, C + + q = q * self.scale + attn = q @ k.transpose(-2, -1) + + relative_position_bias = self.relative_position_bias_table[ + self.relative_position_index[:N, :N].reshape(-1) + ].reshape( + N, N, -1 + ) # Wd*Wh*Ww,Wd*Wh*Ww,nH + relative_position_bias = relative_position_bias.permute( + 2, 0, 1 + ).contiguous() # nH, Wd*Wh*Ww, Wd*Wh*Ww + attn = attn + relative_position_bias.unsqueeze(0) # B_, nH, N, N + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze( + 1 + ).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + attn = self.softmax(attn) + else: + attn = self.softmax(attn) + + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class SwinTransformerBlock3D(nn.Module): + """Swin Transformer Block. + + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (tuple[int]): Window size. + shift_size (tuple[int]): Shift size for SW-MSA. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float, optional): Stochastic depth rate. Default: 0.0 + act_layer (nn.Module, optional): Activation layer. Default: nn.GELU + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__( + self, + dim, + num_heads, + window_size=(2, 7, 7), + shift_size=(0, 0, 0), + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + use_checkpoint=False, + ): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.window_size = window_size + self.shift_size = shift_size + self.mlp_ratio = mlp_ratio + self.use_checkpoint = use_checkpoint + + assert ( + 0 <= self.shift_size[0] < self.window_size[0] + ), "shift_size must in 0-window_size" + assert ( + 0 <= self.shift_size[1] < self.window_size[1] + ), "shift_size must in 0-window_size" + assert ( + 0 <= self.shift_size[2] < self.window_size[2] + ), "shift_size must in 0-window_size" + + self.norm1 = norm_layer(dim) + self.attn = WindowAttention3D( + dim, + window_size=self.window_size, + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=drop, + ) + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop, + ) + + def forward_part1(self, x, mask_matrix): + B, D, H, W, C = x.shape + window_size, shift_size = get_window_size( + (D, H, W), self.window_size, self.shift_size + ) + + x = self.norm1(x) + # pad feature maps to multiples of window size + pad_l = pad_t = pad_d0 = 0 + pad_d1 = (window_size[0] - D % window_size[0]) % window_size[0] + pad_b = (window_size[1] - H % window_size[1]) % window_size[1] + pad_r = (window_size[2] - W % window_size[2]) % window_size[2] + x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b, pad_d0, pad_d1)) + _, Dp, Hp, Wp, _ = x.shape + # cyclic shift + if any(i > 0 for i in shift_size): + shifted_x = torch.roll( + x, + shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), + dims=(1, 2, 3), + ) + attn_mask = mask_matrix + else: + shifted_x = x + attn_mask = None + # partition windows + x_windows = window_partition(shifted_x, window_size) # B*nW, Wd*Wh*Ww, C + # W-MSA/SW-MSA + attn_windows = self.attn(x_windows, mask=attn_mask) # B*nW, Wd*Wh*Ww, C + # merge windows + attn_windows = attn_windows.view(-1, *(window_size + (C,))) + shifted_x = window_reverse( + attn_windows, window_size, B, Dp, Hp, Wp + ) # B D' H' W' C + # reverse cyclic shift + if any(i > 0 for i in shift_size): + x = torch.roll( + shifted_x, + shifts=(shift_size[0], shift_size[1], shift_size[2]), + dims=(1, 2, 3), + ) + else: + x = shifted_x + + if pad_d1 > 0 or pad_r > 0 or pad_b > 0: + x = x[:, :D, :H, :W, :].contiguous() + return x + + def forward_part2(self, x): + return self.drop_path(self.mlp(self.norm2(x))) + + def forward(self, x, mask_matrix): + """Forward function. + + Args: + x: Input feature, tensor size (B, D, H, W, C). + mask_matrix: Attention mask for cyclic shift. + """ + + shortcut = x + if self.use_checkpoint: + x = checkpoint.checkpoint(self.forward_part1, x, mask_matrix) + else: + x = self.forward_part1(x, mask_matrix) + x = shortcut + self.drop_path(x) + + if self.use_checkpoint: + x = x + checkpoint.checkpoint(self.forward_part2, x) + else: + x = x + self.forward_part2(x) + + return x + + +class PatchMerging(nn.Module): + """Patch Merging Layer + + Args: + dim (int): Number of input channels. + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__(self, dim, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(4 * dim) + + def forward(self, x): + """Forward function. + + Args: + x: Input feature, tensor size (B, D, H, W, C). + """ + B, D, H, W, C = x.shape + + # padding + pad_input = (H % 2 == 1) or (W % 2 == 1) + if pad_input: + x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2)) + + x0 = x[:, :, 0::2, 0::2, :] # B D H/2 W/2 C + x1 = x[:, :, 1::2, 0::2, :] # B D H/2 W/2 C + x2 = x[:, :, 0::2, 1::2, :] # B D H/2 W/2 C + x3 = x[:, :, 1::2, 1::2, :] # B D H/2 W/2 C + x = torch.cat([x0, x1, x2, x3], -1) # B D H/2 W/2 4*C + + x = self.norm(x) + x = self.reduction(x) + + return x + + +# cache each stage results +@lru_cache() +def compute_mask(D, H, W, window_size, shift_size, device): + img_mask = torch.zeros((1, D, H, W, 1), device=device) # 1 Dp Hp Wp 1 + cnt = 0 + for d in ( + slice(-window_size[0]), + slice(-window_size[0], -shift_size[0]), + slice(-shift_size[0], None), + ): + for h in ( + slice(-window_size[1]), + slice(-window_size[1], -shift_size[1]), + slice(-shift_size[1], None), + ): + for w in ( + slice(-window_size[2]), + slice(-window_size[2], -shift_size[2]), + slice(-shift_size[2], None), + ): + img_mask[:, d, h, w, :] = cnt + cnt += 1 + mask_windows = window_partition(img_mask, window_size) # nW, ws[0]*ws[1]*ws[2], 1 + mask_windows = mask_windows.squeeze(-1) # nW, ws[0]*ws[1]*ws[2] + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill( + attn_mask == 0, float(0.0) + ) + return attn_mask + + +class BasicLayer(nn.Module): + """A basic Swin Transformer layer for one stage. + + Args: + dim (int): Number of feature channels + depth (int): Depths of this stage. + num_heads (int): Number of attention head. + window_size (tuple[int]): Local window size. Default: (1,7,7). + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None + """ + + def __init__( + self, + dim, + depth, + num_heads, + window_size=(1, 7, 7), + mlp_ratio=4.0, + qkv_bias=False, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + norm_layer=nn.LayerNorm, + downsample=None, + use_checkpoint=False, + ): + super().__init__() + self.window_size = window_size + self.shift_size = tuple(i // 2 for i in window_size) + self.depth = depth + self.use_checkpoint = use_checkpoint + + # build blocks + self.blocks = nn.ModuleList( + [ + SwinTransformerBlock3D( + dim=dim, + num_heads=num_heads, + window_size=window_size, + shift_size=(0, 0, 0) if (i % 2 == 0) else self.shift_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + drop_path=( + drop_path[i] if isinstance(drop_path, list) else drop_path + ), + norm_layer=norm_layer, + use_checkpoint=use_checkpoint, + ) + for i in range(depth) + ] + ) + + self.downsample = downsample + if self.downsample is not None: + self.downsample = downsample(dim=dim, norm_layer=norm_layer) + + def forward(self, x): + """Forward function. + + Args: + x: Input feature, tensor size (B, C, D, H, W). + """ + # calculate attention mask for SW-MSA + B, C, D, H, W = x.shape + window_size, shift_size = get_window_size( + (D, H, W), self.window_size, self.shift_size + ) + x = rearrange(x, "b c d h w -> b d h w c") + Dp = int(np.ceil(D / window_size[0])) * window_size[0] + Hp = int(np.ceil(H / window_size[1])) * window_size[1] + Wp = int(np.ceil(W / window_size[2])) * window_size[2] + attn_mask = compute_mask(Dp, Hp, Wp, window_size, shift_size, x.device) + for blk in self.blocks: + x = blk(x, attn_mask) + x = x.view(B, D, H, W, -1) + + if self.downsample is not None: + x = self.downsample(x) + x = rearrange(x, "b d h w c -> b c d h w") + return x + + +class PatchEmbed3D(nn.Module): + """Video to Patch Embedding. + + Args: + patch_size (int): Patch token size. Default: (2,4,4). + in_chans (int): Number of input video channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__(self, patch_size=(2, 4, 4), in_chans=3, embed_dim=96, norm_layer=None): + super().__init__() + self.patch_size = patch_size + + self.in_chans = in_chans + self.embed_dim = embed_dim + + self.proj = nn.Conv3d( + in_chans, embed_dim, kernel_size=patch_size, stride=patch_size + ) + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = None + + def forward(self, x): + """Forward function.""" + # padding + _, _, D, H, W = x.size() + if W % self.patch_size[2] != 0: + x = F.pad(x, (0, self.patch_size[2] - W % self.patch_size[2])) + if H % self.patch_size[1] != 0: + x = F.pad(x, (0, 0, 0, self.patch_size[1] - H % self.patch_size[1])) + if D % self.patch_size[0] != 0: + x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - D % self.patch_size[0])) + + x = self.proj(x) # B C D Wh Ww + if self.norm is not None: + D, Wh, Ww = x.size(2), x.size(3), x.size(4) + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + x = x.transpose(1, 2).view(-1, self.embed_dim, D, Wh, Ww) + + return x + + +@BACKBONES.register_module() +class SwinTransformer3D(BaseBackbone): + """Swin Transformer backbone. + A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - + https://arxiv.org/pdf/2103.14030 + + Args: + patch_size (int | tuple(int)): Patch size. Default: (4,4,4). + in_chans (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + depths (tuple[int]): Depths of each Swin Transformer stage. + num_heads (tuple[int]): Number of attention head of each stage. + window_size (int): Window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: Truee + qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. + drop_rate (float): Dropout rate. + attn_drop_rate (float): Attention dropout rate. Default: 0. + drop_path_rate (float): Stochastic depth rate. Default: 0.2. + norm_layer: Normalization layer. Default: nn.LayerNorm. + patch_norm (bool): If True, add normalization after patch embedding. Default: False. + frozen_stages (int): Stages to be frozen (stop grad and set eval mode). + -1 means not freezing any parameters. + """ + + def __init__( + self, + pretrained=None, + pretrained2d=True, + patch_size=(4, 4, 4), + in_chans=3, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=(2, 7, 7), + out_indices=(0, 1, 2, 3), + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.2, + norm_layer=nn.LayerNorm, + norm_cfg=dict(type="LN"), + patch_norm=False, + frozen_stages=-1, + use_checkpoint=False, + ): + super(SwinTransformer3D, self).__init__() + + self.pretrained = pretrained + self.pretrained2d = pretrained2d + self.num_layers = len(depths) + self.embed_dim = embed_dim + self.patch_norm = patch_norm + self.frozen_stages = frozen_stages + self.window_size = window_size + self.patch_size = patch_size + + # split image into non-overlapping patches + self.patch_embed = PatchEmbed3D( + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None, + ) + + self.pos_drop = nn.Dropout(p=drop_rate) + + # stochastic depth + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, sum(depths)) + ] # stochastic depth decay rule + + # build layers + self.layers = nn.ModuleList() + for i_layer in range(self.num_layers): + layer = BasicLayer( + dim=int(embed_dim * 2**i_layer), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])], + norm_layer=norm_layer, + downsample=PatchMerging if i_layer < self.num_layers - 1 else None, + use_checkpoint=use_checkpoint, + ) + self.layers.append(layer) + + self.num_features = [int(embed_dim * 2**i) for i in range(self.num_layers)] + + # Add a norm layer for each output + for i in out_indices: + layer = build_norm_layer(norm_cfg, self.num_features[i])[1] + layer_name = f"norm{i}" + self.add_module(layer_name, layer) + + def _freeze_stages(self): + if self.frozen_stages >= 0: + self.patch_embed.eval() + for param in self.patch_embed.parameters(): + param.requires_grad = False + + if self.frozen_stages >= 1: + self.pos_drop.eval() + for i in range(0, self.frozen_stages): + m = self.layers[i] + m.eval() + for param in m.parameters(): + param.requires_grad = False + + def inflate_weights(self, logger): + """Inflate the swin2d parameters to swin3d. + + The differences between swin3d and swin2d mainly lie in an extra + axis. To utilize the pretrained parameters in 2d model, + the weight of swin2d models should be inflated to fit in the shapes of + the 3d counterpart. + + Args: + logger (logging.Logger): The logger used to print + debugging infomation. + """ + checkpoint = torch.load(self.pretrained, map_location="cpu") + state_dict = checkpoint["model"] + + # delete relative_position_index since we always re-init it + relative_position_index_keys = [ + k for k in state_dict.keys() if "relative_position_index" in k + ] + for k in relative_position_index_keys: + del state_dict[k] + + # delete attn_mask since we always re-init it + attn_mask_keys = [k for k in state_dict.keys() if "attn_mask" in k] + for k in attn_mask_keys: + del state_dict[k] + + state_dict["patch_embed.proj.weight"] = ( + state_dict["patch_embed.proj.weight"] + .unsqueeze(2) + .repeat(1, 1, self.patch_size[0], 1, 1) + / self.patch_size[0] + ) + + # bicubic interpolate relative_position_bias_table if not match + relative_position_bias_table_keys = [ + k for k in state_dict.keys() if "relative_position_bias_table" in k + ] + for k in relative_position_bias_table_keys: + relative_position_bias_table_pretrained = state_dict[k] + relative_position_bias_table_current = self.state_dict()[k] + L1, nH1 = relative_position_bias_table_pretrained.size() + L2, nH2 = relative_position_bias_table_current.size() + L2 = (2 * self.window_size[1] - 1) * (2 * self.window_size[2] - 1) + wd = self.window_size[0] + if nH1 != nH2: + logger.warning(f"Error in loading {k}, passing") + else: + if L1 != L2: + S1 = int(L1**0.5) + relative_position_bias_table_pretrained_resized = ( + torch.nn.functional.interpolate( + relative_position_bias_table_pretrained.permute(1, 0).view( + 1, nH1, S1, S1 + ), + size=( + 2 * self.window_size[1] - 1, + 2 * self.window_size[2] - 1, + ), + mode="bicubic", + ) + ) + relative_position_bias_table_pretrained = ( + relative_position_bias_table_pretrained_resized.view( + nH2, L2 + ).permute(1, 0) + ) + state_dict[k] = relative_position_bias_table_pretrained.repeat( + 2 * wd - 1, 1 + ) + + msg = self.load_state_dict(state_dict, strict=False) + logger.info(msg) + logger.info(f"=> loaded successfully '{self.pretrained}'") + del checkpoint + torch.cuda.empty_cache() + + def init_weights(self, pretrained=None): + """Initialize the weights in backbone. + + Args: + pretrained (str, optional): Path to pre-trained weights. + Defaults to None. + """ + + def _init_weights(m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=0.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + if pretrained: + self.pretrained = pretrained + if isinstance(self.pretrained, str): + self.apply(_init_weights) + logger = get_root_logger() + logger.info(f"load model from: {self.pretrained}") + + if self.pretrained2d: + # Inflate 2D model into 3D model. + self.inflate_weights(logger) + else: + # Directly load 3D model. + load_checkpoint(self, self.pretrained, strict=False, logger=logger) + elif self.pretrained is None: + self.apply(_init_weights) + else: + raise TypeError("pretrained must be a str or None") + + def forward(self, x): + """Forward function.""" + x = self.patch_embed(x) + + x = self.pos_drop(x) + + outputs = [] + for i, layer in enumerate(self.layers): + x = layer(x.contiguous()) + x = rearrange(x, "n c d h w -> n d h w c") + if i == len(self.layers) - 1: + norm_layer = getattr(self, f"norm{i}") + else: + norm_layer = getattr(self, f"norm{i+1}") + x = norm_layer(x) + x = rearrange(x, "n d h w c -> n c d h w") + outputs.append(x) + + res = {} + res["embed"] = x + res["outputs"] = outputs + + return res + + def train(self, mode=True): + """Convert the model into training mode while keep layers freezed.""" + super(SwinTransformer3D, self).train(mode) + self._freeze_stages() + + +if __name__ == "__main__": + model = SwinTransformer3D( + embed_dim=128, + patch_size=(4, 4, 4), + depths=[2, 2, 18, 2], + num_heads=[4, 8, 16, 32], + window_size=(8, 7, 7), + mlp_ratio=4, + qkv_bias=True, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.2, + patch_norm=True, + ) + inputs = torch.randn(1, 3, 32, 224, 224) + outputs = model(inputs) + pretrained_weights = "pretrained/swin_base_patch4_window7_224_22k.pth" + print(f"Loading pretrained from --- {pretrained_weights}") + model.init_weights(pretrained=pretrained_weights) + + for k, v in outputs.items(): + print(f"Output length --- {len(outputs)}, Output shape --- {outputs[k].shape}") diff --git a/video/fake-stormer/model_code/models/networks/backbones/vit.py b/video/fake-stormer/model_code/models/networks/backbones/vit.py new file mode 100644 index 0000000000000000000000000000000000000000..40fa725b365eb93d5abf26143a722b2c051c831c --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/backbones/vit.py @@ -0,0 +1,1271 @@ +# -*- coding: utf-8 -*- +# Copyright (c) OpenMMLab. All rights reserved. +import math +from functools import partial +from typing import Dict, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +from einops import rearrange, reduce, repeat +from models.builder import BACKBONES +from timm.models.layers import drop_path, to_2tuple, trunc_normal_ + +from ..common import BN3D_MOMENTUM +from ..pose_efficientNet import EfficientNet +from .base import BaseBackbone +from .efficientNet import get_model_params + + +def get_abs_pos(abs_pos, h, w, ori_h, ori_w, has_cls_token=True): + """ + Calculate absolute positional embeddings. If needed, resize embeddings and remove cls_token + dimension for the original embeddings. + Args: + abs_pos (Tensor): absolute positional embeddings with (1, num_position, C). + has_cls_token (bool): If true, has 1 embedding in abs_pos for cls token. + hw (Tuple): size of input image tokens. + + Returns: + Absolute positional embeddings after processing with shape (1, H, W, C) + """ + cls_token = None + B, L, C = abs_pos.shape + if has_cls_token: + cls_token = abs_pos[:, 0:1] + abs_pos = abs_pos[:, 1:] + + if ori_h != h or ori_w != w: + new_abs_pos = ( + F.interpolate( + abs_pos.reshape(1, ori_h, ori_w, -1).permute(0, 3, 1, 2), + size=(h, w), + mode="bicubic", + align_corners=False, + ) + .permute(0, 2, 3, 1) + .reshape(B, -1, C) + ) + + else: + new_abs_pos = abs_pos + + if cls_token is not None: + new_abs_pos = torch.cat([cls_token, new_abs_pos], dim=1) + return new_abs_pos + + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob=None): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training) + + def extra_repr(self): + return "p={}".format(self.drop_prob) + + +class Mlp(nn.Module): + def __init__( + self, + in_features, + hidden_features=None, + out_features=None, + act_layer=nn.GELU, + drop=0.0, + ): + 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) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +class Attention(nn.Module): + def __init__( + self, + dim, + num_heads=8, + qkv_bias=False, + qk_scale=None, + attn_drop=0.0, + proj_drop=0.0, + attn_head_dim=None, + att_dimension=None, + ): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.dim = dim + + if attn_head_dim is not None: + head_dim = attn_head_dim + all_head_dim = head_dim * self.num_heads + + self.scale = qk_scale or head_dim**-0.5 + + self.qkv = nn.Linear(dim, all_head_dim * 3, bias=qkv_bias) + + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(all_head_dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + self.attn_cam = None + self.attn = None + self.v = None + self.v_cam = None + self.attn_gradients = None + + def get_attn(self): + return self.attn + + def save_attn(self, attn): + self.attn = attn + + def save_attn_cam(self, cam): + self.attn_cam = cam + + def get_attn_cam(self): + return self.attn_cam + + def get_v(self): + return self.v + + def save_v(self, v): + self.v = v + + def save_v_cam(self, cam): + self.v_cam = cam + + def get_v_cam(self): + return self.v_cam + + def save_attn_gradients(self, attn_gradients): + self.attn_gradients = attn_gradients + + def get_attn_gradients(self): + return self.attn_gradients + + def forward(self, x, b_size, **kwargs): + B, N, C = x.shape + # mask = None + + # if 'maskout_pes' in kwargs.keys(): + # mask = kwargs['maskout_pes'] + # T = mask.shape[1] + # num_tokens = mask.shape[2]*mask.shape[3] + + # if num_tokens == (N-1): + # mask = rearrange(mask, 'b t h w -> (b t) (h w)', b=b_size, t=T, h=mask.shape[2], w= mask.shape[3]).unsqueeze(1) + # mask = mask.unsqueeze(3) + # mask = mask.repeat(1, self.num_heads, 1, num_tokens+1) + # mask = torch.cat((torch.ones((b_size*T, self.num_heads, 1, num_tokens+1)).cuda(), mask), 2) + # else: + # mask = rearrange(mask, 'b t h w -> (b h w) t', b=b_size, t=T, h=mask.shape[2], w= mask.shape[3]).unsqueeze(1) + # mask = mask.unsqueeze(3) + # mask = mask.repeat(1, self.num_heads, 1, T+1) + # mask = torch.cat((torch.ones((b_size*num_tokens, self.num_heads, 1, T+1)).cuda(), mask), 2) + + qkv = self.qkv(x) + qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute( + 2, 0, 3, 1, 4 + ) # 3, B (BxT or BxN), H, N, C + q, k, v = ( + qkv[0], + qkv[1], + qkv[2], + ) # make torchscript happy (cannot use tensor as tuple) + + # self.save_v(v) + + q = q * self.scale + attn = q @ k.transpose(-2, -1) + + # if mask is not None: + # attn = attn.masked_fill(mask == 0, float("-1e20")) + + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + # self.save_attn(attn) + # attn.register_hook(self.save_attn_gradients) + + x = (attn @ v).transpose(1, 2).reshape(B, N, -1) + x = self.proj(x) + x = self.proj_drop(x) + + return x + + +class CrossAttention(nn.Module): + def __init__( + self, + dim, + num_heads=8, + qkv_bias=False, + qk_scale=None, + attn_drop=0.0, + proj_drop=0.0, + attn_head_dim=None, + att_dimension="spatial", + ): + """ + Implementation for cross attention through dimension + """ + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.dim = dim + + if attn_head_dim is not None: + head_dim = attn_head_dim + all_head_dim = head_dim * self.num_heads + + self.scale = qk_scale or head_dim**-0.5 + + self.qkv = nn.Linear(dim, all_head_dim * 3, bias=qkv_bias) + + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(all_head_dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + assert att_dimension in ["spatial", "temporal"] + self.att_dimension = att_dimension + + def forward(self, x, b_size): + B, N, C = ( + x.shape + ) # N can be T or HW // P**2, B can be previous reshaped from self.batch_size * T|B + T = B // b_size + + qkv = self.qkv(x) + qkv = qkv.reshape(b_size, T, N, 3, self.num_heads, -1).permute(3, 0, 4, 1, 2, 5) + q, k, v = ( + qkv[0], + qkv[1], + qkv[2], + ) # make torchscript happy (cannot use tensor as tuple) + + # Define the window + window = torch.tensor([-2, -1, 0, 1, 2]).cuda() + + # Start computing locally cross self-attention + q_scale = q * self.scale + + # Initialize attn weights + attn = torch.zeros((b_size, self.num_heads, T, N, N), dtype=torch.float).cuda() + + if self.att_dimension == "spatial": + sqrt_N = int(math.sqrt(N)) + for w in window: + i_indices = torch.arange(1, T).cuda() + j_indices = torch.arange(0, N).cuda() + + i_valid = (i_indices + w >= 0) & (i_indices + w < T) + j_valid = (j_indices + sqrt_N * w >= 0) & (j_indices + sqrt_N * w < N) + + i_indices = i_indices[i_valid] + j_indices = j_indices[j_valid] + + attn[ + :, :, torch.cat((torch.zeros(1, dtype=int).cuda(), i_indices + w)) + ][:, :, :, j_indices][..., j_indices + sqrt_N * w] = q_scale[ + :, :, torch.cat((torch.zeros(1, dtype=int).cuda(), i_indices)) + ][ + :, :, :, j_indices + ] @ k[ + :, :, torch.cat((torch.zeros(1, dtype=int).cuda(), i_indices + w)) + ][ + :, :, :, j_indices + sqrt_N * w + ].transpose( + -2, -1 + ) + else: + sqrt_T = int(math.sqrt(T)) + + for w in window: + i_indices = torch.arange(0, T).cuda() + j_indices = torch.arange(1, N).cuda() + + i_valid = (i_indices + sqrt_T * w >= 0) & (i_indices + sqrt_T * w < T) + j_valid = (j_indices + w >= 1) & (j_indices + w < N) + + i_indices = i_indices[i_valid] + j_indices = j_indices[j_valid] + + attn[:, :, i_indices + sqrt_T * w][ + :, :, :, torch.cat((torch.zeros(1, dtype=int).cuda(), j_indices)) + ][ + ..., torch.cat((torch.zeros(1, dtype=int).cuda(), j_indices + w)) + ] = q_scale[ + :, :, i_indices + ][ + :, :, :, torch.cat((torch.zeros(1, dtype=int).cuda(), j_indices)) + ] @ k[ + :, :, i_indices + sqrt_T * w + ][ + :, + :, + :, + torch.cat((torch.zeros(1, dtype=int).cuda(), j_indices + w)), + ].transpose( + -2, -1 + ) + + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, -1) + x = self.proj(x) + x = self.proj_drop(x) + + return x + + +class Block(nn.Module): + def __init__( + self, + dim, + num_heads, + mlp_ratio=4.0, + qkv_bias=False, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + attn_head_dim=None, + ): + super().__init__() + + self.norm1 = norm_layer(dim) + self.attn = Attention( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=drop, + attn_head_dim=attn_head_dim, + ) + + # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop, + ) + + def forward(self, x, b_size, **kwargs): + x = x + self.drop_path(self.attn(self.norm1(x), b_size)) + x = x + self.drop_path(self.mlp(self.norm2(x))) + return x + + +class Block2D(nn.Module): + def __init__( + self, + dim, + num_heads, + mlp_ratio=4.0, + qkv_bias=False, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.1, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + attention_type="divided_space_time", + **kwargs, + ): + super().__init__() + self.attention_type = attention_type + assert attention_type in [ + "divided_space_time", + "space_only", + "joint_space_time", + ] + self.register_token = kwargs.get("register_token") + self.temp_token = kwargs.get("temp_token") + self.return_s_cls_token = kwargs.get("return_s_cls_token") or False + + self.norm1 = norm_layer(dim) + self.attn = Attention( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=drop, + ) + + ## Temporal Attention Parameters + if self.attention_type == "divided_space_time": + self.temporal_norm1 = norm_layer(dim) + self.temporal_attn = Attention( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=drop, + att_dimension="temporal", + ) + self.temporal_fc = nn.Linear(dim, dim) + + ## drop path + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + self.norm2 = norm_layer(dim) + # self.norm3 = norm_layer(dim) + # self.norm4 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop, + ) + + def forward(self, x, B, T, W, **kwargs): + if self.temp_token: + num_spatial_tokens = (x.size(1) - 2) // T + else: + num_spatial_tokens = (x.size(1) - 1) // T + H = num_spatial_tokens // W + + if self.attention_type in ["space_only", "joint_space_time"]: + x = x + self.drop_path(self.attn(self.norm1(x))) + x = x + self.drop_path(self.mlp(self.norm2(x))) + return x, None + elif self.attention_type == "divided_space_time": + # init class token + init_cls_token = x[:, 0, :].unsqueeze(1) + if self.temp_token: + init_temp_token = x[:, -1, :].unsqueeze(1) + # init_register_token = x[:, -1, :].unsqueeze(1) + + ## Temporal + if self.temp_token: + t_cls_token = init_temp_token.repeat(1, num_spatial_tokens, 1) + else: + t_cls_token = init_cls_token.repeat(1, num_spatial_tokens, 1) + t_cls_token = rearrange( + t_cls_token, "b (h w) m -> (b h w) m", b=B, h=H, w=W + ).unsqueeze(1) + + if self.temp_token: + xt = x[:, 1:-1, :] + else: + xt = x[:, 1:, :] + + xt = rearrange(xt, "b (h w t) m -> (b h w) t m", b=B, h=H, w=W, t=T) + xt = torch.cat((xt, t_cls_token), 1) + + res_temporal = self.drop_path( + self.temporal_attn(self.temporal_norm1(xt), b_size=B, **kwargs) + ) # Processing temporal att. + res_temporal = self.temporal_fc(res_temporal) + res_temporal, t_cls_token = res_temporal[:, :-1, :], res_temporal[:, -1, :] + res_temporal = rearrange( + res_temporal, "(b h w) t m -> b (h w t) m", b=B, h=H, w=W, t=T + ) + t_cls_token = rearrange( + t_cls_token, "(b h w) m -> b (h w) m", b=B, h=H, w=W + ) + # t_cls_token_avg = torch.mean(t_cls_token, 1, True) ## average for every temporal patch + xt = x[:, 1:-1, :] + res_temporal + + ## Spatial + cls_token = init_cls_token.repeat(1, T, 1) + cls_token = rearrange(cls_token, "b t m -> (b t) m", b=B, t=T).unsqueeze(1) + # register_token = init_register_token.repeat(1, T, 1) + # register_token = rearrange(register_token, 'b t m -> (b t) m', b=B, t=T).unsqueeze(1) + xs = xt + xs = rearrange(xs, "b (h w t) m -> (b t) (h w) m", b=B, h=H, w=W, t=T) + xs = torch.cat((cls_token, xs), 1) + # xs = torch.cat((xs, register_token), 1) + res_spatial = self.drop_path(self.attn(self.norm1(xs), b_size=B, **kwargs)) + + ### Taking care of TEMP token + t_cls_token_avg = torch.mean( + t_cls_token, 1, True + ) ## average for every temporal patch + + ### Taking care of CLS token + cls_token = res_spatial[:, 0, :] + cls_token = rearrange(cls_token, "(b t) m -> b t m", b=B, t=T) + cls_token_avg = torch.mean(cls_token, 1, True) ## averaging for every frame + # register_token = res_spatial[:, -1, :] + # register_token = rearrange(register_token, '(b t) m -> b t m', b=B, t=T) + # register_token_avg = torch.mean(register_token, 1, True) ## averaging for every frame + + res_spatial = res_spatial[:, 1:, :] + res_spatial = rearrange( + res_spatial, "(b t) (h w) m -> b (h w t) m", b=B, h=H, w=W, t=T + ) + res = res_spatial + x = xt + + ## Mlp + # x = torch.cat((t_cls_token_avg, x), 1) + torch.cat((cls_token_avg, res), 1) + xt_ = torch.cat((init_cls_token, x), 1) + if self.temp_token: + xt_ = torch.cat((xt_, t_cls_token_avg), 1) + xs_ = torch.cat((cls_token_avg, res), 1) + if self.temp_token: + xs_ = torch.cat((xs_, init_temp_token), 1) + x = xt_ + xs_ + x = x + self.drop_path(self.mlp(self.norm2(x))) + + # Adding MLP for spatial_cls_token, ttemp_cls_token + # t_cls_token = self.drop_path(self.mlp(self.norm3(t_cls_token))) + # s_cls_token = self.drop_path(self.mlp(self.norm4(cls_token))) + + if self.return_s_cls_token: + return x, cls_token, None + else: + return x, None, None + + +class PatchEmbed(nn.Module): + """Image to Patch Embedding""" + + def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, ratio=1): + super().__init__() + img_size = to_2tuple(img_size) + patch_size = to_2tuple(patch_size) + num_patches = ( + (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) * (ratio**2) + ) + self.patch_shape = ( + int(img_size[0] // patch_size[0] * ratio), + int(img_size[1] // patch_size[1] * ratio), + ) + self.origin_patch_shape = ( + int(img_size[0] // patch_size[0]), + int(img_size[1] // patch_size[1]), + ) + self.img_size = img_size + self.patch_size = patch_size + self.num_patches = num_patches + + self.proj = nn.Conv2d( + in_chans, + embed_dim, + kernel_size=patch_size, + stride=(patch_size[0] // ratio), + padding=4 + 2 * (ratio // 2 - 1), + ) + + def forward(self, x, **kwargs): + B, C, H, W = x.shape + x = self.proj(x) + Hp, Wp = x.shape[2], x.shape[3] + + x = x.flatten(2).transpose(1, 2) + return x, (Hp, Wp) + + +class PatchEmbed3D(nn.Module): + """Images to Patch Embedding""" + + def __init__( + self, + img_size=224, + patch_size=16, + in_chans=3, + embed_dim=768, + ratio=1, + low_level=False, + **override_params, + ): + super().__init__() + img_size = to_2tuple(img_size) + patch_size = to_2tuple(patch_size) + num_patches = ( + (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) * (ratio**2) + ) + self.patch_shape = ( + int(img_size[0] // patch_size[0] * ratio), + int(img_size[1] // patch_size[1] * ratio), + ) + self.origin_patch_shape = ( + int(img_size[0] // patch_size[0]), + int(img_size[1] // patch_size[1]), + ) + self.img_size = img_size + self.patch_size = patch_size + self.num_patches = num_patches + self.low_level = low_level + + if not self.low_level: + self.proj = nn.Conv2d( + in_chans, embed_dim, kernel_size=patch_size, stride=patch_size + ) + else: + model_name = "efficientnet-b4" + self.proj = EfficientNet.from_pretrained( + model_name, advprop=True, **override_params + ) + self.fc = nn.Linear(160, embed_dim) + + def forward(self, x): + B, C, T, H, W = x.shape + x = rearrange(x, "b c t h w -> (b t) c h w") + + if not self.low_level: + x = self.proj(x) + else: + endpoints = self.proj.extract_endpoints(x) + x1 = endpoints["reduction_6"] + x2 = endpoints["reduction_5"] + x3 = endpoints["reduction_4"] + x4 = endpoints["reduction_3"] + x5 = endpoints["reduction_2"] + x = x3 + x = x.permute(0, 2, 3, 1) + x = self.fc(x) + x = x.permute(0, 3, 1, 2) + + Hp, Wp = x.shape[2], x.shape[3] + x = x.flatten(2).transpose(1, 2) + return x, T, (Hp, Wp) + + +class HybridEmbed(nn.Module): + """CNN Feature Map Embedding + Extract feature map from CNN, flatten, project to embedding dim. + """ + + def __init__( + self, backbone, img_size=224, feature_size=None, in_chans=3, embed_dim=768 + ): + super().__init__() + assert isinstance(backbone, nn.Module) + img_size = to_2tuple(img_size) + self.img_size = img_size + self.backbone = backbone + if feature_size is None: + with torch.no_grad(): + training = backbone.training + if training: + backbone.eval() + o = self.backbone(torch.zeros(1, in_chans, img_size[0], img_size[1]))[ + -1 + ] + feature_size = o.shape[-2:] + feature_dim = o.shape[1] + backbone.train(training) + else: + feature_size = to_2tuple(feature_size) + feature_dim = self.backbone.feature_info.channels()[-1] + self.num_patches = feature_size[0] * feature_size[1] + self.proj = nn.Linear(feature_dim, embed_dim) + + def forward(self, x): + x = self.backbone(x)[-1] + x = x.flatten(2).transpose(1, 2) + x = self.proj(x) + return x + + +@BACKBONES.register_module() +class ViT(BaseBackbone): + def __init__( + self, + img_size=224, + patch_size=16, + in_chans=3, + num_classes=80, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4.0, + qkv_bias=False, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.0, + hybrid_backbone=None, + norm_layer=None, + use_checkpoint=False, + frozen_stages=-1, + ratio=1, + last_norm=True, + class_token=True, + attention_type="space", + patch_padding="pad", + freeze_attn=False, + freeze_ffn=False, + **kwargs, + ): + # Protect mutable default arguments + super(ViT, self).__init__() + norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) + self.num_classes = num_classes + self.num_features = self.embed_dim = ( + embed_dim # num_features for consistency with other models + ) + self.frozen_stages = frozen_stages + self.use_checkpoint = use_checkpoint + self.patch_padding = patch_padding + self.freeze_attn = freeze_attn + self.freeze_ffn = freeze_ffn + self.depth = depth + self.attention_type = attention_type + + if hybrid_backbone is not None: + self.patch_embed = HybridEmbed( + hybrid_backbone, + img_size=img_size, + in_chans=in_chans, + embed_dim=embed_dim, + ) + else: + self.patch_embed = PatchEmbed( + img_size=img_size, + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + ratio=ratio, + ) + num_patches = self.patch_embed.num_patches + + self.cls_token = ( + nn.Parameter(torch.zeros(1, 1, embed_dim)) if class_token else None + ) + + # since the pretraining model has class token + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) + + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, depth) + ] # stochastic depth decay rule + + self.blocks = nn.ModuleList( + [ + Block( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[i], + norm_layer=norm_layer, + ) + for i in range(depth) + ] + ) + + self.norm = norm_layer(embed_dim) if last_norm else nn.Identity() + + if self.pos_embed is not None: + trunc_normal_(self.pos_embed, std=0.02) + + self._freeze_stages() + + def _freeze_stages(self): + """Freeze parameters.""" + if self.frozen_stages >= 0: + self.patch_embed.eval() + for param in self.patch_embed.parameters(): + param.requires_grad = False + + for i in range(1, self.frozen_stages + 1): + m = self.blocks[i] + m.eval() + for param in m.parameters(): + param.requires_grad = False + + if self.freeze_attn: + for i in range(0, self.depth): + m = self.blocks[i] + m.attn.eval() + m.norm1.eval() + for param in m.attn.parameters(): + param.requires_grad = False + for param in m.norm1.parameters(): + param.requires_grad = False + + if self.freeze_ffn: + self.pos_embed.requires_grad = False + self.patch_embed.eval() + for param in self.patch_embed.parameters(): + param.requires_grad = False + for i in range(0, self.depth): + m = self.blocks[i] + m.mlp.eval() + m.norm2.eval() + for param in m.mlp.parameters(): + param.requires_grad = False + for param in m.norm2.parameters(): + param.requires_grad = False + + def init_weights(self, pretrained=None): + """Initialize the weights in backbone. + Args: + pretrained (str, optional): Path to pre-trained weights. + Defaults to None. + """ + super().init_weights(pretrained, patch_padding=self.patch_padding) + + if pretrained is None: + + def _init_weights(m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=0.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + if self.cls_token is not None: + nn.init.normal_(self.cls_token, std=1e-6) + + self.apply(_init_weights) + + def get_num_layers(self): + return len(self.blocks) + + @torch.jit.ignore + def no_weight_decay(self): + return {"pos_embed", "cls_token"} + + def forward_features(self, x, **kwargs): + B, C, H, W = x.shape + x, (Hp, Wp) = self.patch_embed(x) + + if self.cls_token is not None: + cls_token = self.cls_token.expand(x.shape[0], -1, -1) + x = torch.cat((cls_token, x), dim=1) + + if self.pos_embed is not None: + # fit for multiple GPU training + # since the first element for pos embed (sin-cos manner) is zero, it will cause no difference + # x = x + self.pos_embed[:, 1:] + self.pos_embed[:, :1] + x = x + self.pos_embed + + for blk in self.blocks: + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x) + else: + x = blk(x, B) + + x = self.norm(x) + res = {} + + if self.cls_token is not None: + x_cls = x[:, :1] + else: + x_cls = torch.mean(x[:, 1:], 1, False) + res["cls"] = x_cls + + xp = x[:, 1:] + xp = xp.permute(0, 2, 1).reshape(B, -1, Hp, Wp).contiguous() + res["embed"] = xp + + return res + + def forward(self, x, **kwargs) -> Union[torch.tensor, Dict[str, torch.tensor]]: + x = self.forward_features(x, **kwargs) + return x + + def train(self, mode=True): + """Convert the model into training mode.""" + super().train(mode) + self._freeze_stages() + + +@BACKBONES.register_module() +class TimeViT(ViT): + def __init__( + self, + img_size=224, + patch_size=16, + in_chans=3, + num_classes=80, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4, + qkv_bias=False, + qk_scale=None, + drop_rate=0, + attn_drop_rate=0, + drop_path_rate=0, + hybrid_backbone=None, + norm_layer=None, + use_checkpoint=False, + frozen_stages=-1, + ratio=1, + last_norm=True, + class_token=True, + attention_type="space_only", + patch_padding="pad", + freeze_attn=False, + freeze_ffn=False, + num_frames=4, + **kwargs, + ): + super().__init__( + img_size, + patch_size, + in_chans, + num_classes, + embed_dim, + depth, + num_heads, + mlp_ratio, + qkv_bias, + qk_scale, + drop_rate, + attn_drop_rate, + drop_path_rate, + hybrid_backbone, + norm_layer, + use_checkpoint, + frozen_stages, + ratio, + last_norm, + class_token, + attention_type, + patch_padding, + freeze_attn, + freeze_ffn, + **kwargs, + ) + + norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) + register_token = kwargs.get("register_token") or False + temp_token = kwargs.get("temp_token") or False + self.low_level_enhanced = kwargs.get("low_level_enhanced") or False + self.patch_size = patch_size + + # Temporary + # self.patch_embed = PatchEmbed3D( + # img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, ratio=ratio, + # low_level=True, include_top=False, include_hm_decoder=False) + self.patch_embed = PatchEmbed3D( + img_size=img_size, + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + ratio=ratio, + ) + + self.register_token = ( + nn.Parameter(torch.zeros(1, 1, embed_dim)) if register_token else None + ) + self.temp_token = ( + nn.Parameter(torch.zeros(1, 1, embed_dim)) if temp_token else None + ) + + if self.attention_type != "space_only": + if temp_token: + self.time_embed = nn.Parameter( + torch.zeros(1, num_frames + 1, embed_dim) + ) + else: + self.time_embed = nn.Parameter(torch.zeros(1, num_frames, embed_dim)) + self.time_drop = nn.Dropout(p=drop_rate) + + self.norm_s_cls = norm_layer(embed_dim) + self.norm_t_cls = norm_layer(embed_dim) + + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, depth) + ] # stochastic depth decay rule + self.blocks = nn.ModuleList( + [ + Block2D( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[i], + norm_layer=norm_layer, + attention_type=self.attention_type, + temp_token=temp_token, + register_token=register_token, + return_s_cls_token=True, + ) + for i in range(depth) + ] + ) + + if self.low_level_enhanced: + self.lowconv3d_0 = nn.Conv3d( + 3, + embed_dim // 16, + kernel_size=(3, 4, 4), + stride=(1, 4, 4), + padding=(1, 0, 0), + bias=False, + ) # B x 48 x 4 x 56 x 56 + self.bn_low0 = nn.BatchNorm3d(embed_dim // 16, momentum=BN3D_MOMENTUM) + self.act_low0 = nn.GELU() + self.lowconv3d_01 = nn.Conv3d( + embed_dim // 16, + embed_dim // 4, + kernel_size=(3, 4, 4), + stride=(1, 4, 4), + padding=(1, 0, 0), + bias=False, + ) # B x 48 x 4 x 14 x 14 + self.bn_low01 = nn.BatchNorm3d(embed_dim // 4, momentum=BN3D_MOMENTUM) + self.act_low01 = nn.GELU() + # self.lowmaxpool3d_0 = nn.MaxPool3d(kernel_size=(1, 4, 4), stride=(1, 4, 4), padding=0) # B x 48 x 4 x 14 x 14 + self.lowfc_0 = nn.Linear(embed_dim // 4, embed_dim) + + self.lowconv3d_1 = nn.Conv3d( + embed_dim // 16, + embed_dim // 4, + kernel_size=(3, 2, 2), + stride=(1, 2, 2), + padding=(1, 0, 0), + bias=False, + ) # B x 192 x 4 x 28 x 28 + self.bn_low1 = nn.BatchNorm3d(embed_dim // 4, momentum=BN3D_MOMENTUM) + self.act_low1 = nn.GELU() + self.lowconv3d_11 = nn.Conv3d( + embed_dim // 4, + embed_dim, + kernel_size=(3, 2, 2), + stride=(1, 2, 2), + padding=(1, 0, 0), + bias=False, + ) # B x 192 x 4 x 14 x 14 + self.bn_low11 = nn.BatchNorm3d(embed_dim, momentum=BN3D_MOMENTUM) + self.act_low11 = nn.GELU() + # self.lowmaxpool3d_1 = nn.MaxPool3d(kernel_size=(1,2,2), stride=(1,2,2), padding=0) + self.lowfc_1 = nn.Linear(embed_dim, embed_dim) + + self.lowconv3d_2 = nn.Conv3d( + embed_dim // 4, + embed_dim, + kernel_size=(3, 2, 2), + stride=(1, 2, 2), + padding=(1, 0, 0), + bias=False, + ) # B x 768 x 4 x 14 x 14 + self.bn_low2 = nn.BatchNorm3d(embed_dim, momentum=BN3D_MOMENTUM) + self.act_low2 = nn.GELU() + self.lowconv3d_21 = nn.Conv3d( + embed_dim, + embed_dim, + kernel_size=(3, 3, 3), + stride=(1, 1, 1), + padding=(1, 1, 1), + bias=False, + ) # B x 768 x 4 x 14 x 14 + self.bn_low21 = nn.BatchNorm3d(embed_dim, momentum=BN3D_MOMENTUM) + self.act_low21 = nn.GELU() + # self.lowmaxpool3d_2 = nn.MaxPool3d(kernel_size=(1,1,1), stride=(1,1,1), padding=0) + self.lowfc_2 = nn.Linear(embed_dim, embed_dim) + + self.lowconv3d_3 = nn.Conv3d( + embed_dim, + embed_dim, + kernel_size=(3, 1, 1), + stride=(1, 1, 1), + padding=(1, 0, 0), + bias=False, + ) # B x 768 x 4 x 14 x 14 + self.bn_low3 = nn.BatchNorm3d(embed_dim, momentum=BN3D_MOMENTUM) + self.act_low3 = nn.GELU() + self.lowconv3d_31 = nn.Conv3d( + embed_dim, + embed_dim, + kernel_size=(3, 1, 1), + stride=(1, 1, 1), + padding=(1, 0, 0), + bias=False, + ) # B x 768 x 4 x 14 x 14 + self.bn_low31 = nn.BatchNorm3d(embed_dim, momentum=BN3D_MOMENTUM) + self.act_low31 = nn.GELU() + # self.lowmaxpool3d_3 = nn.MaxPool3d(kernel_size=(1,1,1), stride=(1,1,1), padding=0) + self.lowfc_3 = nn.Linear(embed_dim, embed_dim) + + def init_weights(self, pretrained=None): + super().init_weights(pretrained) + + ## initialization of temporal attention weights + if self.attention_type == "divided_space_time": + print("Initializing weights for temporal FC...") + i = 0 + for m in self.blocks.modules(): + m_str = str(m) + if "Block" in m_str: + if i > 0: + nn.init.constant_(m.temporal_fc.weight, 0) + nn.init.constant_(m.temporal_fc.bias, 0) + i += 1 + + if self.temp_token is not None: + nn.init.normal_(self.temp_token, std=1e-6) + + @torch.jit.ignore + def no_weight_decay(self): + return {"pos_embed", "cls_token", "time_embed"} + + def forward_features(self, x, **kwargs): + B, C, T, H, W = x.shape + + if self.low_level_enhanced: + x_low0 = self.lowconv3d_0(x) + x_low0 = self.bn_low0(x_low0) + x_low0 = self.act_low0(x_low0) + # x_low0_ = self.lowmaxpool3d_0(x_low0) + x_low0_ = self.act_low01(self.bn_low01(self.lowconv3d_01(x_low0))) + x_low0_ = x_low0_.flatten(2).transpose(1, 2) + x_low0_ = self.lowfc_0(x_low0_) + x_low0_ = rearrange( + x_low0_, + "b (t h w) m -> (b t) (h w) m", + b=B, + t=T, + h=H // self.patch_size, + w=W // self.patch_size, + ) + x_low0_ = x_low0_.sigmoid() + x_prev = x_low0 + + x, T, (Hp, Wp) = self.patch_embed(x) + + if self.low_level_enhanced: + x = x_low0_ * x + + if self.cls_token is not None: + cls_token = self.cls_token.expand(x.shape[0], -1, -1) + x = torch.cat((cls_token, x), dim=1) + + # if self.register_token is not None: + # register_token = self.register_token.expand(x.shape[0], -1, -1) + # x = torch.cat((x, register_token), dim=1) + + if self.pos_embed is not None: + # fit for multiple GPU training + # since the first element for pos embed (sin-cos manner) is zero, it will cause no difference + # x = x + self.pos_embed[:, 1:] + self.pos_embed[:, :1] + if x.size(1) != self.pos_embed.size(1): + # Resizing the pos embeds in case they do not match the input at inference + pos_embed = self.pos_embed + cls_pos_embed = pos_embed[0, 0, :].unsqueeze(0).unsqueeze(1) + other_pos_embed = pos_embed[0, 1:, :].unsqueeze(0).transpose(1, 2) + P = int(other_pos_embed.size(2) ** 0.5) + H = x.size(1) // W + other_pos_embed = other_pos_embed.reshape(1, x.size(2), P, P) + new_pos_embed = F.interpolate( + other_pos_embed, size=(H, W), mode="nearest" + ) + new_pos_embed = new_pos_embed.flatten(2) + new_pos_embed = new_pos_embed.transpose(1, 2) + new_pos_embed = torch.cat((cls_pos_embed, new_pos_embed), 1) + x = x + new_pos_embed + else: + x = x + self.pos_embed + + ## Time Embeddings + if self.attention_type != "space_only": + cls_tokens = x[:B, 0, :].unsqueeze(1) + x = x[:, 1:] + x = rearrange(x, "(b t) n m -> (b n) t m", b=B, t=T) + + if self.temp_token is not None: + temp_token = self.temp_token.expand(x.shape[0], -1, -1) + x = torch.cat((x, temp_token), dim=1) + + ## Resizing time embeddings in case they don't match + if T != self.time_embed.size(1) and self.temp_token is None: + time_embed = self.time_embed.transpose(1, 2) + new_time_embed = F.interpolate(time_embed, size=(T), mode="nearest") + new_time_embed = new_time_embed.transpose(1, 2) + x = x + new_time_embed + else: + x = x + self.time_embed + x = self.time_drop(x) + + temp_tokens = x[:B, 0, :].unsqueeze(1) + x = x[:, :-1] + x = rearrange(x, "(b n) t m -> b (n t) m", b=B, t=T) + x = torch.cat((cls_tokens, x), dim=1) + x = torch.cat((x, temp_tokens), dim=1) + + for idx, blk in enumerate(self.blocks): + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x) + else: + x, s_cls_token, t_cls_token = blk(x, B, T, Wp, **kwargs) + if self.low_level_enhanced: + if (idx + 1) % 4 == 0: + x_low = self.__getattr__(f"lowconv3d_{(idx+1)//4}")(x_prev) + # x_low_ = self.__getattr__(f'lowmaxpool3d_{(idx+1)//4}')(x_low) + x_low = self.__getattr__(f"bn_low{(idx+1)//4}")(x_low) + x_low = self.__getattr__(f"act_low{(idx+1)//4}")(x_low) + x_low_ = self.__getattr__(f"lowconv3d_{(idx+1)//4}1")(x_low) + x_low_ = self.__getattr__(f"bn_low{(idx+1)//4}1")(x_low_) + x_low_ = self.__getattr__(f"act_low{(idx+1)//4}1")(x_low_) + + x_low_ = x_low_.flatten(2).transpose(1, 2) + x_low_ = self.__getattr__(f"lowfc_{(idx+1)//4}")(x_low_) + x_low_ = x_low_.sigmoid() + x[:, 1:-1, :] = x[:, 1:-1, :].clone() * x_low_ + x_prev = x_low + + ### Predictions for space-only baseline + if self.attention_type == "space_only": + x = rearrange(x, "(b t) n m -> b t n m", b=B, t=T) + x = torch.mean(x, 1) # averaging predictions for every frame + + x = self.norm(x) + res = {} + + if self.temp_token is not None: + x_cls = x[:, -1:] + elif self.cls_token is not None: + x_cls = x[:, :1] + else: + x_cls = torch.mean(x[:, 1:-1], 1, False) + res["cls"] = x_cls + + xp = x[:, 1:-1] + xp = xp.permute(0, 2, 1).reshape(B, -1, Hp, Wp).contiguous() + + if self.attention_type != "space_only": + xp = rearrange(xp, "b (m t) h w -> b m t h w", b=B, t=T, h=Hp, w=Wp) + res["embed"] = xp + + if s_cls_token is not None: + s_cls_token = self.norm_s_cls(s_cls_token) + res["s_cls_token"] = s_cls_token + + if t_cls_token is not None: + t_cls_token = self.norm_t_cls(t_cls_token) + res["t_cls_token"] = t_cls_token + + return res diff --git a/video/fake-stormer/model_code/models/networks/backbones/xception.py b/video/fake-stormer/model_code/models/networks/backbones/xception.py new file mode 100644 index 0000000000000000000000000000000000000000..aeb8e5fa5aae9c3f7cd18ca69c9e6db8139412ab --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/backbones/xception.py @@ -0,0 +1,253 @@ +# -*- coding: utf-8 -*- +""" +Creates an Xception Model as defined in: + +Francois Chollet +Xception: Deep Learning with Depthwise Separable Convolutions +https://arxiv.org/pdf/1610.02357.pdf + +This weights ported from the Keras implementation. Achieves the following performance on the validation set: + +Loss:0.9173 Prec@1:78.892 Prec@5:94.292 + +REMEMBER to set your image size to 3x299x299 for both test and validation + +normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], + std=[0.5, 0.5, 0.5]) + +The resize parameter of the validation transform should be 333, and make sure to center crop at 299x299 +""" + +import os +import sys + +if not (os.getcwd()) in sys.path: + sys.path.append(os.getcwd()) +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 ...builder import MODELS +from ..common import BN_MOMENTUM, conv_block + +model_urls = { + "xception": "https://www.dropbox.com/s/1hplpzet9d7dv29/xception-c0a72b38.pth.tar?dl=1" +} + + +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: + 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 + + +@MODELS.register_module() +class Xception(nn.Module): + """ + Xception optimized for the ImageNet dataset, as specified in + https://arxiv.org/pdf/1610.02357.pdf + """ + + def __init__(self, num_classes=1000, **kwargs): + """Constructor + Args: + num_classes: number of classes + """ + 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) + # 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) + + 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) + + # do relu here + self.conv4 = SeparableConv2d(1536, 2048, 3, 1, 1) + self.bn4 = nn.BatchNorm2d(2048) + + # self.fc = nn.Linear(2048, num_classes) + + 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.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) + 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) + x = self.relu(x) + + x = F.adaptive_avg_pool2d(x, (1, 1)) + x = x.view(x.size(0), -1) + # x = self.fc(x) + + res = {} + res["cls"] = x + + return res + + def init_weights(self, pretrained=False): + if not pretrained: + for m in self.modules(): + if isinstance(m, nn.Conv2d): + n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + m.weight.data.normal_(0, math.sqrt(2.0 / n)) + elif isinstance(m, nn.BatchNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() + else: + state_dict = model_zoo.load_url(model_urls["xception"]) + state_dict.pop("fc.weight") + state_dict.pop("fc.bias") + self.load_state_dict(state_dict, strict=False) + + +if __name__ == "__main__": + net = Xception() + input = torch.rand((1, 3, 224, 224)) + out = net(input) + print(out["cls"].shape) diff --git a/video/fake-stormer/model_code/models/networks/common.py b/video/fake-stormer/model_code/models/networks/common.py new file mode 100644 index 0000000000000000000000000000000000000000..4c694fc6bf6459d4e4211eeae68eb9cabb99858f --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/common.py @@ -0,0 +1,286 @@ +# -*- coding: utf-8 -*- +import warnings + +import torch +import torch.nn as nn +import torch.nn.functional as F +from mmcv.cnn import build_upsample_layer + +BN_MOMENTUM = 0.1 +BN3D_MOMENTUM = 0.05 # small for small batch size + + +def point_wise_block(inplanes, outplanes): + return nn.Sequential( + nn.Conv2d( + in_channels=inplanes, + out_channels=outplanes, + kernel_size=1, + padding=0, + stride=1, + bias=False, + ), + nn.BatchNorm2d(outplanes, momentum=BN_MOMENTUM), + nn.ReLU(inplace=True), + ) + + +def conv_block(inplanes, outplanes, kernel_size, stride=1, padding=0): + return nn.Sequential( + nn.Conv2d( + in_channels=inplanes, + out_channels=outplanes, + kernel_size=kernel_size, + padding=padding, + stride=stride, + bias=False, + ), + nn.BatchNorm2d(outplanes, momentum=BN_MOMENTUM), + nn.ReLU(inplace=True), + ) + + +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 conv3d_block( + inplanes, + outplanes, + kernel_size=(3, 1, 1), + stride=(1, 1, 1), + padding=0, + bias=False, + inplace=False, + act=nn.GELU, +): + """ + General conv3d block for handling 3d feature maps + """ + return nn.Sequential( + nn.Conv3d( + inplanes, + outplanes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + bias=bias, + ), + nn.BatchNorm3d(outplanes, momentum=BN3D_MOMENTUM), + act(), + ) + + +def deconv3d_block( + inplanes, + outplanes, + kernel_size=(2, 4, 4), + stride=(2, 2, 2), + padding=(0, 1, 1), + bias=False, + inplace=False, + out_padding=(0, 1, 1), + act=None, +): + """ + General Transpose 3D Convolution for handling 3D feature maps + """ + layers = [] + layers.append( + build_upsample_layer( + dict(type="deconv3d"), + in_channels=inplanes, + out_channels=outplanes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + output_padding=out_padding, + bias=bias, + ) + ) + layers.append(nn.BatchNorm3d(outplanes)) + + if act is not None: + layers.append(act()) + + return nn.Sequential(*layers) + + +def resize( + input, + size=None, + scale_factor=None, + mode="nearest", + align_corners=None, + warning=True, +): + if warning: + if size is not None and align_corners: + input_h, input_w = tuple(int(x) for x in input.shape[2:]) + output_h, output_w = tuple(int(x) for x in size) + if output_h > input_h or output_w > output_h: + if ( + (output_h > 1 and output_w > 1 and input_h > 1 and input_w > 1) + and (output_h - 1) % (input_h - 1) + and (output_w - 1) % (input_w - 1) + ): + warnings.warn( + f"When align_corners={align_corners}, " + "the output would more aligned if " + f"input size {(input_h, input_w)} is `x+1` and " + f"out size {(output_h, output_w)} is `nx+1`" + ) + if isinstance(size, torch.Size): + size = tuple(int(x) for x in size) + return F.interpolate(input, size, scale_factor, mode, align_corners) + + +class InceptionBlock(nn.Module): + def __init__(self, inplanes, outplanes, stride=1, pool_size=3): + self.inplanes = inplanes + self.outplanes = outplanes + self.stride = stride + self.pool_size = pool_size + super(InceptionBlock, self).__init__() + + self.pw_block = point_wise_block(self.inplanes, self.outplanes // 4) + self.mp_layer = nn.MaxPool2d( + kernel_size=self.pool_size, stride=stride, padding=1 + ) + self.conv3_block = conv_block( + self.outplanes // 4, self.outplanes // 4, kernel_size=3, stride=1, padding=1 + ) + self.conv5_block = conv_block( + self.outplanes // 4, self.outplanes // 4, kernel_size=5, stride=1, padding=2 + ) + + def forward(self, x): + x1 = self.pw_block(x) + + x2 = self.pw_block(x) + x2 = self.conv3_block(x2) + + x3 = self.pw_block(x) + x3 = self.conv5_block(x3) + + x4 = self.mp_layer(x) + x4 = self.pw_block(x4) + + x = torch.cat((x1, x2, x3, x4), dim=1) + return x + + +class InceptionBlock3D(nn.Module): + def __init__(self, inplanes, outplanes, stride=1, pool_size=3): + self.inplanes = inplanes + self.outplanes = outplanes + self.stride = stride + self.pool_size = pool_size + super(InceptionBlock3D, self).__init__() + + self.pw_block = conv3d_block( + self.inplanes, + self.outplanes // 4, + kernel_size=(1, 1, 1), + stride=(1, 1, 1), + act=nn.ReLU, + ) + self.mp_layer = nn.MaxPool3d( + kernel_size=(self.pool_size, 1, 1), stride=stride, padding=(1, 0, 0) + ) + self.conv3_block = conv3d_block( + self.outplanes // 4, + self.outplanes // 4, + kernel_size=(3, 1, 1), + stride=1, + padding=(1, 0, 0), + act=nn.ReLU, + ) + self.conv5_block = conv3d_block( + self.outplanes // 4, + self.outplanes // 4, + kernel_size=(5, 1, 1), + stride=1, + padding=(2, 0, 0), + act=nn.ReLU, + ) + + def forward(self, x): + x1 = self.pw_block(x) + + x2 = self.pw_block(x) + x2 = self.conv3_block(x2) + + x3 = self.pw_block(x) + x3 = self.conv5_block(x3) + + x4 = self.mp_layer(x) + x4 = self.pw_block(x4) + + x = torch.cat((x1, x2, x3, x4), dim=1) + return x + + +class SELayer(nn.Module): + def __init__(self, channel, reduction=16): + super(SELayer, self).__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channel, channel // reduction, bias=False), + nn.ReLU(inplace=True), + nn.Linear(channel // reduction, channel, bias=False), + 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.expand_as(x) + + +class Texture_Enhance(nn.Module): + def __init__(self, num_features): + super().__init__() + # self.output_features=num_features + self.output_features = num_features * 4 + self.output_features_d = num_features + self.conv0 = nn.Conv2d(num_features, num_features, 1) + self.conv1 = nn.Conv2d(num_features, num_features, 3, padding=1) + self.bn1 = nn.BatchNorm2d(num_features) + self.conv2 = nn.Conv2d(num_features * 2, num_features, 3, padding=1) + self.bn2 = nn.BatchNorm2d(2 * num_features) + self.conv3 = nn.Conv2d(num_features * 3, num_features, 3, padding=1) + self.bn3 = nn.BatchNorm2d(3 * num_features) + self.conv_last = nn.Conv2d(num_features * 4, num_features * 4, 1) + self.bn4 = nn.BatchNorm2d(4 * num_features) + self.bn_last = nn.BatchNorm2d(num_features * 4) + + def forward(self, feature_maps, attention_maps=(1, 1)): + B, N, H, W = feature_maps.shape + + if type(attention_maps) == tuple: + attention_size = (int(H * attention_maps[0]), int(W * attention_maps[1])) + else: + attention_size = (attention_maps.shape[2], attention_maps.shape[3]) + + feature_maps_d = F.adaptive_avg_pool2d(feature_maps, attention_size) + feature_maps = feature_maps - F.interpolate( + feature_maps_d, + (feature_maps.shape[2], feature_maps.shape[3]), + mode="nearest", + ) + feature_maps0 = self.conv0(feature_maps) + feature_maps1 = self.conv1(F.relu(self.bn1(feature_maps0), inplace=True)) + feature_maps1_ = torch.cat([feature_maps0, feature_maps1], dim=1) + feature_maps2 = self.conv2(F.relu(self.bn2(feature_maps1_), inplace=True)) + feature_maps2_ = torch.cat([feature_maps1_, feature_maps2], dim=1) + feature_maps3 = self.conv3(F.relu(self.bn3(feature_maps2_), inplace=True)) + feature_maps3_ = torch.cat([feature_maps2_, feature_maps3], dim=1) + feature_maps = self.bn_last( + self.conv_last(F.relu(self.bn4(feature_maps3_), inplace=True)) + ) + return feature_maps diff --git a/video/fake-stormer/model_code/models/networks/detectors/__init__.py b/video/fake-stormer/model_code/models/networks/detectors/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b4ae7526c86df36e6ffc0adb335e13975e53def0 --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/detectors/__init__.py @@ -0,0 +1,5 @@ +# -*- coding:utf-8 -*- +from .base import BaseDetector +from .topdown import TopDownDetector + +__all__ = ["BaseDetector", "TopDownDetector"] diff --git a/video/fake-stormer/model_code/models/networks/detectors/base.py b/video/fake-stormer/model_code/models/networks/detectors/base.py new file mode 100644 index 0000000000000000000000000000000000000000..b22c4cb88eb959bf0c3b81811e9561b1a1b85a58 --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/detectors/base.py @@ -0,0 +1,145 @@ +# -*- coding: utf-8 -*- +# Copyright (c) OpenMMLab. All rights reserved. +import os +import sys + +if not os.getcwd() in sys.path: + sys.path.append(os.getcwd()) +from abc import ABCMeta, abstractmethod +from collections import OrderedDict +from typing import Tuple, Union + +import torch +import torch.distributed as dist +import torch.nn as nn +from models.builder import DETECTORS + + +@DETECTORS.register_module() +class BaseDetector(nn.Module, metaclass=ABCMeta): + """Base class for pose detectors. + + All recognizers should subclass it. + All subclass should overwrite: + Methods:`forward_train`, supporting to forward when training. + Methods:`forward_test`, supporting to forward when testing. + + Args: + backbone (dict): Backbone modules to extract feature. + head (dict): Head modules to give output. + train_cfg (dict): Config for training. Default: None. + test_cfg (dict): Config for testing. Default: None. + """ + + @abstractmethod + def forward_train(self, img, img_metas, **kwargs): + """Defines the computation performed at training.""" + + @abstractmethod + def forward_test(self, img, img_metas, **kwargs): + """Defines the computation performed at testing.""" + + @abstractmethod + def forward(self, img, img_metas, return_loss=True, **kwargs): + """Forward function.""" + + @staticmethod + def _parse_losses(losses): + """Parse the raw outputs (losses) of the network. + + Args: + losses (dict): Raw output of the network, which usually contain + losses and other necessary information. + + Returns: + tuple[Tensor, dict]: (loss, log_vars), loss is the loss tensor \ + which may be a weighted sum of all losses, log_vars \ + contains all the variables to be sent to the logger. + """ + log_vars = OrderedDict() + for loss_name, loss_value in losses.items(): + if isinstance(loss_value, torch.Tensor): + log_vars[loss_name] = loss_value.mean() + elif isinstance(loss_value, float): + log_vars[loss_name] = loss_value + elif isinstance(loss_value, list): + log_vars[loss_name] = sum(_loss.mean() for _loss in loss_value) + else: + raise TypeError( + f"{loss_name} is not a tensor or list of tensors or float" + ) + + loss = sum(_value for _key, _value in log_vars.items() if "loss" in _key) + + log_vars["loss"] = loss + for loss_name, loss_value in log_vars.items(): + # reduce loss when distributed training + if not isinstance(loss_value, float): + if dist.is_available() and dist.is_initialized(): + loss_value = loss_value.data.clone() + dist.all_reduce(loss_value.div_(dist.get_world_size())) + log_vars[loss_name] = loss_value.item() + else: + log_vars[loss_name] = loss_value + + return loss, log_vars + + def train_step(self, data_batch, optimizer, **kwargs): + """The iteration step during training. + + This method defines an iteration step during training, except for the + back propagation and optimizer updating, which are done in an optimizer + hook. Note that in some complicated cases or models, the whole process + including back propagation and optimizer updating is also defined in + this method, such as GAN. + + Args: + data_batch (dict): The output of dataloader. + optimizer (:obj:`torch.optim.Optimizer` | dict): The optimizer of + runner is passed to ``train_step()``. This argument is unused + and reserved. + + Returns: + dict: It should contain at least 3 keys: ``loss``, ``log_vars``, + ``num_samples``. + ``loss`` is a tensor for back propagation, which can be a + weighted sum of multiple losses. + ``log_vars`` contains all the variables to be sent to the + logger. + ``num_samples`` indicates the batch size (when the model is + DDP, it means the batch size on each GPU), which is used for + averaging the logs. + """ + losses = self.forward(**data_batch) + + loss, log_vars = self._parse_losses(losses) + + outputs = dict( + loss=loss, + log_vars=log_vars, + num_samples=len(next(iter(data_batch.values()))), + ) + + return outputs + + def val_step(self, data_batch, optimizer, **kwargs): + """The iteration step during validation. + + This method shares the same signature as :func:`train_step`, but used + during val epochs. Note that the evaluation after training epochs is + not implemented with this method, but an evaluation hook. + """ + results = self.forward(return_loss=False, **data_batch) + + outputs = dict(results=results) + + return outputs + + # @abstractmethod + # def show_result(self, **kwargs): + # """Visualize the results.""" + # raise NotImplementedError + + +if __name__ == "__main__": + print(BaseDetector._version) diff --git a/video/fake-stormer/model_code/models/networks/detectors/topdown.py b/video/fake-stormer/model_code/models/networks/detectors/topdown.py new file mode 100644 index 0000000000000000000000000000000000000000..74f523d8c1b22d626d3f8909dc4ddd8a48fc50be --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/detectors/topdown.py @@ -0,0 +1,174 @@ +# -*- coding: utf-8 -*- +# Copyright (c) OpenMMLab. All rights reserved. +import warnings + +import mmcv +import numpy as np + +try: + from mmcv.runner import auto_fp16 +except: + raise ValueError("Please try to install mmcv==1.3.9") + +from ...builder import BACKBONES, DETECTORS, HEADS, NECKS +from .base import BaseDetector + + +@DETECTORS.register_module() +class TopDownDetector(BaseDetector): + """Top-down pose detectors. + + Args: + backbone (dict): Backbone modules to extract feature. + keypoint_head (dict): Keypoint head to process feature. + train_cfg (dict): Config for training. Default: None. + test_cfg (dict): Config for testing. Default: None. + pretrained (str): Path to the pretrained models. + loss_pose (None): Deprecated arguments. Please use + `loss_keypoint` for heads instead. + """ + + def __init__( + self, + backbone, + neck=None, + keypoint_head=None, + train_cfg=None, + test_cfg=None, + loss_pose=None, + **kwargs, + ): + super().__init__() + self.fp16_enabled = False + + self.backbone = BACKBONES.build(backbone) + + self.train_cfg = train_cfg + self.test_cfg = test_cfg + + if neck is not None: + self.neck = NECKS.build(neck) + + if keypoint_head is not None: + keypoint_head["train_cfg"] = train_cfg + keypoint_head["test_cfg"] = test_cfg + + if "loss_keypoint" not in keypoint_head and loss_pose is not None: + warnings.warn( + "`loss_pose` for TopDown is deprecated, " + "use `loss_keypoint` for heads instead. See " + "https://github.com/open-mmlab/mmpose/pull/382" + " for more information.", + DeprecationWarning, + ) + keypoint_head["loss_keypoint"] = loss_pose + + self.keypoint_head = HEADS.build(keypoint_head) + + @property + def with_neck(self): + """Check if has neck.""" + return hasattr(self, "neck") + + @property + def with_keypoint(self): + """Check if has keypoint_head.""" + return hasattr(self, "keypoint_head") + + def init_weights(self, pretrained=None): + """Weight initialization for model.""" + self.backbone.init_weights(pretrained=pretrained) + if self.with_neck: + self.neck.init_weights() + if self.with_keypoint: + self.keypoint_head.init_weights() + + def forward(self, img, **kwargs): + output = self.backbone(img, **kwargs) + + if self.with_neck: + output = self.neck(output, **kwargs) + + if self.with_keypoint: + output = self.keypoint_head(output, **kwargs) + + return output + + def forward_train(self, img, target, target_weight, img_metas, **kwargs): + """Defines the computation performed at every call when training.""" + output = self.backbone(img) + if self.with_neck: + output = self.neck(output) + if self.with_keypoint: + output = self.keypoint_head(output) + + # if return loss + losses = dict() + if self.with_keypoint: + keypoint_losses = self.keypoint_head.get_loss(output, target, target_weight) + losses.update(keypoint_losses) + keypoint_accuracy = self.keypoint_head.get_accuracy( + output, target, target_weight + ) + losses.update(keypoint_accuracy) + + return losses + + def forward_test(self, img, img_metas, return_heatmap=False, **kwargs): + """Defines the computation performed at every call when testing.""" + assert img.size(0) == len(img_metas) + batch_size, _, img_height, img_width = img.shape + if batch_size > 1: + assert "bbox_id" in img_metas[0] + + result = {} + + features = self.backbone(img) + if self.with_neck: + features = self.neck(features) + if self.with_keypoint: + output_heatmap = self.keypoint_head.inference_model( + features, flip_pairs=None + ) + + if self.test_cfg.get("flip_test", True): + img_flipped = img.flip(3) + features_flipped = self.backbone(img_flipped) + if self.with_neck: + features_flipped = self.neck(features_flipped) + if self.with_keypoint: + output_flipped_heatmap = self.keypoint_head.inference_model( + features_flipped, img_metas[0]["flip_pairs"] + ) + output_heatmap = (output_heatmap + output_flipped_heatmap) * 0.5 + + if self.with_keypoint: + keypoint_result = self.keypoint_head.decode( + img_metas, output_heatmap, img_size=[img_width, img_height] + ) + result.update(keypoint_result) + + if not return_heatmap: + output_heatmap = None + + result["output_heatmap"] = output_heatmap + + return result + + def forward_dummy(self, img): + """Used for computing network FLOPs. + + See ``tools/get_flops.py``. + + Args: + img (torch.Tensor): Input image. + + Returns: + Tensor: Output heatmaps. + """ + output = self.backbone(img) + if self.with_neck: + output = self.neck(output) + if self.with_keypoint: + output = self.keypoint_head(output) + return output diff --git a/video/fake-stormer/model_code/models/networks/heads/__init__.py b/video/fake-stormer/model_code/models/networks/heads/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d75583dbcccfdb02545e8181f9bc39cf7b71f34f --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/heads/__init__.py @@ -0,0 +1,5 @@ +# -*- coding: utf-8 -*- +from .hm_base_head import TopdownHeatmapBaseHead +from .hm_simple_head import TopdownHeatmapSimpleHead + +__all__ = ["TopdownHeatmapBaseHead", "TopdownHeatmapSimpleHead"] diff --git a/video/fake-stormer/model_code/models/networks/heads/head_design.py b/video/fake-stormer/model_code/models/networks/heads/head_design.py new file mode 100644 index 0000000000000000000000000000000000000000..2553470a89638b864507cdfac72271148a2ec06e --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/heads/head_design.py @@ -0,0 +1,225 @@ +# -*- coding: utf-8 -*- +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn as nn + +from ..common import conv3d_block, conv_block + + +class ClassificationHead(nn.Module): + def __init__( + self, + in_planes, + out_planes, + stages: Optional[List] = [], + return_prob: bool = False, + last_act: str = "sigmoid", + drop: float = 0.0, + **kwargs, + ): + """ + General classification head + + Args: + stages: indicate a list of outputs of hidden layers between in_planes and out_planes + """ + super().__init__() + self.return_prob = return_prob + self.drop = drop + self.in_planes = in_planes + self.global_avg = kwargs.get("avg_pool") or False + + if self.global_avg: + if kwargs.get("features") == "3D": + self.avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1)) + else: + self.avg_pool = nn.AdaptiveAvgPool3d((1, 1)) + + # Initialize layers + layers = [] + for stage_plane in stages: + layers.append(nn.Linear(in_planes, stage_plane)) + layers.append(nn.GELU()) + layers.append(nn.Dropout(self.drop)) + in_planes = stage_plane + + self.fc1 = nn.Sequential(*layers) if len(layers) else nn.Identity() + self.fc_out = nn.Linear(self.in_planes, out_planes) + + if last_act == "sigmoid": + self.act = nn.Sigmoid() + else: + self.act = nn.Softmax(dim=-1) + + def forward(self, x) -> torch.tensor: + B = x.shape[0] + + if self.global_avg: + num_channels = x.shape[1] + x = self.avg_pool(x).view(B, -1, num_channels) + + assert self.in_planes == x.shape[-1] + x = self.fc1(x) + + if x.ndim == 3: + x = x.reshape((B, -1)) + + x = self.fc_out(x) + + if self.return_prob: + results = self.act(x) + else: + results = x + + return results + + +class RegressionHead(nn.Module): + def __init__( + self, + input_shape, + out_planes, + kernel_size, + padding, + conv_2direction=False, + extra=None, + **kwargs, + ): + """ + General regression head for dense predictions like keypoints regression, etc + + Args: + input_shape: [C, H, W] + conv_2direction: Specially design to compute derivative tx, ty from temporal tokens in video input. i.e. FakeSTormer + """ + super().__init__() + self.conv_2direction = conv_2direction + + conv_channels = input_shape[0] + + layers = [] + if extra is not None: + num_conv_layers = extra.get("num_conv_layers", 0) + num_conv_kernels = extra.get("num_conv_kernels", [1] * num_conv_layers) + + for i in range(num_conv_layers): + layers.append( + conv_block( + inplanes=conv_channels, + outplanes=conv_channels, + kernel_size=kernel_size, + padding=padding, + ) + ) + self.before_proj = nn.Sequential(*layers) if len(layers) else nn.Identity() + + if not conv_2direction: + self.proj = nn.Conv2d( + in_channels=input_shape[0], + out_channels=out_planes, + kernel_size=kernel_size, + padding=padding, + ) + else: + self.d_tx = conv_block( + inplanes=input_shape[1], + outplanes=input_shape[1], + kernel_size=kernel_size, + padding=padding, + ) + + self.d_ty = conv_block( + inplanes=input_shape[2], + outplanes=input_shape[2], + kernel_size=kernel_size, + padding=padding, + ) + + self.proj = nn.Conv2d( + in_channels=int(input_shape[0] * 2), + out_channels=out_planes, + kernel_size=kernel_size, + padding=padding, + ) + + def forward(self, x): + assert x.ndim == 4 + B, C, H, W = x.shape + x = self.before_proj(x) + + if self.conv_2direction: + x_tx = self.d_tx(x.transpose(1, 2)).transpose(2, 1) + x_ty = self.d_ty(x.transpose(1, 3)).transpose(3, 1) + x = torch.cat((x_tx, x_ty), 1) + x = self.proj(x) + + return x + + +class TemporalRegressionHead(nn.Module): + def __init__( + self, + inplanes: int, + outplanes: int, + kernel_size: Union[int, Tuple] = (3, 1, 1), + stride: Union[int, Tuple] = (1, 1, 1), + padding: Union[int, Tuple] = (1, 0, 0), + extra: Optional[dict] = None, + act: nn.Module = nn.GELU, + **kwargs, + ): + super().__init__(**kwargs) + """ + General head design for 3D feature maps + args: + inplanes: number of in channels + outplanes: number of out channels + extra: doing extra operation before predicting final outputs + """ + layers = [] + if extra is not None: + num_conv_layers = extra.get("num_conv_layers", 0) + num_conv_kernels = extra.get("num_conv_kernels", [1] * num_conv_layers) + + for i in range(num_conv_layers): + layers.append( + conv3d_block( + inplanes=inplanes, + outplanes=inplanes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + bias=True, + inplace=True, + act=act, + ) + ) + self.before_proj = nn.Sequential(*layers) if len(layers) else nn.Identity() + self.proj = nn.Sequential( + conv3d_block( + inplanes=inplanes, + outplanes=inplanes // 4, + stride=stride, + kernel_size=kernel_size, + padding=padding, + bias=True, + inplace=True, + act=act, + ), + nn.Conv3d( + inplanes // 4, + outplanes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + bias=True, + ), + ) + + def forward(self, x): + assert x.ndim == 5 + + x = self.before_proj(x) + out = self.proj(x) + return out diff --git a/video/fake-stormer/model_code/models/networks/heads/hm_base_head.py b/video/fake-stormer/model_code/models/networks/heads/hm_base_head.py new file mode 100644 index 0000000000000000000000000000000000000000..6489d4f32097042ea0746403ddf48c938a1647fe --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/heads/hm_base_head.py @@ -0,0 +1,54 @@ +# Copyright (c) OpenMMLab. All rights reserved. +# -*- coding: utf-8 -*- +from abc import ABCMeta, abstractmethod + +import numpy as np +import torch.nn as nn + + +class TopdownHeatmapBaseHead(nn.Module): + """Base class for top-down heatmap heads. + + All top-down heatmap heads should subclass it. + All subclass should overwrite: + + Methods:`loss`, supporting to calculate loss. + Methods:`accuracy`, supporting to calculate accuracy. + Methods:`forward`, supporting to forward model. + Methods:`predict`, supporting to inference model. + """ + + __metaclass__ = ABCMeta + + @abstractmethod + def loss(self, **kwargs): + """Gets the loss.""" + + @abstractmethod + def accuracy(self, **kwargs): + """Gets the accuracy.""" + + @abstractmethod + def forward(self, **kwargs): + """Forward function.""" + + @abstractmethod + def predict(self, **kwargs): + """Inference function.""" + + @staticmethod + def _get_deconv_cfg(deconv_kernel): + """Get configurations for deconv layers.""" + if deconv_kernel == 4: + padding = 1 + output_padding = 0 + elif deconv_kernel == 3: + padding = 1 + output_padding = 1 + elif deconv_kernel == 2: + padding = 0 + output_padding = 0 + else: + raise ValueError(f"Not supported num_kernels ({deconv_kernel}).") + + return deconv_kernel, padding, output_padding diff --git a/video/fake-stormer/model_code/models/networks/heads/hm_simple_head.py b/video/fake-stormer/model_code/models/networks/heads/hm_simple_head.py new file mode 100644 index 0000000000000000000000000000000000000000..c27c331e9e13a12d19e15561aebe50a45f90a4e7 --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/heads/hm_simple_head.py @@ -0,0 +1,378 @@ +# Copyright (c) OpenMMLab. All rights reserved. +# -*- coding: utf-8 -*- +import os +import sys +from typing import Dict, Union + +if not os.getcwd() in sys.path: + sys.path.append(os.getcwd()) + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange, reduce, repeat +from losses import LOSSES, build_losses +from mmcv.cnn import ( + build_conv_layer, + build_norm_layer, + build_upsample_layer, + constant_init, + normal_init, +) +from models.builder import HEADS + +from ..common import BN_MOMENTUM, conv_block, resize +from .head_design import ClassificationHead, RegressionHead, TemporalRegressionHead +from .hm_base_head import TopdownHeatmapBaseHead + + +@HEADS.register_module() +class TopdownHeatmapSimpleHead(TopdownHeatmapBaseHead): + """Top-down heatmap simple head. paper ref: Bin Xiao et al. ``Simple + Baselines for Human Pose Estimation and Tracking``. + + TopdownHeatmapSimpleHead is consisted of (>=0) number of deconv layers + and a simple conv2d layer. + + Args: + in_channels (int): Number of input channels + out_channels (int): Number of output channels + num_deconv_layers (int): Number of deconv layers. + num_deconv_layers should >= 0. Note that 0 means + no deconv layers. + num_deconv_filters (list|tuple): Number of filters. + If num_deconv_layers > 0, the length of + num_deconv_kernels (list|tuple): Kernel sizes. + in_index (int|Sequence[int]): Input feature index. Default: 0 + input_transform (str|None): Transformation type of input features. + Options: 'resize_concat', 'multiple_select', None. + Default: None. + + - 'resize_concat': Multiple feature maps will be resized to the + same size as the first one and then concat together. + Usually used in FCN head of HRNet. + - 'multiple_select': Multiple feature maps will be bundle into + a list and passed into decode head. + - None: Only one select feature map is allowed. + align_corners (bool): align_corners argument of F.interpolate. + Default: False. + loss_keypoint (dict): Config for keypoint loss. Default: None. + """ + + def __init__( + self, + in_channels, + heads, + extra=None, + hm_size=[14, 14], + in_index=0, + input_transform=None, + align_corners=False, + loss_keypoint=None, + train_cfg=None, + test_cfg=None, + upsample=0, + conv_2direction=False, + use_temp_token=False, + features="2D", + act="GELU", + **kwargs, + ): + super().__init__() + + self.in_channels = in_channels + self.loss = build_losses(loss_keypoint, LOSSES) + self.upsample = upsample + self.use_temp_token = use_temp_token + + self.train_cfg = {} if train_cfg is None else train_cfg + self.test_cfg = {} if test_cfg is None else test_cfg + self.target_type = self.test_cfg.get("target_type", "GaussianHeatmap") + + self._init_inputs(in_channels, in_index, input_transform) + self.in_index = in_index + self.align_corners = align_corners + + assert ( + isinstance(heads, dict) and heads is not None + ), "Head config can not be None!" + self.heads = heads + + if extra is not None and not isinstance(extra, dict): + raise TypeError("extra should be dict or None.") + + identity_final_layer = False + if extra is not None and "final_conv_kernel" in extra: + assert extra["final_conv_kernel"] in [0, 1, 3] + if extra["final_conv_kernel"] == 3: + padding = 1 + elif extra["final_conv_kernel"] == 1: + padding = 0 + else: + # 0 for Identity mapping. + identity_final_layer = True + kernel_size = extra["final_conv_kernel"] + else: + kernel_size = 1 + padding = 0 + + final_layer = {} + if not identity_final_layer: + for head, out_channel in self.heads.items(): + if head == "hm" or head == "cstency": + if features == "2D": + layer = RegressionHead( + input_shape=(self.in_channels, hm_size[0], hm_size[1]), + out_planes=out_channel, + kernel_size=kernel_size, + padding=padding, + extra=extra, + conv_2direction=conv_2direction, + ) + elif features == "3D": + act_func = nn.GELU if act == "GELU" else nn.ReLU + layer = TemporalRegressionHead( + inplanes=self.in_channels, + outplanes=out_channel, + extra=extra, + act=act_func, + ) + else: + raise ValueError( + "Only support 2D or 3D features, please check your feature dimension!" + ) + final_layer[head] = layer + elif head == "cls": + layer = ClassificationHead( + in_planes=self.in_channels, + out_planes=out_channel, + return_prob=False, + last_act="sigmoid", + drop=0.2, + features=features, + avg_pool=kwargs.get("avg_pool"), + ) + final_layer[head] = layer + elif head == "temp_loc": + layer = ClassificationHead( + in_planes=self.in_channels, + out_planes=out_channel, + stages=[self.in_channels // out_channel], + return_prob=False, + last_act="sigmoid", + drop=0.2, + ) + final_layer[head] = layer + else: + final_layer["cls"] = nn.Identity() + + for k, l in final_layer.items(): + self.__setattr__(k, l) + + def _init_inputs(self, in_channels, in_index, input_transform): + """Check and initialize input transforms. + + The in_channels, in_index and input_transform must match. + Specifically, when input_transform is None, only single feature map + will be selected. So in_channels and in_index must be of type int. + When input_transform is not None, in_channels and in_index must be + list or tuple, with the same length. + + Args: + in_channels (int|Sequence[int]): Input channels. + in_index (int|Sequence[int]): Input feature index. + input_transform (str|None): Transformation type of input features. + Options: 'resize_concat', 'multiple_select', None. + + - 'resize_concat': Multiple feature maps will be resize to the + same size as first one and than concat together. + Usually used in FCN head of HRNet. + - 'multiple_select': Multiple feature maps will be bundle into + a list and passed into decode head. + - None: Only one select feature map is allowed. + """ + + if input_transform is not None: + assert input_transform in ["resize_concat", "multiple_select"] + self.input_transform = input_transform + self.in_index = in_index + if input_transform is not None: + assert isinstance(in_channels, (list, tuple)) + assert isinstance(in_index, (list, tuple)) + assert len(in_channels) == len(in_index) + if input_transform == "resize_concat": + self.in_channels = sum(in_channels) + else: + self.in_channels = in_channels + else: + assert isinstance(in_channels, int) + assert isinstance(in_index, int) + self.in_channels = in_channels + + def _transform_inputs(self, inputs: Union[torch.tensor, Dict[str, torch.tensor]]): + """Transform inputs for decoder. + + Args: + inputs (list[Tensor] | Tensor): multi-level img features. + + Returns: + Tensor: The transformed inputs + """ + additional_inputs = {} + + if isinstance(inputs, dict): + # assert 'embed' in inputs.keys(), "Embed token must be present in the input dict key" + for k in inputs.keys(): + if k != "embed": + additional_inputs[k] = inputs[k] + inputs = inputs["embed"] if "embed" in inputs.keys() else None + + if not isinstance(inputs, list): + if not isinstance(inputs, list): + if self.upsample > 0: + inputs = resize( + input=F.relu(inputs), + scale_factor=self.upsample, + mode="bilinear", + align_corners=self.align_corners, + ) + return inputs, additional_inputs + + if self.input_transform == "resize_concat": + inputs = [inputs[i] for i in self.in_index] + upsampled_inputs = [ + resize( + input=x, + size=inputs[0].shape[2:], + mode="bilinear", + align_corners=self.align_corners, + ) + for x in inputs + ] + inputs = torch.cat(upsampled_inputs, dim=1) + elif self.input_transform == "multiple_select": + inputs = [inputs[i] for i in self.in_index] + else: + inputs = inputs[self.in_index] + + return inputs, additional_inputs + + def _make_deconv_layer(self, num_layers, num_filters, num_kernels): + """Make deconv layers.""" + if num_layers != len(num_filters): + error_msg = ( + f"num_layers({num_layers}) " + f"!= length of num_filters({len(num_filters)})" + ) + raise ValueError(error_msg) + if num_layers != len(num_kernels): + error_msg = ( + f"num_layers({num_layers}) " + f"!= length of num_kernels({len(num_kernels)})" + ) + raise ValueError(error_msg) + + layers = [] + for i in range(num_layers): + kernel, padding, output_padding = self._get_deconv_cfg(num_kernels[i]) + + planes = num_filters[i] + layers.append( + build_upsample_layer( + dict(type="deconv"), + in_channels=self.in_channels, + out_channels=planes, + kernel_size=kernel, + stride=2, + padding=padding, + output_padding=output_padding, + bias=False, + ) + ) + layers.append(nn.BatchNorm2d(planes)) + layers.append(nn.ReLU(inplace=True)) + self.in_channels = planes + + return nn.Sequential(*layers) + + def init_weights(self): + """Initialize model weights.""" + for head in self.heads.keys(): + for m in self.__getattr__(head).modules(): + if isinstance(m, nn.Conv2d): + normal_init(m, std=0.001, bias=0) + elif isinstance(m, nn.BatchNorm2d): + constant_init(m, 1) + + def forward(self, x, **kwargs): + """Forward function. + The input is multiscale feature maps and the output is the heatmap without post processing + """ + x_embed, additional_inputs = self._transform_inputs(x) + if "temp_loc" in self.heads.keys(): + if not self.use_temp_token: + B, L, T, H, W = x_embed.shape + x_temp_loc = x_embed.reshape(B, T, L, -1) + x_temp_loc = torch.max(x_temp_loc, dim=3, keepdim=False)[0] + else: + x_temp_loc = additional_inputs["s_cls_token"] + + x_outs = {} + # for head in self.heads.keys(): + if "cls" in self.heads.keys(): + assert hasattr( + self, "cls" + ), "There must be a Classification Head, please check the head design!" + if bool(additional_inputs) and "cls" in additional_inputs.keys(): + x_outs["cls"] = self.__getattr__("cls")(additional_inputs["cls"]) + else: + x_outs["cls"] = self.__getattr__("cls")(x_embed) + + # for head in self.heads.keys(): + if "temp_loc" in self.heads.keys(): + assert hasattr( + self, "temp_loc" + ), "There must be a head for temporal localization, please check the head design!" + x_outs["temp_loc"] = self.__getattr__("temp_loc")(x_temp_loc) + + if "hm" in self.heads.keys(): + assert hasattr( + self, "hm" + ), "There must always be a Heatmap Head in the head!" + x_outs["hm"] = self.__getattr__("hm")(x_embed) + + if "cstency" in self.heads.keys(): + assert hasattr( + self, "cstency" + ), "There must always be a Consistency Head in the head!" + x_outs["cstency"] = self.__getattr__("cstency")(x_embed) + + return [x_outs] + + def get_loss(self, output, target, target_weight): + """Calculate top-down keypoint loss. + + Note: + - batch_size: N + - num_keypoints: K + - heatmaps height: H + - heatmaps weight: W + + Args: + output (torch.Tensor[N,K,H,W]): Output heatmaps. + target (torch.Tensor[N,K,H,W]): Target heatmaps. + target_weight (torch.Tensor[N,K,1]): + Weights across different joint types. + """ + + losses = dict() + + assert not isinstance(self.loss, nn.Sequential) + assert target.dim() == 4 and target_weight.dim() == 3 + losses["heatmap_loss"] = self.loss(output, target, target_weight) + + return losses + + +if __name__ == "__main__": + cfg = {} diff --git a/video/fake-stormer/model_code/models/networks/mrsa_resnet.py b/video/fake-stormer/model_code/models/networks/mrsa_resnet.py new file mode 100644 index 0000000000000000000000000000000000000000..38172a90c9faf9be5a10bc1d2adee9cc65a5c4af --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/mrsa_resnet.py @@ -0,0 +1,490 @@ +# -*- coding: utf-8 -*- +from __future__ import absolute_import, division, print_function + +import math +import os + +import torch +import torch.nn as nn +import torch.utils.model_zoo as model_zoo +from torch.nn.modules.activation import ReLU +from torch.nn.modules.batchnorm import BatchNorm2d +from torch.nn.modules.pooling import MaxPool2d + +from ..builder import MODELS, build_model +from .common import ( + BN_MOMENTUM, + InceptionBlock, + conv_block, + point_wise_block, +) + +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 = 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 + + @staticmethod + def __repr__(): + return "BasicBlock" + + +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 + + @staticmethod + def __repr__(): + return "Bottleneck" + + +@MODELS.register_module() +class PoseResNet(nn.Module): + def __init__( + self, + block, + layers, + heads, + head_conv, + dropout_prob, + fpn=False, + cls_based_hm=True, + use_c2=False, + **kwargs, + ): + self.inplanes = 64 + self.deconv_with_bias = False + self.heads = heads + self.fpn = fpn + self.cls_based_hm = cls_based_hm + self.use_c2 = use_c2 + + # Convert Cls name into Cls Object + if isinstance(block, str): + for bl in [BasicBlock, Bottleneck]: + if block == bl.__repr__(): + block = bl + + for k, v in kwargs.items(): + if v is None: + raise ValueError( + f"The {k} argument receive a None value, Please check!" + ) + self.__setattr__(k, v) + + super(PoseResNet, self).__init__() + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False) + self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM) + 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) + + # Custom dropout layer + self.dropout_layer = nn.Dropout(dropout_prob) + + if self.fpn: + # Adding sidmoid layer + self.sigmoid_layer = nn.Sigmoid() + + # Adding pointwise block + self.pw_block_1 = self._point_wise_block(2048, 1024) + + # used for deconv layers + deconv_filters = [256, 128, 256] if self.fpn else [256, 256, 256] + self.deconv_layers = self._make_deconv_layer( + 3, + deconv_filters, + [4, 4, 4], + ) + + # Adding inception block + if self.fpn: + for idx, deconv_layer in enumerate(self.deconv_layers): + self.__setattr__(f"deconv_layer_{idx}", nn.Sequential(deconv_layer)) + self.pw_block_2 = self._point_wise_block(512, 512) + if self.use_c2: + self.pw_block_3 = self._point_wise_block(512, 256) + self.pw_block_c3 = self._point_wise_block(1024, 256) + self.pw_block_c2 = self._point_wise_block(512, 128) + self.inception_block = InceptionBlock(256, 256, stride=1, pool_size=3) + + for head in sorted(self.heads): + num_output = self.heads[head] + if head_conv > 0: + if head != "cls": + fc = nn.Sequential( + nn.Conv2d(256, head_conv, kernel_size=3, padding=1, bias=True), + nn.BatchNorm2d(head_conv), + nn.ReLU(inplace=True), + nn.Conv2d( + head_conv, num_output, kernel_size=1, stride=1, padding=0 + ), + ) + else: + if self.cls_based_hm: + fc = nn.Sequential( + nn.AdaptiveMaxPool2d(head_conv // 4), + nn.Flatten(), + nn.Linear( + num_output * ((head_conv // 4) ** 2), + head_conv, + bias=True, + ), + nn.BatchNorm1d(head_conv, momentum=BN_MOMENTUM), + nn.ReLU(inplace=True), + nn.Linear(head_conv, 1, bias=True), + nn.Sigmoid(), + ) + else: + fc = nn.Sequential( + nn.Conv2d( + 256, head_conv, kernel_size=3, padding=1, bias=True + ), + nn.BatchNorm2d(head_conv, momentum=BN_MOMENTUM), + nn.ReLU(inplace=True), + # nn.Conv2d(head_conv, num_output, kernel_size=1, + # stride=1, padding=0, bias=True), + # nn.BatchNorm2d(num_output), + # nn.ReLU(inplace=True), + # nn.AdaptiveMaxPool2d(head_conv//4), + nn.AdaptiveAvgPool2d(1), + nn.Flatten(), + # nn.Linear((head_conv//4)**2, head_conv, bias=True), + # nn.BatchNorm1d(head_conv, momentum=BN_MOMENTUM), + # nn.ReLU(inplace=True), + nn.Linear(head_conv, 1, bias=True), + # nn.Sigmoid() + ) + else: + fc = nn.Conv2d( + in_channels=256, + out_channels=num_output, + kernel_size=1, + stride=1, + padding=0, + ) + self.__setattr__(head, fc) + + def _point_wise_block(self, inplanes, outplanes): + self.inplanes = outplanes + module = point_wise_block(inplanes, outplanes) + return module + + def _conv_block(self, inplanes, outplanes, kernel_size, stride=1): + self.inplanes = outplanes + module = conv_block(inplanes, outplanes, kernel_size=kernel_size, stride=stride) + return module + + 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, momentum=BN_MOMENTUM), + ) + + 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 _get_deconv_cfg(self, deconv_kernel, index): + if deconv_kernel == 4: + padding = 1 + output_padding = 0 + elif deconv_kernel == 3: + padding = 1 + output_padding = 1 + elif deconv_kernel == 2: + padding = 0 + output_padding = 0 + + return deconv_kernel, padding, output_padding + + def _make_deconv_layer(self, num_layers, num_filters, num_kernels): + assert num_layers == len( + num_filters + ), "ERROR: num_deconv_layers is different len(num_deconv_filters)" + assert num_layers == len( + num_kernels + ), "ERROR: num_deconv_layers is different len(num_deconv_filters)" + + layers = [] + for i in range(num_layers): + kernel, padding, output_padding = self._get_deconv_cfg(num_kernels[i], i) + + planes = num_filters[i] + layers.append( + nn.Sequential( + nn.ConvTranspose2d( + in_channels=self.inplanes, + out_channels=planes, + kernel_size=kernel, + stride=2, + padding=padding, + output_padding=output_padding, + bias=self.deconv_with_bias, + ), + nn.BatchNorm2d(planes, momentum=BN_MOMENTUM), + ) + ) + if not self.fpn: + layers.append(nn.ReLU(inplace=True)) + + self.inplanes = planes if not self.fpn else planes * 2 + + if self.fpn: + return layers + else: + return nn.Sequential(*layers) + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) + + x1 = self.layer1(x) # 256 x 64 x 64 + x2 = self.layer2(x1) # 512 x 32 x 32 + x3 = self.layer3(x2) # 1024 x 16 x 16 + x4 = self.layer4(x3) # 2048 x 8 x 8 + + # Custom dropout layer + x = self.dropout_layer(x4) # B x 8 x 8 x 2048 + x3 = self.dropout_layer(x3) + x2 = self.dropout_layer(x2) + x1 = self.dropout_layer(x1) + + # Custom FPN + if self.fpn: + assert isinstance( + self.deconv_layers, list + ), "To custom FPN, decompose deconv layers as a list!" + x = self.pw_block_1(x) # B x 1024 x 8 x 8 + x = self.deconv_layer_0(x) # B x 256 x 16 x 16 + # x = self.relu(x) # B x 256 x 16 x 16 + + x_weighted = self.sigmoid_layer(x) # B x 256 x 16 x 16 + x_inverse = torch.sub(1, x_weighted, alpha=1) # B x 256 x 16 x 16 + x3 = self.pw_block_c3(x3) # B x 256 x 16 x 16 + x3_ = torch.multiply(x3, x_inverse) # B x 256 x 16 x 16 + x = torch.cat((x, x3_), dim=1) # B x 512 x 16 x 16 + + x = self.pw_block_2(x) # B x 512 x 16 x 16 + x = self.deconv_layer_1(x) # B x 128 x 32 x 32 + # x = self.relu(x) #B x 128 x 32 x 32 + + x_weighted = self.sigmoid_layer(x) # B x 128 x 32 x 32 + x_inverse = torch.sub(1, x_weighted, alpha=1) # B x 128 x 32 x 32 + x2 = self.pw_block_c2(x2) + x2_ = torch.multiply(x2, x_inverse) # B x 128 x 32 x 32 + x = torch.cat((x, x2_), dim=1) # B x 256 x 32 x 32 + + x = self.inception_block(x) # B x 256 x 64 x 64 + x = self.deconv_layer_2(x) # B x 256 x 64 x 64 + + if self.use_c2: + x_weighted = self.sigmoid_layer(x) + x_inverse = torch.sub(1, x_weighted, alpha=1) + x1_ = torch.multiply(x1, x_inverse) + x = torch.cat((x, x1_), dim=1) + x = self.pw_block_3(x) + else: + x = self.relu(x) # B x 256 x 64 x 64 + else: + assert isinstance( + self.deconv_layers, nn.Module + ), "Deconv Layer must be nn Module to compute!" + x = self.deconv_layers(x) + + ret = {} + x1_hm = None + for head in self.heads: + if self.cls_based_hm and head == "cls" and x1_hm is not None: + x = x1_hm + elif head == "hm": + x1_hm = x + + ret[head] = self.__getattr__(head)(x) + + return [ret] + + def init_weights(self, pretrained=True, **kwargs): + num_layers = kwargs.get("num_layers") + if pretrained: + if self.fpn: + for bl in [self.pw_block_1, self.pw_block_2]: + for _, l in bl.named_parameters(): + if isinstance(l, nn.Conv2d): + nn.init.normal_(l.weight, std=0.001) + nn.init.constant_(l.bias, 0) + + for _, l in self.inception_block.named_parameters(): + if isinstance(l, nn.Conv2d): + nn.init.normal_(l.weight, std=0.001) + nn.init.constant_(l.bias, 0) + + # print('=> init resnet deconv weights from normal distribution') + if isinstance(self.deconv_layers, nn.Module): + for _, m in self.deconv_layers.named_modules(): + if isinstance(m, nn.ConvTranspose2d): + # print('=> init {}.weight as normal(0, 0.001)'.format(name)) + # print('=> init {}.bias as 0'.format(name)) + nn.init.normal_(m.weight, std=0.001) + if self.deconv_with_bias: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.BatchNorm2d): + # print('=> init {}.weight as 1'.format(name)) + # print('=> init {}.bias as 0'.format(name)) + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + else: + for layer in [ + self.deconv_layer_0, + self.deconv_layer_1, + self.deconv_layer_2, + ]: + for _, m in layer.named_modules(): + if isinstance(m, nn.ConvTranspose2d): + # print('=> init {}.weight as normal(0, 0.001)'.format(name)) + # print('=> init {}.bias as 0'.format(name)) + nn.init.normal_(m.weight, std=0.001) + if self.deconv_with_bias: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.BatchNorm2d): + # print('=> init {}.weight as 1'.format(name)) + # print('=> init {}.bias as 0'.format(name)) + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + # print('=> init final conv weights from normal distribution') + for head in self.heads: + final_layer = self.__getattr__(head) + for i, m in enumerate(final_layer.modules()): + if isinstance(m, nn.Conv2d): + # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + # print('=> init {}.weight as normal(0, 0.001)'.format(name)) + # print('=> init {}.bias as 0'.format(name)) + if m.weight.shape[0] == self.heads[head]: + if "hm" in head: + nn.init.constant_(m.bias, -2.19) + else: + nn.init.normal_(m.weight, std=0.001) + nn.init.constant_(m.bias, 0) + # if isinstance(m, nn.Linear): + # if m.weight.shape[0] == self.heads[head]: + # prior = 1/71 + # nn.init.constant_(m.bias, -math.log((1-prior)/prior)) + # else: + # nn.init.normal_(m.weight, std=0.001) + # nn.init.constant_(m.bias, 0) + + # pretrained_state_dict = torch.load(pretrained) + url = model_urls["resnet{}".format(num_layers)] + pretrained_state_dict = model_zoo.load_url(url) + print("=> loading pretrained model {}".format(url)) + self.load_state_dict(pretrained_state_dict, strict=False) + else: + print("=> imagenet pretrained model dose not exist") + print("=> please download it first") + raise ValueError("imagenet pretrained model does not exist") + + +resnet_spec = { + 18: (BasicBlock, [2, 2, 2, 2]), + 34: (BasicBlock, [3, 4, 6, 3]), + 50: (Bottleneck, [3, 4, 6, 3]), + 101: (Bottleneck, [3, 4, 23, 3]), + 152: (Bottleneck, [3, 8, 36, 3]), +} diff --git a/video/fake-stormer/model_code/models/networks/necks/__init__.py b/video/fake-stormer/model_code/models/networks/necks/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fd9cea73d48be2870de0d4fc54500d559133033d --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/necks/__init__.py @@ -0,0 +1,4 @@ +# -*- coding: utf-8 -*- +from .efpn3D import EFPN3D + +__all__ = ["EFPN3D"] diff --git a/video/fake-stormer/model_code/models/networks/necks/base.py b/video/fake-stormer/model_code/models/networks/necks/base.py new file mode 100644 index 0000000000000000000000000000000000000000..e4aa3766502e681f3bce863dee33118f707178fc --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/necks/base.py @@ -0,0 +1,40 @@ +# -*- coding: utf-8 -*- +from abc import ABCMeta, abstractmethod +from typing import Dict + +import torch +import torch.nn as nn + + +class BaseNeck(nn.Module, metaclass=ABCMeta): + def __init__(self, **kwargs) -> None: + return super().__init__() + + @abstractmethod + def forward(self, x: Dict[str, torch.tensor], **kwargs): + return NotImplemented + + @abstractmethod + def preprocess_inputs(self, x: Dict[str, torch.tensor], **kwargs): + return NotImplemented + + @abstractmethod + def init_weights(self, pretrained=None): + return NotImplemented + + @staticmethod + def _get_deconv_cfg(deconv_kernel: int): + """Get configurations for deconv layers.""" + if deconv_kernel == 4: + padding = 1 + output_padding = 0 + elif deconv_kernel == 3: + padding = 1 + output_padding = 1 + elif deconv_kernel in [1, 2]: + padding = 0 + output_padding = 0 + else: + raise ValueError(f"Not supported num_kernels ({deconv_kernel}).") + + return deconv_kernel, padding, output_padding diff --git a/video/fake-stormer/model_code/models/networks/necks/efpn3D.py b/video/fake-stormer/model_code/models/networks/necks/efpn3D.py new file mode 100644 index 0000000000000000000000000000000000000000..0794dea46ce1928e0fe041fa23ed2514d83b51eb --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/necks/efpn3D.py @@ -0,0 +1,173 @@ +# -*- coding: utf-8 -*- +from typing import Dict, List, Sequence + +import torch +import torch.nn as nn +from mmengine.model import xavier_init + +from ...builder import NECKS +from ..common import InceptionBlock3D, conv3d_block, deconv3d_block +from .base import BaseNeck + + +@NECKS.register_module() +class EFPN3D(BaseNeck): + def __init__( + self, + in_channels: int, + num_deconv_layers: int, + num_deconv_filters: Sequence[int], + num_deconv_kernels: Sequence[int], + num_deconv_strides: Sequence[int], + efpn: bool = True, + **kwargs, + ) -> None: + super().__init__(**kwargs) + + self.in_channels = in_channels + self.efpn = efpn + self.num_deconv_layers = num_deconv_layers + + if num_deconv_layers > 0: + self.deconv_layers = self._make_3ddeconv_layer( + num_deconv_filters[0], + num_deconv_layers, + num_deconv_filters, + num_deconv_kernels, + num_deconv_strides, + ) + elif num_deconv_layers == 0: + self.deconv_layers = nn.Identity() + else: + raise ValueError(f"num_deconv_layers ({num_deconv_layers}) should >= 0.") + + if self.efpn: + in_filter = self.in_channels + for idx, out_filter in enumerate(num_deconv_filters): + deconv_block = nn.Sequential( + conv3d_block( + inplanes=in_filter, + outplanes=out_filter, + kernel_size=(1, 1, 1), + stride=(1, 1, 1), + bias=False, + ), + self.deconv_layers[idx], + ) + in_filter = out_filter * 2 + self.__setattr__(f"deconv_block_{idx}", deconv_block) + self.inception_block3d = InceptionBlock3D( + inplanes=num_deconv_filters[-1], + outplanes=num_deconv_filters[-1], + stride=1, + pool_size=3, + ) + + def preprocess_inputs(self, x: Dict[str, torch.tensor], **kwargs): + assert "embed" in x.keys(), "The input dict must contain embedding features!" + inputs = x["embed"] + additional_inputs = {} + + for k, v in x.items(): + if k != "embed": + additional_inputs[k] = v + + return inputs, additional_inputs + + def _make_3ddeconv_layer( + self, + in_channels, + num_layers: int, + num_filters: List[int], + num_kernels: Sequence[int], + num_strides: Sequence[int], + ) -> list: + """Make deconv layers.""" + if num_layers != len(num_filters): + error_msg = ( + f"num_layers({num_layers}) " + f"!= length of num_filters({len(num_filters)})" + ) + raise ValueError(error_msg) + if num_layers != len(num_kernels): + error_msg = ( + f"num_layers({num_layers}) " + f"!= length of num_kernels({len(num_kernels)})" + ) + raise ValueError(error_msg) + + layers = [] + in_planes = in_channels + for i in range(num_layers): + kernels = [] + paddings = [] + output_paddings = [] + + for j in range(len(num_kernels[i])): + kernel, padding, output_padding = self._get_deconv_cfg( + num_kernels[i][j] + ) + kernels.append(kernel) + paddings.append(padding) + output_paddings.append(output_padding) + + outplanes = num_filters[i] + layers.append( + deconv3d_block( + inplanes=in_planes, + outplanes=outplanes, + kernel_size=kernels, + stride=num_strides[i], + padding=paddings, + bias=False, + out_padding=output_paddings, + ) + ) + + # This condition to match n_filters after convolution and optimize number of EFPN parameters + if self.efpn: + in_planes = ( + num_filters[i + 1] + if (i + 1 < num_layers) + else num_filters[num_layers - 1] + ) + else: + in_planes = outplanes + + return layers + + def forward(self, x: Dict[str, torch.tensor], **kwargs) -> dict: + inputs, additional_inputs = self.preprocess_inputs(x) + trapezoids = None + if "outputs" in additional_inputs.keys(): + trapezoids = additional_inputs["outputs"] + # x2, x3, x4, x5 = trapezoids[0], trapezoids[1], trapezoids[2], trapezoids[3] + + if not self.efpn: + x_embed = self.deconv_layers(inputs) + else: + assert ( + "outputs" in additional_inputs.keys() + ), "EFPN requires multiscale feature outputs!" + x_embed = inputs + for i in range(self.num_deconv_layers): + x_embed = self.__getattr__(f"deconv_block_{i}")(x_embed) + if i < self.num_deconv_layers - 1: + x_weights = x_embed.sigmoid_() + x_inv = torch.sub(1, x_weights, alpha=1) + x_ = torch.multiply( + x_inv, trapezoids[self.num_deconv_layers - (i + 2)] + ) + x_embed = torch.cat([x_embed, x_], dim=1) + + # Last layer to capture multi-scale artifacts + x_embed = self.inception_block3d(x_embed) + + res = {} + res["embed"] = x_embed + + return res + + def init_weights(self) -> None: + for layer in self.deconv_layers: + xavier_init(layer, distribution="uniform") diff --git a/video/fake-stormer/model_code/models/networks/pose_efficientNet.py b/video/fake-stormer/model_code/models/networks/pose_efficientNet.py new file mode 100644 index 0000000000000000000000000000000000000000..d6c3cf28fc7c1170289d3039330c4798b843198e --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/pose_efficientNet.py @@ -0,0 +1,978 @@ +# -*- coding: utf-8 -*- +import math +import os +import sys + +if not os.getcwd() in sys.path: + sys.path.append(os.getcwd()) + +import torch +from torch import nn +from torch.nn import functional as F +from torch.utils import model_zoo + +from ..builder import MODELS, build_model +from .backbones.efficientNet import ( + MemoryEfficientSwish, + Swish, + calculate_output_image_size, + drop_connect, + efficientnet_params, + get_model_params, + get_same_padding_conv2d, + load_pretrained_weights, + round_filters, + round_repeats, + url_map, + url_map_advprop, +) +from .common import ( + BN_MOMENTUM, + InceptionBlock, + SELayer, + Texture_Enhance, + conv_block, + point_wise_block, +) + +VALID_MODELS = ( + "efficientnet-b0", + "efficientnet-b1", + "efficientnet-b2", + "efficientnet-b3", + "efficientnet-b4", + "efficientnet-b5", + "efficientnet-b6", + "efficientnet-b7", + "efficientnet-b8", + # Support the construction of 'efficientnet-l2' without pretrained weights + "efficientnet-l2", +) + + +class MBConvBlock(nn.Module): + """Mobile Inverted Residual Bottleneck Block. + Args: + block_args (namedtuple): BlockArgs, defined in utils.py. + global_params (namedtuple): GlobalParam, defined in utils.py. + image_size (tuple or list): [image_height, image_width]. + References: + [1] https://arxiv.org/abs/1704.04861 (MobileNet v1) + [2] https://arxiv.org/abs/1801.04381 (MobileNet v2) + [3] https://arxiv.org/abs/1905.02244 (MobileNet v3) + """ + + def __init__(self, block_args, global_params, image_size=None): + super().__init__() + self._block_args = block_args + self._bn_mom = ( + 1 - global_params.batch_norm_momentum + ) # pytorch's difference from tensorflow + self._bn_eps = global_params.batch_norm_epsilon + self.has_se = (self._block_args.se_ratio is not None) and ( + 0 < self._block_args.se_ratio <= 1 + ) + self.id_skip = ( + block_args.id_skip + ) # whether to use skip connection and drop connect + + # Expansion phase (Inverted Bottleneck) + inp = self._block_args.input_filters # number of input channels + oup = ( + self._block_args.input_filters * self._block_args.expand_ratio + ) # number of output channels + if self._block_args.expand_ratio != 1: + Conv2d = get_same_padding_conv2d(image_size=image_size) + self._expand_conv = Conv2d( + in_channels=inp, out_channels=oup, kernel_size=1, bias=False + ) + self._bn0 = nn.BatchNorm2d( + num_features=oup, momentum=self._bn_mom, eps=self._bn_eps + ) + # image_size = calculate_output_image_size(image_size, 1) <-- this wouldn't modify image_size + + # Depthwise convolution phase + k = self._block_args.kernel_size + s = self._block_args.stride + Conv2d = get_same_padding_conv2d(image_size=image_size) + self._depthwise_conv = Conv2d( + in_channels=oup, + out_channels=oup, + groups=oup, # groups makes it depthwise + kernel_size=k, + stride=s, + bias=False, + ) + self._bn1 = nn.BatchNorm2d( + num_features=oup, momentum=self._bn_mom, eps=self._bn_eps + ) + image_size = calculate_output_image_size(image_size, s) + + # Squeeze and Excitation layer, if desired + if self.has_se: + Conv2d = get_same_padding_conv2d(image_size=(1, 1)) + num_squeezed_channels = max( + 1, int(self._block_args.input_filters * self._block_args.se_ratio) + ) + self._se_reduce = Conv2d( + in_channels=oup, out_channels=num_squeezed_channels, kernel_size=1 + ) + self._se_expand = Conv2d( + in_channels=num_squeezed_channels, out_channels=oup, kernel_size=1 + ) + + # Pointwise convolution phase + final_oup = self._block_args.output_filters + Conv2d = get_same_padding_conv2d(image_size=image_size) + self._project_conv = Conv2d( + in_channels=oup, out_channels=final_oup, kernel_size=1, bias=False + ) + self._bn2 = nn.BatchNorm2d( + num_features=final_oup, momentum=self._bn_mom, eps=self._bn_eps + ) + self._swish = MemoryEfficientSwish() + + def forward(self, inputs, drop_connect_rate=None): + """MBConvBlock's forward function. + Args: + inputs (tensor): Input tensor. + drop_connect_rate (bool): Drop connect rate (float, between 0 and 1). + Returns: + Output of this block after processing. + """ + + # Expansion and Depthwise Convolution + x = inputs + if self._block_args.expand_ratio != 1: + x = self._expand_conv(inputs) + x = self._bn0(x) + x = self._swish(x) + + x = self._depthwise_conv(x) + x = self._bn1(x) + x = self._swish(x) + + # Squeeze and Excitation + if self.has_se: + x_squeezed = F.adaptive_avg_pool2d(x, 1) + x_squeezed = self._se_reduce(x_squeezed) + x_squeezed = self._swish(x_squeezed) + x_squeezed = self._se_expand(x_squeezed) + x = torch.sigmoid(x_squeezed) * x + + # Pointwise Convolution + x = self._project_conv(x) + x = self._bn2(x) + + # Skip connection and drop connect + input_filters, output_filters = ( + self._block_args.input_filters, + self._block_args.output_filters, + ) + if ( + self.id_skip + and self._block_args.stride == 1 + and input_filters == output_filters + ): + # The combination of skip connection and drop connect brings about stochastic depth. + if drop_connect_rate: + x = drop_connect(x, p=drop_connect_rate, training=self.training) + x = x + inputs # skip connection + return x + + def set_swish(self, memory_efficient=True): + """Sets swish function as memory efficient (for training) or standard (for export). + Args: + memory_efficient (bool): Whether to use memory-efficient version of swish. + """ + self._swish = MemoryEfficientSwish() if memory_efficient else Swish() + + +@MODELS.register_module() +class EfficientNet(nn.Module): + """EfficientNet model. + Most easily loaded with the .from_name or .from_pretrained methods. + Args: + blocks_args (list[namedtuple]): A list of BlockArgs to construct blocks. + global_params (namedtuple): A set of GlobalParams shared between blocks. + References: + [1] https://arxiv.org/abs/1905.11946 (EfficientNet) + Example: + >>> import torch + >>> from efficientnet.model import EfficientNet + >>> inputs = torch.rand(1, 3, 224, 224) + >>> model = EfficientNet.from_pretrained('efficientnet-b0') + >>> model.eval() + >>> outputs = model(inputs) + """ + + def __init__(self, blocks_args=None, global_params=None): + super().__init__() + assert isinstance(blocks_args, list), "blocks_args should be a list" + assert len(blocks_args) > 0, "block args must be greater than 0" + self._global_params = global_params + self._blocks_args = blocks_args + + # Batch norm parameters + bn_mom = 1 - self._global_params.batch_norm_momentum + bn_eps = self._global_params.batch_norm_epsilon + + # Get stem static or dynamic convolution depending on image size + image_size = global_params.image_size + Conv2d = get_same_padding_conv2d(image_size=image_size) + + # Stem + in_channels = 3 # rgb + out_channels = round_filters( + 32, self._global_params + ) # number of output channels + self._conv_stem = Conv2d( + in_channels, out_channels, kernel_size=3, stride=2, bias=False + ) + self._bn0 = nn.BatchNorm2d( + num_features=out_channels, momentum=bn_mom, eps=bn_eps + ) + image_size = calculate_output_image_size(image_size, 2) + + # Build blocks + self._blocks = nn.ModuleList([]) + for block_args in self._blocks_args: + + # Update block input and output filters based on depth multiplier. + block_args = block_args._replace( + input_filters=round_filters( + block_args.input_filters, self._global_params + ), + output_filters=round_filters( + block_args.output_filters, self._global_params + ), + num_repeat=round_repeats(block_args.num_repeat, self._global_params), + ) + + # The first block needs to take care of stride and filter size increase. + self._blocks.append( + MBConvBlock(block_args, self._global_params, image_size=image_size) + ) + image_size = calculate_output_image_size(image_size, block_args.stride) + if block_args.num_repeat > 1: # modify block_args to keep same output size + block_args = block_args._replace( + input_filters=block_args.output_filters, stride=1 + ) + for _ in range(block_args.num_repeat - 1): + self._blocks.append( + MBConvBlock(block_args, self._global_params, image_size=image_size) + ) + # image_size = calculate_output_image_size(image_size, block_args.stride) # stride = 1 + + # Head + in_channels = block_args.output_filters # output of final block + out_channels = round_filters(1280, self._global_params) + Conv2d = get_same_padding_conv2d(image_size=image_size) + self._conv_head = Conv2d(in_channels, out_channels, kernel_size=1, bias=False) + self._bn1 = nn.BatchNorm2d( + num_features=out_channels, momentum=bn_mom, eps=bn_eps + ) + + # Final linear layer + self._avg_pooling = nn.AdaptiveAvgPool2d(1) + if self._global_params.include_top: + self._dropout = nn.Dropout(self._global_params.dropout_rate) + self._fc = nn.Linear(out_channels, self._global_params.num_classes) + + # Heatmap Decoder Construction + if self._global_params.include_hm_decoder: + print("Constructing the heatmap Decoder!") + self.efpn = self._global_params.efpn + self.tfpn = self._global_params.tfpn + + assert not ( + self.efpn and self.tfpn + ), "Only one of E-FPN or FPN is intergrated!" + + self.se_layer = self._global_params.se_layer + # self.hm_decoder_filters = [1792, 448, 160, 56] if self.fpn else [1792, 256, 256, 128] + self.hm_decoder_filters = [1792, 448, 160, 56] + num_kernels = [4, 4, 4, 4] if (self.efpn or self.tfpn) else [4, 4, 4] + self._dropout = nn.Dropout(self._global_params.dropout_rate) + self._sigmoid = nn.Sigmoid() + self._relu = nn.ReLU(inplace=True) + self._relu1 = nn.ReLU(inplace=False) + self.deconv_with_bias = False + + # if self._global_params.use_c2: + # self.inception_block_4 = InceptionBlock(224, 224, stride=1, pool_size=3) + if self._global_params.use_c3: + self.inception_block = InceptionBlock(112, 112, stride=1, pool_size=3) + else: + self.inception_block = InceptionBlock(56, 56, stride=1, pool_size=3) + + self.heads = self._global_params.heads + n_deconv = len(self.hm_decoder_filters) + self.fpn_layers = [ + self._global_params.use_c51, + self._global_params.use_c4, + self._global_params.use_c3, + ] + + if self.efpn or self.tfpn: + for idx in range(n_deconv): + in_decod_filters = self.hm_decoder_filters[idx] + + if idx == 0: + out_decod_filters = self.hm_decoder_filters[idx + 1] + deconv = nn.Sequential( + conv_block( + in_decod_filters, + out_decod_filters, + (3, 3), + stride=1, + padding=1, + ), + ) + else: + in_decod_filters = ( + in_decod_filters * 2 + if self.fpn_layers[idx - 1] + else in_decod_filters + ) + kernel, padding, output_padding = self._get_deconv_cfg( + num_kernels[idx] + ) + + if idx + 1 < n_deconv: + out_decod_filters = self.hm_decoder_filters[idx + 1] + deconv = nn.Sequential( + conv_block( + in_decod_filters, + out_decod_filters, + (3, 3), + stride=1, + padding=1, + ), + nn.ConvTranspose2d( + in_channels=out_decod_filters, + out_channels=out_decod_filters, + kernel_size=kernel, + stride=2, + padding=padding, + output_padding=output_padding, + bias=self.deconv_with_bias, + ), + nn.BatchNorm2d(out_decod_filters, momentum=BN_MOMENTUM), + ) + else: + out_decod_filters = in_decod_filters + deconv = nn.Sequential( + self.inception_block, + # conv_block(in_decod_filters, out_decod_filters, (3,3), stride=1, padding=1), + nn.ConvTranspose2d( + in_channels=out_decod_filters, + out_channels=out_decod_filters, + kernel_size=kernel, + stride=2, + padding=padding, + output_padding=output_padding, + bias=self.deconv_with_bias, + ), + nn.BatchNorm2d(out_decod_filters, momentum=BN_MOMENTUM), + ) + + # In case of using C2, this conv to apply to C2 features to get the same filters of the last deconv + if self._global_params.use_c2: + if self._global_params.norm_c2: + self.texture_enhance = Texture_Enhance(32) + # self.conv_c2 = point_wise_block(128, out_decod_filters) + self.conv_c2 = conv_block( + 128, + out_decod_filters, + (3, 3), + stride=1, + padding=1, + ) + else: + self.conv_c2 = conv_block( + 32, + out_decod_filters, + (3, 3), + stride=1, + padding=1, + ) + if self.se_layer: + se = SELayer(channel=out_decod_filters * 2) + self.__setattr__(f"se_layer_{idx+1}", se) + + self.__setattr__(f"deconv_{idx+1}", deconv) + else: + self.deconv_layers = self._make_deconv_layer( + len(num_kernels), + self.hm_decoder_filters, + num_kernels, + ) + + for head, num_output in self.heads.items(): + head_conv = int(self._global_params.head_conv) + num_output = int(num_output) + if self._global_params.use_c2: + assert ( + self._global_params.efpn or self._global_params.tfpn + ), "FPN Design must be set active!" + assert ( + self._global_params.use_c3 + ), "C3 must be utilized for FPN intergration of C2" + in_head_filters = self.hm_decoder_filters[-1] * 4 + elif self._global_params.use_c3: + in_head_filters = self.hm_decoder_filters[-1] * 2 + else: + in_head_filters = self.hm_decoder_filters[-1] + + if head_conv > 0: + if head != "cls": + fc = nn.Sequential( + nn.Conv2d( + in_head_filters, + head_conv, + kernel_size=3, + padding=1, + bias=True, + ), + nn.BatchNorm2d(head_conv), + nn.ReLU(inplace=True), + nn.Conv2d( + head_conv, + num_output, + kernel_size=1, + stride=1, + padding=0, + ), + ) + else: + fc = nn.Sequential( + nn.Conv2d( + in_head_filters, + head_conv, + kernel_size=3, + padding=1, + bias=True, + ), + nn.BatchNorm2d(head_conv, momentum=BN_MOMENTUM), + nn.ReLU(inplace=True), + # nn.Conv2d(head_conv, num_output, kernel_size=1, + # stride=1, padding=0, bias=True), + # nn.BatchNorm2d(num_output), + # nn.ReLU(inplace=True), + # nn.AdaptiveMaxPool2d(head_conv//4), + nn.AdaptiveAvgPool2d(1), + nn.Flatten(), + # nn.Linear((head_conv//4)**2, head_conv, bias=True), + # nn.BatchNorm1d(head_conv, momentum=BN_MOMENTUM), + # nn.ReLU(inplace=True), + nn.Linear(head_conv, num_output, bias=True), + # nn.Sigmoid(), + # nn.Softmax(dim=-1) + ) + else: + fc = nn.Conv2d( + in_channels=in_head_filters, + out_channels=num_output, + kernel_size=1, + stride=1, + padding=0, + ) + self.__setattr__(head, fc) + + # set activation to memory efficient swish by default + self._swish = MemoryEfficientSwish() + + def _get_deconv_cfg(self, deconv_kernel): + if deconv_kernel == 4: + padding = 1 + output_padding = 0 + elif deconv_kernel == 3: + padding = 1 + output_padding = 1 + elif deconv_kernel == 2: + padding = 0 + output_padding = 0 + + return deconv_kernel, padding, output_padding + + def _make_deconv_layer(self, num_layers, num_filters, num_kernels): + assert num_layers == ( + len(num_filters) - 1 + ), "ERROR: num_deconv_layers is different len(num_deconv_filters)" + assert num_layers == len( + num_kernels + ), "ERROR: num_deconv_layers is different len(num_deconv_filters)" + + layers = [] + for i in range(num_layers): + kernel, padding, output_padding = self._get_deconv_cfg(num_kernels[i]) + + in_planes = num_filters[i] + out_planes = num_filters[i + 1] + + layers.append( + nn.Sequential( + nn.ConvTranspose2d( + in_channels=in_planes, + out_channels=out_planes, + kernel_size=kernel, + stride=2, + padding=padding, + output_padding=output_padding, + bias=self.deconv_with_bias, + ), + nn.BatchNorm2d(out_planes, momentum=BN_MOMENTUM), + nn.ReLU(inplace=True), + ) + ) + + return nn.Sequential(*layers) + + def set_swish(self, memory_efficient=True): + """Sets swish function as memory efficient (for training) or standard (for export). + Args: + memory_efficient (bool): Whether to use memory-efficient version of swish. + """ + self._swish = MemoryEfficientSwish() if memory_efficient else Swish() + for block in self._blocks: + block.set_swish(memory_efficient) + + def extract_endpoints(self, inputs): + """Use convolution layer to extract features + from reduction levels i in [1, 2, 3, 4, 5]. + Args: + inputs (tensor): Input tensor. + Returns: + Dictionary of last intermediate features + with reduction levels i in [1, 2, 3, 4, 5]. + Example: + >>> import torch + >>> from efficientnet.model import EfficientNet + >>> inputs = torch.rand(1, 3, 224, 224) + >>> model = EfficientNet.from_pretrained('efficientnet-b0') + >>> endpoints = model.extract_endpoints(inputs) + >>> print(endpoints['reduction_1'].shape) # torch.Size([1, 16, 112, 112]) + >>> print(endpoints['reduction_2'].shape) # torch.Size([1, 24, 56, 56]) + >>> print(endpoints['reduction_3'].shape) # torch.Size([1, 40, 28, 28]) + >>> print(endpoints['reduction_4'].shape) # torch.Size([1, 112, 14, 14]) + >>> print(endpoints['reduction_5'].shape) # torch.Size([1, 320, 7, 7]) + >>> print(endpoints['reduction_6'].shape) # torch.Size([1, 1280, 7, 7]) + """ + endpoints = dict() + + # Stem + x = self._swish(self._bn0(self._conv_stem(inputs))) + prev_x = x + + # Blocks + for idx, block in enumerate(self._blocks): + drop_connect_rate = self._global_params.drop_connect_rate + if drop_connect_rate: + drop_connect_rate *= float(idx) / len( + self._blocks + ) # scale drop connect_rate + x = block(x, drop_connect_rate=drop_connect_rate) + # print('Prev', prev_x.size()) + # print('X', x.size()) + if prev_x.size(2) > x.size(2): + endpoints["reduction_{}".format(len(endpoints) + 1)] = prev_x + elif idx == len(self._blocks) - 1: + endpoints["reduction_{}".format(len(endpoints) + 1)] = x + prev_x = x + + # Head + x = self._swish(self._bn1(self._conv_head(x))) + endpoints["reduction_{}".format(len(endpoints) + 1)] = x + + return endpoints + + def extract_features(self, inputs): + """use convolution layer to extract feature . + Args: + inputs (tensor): Input tensor. + Returns: + Output of the final convolution + layer in the efficientnet model. + """ + # Stem + x = self._swish(self._bn0(self._conv_stem(inputs))) + + # Blocks + for idx, block in enumerate(self._blocks): + drop_connect_rate = self._global_params.drop_connect_rate + if drop_connect_rate: + drop_connect_rate *= float(idx) / len( + self._blocks + ) # scale drop connect_rate + x = block(x, drop_connect_rate=drop_connect_rate) + + # Head + x = self._swish(self._bn1(self._conv_head(x))) + + return x + + def forward(self, inputs): + """EfficientNet's forward function. + Calls extract_features to extract features, applies final linear layer, and returns logits. + Args: + inputs (tensor): Input tensor. + Returns: + Output of this model after processing. + """ + # Convolution layers + # x = self.extract_features(inputs) + endpoints = self.extract_endpoints(inputs) + x1 = endpoints["reduction_6"] + x2 = endpoints["reduction_5"] + x3 = endpoints["reduction_4"] + x4 = endpoints["reduction_3"] + x5 = endpoints["reduction_2"] + x = x1 + + if self._global_params.include_top: + # Pooling and final linear layer + x = self._avg_pooling(x) + + x = x.flatten(start_dim=1) + x = self._dropout(x) + x = self._fc(x) + res = {} + res["cls"] = x + + return res + + if self._global_params.include_hm_decoder: + x1 = self._dropout(x1) + x2 = self._dropout(x2) + x3 = self._dropout(x3) + x4 = self._dropout(x4) + + if self.efpn: + assert ( + self._global_params.use_c51 + ), "C51 must be utilized for FPN intergration" + + x = self.__getattr__("deconv_1")(x1) + + if self._global_params.use_c51: + x_weighted = self._sigmoid(x) + x_inv = torch.sub(1, x_weighted, alpha=1) + x2_ = torch.multiply(x_inv, x2) + x = torch.cat([x, x2_], dim=1) + + if self.se_layer: + x = self.__getattr__("se_layer_1")(x) + else: + x = self._relu(x) + + x = self.__getattr__("deconv_2")(x) + + if self._global_params.use_c4: + x_weighted = self._sigmoid(x) + x_inv = torch.sub(1, x_weighted, alpha=1) + x3_ = torch.multiply(x_inv, x3) + x = torch.cat([x, x3_], dim=1) + + if self.se_layer: + x = self.__getattr__("se_layer_2")(x) + else: + x = self._relu(x) + + x = self.__getattr__("deconv_3")(x) + + if self._global_params.use_c3: + assert ( + self._global_params.use_c4 + ), "C4 must be utilized for FPN intergration of C3" + + x_weighted = self._sigmoid(x) + x_inv = torch.sub(1, x_weighted, alpha=1) + x4_ = torch.multiply(x_inv, x4) + x = torch.cat([x, x4_], dim=1) + + if self.se_layer: + x = self.__getattr__("se_layer_3")(x) + else: + x = self._relu(x) + + x = self.__getattr__("deconv_4")(x) + + if not self._global_params.use_c2: + x = self._relu(x) + else: + assert ( + self._global_params.use_c3 + ), "C3 must be utilized for FPN intergration of C2" + + x5 = self._dropout(x5) + if self._global_params.norm_c2: + x5_ = self.texture_enhance(x5, (16, 16)) + x5_ = self.conv_c2(x5_) + else: + x5_ = self.conv_c2(x5) + + x_weighted = self._sigmoid(x) + x_inv = torch.sub(1, x_weighted, alpha=1) + x5_ = torch.multiply(x_inv, x5_) + x = torch.cat([x, x5_], dim=1) + + # # Adding multi receptive fields + # x = self.inception_block_4(x) + + if self.se_layer: + x = self.__getattr__("se_layer_4")(x) + elif self.tfpn: + assert ( + self._global_params.use_c51 + ), "C51 must be utilized for FPN intergration" + x = self.__getattr__("deconv_1")(x1) + x = self._relu1(x) + x = torch.cat([x, x2], dim=1) + + x = self.__getattr__("deconv_2")(x) + if not self._global_params.use_c4: + x = self._relu1(x) + else: + x = torch.cat([x, x3], dim=1) + + x = self.__getattr__("deconv_3")(x) + if not self._global_params.use_c3: + x = self._relu1(x) + else: + assert ( + self._global_params.use_c4 + ), "C4 must be utilized for FPN intergration of C3" + x = torch.cat([x, x4], dim=1) + + x = self.__getattr__("deconv_4")(x) + if not self._global_params.use_c2: + x = self._relu(x) + else: + assert ( + self._global_params.use_c3 + ), "C3 must be utilized for FPN intergration of C2" + x5 = self._dropout(x5) + x5 = self.conv_c2(x5) + x = self._relu1(x) + x = torch.cat([x, x5], dim=1) + else: + x = self.deconv_layers(x1) + + ret = {} + for head in self.heads: + ret[head] = self.__getattr__(head)(x) + + return [ret] + + @classmethod + def from_name(cls, model_name, in_channels=3, **override_params): + """Create an efficientnet model according to name. + Args: + model_name (str): Name for efficientnet. + in_channels (int): Input data's channel number. + override_params (other key word params): + Params to override model's global_params. + Optional key: + 'width_coefficient', 'depth_coefficient', + 'image_size', 'dropout_rate', + 'num_classes', 'batch_norm_momentum', + 'batch_norm_epsilon', 'drop_connect_rate', + 'depth_divisor', 'min_depth' + Returns: + An efficientnet model. + """ + cls._check_model_name_is_valid(model_name) + blocks_args, global_params = get_model_params(model_name, override_params) + model = cls(blocks_args, global_params) + model._change_in_channels(in_channels) + return model + + @classmethod + def from_pretrained( + cls, + model_name, + weights_path=None, + advprop=False, + in_channels=3, + num_classes=1000, + **override_params, + ): + """Create an efficientnet model according to name. + Args: + model_name (str): Name for efficientnet. + weights_path (None or str): + str: path to pretrained weights file on the local disk. + None: use pretrained weights downloaded from the Internet. + advprop (bool): + Whether to load pretrained weights + trained with advprop (valid when weights_path is None). + in_channels (int): Input data's channel number. + num_classes (int): + Number of categories for classification. + It controls the output size for final linear layer. + override_params (other key word params): + Params to override model's global_params. + Optional key: + 'width_coefficient', 'depth_coefficient', + 'image_size', 'dropout_rate', + 'batch_norm_momentum', + 'batch_norm_epsilon', 'drop_connect_rate', + 'depth_divisor', 'min_depth' + Returns: + A pretrained efficientnet model. + """ + model = cls.from_name(model_name, num_classes=num_classes, **override_params) + load_pretrained_weights( + model, + model_name, + weights_path=weights_path, + load_fc=((num_classes == 1000) and (model._global_params.include_top)), + advprop=advprop, + ) + model._change_in_channels(in_channels) + return model + + @classmethod + def get_image_size(cls, model_name): + """Get the input image size for a given efficientnet model. + Args: + model_name (str): Name for efficientnet. + Returns: + Input image size (resolution). + """ + cls._check_model_name_is_valid(model_name) + _, _, res, _ = efficientnet_params(model_name) + return res + + @classmethod + def _check_model_name_is_valid(cls, model_name): + """Validates model name. + Args: + model_name (str): Name for efficientnet. + Returns: + bool: Is a valid name or not. + """ + if model_name not in VALID_MODELS: + raise ValueError("model_name should be one of: " + ", ".join(VALID_MODELS)) + + def _change_in_channels(self, in_channels): + """Adjust model's first convolution layer to in_channels, if in_channels not equals 3. + Args: + in_channels (int): Input data's channel number. + """ + if in_channels != 3: + Conv2d = get_same_padding_conv2d(image_size=self._global_params.image_size) + out_channels = round_filters(32, self._global_params) + self._conv_stem = Conv2d( + in_channels, out_channels, kernel_size=3, stride=2, bias=False + ) + + +@MODELS.register_module() +class PoseEfficientNet(EfficientNet): + def __init__(self, model_name, in_channels=3, **override_params): + self.model_name = model_name + self.in_channels = in_channels + + # Initialize Parent Class + super()._check_model_name_is_valid(model_name) + blocks_args, global_params = get_model_params(model_name, override_params) + super().__init__(blocks_args, global_params) + + @classmethod + def from_name(cls, model_name, in_channels, **override_params): + return NotImplemented + + @classmethod + def from_pretrained( + cls, + model_name, + weights_path, + advprop, + in_channels, + num_classes, + **override_params, + ): + return NotImplemented + + def _change_in_channels(self, in_channels): + return NotImplemented + + def init_weights(self, pretrained=False, advprop=False, verbose=True): + if pretrained: + url_map_ = url_map_advprop if advprop else url_map + state_dict = model_zoo.load_url(url_map_[self.model_name]) + state_dict.pop("_fc.weight") + state_dict.pop("_fc.bias") + self.load_state_dict(state_dict, strict=False) + + # Initialize weights for Deconvolution Layer + if self._global_params.include_hm_decoder: + if self.efpn or self.tfpn: + deconv_layers = [ + self.deconv_1, + self.deconv_2, + self.deconv_3, + self.deconv_4, + ] + else: + deconv_layers = self.deconv_layers + + for layer in deconv_layers: + for _, m in layer.named_modules(): + if isinstance(m, nn.ConvTranspose2d): + n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + m.weight.data.normal_(0, math.sqrt(2.0 / n)) + if self.deconv_with_bias: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + # Init head parameters + for head in self.heads: + final_layer = self.__getattr__(head) + for i, m in enumerate(final_layer.modules()): + if isinstance(m, nn.Conv2d): + if m.weight.shape[0] == self.heads[head]: + if "hm" in head: + nn.init.constant_(m.bias, -2.19) + else: + # nn.init.normal_(m.weight, std=0.001) + n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + m.weight.data.normal_(0, math.sqrt(2.0 / n)) + nn.init.constant_(m.bias, 0) + + self._change_in_channels(in_channels=self.in_channels) + if verbose: + print("Loaded pretrained weights for {}".format(self.model_name)) + + +if __name__ == "__main__": + cfg = dict( + type="PoseEfficientNet", + model_name="efficientnet-b4", + include_top=False, + include_hm_decoder=True, + head_conv=64, + heads={"hm": 1, "cls": 1, "cstency": 256}, + use_c2=True, + ) + model = build_model(cfg, MODELS) + model.init_weights(pretrained=True) + model.eval() + inputs = torch.rand((1, 3, 384, 384)) + + for i, (n, p) in enumerate(model.named_parameters()): + print(i, n) + + # To show the whole pose EFN model outputs shape + x = model(inputs)[0] + for head in x.keys(): + print(f"{head} shape is --- {x[head].shape}") + + # To show the endpoints features shape + # endpoints = model.extract_endpoints(inputs) + # for k in endpoints.keys(): + # print(endpoints[k].shape) diff --git a/video/fake-stormer/model_code/models/networks/pose_hrnet.py b/video/fake-stormer/model_code/models/networks/pose_hrnet.py new file mode 100644 index 0000000000000000000000000000000000000000..0c7f6bfc240320820f070a12e57687e30b608b9b --- /dev/null +++ b/video/fake-stormer/model_code/models/networks/pose_hrnet.py @@ -0,0 +1,540 @@ +# -*- coding: utf-8 -*- +from __future__ import absolute_import, division, print_function + +import logging +import os +import re + +import torch +import torch.nn as nn + +from ..builder import MODELS +from .common import BN_MOMENTUM, conv3x3 + + +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(True) + + 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]), + 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), + ) + ) + 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), + nn.ReLU(True), + ) + ) + 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} + + +@MODELS.register_module() +class PoseHighResolutionNet(nn.Module): + def __init__(self, cfg, **kwargs): + self.inplanes = 64 + extra = cfg.MODEL.EXTRA + self.cls_based_hm = cfg.MODEL.cls_based_hm + self.heads = cfg.MODEL.heads + super(PoseHighResolutionNet, self).__init__() + + # stem net + 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.layer1 = self._make_layer(Bottleneck, 64, 4) + + 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([256], 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=False + ) + + self.final_layer = nn.Conv2d( + in_channels=pre_stage_channels[0], + out_channels=cfg.MODEL.NUM_JOINTS, + kernel_size=extra.FINAL_CONV_KERNEL, + stride=1, + padding=1 if extra.FINAL_CONV_KERNEL == 3 else 0, + ) + + self.final_layer_cls = nn.Sequential( + nn.BatchNorm2d(cfg.MODEL.NUM_JOINTS, momentum=BN_MOMENTUM), + nn.AdaptiveMaxPool2d(cfg.MODEL.HEATMAP_SIZE[0] // 4), + nn.Flatten(), + nn.Linear( + (cfg.MODEL.HEATMAP_SIZE[0] // 4) ** 2, cfg.MODEL.NUM_JOINTS, bias=True + ), + nn.Sigmoid(), + ) + + self.pretrained_layers = cfg["MODEL"]["EXTRA"]["PRETRAINED_LAYERS"] + + 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]), + 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), + nn.ReLU(inplace=True), + ) + ) + transition_layers.append(nn.Sequential(*conv3x3s)) + + return nn.ModuleList(transition_layers) + + 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, momentum=BN_MOMENTUM), + ) + + 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 _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) + + x = self.final_layer(y_list[0]) + + ret = {} + for head in self.heads.keys(): + if head == "hm": + ret[head] = x + else: + x1 = self.final_layer_cls(x) + ret[head] = x1 + return [ret] + + def init_weights(self, pretrained="", **kwargs): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + nn.init.normal_(m.weight, std=0.001) + for name, _ in m.named_parameters(): + if name in ["bias"]: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.ConvTranspose2d): + nn.init.normal_(m.weight, std=0.001) + for name, _ in m.named_parameters(): + if name in ["bias"]: + nn.init.constant_(m.bias, 0) + + if os.path.isfile(pretrained): + pretrained_state_dict = torch.load( + pretrained, map_location=torch.device("cpu") + ) + + need_init_state_dict = {} + for name, m in pretrained_state_dict.items(): + if ( + name.split(".")[0] in self.pretrained_layers + or self.pretrained_layers[0] == "*" + ): + need_init_state_dict[name] = m + self.load_state_dict(need_init_state_dict, strict=False) + elif pretrained: + raise ValueError("{} is not exist!".format(pretrained)) + + +def get_pose_net(cfg, is_train, **kwargs): + model = PoseHighResolutionNet(cfg, **kwargs) + + if is_train and cfg.MODEL.INIT_WEIGHTS: + model.init_weights(cfg.MODEL.PRETRAINED) + + return model + + +if __name__ == "__main__": + from builder import build_model + from configs.get_config import load_config + + cfg = load_config("configs/hrnet_sbi.yaml") + + hrnet = build_model(cfg.MODEL, MODELS, default_args=dict(cfg=cfg)) + print(hrnet) diff --git a/video/fake-stormer/model_code/models/utils/__init__.py b/video/fake-stormer/model_code/models/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f469208ff760535cac619667e3c040138c84d20e --- /dev/null +++ b/video/fake-stormer/model_code/models/utils/__init__.py @@ -0,0 +1,28 @@ +# -*- coding:utf-8 -*- +from .check_and_update_config import check_and_update_config +from .utils import ( + freeze_backbone, + load_checkpoint, + load_model, + load_pretrained, + n_param_model, + preset_model, + save_model, + swin_converter, + unfreeze_backbone, +) + +__all__ = [ + "check_and_update_config", + "build_model", + "load_pretrained", + "freeze_backbone", + "resnet_spec", + "load_model", + "save_model", + "unfreeze_backbone", + "preset_model", + "load_checkpoint", + "n_param_model", + "swin_converter", +] diff --git a/video/fake-stormer/model_code/models/utils/check_and_update_config.py b/video/fake-stormer/model_code/models/utils/check_and_update_config.py new file mode 100644 index 0000000000000000000000000000000000000000..7541a1bb8dc4a32ad4ae9f7f73ac498041e69a1a --- /dev/null +++ b/video/fake-stormer/model_code/models/utils/check_and_update_config.py @@ -0,0 +1,262 @@ +# -*- coding:utf-8 -*- +# Copyright (c) OpenMMLab. All rights reserved. +from typing import Dict, Optional, Tuple, Union + +from mmengine.config import Config, ConfigDict +from mmengine.dist import master_only +from mmengine.logging import MMLogger + +ConfigType = Union[Config, ConfigDict] + + +def process_input_transform( + input_transform: str, + head: Dict, + head_new: Dict, + head_deleted_dict: Dict, + head_append_dict: Dict, + neck_new: Dict, + input_index: Tuple[int], + align_corners: bool, +) -> None: + """Process the input_transform field and update head and neck + dictionaries.""" + if input_transform == "resize_concat": + in_channels = head_new.pop("in_channels") + head_deleted_dict["in_channels"] = str(in_channels) + in_channels = sum([in_channels[i] for i in input_index]) + head_new["in_channels"] = in_channels + head_append_dict["in_channels"] = str(in_channels) + + neck_new.update( + dict( + type="FeatureMapProcessor", + concat=True, + select_index=input_index, + ) + ) + if align_corners: + neck_new["align_corners"] = align_corners + + elif input_transform == "select": + if input_index != (-1,): + neck_new.update(dict(type="FeatureMapProcessor", select_index=input_index)) + if isinstance(head["in_channels"], tuple): + in_channels = head_new.pop("in_channels") + head_deleted_dict["in_channels"] = str(in_channels) + if isinstance(input_index, int): + in_channels = in_channels[input_index] + else: + in_channels = tuple([in_channels[i] for i in input_index]) + head_new["in_channels"] = in_channels + head_append_dict["in_channels"] = str(in_channels) + if align_corners: + neck_new["align_corners"] = align_corners + + else: + raise ValueError( + f"model.head get invalid value for argument " + f"input_transform: {input_transform}" + ) + + +def process_extra_field( + extra: Dict, + head_new: Dict, + head_deleted_dict: Dict, + head_append_dict: Dict, + neck_new: Dict, +) -> None: + """Process the extra field and update head and neck dictionaries.""" + head_deleted_dict["extra"] = "dict(" + for key, value in extra.items(): + head_deleted_dict["extra"] += f"{key}={value}," + head_deleted_dict["extra"] = head_deleted_dict["extra"][:-1] + ")" + if "final_conv_kernel" in extra: + kernel_size = extra["final_conv_kernel"] + if kernel_size > 1: + padding = kernel_size // 2 + head_new["final_layer"] = dict(kernel_size=kernel_size, padding=padding) + head_append_dict["final_layer"] = ( + f"dict(kernel_size={kernel_size}, " f"padding={padding})" + ) + else: + head_new["final_layer"] = dict(kernel_size=kernel_size) + head_append_dict["final_layer"] = f"dict(kernel_size={kernel_size})" + if "upsample" in extra: + neck_new.update( + dict( + type="FeatureMapProcessor", + scale_factor=float(extra["upsample"]), + apply_relu=True, + ) + ) + + +def process_has_final_layer( + has_final_layer: bool, + head_new: Dict, + head_deleted_dict: Dict, + head_append_dict: Dict, +) -> None: + """Process the has_final_layer field and update the head dictionary.""" + head_deleted_dict["has_final_layer"] = str(has_final_layer) + if not has_final_layer: + if "final_layer" not in head_new: + head_new["final_layer"] = None + head_append_dict["final_layer"] = "None" + + +def check_and_update_config( + neck: Optional[ConfigType], head: ConfigType +) -> Tuple[Optional[Dict], Dict]: + """Check and update the configuration of the head and neck components. + Args: + neck (Optional[ConfigType]): Configuration for the neck component. + head (ConfigType): Configuration for the head component. + + Returns: + Tuple[Optional[Dict], Dict]: Updated configurations for the neck + and head components. + """ + head_new, neck_new = head.copy(), neck.copy() if isinstance(neck, dict) else {} + head_deleted_dict, head_append_dict = {}, {} + + if "input_transform" in head: + input_transform = head_new.pop("input_transform") + head_deleted_dict["input_transform"] = f"'{input_transform}'" + else: + input_transform = "select" + + if "input_index" in head: + input_index = head_new.pop("input_index") + head_deleted_dict["input_index"] = str(input_index) + else: + input_index = (-1,) + + if "align_corners" in head: + align_corners = head_new.pop("align_corners") + head_deleted_dict["align_corners"] = str(align_corners) + else: + align_corners = False + + process_input_transform( + input_transform, + head, + head_new, + head_deleted_dict, + head_append_dict, + neck_new, + input_index, + align_corners, + ) + + if "extra" in head: + extra = head_new.pop("extra") + process_extra_field( + extra, head_new, head_deleted_dict, head_append_dict, neck_new + ) + + if "has_final_layer" in head: + has_final_layer = head_new.pop("has_final_layer") + process_has_final_layer( + has_final_layer, head_new, head_deleted_dict, head_append_dict + ) + + display_modifications(head_deleted_dict, head_append_dict, neck_new) + + neck_new = neck_new if len(neck_new) else None + return neck_new, head_new + + +@master_only +def display_modifications( + head_deleted_dict: Dict, head_append_dict: Dict, neck: Dict +) -> None: + """Display the modifications made to the head and neck configurations. + + Args: + head_deleted_dict (Dict): Dictionary of deleted fields in the head. + head_append_dict (Dict): Dictionary of appended fields in the head. + neck (Dict): Updated neck configuration. + """ + if len(head_deleted_dict) + len(head_append_dict) == 0: + return + + old_model_info, new_model_info = build_model_info( + head_deleted_dict, head_append_dict, neck + ) + + total_info = ( + "\nThe config you are using is outdated. " + "The following section of the config:\n```\n" + ) + total_info += old_model_info + total_info += "```\nshould be updated to\n```\n" + total_info += new_model_info + total_info += ( + "```\nFor more information, please refer to " + "https://mmpose.readthedocs.io/en/latest/" + "guide_to_framework.html#step3-model" + ) + + logger: MMLogger = MMLogger.get_current_instance() + logger.warning(total_info) + + +def build_model_info( + head_deleted_dict: Dict, head_append_dict: Dict, neck: Dict +) -> Tuple[str, str]: + """Build the old and new model information strings. + Args: + head_deleted_dict (Dict): Dictionary of deleted fields in the head. + head_append_dict (Dict): Dictionary of appended fields in the head. + neck (Dict): Updated neck configuration. + + Returns: + Tuple[str, str]: Old and new model information strings. + """ + old_head_info = build_head_info(head_deleted_dict) + new_head_info = build_head_info(head_append_dict) + neck_info = build_neck_info(neck) + + old_model_info = "model=dict(\n" + " " * 4 + "...,\n" + old_head_info + new_model_info = "model=dict(\n" + " " * 4 + "...,\n" + neck_info + new_head_info + + return old_model_info, new_model_info + + +def build_head_info(head_dict: Dict) -> str: + """Build the head information string. + + Args: + head_dict (Dict): Dictionary of fields in the head configuration. + Returns: + str: Head information string. + """ + head_info = " " * 4 + "head=dict(\n" + for key, value in head_dict.items(): + head_info += " " * 8 + f"{key}={value},\n" + head_info += " " * 8 + "...),\n" + return head_info + + +def build_neck_info(neck: Dict) -> str: + """Build the neck information string. + Args: + neck (Dict): Updated neck configuration. + + Returns: + str: Neck information string. + """ + if len(neck) > 0: + neck = neck.copy() + neck_info = ( + " " * 4 + "neck=dict(\n" + " " * 8 + f'type=\'{neck.pop("type")}\',\n' + ) + for key, value in neck.items(): + neck_info += " " * 8 + f"{key}={str(value)},\n" + neck_info += " " * 4 + "),\n" + else: + neck_info = "" + return neck_info diff --git a/video/fake-stormer/model_code/models/utils/utils.py b/video/fake-stormer/model_code/models/utils/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..469f8e4c3403cccd8de8699eb4116cccbd13274a --- /dev/null +++ b/video/fake-stormer/model_code/models/utils/utils.py @@ -0,0 +1,658 @@ +# -*- coding: utf-8 -*- +from __future__ import absolute_import, division, print_function + +import copy +import io +import os +import os.path as osp +import pkgutil +import re +import warnings +from collections import OrderedDict +from importlib import import_module +from tempfile import TemporaryDirectory + +import mmcv +import torch +import torch.nn as nn +import torchvision +from mmcv.parallel import is_module_wrapper +from mmengine.dist import get_dist_info +from mmengine.fileio import FileClient +from mmengine.fileio import load as load_file +from mmengine.utils import mkdir_or_exist +from ptflops import get_model_complexity_info +from torch.utils import model_zoo + +ENV_MMCV_HOME = "MMCV_HOME" +ENV_XDG_CACHE_HOME = "XDG_CACHE_HOME" +DEFAULT_CACHE_DIR = "~/.cache" + + +layers_position = { + "PoseResNet_50": 158, + "PoseResNet_101": 311, + "PoseEfficientNet_B4": 415, +} + + +def m_flops(model, cfg): + if cfg.DATASET.DATA_TYPE != "video": + macs, params = get_model_complexity_info( + model, + (3, cfg.DATASET.IMAGE_SIZE[0], cfg.DATASET.IMAGE_SIZE[0]), + as_strings=True, + verbose=True, + ) + else: + macs, params = get_model_complexity_info( + model, + ( + 3, + cfg.DATASET.DATA.SAMPLES_PER_VIDEO.TRAIN, + cfg.DATASET.IMAGE_SIZE[0], + cfg.DATASET.IMAGE_SIZE[0], + ), + as_strings=True, + verbose=True, + ) + print("{:<30} {:<8}".format("Computational complexity: ", macs)) + print("{:<30} {:<8}".format("Number of parameters: ", params)) + + +def n_param_model(model): + print("Number of parameters", sum(p.numel() for p in model.parameters())) + print( + "Number of trainable parameters", + sum(p.numel() for p in model.parameters() if p.requires_grad), + ) + + +def preset_model(cfg, model, optimizer=None, scaler=None): + # Loading models from config, make sure the pretrained path correct to the model name + start_epoch = 0 + if "pretrained" in cfg.TRAIN and os.path.isfile(cfg.TRAIN.pretrained): + model, optimizer, start_epoch, scaler = load_model( + model, + cfg.TRAIN.pretrained, + optimizer=optimizer, + scaler=scaler, + resume=cfg.TRAIN.resume, + lr=cfg.TRAIN.lr, + lr_step=cfg.TRAIN.lr_scheduler.milestones, + gamma=cfg.TRAIN.lr_scheduler.gamma, + ) + else: + model.init_weights(**cfg.MODEL.INIT_WEIGHTS) + print("Loading model successfully -- {}".format(cfg.MODEL.type)) + + # Showing model FLOPS + m_flops(model, cfg) + + # Freeze backbone if begin_epoch < warm up + if cfg.TRAIN.freeze_backbone and start_epoch < cfg.TRAIN.warm_up: + freeze_backbone(cfg.MODEL, model) + + n_param_model(model) + return model, optimizer, start_epoch, scaler + + +def load_pretrained(model, weight_path): + """ + This function only care about state dict of model + For other modules such as optimizer, resume learning, please refer @load_model + """ + state_dict = torch.load(weight_path, map_location=torch.device("cpu"))["state_dict"] + model.load_state_dict(state_dict, strict=True) + return model + + +def freeze_backbone(cfg, model): + """ + This func to freeze some specific layers to warm up the models + """ + if hasattr(model, "backbone"): + backbone = model.backbone + for param in backbone.parameters(): + param.requires_grad = False + else: + for i, (n, p) in enumerate(model.named_parameters()): + if i <= layers_position[f"{cfg.type}_{cfg.num_layers}"]: + p.requires_grad = False + + +def unfreeze_backbone(model): + """ + This func to unfreeze all model layers + """ + for param in model.parameters(): + if not param.requires_grad: + param.requires_grad = True + + +def load_model( + model, + model_path, + optimizer=None, + scaler=None, + resume=False, + lr=None, + lr_step=None, + gamma=None, +): + start_epoch = 0 + checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage) + print("loaded {}, epoch {}".format(model_path, checkpoint["epoch"])) + state_dict_ = checkpoint["state_dict"] + state_dict = {} + + # convert data_parallal to model + for k in state_dict_: + if k.startswith("module") and not k.startswith("module_list"): + state_dict[k[7:]] = state_dict_[k] + else: + state_dict[k] = state_dict_[k] + model_state_dict = model.state_dict() + + # check loaded parameters and created model parameters + msg = ( + "If you see this, your model does not fully load the " + + "pre-trained weight. Please make sure " + + "you have correctly specified --arch xxx " + + "or set the correct --num_classes for your own dataset." + ) + for k in state_dict: + if k in model_state_dict: + if state_dict[k].shape != model_state_dict[k].shape: + print( + "Skip loading parameter {}, required shape{}, " + "loaded shape{}. {}".format( + k, model_state_dict[k].shape, state_dict[k].shape, msg + ) + ) + state_dict[k] = model_state_dict[k] + else: + print("Drop parameter {}.".format(k) + msg) + for k in model_state_dict: + if not (k in state_dict): + print("No param {}.".format(k) + msg) + state_dict[k] = model_state_dict[k] + model.load_state_dict(state_dict, strict=False) + + # resume optimizer parameters + if optimizer is not None and resume: + if "optimizer" in checkpoint: + optimizer.load_state_dict(checkpoint["optimizer"]) + if "scaler" in checkpoint: + scaler.load_state_dict(checkpoint["scaler"]) + + start_epoch = checkpoint["epoch"] + 1 + start_lr = lr + for step in lr_step: + if start_epoch >= step: + start_lr *= gamma + for param_group in optimizer.param_groups: + param_group["lr"] = start_lr + print("Resumed optimizer with start lr", start_lr) + else: + print("No optimizer parameters in checkpoint.") + + return model, optimizer, start_epoch, scaler + + +def save_model(path, epoch, model, optimizer=None, scaler=None): + if isinstance(model, torch.nn.DataParallel): + state_dict = model.module.state_dict() + else: + state_dict = model.state_dict() + data = {"epoch": epoch, "state_dict": state_dict} + + if not (optimizer is None): + data["optimizer"] = optimizer.state_dict() + + if not (scaler is None): + data["scaler"] = scaler.state_dict() + + torch.save(data, path) + + +def _get_mmcv_home(): + mmcv_home = os.path.expanduser( + os.getenv( + ENV_MMCV_HOME, + os.path.join(os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), "mmcv"), + ) + ) + + mkdir_or_exist(mmcv_home) + return mmcv_home + + +def load_url_dist(url, model_dir=None, map_location="cpu"): + """In distributed setting, this function only download checkpoint at local + rank 0.""" + rank, world_size = get_dist_info() + rank = int(os.environ.get("LOCAL_RANK", rank)) + if rank == 0: + checkpoint = model_zoo.load_url( + url, model_dir=model_dir, map_location=map_location + ) + if world_size > 1: + torch.distributed.barrier() + if rank > 0: + checkpoint = model_zoo.load_url( + url, model_dir=model_dir, map_location=map_location + ) + return checkpoint + + +def get_torchvision_models(): + model_urls = dict() + for _, name, ispkg in pkgutil.walk_packages(torchvision.models.__path__): + if ispkg: + continue + _zoo = import_module(f"torchvision.models.{name}") + if hasattr(_zoo, "model_urls"): + _urls = getattr(_zoo, "model_urls") + model_urls.update(_urls) + return model_urls + + +def get_external_models(): + mmcv_home = _get_mmcv_home() + default_json_path = osp.join(mmcv.__path__[0], "model_zoo/open_mmlab.json") + default_urls = load_file(default_json_path) + assert isinstance(default_urls, dict) + external_json_path = osp.join(mmcv_home, "open_mmlab.json") + if osp.exists(external_json_path): + external_urls = load_file(external_json_path) + assert isinstance(external_urls, dict) + default_urls.update(external_urls) + + return default_urls + + +def get_deprecated_model_names(): + deprecate_json_path = osp.join(mmcv.__path__[0], "model_zoo/deprecated.json") + deprecate_urls = load_file(deprecate_json_path) + assert isinstance(deprecate_urls, dict) + + return deprecate_urls + + +def get_mmcls_models(): + mmcls_json_path = osp.join(mmcv.__path__[0], "model_zoo/mmcls.json") + mmcls_urls = load_file(mmcls_json_path) + + return mmcls_urls + + +def _process_mmcls_checkpoint(checkpoint): + state_dict = checkpoint["state_dict"] + new_state_dict = OrderedDict() + for k, v in state_dict.items(): + if k.startswith("backbone."): + new_state_dict[k[9:]] = v + new_checkpoint = dict(state_dict=new_state_dict) + + return new_checkpoint + + +def load_pavimodel_dist(model_path, map_location=None): + """In distributed setting, this function only download checkpoint at local + rank 0.""" + try: + from pavi import modelcloud + except ImportError: + raise ImportError("Please install pavi to load checkpoint from modelcloud.") + rank, world_size = get_dist_info() + rank = int(os.environ.get("LOCAL_RANK", rank)) + if rank == 0: + model = modelcloud.get(model_path) + with TemporaryDirectory() as tmp_dir: + downloaded_file = osp.join(tmp_dir, model.name) + model.download(downloaded_file) + checkpoint = torch.load(downloaded_file, map_location=map_location) + if world_size > 1: + torch.distributed.barrier() + if rank > 0: + model = modelcloud.get(model_path) + with TemporaryDirectory() as tmp_dir: + downloaded_file = osp.join(tmp_dir, model.name) + model.download(downloaded_file) + checkpoint = torch.load(downloaded_file, map_location=map_location) + return checkpoint + + +def load_fileclient_dist(filename, backend, map_location): + """In distributed setting, this function only download checkpoint at local + rank 0.""" + rank, world_size = get_dist_info() + rank = int(os.environ.get("LOCAL_RANK", rank)) + allowed_backends = ["ceph"] + if backend not in allowed_backends: + raise ValueError(f"Load from Backend {backend} is not supported.") + if rank == 0: + fileclient = FileClient(backend=backend) + buffer = io.BytesIO(fileclient.get(filename)) + checkpoint = torch.load(buffer, map_location=map_location) + if world_size > 1: + torch.distributed.barrier() + if rank > 0: + fileclient = FileClient(backend=backend) + buffer = io.BytesIO(fileclient.get(filename)) + checkpoint = torch.load(buffer, map_location=map_location) + return checkpoint + + +def load_state_dict(module, state_dict, strict=False, logger=None): + """Load state_dict to a module. + + This method is modified from :meth:`torch.nn.Module.load_state_dict`. + Default value for ``strict`` is set to ``False`` and the message for + param mismatch will be shown even if strict is False. + + Args: + module (Module): Module that receives the state_dict. + state_dict (OrderedDict): Weights. + strict (bool): whether to strictly enforce that the keys + in :attr:`state_dict` match the keys returned by this module's + :meth:`~torch.nn.Module.state_dict` function. Default: ``False``. + logger (:obj:`logging.Logger`, optional): Logger to log the error + message. If not specified, print function will be used. + """ + unexpected_keys = [] + all_missing_keys = [] + err_msg = [] + + metadata = getattr(state_dict, "_metadata", None) + state_dict = state_dict.copy() + if metadata is not None: + state_dict._metadata = metadata + + # use _load_from_state_dict to enable checkpoint version control + def load(module, prefix=""): + # recursively check parallel module in case that the model has a + # complicated structure, e.g., nn.Module(nn.Module(DDP)) + if is_module_wrapper(module): + module = module.module + local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {}) + module._load_from_state_dict( + state_dict, + prefix, + local_metadata, + True, + all_missing_keys, + unexpected_keys, + err_msg, + ) + for name, child in module._modules.items(): + if child is not None: + load(child, prefix + name + ".") + + load(module) + load = None # break load->load reference cycle + + # ignore "num_batches_tracked" of BN layers + missing_keys = [key for key in all_missing_keys if "num_batches_tracked" not in key] + + if unexpected_keys: + err_msg.append( + "unexpected key in source " f'state_dict: {", ".join(unexpected_keys)}\n' + ) + if missing_keys: + err_msg.append( + f'missing keys in source state_dict: {", ".join(missing_keys)}\n' + ) + + rank, _ = get_dist_info() + if len(err_msg) > 0 and rank == 0: + err_msg.insert(0, "The model and loaded state dict do not match exactly\n") + err_msg = "\n".join(err_msg) + if strict: + raise RuntimeError(err_msg) + elif logger is not None: + logger.warning(err_msg) + else: + print(err_msg) + + +def _load_checkpoint(filename, map_location=None): + """Load checkpoint from somewhere (modelzoo, file, url). + + Args: + filename (str): Accept local filepath, URL, ``torchvision://xxx``, + ``open-mmlab://xxx``. Please refer to ``docs/model_zoo.md`` for + details. + map_location (str | None): Same as :func:`torch.load`. Default: None. + + Returns: + dict | OrderedDict: The loaded checkpoint. It can be either an + OrderedDict storing model weights or a dict containing other + information, which depends on the checkpoint. + """ + if filename.startswith("modelzoo://"): + warnings.warn( + 'The URL scheme of "modelzoo://" is deprecated, please ' + 'use "torchvision://" instead' + ) + model_urls = get_torchvision_models() + model_name = filename[11:] + checkpoint = load_url_dist(model_urls[model_name]) + elif filename.startswith("torchvision://"): + model_urls = get_torchvision_models() + model_name = filename[14:] + checkpoint = load_url_dist(model_urls[model_name]) + elif filename.startswith("open-mmlab://"): + model_urls = get_external_models() + model_name = filename[13:] + deprecated_urls = get_deprecated_model_names() + if model_name in deprecated_urls: + warnings.warn( + f"open-mmlab://{model_name} is deprecated in favor " + f"of open-mmlab://{deprecated_urls[model_name]}" + ) + model_name = deprecated_urls[model_name] + model_url = model_urls[model_name] + # check if is url + if model_url.startswith(("http://", "https://")): + checkpoint = load_url_dist(model_url) + else: + filename = osp.join(_get_mmcv_home(), model_url) + if not osp.isfile(filename): + raise IOError(f"{filename} is not a checkpoint file") + checkpoint = torch.load(filename, map_location=map_location) + elif filename.startswith("mmcls://"): + model_urls = get_mmcls_models() + model_name = filename[8:] + checkpoint = load_url_dist(model_urls[model_name]) + checkpoint = _process_mmcls_checkpoint(checkpoint) + elif filename.startswith(("http://", "https://")): + checkpoint = load_url_dist(filename) + elif filename.startswith("pavi://"): + model_path = filename[7:] + checkpoint = load_pavimodel_dist(model_path, map_location=map_location) + elif filename.startswith("s3://"): + checkpoint = load_fileclient_dist( + filename, backend="ceph", map_location=map_location + ) + else: + if not osp.isfile(filename): + raise IOError(f"{filename} is not a checkpoint file") + checkpoint = torch.load(filename, map_location=map_location) + return checkpoint + + +def load_checkpoint( + model, + filename, + map_location="cpu", + strict=False, + logger=None, + patch_padding="pad", + part_features=None, +): + """Load checkpoint from a file or URI. + + Args: + model (Module): Module to load checkpoint. + filename (str): Accept local filepath, URL, ``torchvision://xxx``, + ``open-mmlab://xxx``. Please refer to ``docs/model_zoo.md`` for + details. + map_location (str): Same as :func:`torch.load`. + strict (bool): Whether to allow different params for the model and + checkpoint. + logger (:mod:`logging.Logger` or None): The logger for error message. + patch_padding (str): 'pad' or 'bilinear' or 'bicubic', used for interpolate patch embed from 14x14 to 16x16 + + Returns: + dict or OrderedDict: The loaded checkpoint. + """ + checkpoint = _load_checkpoint(filename, map_location) + # OrderedDict is a subclass of dict + if not isinstance(checkpoint, dict): + raise RuntimeError(f"No state_dict found in checkpoint file {filename}") + # get state_dict from checkpoint + if "state_dict" in checkpoint: + state_dict = checkpoint["state_dict"] + elif "model" in checkpoint: + state_dict = checkpoint["model"] + elif "module" in checkpoint: + state_dict = checkpoint["module"] + else: + state_dict = checkpoint + # strip prefix of state_dict + if list(state_dict.keys())[0].startswith("module."): + state_dict = {k[7:]: v for k, v in state_dict.items()} + + # for MoBY, load model of online branch + if sorted(list(state_dict.keys()))[0].startswith("encoder"): + state_dict = { + k.replace("encoder.", ""): v + for k, v in state_dict.items() + if k.startswith("encoder.") + } + + rank, _ = get_dist_info() + + layer_names = [name for name, param in model.named_parameters()] + + if ( + "patch_embed.proj.weight" in state_dict + and "patch_embed.proj._conv_stem.weight" not in layer_names + ): + proj_weight = state_dict["patch_embed.proj.weight"] + orig_size = proj_weight.shape[2:] + current_size = model.patch_embed.proj.weight.shape[2:] + padding_size = current_size[0] - orig_size[0] + padding_l = padding_size // 2 + padding_r = padding_size - padding_l + if orig_size != current_size: + if "pad" in patch_padding: + proj_weight = torch.nn.functional.pad( + proj_weight, (padding_l, padding_r, padding_l, padding_r) + ) + elif "bilinear" in patch_padding: + proj_weight = torch.nn.functional.interpolate( + proj_weight, size=current_size, mode="bilinear", align_corners=False + ) + elif "bicubic" in patch_padding: + proj_weight = torch.nn.functional.interpolate( + proj_weight, size=current_size, mode="bicubic", align_corners=False + ) + state_dict["patch_embed.proj.weight"] = proj_weight + + if "pos_embed" in state_dict: + pos_embed_checkpoint = state_dict["pos_embed"] + # pos_embed_checkpoint = pos_embed_checkpoint[:, :-1, :] + embedding_size = pos_embed_checkpoint.shape[-1] + H, W = model.patch_embed.patch_shape + num_patches = model.patch_embed.num_patches + num_extra_tokens = model.pos_embed.shape[-2] - num_patches + # height (== width) for the checkpoint position embedding + orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5) + if rank == 0: + print( + "Position interpolate from %dx%d to %dx%d" + % (orig_size, orig_size, H, W) + ) + extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens] + # only the position tokens are interpolated + pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:] + pos_tokens = pos_tokens.reshape( + -1, orig_size, orig_size, embedding_size + ).permute(0, 3, 1, 2) + pos_tokens = torch.nn.functional.interpolate( + pos_tokens, size=(H, W), mode="bicubic", align_corners=False + ) + pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2) + new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1) + state_dict["pos_embed"] = new_pos_embed + + new_state_dict = copy.deepcopy(state_dict) + if part_features is not None: + current_keys = list(model.state_dict().keys()) + for key in current_keys: + if "mlp.experts" in key: + source_key = re.sub(r"experts.\d+.", "fc2.", key) + new_state_dict[key] = state_dict[source_key][-part_features:] + elif "fc2" in key: + new_state_dict[key] = state_dict[key][:-part_features] + + # load state_dict + load_state_dict(model, new_state_dict, strict, logger) + return checkpoint + + +def swin_converter(ckpt): + + new_ckpt = OrderedDict() + + def correct_unfold_reduction_order(x): + out_channel, in_channel = x.shape + x = x.reshape(out_channel, 4, in_channel // 4) + x = x[:, [0, 2, 1, 3], :].transpose(1, 2).reshape(out_channel, in_channel) + return x + + def correct_unfold_norm_order(x): + in_channel = x.shape[0] + x = x.reshape(4, in_channel // 4) + x = x[[0, 2, 1, 3], :].transpose(0, 1).reshape(in_channel) + return x + + for k, v in ckpt.items(): + if k.startswith("head"): + continue + elif k.startswith("layers"): + new_v = v + if "attn." in k: + new_k = k.replace("attn.", "attn.w_msa.") + elif "mlp." in k: + if "mlp.fc1." in k: + new_k = k.replace("mlp.fc1.", "ffn.layers.0.0.") + elif "mlp.fc2." in k: + new_k = k.replace("mlp.fc2.", "ffn.layers.1.") + else: + new_k = k.replace("mlp.", "ffn.") + elif "downsample" in k: + new_k = k + if "reduction." in k: + new_v = correct_unfold_reduction_order(v) + elif "norm." in k: + new_v = correct_unfold_norm_order(v) + else: + new_k = k + new_k = new_k.replace("layers", "stages", 1) + elif k.startswith("patch_embed"): + new_v = v + if "proj" in k: + new_k = k.replace("proj", "projection") + else: + new_k = k + else: + new_v = v + new_k = k + + new_ckpt["backbone." + new_k] = new_v + + return new_ckpt diff --git a/video/fake-stormer/model_code/package_utils/__init__.py b/video/fake-stormer/model_code/package_utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..40a96afc6ff09d58a702b76e3f7dd412fe975e26 --- /dev/null +++ b/video/fake-stormer/model_code/package_utils/__init__.py @@ -0,0 +1 @@ +# -*- coding: utf-8 -*- diff --git a/video/fake-stormer/model_code/package_utils/_typing.py b/video/fake-stormer/model_code/package_utils/_typing.py new file mode 100644 index 0000000000000000000000000000000000000000..37b26c25145f281bc23b6f98915deb58f886ffb1 --- /dev/null +++ b/video/fake-stormer/model_code/package_utils/_typing.py @@ -0,0 +1,22 @@ +# -*- coding: utf-8 -*- +from typing import Dict, List, Optional, Tuple, Union + +from mmengine.config import ConfigDict +from mmengine.structures import InstanceData, PixelData +from torch import Tensor + +# Type hint of config data +ConfigType = Union[ConfigDict, dict] +OptConfigType = Optional[ConfigType] + +# Type hint of one or more config data +MultiConfig = Union[ConfigType, List[ConfigType]] +OptMultiConfig = Optional[MultiConfig] + +# Type hint of data samples +InstanceList = List[InstanceData] +PixelDataList = List[PixelData] +Predictions = Union[InstanceList, Tuple[InstanceList, PixelDataList]] + +# Type hint of features +Features = Union[Tuple[Tensor], List[Tuple[Tensor]], List[List[Tuple[Tensor]]]] diff --git a/video/fake-stormer/model_code/package_utils/bi_online_generation.py b/video/fake-stormer/model_code/package_utils/bi_online_generation.py new file mode 100644 index 0000000000000000000000000000000000000000..dfd4a95891008f049b77eced6edb96a21e710728 --- /dev/null +++ b/video/fake-stormer/model_code/package_utils/bi_online_generation.py @@ -0,0 +1,551 @@ +# -*- coding: utf-8 -*- +import os +import random +import sys + +if os.getcwd() not in sys.path: + sys.path.insert(0, os.getcwd()) + +import argparse +import multiprocessing as mp +import queue +import time +from threading import Thread + +import cv2 +import numpy as np +from imgaug import augmenters as iaa +from matplotlib import pyplot as plt +from package_utils.deepfake_mask import random_get_hull +from package_utils.utils import load_file, save_file +from PIL import Image +from skimage import io +from skimage import transform as sktransform +from tqdm import tqdm + +# from datasets.sbi.utils import gen_target + +IMAGE_H, IMAGE_W, IMAGE_CHANNEL = 256, 256, 3 +COMPRESSION = "c0" +SPLIT = "train" +DATA_TYPE = "frames" +IMAGE_ROOT = f"/data/deepfake_cluster/datasets_df/FaceForensics++/{COMPRESSION}/" +ANNO_FILE = "processed_data/train_faceforensics_processed.json" +DEST_DIR = "FaceXRay" +LABEL_FILE = "train_FF_FaceXRay.json" +NUMBER_OF_PROCESS = 4 +BLENDING_TYPE = "BI" +MARGIN = 20 + + +def args_parse(args=None): + args_parser = argparse.ArgumentParser("Blending Image Processing Hub...") + args_parser.add_argument("--task", "-t", help="Defining task!") + args_parser.add_argument("--anno_file", "-f", help="Pre annotation file!") + args_parser.add_argument("--fake_type", "-ft", help="Faketype to manipulation!") + args_parser.add_argument( + "--mp", "-m", help="Apply multiprocessing", action="store_true" + ) + args = args_parser.parse_args(args) + return args + + +def name_resolve(path): + if COMPRESSION == "c23": + name = os.path.splitext(os.path.basename(path))[0] + vid_id, frame_id = name.split("_")[0:2] + else: + name = path.split("/") + vid_id, frame_id = name[-2], os.path.splitext(name[-1])[0] + return vid_id, frame_id + + +def gen_real_fimg_mask(img_path): + face_img = io.imread(img_path) + mask = np.zeros((face_img.shape[0], face_img.shape[1], 3)) + mask = (mask * 255).astype(np.uint8) + return face_img, mask + + +def total_euclidean_distance(a, b): + assert len(a.shape) == 2 + return np.sum(np.linalg.norm(a - b, axis=1)) + + +def random_erode_dilate(mask, **kwargs): + ksize = kwargs.get("ksize") + rand_erode = kwargs.get("rand_erode") or random.random() + + if rand_erode > 0.5: + if ksize is None: + ksize = random.randint(1, 21) + if ksize % 2 == 0: + ksize += 1 + mask = np.array(mask).astype(np.uint8) * 255 + kernel = np.ones((ksize, ksize), np.uint8) + mask = cv2.erode(mask, kernel, 1) / 255 + else: + if ksize is None: + ksize = random.randint(1, 5) + if ksize % 2 == 0: + ksize += 1 + mask = np.array(mask).astype(np.uint8) * 255 + kernel = np.ones((ksize, ksize), np.uint8) + mask = cv2.dilate(mask, kernel, 1) / 255 + return mask, ksize, rand_erode + + +# borrow from https://github.com/MarekKowalski/FaceSwap +def blendImages(src, dst, mask, featherAmount=0.2, **kwargs): + if kwargs.get("blend_ratio") is None: + blend_list = [0.25, 0.5, 0.75, 1, 1, 1] + blend_ratio = blend_list[np.random.randint(len(blend_list))] + else: + blend_ratio = kwargs.get("blend_ratio") + + # mask = blend_ratio * mask # Applying a blending ratio from SBI to BI blending weights + maskIndices = np.where(mask != 0) + + src_mask = np.ones_like(mask) + dst_mask = np.zeros_like(mask) + + maskPts = np.hstack((maskIndices[1][:, np.newaxis], maskIndices[0][:, np.newaxis])) + faceSize = np.max(maskPts, axis=0) - np.min(maskPts, axis=0) + featherAmount = featherAmount * np.max(faceSize) + + hull = cv2.convexHull(maskPts) + dists = np.zeros(maskPts.shape[0]) + + for i in range(maskPts.shape[0]): + dists[i] = cv2.pointPolygonTest(hull, (maskPts[i, 0], maskPts[i, 1]), True) + + weights = np.clip(dists / featherAmount, 0, 1) + + composedImg = np.copy(dst) + composedImg[maskIndices[0], maskIndices[1]] = ( + weights[:, np.newaxis] * src[maskIndices[0], maskIndices[1]] + + (1 - weights[:, np.newaxis]) * dst[maskIndices[0], maskIndices[1]] + ) + + composedMask = np.copy(dst_mask) + composedMask[maskIndices[0], maskIndices[1]] = ( + weights[:, np.newaxis] * src_mask[maskIndices[0], maskIndices[1]] + + (1 - weights[:, np.newaxis]) * dst_mask[maskIndices[0], maskIndices[1]] + ) + + blend_params = {"blend_ratio": blend_ratio} + + return composedImg, composedMask, blend_params + + +# borrow from https://github.com/MarekKowalski/FaceSwap +def colorTransfer(src, dst, mask): + transferredDst = np.copy(dst) + + maskIndices = np.where(mask != 0) + + maskedSrc = src[maskIndices[0], maskIndices[1]].astype(np.int32) + maskedDst = dst[maskIndices[0], maskIndices[1]].astype(np.int32) + + meanSrc = np.mean(maskedSrc, axis=0) + meanDst = np.mean(maskedDst, axis=0) + + maskedDst = maskedDst - meanDst + maskedDst = maskedDst + meanSrc + maskedDst = np.clip(maskedDst, 0, 255) + + transferredDst[maskIndices[0], maskIndices[1]] = maskedDst + + return transferredDst + + +class BIOnlineGeneration: + def __init__( + self, data_record, queue_size=1024, mlprocess=False, number=1, fake_type=None + ): + self.landmarks_record = {} + self.data_record = data_record["data"] + self.mlprocess = mlprocess + self.number = number + self.fake_type = fake_type + + if self.fake_type is not None: + self.data_record = self._filter_data() + if not len(self.data_record): + raise ValueError("DataList can not be Empty!") + + print( + f"You are generating data for --- {self.fake_type} --- {len(self.data_record)} images" + ) + + for item in self.data_record: + if "aligned_lms" in item.keys() and len(item["aligned_lms"]): + self.landmarks_record[item["image_path"]] = np.array( + item["aligned_lms"] + ) + else: + self.landmarks_record[item["image_path"]] = np.array(item["orig_lms"]) + + # extract all frame from all video in the name of {videoid}_{frameid} + self.data_list = [item["image_path"] for item in self.data_record] + self.file_names = [item["file_name"] for item in self.data_record] + + if COMPRESSION != "c23": + self.labels = [ + item["image_path"].split("/")[-3] for item in self.data_record + ] + else: + self.labels = [ + item["image_path"].split("/")[-2] for item in self.data_record + ] + self.vid_ids = [item["image_path"].split("/")[-2] for item in self.data_record] + + # predefine mask distortion + self.distortion = iaa.Sequential([iaa.PiecewiseAffine(scale=(0.01, 0.15))]) + + self.result_queue = mp.Queue(maxsize=queue_size) + self.final_results = [] + + def register_task(self, target): + if self.mlprocess: + p = mp.Process(target=target, args=()) + else: + p = Thread(target=target, args=()) + return p + + def start(self, p): + self.result_worker = p + self.result_worker.start() + + def wait_n_put(self, item): + self.result_queue.put(item) + + def wait_n_get(self): + return self.result_queue.get() + + def count(self): + return self.result_queue.qsize() + + def stop(self): + self.result_worker.join() + + def running(self): + return not self.result_queue.empty() + + def terminate(self): + self.result_worker.terminate() + + def clear(self): + while not self.result_queue.empty(): + self.result_queue.get() + + def clear_sequences(self): + self.clear() + + def get_results(self): + all_objs = [] + + while not self.final_results.empty(): + all_objs.append(self.final_results.get()) + return all_objs + + def _filter_data(self): + assert self.fake_type is not None, "Fake type is require to filter data!" + assert self.fake_type in [ + "Deepfakes", + "Face2Face", + "FaceSwap", + "NeuralTextures", + ] + + self.data_record = [ + item for item in self.data_record if item["fake_type"] == self.fake_type + ] + return self.data_record + + def gen_one_datapoint(self, idx): + background_face_path = self.data_list[idx] + label = self.labels[idx] + + # Choose Blending type + if BLENDING_TYPE != "SBI": + data_type = "real" if random.randint(0, 1) else "fake" + else: + data_type = "fake" + + # Handle for the cases of blank landmarks, auto real if real label is real image, otherwise return + if not self.landmarks_record[background_face_path].any(): + if not ("fake" == label): + data_type = "real" + else: + return None, None, None + + if data_type == "fake": + if BLENDING_TYPE != "SBI": + face_img, mask = self.get_blended_face( + background_face_path, self.landmarks_record[background_face_path] + ) + else: + face_img, mask, face_r, mask_r = gen_target( + os.path.join(IMAGE_ROOT, background_face_path), + self.landmarks_record[background_face_path], + margin=[MARGIN, MARGIN], + ) + else: + face_img, mask = gen_real_fimg_mask( + os.path.join(IMAGE_ROOT, background_face_path) + ) + + face_img = face_img[MARGIN : IMAGE_H - MARGIN, MARGIN : IMAGE_W - MARGIN, :] + mask = mask[MARGIN : IMAGE_H - MARGIN, MARGIN : IMAGE_W - MARGIN, :] + + return face_img, mask, data_type + + def get_blended_face(self, background_face_path, background_landmark): + background_face = io.imread(os.path.join(IMAGE_ROOT, background_face_path)) + + foreground_face_path = self.search_similar_face( + background_landmark, background_face_path, get_best=True + ) + foreground_face = io.imread(os.path.join(IMAGE_ROOT, foreground_face_path)) + + # down sample before blending + img_h, img_w = background_face.shape[:2] + aug_size = random.randint(img_h // 2, img_h) + background_landmark = background_landmark * (aug_size / img_h) + + foreground_face = sktransform.resize( + foreground_face, (aug_size, aug_size), preserve_range=True + ).astype(np.uint8) + background_face = sktransform.resize( + background_face, (aug_size, aug_size), preserve_range=True + ).astype(np.uint8) + + # get random type of initial blending mask + mask = random_get_hull(background_landmark, background_face) + + # random deform mask + mask = self.distortion.augment_image(mask) + mask = random_erode_dilate(mask) + + # filte empty mask after deformation + if np.sum(mask) == 0: + print( + f"There was an issue when doing blending with Image -- {background_face_path}" + ) + print(f"Reverting by returning a real image and mask...") + face_img, mask = gen_real_fimg_mask( + os.path.join(IMAGE_ROOT, background_face_path) + ) + return face_img, mask + + # apply color transfer + foreground_face = colorTransfer(background_face, foreground_face, mask * 255) + + # blend two face + blended_face, mask = blendImages(foreground_face, background_face, mask * 255) + blended_face = blended_face.astype(np.uint8) + + # resize back to default resolution + blended_face = sktransform.resize( + blended_face, (img_h, img_w), preserve_range=True + ).astype(np.uint8) + mask = sktransform.resize(mask, (img_h, img_w), preserve_range=True) + mask = mask[:, :, 0:1] + mask = (1 - mask) * mask * 4 + mask = np.repeat(mask, 3, 2) + mask = (mask * 255).astype(np.uint8) + + # randomly downsample after BI pipeline + face_img = Image.fromarray(blended_face) + if random.randint(0, 1): + aug_size = random.randint(img_h // 4, img_h) + if random.randint(0, 1): + face_img = face_img.resize((aug_size, aug_size), Image.BILINEAR) + else: + face_img = face_img.resize((aug_size, aug_size), Image.NEAREST) + face_img = face_img.resize((img_h, img_w), Image.BILINEAR) + face_img = np.array(face_img) + + # # random jpeg compression after BI pipeline + # if random.randint(0,1): + # quality = random.randint(60, 100) + # encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality] + # face_img_encode = cv2.imencode('.jpg', face_img, encode_param)[1] + # face_img = cv2.imdecode(face_img_encode, cv2.IMREAD_COLOR) + + # # random flip + # if random.randint(0,1): + # face_img = np.flip(face_img,1) + # mask = np.flip(mask,1) + + return face_img, mask + + def search_similar_face(self, this_landmark, background_face_path, get_best=False): + vid_id, frame_id = name_resolve(background_face_path) + min_dist = 99999999 + + # random sample 5000 frame from all frams: + all_candidate_path = random.sample(self.data_list, k=10000) + + # filter all frame that comes from the same video as background face + all_candidate_path = filter( + lambda k: name_resolve(k)[0] != vid_id, all_candidate_path + ) + all_candidate_path = list(all_candidate_path) + candidate_distance_list = {} + + # loop throungh all candidates frame to get best match + for candidate_path in all_candidate_path: + candidate_landmark = self.landmarks_record[candidate_path].astype( + np.float32 + ) + + if not candidate_landmark.any(): + continue + + candidate_distance = total_euclidean_distance( + candidate_landmark, this_landmark + ) + if candidate_distance < min_dist: + min_dist = candidate_distance + min_path = candidate_path + candidate_distance_list[candidate_path] = candidate_distance + if not get_best: + return candidate_distance_list + else: + return min_path + + def search_similar_faces(self): + while True: + item = self.wait_n_get() + if item is None: + obj_list = self.final_results + data = {"data": obj_list} + save_file( + data, f"processed_data/{COMPRESSION}/dynamic_trainBI_FFv4.json" + ) + return True + bg_path = item["image_path"] + if "aligned_lms" in item.keys() and len(item["aligned_lms"]): + f_lms = np.array(item["aligned_lms"]) + else: + f_lms = np.array(item["orig_lms"]) + best_match_paths = [] + + if f_lms.any(): + candidate_list = self.search_similar_face(f_lms, bg_path) + best_match_paths = sorted( + candidate_list.items(), key=lambda x: x[1], reverse=False + )[: self.number] + best_match_paths = [it[0] for it in best_match_paths] + item["best_match"] = best_match_paths + self.final_results.append(item) + + +if __name__ == "__main__": + if sys.argv[1:] is not None: + args = args_parse(sys.argv[1:]) + else: + args = sys.argv[:-1] + + task = args.task + anno_file = args.anno_file + mp_ = args.mp + fake_type = args.fake_type + assert len(anno_file), "Annotation file path can not be empty!" + assert os.access( + anno_file, os.R_OK + ), "Annotation file path must be valid to access!" + + print("Starting to load processed data...") + start = time.time() + data_record = load_file(anno_file) + print("Loading time --- {}".format(time.time() - start)) + ds = BIOnlineGeneration(data_record, mlprocess=mp_, number=30, fake_type=fake_type) + data = {} + + assert task in ["save_blending", "search_similar_lms"] + if task == "save_blending": + all_object = [] + + for i in tqdm(range(len(ds.data_list))): + img, mask, label = ds.gen_one_datapoint(i) + if img is None and mask is None and label is None: + continue + + if COMPRESSION == "c23": + image_path = os.path.join( + IMAGE_ROOT, DEST_DIR, "images", ds.file_names[i] + ) + mask_path = os.path.join( + IMAGE_ROOT, DEST_DIR, "masks", ds.file_names[i] + ) + else: + image_path = os.path.join( + IMAGE_ROOT, + SPLIT, + DATA_TYPE, + DEST_DIR, + "images", + f"{ds.vid_ids[i]}_{ds.file_names[i]}", + ) + mask_path = os.path.join( + IMAGE_ROOT, + SPLIT, + DATA_TYPE, + DEST_DIR, + "masks", + f"{ds.vid_ids[i]}_{ds.file_names[i]}", + ) + + try: + mask_pil = Image.fromarray(mask) + mask_pil.save(mask_path) + + image = Image.fromarray(img) + image.save(image_path) + except Exception as e: + print(e) + continue + + all_object.append( + { + "id": i, + "image_path": image_path, + "mask_path": mask_path, + "label": label, + } + ) + data["data"] = all_object + save_file( + data, + file_path=os.path.join(IMAGE_ROOT, SPLIT, DATA_TYPE, DEST_DIR, LABEL_FILE), + ) + elif task == "search_similar_lms": + p = ds.register_task(ds.search_similar_faces) + ds.start(p) + + try: + for i, item in enumerate(tqdm(ds.data_record, dynamic_ncols=True)): + ds.wait_n_put(item) + ds.wait_n_put(None) + + while ds.running(): + time.sleep(1) + print( + "===============> Rendering remaining " + + str(ds.count()) + + " images in the queue...", + end="\r", + ) + ds.stop() + except Exception as e: + print(repr(e)) + print("There is an exception during process! Please check it") + except KeyboardInterrupt: + ds.terminate() + ds.clear_sequences() + exit(0) + else: + raise ValueError("This task {} is not supported at the moment!") diff --git a/video/fake-stormer/model_code/package_utils/cam_vis.py b/video/fake-stormer/model_code/package_utils/cam_vis.py new file mode 100644 index 0000000000000000000000000000000000000000..52c9b52a1fcd42051c2a0e2bd65883b91e5c853e --- /dev/null +++ b/video/fake-stormer/model_code/package_utils/cam_vis.py @@ -0,0 +1,269 @@ +# -*-coding: utf-8 -*- +import argparse +import math +import os +import sys +from copy import deepcopy + +if not os.getcwd() in sys.path: + sys.path.append(os.getcwd()) + +from glob import glob + +import matplotlib.pyplot as plt +import numpy as np +import torch +import torch.nn.functional as F +from configs.get_config import load_config +from image_utils import overlay_mask +from logs.logger import Logger +from models import MODELS, build_model +from models.utils import load_pretrained +from natsort import natsorted +from PIL import Image +from torchcam import methods +from torchvision.transforms.functional import resize, to_pil_image +from transform import final_transform + + +def main(): + argparser = argparse.ArgumentParser("Arguments for CAM visualization...") + argparser.add_argument("--cfg", help="Specify config to load", required=True) + argparser.add_argument( + "--target_layer", + "-t", + help="Specify layer names to visualize CAM", + required=False, + ) + argparser.add_argument( + "--method", "-m", type=str, default="GradCAM", help="CAM method to use" + ) + argparser.add_argument( + "--mode", + type=str, + choices=["image", "video"], + default="image", + help="Mode to visualize gradCAM", + ) + argparser.add_argument( + "--image", "-i", help="Specify image to overlay CAM", required=False + ) + argparser.add_argument( + "--video", "-v", help="Specify video to overlay CAM", required=False + ) + argparser.add_argument( + "--savefig", type=str, default=None, help="Path to save figure" + ) + argparser.add_argument( + "--rows", type=int, default=1, help="Number of rows for the layout" + ) + argparser.add_argument( + "--class-idx", type=int, default=0, help="Index of the class to inspect" + ) + argparser.add_argument( + "--alpha", type=float, default=0.5, help="Transparency of the heatmap" + ) + argparser.add_argument("--cuda", action="store_true", help="Running CAM with cuda") + argparser.add_argument( + "--save_inverse", action="store_true", help="Saving the inverse of CAM" + ) + args = argparser.parse_args() + print(args) + + # Loading configs + cfg = load_config(args.cfg) + + # Logger + logger = Logger(task="CAM_vis") + + # Loading model based on the config + model = build_model(cfg.MODEL, MODELS).to(torch.float) + logger.info("Loading weight ... {}".format(cfg.TEST.pretrained)) + model = load_pretrained(model, cfg.TEST.pretrained) + + if args.cuda: + model = model.cuda() + model.eval() + + # Freeze the model + for p in model.parameters(): + p.requires_grad_(False) + + # Loading image + img_list = [] + if args.mode == "image": + assert os.path.exists( + args.image + ), "Image path must be valid, please check the path again!" + img = Image.open(args.image) + H, W = img.size + img = img.crop((0, 0, W - 0, H - 0)) + img_list.append(img) + elif args.mode == "video": + assert os.path.exists( + args.video + ), "Video path must be valid, please check the path again!" + n_frames = cfg.DATASET.DATA.SAMPLES_PER_VIDEO.NUM_FRAMES + assert n_frames is not None, "Number of video frames can not be None!" + # Load first n_frames inside the video + img_paths = glob(f"{args.video}/*.png") + img_paths = natsorted(img_paths) # correct the order of image paths + img_paths = img_paths[:n_frames] + for img_path in img_paths: + img = Image.open(img_path) + H, W = img.size + img = img.crop((0, 0, W - 0, H - 0)) + img_list.append(img) + else: + raise ValueError( + "We only support GradCAM for image or video mode at the moment!" + ) + + # Preprocess image + transform = final_transform(cfg.DATASET) + image_size = (cfg.DATASET.IMAGE_SIZE[0], cfg.DATASET.IMAGE_SIZE[1]) + + # Transform images + transformed_imgs = torch.tensor([]) + for _i in img_list: + img_resize = _i.resize(image_size) + img_resize = np.array(img_resize) / 255 + img_tensor = transform(img_resize).to(torch.float) + if args.cuda: + img_tensor = img_tensor.cuda() + img_tensor.requires_grad_(True) + transformed_imgs = torch.cat((transformed_imgs, img_tensor.unsqueeze(0)), 0) + + # Hook the corresponding layer in the model + if isinstance(args.method, str): + cam_methods = args.method.split(",") + else: + cam_methods = [ + "CAM", + "GradCAM", + "GradCAMpp", + "SmoothGradCAMpp", + "ScoreCAM", + "SSCAM", + "ISCAM", + "XGradCAM", + "LayerCAM", + ] + cam_extractors = [ + methods.__dict__[name]( + model, target_layer=args.target_layer, enable_hooks=False + ) + for name in cam_methods + ] + + if args.mode == "image": + num_rows = args.rows + num_cols = math.ceil((len(cam_extractors)) / num_rows) + 1 + else: + num_cols = n_frames + num_rows = len(cam_extractors) + 1 + + _, axes = plt.subplots(num_rows, num_cols, figsize=(6, 4)) + # Display input + for idx, _i in enumerate(img_list): + ax = axes[0][idx] if num_rows > 1 else axes[0] if num_cols > 1 else axes + ax.imshow(_i) + ax.set_title("Input", size=8) + + for idx, extractor in zip(range(1, len(cam_extractors) + 1), cam_extractors): + extractor._hooks_enabled = True + model.zero_grad() + if args.mode == "image": + scores = model(transformed_imgs)[0]["cls"].sigmoid() + else: + transformed_imgs = transformed_imgs.transpose(0, 1).unsqueeze(0) + scores = ( + model(transformed_imgs)[0]["hm"] + .sigmoid() + .view(1, -1) + .max(1, keepdim=True) + .values + ) + # output, attn = model(transformed_imgs) # For visualizing the attn scores, will comeback later + # scores = output[0]['temp_loc'].sigmoid() + print("Classification Score -- {}".format(scores)) + + # Select the class index + class_idx = ( + scores.squeeze(0).argmax().item() + if args.class_idx is None + else args.class_idx + ) + # class_idx = img_idx + + # Use the hooked data to compute activation map + activation_map = ( + extractor(class_idx, scores)[0].to(torch.float).squeeze(0).cpu() + ) + # activation_map = torch.cat((activation_map, torch.zeros(4)), 0) + # activation_map = F.adaptive_avg_pool1d(activation_map.unsqueeze(0), 196).squeeze(0) + # activation_map = activation_map[class_idx, 1:].reshape((14, 14)) + + # Clean data + extractor.remove_hooks() + extractor._hooks_enabled = False + + for img_idx, i_ in enumerate(img_list): + # Convert it to PL image + # The indexing below means first image in batch + heatmap = to_pil_image(activation_map[img_idx].unsqueeze(0), mode="F") + # activation_map = attn[img_idx].mean(0)[0, 1:] + # activation_map = activation_map.reshape((14, 14)).detach() + # activation_map = (activation_map - activation_map.min()) / (activation_map.max() - activation_map.min()) + # heatmap = to_pil_image(activation_map.unsqueeze(0), mode='F') + + # Plot the result + result = overlay_mask(deepcopy(i_), heatmap, alpha=args.alpha) + + ax = ( + axes[idx][img_idx] + if num_rows > 1 + else axes[idx] if num_cols > 1 else axes + ) + + ax.imshow(result) + ax.set_title(extractor.__class__.__name__, size=8) + + # Compute the inverse heatmap + if args.save_inverse: + inverse_activation_map = torch.sub( + 1, activation_map[img_idx].unsqueeze(0) + ) + inverse_heatmap = to_pil_image(inverse_activation_map, mode="F") + result_inverse = overlay_mask( + deepcopy(img), inverse_heatmap, alpha=args.alpha + ) + ax = ( + axes[idx][img_idx] + if args.rows > 1 + else axes[idx] if num_cols > 1 else axes + ) + ax.imshow(result_inverse) + ax.set_title(f"{extractor.__class__.__name__}_inverse", size=8) + + # Clear axes + if num_cols > 1: + for _axes in axes: + if num_rows > 1: + for ax in _axes: + ax.axis("off") + else: + _axes.axis("off") + + else: + axes.axis("off") + + plt.tight_layout() + if args.savefig: + plt.savefig( + args.savefig, dpi=200, transparent=True, bbox_inches="tight", pad_inches=0 + ) + + +if __name__ == "__main__": + main() diff --git a/video/fake-stormer/model_code/package_utils/deepfake_mask.py b/video/fake-stormer/model_code/package_utils/deepfake_mask.py new file mode 100644 index 0000000000000000000000000000000000000000..914376b2e72de3b82d01d9ae22ee9fbb00e1e145 --- /dev/null +++ b/video/fake-stormer/model_code/package_utils/deepfake_mask.py @@ -0,0 +1,303 @@ +# -*- coding: utf-8 -*- +"""Masks functions for faceswap.py""" + +import inspect +import logging +import random +import sys + +import cv2 +import numpy as np + +logger = logging.getLogger(__name__) # pylint: disable=invalid-name + + +def get_available_masks(): + """Return a list of the available masks for cli""" + masks = sorted( + [ + name + for name, obj in inspect.getmembers(sys.modules[__name__]) + if inspect.isclass(obj) and name != "Mask" + ] + ) + masks.append("none") + logger.debug(masks) + return masks + + +def get_default_mask(): + """Set the default mask for cli""" + masks = get_available_masks() + default = "dfl_full" + default = default if default in masks else masks[0] + logger.debug(default) + return default + + +class Mask: + """Parent class for masks + the output mask will be .mask + channels: 1, 3 or 4: + 1 - Returns a single channel mask + 3 - Returns a 3 channel mask + 4 - Returns the original image with the mask in the alpha channel""" + + def __init__(self, landmarks, face, channels=4): + # logger.info("Initializing %s: (face_shape: %s, channels: %s, landmarks: %s)", + # self.__class__.__name__, face.shape, channels, landmarks) + self.landmarks = landmarks + self.face = face + self.channels = channels + + mask = self.build_mask() + self.mask = self.merge_mask(mask) + # logger.info("Initialized %s", self.__class__.__name__) + + def build_mask(self): + """Override to build the mask""" + raise NotImplementedError + + def merge_mask(self, mask): + """Return the mask in requested shape""" + # logger.info("mask_shape: %s", mask.shape) + assert self.channels in (1, 3, 4), "Channels should be 1, 3 or 4" + assert ( + mask.shape[2] == 1 and mask.ndim == 3 + ), "Input mask be 3 dimensions with 1 channel" + + if self.channels == 3: + retval = np.tile(mask, 3) + elif self.channels == 4: + retval = np.concatenate((self.face, mask), -1) + else: + retval = mask + + # logger.info("Final mask shape: %s", retval.shape) + return retval + + +class dfl_full(Mask): # pylint: disable=invalid-name + """DFL facial mask""" + + def build_mask(self): + mask = np.zeros(self.face.shape[0:2] + (1,), dtype=np.float32) + + nose_ridge = (self.landmarks[27:31], self.landmarks[33:34]) + jaw = ( + self.landmarks[0:17], + self.landmarks[48:68], + self.landmarks[0:1], + self.landmarks[8:9], + self.landmarks[16:17], + ) + eyes = ( + self.landmarks[17:27], + self.landmarks[0:1], + self.landmarks[27:28], + self.landmarks[16:17], + self.landmarks[33:34], + ) + parts = [jaw, nose_ridge, eyes] + + for item in parts: + merged = np.concatenate(item) + cv2.fillConvexPoly( + mask, cv2.convexHull(merged), 255.0 + ) # pylint: disable=no-member + return mask + + +class components(Mask): # pylint: disable=invalid-name + """Component model mask""" + + def build_mask(self): + mask = np.zeros(self.face.shape[0:2] + (1,), dtype=np.float32) + + r_jaw = (self.landmarks[0:9], self.landmarks[17:18]) + l_jaw = (self.landmarks[8:17], self.landmarks[26:27]) + r_cheek = (self.landmarks[17:20], self.landmarks[8:9]) + l_cheek = (self.landmarks[24:27], self.landmarks[8:9]) + nose_ridge = ( + self.landmarks[19:25], + self.landmarks[8:9], + ) + r_eye = ( + self.landmarks[17:22], + self.landmarks[27:28], + self.landmarks[31:36], + self.landmarks[8:9], + ) + l_eye = ( + self.landmarks[22:27], + self.landmarks[27:28], + self.landmarks[31:36], + self.landmarks[8:9], + ) + nose = (self.landmarks[27:31], self.landmarks[31:36]) + parts = [r_jaw, l_jaw, r_cheek, l_cheek, nose_ridge, r_eye, l_eye, nose] + + for item in parts: + merged = np.concatenate(item) + cv2.fillConvexPoly( + mask, cv2.convexHull(merged), 255.0 + ) # pylint: disable=no-member + return mask + + +class extended(Mask): # pylint: disable=invalid-name + """Extended mask + Based on components mask. Attempts to extend the eyebrow points up the forehead + """ + + def build_mask(self): + mask = np.zeros(self.face.shape[0:2] + (1,), dtype=np.float32) + + landmarks = self.landmarks.copy() + # mid points between the side of face and eye point + ml_pnt = (landmarks[36] + landmarks[0]) // 2 + mr_pnt = (landmarks[16] + landmarks[45]) // 2 + + # mid points between the mid points and eye + ql_pnt = (landmarks[36] + ml_pnt) // 2 + qr_pnt = (landmarks[45] + mr_pnt) // 2 + + # Top of the eye arrays + bot_l = np.array( + (ql_pnt, landmarks[36], landmarks[37], landmarks[38], landmarks[39]) + ) + bot_r = np.array( + (landmarks[42], landmarks[43], landmarks[44], landmarks[45], qr_pnt) + ) + + # Eyebrow arrays + top_l = landmarks[17:22] + top_r = landmarks[22:27] + + # Adjust eyebrow arrays + landmarks[17:22] = top_l + ((top_l - bot_l) // 2) + landmarks[22:27] = top_r + ((top_r - bot_r) // 2) + + r_jaw = (landmarks[0:9], landmarks[17:18]) + l_jaw = (landmarks[8:17], landmarks[26:27]) + r_cheek = (landmarks[17:20], landmarks[8:9]) + l_cheek = (landmarks[24:27], landmarks[8:9]) + nose_ridge = ( + landmarks[19:25], + landmarks[8:9], + ) + r_eye = (landmarks[17:22], landmarks[27:28], landmarks[31:36], landmarks[8:9]) + l_eye = (landmarks[22:27], landmarks[27:28], landmarks[31:36], landmarks[8:9]) + nose = (landmarks[27:31], landmarks[31:36]) + parts = [r_jaw, l_jaw, r_cheek, l_cheek, nose_ridge, r_eye, l_eye, nose] + + for item in parts: + merged = np.concatenate(item) + cv2.fillConvexPoly( + mask, cv2.convexHull(merged), 255.0 + ) # pylint: disable=no-member + return mask + + +class facehull(Mask): # pylint: disable=invalid-name + """Basic face hull mask""" + + def build_mask(self): + mask = np.zeros(self.face.shape[0:2] + (1,), dtype=np.float32) + hull = cv2.convexHull( # pylint: disable=no-member + np.array(self.landmarks).reshape((-1, 2)) + ) + cv2.fillConvexPoly( + mask, hull, 255.0, lineType=cv2.LINE_AA + ) # pylint: disable=no-member + return mask + + +def random_get_hull(landmark, img1, hull_type=None): + if hull_type is None: + hull_type = random.choice([0, 1, 2, 3]) + + if hull_type == 0: + mask = dfl_full(landmarks=landmark.astype("int32"), face=img1, channels=3).mask + return mask / 255, hull_type + elif hull_type == 1: + mask = extended(landmarks=landmark.astype("int32"), face=img1, channels=3).mask + return mask / 255, hull_type + elif hull_type == 2: + mask = components( + landmarks=landmark.astype("int32"), face=img1, channels=3 + ).mask + return mask / 255, hull_type + elif hull_type == 3: + mask = facehull(landmarks=landmark.astype("int32"), face=img1, channels=3).mask + return mask / 255, hull_type + + +def dynamic_blend(source, target, mask, **kwargs): + mask_blured, size_h, size_w, kernel_1, kernel_2, sigma_rand = get_blend_mask( + mask, **kwargs + ) + + if kwargs.get("blend_ratio") is None: + blend_list = [0.25, 0.5, 0.75, 1, 1, 1] + blend_ratio = blend_list[np.random.randint(len(blend_list))] + else: + blend_ratio = kwargs.get("blend_ratio") + + mask_blured_ = mask_blured * blend_ratio + + img_blended = mask_blured_ * source + (1 - mask_blured_) * target + + blend_params = { + "blend_ratio": blend_ratio, + "size_h": size_h, + "size_w": size_w, + "kernel_1": kernel_1, + "kernel_2": kernel_2, + "sigma_rand": sigma_rand, + } + return img_blended, mask_blured, blend_params + + +def get_blend_mask(mask, **kwargs): + H, W = mask.shape + + if kwargs.get("size_h") is None and kwargs.get("size_w") is None: + size_h = np.random.randint(H * 0.8, H / 0.8) + size_w = np.random.randint(W * 0.8, W / 0.8) + else: + size_h = kwargs.get("size_h") + size_w = kwargs.get("size_w") + + mask = cv2.resize(mask, (size_w, size_h)) + + if kwargs.get("kernel_1") is None and kwargs.get("kernel_2") is None: + kernel_1 = random.randrange(5, 26, 2) + kernel_2 = random.randrange(5, 26, 2) + kernel_1 = (kernel_1, kernel_1) + kernel_2 = (kernel_2, kernel_2) + else: + kernel_1 = kwargs.get("kernel_1") + kernel_2 = kwargs.get("kernel_2") + + mask_blured = cv2.GaussianBlur(mask, kernel_1, 0) + mask_blured = mask_blured / (mask_blured.max()) + mask_blured[mask_blured < 1] = 0 + + if kwargs.get("sigma_rand") is None: + sigma_rand = np.random.randint(5, 46) + else: + sigma_rand = kwargs.get("sigma_rand") + mask_blured = cv2.GaussianBlur(mask_blured, kernel_2, sigma_rand) + mask_blured = mask_blured / (mask_blured.max()) + + mask_blured = cv2.resize(mask_blured, (W, H)) + + return ( + mask_blured.reshape((mask_blured.shape + (1,))), + size_h, + size_w, + kernel_1, + kernel_2, + sigma_rand, + ) diff --git a/video/fake-stormer/model_code/package_utils/geo_landmarks_extraction.py b/video/fake-stormer/model_code/package_utils/geo_landmarks_extraction.py new file mode 100644 index 0000000000000000000000000000000000000000..50ea211ad696c36fb969c0c36f3bf6022c5dd556 --- /dev/null +++ b/video/fake-stormer/model_code/package_utils/geo_landmarks_extraction.py @@ -0,0 +1,348 @@ +# -*- coding: utf-8 -*- +import argparse +import math +import os +import sys +import time + +if not os.getcwd() in sys.path: + sys.path.append(os.getcwd()) +import random +from glob import glob + +import cv2 +import dlib +import numpy as np +import simplejson as json +from box import Box as edict +from configs.get_config import load_config +from imutils import face_utils +from tqdm import tqdm +from transform import affine_transform + +from utils import draw_landmarks + + +class LandmarkUtility(object): + def __init__(self, cfg, load_imgs=False, **kwargs): + super().__init__() + + assert "DATASET" in cfg, "Dataset can not be None!" + assert "ROOT" in cfg, "Image Directory need to be provided!" + + if not isinstance(cfg, edict): + cfg = edict(cfg) + + self.load_imgs = load_imgs + self.image_root = cfg.ROOT + self.image_suffix = cfg.IMAGE_SUFFIX or "jpg" + self.dataset = cfg.DATASET + self.split = cfg.SPLIT or "train" + self.data_type = cfg.DATA_TYPE or "images" + self.fol_label = cfg.LABEL or ["real"] + self.debug = cfg.DEBUG + self.fake_types = cfg.FAKETYPE + self.compression = cfg.COMPRESSION + + if kwargs is not None: + for k, v in kwargs.items(): + if v is None: + raise ValueError(f"{k}:{v} recieve a None value!") + self.__setattr__(k, v) + + def __contain__(self, key): + return hasattr(self, key) + + def _load_data(self): + img_paths = [] + file_names = [] + + print(f"Loading data from dataset --- {self.dataset}") + if self.load_imgs: + img_paths, file_names = self._load_data_from_path() + else: + assert self.__contain__( + "file_path" + ), "Loading data from file need a file path" + img_paths, file_names = self._load_data_from_file( + self.__getattribute__("file_path") + ) + + assert ( + len(img_paths) != 0 + ), "Image paths have not been loaded! Please check image directory!" + assert ( + len(file_names) != 0 + ), "Image files have not been loaded! Please check image suffixes!" + return img_paths, file_names + + def _load_data_from_path(self): + """ + Currenly, Using Glob for loading file with regex + It might be changed for better performance in large datasets + """ + assert os.path.exists(self.image_root), "Root path to dataset can not be None!" + data_type = self.data_type + fake_types = self.fake_types + img_paths = [] + + # Load image data for each type of fake techniques + for idx, ft in enumerate(fake_types): + data_dir = os.path.join(self.image_root, self.split, data_type, ft) + if not os.path.exists(data_dir): + raise ValueError("Data Directory can not be invalid!") + + for sub_dir in os.listdir(data_dir): + sub_dir_path = os.path.join(data_dir, sub_dir) + img_paths_ = glob(f"{sub_dir_path}/*.{self.image_suffix}") + + img_paths.extend(img_paths_) + + print( + "{} image paths have been loaded from {}!".format( + len(img_paths), self.dataset + ) + ) + file_names = [ip.split("/")[-1] for ip in img_paths] + + return img_paths, file_names + + def _load_data_from_file(self, file_path): + """ + Each extension will be treated with particular extension loader + """ + filename, file_extension = os.path.splitext(file_path) + img_paths, file_names = [], [] + if file_extension == ".json": + f = open(file_path) + data = json.load(f) + obj_data = data["data"] + + for item in obj_data: + img_paths.append(item["image_path"]) + file_names.append(item["file_name"]) + return img_paths, file_names + + def _img_obj(self, img_path, file_name, **kwargs): + img_path = img_path.replace(self.image_root, "") + obj = dict(image_path=img_path, file_name=file_name, **kwargs) + return obj + + def _load_image(self, img_path): + image = cv2.imread(os.path.join(self.image_root, img_path)) + return image + + def _facial_landmark(self, image, detector, lm_predictor): + try: + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + except: + gray = image + + f_rect = detector(gray, 1) + if len(f_rect) > 0: + f_lms = lm_predictor(gray, f_rect[0]) + f_lms = face_utils.shape_to_np(f_lms) + return f_lms + else: + return None + + def _align_face(self, image, f_lms): + assert f_lms is not None, "Facial Landmarks can not be None!" + eyepoints = f_lms[39], f_lms[42] + le_x, le_y = eyepoints[0] + re_x, re_y = eyepoints[1] + + angle = math.atan((le_y - re_y) / (le_x - re_x)) * (180 / math.pi) + origin_point = tuple(np.array(image.shape[1::-1]) / 2) + + rot_mat = cv2.getRotationMatrix2D(origin_point, angle, 1.0) + rot_img = cv2.warpAffine( + image, rot_mat, image.shape[1::-1], flags=cv2.INTER_LINEAR + ) + + # Aligning face landmarks by the rotation matrix + rot_f_lms = None + if f_lms is not None: + rot_f_lms = np.empty_like(f_lms) + for i, p in enumerate(f_lms): + rot_f_lms[i] = affine_transform(p, rot_mat) + + return rot_img, f_lms, rot_f_lms + + def facial_landmarks(self, img_paths, detector, lm_predictor): + rot_imgs, f_lmses, rot_f_lmses = [], [], [] + + for i, ip in enumerate(tqdm(img_paths, dynamic_ncols=True)): + image = self._load_image(ip) + + # Checking time processing for each item + s_t = time.time() + f_lms = None + try: + f_lms = self._facial_landmark(image, detector, lm_predictor) + if f_lms is None: + if self.debug: + cv2.imwrite(f"samples/exception_img_{i}.jpg", image) + print(f"Image {i}--{ip} did not find any landmarks!") + except Exception as e: + print(e) + + if i == 1: + print( + "Landmark detection processing time ---- {}".format( + time.time() - s_t + ) + ) + + if f_lms is not None: + rot_img, _f_lms, rot_f_lms = self._align_face(image, f_lms) + else: + rot_img, _f_lms, rot_f_lms = image, [], [] + rot_imgs.append(rot_img) + f_lmses.append(_f_lms) + rot_f_lmses.append(rot_f_lms) + + # Visualizing landmarks to test + if i < 10 and self.debug: + rot_img = draw_landmarks(rot_img, rot_f_lms) + cv2.imwrite(f"samples/test_{i}.jpg", rot_img) + + if i % 100 == 0: + print(f"Landmarks have been detected for {i} images") + return rot_imgs, f_lmses, rot_f_lmses + + def build_data(self, img_paths, file_names, **kwargs): + data = dict(data=[]) + + if "orig_lmses" in kwargs.keys(): + if not bool(kwargs["orig_lmses"]): + raise ValueError("Original Landmarks cannot be None!") + else: + orig_lmses = kwargs["orig_lmses"] + assert len(orig_lmses) == len( + img_paths + ), "The length of images and landmarks is not compatible!" + + if "aligned_lmses" in kwargs.keys(): + if not bool(kwargs["aligned_lmses"]): + raise ValueError("Aligned Landmarks cannot be None!") + else: + aligned_lmses = kwargs["aligned_lmses"] + assert len(aligned_lmses) == len( + img_paths + ), "The length of images and aligned landmarks is not compatible!" + + for i, (p, f) in enumerate(zip(img_paths, file_names)): + fake_type = ( + p.split("/")[-2] + if self.fake_types != ["original"] + else self.fake_types[0] + ) + img_obj = self._img_obj(p, f, id=i, fake_type=fake_type) + + if "orig_lmses" in kwargs.keys(): + img_obj["orig_lms"] = ( + orig_lmses[i].tolist() + if isinstance(orig_lmses[i], np.ndarray) + else orig_lmses[i] + ) # To save to JSON + if "aligned_lmses" in kwargs.keys(): + img_obj["aligned_lms"] = ( + aligned_lmses[i].tolist() + if isinstance(aligned_lmses[i], np.ndarray) + else aligned_lmses[i] + ) # To save to JSON + data["data"].append(img_obj) + return data + + def save2json(self, data, fn="faceforensics_processed.json"): + assert len(data), "Data can not be empty!" + target = "processed_data/{}".format(self.compression) + if not os.path.exists(target): + os.mkdir(target) + fp = os.path.join(target, fn) + with open(fp, "w") as f: + json.dump(data, f) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Landmarks preprocessing!") + parser.add_argument("--config", help="Config file to proceed preprocessing") + parser.add_argument("--file_path", help="File to load processed data") + parser.add_argument( + "--extract_landmark", help="Use Dlib to extract landmarks", action="store_true" + ) + parser.add_argument( + "--save_aligned", help="Save aligned images", action="store_true" + ) + args = parser.parse_args() + print(args) + + cfg = load_config(args.config) + extract_landmark = args.extract_landmark + save_aligned = args.save_aligned + + kwargs = {} + if extract_landmark: + kwargs["extract_landmark"] = extract_landmark + + # Initialize Landmark Utility instance + if args.file_path: + lm_ins = LandmarkUtility( + cfg.PREPROCESSING, load_imgs=False, file_path=args.file_path, **kwargs + ) + else: + lm_ins = LandmarkUtility(cfg.PREPROCESSING, load_imgs=True, **kwargs) + img_paths, file_names = lm_ins._load_data() + print(f"{len(img_paths)} images have been loaded for processing!") + + if extract_landmark: + assert ( + cfg.PREPROCESSING.facial_lm_pretrained is not None + ), "Landmark pretrained can not be None!" + f_detector = dlib.get_frontal_face_detector() + f_lm_detector = dlib.shape_predictor(cfg.PREPROCESSING.facial_lm_pretrained) + rot_imgs, f_lmses, rot_f_lmses = lm_ins.facial_landmarks( + img_paths, f_detector, f_lm_detector + ) + + if save_aligned: + os.makedirs( + f"{lm_ins.image_root}{lm_ins.split}/{lm_ins.data_type}/aligned_{lm_ins.fake_types[0]}_{cfg.PREPROCESSING.N_LANDMARKS}", + exist_ok=True, + ) + for i, img_p in enumerate(tqdm(img_paths, dynamic_ncols=True)): + rot_img = rot_imgs[i] + fn = file_names[i] + vid_id = img_p.split("/")[-2] + os.makedirs( + f"{lm_ins.image_root}{lm_ins.compression}/{lm_ins.split}/{lm_ins.data_type}/aligned_{lm_ins.fake_types[0]}_{cfg.PREPROCESSING.N_LANDMARKS}/{vid_id}", + exist_ok=True, + ) + + aligned_img_p = img_p.replace( + lm_ins.fake_types[0], + f"aligned_{lm_ins.fake_types[0]}_{cfg.PREPROCESSING.N_LANDMARKS}", + ) + cv2.imwrite(os.path.join(lm_ins.image_root, aligned_img_p), rot_img) + img_paths[i] = aligned_img_p + + print("All landmarks have been detected and stored in memory!") + print("Ready to save to file...") + + if args.file_path is None: + if extract_landmark: + data = lm_ins.build_data( + img_paths, file_names, orig_lmses=f_lmses, aligned_lmses=rot_f_lmses + ) + else: + data = lm_ins.build_data(img_paths, file_names) + + try: + lm_ins.save2json( + data, + fn=f"{lm_ins.split}_{lm_ins.dataset}_{lm_ins.data_type}_{cfg.PREPROCESSING.N_LANDMARKS}.json", + ) + except Exception as e: + print(e) + print("Processed Data has been saved successfully!") diff --git a/video/fake-stormer/model_code/package_utils/image_augmentation.py b/video/fake-stormer/model_code/package_utils/image_augmentation.py new file mode 100644 index 0000000000000000000000000000000000000000..1354db8602ce5a207d7417ca8f61eacb0229c4e1 --- /dev/null +++ b/video/fake-stormer/model_code/package_utils/image_augmentation.py @@ -0,0 +1,96 @@ +# -*- coding: utf-8 -*- +import argparse +import glob +import os +import random + +import numpy as np +import torch +from image_utils import ( + block_wise, + color_contrast, + color_saturation, + gaussian_blur, + gaussian_noise_color, + jpeg_compression, + load_image, + video_compression, +) +from PIL import Image + +# DIST_LEVEL = 3 + + +def get_distortion_parameter(type, level): + param_dict = dict() # a dict of list + param_dict["CS"] = [0.4, 0.3, 0.2, 0.1, 0.0] # smaller, worse + param_dict["CC"] = [0.85, 0.725, 0.6, 0.475, 0.35] # smaller, worse + param_dict["BW"] = [16, 32, 48, 64, 80] # larger, worse + param_dict["GNC"] = [0.001, 0.002, 0.005, 0.01, 0.05] # larger, worse + param_dict["GB"] = [7, 9, 13, 17, 21] # larger, worse + param_dict["JPEG"] = [2, 3, 4, 5, 6] # larger, worse + param_dict["VC"] = [30, 32, 35, 38, 40] # larger, worse + + # level starts from 1, list starts from 0 + return param_dict[type][level - 1] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("-i", dest="path", type=str, default="") + parser.add_argument( + "-t", + dest="task", + choices=[ + "noise", + "block", + "saturation", + "contrast", + "blur", + "pixel", + "compression", + ], + default="noise", + ) + args = parser.parse_args() + # Setting device + device = torch.device("cuda") + dest = args.path + "_" + args.task + "_" + "random" + "/" + + if not os.path.exists(dest): + os.makedirs(dest) + + for dirpath, dirnames, filenames in os.walk(args.path): + possible_files = os.path.join(dirpath, "*.png") + + for file in glob.glob(possible_files): + img = load_image(file) + dist_level = random.randint(1, 5) + + if args.task == "noise": + params = get_distortion_parameter("GNC", dist_level) + img = gaussian_noise_color(img, params) + elif args.task == "block": + params = get_distortion_parameter("BW", dist_level) + img = block_wise(img, params) + elif args.task == "saturation": + params = get_distortion_parameter("CS", dist_level) + img = color_saturation(img, params) + elif args.task == "contrast": + params = get_distortion_parameter("CC", dist_level) + img = color_contrast(img, params) + elif args.task == "blur": + params = get_distortion_parameter("GB", dist_level) + img = gaussian_blur(img, params) + elif args.task == "pixel": + params = get_distortion_parameter("JPEG", dist_level) + img = jpeg_compression(img, params) + elif args.task == "compression": + params = get_distortion_parameter("VC", dist_level) + img = video_compression(img, params) + res = dest + file.split("/")[-2] + + if not os.path.exists(res): + os.makedirs(res) + # print(dest+('/').join(file.split('/')[-2:])) + Image.fromarray(img).save(res + "/" + file.split("/")[-1]) diff --git a/video/fake-stormer/model_code/package_utils/image_utils.py b/video/fake-stormer/model_code/package_utils/image_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..126d944f958e285e6c50cc84f98b2ef25a51c85e --- /dev/null +++ b/video/fake-stormer/model_code/package_utils/image_utils.py @@ -0,0 +1,214 @@ +# -*- coding: utf-8 -*- +import math +import os +import random + +import cv2 +import numpy as np +from matplotlib import cm +from PIL import Image + + +def load_image(image_path): + """Loading image""" + img = Image.open(image_path) + # Fix bug RGBA + if img.mode != "RGB": + img = img.convert("RGB") + img = np.array(img) + return img + + +def crop_by_margin(image, margin=[0, 0]): + """Cropping images by margins as a step of preprocessing""" + H, W = image.shape[:2] + margin_x, margin_y = margin + image = image[margin_y : H - margin_y, margin_x : W - margin_x, :] + return image + + +def gaussian_radius(det_size, min_overlap=0.7): + """Calculating gaussian radius to compute std for Unnormalized Gaussian Mask""" + height, width = det_size + + a1 = 1 + b1 = height + width + c1 = width * height * (1 - min_overlap) / (1 + min_overlap) + sq1 = np.sqrt(b1**2 - 4 * a1 * c1) + r1 = (b1 + sq1) / 2 + + a2 = 4 + b2 = 2 * (height + width) + c2 = (1 - min_overlap) * width * height + sq2 = np.sqrt(b2**2 - 4 * a2 * c2) + r2 = (b2 + sq2) / 2 + + a3 = 4 * min_overlap + b3 = -2 * min_overlap * (height + width) + c3 = (min_overlap - 1) * width * height + sq3 = np.sqrt(b3**2 - 4 * a3 * c3) + r3 = (b3 + sq3) / 2 + return min(r1, r2, r3) + + +def cal_mask_wh(p, mask): + """Adaptively calculating blending mask W, H at the most vulnerable points perspective""" + cy, cx = p + mask_h, mask_w = mask.shape + w = 0 + h = 0 + + for i in [-1, 1]: + shift_y = 0 + while ( + (cy + shift_y > -mask_h) + and (cy + shift_y < mask_h) + and (mask[cy + shift_y, cx] > 128) + ): + w += 1 + shift_y += i + + shift_x = 0 + while ( + (cx + shift_x > -mask_w) + and (cx + shift_x < mask_w) + and (mask[cy, cx + shift_x] > 128) + ): + h += 1 + shift_x += i + + return w, h + + +def overlay_mask( + img: Image.Image, mask: Image.Image, colormap: str = "jet", alpha: float = 0.7 +) -> Image.Image: + """Overlay a colormapped mask on a background image + + >>> from PIL import Image + >>> import matplotlib.pyplot as plt + >>> from torchcam.utils import overlay_mask + >>> img = ... + >>> cam = ... + >>> overlay = overlay_mask(img, cam) + + Args: + img: background image + mask: mask to be overlayed in grayscale + colormap: colormap to be applied on the mask + alpha: transparency of the background image + + Returns: + overlayed image + + Raises: + TypeError: when the arguments have invalid types + ValueError: when the alpha argument has an incorrect value + """ + + if not isinstance(img, Image.Image) or not isinstance(mask, Image.Image): + raise TypeError("img and mask arguments need to be PIL.Image") + + if not isinstance(alpha, float) or alpha < 0 or alpha >= 1: + raise ValueError( + "alpha argument is expected to be of type float between 0 and 1" + ) + + cmap = cm.get_cmap(colormap) + # Resize mask and apply colormap + overlay = mask.resize(img.size, resample=Image.BICUBIC) + overlay = (255 * cmap(np.asarray(overlay) ** 1)[:, :, :3]).astype(np.uint8) + # Overlay the image with the mask + overlayed_img = Image.fromarray( + (alpha * np.asarray(img) + (1 - alpha) * overlay).astype(np.uint8) + ) + + return overlayed_img + + +def bgr2ycbcr(img_bgr): + img_bgr = img_bgr.astype(np.float32) + img_ycrcb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2YCR_CB) + img_ycbcr = img_ycrcb[:, :, (0, 2, 1)].astype(np.float32) + # to [16/255, 235/255] + img_ycbcr[:, :, 0] = (img_ycbcr[:, :, 0] * (235 - 16) + 16) / 255.0 + # to [16/255, 240/255] + img_ycbcr[:, :, 1:] = (img_ycbcr[:, :, 1:] * (240 - 16) + 16) / 255.0 + + return img_ycbcr + + +def ycbcr2bgr(img_ycbcr): + img_ycbcr = img_ycbcr.astype(np.float32) + # to [0, 1] + img_ycbcr[:, :, 0] = (img_ycbcr[:, :, 0] * 255.0 - 16) / (235 - 16) + # to [0, 1] + img_ycbcr[:, :, 1:] = (img_ycbcr[:, :, 1:] * 255.0 - 16) / (240 - 16) + img_ycrcb = img_ycbcr[:, :, (0, 2, 1)].astype(np.float32) + img_bgr = cv2.cvtColor(img_ycrcb, cv2.COLOR_YCR_CB2BGR) + + return img_bgr + + +def gaussian_noise_color(img, param=None): + if param is None: + param = [0.001, 0.002, 0.005, 0.01, 0.05] + ycbcr = bgr2ycbcr(img) / 255 + size_a = ycbcr.shape + b = ( + ycbcr + math.sqrt(param) * np.random.randn(size_a[0], size_a[1], size_a[2]) + ) * 255 + b = ycbcr2bgr(b) + img = np.clip(b, 0, 255).astype(np.uint8) + return img + + +def block_wise(img, param): + width = 8 + block = np.ones((width, width, 3)).astype(int) * 128 + param = min(img.shape[0], img.shape[1]) // 256 * param + for i in range(param): + r_w = random.randint(0, img.shape[1] - 1 - width) + r_h = random.randint(0, img.shape[0] - 1 - width) + img[r_h : r_h + width, r_w : r_w + width, :] = block + + return img + + +def color_saturation(img, param): + ycbcr = bgr2ycbcr(img) + ycbcr[:, :, 1] = 0.5 + (ycbcr[:, :, 1] - 0.5) * param + ycbcr[:, :, 2] = 0.5 + (ycbcr[:, :, 2] - 0.5) * param + img = ycbcr2bgr(ycbcr).astype(np.uint8) + + return img + + +def color_contrast(img, param): + img = img.astype(np.float32) * param + img = img.astype(np.uint8) + + return img + + +def gaussian_blur(img, param): + img = cv2.GaussianBlur(img, (param, param), param * 1.0 / 6) + + return img + + +def jpeg_compression(img, param): + h, w, _ = img.shape + s_h = h // param + s_w = w // param + img = cv2.resize(img, (s_w, s_h)) + img = cv2.resize(img, (w, h)) + + return img + + +def video_compression(vid_in, vid_out, param): + cmd = f"ffmpeg -i {vid_in} -crf {param} -y {vid_out}" + os.system(cmd) + + return diff --git a/video/fake-stormer/model_code/package_utils/images_crop.py b/video/fake-stormer/model_code/package_utils/images_crop.py new file mode 100644 index 0000000000000000000000000000000000000000..e01e095d2d93e71972fee0231e8b9a447020e918 --- /dev/null +++ b/video/fake-stormer/model_code/package_utils/images_crop.py @@ -0,0 +1,444 @@ +# -*- coding: utf-8 -*- +import argparse +import csv +import json +import os +import shutil +from glob import glob + +import cv2 +import numpy as np +import pandas as pd +import torch +from imutils import face_utils +from retinaface.pre_trained_models import get_model +from retinaface.utils import vis_annotations +from sklearn.utils import shuffle +from tqdm import tqdm + +ROOT = "/data/deepfake_cluster/datasets_df" +SAVE_DIR = f"{ROOT}/FGN" +IMAGE_H, IMAGE_W, IMAGE_C = 256, 256, 3 +PADDING = 0.25 +DATA_TYPE = "frames" # frames or videos + + +def facecrop( + model, + org_path, + save_path, + period=1, + num_frames=10, + dataset="original", + label=None, + mask_path=None, + padding=PADDING, +): + print(f"Processing video --- {org_path}") + cap_org = cv2.VideoCapture(org_path) + if mask_path is not None: + mask_cap = cv2.VideoCapture(mask_path) + croppedfaces = [] + frame_count_org = int(cap_org.get(cv2.CAP_PROP_FRAME_COUNT)) + print("N frame count --- ", frame_count_org) + + if label is not None: + # Custom org_path for ForgeryNet + if dataset == "ForgeryNet": + org_path_ = org_path.split("/") + org_path = "/".join(org_path_[:7] + ["_".join(org_path_[8:])]) + save_path_ = ( + save_path + + f"/{DATA_TYPE}/{str(label)}/" + + os.path.basename(org_path).replace(".mp4", "/") + ) + else: + save_path_ = ( + save_path + + f"/{DATA_TYPE}/{dataset}/" + + os.path.basename(org_path).replace(".mp4", "/") + ) + os.makedirs(save_path_, exist_ok=True) + + if mask_path is not None: + save_mask_path_ = ( + save_path + + f"/masks/{dataset}/" + + os.path.basename(mask_path).replace(".mp4", "/") + ) + os.makedirs(save_mask_path_, exist_ok=True) + + if DATA_TYPE == "frames": + frame_idxs = np.linspace( + 0, frame_count_org - 1, num_frames, endpoint=True, dtype=np.int64 + ) + else: + frame_idxs = range(0, num_frames) + + for cnt_frame in range(frame_count_org): + image_path = save_path_ + str(cnt_frame).zfill(3) + ".png" + if os.path.isfile(image_path): + continue + if mask_path is not None: + mask_f_path = save_mask_path_ + str(cnt_frame).zfill(3) + ".png" + if os.path.isfile(mask_f_path): + continue + + try: + ret_org, frame_org = cap_org.read() + if mask_path is not None: + ret_m_org, mask_org = mask_cap.read() + height, width = frame_org.shape[:-1] + if not ret_org: + tqdm.write( + "Frame read {} Error! : {}".format( + cnt_frame, os.path.basename(org_path) + ) + ) + continue + + if cnt_frame not in frame_idxs: + continue + + frame = cv2.cvtColor(frame_org, cv2.COLOR_BGR2RGB) + faces = model.predict_jsons(frame) + try: + if len(faces) == 0: + print(faces) + tqdm.write( + "No faces in {}:{}".format( + cnt_frame, os.path.basename(org_path) + ) + ) + continue + + face_s_max = -1 + landmarks = [] + face_crop = None + score_max = -1 + for face_idx in range(len(faces)): + x0, y0, x1, y1 = faces[face_idx]["bbox"] + # landmark = np.array([[x0,y0],[x1,y1]] + faces[face_idx]['landmarks']) + face_w = x1 - x0 + face_h = y1 - y0 + face_s = face_w * face_h + score = faces[face_idx]["score"] + + if face_s > face_s_max and score > score_max: + f_c_x0 = max(0, x0 - int(face_w * padding)) + f_c_x1 = min(width, x1 + int(face_w * padding)) + f_c_y0 = max(0, y0 - int(face_h * padding)) + f_c_y1 = min(height, y1 + int(face_h * padding)) + + face_crop = frame_org[f_c_y0:f_c_y1, f_c_x0:f_c_x1, :] + if mask_path is not None: + mask_crop = mask_org[f_c_y0:f_c_y1, f_c_x0:f_c_x1, :] + face_s_max = face_s + score_max = score + # size_list.append(face_s) + # # landmarks.append(landmark) + except Exception as e: + print(f"error in {cnt_frame}:{org_path}") + print(e) + continue + except Exception as e1: + print(e1) + continue + + # landmarks=np.concatenate(landmarks).reshape((len(size_list),) + landmark.shape) + # landmarks=landmarks[np.argsort(np.array(size_list))[::-1]] + + # land_path=save_path_+str(cnt_frame).zfill(3) + # land_path=land_path.replace('/frames','/retina') + # os.makedirs(os.path.dirname(land_path),exist_ok=True) + # np.save(land_path, landmarks) + # if not os.path.isfile(image_path): + face_crop = cv2.resize( + face_crop, (IMAGE_H, IMAGE_W), interpolation=cv2.INTER_LINEAR + ) + cv2.imwrite(image_path, face_crop) + + if mask_path is not None: + mask_crop = cv2.resize( + mask_crop, (IMAGE_H, IMAGE_W), interpolation=cv2.INTER_LINEAR + ) + cv2.imwrite(mask_f_path, mask_crop) + + cap_org.release() + return + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "-d", + dest="dataset", + choices=[ + "FaceShifter", + "Face2Face", + "Deepfakes", + "FaceSwap", + "NeuralTextures", + "Original", + "Celeb-real", + "Celeb-synthesis", + "YouTube-real", + "DFDC", + "DFDCP", + "method_A", + "method_B", + "original_videos", + "DFo_source_videos", + "DFo_manipulated_videos", + "ForgeryNet", + ], + ) + parser.add_argument("-c", dest="comp", choices=["raw", "c23", "c40"], default="raw") + parser.add_argument("-n", dest="num_frames", type=int, default=32) + parser.add_argument( + "-t", dest="task", choices=["train", "val", "test"], default="train" + ) + parser.add_argument("--save_mask", "-sm", action="store_true") + parser.add_argument( + "--alloc_mem", "-a", help="Pre allocating GPU memory", action="store_true" + ) + args = parser.parse_args() + + # Allocate memory + if args.alloc_mem: + mem_all_tensors = torch.rand(60, 10000, 10000) + mem_all_tensors.to("cuda:0") + + # Setting device + device = torch.device("cuda") + + # Setting the dataset path based on the dataset name + if args.dataset == "Original": + dataset_path = ( + "{}/FaceForensics++/original_download/original_sequences/youtube/".format( + ROOT + ) + ) + elif args.dataset == "DeepFakeDetection_original": + dataset_path = "/data/FaceForensics++/original_sequences/actors/{}/".format( + args.comp + ) + elif args.dataset in [ + "DeepFakeDetection", + "FaceShifter", + "Face2Face", + "Deepfakes", + "FaceSwap", + "NeuralTextures", + ]: + dataset_path = ( + "{}/FaceForensics++/original_download/manipulated_sequences/{}/".format( + ROOT, args.dataset + ) + ) + elif args.dataset in ["Celeb-real", "Celeb-synthesis", "YouTube-real"]: + if "v1" in SAVE_DIR: + dataset_path = "{}/Celeb-DFv1/".format(ROOT) + else: + dataset_path = "{}/Celeb-DFv2/Celeb-DF-v2/".format(ROOT) + elif args.dataset in ["method_A", "method_B", "original_videos"]: + dataset_path = "{}/DFDCP/".format(ROOT) + elif args.dataset in ["DFDC"]: + dataset_path = "{}/DFDC/".format(ROOT) + elif args.dataset in ["DFo_source_videos", "DFo_manipulated_videos"]: + dataset_path = "{}/DeeperForensics/".format(ROOT) + elif args.dataset in ["ForgeryNet"]: + dataset_path = "{}/FGN/".format(ROOT) + else: + raise NotImplementedError + + # Loading model + model = get_model("resnet50_2020-07-20", max_size=2048, device=device) + model.eval() + + labels = [] + if args.dataset in [ + "Original", + "DeepFakeDetection", + "FaceShifter", + "Face2Face", + "Deepfakes", + "FaceSwap", + "NeuralTextures", + ]: + movies_path = os.path.join(dataset_path, args.comp, "videos/") + mask_mov_paths = os.path.join(dataset_path, "masks", "videos/") + + # Annotation file for FF++ + with open(f"{ROOT}/FaceForensics++/original_download/{args.task}.json") as f: + vid_ids = json.load(f) + elif args.dataset in ["Celeb-real", "Celeb-synthesis", "YouTube-real"]: + if "v1" in SAVE_DIR: + movies_path = dataset_path + else: + movies_path = os.path.join(dataset_path, args.dataset, "videos") + + # Annotation file for Celeb-DF + with open(f"{dataset_path}List_of_{args.task}ing_videos.txt") as f: + vid_ids = pd.read_csv(f).values.reshape(-1) + elif args.dataset in ["DFDC"]: + movies_path = os.path.join(dataset_path, args.task, "download_videos") + + # Annotation file for DFDC + with open(os.path.join(dataset_path, args.task, "labels.csv")) as f: + df = pd.read_csv(f) + # df['path'] = df['label'].astype(str) + '/' + df['filename'] + vid_ids = df["filename"].values.reshape(-1) + labels = df["label"].values.reshape(-1) + elif args.dataset in ["DFo_source_videos", "DFo_manipulated_videos"]: + movies_path = dataset_path + + with open(f"{dataset_path}lists/splits/{args.task}.txt") as f: + vid_ids = pd.read_csv(f).values.reshape(-1) + + if args.dataset == "DFo_source_videos": + vid_ids = list( + set([vid_id.split("_")[1].split(".")[0] for vid_id in vid_ids]) + ) + elif args.dataset in ["ForgeryNet"]: + movies_path = os.path.join(dataset_path, args.task, "video") + + # Annotation file for FGN + vid_ids = [] + with open( + os.path.join(dataset_path, args.task, "video_list.txt"), encoding="utf-8" + ) as f: + df = csv.reader(f, delimiter="\n") + + for idx, row in enumerate(df): + row_data = row[0].split(" ") + if int(row_data[-1]) in [0, 7]: + vid_ids.append(row_data[1]) + label = ( + f"fake_{args.task}" + if int(row_data[-1]) == 7 + else f"real_{args.task}" + ) + labels.append(label) + else: + movies_path = dataset_path + + # Annotation file for DFDCP + with open(f"{ROOT}/DFDCP/dataset.json") as f: + movie_data = json.load(f) + vid_ids = [] + for mv_id, item_data in movie_data.items(): + if item_data["set"] == args.task: + vid_ids.append(mv_id) + + movies_path_list = [] + mask_mov_path_list = [] + file_list = [] + file_path = None + vid_id_count = {} + + # Loading the list of specific video's names for an invidual task 'train/val/test + for i in range(len(vid_ids)): + if args.dataset == "Original": + file_list += vid_ids[i] + elif args.dataset in [ + "Face2Face", + "Deepfakes", + "FaceSwap", + "NeuralTextures", + "FaceShifter", + ]: + file_list.append("_".join([vid_ids[i][0], vid_ids[i][1]])) + file_list.append("_".join([vid_ids[i][1], vid_ids[i][0]])) + elif args.dataset in ["method_A", "method_B", "original_videos"]: + if args.dataset in vid_ids[i]: + file_list.append(vid_ids[i]) + elif args.dataset in ["DFDC", "ForgeryNet"]: + file_list.append(vid_ids[i]) + elif args.dataset in ["DFo_source_videos", "DFo_manipulated_videos"]: + sub_dataset = args.dataset.replace("DFo_", "") + file_list_path = os.path.join(dataset_path, "lists", f"{sub_dataset}_lists") + + if i == 0: + if sub_dataset == "source_videos": + file_path = f"{file_list_path}/{sub_dataset}_list.txt" + else: + file_path = f"{file_list_path}/{sub_dataset}_end_to_end_list.txt" + with open(file_path) as f: + full_file_list = pd.read_csv(f).values.reshape(-1) + full_file_list = shuffle(full_file_list, random_state=259) + + for item in full_file_list: + if vid_ids[i] in vid_id_count.keys() and vid_id_count[vid_ids[i]] > 10: + break + if vid_ids[i] in item: + file_list.append(item) + + if vid_ids[i] in vid_id_count.keys(): + vid_id_count[vid_ids[i]] += 1 + else: + vid_id_count[vid_ids[i]] = 1 + else: + if args.dataset in vid_ids[i]: + file_list.append(vid_ids[i].split(" ")[-1]) + + # movies_path_list = sorted(glob(movies_path+'*.mp4')) + if args.dataset in [ + "Original", + "DeepFakeDetection", + "FaceShifter", + "Face2Face", + "Deepfakes", + "FaceSwap", + "NeuralTextures", + ]: + [movies_path_list.append(movies_path + i + ".mp4") for i in file_list] + if args.save_mask: + [mask_mov_path_list.append(mask_mov_paths + i + ".mp4") for i in file_list] + else: + if "v2" in SAVE_DIR: + [ + movies_path_list.append(os.path.join(movies_path, i.split("/")[-1])) + for i in file_list + ] + else: + [movies_path_list.append(os.path.join(movies_path, i)) for i in file_list] + + print("{} : videos are exist in {}".format(len(movies_path_list), args.dataset)) + n_sample = len(movies_path_list) + print(f"number of video samples -- {n_sample}") + + # Defining the path to store the images + save_path = os.path.join(SAVE_DIR, args.task) + os.makedirs(save_path, exist_ok=True) + + for i in tqdm(range(0, n_sample)): + # folder_path=movies_path_list[i].replace('videos/','frames/').replace('.mp4','/') + # if len(glob(folder_path.replace('/frames/','/retina/')+'*.npy')) < args.num_frames: + if len(labels): + facecrop( + model, + movies_path_list[i], + save_path=save_path, + num_frames=args.num_frames, + dataset=args.dataset, + label=labels[i], + ) + else: + if not args.save_mask: + facecrop( + model, + movies_path_list[i], + save_path=save_path, + num_frames=args.num_frames, + dataset=args.dataset, + ) + else: + facecrop( + model, + movies_path_list[i], + save_path=save_path, + num_frames=args.num_frames, + dataset=args.dataset, + mask_path=mask_mov_path_list[i], + ) diff --git a/video/fake-stormer/model_code/package_utils/metrics_based_preds.py b/video/fake-stormer/model_code/package_utils/metrics_based_preds.py new file mode 100644 index 0000000000000000000000000000000000000000..a821c48a0d7ad367027efb4680c1c206ead6662c --- /dev/null +++ b/video/fake-stormer/model_code/package_utils/metrics_based_preds.py @@ -0,0 +1,149 @@ +# -*- coding: utf-8 -*- +import os +import sys + +if not os.getcwd() in sys.path: + sys.path.insert(0, os.getcwd()) +import argparse + +import numpy as np +import torch +from lib.metrics import ( + apply_cdf_transform, + bin_calculate_auc_ap_ar, + get_acc_mesure_func, +) +from scipy.stats import wasserstein_distance + +from utils import load_file + + +def sigmoid(z): + return 1 / (1 + np.exp(-z)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Reporting metrics based on saved predictions!" + ) + parser.add_argument("--saved_preds", "-f", help="Path to saved prediction file") + parser.add_argument( + "--model_base", + "-m", + help="The input model is image-level or video-level", + default="image", + ) + parser.add_argument( + "--metric_level", + "-ml", + help="Report metrics at image-level or video level", + default="image", + ) + parser.add_argument( + "--apr", + help="Report average metrics instead of simple ones", + action="store_true", + ) + args = parser.parse_args() + + file = args.saved_preds + model_base = args.model_base + metric_level = args.metric_level + apr = args.apr + # Load data + data = load_file(file_path=file) + + if "data" in data.keys(): + data = data["data"] + + total_preds = [] + total_labels = [] + neg_preds = [] + pos_preds = [] + vid_preds = {} + vid_labels = {} + + if model_base == "image": + for ip in data.keys(): + vid_id = os.path.dirname(ip) + pred = data[ip][0] + label = data[ip][1] + + # Just append to total preds + total_preds.append(pred) + total_labels.append(label) + + if vid_id in vid_preds.keys(): + vid_preds[vid_id].append(pred) + else: + vid_preds[vid_id] = [pred] + vid_labels[vid_id] = [label] + + if metric_level == "video": + total_preds = [ + np.mean(vid_preds[k], keepdims=True) for k in vid_preds.keys() + ] + total_labels = [vid_labels[k] for k in vid_labels.keys()] + else: + for ip in data.keys(): + vid_id = ip.split("/")[-1] + pred = [v for idx, v in enumerate(data[ip]) if idx % 2 == 0] + pred = [np.array(pred).mean()] + label = [data[ip][1]] + + # Just append to total preds + total_preds.append(pred) + total_labels.append(label) + + total_preds = sigmoid(np.array(total_preds)) + total_labels = np.array(total_labels) + + # Assigning predictions to neg/pos groups + neg_preds = total_preds[total_labels < 1].squeeze() + pos_preds = total_preds[total_labels == 1].squeeze() + + # Computing metric section + acc_measure = get_acc_mesure_func("binary") + acc_ = acc_measure(total_preds, targets=None, labels=total_labels) + metrics = bin_calculate_auc_ap_ar(total_preds, total_labels, apr=apr) + best_thr = metrics["best_thr"] + thr_var = metrics["thr_var"] + + if apr: + auc_, ap_, ar_, mf1_ = ( + metrics["auc"], + metrics["ap"], + metrics["ar"], + metrics["mf1"], + ) + print( + f"Current ACC, AUC, AP, AR, mF1, THR --- {acc_*100} -- {auc_*100} -- {ap_*100} -- {ar_*100} -- {mf1_*100} -- {best_thr}" + ) + else: + bacc_, auc_, p_, r_, s_, f1_, eer_ = ( + metrics["bacc"], + metrics["auc"], + metrics["p"], + metrics["r"], + metrics["s"], + metrics["f1"], + metrics["eer"], + ) + print( + f"Current ACC, BACC, AUC, P, R, S, F1, EER, THR, THR_VAR -- {acc_*100} -- {bacc_*100} -- {auc_*100} -- {p_*100} -- {r_*100} -- {s_*100} -- {f1_*100} -- {eer_*100} -- {best_thr:.3f} -- {thr_var:.6f}" + ) + + # Computing Wasserstein distance + for cdf_type in ["empirical", "kde", "para", "quantile"]: + cdf1, cdf2 = apply_cdf_transform( + neg_data=neg_preds, pos_data=pos_preds, cdf_type=cdf_type + ) + + # Customize the calculation of quantile + if cdf_type != "quantile": + wd = wasserstein_distance(cdf1, cdf2) + else: + q = np.linspace(0, 1, min(len(neg_preds), len(pos_preds))) + wd = np.trapz(np.abs(cdf1 - cdf2), q) + + print(f"Wasserstein Distance --- {cdf_type} --- {wd:.6f}") diff --git a/video/fake-stormer/model_code/package_utils/misc.py b/video/fake-stormer/model_code/package_utils/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..2cf8addd71b82a0d60864af10a73536e2c4cf377 --- /dev/null +++ b/video/fake-stormer/model_code/package_utils/misc.py @@ -0,0 +1,74 @@ +# -*- coding:utf-8 -*- +import torch +from torch._six import inf + + +class NativeScalerWithGradNormCount: + state_dict_key = "amp_scaler" + + def __init__(self): + self._scaler = torch.cuda.amp.GradScaler() + + def __call__( + self, + cfg, + loss, + optimizer, + clip_grad=None, + parameters=None, + create_graph=False, + update_grad=True, + step=0, + ): + 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) + + if cfg.TRAIN.optimizer != "SAM": + self._scaler.step(optimizer) + else: + if step == 0: + optimizer.first_step(zero_grad=True) + else: + self._scaler = optimizer.second_step( + zero_grad=False, scaler=self._scaler + ) + 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.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 diff --git a/video/fake-stormer/model_code/package_utils/tensors.py b/video/fake-stormer/model_code/package_utils/tensors.py new file mode 100644 index 0000000000000000000000000000000000000000..046344eca868ba6ee2da75bb6ef885ad1e728a91 --- /dev/null +++ b/video/fake-stormer/model_code/package_utils/tensors.py @@ -0,0 +1,93 @@ +# -*- coding:utf-8 -*- +import cv2 +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + + +def masked_inputs( + inputs: torch.tensor, + hm_preds: torch.tensor, + cfg: dict, + patch_size: int = 16, + prev_pos_mask: torch.tensor = None, + debug=False, + shot=0, + vid_ids=None, +): + """ + Masked out inputs given positions extracted from Heatmap_preds for multi-shot predictions. + args: + inputs of shape (B, 3, T, H, W) + hm of shape (B, 1, T, Hp, Wp) where Hp, Wp = H//patch_size, W//patch_size + """ + # Prepare positions to mask + batch_size = inputs.shape[0] + Hp, Wp = hm_preds.shape[-2:] + H, W = inputs.shape[-2:] + + hm_preds = hm_preds.reshape(batch_size, hm_preds.shape[1], hm_preds.shape[2], -1) + max_vals = hm_preds.max(axis=-1, keepdim=True)[0] + # pos_ = hm_preds.eq(max_vals).float()[:,:,0,:].unsqueeze(2).repeat(1,1,hm_preds.shape[2],1).unsqueeze(-1) #Take the first pos matrix + # pos_ = pos_.repeat(1,1,1,1,patch_size*patch_size).reshape(batch_size, pos_.shape[1], pos_.shape[2], H, W) + pos_ = ( + hm_preds.eq(max_vals) + .float() + .reshape(batch_size, hm_preds.shape[1], hm_preds.shape[2], Hp, Wp) + ) + pos_ = pos_[:, :, 1, :, :].unsqueeze(2).repeat(1, 1, hm_preds.shape[2], 1, 1) + pos_ = F.interpolate(pos_, size=(hm_preds.shape[2], H, W), mode="nearest") + + # Prepare values to fill up + hard_vals = torch.zeros((1, 1, 1, 1, 3), dtype=torch.float) + mean = torch.as_tensor(cfg.TRANSFORM.normalize.mean) + std = torch.as_tensor(cfg.TRANSFORM.normalize.std) + normalized_vals = hard_vals.sub_(mean).div_(std).to(dtype=torch.float).cuda() + + inputs_ = inputs.permute(0, 2, 3, 4, 1) + pos = pos_.permute(0, 2, 3, 4, 1) + + if prev_pos_mask is not None: + pos = torch.logical_or(pos, prev_pos_mask).int() + + outs = torch.where(pos.repeat(1, 1, 1, 1, 3).bool(), normalized_vals, inputs_) + inputs = outs.permute(0, 4, 1, 2, 3) + + if debug: + if vid_ids is not None: + vis_input_tensor(inputs[0], file_name=f"test_{vid_ids[0]}_{shot+1}.png") + else: + vis_input_tensor(inputs[0], file_name=f"test_{shot+1}.png") + + return inputs, pos + + +def vis_input_tensor(tensor, file_name, normalize=True): + """ + Visualize input tensor for debugging + args: + tensor of shape (3, H, W) or (3, T, H, W) + """ + assert tensor.ndim in [3, 4] + H, W = tensor.shape[-2:] + + if normalize: + tensor = tensor.clone() + min = float(tensor.min()) + max = float(tensor.max()) + tensor.add_(-min).div_(max - min + 1e-5) + + if tensor.ndim == 4: + depth = tensor.shape[1] + inputs = tensor.mul(255).clamp(0, 255).byte().permute(1, 2, 3, 0).cpu().numpy() + else: + inputs = tensor.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy() + depth = 1 + + grid_image = np.zeros((H, depth * W, 3), dtype=np.uint8) + for i in range(depth): + image = inputs[i] + grid_image[0:H, i * W : (i + 1) * W, :] = image + + Image.fromarray(grid_image).save(file_name) diff --git a/video/fake-stormer/model_code/package_utils/transform.py b/video/fake-stormer/model_code/package_utils/transform.py new file mode 100644 index 0000000000000000000000000000000000000000..76b44718c4bb5e1a648ff578e6448817be11c87d --- /dev/null +++ b/video/fake-stormer/model_code/package_utils/transform.py @@ -0,0 +1,167 @@ +# -*- coding: utf-8 -*- +from copy import deepcopy + +import albumentations as alb +import cv2 +import numpy as np +from torchvision import transforms + + +def get_dir(src_point, rot_rad): + sn, cs = np.sin(rot_rad), np.cos(rot_rad) + + src_result = [0, 0] + src_result[0] = src_point[0] * cs - src_point[1] * sn + src_result[1] = src_point[0] * sn + src_point[1] * cs + + return src_result + + +def get_3rd_point(a, b): + direct = a - b + return b + np.array([-direct[1], direct[0]], dtype=np.float32) + + +def get_affine_transform( + center, + scale, + rot, + output_size, + shift=np.array([0, 0], dtype=np.float32), + inv=0, + pixel_std=200, +): + if not isinstance(scale, np.ndarray) and not isinstance(scale, list): + print(scale) + scale = np.array([scale, scale]) + + scale_tmp = scale * pixel_std + src_w = scale_tmp[0] + dst_w = output_size[0] + dst_h = output_size[1] + + rot_rad = np.pi * rot / 180 + + src_dir = get_dir([0, (src_w - 1) * -0.5], rot_rad) + dst_dir = np.array([0, (dst_w - 1) * -0.5], np.float32) + src = np.zeros((3, 2), dtype=np.float32) + dst = np.zeros((3, 2), dtype=np.float32) + src[0, :] = center + scale_tmp * shift + src[1, :] = center + src_dir + scale_tmp * shift + dst[0, :] = [(dst_w - 1) * 0.5, (dst_h - 1) * 0.5] + dst[1, :] = np.array([(dst_w - 1) * 0.5, (dst_h - 1) * 0.5]) + dst_dir + + src[2:, :] = get_3rd_point(src[0, :], src[1, :]) + dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :]) + + if inv: + trans = cv2.getAffineTransform(np.float32(dst), np.float32(src)) + else: + trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) + + return trans + + +def affine_transform(pt, t): + """ + This function apply the affine transform to each point given by an affine matrix + """ + new_pt = np.array([pt[0], pt[1], 1.0]).T + new_pt = np.dot(t, new_pt) + return new_pt[:2] + + +def get_center_scale(shape, aspect_ratio, pixel_std=200): + h, w = shape[0], shape[1] + center = np.zeros((2), dtype=np.float32) + center[0] = (shape[1] - 1) / 2 + center[1] = (shape[0] - 1) / 2 + + if w > h * aspect_ratio: + h = w * 1.0 / aspect_ratio + else: + w = h * 1.0 / aspect_ratio + scale = np.array([w * 1.0 / pixel_std, h * 1.0 / pixel_std], dtype=np.float32) + + return center, scale + + +def final_transform(_cfg): + return transforms.Compose( + [ + transforms.ToTensor(), + transforms.Normalize( + mean=_cfg.TRANSFORM.normalize.mean, + std=_cfg.TRANSFORM.normalize.std, + ), + ] + ) + + +def randaffine(img, mask, index=0, data_type="image", **kwargs): + assert data_type in ["image", "video"] + assert mask.ndim == 2 + + if data_type == "image": + f = alb.Affine( + translate_percent={"x": (-0.03, 0.03), "y": (-0.015, 0.015)}, + scale=[0.95, 1 / 0.95], + fit_output=False, + p=1, + ) + + g = alb.ElasticTransform(alpha=50, sigma=7, alpha_affine=0, p=1) + else: + f = alb.ReplayCompose( + [ + alb.Affine( + translate_percent={"x": (-0.03, 0.03), "y": (-0.015, 0.015)}, + scale=[0.95, 1 / 0.95], + fit_output=False, + p=1, + ) + ], + p=1, + ) + + g = alb.ReplayCompose( + [alb.ElasticTransform(alpha=50, sigma=7, alpha_affine=0, p=1)], p=1 + ) + + if index == 0 or data_type == "image": + data_f = f(image=img, mask=mask) + img = data_f["image"] + mask = data_f["mask"] + + data_g = g(image=img, mask=mask) + mask = data_g["mask"] + + if data_type == "image": + return img, mask, None + else: + f_replay_params = data_f["replay"] + g_replay_params = data_g["replay"] + return ( + img, + mask, + { + "f_replay_params": f_replay_params, + "g_replay_params": g_replay_params, + }, + ) + else: + f_replay_params = kwargs.get("f_replay_params") + g_replay_params = kwargs.get("g_replay_params") + assert f_replay_params is not None and g_replay_params is not None + + data_f = alb.ReplayCompose.replay(f_replay_params, image=img, mask=mask) + img = data_f["image"] + mask = data_f["mask"] + + data_g = alb.ReplayCompose.replay(g_replay_params, image=img, mask=mask) + mask = data_g["mask"] + return ( + img, + mask, + {"f_replay_params": f_replay_params, "g_replay_params": g_replay_params}, + ) diff --git a/video/fake-stormer/model_code/package_utils/utils.py b/video/fake-stormer/model_code/package_utils/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..637ac9452b89aa6074fd6ea46ae11e3940a50870 --- /dev/null +++ b/video/fake-stormer/model_code/package_utils/utils.py @@ -0,0 +1,267 @@ +# -*- coding: utf-8 -*- +import logging +import os +from copy import deepcopy + +import cv2 +import numpy as np +import plotly.graph_objects as go +import simplejson as json +import torch +from losses.losses import _sigmoid + + +def file_extention(file_path): + f_name, f_extension = os.path.splitext(file_path) + return f_name, f_extension + + +def make_dir(dir_path): + if not os.path.exists(dir_path): + os.mkdir(dir_path) + + +def vis_heatmap(images, heatmaps, file_name, **kwargs): + temp_locs = kwargs.get("temp_loc_preds") + + # hm_h, hm_w = heatmaps.shape[1:] + hm_h, hm_w = np.array(images[0]).shape[:2] + + masked_image = np.zeros((hm_h, hm_w * heatmaps.shape[0], 3), dtype=np.uint8) + + for i in range(heatmaps.shape[0]): + heatmap = heatmaps[i] + heatmap = np.clip(heatmap * 255, 0, 255).astype(np.uint8) + heatmap = np.squeeze(heatmap) + + # heatmap_h = heatmap.shape[0] + # heatmap_w = heatmap.shape[1] + + if isinstance(images, list): + # resized_image = cv2.resize(np.array(images[i]), (int(heatmap_h), int(heatmap_w))) + heatmap = cv2.resize( + heatmap, np.array(images[i]).shape[:2], interpolation=cv2.INTER_LINEAR + ) + colored_heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) + masked_image[:, hm_w * i : hm_w * (i + 1), :] = ( + colored_heatmap * 0.7 + np.array(images[i]) * 0.3 + ) + else: + # resized_image = cv2.resize(images, (int(heatmap_h), int(heatmap_w))) + heatmap = cv2.resize( + heatmap, images.shape[:2], interpolation=cv2.INTER_LINEAR + ) + colored_heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) + masked_image[:, hm_w * i : hm_w * (i + 1), :] = ( + colored_heatmap * 0.7 + images * 0.3 + ) + + if temp_locs is not None: + font = cv2.FONT_HERSHEY_SIMPLEX + font_scale = 1 + color = (255, 255, 255) # White color in BGR + thickness = 2 + position = (20, 30) + + temp_loc = temp_locs[i] + text = f"{temp_loc:.5f}" + cv2.putText( + masked_image[:, hm_w * i : hm_w * (i + 1), :], + text, + position, + font, + font_scale, + color, + thickness, + ) + cv2.imwrite(file_name, masked_image) + + +def vis_3d_heatmap(heatmap, file_name): + # Define the dimensions of the cuboid + z_dim, y_dim, x_dim = heatmap.shape + + Z, Y, X = np.mgrid[:z_dim, :y_dim, :x_dim] + + fig = go.Figure( + data=go.Volume( + x=X.flatten(), + y=Y.flatten(), + z=Z.flatten(), + value=heatmap.flatten(), + isomin=0.0, + isomax=0.999, + opacity=0.1, + surface_count=25, + ) + ) + fig.update_layout( + scene_xaxis_showticklabels=False, + scene_yaxis_showticklabels=False, + scene_zaxis_showticklabels=False, + ) + fig.write_image(file_name) + + +def save_batch_heatmaps( + batch_image, batch_heatmaps, file_name, normalize=True, batch_cls=None +): + """ + batch_image: [batch_size, channel, height, width] + batch_heatmaps: ['batch_size, num_joints, height, width] + batch_cls: ['batch_size, num_joints, 1] + file_name: saved file name + """ + if normalize: + batch_image = batch_image.clone() + min = float(batch_image.min()) + max = float(batch_image.max()) + + batch_image.add_(-min).div_(max - min + 1e-5) + + batch_size = batch_heatmaps.size(0) + num_joints = batch_heatmaps.size(1) + heatmap_height = batch_heatmaps.size(2) + heatmap_width = batch_heatmaps.size(3) + + grid_image = np.zeros( + (batch_size * heatmap_height, (num_joints + 1) * heatmap_width, 3), + dtype=np.uint8, + ) + + for i in range(batch_size): + if batch_image.dim() == 4: + image = ( + batch_image[i] + .mul(255) + .clamp(0, 255) + .byte() + .permute(1, 2, 0) + .cpu() + .numpy() + ) + else: + image = ( + batch_image[i] + .mul(255) + .clamp(0, 255) + .byte() + .permute(1, 2, 3, 0) + .cpu() + .numpy() + ) + heatmaps = batch_heatmaps[i].mul(255).clamp(0, 255).byte().cpu().numpy() + + height_begin = heatmap_height * i + height_end = heatmap_height * (i + 1) + for j in range(num_joints): + if image.ndim == 4: + image = image[j, :, :, :] + + resized_image = cv2.resize(image, (int(heatmap_width), int(heatmap_height))) + + heatmap = heatmaps[j, :, :] + colored_heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) + + if batch_cls is not None: + cls = batch_cls[i][j].detach().cpu().numpy() + colored_heatmap = cv2.putText( + colored_heatmap, + f"Cls Pred: {cls}", + (heatmap_width * (j + 1) - 15, 10), + cv2.FONT_HERSHEY_SIMPLEX, + 1, + 1, + cv2.LINE_AA, + ) + masked_image = colored_heatmap * 0.7 + resized_image * 0.3 + + width_begin = heatmap_width * (j + 1) + width_end = heatmap_width * (j + 2) + + grid_image[height_begin:height_end, width_begin:width_end, :] = masked_image + + grid_image[height_begin:height_end, 0:heatmap_width, :] = resized_image + cv2.imwrite(file_name, grid_image) + + +def debugging_panel( + debug_cfg, + batch_image, + batch_heatmaps_gt, + batch_heatmaps_pred, + idx, + normalize=True, + batch_cls_gt=None, + batch_cls_pred=None, + split="train", +): + if debug_cfg.save_hm_gt: + save_batch_heatmaps( + batch_image, + batch_heatmaps_gt, + f"samples/{split}_debugs/hm_gt_{idx}.jpg", + normalize=normalize, + ) + + if debug_cfg.save_hm_pred: + batch_heatmaps_pred_ = _sigmoid(batch_heatmaps_pred.clone()) + save_batch_heatmaps( + batch_image, + batch_heatmaps_pred_, + f"samples/{split}_debugs/hm_pred_{idx}.jpg", + normalize=normalize, + ) + + +def save_file(data, file_path): + f_name, f_extention = file_extention(file_path) + + if f_extention == ".json": + with open(file_path, "w") as f: + json.dump(data, f) + print(f"Data has been saved to --- {file_path}") + else: + raise ValueError(f"{f_extention} is not supported now!") + + +def load_file(file_path): + f_name, f_extention = file_extention(file_path) + + if f_extention == ".json": + with open(file_path, "r") as f: + data = json.load(f) + print(f"Data has been loaded from --- {file_path}") + else: + raise ValueError(f"{f_extention} is not supported now!") + + return data + + +def draw_landmarks(image, landmarks): + """This function is to draw facial landmarks into transformed images""" + assert landmarks is not None, "Landmarks can not be None!" + + img_cp = deepcopy(image) + landmarks = landmarks.astype(int) + + for i, p in enumerate(landmarks): + img_cp = cv2.circle(img_cp, (p[0], p[1]), 2, (0, 255, 0), 1) + + return img_cp + + +def draw_most_vul_points(blended_mask): + """Detecting and Drawing the most vulnerable points for visualization purpose""" + b_mask_cp = deepcopy(blended_mask) + target_H, target_W, target_C = b_mask_cp.shape + + max_val = b_mask_cp[..., 0].max() + max_val = max_val if max_val > 0 else 1 + + m_v_indices = np.where(b_mask_cp == max_val) + + for j, i in zip(m_v_indices[0], m_v_indices[1]): + b_mask_cp[j, i] = (255, 0, 0) + + return b_mask_cp diff --git a/video/fake-stormer/model_code/register/misc.py b/video/fake-stormer/model_code/register/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..179c8a1bcbb078350cf84a3a170ab61fd174661d --- /dev/null +++ b/video/fake-stormer/model_code/register/misc.py @@ -0,0 +1,359 @@ +# -*- coding: utf-8 -*- +import collections.abc +import functools +import itertools +import subprocess +import warnings +from collections import abc +from importlib import import_module +from inspect import getfullargspec +from itertools import repeat + + +# From PyTorch internals +def _ntuple(n): + + def parse(x): + if isinstance(x, collections.abc.Iterable): + return x + return tuple(repeat(x, n)) + + return parse + + +to_1tuple = _ntuple(1) +to_2tuple = _ntuple(2) +to_3tuple = _ntuple(3) +to_4tuple = _ntuple(4) +to_ntuple = _ntuple + + +def is_str(x): + """Whether the input is an string instance. + Note: This method is deprecated since python 2 is no longer supported. + """ + return isinstance(x, str) + + +def import_modules_from_strings(imports, allow_failed_imports=False): + """Import modules from the given list of strings. + Args: + imports (list | str | None): The given module names to be imported. + allow_failed_imports (bool): If True, the failed imports will return + None. Otherwise, an ImportError is raise. Default: False. + Returns: + list[module] | module | None: The imported modules. + Examples: + >>> osp, sys = import_modules_from_strings( + ... ['os.path', 'sys']) + >>> import os.path as osp_ + >>> import sys as sys_ + >>> assert osp == osp_ + >>> assert sys == sys_ + """ + if not imports: + return + single_import = False + if isinstance(imports, str): + single_import = True + imports = [imports] + if not isinstance(imports, list): + raise TypeError(f"custom_imports must be a list but got type {type(imports)}") + imported = [] + for imp in imports: + if not isinstance(imp, str): + raise TypeError(f"{imp} is of type {type(imp)} and cannot be imported.") + try: + imported_tmp = import_module(imp) + except ImportError: + if allow_failed_imports: + warnings.warn(f"{imp} failed to import and is ignored.", UserWarning) + imported_tmp = None + else: + raise ImportError + imported.append(imported_tmp) + if single_import: + imported = imported[0] + return imported + + +def iter_cast(inputs, dst_type, return_type=None): + """Cast elements of an iterable object into some type. + Args: + inputs (Iterable): The input object. + dst_type (type): Destination type. + return_type (type, optional): If specified, the output object will be + converted to this type, otherwise an iterator. + Returns: + iterator or specified type: The converted object. + """ + if not isinstance(inputs, abc.Iterable): + raise TypeError("inputs must be an iterable object") + if not isinstance(dst_type, type): + raise TypeError('"dst_type" must be a valid type') + + out_iterable = map(dst_type, inputs) + + if return_type is None: + return out_iterable + else: + return return_type(out_iterable) + + +def list_cast(inputs, dst_type): + """Cast elements of an iterable object into a list of some type. + A partial method of :func:`iter_cast`. + """ + return iter_cast(inputs, dst_type, return_type=list) + + +def tuple_cast(inputs, dst_type): + """Cast elements of an iterable object into a tuple of some type. + A partial method of :func:`iter_cast`. + """ + return iter_cast(inputs, dst_type, return_type=tuple) + + +def is_seq_of(seq, expected_type, seq_type=None): + """Check whether it is a sequence of some type. + Args: + seq (Sequence): The sequence to be checked. + expected_type (type): Expected type of sequence items. + seq_type (type, optional): Expected sequence type. + Returns: + bool: Whether the sequence is valid. + """ + if seq_type is None: + exp_seq_type = abc.Sequence + else: + assert isinstance(seq_type, type) + exp_seq_type = seq_type + if not isinstance(seq, exp_seq_type): + return False + for item in seq: + if not isinstance(item, expected_type): + return False + return True + + +def is_list_of(seq, expected_type): + """Check whether it is a list of some type. + A partial method of :func:`is_seq_of`. + """ + return is_seq_of(seq, expected_type, seq_type=list) + + +def is_tuple_of(seq, expected_type): + """Check whether it is a tuple of some type. + A partial method of :func:`is_seq_of`. + """ + return is_seq_of(seq, expected_type, seq_type=tuple) + + +def slice_list(in_list, lens): + """Slice a list into several sub lists by a list of given length. + Args: + in_list (list): The list to be sliced. + lens(int or list): The expected length of each out list. + Returns: + list: A list of sliced list. + """ + if isinstance(lens, int): + assert len(in_list) % lens == 0 + lens = [lens] * int(len(in_list) / lens) + if not isinstance(lens, list): + raise TypeError('"indices" must be an integer or a list of integers') + elif sum(lens) != len(in_list): + raise ValueError( + "sum of lens and list length does not " + f"match: {sum(lens)} != {len(in_list)}" + ) + out_list = [] + idx = 0 + for i in range(len(lens)): + out_list.append(in_list[idx : idx + lens[i]]) + idx += lens[i] + return out_list + + +def concat_list(in_list): + """Concatenate a list of list into a single list. + Args: + in_list (list): The list of list to be merged. + Returns: + list: The concatenated flat list. + """ + return list(itertools.chain(*in_list)) + + +def check_prerequisites( + prerequisites, + checker, + msg_tmpl='Prerequisites "{}" are required in method "{}" but not ' + "found, please install them first.", +): # yapf: disable + """A decorator factory to check if prerequisites are satisfied. + Args: + prerequisites (str of list[str]): Prerequisites to be checked. + checker (callable): The checker method that returns True if a + prerequisite is meet, False otherwise. + msg_tmpl (str): The message template with two variables. + Returns: + decorator: A specific decorator. + """ + + def wrap(func): + + @functools.wraps(func) + def wrapped_func(*args, **kwargs): + requirements = ( + [prerequisites] if isinstance(prerequisites, str) else prerequisites + ) + missing = [] + for item in requirements: + if not checker(item): + missing.append(item) + if missing: + print(msg_tmpl.format(", ".join(missing), func.__name__)) + raise RuntimeError("Prerequisites not meet.") + else: + return func(*args, **kwargs) + + return wrapped_func + + return wrap + + +def _check_py_package(package): + try: + import_module(package) + except ImportError: + return False + else: + return True + + +def _check_executable(cmd): + if subprocess.call(f"which {cmd}", shell=True) != 0: + return False + else: + return True + + +def requires_package(prerequisites): + """A decorator to check if some python packages are installed. + Example: + >>> @requires_package('numpy') + >>> func(arg1, args): + >>> return numpy.zeros(1) + array([0.]) + >>> @requires_package(['numpy', 'non_package']) + >>> func(arg1, args): + >>> return numpy.zeros(1) + ImportError + """ + return check_prerequisites(prerequisites, checker=_check_py_package) + + +def requires_executable(prerequisites): + """A decorator to check if some executable files are installed. + Example: + >>> @requires_executable('ffmpeg') + >>> func(arg1, args): + >>> print(1) + 1 + """ + return check_prerequisites(prerequisites, checker=_check_executable) + + +def deprecated_api_warning(name_dict, cls_name=None): + """A decorator to check if some arguments are deprecate and try to replace + deprecate src_arg_name to dst_arg_name. + Args: + name_dict(dict): + key (str): Deprecate argument names. + val (str): Expected argument names. + Returns: + func: New function. + """ + + def api_warning_wrapper(old_func): + + @functools.wraps(old_func) + def new_func(*args, **kwargs): + # get the arg spec of the decorated method + args_info = getfullargspec(old_func) + # get name of the function + func_name = old_func.__name__ + if cls_name is not None: + func_name = f"{cls_name}.{func_name}" + if args: + arg_names = args_info.args[: len(args)] + for src_arg_name, dst_arg_name in name_dict.items(): + if src_arg_name in arg_names: + warnings.warn( + f'"{src_arg_name}" is deprecated in ' + f'`{func_name}`, please use "{dst_arg_name}" ' + "instead", + DeprecationWarning, + ) + arg_names[arg_names.index(src_arg_name)] = dst_arg_name + if kwargs: + for src_arg_name, dst_arg_name in name_dict.items(): + if src_arg_name in kwargs: + + assert dst_arg_name not in kwargs, ( + f"The expected behavior is to replace " + f"the deprecated key `{src_arg_name}` to " + f"new key `{dst_arg_name}`, but got them " + f"in the arguments at the same time, which " + f"is confusing. `{src_arg_name} will be " + f"deprecated in the future, please " + f"use `{dst_arg_name}` instead." + ) + + warnings.warn( + f'"{src_arg_name}" is deprecated in ' + f'`{func_name}`, please use "{dst_arg_name}" ' + "instead", + DeprecationWarning, + ) + kwargs[dst_arg_name] = kwargs.pop(src_arg_name) + + # apply converted arguments to the decorated method + output = old_func(*args, **kwargs) + return output + + return new_func + + return api_warning_wrapper + + +def is_method_overridden(method, base_class, derived_class): + """Check if a method of base class is overridden in derived class. + Args: + method (str): the method name to check. + base_class (type): the class of the base class. + derived_class (type | Any): the class or instance of the derived class. + """ + assert isinstance( + base_class, type + ), "base_class doesn't accept instance, Please pass class instead." + + if not isinstance(derived_class, type): + derived_class = derived_class.__class__ + + base_method = getattr(base_class, method) + derived_method = getattr(derived_class, method) + return derived_method != base_method + + +def has_method(obj: object, method: str) -> bool: + """Check whether the object has a method. + Args: + method (str): The method name to check. + obj (object): The object to check. + Returns: + bool: True if the object has the method else False. + """ + return hasattr(obj, method) and callable(getattr(obj, method)) diff --git a/video/fake-stormer/model_code/register/register.py b/video/fake-stormer/model_code/register/register.py new file mode 100644 index 0000000000000000000000000000000000000000..e3d934f7309e3d3a22c3616804707a0e90cf2f9c --- /dev/null +++ b/video/fake-stormer/model_code/register/register.py @@ -0,0 +1,323 @@ +# -*- coding: utf-8 -*- +import inspect +import warnings +from functools import partial +from typing import Any, Dict, Optional + +from .misc import deprecated_api_warning, is_seq_of + + +def build_from_cfg( + cfg: Dict, registry: "Registry", default_args: Optional[Dict] = None +) -> Any: + """Build a module from config dict when it is a class configuration, or + call a function from config dict when it is a function configuration. + Example: + >>> MODELS = Registry('models') + >>> @MODELS.register_module() + >>> class ResNet: + >>> pass + >>> resnet = build_from_cfg(dict(type='Resnet'), MODELS) + >>> # Returns an instantiated object + >>> @MODELS.register_module() + >>> def resnet50(): + >>> pass + >>> resnet = build_from_cfg(dict(type='resnet50'), MODELS) + >>> # Return a result of the calling function + Args: + cfg (dict): Config dict. It should at least contain the key "type". + registry (:obj:`Registry`): The registry to search the type from. + default_args (dict, optional): Default initialization arguments. + Returns: + object: The constructed object. + """ + if not isinstance(cfg, dict): + raise TypeError(f"cfg must be a dict, but got {type(cfg)}") + if "type" not in cfg: + if default_args is None or "type" not in default_args: + raise KeyError( + '`cfg` or `default_args` must contain the key "type", ' + f"but got {cfg}\n{default_args}" + ) + if not isinstance(registry, Registry): + raise TypeError( + "registry must be an mmcv.Registry object, " f"but got {type(registry)}" + ) + if not (isinstance(default_args, dict) or default_args is None): + raise TypeError( + "default_args must be a dict or None, " f"but got {type(default_args)}" + ) + + args = cfg.copy() + + if default_args is not None: + for name, value in default_args.items(): + args.setdefault(name, value) + + obj_type = args.pop("type") + if isinstance(obj_type, str): + obj_cls = registry.get(obj_type) + if obj_cls is None: + raise KeyError(f"{obj_type} is not in the {registry.name} registry") + elif inspect.isclass(obj_type) or inspect.isfunction(obj_type): + obj_cls = obj_type + else: + raise TypeError(f"type must be a str or valid type, but got {type(obj_type)}") + try: + return obj_cls(**args) + except Exception as e: + # Normal TypeError does not print class name. + raise type(e)(f"{obj_cls.__name__}: {e}") + + +class Registry: + """A registry to map strings to classes or functions. + Registered object could be built from registry. Meanwhile, registered + functions could be called from registry. + Example: + >>> MODELS = Registry('models') + >>> @MODELS.register_module() + >>> class ResNet: + >>> pass + >>> resnet = MODELS.build(dict(type='ResNet')) + >>> @MODELS.register_module() + >>> def resnet50(): + >>> pass + >>> resnet = MODELS.build(dict(type='resnet50')) + Please refer to + https://mmcv.readthedocs.io/en/latest/understand_mmcv/registry.html for + advanced usage. + Args: + name (str): Registry name. + build_func(func, optional): Build function to construct instance from + Registry, func:`build_from_cfg` is used if neither ``parent`` or + ``build_func`` is specified. If ``parent`` is specified and + ``build_func`` is not given, ``build_func`` will be inherited + from ``parent``. Default: None. + parent (Registry, optional): Parent registry. The class registered in + children registry could be built from parent. Default: None. + scope (str, optional): The scope of registry. It is the key to search + for children registry. If not specified, scope will be the name of + the package where class is defined, e.g. mmdet, mmcls, mmseg. + Default: None. + """ + + def __init__(self, name, build_func=None, parent=None, scope=None): + self._name = name + self._module_dict = dict() + self._children = dict() + self._scope = self.infer_scope() if scope is None else scope + + # self.build_func will be set with the following priority: + # 1. build_func + # 2. parent.build_func + # 3. build_from_cfg + if build_func is None: + if parent is not None: + self.build_func = parent.build_func + else: + self.build_func = build_from_cfg + else: + self.build_func = build_func + if parent is not None: + assert isinstance(parent, Registry) + parent._add_children(self) + self.parent = parent + else: + self.parent = None + + def __len__(self): + return len(self._module_dict) + + def __contains__(self, key): + return self.get(key) is not None + + def __repr__(self): + format_str = ( + self.__class__.__name__ + f"(name={self._name}, " + f"items={self._module_dict})" + ) + return format_str + + @staticmethod + def infer_scope(): + """Infer the scope of registry. + The name of the package where registry is defined will be returned. + Example: + >>> # in mmdet/models/backbone/resnet.py + >>> MODELS = Registry('models') + >>> @MODELS.register_module() + >>> class ResNet: + >>> pass + The scope of ``ResNet`` will be ``mmdet``. + Returns: + str: The inferred scope name. + """ + # We access the caller using inspect.currentframe() instead of + # inspect.stack() for performance reasons. See details in PR #1844 + frame = inspect.currentframe() + # get the frame where `infer_scope()` is called + infer_scope_caller = frame.f_back.f_back + filename = inspect.getmodule(infer_scope_caller).__name__ + split_filename = filename.split(".") + return split_filename[0] + + @staticmethod + def split_scope_key(key): + """Split scope and key. + The first scope will be split from key. + Examples: + >>> Registry.split_scope_key('mmdet.ResNet') + 'mmdet', 'ResNet' + >>> Registry.split_scope_key('ResNet') + None, 'ResNet' + Return: + tuple[str | None, str]: The former element is the first scope of + the key, which can be ``None``. The latter is the remaining key. + """ + split_index = key.find(".") + if split_index != -1: + return key[:split_index], key[split_index + 1 :] + else: + return None, key + + @property + def name(self): + return self._name + + @property + def scope(self): + return self._scope + + @property + def module_dict(self): + return self._module_dict + + @property + def children(self): + return self._children + + def get(self, key): + """Get the registry record. + Args: + key (str): The class name in string format. + Returns: + class: The corresponding class. + """ + scope, real_key = self.split_scope_key(key) + if scope is None or scope == self._scope: + # get from self + if real_key in self._module_dict: + return self._module_dict[real_key] + else: + # get from self._children + if scope in self._children: + return self._children[scope].get(real_key) + else: + # goto root + parent = self.parent + while parent.parent is not None: + parent = parent.parent + return parent.get(key) + + def build(self, *args, **kwargs): + return self.build_func(*args, **kwargs, registry=self) + + def _add_children(self, registry): + """Add children for a registry. + The ``registry`` will be added as children based on its scope. + The parent registry could build objects from children registry. + Example: + >>> models = Registry('models') + >>> mmdet_models = Registry('models', parent=models) + >>> @mmdet_models.register_module() + >>> class ResNet: + >>> pass + >>> resnet = models.build(dict(type='mmdet.ResNet')) + """ + + assert isinstance(registry, Registry) + assert registry.scope is not None + assert ( + registry.scope not in self.children + ), f"scope {registry.scope} exists in {self.name} registry" + self.children[registry.scope] = registry + + @deprecated_api_warning(name_dict=dict(module_class="module")) + def _register_module(self, module, module_name=None, force=False): + if not inspect.isclass(module) and not inspect.isfunction(module): + raise TypeError( + "module must be a class or a function, " f"but got {type(module)}" + ) + + if module_name is None: + module_name = module.__name__ + if isinstance(module_name, str): + module_name = [module_name] + for name in module_name: + if not force and name in self._module_dict: + raise KeyError(f"{name} is already registered " f"in {self.name}") + self._module_dict[name] = module + + def deprecated_register_module(self, cls=None, force=False): + warnings.warn( + "The old API of register_module(module, force=False) " + "is deprecated and will be removed, please use the new API " + "register_module(name=None, force=False, module=None) instead.", + DeprecationWarning, + ) + if cls is None: + return partial(self.deprecated_register_module, force=force) + self._register_module(cls, force=force) + return cls + + def register_module(self, name=None, force=False, module=None): + """Register a module. + A record will be added to `self._module_dict`, whose key is the class + name or the specified name, and value is the class itself. + It can be used as a decorator or a normal function. + Example: + >>> backbones = Registry('backbone') + >>> @backbones.register_module() + >>> class ResNet: + >>> pass + >>> backbones = Registry('backbone') + >>> @backbones.register_module(name='mnet') + >>> class MobileNet: + >>> pass + >>> backbones = Registry('backbone') + >>> class ResNet: + >>> pass + >>> backbones.register_module(ResNet) + Args: + name (str | None): The module name to be registered. If not + specified, the class name will be used. + force (bool, optional): Whether to override an existing class with + the same name. Default: False. + module (type): Module class or function to be registered. + """ + if not isinstance(force, bool): + raise TypeError(f"force must be a boolean, but got {type(force)}") + # NOTE: This is a walkaround to be compatible with the old api, + # while it may introduce unexpected bugs. + if isinstance(name, type): + return self.deprecated_register_module(name, force=force) + + # raise the error ahead of time + if not (name is None or isinstance(name, str) or is_seq_of(name, str)): + raise TypeError( + "name must be either of None, an instance of str or a sequence" + f" of str, but got {type(name)}" + ) + + # use it as a normal method: x.register_module(module=SomeClass) + if module is not None: + self._register_module(module=module, module_name=name, force=force) + return module + + # use it as a decorator: @x.register_module() + def _register(module): + self._register_module(module=module, module_name=name, force=force) + return module + + return _register diff --git a/video/fake-stormer/model_code/scripts/efn_sbi.sh b/video/fake-stormer/model_code/scripts/efn_sbi.sh new file mode 100644 index 0000000000000000000000000000000000000000..4dabcee388cf55154799f3fe51a624df0fa62bbc --- /dev/null +++ b/video/fake-stormer/model_code/scripts/efn_sbi.sh @@ -0,0 +1,3 @@ +#! /bin/bash + +CUDA_VISIBLE_DEVICES=0 python scripts/train.py --cfg configs/spatial/efn4_fpn_sbi_adv.yaml diff --git a/video/fake-stormer/model_code/scripts/fakesformer_sbi.sh b/video/fake-stormer/model_code/scripts/fakesformer_sbi.sh new file mode 100644 index 0000000000000000000000000000000000000000..fe62bbe2b49973e65e530eb01ddf04d9a21346e6 --- /dev/null +++ b/video/fake-stormer/model_code/scripts/fakesformer_sbi.sh @@ -0,0 +1,8 @@ +#! /bin/bash + +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/temporal/FakeSFormer_base_c23.yaml +CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/temporal/FakeSFormer_base_c0.yaml +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/temporal/FakeSFormer_base_c40.yaml +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/temporal/FakeSFormer_large_c23.yaml +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/temporal/FakeSwin3D_base_c23.yaml +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/temporal/FakeSFormer_base_c23_224p8.yaml diff --git a/video/fake-stormer/model_code/scripts/laanet_temporal.sh b/video/fake-stormer/model_code/scripts/laanet_temporal.sh new file mode 100644 index 0000000000000000000000000000000000000000..a867e6f9c088b190c823ada306dd4c009a35b05f --- /dev/null +++ b/video/fake-stormer/model_code/scripts/laanet_temporal.sh @@ -0,0 +1,5 @@ +#! /bin/bash + +CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/temporal/ResNet3D_EFPN3D_hm3D_c23.yaml +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/temporal/FakeSwin3D_base_c23.yaml +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/temporal/FakeSwin3D_base_c0.yaml diff --git a/video/fake-stormer/model_code/scripts/swin_bi.sh b/video/fake-stormer/model_code/scripts/swin_bi.sh new file mode 100644 index 0000000000000000000000000000000000000000..253b2ef1a09315cb71c8d31a8c7fc93c7ae72027 --- /dev/null +++ b/video/fake-stormer/model_code/scripts/swin_bi.sh @@ -0,0 +1,5 @@ +#! /bin/bash + +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/spatial/swin_sbi_base.yaml +CUDA_VISIBLE_DEVICES=0 python scripts/train.py --cfg configs/spatial/swin_bi_small.yaml +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/spatial/swin_sbi_tiny.yaml diff --git a/video/fake-stormer/model_code/scripts/swin_sbi.sh b/video/fake-stormer/model_code/scripts/swin_sbi.sh new file mode 100644 index 0000000000000000000000000000000000000000..96331da07787a28fb043536ef4aa09c3e8aa5bcf --- /dev/null +++ b/video/fake-stormer/model_code/scripts/swin_sbi.sh @@ -0,0 +1,5 @@ +#! /bin/bash + +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/spatial/swin_sbi_base.yaml +CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/spatial/swin_sbi_small.yaml +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/spatial/swin_sbi_tiny.yaml diff --git a/video/fake-stormer/model_code/scripts/test.py b/video/fake-stormer/model_code/scripts/test.py new file mode 100644 index 0000000000000000000000000000000000000000..3bed76c022f1b0379e47c8d3c7cd6b54dafc2d41 --- /dev/null +++ b/video/fake-stormer/model_code/scripts/test.py @@ -0,0 +1,411 @@ +# -*- coding: utf-8 -*- +import argparse +import os +import sys +import time + +if not os.getcwd() in sys.path: + sys.path.append(os.getcwd()) +import math +import random +from glob import glob + +import cv2 +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from configs.get_config import load_config +from datasets import DATASETS, build_dataset +from lib.core_function import AverageMeter +from lib.metrics import bin_calculate_auc_ap_ar, get_acc_mesure_func +from logs.logger import LOG_DIR, Logger +from losses.losses import _sigmoid +from models import * +from natsort import natsorted +from package_utils.image_utils import crop_by_margin, load_image +from package_utils.tensors import masked_inputs +from package_utils.transform import ( + final_transform, + get_affine_transform, + get_center_scale, +) +from package_utils.utils import save_file, vis_heatmap +from PIL import Image +from torch.utils.data import DataLoader +from tqdm import tqdm + + +def parse_args(args=None): + arg_parser = argparse.ArgumentParser("Processing testing...") + arg_parser.add_argument("--cfg", "-c", help="Config file", required=True) + arg_parser.add_argument( + "--image", "-i", type=str, help="Image for the single testing mode!" + ) + arg_parser.add_argument( + "--video", "-v", type=str, help="Video for the single testing mode!" + ) + args = arg_parser.parse_args(args) + + return args + + +if __name__ == "__main__": + if sys.argv[1:] is not None: + args = sys.argv[1:] + else: + args = sys.argv[:-1] + args = parse_args(args) + + # Loading config file + cfg = load_config(args.cfg) + logger = Logger(task="testing") + + # Seed + seed = cfg.SEED + random.seed(seed) + torch.manual_seed(seed) + np.random.seed(seed) + torch.cuda.manual_seed(seed) + + task = cfg.TEST.subtask + flip_test = cfg.TEST.flip_test + logger.info("Flip Test is used --- {}".format(flip_test)) + + save_preds = cfg.TEST.save_preds + pred_file = cfg.TEST.pred_file + + if task == "test_img": + assert ( + args.image is not None + ), "Image can not be None with single image test mode!" + logger.info("Turning on single image test mode...") + if task == "test_vid": + assert ( + args.video is not None + ), "Video can not be None with single video test mode!" + assert os.path.exists( + args.video + ), "Video path must be valid, please check the path again!" + logger.info("Turning on single video test mode...") + else: + logger.info("Turning on evaluation mode...") + if task == "eval" and cfg.DATASET.DATA.TEST.FROM_FILE: + assert ( + cfg.DATASET.DATA.TEST.ANNO_FILE is not None + ), "Annotation file can not be None with evaluation test mode!" + assert len( + cfg.DATASET.DATA.TEST.ANNO_FILE + ), "Annotation file can not be empty with evaluation test mode!" + # assert os.access(cfg.DATASET.DATA.TEST.ANNO_FILE, os.R_OK), "Annotation file must be valid with evaluation test mode!" + device_count = torch.cuda.device_count() + + # build and load/initiate pretrained model + model = build_model(cfg.MODEL, MODELS).to(torch.float) + logger.info("Loading weight ... {}".format(cfg.TEST.pretrained)) + model = load_pretrained(model, cfg.TEST.pretrained) + + if device_count >= 1: + model = nn.DataParallel(model, device_ids=cfg.TEST.gpus).cuda() + else: + model = model.cuda() + + # Define essential variables + image = args.image + vid = args.video + test_file = cfg.TEST.test_file + video_level = cfg.TEST.video_level + aspect_ratio = cfg.DATASET.IMAGE_SIZE[1] * 1.0 / cfg.DATASET.IMAGE_SIZE[0] + pixel_std = 200 + rot = 0 + transforms = final_transform(cfg.DATASET) + metrics_base = cfg.METRICS_BASE + acc_measure = get_acc_mesure_func(metrics_base) + no_shot_preds = cfg.TEST.no_shot_preds or 1 + + model.eval() + if image is not None and task == "test_img": + img = load_image(image) + img = cv2.resize(img, (317, 317)) + img = img[18 : (317 - 18), 18 : (317 - 18), :] + c, s = get_center_scale(img.shape[:2], aspect_ratio, pixel_std) + trans = get_affine_transform(c, s, rot, cfg.DATASET.IMAGE_SIZE) + input = cv2.warpAffine( + img, + trans, + (int(cfg.DATASET.IMAGE_SIZE[0]), int(cfg.DATASET.IMAGE_SIZE[1])), + flags=cv2.INTER_LINEAR, + ) + with torch.no_grad(): + st = time.time() + img_trans = transforms(input / 255).to(torch.float) + img_trans = torch.unsqueeze(img_trans, 0) + if device_count > 0: + img_trans = img_trans.cuda(non_blocking=True) + + outputs = model(img_trans) + hm_outputs = outputs[0]["hm"] + cls_outputs = outputs[0]["cls"].sigmoid() + hm_preds = _sigmoid(hm_outputs).cpu().numpy() + if cfg.TEST.vis_hm: + print(f"Heatmap max value --- {hm_preds.max()}") + vis_heatmap(img, hm_preds[0], "output_pred.jpg") + label_pred = cls_outputs.cpu().numpy() + label = "Fake" if label_pred[0][-1] > cfg.TEST.threshold else "Real" + logger.info("Inferencing time --- {}".format(time.time() - st)) + logger.info("{} --- {}".format(label, label_pred[0][-1])) + logger.info("-----------------***--------------------") + if vid is not None and task == "test_vid": + print(vid) + img_list = [] + n_frames = cfg.DATASET.DATA.SAMPLES_PER_VIDEO.NUM_FRAMES + assert n_frames is not None, "Number of video frames can not be None!" + # Load first n_frames inside the video + img_paths = glob(f"{args.video}/*.png") + img_paths = natsorted(img_paths) # correct the order of image paths + img_paths = img_paths[:n_frames] + for img_path in img_paths: + img = Image.open(img_path) + H, W = img.size + img = img.crop((15, 15, W - 15, H - 15)) + img_list.append(img) + + # Transform images + transformed_imgs = torch.tensor([]).cuda() + for _i in img_list: + img_resize = _i.resize( + (int(cfg.DATASET.IMAGE_SIZE[0]), int(cfg.DATASET.IMAGE_SIZE[1])) + ) + img_resize = np.array(img_resize) / 255 + img_tensor = transforms(img_resize).to(torch.float) + if device_count > 0: + img_tensor = img_tensor.cuda(non_blocking=True) + transformed_imgs = torch.cat((transformed_imgs, img_tensor.unsqueeze(0)), 0) + + with torch.no_grad(): + st = time.time() + transformed_imgs = transformed_imgs.transpose(0, 1).unsqueeze(0) + outputs = model(transformed_imgs) + + hm_outputs = outputs[0]["hm"] + cls_outputs = outputs[0]["cls"].sigmoid() + temp_loc_outputs = outputs[0]["temp_loc"].sigmoid() + temp_loc_preds = temp_loc_outputs.cpu().numpy() + hm_preds = _sigmoid(hm_outputs).cpu().numpy() + + if cfg.TEST.vis_hm: + print(f"Heatmap max value --- {hm_preds.max()}") + print(f"Heatmap min value --- {hm_preds.min()}") + vis_heatmap( + img_list, + hm_preds[0], + "output_pred.jpg", + temp_loc_preds=temp_loc_preds[0], + ) + + label = "Fake" if temp_loc_preds[0][-1] > cfg.TEST.threshold else "Real" + logger.info("Inferencing time --- {}".format(time.time() - st)) + logger.info("{} --- {}".format(label, temp_loc_preds[0])) + logger.info("-----------------***--------------------") + if task == "eval": + logger.info(f"Using metric-base {metrics_base} for evaluation!") + logger.info(f"Video level evaluation mode: {video_level}") + st = time.time() + test_dataset = build_dataset( + cfg.DATASET, DATASETS, default_args=dict(split="test", config=cfg.DATASET) + ) + test_dataloader = DataLoader( + test_dataset, + batch_size=cfg.TRAIN.batch_size * len(cfg.TRAIN.gpus), + shuffle=True, + num_workers=cfg.DATASET.NUM_WORKERS, + ) + logger.info("Dataset loading time --- {}".format(time.time() - st)) + + apr = cfg.TEST.apr + test_dataloader = tqdm(test_dataloader, dynamic_ncols=True) + with torch.no_grad(): + # Make sure all tensors in same device + total_preds = torch.tensor([]).cuda().to(dtype=torch.float) + total_labels = torch.tensor([]).cuda().to(dtype=torch.float) + vid_preds = {} + vid_labels = {} + + # Achieving frame-level predictions to save into file + pred_meta = {} + + for b, (inputs, labels, meta) in enumerate(test_dataloader): + i_st = time.time() + prev_pos_mask = None + prev_hm_preds = None + b_vid_ids = [vid for vid in meta["vid_id"]] + + if "img_path" in meta.keys(): + b_data_paths = [ip for ip in meta["img_path"]] + elif "vid_path" in meta.keys(): + b_data_paths = [ip for ip in meta["vid_path"]] + else: + if save_preds: + raise ValueError("There is no img or vid data for saving!") + + if device_count > 0: + inputs = inputs.to(dtype=torch.float).cuda() + labels = labels.to(dtype=torch.float).cuda() + + for i_shot in range(no_shot_preds): # multi-shot predictions + logger.info(f"Running the {i_shot} shot of predictions") + if i_shot > 0: + new_inputs, pos_mask = masked_inputs( + inputs=inputs, + hm_preds=prev_hm_preds, + prev_pos_mask=prev_pos_mask, + cfg=cfg.DATASET, + patch_size=16, + shot=i_shot, + debug=False, + vid_ids=b_vid_ids, + ) + outputs = model(new_inputs) + prev_pos_mask = pos_mask + else: + outputs = model(inputs) + + # Applying Flip test + if flip_test: + if inputs.dim() == 4: + outputs_1 = model(inputs.flip(dims=(3,))) + else: + outputs_1 = model(inputs.flip(dims=(4,))) + + if isinstance(outputs, list): + outputs = outputs[0] + if flip_test: + outputs_1 = outputs_1[0] + + # In case outputs contain a dict key + if isinstance(outputs, dict): + if flip_test: + hm_outputs = ( + (outputs["hm"] + outputs_1["hm"]) / 2 + if "hm" in outputs.keys() + else None + ) + cls_outputs = (outputs["cls"] + outputs_1["cls"]) / 2 + outputs_temp_loc = ( + (outputs["temp_loc"] + outputs_1["temp_loc"]) / 2 + if "temp_loc" in outputs.keys() + else None + ) + else: + hm_outputs = ( + outputs["hm"] if "hm" in outputs.keys() else None + ) + cls_outputs = outputs["cls"] + outputs_temp_loc = ( + outputs["temp_loc"] + if "temp_loc" in outputs.keys() + else None + ) + prev_hm_preds = hm_outputs + logger.info("Inferencing time --- {}".format(time.time() - st)) + + # Grisping data item + for b_i in range(len(b_data_paths)): + if b_data_paths[b_i] not in pred_meta.keys(): + pred_meta[b_data_paths[b_i]] = list( + ( + cls_outputs[b_i].clone().detach().item(), + labels[b_i].clone().detach().item(), + ) + ) + else: + pred_meta[b_data_paths[b_i]].extend( + list( + ( + cls_outputs[b_i].clone().detach().item(), + labels[b_i].clone().detach().item(), + ) + ) + ) + + if i_shot == (no_shot_preds - 1): + if not video_level: + total_preds = torch.cat((total_preds, cls_outputs), 0) + total_labels = torch.cat((total_labels, labels), 0) + else: + for idx, vid_id in enumerate(b_vid_ids): + if vid_id in vid_preds.keys(): + vid_preds[vid_id] = torch.cat( + ( + vid_preds[vid_id], + torch.unsqueeze(cls_outputs[idx], 0), + ), + 0, + ) + else: + vid_preds[vid_id] = ( + torch.unsqueeze( + cls_outputs[idx].clone().detach(), 0 + ) + .cuda() + .to(dtype=torch.float) + ) + vid_labels[vid_id] = ( + torch.unsqueeze(labels[idx].clone().detach(), 0) + .cuda() + .to(dtype=torch.float) + ) + + if video_level: + for k in vid_preds.keys(): + total_preds = torch.cat( + (total_preds, torch.mean(vid_preds[k], 0, keepdim=True)), 0 + ) + total_labels = torch.cat((total_labels, vid_labels[k]), 0) + + acc_ = acc_measure( + total_preds, + targets=None, + labels=total_labels, + threshold=cfg.TEST.threshold, + ) + metrics = bin_calculate_auc_ap_ar( + total_preds, + total_labels, + metrics_base=metrics_base, + threshold=cfg.TEST.threshold, + apr=apr, + ) + best_thr = metrics["best_thr"] + + if apr: + auc_, ap_, ar_, mf1_ = ( + metrics["auc"], + metrics["ap"], + metrics["ar"], + metrics["mf1"], + ) + + logger.info( + f"Current ACC, AUC, AP, AR, mF1, THR for {cfg.DATASET.DATA.TEST.FAKETYPE} --- {cfg.DATASET.DATA.TEST.LABEL_FOLDER} -- \ + {acc_*100} -- {auc_*100} -- {ap_*100} -- {ar_*100} -- {mf1_*100} -- {best_thr}" + ) + else: + bacc_, auc_, p_, r_, s_, f1_, eer_ = ( + metrics["bacc"], + metrics["auc"], + metrics["p"], + metrics["r"], + metrics["s"], + metrics["f1"], + metrics["eer"], + ) + + logger.info( + f"Current ACC, BACC, AUC, P, R, S, F1, EER, THR for {cfg.DATASET.DATA.TEST.FAKETYPE} --- {cfg.DATASET.DATA.TEST.LABEL_FOLDER} -- \ + {acc_*100} -- {bacc_*100} -- {auc_*100} -- {p_*100} -- {r_*100} -- {s_*100} -- {f1_*100} -- {eer_*100} -- {best_thr}" + ) + + if save_preds: + logger.info(f"Preditions will be saved into -- {pred_file}") + save_file(data=pred_meta, file_path=pred_file) diff --git a/video/fake-stormer/model_code/scripts/test.sh b/video/fake-stormer/model_code/scripts/test.sh new file mode 100644 index 0000000000000000000000000000000000000000..b6038464d4dd1b2463dd178e685e6c57a6dee906 --- /dev/null +++ b/video/fake-stormer/model_code/scripts/test.sh @@ -0,0 +1,10 @@ +#! /bin/bash + +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/test.py --cfg configs/temporal/bin_cls/ResNet3D_c23.yaml \ +# -i 447.png + +# CUDA_VISIBLE_DEVICES=0 python scripts/test.py --cfg configs/temporal/FakeSFormer_base_c23.yaml \ +# -v /data/deepfake_cluster/datasets_df/DFW/test/frames/fake_test/fake_98_187 + +CUDA_VISIBLE_DEVICES=0 python scripts/test.py --cfg configs/spatial/xception_sbi.yaml \ + -i 447.png diff --git a/video/fake-stormer/model_code/scripts/test_bi.sh b/video/fake-stormer/model_code/scripts/test_bi.sh new file mode 100644 index 0000000000000000000000000000000000000000..b9348cd6ceff0a95e2471e075877e4bd5ed5e049 --- /dev/null +++ b/video/fake-stormer/model_code/scripts/test_bi.sh @@ -0,0 +1,10 @@ +#! /bin/bash + +# CUDA_VISIBLE_DEVICES=0 python scripts/test.py --cfg configs/spatial/vit_bi_small.yaml \ +# -i /data/deepfake_cluster/datasets_df/FaceForensics++/c0/test/frames/Deepfakes/000_003/012.png + +# CUDA_VISIBLE_DEVICES=0 python scripts/test.py --cfg configs/spatial/swin_bi_small.yaml \ +# -i ~/data/FaceForensics++/c0/test/frames/NeuralTextures/035_036/000.png + +CUDA_VISIBLE_DEVICES=0 python scripts/test.py --cfg configs/spatial/efn4_fpn_hm_adv.yaml \ + -i ~/data/FaceForensics++/c0/test/frames/NeuralTextures/035_036/000.png diff --git a/video/fake-stormer/model_code/scripts/test_efn.sh b/video/fake-stormer/model_code/scripts/test_efn.sh new file mode 100644 index 0000000000000000000000000000000000000000..3385c4b5532508a1c401c0f674ad2f2dba5811cc --- /dev/null +++ b/video/fake-stormer/model_code/scripts/test_efn.sh @@ -0,0 +1,4 @@ +#! /bin/bash + +CUDA_VISIBLE_DEVICES=0 python scripts/test.py --cfg configs/spatial/efn4_fpn_hm_adv.yaml \ + -i 447.png diff --git a/video/fake-stormer/model_code/scripts/test_fakestormer.sh b/video/fake-stormer/model_code/scripts/test_fakestormer.sh new file mode 100644 index 0000000000000000000000000000000000000000..d7ab6f1003261b9b307da6e1157bcfd606a88c50 --- /dev/null +++ b/video/fake-stormer/model_code/scripts/test_fakestormer.sh @@ -0,0 +1,10 @@ +#! /bin/bash + +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/test.py --cfg configs/temporal/FakeSFormer_base_c23.yaml \ +# -i samples/debugs/affine_f_2883.jpg + +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/test.py --cfg configs/temporal/FakeSFormer_base_c40.yaml \ +# -i samples/debugs/affine_f_2883.jpg + +CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/test.py --cfg configs/temporal/FakeSFormer_base_c0.yaml \ + -i samples/debugs/affine_f_2883.jpg diff --git a/video/fake-stormer/model_code/scripts/test_laanet_temporal.sh b/video/fake-stormer/model_code/scripts/test_laanet_temporal.sh new file mode 100644 index 0000000000000000000000000000000000000000..f13d75df68c624335232c1d786f3ec8e607471fd --- /dev/null +++ b/video/fake-stormer/model_code/scripts/test_laanet_temporal.sh @@ -0,0 +1,4 @@ +#! /bin/bash + +CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/test.py --cfg configs/temporal/FakeSwin3D_base_c23.yaml \ + -i samples/debugs/affine_f_2883.jpg diff --git a/video/fake-stormer/model_code/scripts/test_sbi.sh b/video/fake-stormer/model_code/scripts/test_sbi.sh new file mode 100644 index 0000000000000000000000000000000000000000..de415980095ffcf1aa7e364376d3b387d416b74b --- /dev/null +++ b/video/fake-stormer/model_code/scripts/test_sbi.sh @@ -0,0 +1,28 @@ +#! /bin/bash + +# CUDA_VISIBLE_DEVICES=0 python scripts/test.py --cfg configs/spatial/efn4_fpn_sbi_adv.yaml \ +# -i samples/debugs/affine_f_2883.jpg + +# CUDA_VISIBLE_DEVICES=0 python scripts/test.py --cfg configs/spatial/vit_sbi_base.yaml \ +# -i samples/debugs/affine_f_2883.jpg + +CUDA_VISIBLE_DEVICES=0 python scripts/test.py --cfg configs/spatial/vit_sbi_small.yaml \ + -i /home/users/XXX/data/FaceForensics++/c0/test/frames/Deepfakes/000_003/000.png + +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/test.py --cfg configs/spatial/vit_sbi_large.yaml \ +# -i samples/debugs/affine_f_2883.jpg + +# CUDA_VISIBLE_DEVICES=0 python scripts/test.py --cfg configs/spatial/swin_sbi_base.yaml \ +# -i samples/debugs/affine_f_2883.jpg + +# CUDA_VISIBLE_DEVICES=0 python scripts/test.py --cfg configs/spatial/swin_sbi_small.yaml \ +# -i samples/debugs/affine_f_2883.jpg + +# CUDA_VISIBLE_DEVICES=0 python scripts/test.py --cfg configs/spatial/swin_bi_small.yaml \ +# -i samples/debugs/affine_f_2883.jpg + +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/test.py --cfg configs/spatial/swin_sbi_tiny.yaml \ +# -i samples/debugs/affine_f_2883.jpg + +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/test.py --cfg configs/temporal/FakeSFormer_base.yaml \ +# -i samples/debugs/affine_f_2883.jpg diff --git a/video/fake-stormer/model_code/scripts/train.py b/video/fake-stormer/model_code/scripts/train.py new file mode 100644 index 0000000000000000000000000000000000000000..0f994354bdd5139cf9d8a9f6b28dc153148f5736 --- /dev/null +++ b/video/fake-stormer/model_code/scripts/train.py @@ -0,0 +1,243 @@ +# -*- coding: utf-8 -*- +from __future__ import absolute_import + +import os +import sys +import time + +if os.getcwd() not in sys.path: + sys.path.append(os.getcwd()) +import argparse +import random +from datetime import datetime + +import numpy as np +import torch +import torch.nn as nn +import torch.optim as optim +from configs.get_config import load_config +from datasets import * +from lib.core_function import test, train, validate +from lib.optimizers.sam import SAM +from lib.scheduler.linear_decay import LinearDecayLR +from logs.logger import LOG_DIR, Logger +from losses import * +from models import * +from package_utils.misc import NativeScalerWithGradNormCount as NativeScaler +from tensorboardX import SummaryWriter +from torch.utils.data import DataLoader + + +def args_parser(args=None): + parser = argparse.ArgumentParser("Training process...") + parser.add_argument("--cfg", help="Config file", required=True) + parser.add_argument( + "--alloc_mem", "-a", help="Pre allocating GPU memory", action="store_true" + ) + return parser.parse_args(args) + + +if __name__ == "__main__": + if len(sys.argv[1:]): + args = sys.argv[1:] + else: + args = None + + args = args_parser(args) + cfg = load_config(args.cfg) + logger = Logger(task=f"training_{cfg.TASK}") + + # Seed + seed = cfg.SEED + random.seed(seed) + torch.manual_seed(seed) + np.random.seed(seed) + torch.cuda.manual_seed(seed) + + # Allocate memory + if args.alloc_mem: + mem_all_tensors = torch.rand(60, 10000, 10000) + mem_all_tensors.to("cuda:0") + + # Configuing GPU devices + devices = torch.device("cpu") + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + + if "gpus" in cfg.TRAIN.gpus and cfg.TRAIN.gpus is not None: + # Only support a single gpu for training now + devices = torch.device("cuda:1") + model = build_model(cfg.MODEL, MODELS).cuda() + + # Loading Dataloader + start_loading = time.time() + val_dataset = build_dataset( + cfg.DATASET, DATASETS, default_args=dict(split="val", config=cfg.DATASET) + ) + val_dataloader = DataLoader( + val_dataset, + batch_size=cfg.TRAIN.batch_size * len(cfg.TRAIN.gpus), + shuffle=False, + pin_memory=cfg.DATASET.PIN_MEMORY, + num_workers=cfg.DATASET.NUM_WORKERS, + worker_init_fn=val_dataset.train_worker_init_fn, + collate_fn=val_dataset.train_collate_fn, + ) + logger.info( + "Loading val dataloader successfully! -- {}".format(time.time() - start_loading) + ) + + start_loading = time.time() + train_dataset = build_dataset( + cfg.DATASET, DATASETS, default_args=dict(split="train", config=cfg.DATASET) + ) + train_dataloader = DataLoader( + train_dataset, + batch_size=cfg.TRAIN.batch_size * len(cfg.TRAIN.gpus), + shuffle=True, + pin_memory=cfg.DATASET.PIN_MEMORY, + num_workers=cfg.DATASET.NUM_WORKERS, + worker_init_fn=train_dataset.train_worker_init_fn, + collate_fn=train_dataset.train_collate_fn, + ) + logger.info( + "Loading Train dataloader successfully! -- {}".format( + time.time() - start_loading + ) + ) + + # Defining Loss function and Optimizer + critetion = build_losses( + cfg.TRAIN.loss, LOSSES, default_args=dict(cfg=cfg.TRAIN.loss) + ).cuda() + + if cfg.TRAIN.use_amp: + eff_lr = ( + cfg.TRAIN.lr + * cfg.TRAIN.accumulation_steps + * cfg.TRAIN.batch_size + * len(cfg.TRAIN.gpus) + / 64 + ) # 16*4=64 as default, might change + else: + eff_lr = cfg.TRAIN.lr + + if cfg.TRAIN.optimizer == "Adam": + optimizer = optim.Adam(model.parameters(), lr=eff_lr, weight_decay=1e-4) + elif cfg.TRAIN.optimizer == "AdamW": + optimizer = optim.AdamW( + model.parameters(), lr=eff_lr, betas=(0.9, 0.999), weight_decay=1e-4 + ) + elif cfg.TRAIN.optimizer == "SAM": + # optimizer = SAM(model.parameters(), optim.Adam, lr=cfg.TRAIN.lr, weight_decay=1e-4) + optimizer = SAM( + model.parameters(), + optim.Adam, + lr=eff_lr, + betas=(0.9, 0.995), + weight_decay=1e-4, + ) + else: + optimizer = optim.SGD( + model.parameters(), lr=eff_lr, weight_decay=1e-5, momentum=0.9 + ) + + # Defining scaler + scaler = NativeScaler() if cfg.TRAIN.use_amp else None + + # Loading model + model, optimizer, start_epoch, scaler = preset_model( + cfg, model, optimizer=optimizer, scaler=scaler + ) + if len(cfg.TRAIN.gpus) > 0: + model = nn.DataParallel(model, device_ids=cfg.TRAIN.gpus).cuda() + else: + model = model.cuda() + + # Learning rate Scheduler + if cfg.TRAIN.lr_scheduler == "MultiStepLR": + lr_scheduler = optim.lr_scheduler.MultiStepLR( + optimizer, **cfg.TRAIN.lr_scheduler + ) + else: + lr_scheduler = LinearDecayLR( + optimizer, + cfg.TRAIN.epochs, + cfg.TRAIN.epochs // cfg.TRAIN.start_decay, + last_epoch=cfg.TRAIN.begin_epoch, + booster=cfg.TRAIN.booster, + ) + + # Enabling tensorboard + writer = SummaryWriter( + ".tensorboard/{}_{}".format(datetime.today().strftime("%Y-%m-%d"), cfg.TASK) + ) + + trainIters = 0 + valIters = 0 + min_val_loss = 1e10 + max_val_acc = 0 + max_test_auc = 0 + metrics_base = ( + cfg.METRICS_BASE + ) # Combine heatmap + cls prediction to calculate accuracy + + # Starting training process + logger.info("Starting training process...") + for epoch in range(start_epoch, cfg.TRAIN.epochs): + # Unfreezin backbone to update weights + if cfg.TRAIN.freeze_backbone and epoch == cfg.TRAIN.warm_up: + unfreeze_backbone(model) + + np.random.seed(seed + epoch) + if epoch > 0 and cfg.DATA_RELOAD: + logger.info(f"Reloading data for epoch {epoch}...") + train_dataset._reload_data(epoch=epoch) + train_dataloader = DataLoader( + train_dataset, + batch_size=cfg.TRAIN.batch_size * len(cfg.TRAIN.gpus), + shuffle=True, + pin_memory=cfg.DATASET.PIN_MEMORY, + num_workers=cfg.DATASET.NUM_WORKERS, + worker_init_fn=train_dataset.train_worker_init_fn, + collate_fn=train_dataset.train_collate_fn, + ) + + loss_avg, acc_avg, trainIters = train( + cfg, + model, + critetion, + optimizer, + epoch, + train_dataloader, + logger, + writer, + devices, + trainIters, + metrics_base=metrics_base, + scaler=scaler, + ) + if epoch % cfg.TRAIN.every_val_epochs == 0: + loss_val, acc_val, valIters = validate( + cfg, + model, + critetion, + epoch, + val_dataloader, + logger, + writer, + devices, + valIters, + metrics_base=metrics_base, + ) + + if acc_val.avg > max_val_acc: + # Saving checkpoint + ckp_path = os.path.join( + LOG_DIR, "{}_{}_model_best.pth".format(cfg.MODEL.type, cfg.TASK) + ) + save_model(path=ckp_path, epoch=epoch, model=model, optimizer=optimizer) + min_val_loss = loss_val.avg + max_val_acc = acc_val.avg + logger.info(f"Saved best model at epoch --- {epoch}") + lr_scheduler.step() diff --git a/video/fake-stormer/model_code/scripts/train.sh b/video/fake-stormer/model_code/scripts/train.sh new file mode 100644 index 0000000000000000000000000000000000000000..e1aefaffa1f1fc9523dda6436fb6cff04dd84e2e --- /dev/null +++ b/video/fake-stormer/model_code/scripts/train.sh @@ -0,0 +1,3 @@ +#! /bin/bash + +CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/temporal/bin_cls/ResNet3D_c23.yaml diff --git a/video/fake-stormer/model_code/scripts/train_efn.sh b/video/fake-stormer/model_code/scripts/train_efn.sh new file mode 100644 index 0000000000000000000000000000000000000000..aaa8961777e33d0129939acb502a566cbc6c8f47 --- /dev/null +++ b/video/fake-stormer/model_code/scripts/train_efn.sh @@ -0,0 +1,3 @@ +#! /bin/bash + +CUDA_VISIBLE_DEVICES=0 python scripts/train.py --cfg configs/spatial/binary_cls/efns/efn_4.yaml diff --git a/video/fake-stormer/model_code/scripts/vit_bi.sh b/video/fake-stormer/model_code/scripts/vit_bi.sh new file mode 100644 index 0000000000000000000000000000000000000000..aa674fa2bc529158db29e7c8475dff953c2c97a5 --- /dev/null +++ b/video/fake-stormer/model_code/scripts/vit_bi.sh @@ -0,0 +1,4 @@ +#! /bin/bash + +CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/spatial/vit_bi_small.yaml +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/spatial/vit_sbi_large.yaml diff --git a/video/fake-stormer/model_code/scripts/vit_sbi.sh b/video/fake-stormer/model_code/scripts/vit_sbi.sh new file mode 100644 index 0000000000000000000000000000000000000000..2dccd0f9ef3cf1da14918a34afafbb2768431291 --- /dev/null +++ b/video/fake-stormer/model_code/scripts/vit_sbi.sh @@ -0,0 +1,5 @@ +#! /bin/bash + +CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/spatial/vit_sbi_small.yaml +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/spatial/vit_sbi_base.yaml +# CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/train.py --cfg configs/spatial/vit_sbi_large.yaml diff --git a/video/fake-stormer/requirements.txt b/video/fake-stormer/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..eae304e0719eeb1832984c2d5e640dbc7f9cb1d8 --- /dev/null +++ b/video/fake-stormer/requirements.txt @@ -0,0 +1,24 @@ +fastapi +uvicorn +pydantic +python-multipart +torch>=1.8.0 +torchvision>=0.9.0 +opencv-python-headless +numpy<2.0.0 +Pillow +PyYAML +natsort +tqdm +scikit-image +albumentations==1.1.0 +imgaug==0.4.0 +tensorboardX==2.5.1 +plotly +simplejson +ptflops +mmengine +einops +timm +python-box +mmcv==1.6.1 diff --git a/video/lipfd/Dockerfile b/video/lipfd/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..a809444988f9417b12da2ac43474c7f40ca4385f --- /dev/null +++ b/video/lipfd/Dockerfile @@ -0,0 +1,58 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +WORKDIR /app + +# Install Python 3.10 and system dependencies for OpenCV +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3 python3-pip \ + libgl1 \ + libglib2.0-0 \ + libsm6 \ + libxext6 \ + libxrender-dev \ + && rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +# Install PyTorch with CUDA 12.1 +RUN pip install --no-cache-dir \ + torch==2.5.1 torchvision==0.20.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +# Copy requirements and install (facenet-pytorch needs --no-deps due to torch<2.3 pin) +COPY requirements.txt . +RUN pip install --no-cache-dir --no-deps facenet-pytorch && \ + pip install --no-cache-dir -r requirements.txt + +# Create logs and weights directories +RUN mkdir -p logs weights + +# Copy model code (CLIP + LipFD architecture) +COPY model_code/ /app/model_code/ + +# Copy weights +COPY weights/ /app/weights/ + +# Copy application code +COPY app.py . + +# Environment variables +ENV MODEL_PORT=7006 +ENV PRELOAD_MODEL=false +ENV MODEL_TIMEOUT=1800 +ENV WEIGHTS_PATH=/app/weights/lipfd_checkpoint.pth +ENV MODEL_CODE_DIR=/app/model_code + +# Expose port +EXPOSE 7006 + +# Drop root privileges +RUN adduser --disabled-password --gecos '' appuser && \ + chown -R appuser:appuser /app/logs /app/weights +USER appuser + +# Run the service +CMD ["python", "app.py"] diff --git a/video/lipfd/app.py b/video/lipfd/app.py new file mode 100644 index 0000000000000000000000000000000000000000..22156690ebc6d225c6454a5128d19c7e1514dbc2 --- /dev/null +++ b/video/lipfd/app.py @@ -0,0 +1,522 @@ +"""LipFD video deepfake detection service. + +Wraps the LipFD (NeurIPS 2024) lip forgery detection model with a +FastAPI endpoint. Uses CLIP ViT-L/14 as a global feature extractor +and a ResNet-50-based region-aware classifier that operates on +multi-scale crops of detected faces. + +Reference: Hou et al., "LipFD: Lip Forgery Detection via +Region-Aware CLIP", NeurIPS 2024. +""" + +import base64 +import gc +import logging +import os +import platform +import sys +import tempfile +import threading +import time +from typing import Any, Dict, List, Optional, Tuple + +import cv2 +import numpy as np +import torch +import torchvision.transforms as transforms +import uvicorn +from facenet_pytorch import MTCNN +from fastapi import FastAPI, HTTPException +from PIL import Image +from pydantic import BaseModel, ConfigDict, Field + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +MODEL_PORT = int(os.environ.get("MODEL_PORT", 7006)) +PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true" +MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", 1800)) +WEIGHTS_PATH = os.environ.get("WEIGHTS_PATH", "/app/weights/lipfd_checkpoint.pth") + +# Add model_code to path so LipFD's internal imports resolve +_MODEL_CODE_DIR = os.environ.get( + "MODEL_CODE_DIR", + os.path.join(os.path.dirname(os.path.abspath(__file__)), "model_code"), +) +if _MODEL_CODE_DIR not in sys.path: + sys.path.insert(0, _MODEL_CODE_DIR) + +# Number of frames to uniformly sample from the video +NUM_FRAMES = 32 +# Face bounding-box margin factor (fraction of bbox dimension) +MARGIN_FACTOR = 0.5 +# CLIP normalization constants (from OpenAI CLIP) +CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073] +CLIP_STD = [0.26862954, 0.26130258, 0.27577711] +# Multi-scale crop indices (from LipFD datasets.py) +# Scale 0 = full 224x224 crop, scale 1 = inner 168x168, scale 2 = inner 102x102 +CROP_IDX = [(28, 196), (61, 163)] +# Number of spatial sub-crops per face (horizontal slices of lower-face) +NUM_SUBCROP_POSITIONS = 5 + + +# ── Device selection ─────────────────────────────────────────────────────── + + +def _get_device() -> torch.device: + """Select optimal device: CUDA > MPS (Apple Silicon) > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + return torch.device("cpu") + + +# ── Global state ─────────────────────────────────────────────────────────── + +_model = None +_face_detector: Optional[MTCNN] = None +_device: Optional[torch.device] = None +_load_lock = threading.Lock() + + +def _load_models() -> None: + """Load LipFD model and MTCNN face detector (thread-safe).""" + global _model, _face_detector, _device + + if _model is not None: + return + + with _load_lock: + if _model is not None: + return + + _device = _get_device() + if _device.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + if _device.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) + else: + logger.warning( + "Device: %s (no CUDA available -- check nvidia-container-toolkit)", + _device, + ) + logger.info("Loading LipFD model on %s ...", _device) + + # ── Face detector (MTCNN) ────────────────────────────────────── + _face_detector = MTCNN( + keep_all=True, + device=_device, + post_process=False, + ) + + # ── LipFD classifier ────────────────────────────────────────── + from models import build_model # noqa: E402 + + model = build_model("CLIP:ViT-L/14") + + if not os.path.exists(WEIGHTS_PATH): + raise FileNotFoundError(f"LipFD weights not found at {WEIGHTS_PATH}") + + checkpoint = torch.load(WEIGHTS_PATH, map_location="cpu", weights_only=False) + model.load_state_dict(checkpoint["model"]) + model = model.to(_device) + model.eval() + + _model = model + logger.info("LipFD model loaded successfully.") + + +def _is_model_loaded() -> bool: + """Return True if both the classifier and face detector are loaded.""" + return _model is not None and _face_detector is not None + + +# ── FastAPI app ──────────────────────────────────────────────────────────── + +app = FastAPI( + title="LipFD Detection Service", + description=( + "Lip Forgery Detection via Region-Aware CLIP " "(ViT-L/14, NeurIPS 2024)" + ), + version="1.0.0", +) + + +class PredictRequest(BaseModel): + """Incoming prediction request.""" + + video_data: str # Base64-encoded video bytes + threshold: float = 0.5 + + +class PredictResponse(BaseModel): + """Outgoing prediction result.""" + + model_config = ConfigDict(populate_by_name=True) + + model: str = "lipfd_detection" + probability: float + prediction: int + class_name: str = Field(..., alias="class") + inference_time: float + metadata: Dict[str, Any] + + +@app.on_event("startup") +async def startup_event(): + """Optionally preload model at startup.""" + if PRELOAD_MODEL: + _load_models() + + +@app.get("/") +def root(): + """Service info endpoint.""" + return { + "service": "lipfd_detection", + "port": MODEL_PORT, + "model_loaded": _is_model_loaded(), + "device": str(_device) if _device else "unknown", + } + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and _device is not None and _device.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +def health(): + """Health check endpoint.""" + return { + "status": "healthy", + "model": "lipfd_detection", + "device": str(_device) if _device else "cpu", + "model_loaded": _is_model_loaded(), + "weights_exist": os.path.exists(WEIGHTS_PATH), + **_gpu_health_info(), + } + + +# ── Video / face utilities ───────────────────────────────────────────────── + + +def _extract_frames(video_path: str, num_frames: int = NUM_FRAMES) -> List[np.ndarray]: + """Uniformly sample *num_frames* RGB frames from a video file. + + Args: + video_path: Path to the video on disk. + num_frames: Number of frames to extract. + + Returns: + List of RGB uint8 numpy arrays (H, W, 3). + """ + cap = cv2.VideoCapture(video_path) + total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + + if total <= 0: + cap.release() + return [] + + indices = np.linspace(0, total - 1, num_frames, endpoint=True, dtype=int) + frames: List[np.ndarray] = [] + + for idx in indices: + cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx)) + ret, frame = cap.read() + if ret: + frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + + cap.release() + return frames + + +def _crop_face( + img: np.ndarray, + bbox: Tuple[float, float, float, float], + margin: float = MARGIN_FACTOR, +) -> np.ndarray: + """Crop a face region with a relative margin around the bbox. + + Args: + img: RGB image array (H, W, 3). + bbox: (x0, y0, x1, y1) face bounding box. + margin: Fraction of bbox dimension to add as padding. + + Returns: + Cropped face region as a numpy array. + """ + h_img, w_img = img.shape[:2] + x0, y0, x1, y1 = bbox + w = x1 - x0 + h = y1 - y0 + + x0_new = max(0, int(x0 - w * margin / 2)) + x1_new = min(w_img, int(x1 + w * margin / 2) + 1) + y0_new = max(0, int(y0 - h * margin / 2)) + y1_new = min(h_img, int(y1 + h * margin / 2) + 1) + + return img[y0_new:y1_new, x0_new:x1_new] + + +def _detect_faces(frame: np.ndarray) -> List[np.ndarray]: + """Detect faces in a frame and return margin-expanded crops. + + Args: + frame: RGB image array (H, W, 3). + + Returns: + List of face-crop arrays (variable size, RGB uint8). + """ + assert _face_detector is not None + + pil_img = Image.fromarray(frame) + boxes, _ = _face_detector.detect(pil_img) + + if boxes is None or len(boxes) == 0: + return [] + + crops: List[np.ndarray] = [] + for box in boxes: + x0, y0, x1, y1 = box.tolist() + face = _crop_face(frame, (x0, y0, x1, y1)) + if face.size == 0: + continue + crops.append(face) + + return crops + + +# ── LipFD preprocessing ─────────────────────────────────────────────────── + +_resize_224 = transforms.Resize((224, 224)) +_resize_1120 = transforms.Resize((1120, 1120)) +_clip_normalize = transforms.Normalize(mean=CLIP_MEAN, std=CLIP_STD) + + +def _prepare_lipfd_inputs( + face_crop: np.ndarray, +) -> Tuple[torch.Tensor, List[List[torch.Tensor]]]: + """Convert a single face crop into LipFD model inputs. + + Mirrors the preprocessing in LipFD's AVLip dataset: + 1. Convert face to float32 tensor (C, H, W), no /255 normalization + (the original dataset uses raw cv2.imread pixel values). + 2. CLIP-normalize, then create 5 horizontal sub-crops of the + lower-face region, each at 3 zoom scales (1.0x, 0.65x, 0.45x). + 3. Resize the raw image to 1120x1120 for the CLIP encoder. + + Args: + face_crop: RGB uint8 numpy array of a detected face (H, W, 3). + + Returns: + img_1120: Tensor of shape (3, 1120, 1120) for CLIP feature + extraction (raw pixel values, no normalization). + crops: List of 3 scale-lists, each containing 5 tensors of + shape (3, 224, 224) -- CLIP-normalized. + """ + # Raw float tensor (same as original: torch.tensor(cv2.imread(...))) + # cv2.imread returns BGR but we already have RGB from our pipeline. + # The original code reads BGR from cv2 and feeds it directly, so + # we convert back to BGR to match training distribution. + face_bgr = cv2.cvtColor(face_crop, cv2.COLOR_RGB2BGR) + img_tensor = torch.tensor(face_bgr, dtype=torch.float32).permute(2, 0, 1) + + # CLIP-normalize for the crop pathway + img_norm = _clip_normalize(img_tensor) + + # Build multi-scale crops from the lower-face region. + # Original dataset: face images are 1000x(variable), crops come from + # img[:, 500:, i:i+500] for 5 horizontal positions. + # We adapt to arbitrary face dimensions: take the lower half and split + # into 5 overlapping horizontal positions. + _, h, w = img_norm.shape + half_h = h // 2 + lower_face = img_norm[:, half_h:, :] # lower half of the face + + lf_h, lf_w = lower_face.shape[1], lower_face.shape[2] + # Create 5 square-ish sub-crops across the lower face width + crop_size = min(lf_h, lf_w) + if crop_size < 2: + crop_size = max(lf_h, lf_w, 2) + + # Positions spread across the width + if lf_w > crop_size: + positions = np.linspace(0, lf_w - crop_size, NUM_SUBCROP_POSITIONS, dtype=int) + else: + positions = [0] * NUM_SUBCROP_POSITIONS + + # Scale 0 (1.0x): 5 crops from lower face, each resized to 224x224 + crops: List[List[torch.Tensor]] = [[], [], []] + for pos in positions: + patch = lower_face[:, :crop_size, pos : pos + crop_size] + crop_224 = _resize_224(patch) + crops[0].append(crop_224) + + # Scale 1 (0.65x): inner crop of the 224x224 + crops[1].append( + _resize_224( + crop_224[ + :, + CROP_IDX[0][0] : CROP_IDX[0][1], + CROP_IDX[0][0] : CROP_IDX[0][1], + ] + ) + ) + # Scale 2 (0.45x): tighter inner crop + crops[2].append( + _resize_224( + crop_224[ + :, + CROP_IDX[1][0] : CROP_IDX[1][1], + CROP_IDX[1][0] : CROP_IDX[1][1], + ] + ) + ) + + # Full image resized to 1120x1120 for CLIP global feature extraction + img_1120 = _resize_1120(img_tensor) + + return img_1120, crops + + +# ── Prediction endpoint ─────────────────────────────────────────────────── + + +@app.post("/predict", response_model=PredictResponse) +async def predict(request: PredictRequest): + """Run LipFD face-forgery detection on a base64-encoded video. + + Pipeline: + 1. Decode video and write to temp file. + 2. Extract uniformly-sampled frames. + 3. Detect and crop faces per frame (MTCNN). + 4. For each face, create multi-scale crops and CLIP features. + 5. Run the region-aware classifier. + 6. Per-frame max probability, then average across frames. + + Returns probability=0.5 (undetermined) if no faces are detected. + """ + if not _is_model_loaded(): + _load_models() + + start_time = time.time() + + # ── Decode video ─────────────────────────────────────────────── + with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp: + try: + video_bytes = base64.b64decode(request.video_data) + tmp.write(video_bytes) + tmp_path = tmp.name + except Exception as e: + raise HTTPException(status_code=400, detail=f"Failed to decode video: {e}") + + try: + # ── Extract frames ───────────────────────────────────────── + frames = _extract_frames(tmp_path, NUM_FRAMES) + if not frames: + raise HTTPException( + status_code=400, + detail="Could not extract frames from video.", + ) + + # ── Detect faces and classify ────────────────────────────── + per_frame_max: List[float] = [] + total_faces = 0 + + for frame in frames: + face_crops = _detect_faces(frame) + if not face_crops: + continue + + frame_scores: List[float] = [] + + for face_crop in face_crops: + img_1120, crops = _prepare_lipfd_inputs(face_crop) + + # Batch dim = 1 for single-face inference + img_batch = img_1120.unsqueeze(0).to(_device) + + # Build crop tensors: each scale-list element gets batch dim + crops_batch = [ + [c.unsqueeze(0).to(_device) for c in scale_crops] + for scale_crops in crops + ] + + with torch.no_grad(): + # Get CLIP global features + features = _model.get_features(img_batch).to(_device) + # Run region-aware classifier + pred_score, _, _ = _model(crops_batch, features) + prob = pred_score.sigmoid().item() + + frame_scores.append(prob) + total_faces += 1 + + if frame_scores: + per_frame_max.append(max(frame_scores)) + + # ── Aggregate ────────────────────────────────────────────── + if per_frame_max: + probability = float(np.mean(per_frame_max)) + else: + probability = 0.5 + + prediction = 1 if probability >= request.threshold else 0 + class_name = "fake" if prediction == 1 else "real" + + return PredictResponse( + probability=probability, + prediction=prediction, + class_name=class_name, + inference_time=time.time() - start_time, + metadata={ + "frames_sampled": len(frames), + "frames_with_faces": len(per_frame_max), + "total_faces_detected": total_faces, + "device": str(_device), + }, + ) + + except HTTPException: + raise + except Exception as e: + logger.exception("Error during LipFD prediction") + raise HTTPException(status_code=500, detail=str(e)) + finally: + if os.path.exists(tmp_path): + os.remove(tmp_path) + gc.collect() + + +if __name__ == "__main__": + uvicorn.run(app, host="0.0.0.0", port=MODEL_PORT) diff --git a/video/lipfd/requirements.txt b/video/lipfd/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..2a8cb380490873540f836e9f3ad2670913e72c27 --- /dev/null +++ b/video/lipfd/requirements.txt @@ -0,0 +1,13 @@ +fastapi +uvicorn +pydantic +python-multipart +torch>=1.8.0 +torchvision>=0.9.0 +facenet-pytorch +opencv-python-headless +numpy<2.0.0 +Pillow +ftfy +regex +tqdm diff --git a/video/mintime/Dockerfile b/video/mintime/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..b1ce880fbf1303d4e1ee94da3c35803f033bd7f8 --- /dev/null +++ b/video/mintime/Dockerfile @@ -0,0 +1,58 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +WORKDIR /app + +# Install Python 3.10 and system dependencies for OpenCV +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3 python3-pip \ + libgl1 \ + libglib2.0-0 \ + libsm6 \ + libxext6 \ + libxrender-dev \ + && rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +# Install PyTorch with CUDA 12.1 +RUN pip install --no-cache-dir \ + torch==2.5.1 torchvision==0.20.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +# Copy requirements and install (facenet-pytorch needs --no-deps due to torch<2.3 pin) +COPY requirements.txt . +RUN pip install --no-cache-dir --no-deps facenet-pytorch && \ + pip install --no-cache-dir -r requirements.txt + +# Create logs and weights directories +RUN mkdir -p logs weights + +# Copy model code (vendored from MINTIME repo) +COPY model_code/ /app/model_code/ + +# Copy weights +COPY weights/ /app/weights/ + +# Copy application code +COPY app.py . + +# Environment variables +ENV MODEL_PORT=7008 +ENV PRELOAD_MODEL=false +ENV MODEL_TIMEOUT=1800 +ENV EXTRACTOR_WEIGHTS_PATH=/app/weights/MINTIME/MINTIME_XC_Extractor_checkpoint30 +ENV MODEL_WEIGHTS_PATH=/app/weights/MINTIME/MINTIME_XC_Model_checkpoint30 + +# Expose port +EXPOSE 7008 + +# Drop root privileges +RUN adduser --disabled-password --gecos '' appuser && \ + chown -R appuser:appuser /app/logs /app/weights +USER appuser + +# Run the service +CMD ["python", "app.py"] diff --git a/video/mintime/app.py b/video/mintime/app.py new file mode 100644 index 0000000000000000000000000000000000000000..073ea658865ffe5dd29aaf4b61022a0b62d30f60 --- /dev/null +++ b/video/mintime/app.py @@ -0,0 +1,889 @@ +"""MINTIME video deepfake detection service. + +Wraps the MINTIME model (IEEE T-IFS 2024, Multi-Identity-size-iNvariant +TIMEsformer) with a FastAPI endpoint. Uses Xception as a feature +extractor and a Size-Invariant TimeSformer for temporal classification +with identity-aware attention masks. + +The pipeline is: + 1. Extract frames from the video. + 2. Detect faces per frame with MTCNN. + 3. Crop faces, cluster them by identity (InceptionResnetV1 embeddings). + 4. Build identity-ordered sequences with size embeddings and masks. + 5. Extract Xception features, run the TimeSformer, return sigmoid score. + +Reference: Coccomini et al., "MINTIME: Multi-Identity Size-Invariant +Video Deepfake Detection", IEEE T-IFS 2024. +""" + +import base64 +import gc +import logging +import os +import platform +import sys +import tempfile +import threading +import time +from statistics import mean +from typing import Any, Dict, List, Optional, Tuple + +import cv2 +import networkx as nx +import numpy as np +import torch +import torch.nn.functional as F +import uvicorn +from einops import rearrange +from facenet_pytorch import InceptionResnetV1, fixed_image_standardization +from fastapi import FastAPI, HTTPException +from PIL import Image +from pydantic import BaseModel, ConfigDict, Field +from torch import nn + +# Prepend model_code to sys.path so vendored modules resolve correctly. +_MODEL_CODE_DIR = os.path.join(os.path.dirname(__file__), "model_code") +sys.path.insert(0, _MODEL_CODE_DIR) + +# The vendored transforms/albu.py imports crop from an old albumentations +# path that was removed in v2.x. The function is never actually called +# (IsotropicResize only uses cv2.resize), so inject a harmless stub. +import types as _types + +_compat = _types.ModuleType("albumentations.augmentations.functional") +_compat.crop = None +sys.modules.setdefault("albumentations.augmentations.functional", _compat) + +from models.size_invariant_timesformer import SizeInvariantTimeSformer # noqa: E402 +from models.xception import xception # noqa: E402 +from transforms.albu import IsotropicResize # noqa: E402 + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +MODEL_PORT = int(os.environ.get("MODEL_PORT", 7008)) +PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true" +MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", 1800)) + +# Paths -- extractor and SizeInvariantTimeSformer checkpoints +EXTRACTOR_WEIGHTS_PATH = os.environ.get( + "EXTRACTOR_WEIGHTS_PATH", + "/app/weights/MINTIME/MINTIME_XC_Extractor_checkpoint30", +) +MODEL_WEIGHTS_PATH = os.environ.get( + "MODEL_WEIGHTS_PATH", + "/app/weights/MINTIME/MINTIME_XC_Model_checkpoint30", +) + +# Model hyperparameters from size_invariant_timesformer.yaml +NUM_FRAMES = 16 +IMAGE_SIZE = 224 +NUM_PATCHES = 49 # 7 x 7 spatial from Xception feature map +MAX_IDENTITIES = 2 +RANGE_SIZE = 5 +SIZE_EMB_DICT = [ + (1 + i * RANGE_SIZE, (i + 1) * RANGE_SIZE) if i != 0 else (0, RANGE_SIZE) + for i in range(20) +] + +# Model config dict matching size_invariant_timesformer.yaml +MODEL_CONFIG = { + "model": { + "image-size": IMAGE_SIZE, + "patch-size": 1, + "num-classes": 1, + "num-patches": NUM_PATCHES, + "num-frames": NUM_FRAMES, + "max-identities": MAX_IDENTITIES, + "dim": 512, + "depth": 9, + "dim-head": 64, + "channels": 2048, + "heads": 8, + "attn-dropout": 0.0, + "ff-dropout": 0.0, + "shift-tokens": False, + "enable-size-emb": True, + "enable-pos-emb": True, + "enable-identity-attention": True, + } +} + + +# ── Helpers ────────────────────────────────────────────────────────────── + + +def _generate_connected_components( + similarities: np.ndarray, + similarity_threshold: float = 0.80, +) -> List[List[int]]: + """Build a similarity graph and return connected components.""" + graph = nx.Graph() + n = len(similarities) + for i in range(n): + for j in range(i + 1, n): + if similarities[i, j] > similarity_threshold: + graph.add_edge(i, j) + components = [sorted(c) for c in nx.connected_components(graph)] + # Include isolated nodes (faces not similar to any other) + all_in_components = set() + for c in components: + all_in_components.update(c) + for i in range(n): + if i not in all_in_components: + components.append([i]) + return components + + +def _preprocess_face_for_clustering(img: Image.Image) -> np.ndarray: + """Resize a PIL face crop to 128x128 for embedding extraction.""" + from torchvision.transforms import Resize + + return np.asarray(Resize([128, 128])(img)) + + +# ── Device selection ───────────────────────────────────────────────────── + + +def _get_device() -> torch.device: + """Select optimal device: CUDA > MPS > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + return torch.device("cpu") + + +# ── Global state ───────────────────────────────────────────────────────── + +_extractor: Optional[nn.Module] = None +_model: Optional[nn.Module] = None +_embedding_model: Optional[nn.Module] = None +_device: Optional[torch.device] = None +_load_lock = threading.Lock() + + +def _strip_module_prefix(state_dict: dict) -> dict: + """Remove 'module.' prefix from DataParallel state dicts.""" + new_sd = {} + for k, v in state_dict.items(): + new_key = k.replace("module.", "", 1) if k.startswith("module.") else k + new_sd[new_key] = v + return new_sd + + +def _load_models() -> None: + """Load Xception extractor, SizeInvariantTimeSformer, and + InceptionResnetV1 for identity clustering (thread-safe).""" + global _extractor, _model, _embedding_model, _device + + if _model is not None: + return + + with _load_lock: + if _model is not None: + return + + _device = _get_device() + if _device.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) + else: + logger.warning( + "Device: %s (no CUDA -- inference may be slow)", + _device, + ) + + logger.info("Loading MINTIME models on %s ...", _device) + + # ── Xception feature extractor ────────────────────────────── + if not os.path.exists(EXTRACTOR_WEIGHTS_PATH): + raise FileNotFoundError( + f"Extractor weights not found: {EXTRACTOR_WEIGHTS_PATH}" + ) + feat_ext = xception(num_classes=1, pretrain_path=None) + ext_sd = torch.load( + EXTRACTOR_WEIGHTS_PATH, map_location="cpu", weights_only=False + ) + feat_ext.load_state_dict(_strip_module_prefix(ext_sd)) + feat_ext = feat_ext.to(_device) + feat_ext.train(False) + _extractor = feat_ext + + # ── SizeInvariantTimeSformer ──────────────────────────────── + if not os.path.exists(MODEL_WEIGHTS_PATH): + raise FileNotFoundError(f"Model weights not found: {MODEL_WEIGHTS_PATH}") + sit = SizeInvariantTimeSformer(config=MODEL_CONFIG, require_attention=False) + model_sd = torch.load( + MODEL_WEIGHTS_PATH, map_location="cpu", weights_only=False + ) + sit.load_state_dict(_strip_module_prefix(model_sd)) + sit = sit.to(_device) + sit.train(False) + _model = sit + + # ── InceptionResnetV1 for identity clustering ─────────────── + emb = InceptionResnetV1(pretrained="vggface2").to(_device) + emb.train(False) + _embedding_model = emb + + logger.info("MINTIME models loaded successfully.") + + +def _is_model_loaded() -> bool: + """Return True if all three models are loaded.""" + return ( + _model is not None and _extractor is not None and _embedding_model is not None + ) + + +# ── FastAPI app ────────────────────────────────────────────────────────── + +app = FastAPI( + title="MINTIME Detection Service", + description=( + "Multi-Identity Size-Invariant Video Deepfake Detection " + "(Xception + TimeSformer, IEEE T-IFS 2024)" + ), + version="1.0.0", +) + + +class PredictRequest(BaseModel): + """Incoming prediction request.""" + + video_data: str # Base64-encoded video bytes + threshold: float = 0.5 + + +class PredictResponse(BaseModel): + """Outgoing prediction result.""" + + model_config = ConfigDict(populate_by_name=True) + + model: str = "mintime_detection" + probability: float + prediction: int + class_name: str = Field(..., alias="class") + inference_time: float + metadata: Dict[str, Any] + + +@app.on_event("startup") +async def startup_event(): + """Optionally preload model at startup.""" + if PRELOAD_MODEL: + _load_models() + + +@app.get("/") +def root(): + """Service info endpoint.""" + return { + "service": "mintime_detection", + "port": MODEL_PORT, + "model_loaded": _is_model_loaded(), + "device": str(_device) if _device else "unknown", + } + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and _device is not None and _device.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +def health(): + """Health check endpoint.""" + return { + "status": "healthy", + "model": "mintime_detection", + "device": str(_device) if _device else "cpu", + "model_loaded": _is_model_loaded(), + "extractor_weights_exist": os.path.exists(EXTRACTOR_WEIGHTS_PATH), + "model_weights_exist": os.path.exists(MODEL_WEIGHTS_PATH), + **_gpu_health_info(), + } + + +# ── Video processing utilities ─────────────────────────────────────────── + + +def _extract_frames( + video_path: str, +) -> Tuple[List[np.ndarray], int, int, int]: + """Extract all frames from a video file. + + Returns: + Tuple of (all_frames, fps, width, height). + """ + cap = cv2.VideoCapture(video_path) + fps = max(int(cap.get(cv2.CAP_PROP_FPS)), 1) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + frames: List[np.ndarray] = [] + while True: + ret, frame = cap.read() + if not ret: + break + frames.append(frame) + cap.release() + return frames, fps, width, height + + +def _detect_faces_mtcnn( + frames: List[np.ndarray], + fps: int, +) -> Dict[str, Any]: + """Detect faces in sampled frames using MTCNN. + + Samples one frame per second (every fps frames) and runs MTCNN + on half-resolution PIL images, matching the original preprocessing. + + Returns: + Dict mapping str(frame_index) to list of bboxes, or None. + """ + from facenet_pytorch import MTCNN + + mtcnn = MTCNN( + device=_device, + thresholds=[0.85, 0.95, 0.95], + margin=0, + ) + + bboxes_dict: Dict[str, Any] = {} + # Sample frames at ~1 per second + indices = list(range(0, len(frames), max(fps, 1))) + if not indices: + indices = [0] if frames else [] + + for idx in indices: + frame = frames[idx] + rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + pil_img = Image.fromarray(rgb) + # Half resolution, matching original VideoDataset + pil_img = pil_img.resize([s // 2 for s in pil_img.size]) + boxes, _ = mtcnn.detect(pil_img) + if boxes is not None: + bboxes_dict[str(idx)] = boxes.tolist() + else: + bboxes_dict[str(idx)] = None + + return bboxes_dict + + +def _extract_crops( + frames: List[np.ndarray], + bboxes_dict: Dict[str, Any], + fps: int, +) -> List[Tuple[int, Image.Image, list]]: + """Extract face crops from video frames using detected bboxes. + + Follows the original extract_crops logic: iterate per-second windows, + find the nearest frame with bboxes, crop with padding, make square. + + Returns: + List of (frame_index, pil_crop, bbox). + """ + frames_num = len(frames) + crops = [] + + for i in range(0, frames_num, fps): + # Find nearest frame with valid bboxes in this window + idx = i + limit = min(i + fps - 1, frames_num - 1) + + # Walk forward to find a frame with bboxes + while str(idx) not in bboxes_dict or bboxes_dict.get(str(idx)) is None: + if idx >= limit: + break + idx += 1 + + if str(idx) not in bboxes_dict or bboxes_dict.get(str(idx)) is None: + continue + + bboxes = bboxes_dict[str(idx)] + frame = frames[i] if i < frames_num else frames[-1] + + for bbox in bboxes: + xmin, ymin, xmax, ymax = [int(b * 2) for b in bbox] + w = xmax - xmin + h = ymax - ymin + + if w <= 0 or h <= 0: + continue + + # Padding + p_h = h // 3 + p_w = w // 3 + + crop_h = (ymax + p_h) - max(ymin - p_h, 0) + crop_w = (xmax + p_w) - max(xmin - p_w, 0) + + # Make square + if crop_h > crop_w: + p_h -= int((crop_h - crop_w) / 2) + else: + p_w -= int((crop_w - crop_h) / 2) + + crop = frame[ + max(ymin - p_h, 0) : ymax + p_h, + max(xmin - p_w, 0) : xmax + p_w, + ] + + h_c, w_c = crop.shape[:2] + if h_c <= 0 or w_c <= 0: + continue + + # Final square trim + if h_c > w_c: + diff = (h_c - w_c) // 2 + if diff > 0: + crop = crop[diff:-diff, :] + else: + crop = crop[1:, :] + elif h_c < w_c: + diff = (w_c - h_c) // 2 + if diff > 0: + crop = crop[:, diff:-diff] + else: + crop = crop[:, :-1] + + if crop.size == 0: + continue + + rgb_crop = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB) + crops.append((i, Image.fromarray(rgb_crop), bbox)) + + return crops + + +def _cluster_faces( + crops: List[Tuple[int, Image.Image, list]], + similarity_threshold: float = 0.45, +) -> Dict[int, List]: + """Cluster face crops by identity using InceptionResnetV1 embeddings. + + Returns: + Dict mapping identity_index to list of (frame_idx, pil_img, bbox). + """ + if not crops: + return {} + + crops_images = [row[1] for row in crops] + + # Prepare face tensors for embedding extraction + faces = [_preprocess_face_for_clustering(face) for face in crops_images] + faces = np.stack([np.uint8(f) for f in faces]) + faces_tensor = torch.as_tensor(faces).permute(0, 3, 1, 2).float() + faces_tensor = fixed_image_standardization(faces_tensor) + faces_tensor = faces_tensor.to(_device) + + with torch.no_grad(): + embeddings = _embedding_model(faces_tensor).cpu().numpy() + + # Cosine similarity matrix + similarities = np.dot(embeddings, embeddings.T) + + components = _generate_connected_components( + similarities, similarity_threshold=similarity_threshold + ) + + clustered_faces: Dict[int, List] = {} + for identity_index, component in enumerate(components): + clustered_faces[identity_index] = [crops[fi] for fi in component] + + return clustered_faces + + +def _get_sorted_identities( + identities: Dict[int, List], + num_frames: int = NUM_FRAMES, + max_identities: int = MAX_IDENTITIES, +) -> List[list]: + """Sort identities by face size and allocate frame slots. + + Returns list of [identity_id, mean_side, num_faces, faces_list]. + """ + sorted_ids = [] + for identity in identities: + faces = identities[identity] + mean_side = mean([row[1].size[0] for row in faces]) + sorted_ids.append([identity, mean_side, len(faces), faces]) + + # Sort by face size descending (largest first) + sorted_ids.sort(key=lambda x: x[1], reverse=True) + + if len(sorted_ids) > max_identities: + sorted_ids = sorted_ids[:max_identities] + + identities_number = len(sorted_ids) + available_additional = [] + + if identities_number > 1: + max_faces_map = { + 1: [num_frames], + 2: [num_frames // 2, num_frames // 2], + 3: [num_frames // 3, num_frames // 3, num_frames // 4], + 4: [num_frames // 3, num_frames // 3, num_frames // 8, num_frames // 8], + } + alloc = max_faces_map[identities_number] + + for i in range(identities_number): + if sorted_ids[i][2] < alloc[i] and i < identities_number - 1: + sorted_ids[i + 1][2] += alloc[i] - sorted_ids[i][2] + available_additional.append(0) + elif sorted_ids[i][2] > alloc[i]: + available_additional.append(sorted_ids[i][2] - alloc[i]) + sorted_ids[i][2] = alloc[i] + else: + available_additional.append(0) + else: + sorted_ids[0][2] = num_frames + available_additional.append(0) + + # Fill remaining slots if needed + input_len = sum(r[2] for r in sorted_ids) + if input_len < num_frames: + for i in range(identities_number): + needed = num_frames - input_len + if available_additional[i] > 0: + added = min(available_additional[i], needed) + sorted_ids[i][2] += added + input_len += added + if input_len == num_frames: + break + if input_len < num_frames: + sorted_ids[-1][2] += num_frames - input_len + + return sorted_ids + + +def _create_val_transform(size: int, additional_targets: dict): + """Create the validation-time albumentations transform.""" + from albumentations import Compose, PadIfNeeded, Resize + + return Compose( + [ + IsotropicResize( + max_side=size, + interpolation_down=cv2.INTER_AREA, + interpolation_up=cv2.INTER_CUBIC, + ), + PadIfNeeded( + min_height=size, + min_width=size, + border_mode=cv2.BORDER_CONSTANT, + ), + Resize(height=size, width=size), + ], + additional_targets=additional_targets, + ) + + +def _generate_masks( + video_width: int, + video_height: int, + identities: List[list], + num_frames: int = NUM_FRAMES, + image_size: int = IMAGE_SIZE, + num_patches: int = NUM_PATCHES, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, list]: + """Build input tensors, masks, size embeddings, and positions. + + Mirrors generate_masks from predict.py. + """ + mask = [] + sequence = [] + size_embeddings = [] + images_frames = [] + video_area = video_width * video_height / 2 + + for identity in identities: + max_faces = identity[2] + identity_images = identity[3] + + # Uniform sampling if too many faces + if len(identity_images) > max_faces: + idx = np.round(np.linspace(0, len(identity_images) - 2, max_faces)).astype( + int + ) + identity_images = list(np.asarray(identity_images, dtype=object)[idx]) + + images_frames.extend(img[0] for img in identity_images) + pil_images = [img[1] for img in identity_images] + + # Size embeddings based on face-frame area ratio + identity_size_embs = [] + for img in pil_images: + face_area = img.size[0] * img.size[1] + ratio = int(face_area * 100 / max(video_area, 1)) + side_ranges = list( + map( + lambda a_: ratio in range(a_[0], a_[1] + 1), + SIZE_EMB_DICT, + ) + ) + matches = np.where(side_ranges)[0] + identity_size_embs.append(int(matches[0] + 1) if len(matches) > 0 else 1) + + # Pad with empty frames if needed + if len(pil_images) < max_faces: + diff = max_faces - len(identity_size_embs) + identity_size_embs = list(identity_size_embs) + [0] * diff + pil_images.extend( + [ + np.zeros((image_size, image_size, 3), dtype=np.uint8) + for _ in range(diff) + ] + ) + mask.extend([1 if i < max_faces - diff else 0 for i in range(max_faces)]) + images_frames.extend([max(images_frames)] * diff) + else: + mask.extend([1] * max_faces) + + size_embeddings.extend(identity_size_embs) + sequence.extend(pil_images) + + # Convert PIL images to numpy arrays + sequence = [np.asarray(img) for img in sequence] + + # Apply albumentations transform to all frames together + additional_targets_keys = [f"image{i}" for i in range(num_frames)] + additional_targets_values = ["image"] * num_frames + additional_targets = dict(zip(additional_targets_keys, additional_targets_values)) + transform = _create_val_transform(image_size, additional_targets) + + # Build transform kwargs + transform_kwargs = {"image": sequence[0]} + for i in range(1, len(sequence)): + transform_kwargs[f"image{i}"] = sequence[i] + + transformed = transform(**transform_kwargs) + sequence = [transformed[k] for k in transformed] + + # Build identities_mask + identities_mask = [] + last_range_end = 0 + for identity in identities: + n_faces = identity[2] + identity_mask = [ + last_range_end <= i < last_range_end + n_faces for i in range(num_frames) + ] + for _ in range(n_faces): + identities_mask.append(identity_mask) + last_range_end += n_faces + + # Coherent temporal-positional embedding + images_frames_positions = { + k: v + 1 for v, k in enumerate(sorted(set(images_frames))) + } + frame_positions = [images_frames_positions[f] for f in images_frames] + + if num_patches is not None: + positions = [] + for fp in frame_positions: + positions.extend( + [i + 1 for i in range((fp - 1) * num_patches, num_patches * fp)] + ) + positions.insert(0, 0) # CLS token position + else: + positions = [] + + tokens_per_identity = [] + for i, ident in enumerate(identities): + if i > 0: + tokens_per_identity.append( + (ident[0], ident[2] * num_patches + identities[i - 1][2] * num_patches) + ) + else: + tokens_per_identity.append((ident[0], ident[2] * num_patches)) + + return ( + torch.tensor([sequence]).float(), + torch.tensor([size_embeddings]).int(), + torch.tensor([mask]).bool(), + torch.tensor([identities_mask]).bool(), + torch.tensor([positions]), + tokens_per_identity, + ) + + +# ── Prediction endpoint ───────────────────────────────────────────────── + + +@app.post("/predict", response_model=PredictResponse) +async def predict(request: PredictRequest): + """Run MINTIME deepfake detection on a base64-encoded video. + + Pipeline: + 1. Decode video, extract all frames. + 2. Detect faces per-second with MTCNN. + 3. Crop faces and cluster by identity. + 4. Build identity-ordered input with masks and size embeddings. + 5. Extract Xception features, run TimeSformer. + 6. Return sigmoid probability. + + Returns probability=0.5 if no faces are detected. + """ + if not _is_model_loaded(): + _load_models() + + start_time = time.time() + + # ── Decode video ──────────────────────────────────────────────── + with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp: + try: + video_bytes = base64.b64decode(request.video_data) + tmp.write(video_bytes) + tmp_path = tmp.name + except Exception as e: + raise HTTPException(status_code=400, detail=f"Failed to decode video: {e}") + + try: + # ── Extract frames ────────────────────────────────────────── + frames, fps, vid_w, vid_h = _extract_frames(tmp_path) + if not frames: + raise HTTPException( + status_code=400, + detail="Could not extract frames from video.", + ) + + # ── Detect faces ──────────────────────────────────────────── + bboxes_dict = _detect_faces_mtcnn(frames, fps) + + # Check if any faces found + has_faces = any(v is not None and len(v) > 0 for v in bboxes_dict.values()) + if not has_faces: + return PredictResponse( + probability=0.5, + prediction=0, + class_name="real", + inference_time=time.time() - start_time, + metadata={ + "frames_extracted": len(frames), + "faces_detected": 0, + "identities": 0, + "note": "no faces detected", + "device": str(_device), + }, + ) + + # ── Extract crops ─────────────────────────────────────────── + crops = _extract_crops(frames, bboxes_dict, fps) + if not crops: + return PredictResponse( + probability=0.5, + prediction=0, + class_name="real", + inference_time=time.time() - start_time, + metadata={ + "frames_extracted": len(frames), + "faces_detected": 0, + "identities": 0, + "note": "no valid face crops", + "device": str(_device), + }, + ) + + # ── Cluster by identity ───────────────────────────────────── + clustered = _cluster_faces(crops) + num_identities = len(clustered) + + # ── Build identity sequence ───────────────────────────────── + sorted_identities = _get_sorted_identities(clustered) + + ( + videos_tensor, + size_embeddings, + mask_tensor, + identities_mask, + positions, + tokens_per_identity, + ) = _generate_masks( + vid_w, + vid_h, + sorted_identities, + ) + + # ── Run inference ─────────────────────────────────────────── + b, f, h, w, c = videos_tensor.shape + videos_tensor = videos_tensor.to(_device) + identities_mask = identities_mask.to(_device) + mask_tensor = mask_tensor.to(_device) + positions = positions.to(_device) + + with torch.no_grad(): + # Feature extraction: (B*F, C, H, W) + video_input = rearrange(videos_tensor, "b f h w c -> (b f) c h w") + features = _extractor(video_input) # (B*F, 2048, 7, 7) + features = rearrange(features, "(b f) c h w -> b f c h w", b=b, f=f) + + # TimeSformer classification + pred = _model( + features, + mask=mask_tensor, + size_embedding=size_embeddings, + identities_mask=identities_mask, + positions=positions, + ) + + probability = float(torch.sigmoid(pred[0]).item()) + prediction = 1 if probability >= request.threshold else 0 + class_name = "fake" if prediction == 1 else "real" + + return PredictResponse( + probability=probability, + prediction=prediction, + class_name=class_name, + inference_time=time.time() - start_time, + metadata={ + "frames_extracted": len(frames), + "faces_detected": len(crops), + "identities": num_identities, + "device": str(_device), + }, + ) + + except HTTPException: + raise + except Exception as e: + logger.exception("Error during MINTIME prediction") + raise HTTPException(status_code=500, detail=str(e)) + finally: + if os.path.exists(tmp_path): + os.remove(tmp_path) + gc.collect() + + +if __name__ == "__main__": + uvicorn.run(app, host="0.0.0.0", port=MODEL_PORT) diff --git a/video/mintime/requirements.txt b/video/mintime/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..1fbb7aced56711864adef98fae4165608d97564b --- /dev/null +++ b/video/mintime/requirements.txt @@ -0,0 +1,14 @@ +fastapi +uvicorn +pydantic +python-multipart +torch>=1.8.0 +torchvision>=0.9.0 +facenet-pytorch +opencv-python-headless +numpy<2.0.0 +Pillow +einops +albumentations +networkx +pyyaml diff --git a/video/pwtf-dvd/Dockerfile b/video/pwtf-dvd/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..c1ff8a4d2d404c73a3d00db58389f37fd75fc73f --- /dev/null +++ b/video/pwtf-dvd/Dockerfile @@ -0,0 +1,57 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +WORKDIR /app + +# Install Python 3.10 and system dependencies for OpenCV +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3 python3-pip \ + libgl1 \ + libglib2.0-0 \ + libsm6 \ + libxext6 \ + libxrender-dev \ + && rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +# Install PyTorch with CUDA 12.1 +RUN pip install --no-cache-dir \ + torch==2.5.1 torchvision==0.20.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +# Copy requirements and install +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +# Create logs and weights directories +RUN mkdir -p logs weights model_code + +# Copy model code (read-only reference, never modified) +COPY model_code/ /app/model_code/ + +# Copy weights +COPY weights/ /app/weights/ + +# Copy application code +COPY app.py . + +# Environment variables +ENV MODEL_PORT=7005 +ENV PRELOAD_MODEL=false +ENV MODEL_TIMEOUT=1800 +ENV WEIGHTS_PATH=/app/weights/pwtf_dvd_checkpoint.pth +ENV MODEL_CODE_DIR=/app/model_code + +# Expose port +EXPOSE 7005 + +# Drop root privileges +RUN adduser --disabled-password --gecos '' appuser && \ + chown -R appuser:appuser /app/logs /app/weights +USER appuser + +# Run the service +CMD ["python", "app.py"] diff --git a/video/pwtf-dvd/app.py b/video/pwtf-dvd/app.py new file mode 100644 index 0000000000000000000000000000000000000000..92bb0d6a0a71d80331ac923ae74af59d4acd03e1 --- /dev/null +++ b/video/pwtf-dvd/app.py @@ -0,0 +1,520 @@ +"""PwTF-DVD (Pixel-wise Temporal Frequency) deepfake video detection service. + +Wraps the PwTF-DVD (ICCV 2025) face forgery detection model with a +FastAPI endpoint. Uses a dual-stream architecture: an I3D backbone for +spatial features and a ResNet-based attention network for temporal +frequency features, fused via spatial and temporal transformer encoders. + +The preprocessing pipeline performs RetinaFace detection, SORT-based +tracking, face alignment via 68-point landmarks, and temporal FFT +computation on median-filtered residuals. + +Reference: "Pixel-wise Temporal Frequency Domain Video Deepfake +Detection", ICCV 2025. +""" + +import base64 +import gc +import logging +import os +import platform +import sys +import tempfile +import threading +import time +from typing import Any, Dict, List, Optional + +import cv2 +import numpy as np +import torch +import torch.nn.functional as F +import uvicorn +from fastapi import FastAPI, HTTPException +from PIL import Image +from pydantic import BaseModel, ConfigDict, Field +from torchvision.transforms import Compose, Normalize, ToTensor + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +MODEL_PORT = int(os.environ.get("MODEL_PORT", 7005)) +PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true" +MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", 1800)) +WEIGHTS_PATH = os.environ.get("WEIGHTS_PATH", "/app/weights/pwtf_dvd_checkpoint.pth") +MODEL_CODE_DIR = os.environ.get("MODEL_CODE_DIR", "/app/model_code") + +# PwTF-DVD uses 224x224 face crops after alignment +FACE_CROP_SIZE = 224 +# Clip size for temporal analysis (from root_setting.yaml clip_size: 32) +CLIP_SIZE = 32 +# Maximum frames to extract from the video +MAX_FRAMES = 768 + + +def _get_device() -> torch.device: + """Select optimal device: CUDA (NVIDIA) > MPS (Apple) > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + return torch.device("cpu") + + +# ── Global state ──────────────────────────────────────────────────────────── + +_model: Optional[torch.nn.Module] = None +_device: Optional[torch.device] = None +_load_lock = threading.Lock() + +# Lazy-loaded references to model_code modules +_detect_all = None +_grab_all_frames = None +_get_crop_box = None +_multiple_tracking = None +_find_longest = None +_FasterCropAlignXRay = None +_crop_align_func = None + + +def _ensure_model_code_on_path() -> None: + """Add model_code/inference to sys.path so its internal imports work. + + The model code uses bare imports like ``from model.framework import + get_model`` and ``from config_ftcn import config`` which expect the + ``inference/`` directory to be on ``sys.path``. + """ + inference_dir = os.path.join(MODEL_CODE_DIR, "inference") + if inference_dir not in sys.path: + sys.path.insert(0, inference_dir) + + +def _load_models() -> None: + """Load PwTF-DVD model and face detection tools (thread-safe).""" + global _model, _device + global _detect_all, _grab_all_frames, _get_crop_box + global _multiple_tracking, _find_longest + global _FasterCropAlignXRay, _crop_align_func + + if _model is not None: + return + + with _load_lock: + # Double-check after acquiring lock + if _model is not None: + return + + _device = _get_device() + if _device.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + if _device.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) + else: + logger.warning( + "Device: %s (no CUDA -- check nvidia-container-toolkit)", + _device, + ) + + logger.info("Loading PwTF-DVD model on %s ...", _device) + + # ── Import model_code modules ─────────────────────────────── + _ensure_model_code_on_path() + + from model.framework import get_model + 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 + + _detect_all = detect_all + _grab_all_frames = grab_all_frames + _get_crop_box = get_crop_box + _multiple_tracking = multiple_tracking + _find_longest = find_longest + _FasterCropAlignXRay = FasterCropAlignXRay + + # ── Face crop alignment ───────────────────────────────────── + _crop_align_func = FasterCropAlignXRay(FACE_CROP_SIZE) + + # ── PwTF-DVD classifier ───────────────────────────────────── + if not os.path.exists(WEIGHTS_PATH): + raise FileNotFoundError(f"PwTF-DVD weights not found at {WEIGHTS_PATH}") + + model = get_model() + state_dict = torch.load(WEIGHTS_PATH, map_location="cpu", weights_only=False) + model.load_state_dict(state_dict) + model = model.to(_device) + model.eval() + + _model = model + logger.info("PwTF-DVD model loaded successfully.") + + +def _is_model_loaded() -> bool: + """Return True if the model is loaded and ready.""" + return _model is not None + + +# ── FastAPI app ───────────────────────────────────────────────────────────── + +app = FastAPI( + title="PwTF-DVD Detection Service", + description=( + "Pixel-wise Temporal Frequency Domain Video Deepfake Detection " "(ICCV 2025)" + ), + version="1.0.0", +) + + +class PredictRequest(BaseModel): + """Incoming prediction request.""" + + video_data: str # Base64-encoded video bytes + threshold: float = 0.5 + + +class PredictResponse(BaseModel): + """Outgoing prediction result.""" + + model_config = ConfigDict(populate_by_name=True) + + model: str = "pwtf_dvd_detection" + probability: float + prediction: int + class_name: str = Field(..., alias="class") + inference_time: float + metadata: Dict[str, Any] + + +@app.on_event("startup") +async def startup_event(): + """Optionally preload model at startup.""" + if PRELOAD_MODEL: + _load_models() + + +@app.get("/") +def root(): + """Service info endpoint.""" + return { + "service": "pwtf_dvd_detection", + "port": MODEL_PORT, + "model_loaded": _is_model_loaded(), + "device": str(_device) if _device else "unknown", + } + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and _device is not None and _device.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +def health(): + """Health check endpoint.""" + return { + "status": "healthy", + "model": "pwtf_dvd_detection", + "device": str(_device) if _device else "cpu", + "model_loaded": _is_model_loaded(), + "weights_exist": os.path.exists(WEIGHTS_PATH), + **_gpu_health_info(), + } + + +# ── Inference pipeline ────────────────────────────────────────────────────── + + +def _run_inference(video_path: str) -> dict: + """Run the full PwTF-DVD inference pipeline on a video file. + + Follows the same logic as ``model_code/inference/test_on_raw_video.py``: + 1. Detect faces in all frames (RetinaFace). + 2. Track faces across frames (SORT-based tracker). + 3. Generate sliding-window clips of ``CLIP_SIZE`` frames. + 4. For each clip: align faces, compute temporal FFT residuals, + run dual-stream model. + 5. Aggregate per-clip predictions. + + Args: + video_path: Path to the video file on disk. + + Returns: + Dict with ``probability``, ``num_frames``, ``num_tracks``, + ``num_clips``, ``frames_with_faces``. + """ + # ── Step 1: Detect faces in all frames ────────────────────────── + detect_res, all_lm68, frames = _detect_all( + video_path, return_frames=True, max_size=MAX_FRAMES + ) + + if not frames: + return { + "probability": 0.5, + "num_frames": 0, + "num_tracks": 0, + "num_clips": 0, + "frames_with_faces": 0, + } + + shape = frames[0].shape[:2] + + # Merge 68-landmark data into detection results + all_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_faces.append((box, lm5, face_lm68, score)) + all_detect_res.append(new_faces) + detect_res = all_detect_res + + # ── Step 2: Track faces ───────────────────────────────────────── + tracks = _multiple_tracking(detect_res) + tuples = [(0, len(detect_res))] * len(tracks) + + if len(tracks) == 0: + tuples, tracks = _find_longest(detect_res) + + if len(tracks) == 0: + return { + "probability": 0.5, + "num_frames": len(frames), + "num_tracks": 0, + "num_clips": 0, + "frames_with_faces": 0, + } + + # ── Step 3: Extract face crops and landmarks ──────────────────── + data_storage = {} + frame_boxes = {} + super_clips = [] + + for track_i, ((start, end), track) in enumerate(zip(tuples, tracks)): + 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) + + # ── Step 4: Generate sliding-window clips ─────────────────────── + clips_for_video = [] + 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: + post_module = inner_index[1:-1][::-1] + inner_index + l_post = len(post_module) + if l_post == 0: + continue + post_module = post_module * (pad_length // l_post + 1) + post_module = post_module[:pad_length] + if len(post_module) != pad_length: + continue + + pre_module = inner_index + inner_index[1:-1][::-1] + l_pre = len(pre_module) + if l_pre == 0: + continue + pre_module = pre_module * (pad_length // l_pre + 1) + pre_module = pre_module[-pad_length:] + if len(pre_module) != pad_length: + continue + + inner_index = pre_module + inner_index + post_module + + padded_size = len(inner_index) + frame_range = [ + inner_index[i : i + CLIP_SIZE] + for i in range(padded_size) + if i + CLIP_SIZE <= padded_size + ] + + for indices in frame_range: + clip = [(super_clip_idx, t) for t in indices] + clips_for_video.append(clip) + + if not clips_for_video: + return { + "probability": 0.5, + "num_frames": len(frames), + "num_tracks": len(tracks), + "num_clips": 0, + "frames_with_faces": len(frame_boxes), + } + + # ── Step 5: Run inference on clips ────────────────────────────── + preds = [] + test_transform = Compose( + [ + ToTensor(), + Normalize( + mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225], + ), + ] + ) + + for clip in clips_for_video: + images = [data_storage[f"{i}_{j}_img"] for i, j in clip] + landmarks = [data_storage[f"{i}_{j}_ldm"] for i, j in clip] + + # Align and crop faces + landmarks, images = _crop_align_func(landmarks, images) + + # Build image tensor and temporal frequency features + images_tensor = [] + ft_images = [] + for image in images: + image = np.array(image) + img_pil = Image.fromarray(image) + img_tensor = test_transform(img_pil) + images_tensor.append(img_tensor) + + # Median filter residual -> grayscale for FFT + img_filtered = cv2.medianBlur(image.copy(), 5) + residual = cv2.cvtColor((image - img_filtered), cv2.COLOR_RGB2GRAY) + ft_images.append(residual) + + # Temporal FFT: take first half of frequencies + ft_array = np.array(ft_images) + ft_array = np.absolute( + np.fft.fft(ft_array, axis=0)[: CLIP_SIZE // 2] * (1.0 / CLIP_SIZE) + ) + ft_tensor = torch.from_numpy(ft_array).to(_device).unsqueeze(0) + + # Stack image frames: (1, C, T, H, W) + img_stack = torch.stack(images_tensor, dim=1).unsqueeze(0) + img_stack = img_stack.to(_device) + + with torch.no_grad(): + output = _model(img_stack, ft_tensor) + output = torch.sigmoid(output).squeeze(0) + + pred = float(output.item()) + preds.append(pred) + + # ── Step 6: Aggregate ─────────────────────────────────────────── + probability = float(np.mean(preds)) + + return { + "probability": probability, + "num_frames": len(frames), + "num_tracks": len(tracks), + "num_clips": len(clips_for_video), + "frames_with_faces": len(frame_boxes), + } + + +# ── Prediction endpoint ───────────────────────────────────────────────────── + + +@app.post("/predict", response_model=PredictResponse) +async def predict(request: PredictRequest): + """Run PwTF-DVD deepfake detection on a base64-encoded video. + + Pipeline: + 1. Decode video and write to temp file. + 2. Run full PwTF-DVD pipeline (face detection, tracking, + temporal FFT, dual-stream classification). + 3. Aggregate per-clip predictions into a single probability. + + If no faces are detected the service returns probability=0.5 + (undetermined) rather than raising an error. + """ + if not _is_model_loaded(): + _load_models() + + start_time = time.time() + + # ── Decode video ──────────────────────────────────────────────── + with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp: + try: + video_bytes = base64.b64decode(request.video_data) + tmp.write(video_bytes) + tmp_path = tmp.name + except Exception as exc: + raise HTTPException( + status_code=400, + detail=f"Failed to decode video: {exc}", + ) + + try: + # ── Run inference pipeline ────────────────────────────────── + result = _run_inference(tmp_path) + probability = result["probability"] + prediction = 1 if probability >= request.threshold else 0 + class_name = "fake" if prediction == 1 else "real" + + return PredictResponse( + probability=probability, + prediction=prediction, + class_name=class_name, + inference_time=time.time() - start_time, + metadata={ + "frames_sampled": result["num_frames"], + "frames_with_faces": result["frames_with_faces"], + "num_tracks": result["num_tracks"], + "num_clips": result["num_clips"], + "device": str(_device), + }, + ) + + except HTTPException: + raise + except Exception as exc: + logger.exception("Error during PwTF-DVD prediction") + raise HTTPException(status_code=500, detail=str(exc)) + finally: + if os.path.exists(tmp_path): + os.remove(tmp_path) + gc.collect() + + +if __name__ == "__main__": + uvicorn.run(app, host="0.0.0.0", port=MODEL_PORT) diff --git a/video/pwtf-dvd/requirements.txt b/video/pwtf-dvd/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..6aff79022bb8a120c5d0980ca47676a180b1a9fd --- /dev/null +++ b/video/pwtf-dvd/requirements.txt @@ -0,0 +1,16 @@ +fastapi +uvicorn +pydantic +python-multipart +torch>=2.0.0 +torchvision>=0.15.0 +opencv-python-headless +numpy<2.0.0 +Pillow +einops +fvcore +pyyaml +scipy +filterpy +matplotlib +simplejson diff --git a/video/recce/Dockerfile b/video/recce/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..65fbdfa1a5e7bbd9375142c70e0ed4074ed45b44 --- /dev/null +++ b/video/recce/Dockerfile @@ -0,0 +1,57 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +WORKDIR /app + +# Install Python 3.10 and system dependencies for OpenCV +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3 python3-pip \ + libgl1 \ + libglib2.0-0 \ + libsm6 \ + libxext6 \ + libxrender-dev \ + && rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +# Install PyTorch with CUDA 12.1 +RUN pip install --no-cache-dir \ + torch==2.5.1 torchvision==0.20.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +# Copy requirements and install (facenet-pytorch needs --no-deps due to torch<2.3 pin) +COPY requirements.txt . +RUN pip install --no-cache-dir --no-deps facenet-pytorch && \ + pip install --no-cache-dir -r requirements.txt + +# Create logs and weights directories +RUN mkdir -p logs weights + +# Copy model code (vendored RECCE architecture) +COPY model_code/ /app/model_code/ + +# Copy weights +COPY weights/ /app/weights/ + +# Copy application code +COPY app.py . + +# Environment variables +ENV MODEL_PORT=7007 +ENV PRELOAD_MODEL=false +ENV MODEL_TIMEOUT=1800 +ENV WEIGHTS_PATH=/app/weights/recce_checkpoint.pth + +# Expose port +EXPOSE 7007 + +# Drop root privileges +RUN adduser --disabled-password --gecos '' appuser && \ + chown -R appuser:appuser /app/logs /app/weights +USER appuser + +# Run the service +CMD ["python", "app.py"] diff --git a/video/recce/app.py b/video/recce/app.py new file mode 100644 index 0000000000000000000000000000000000000000..821186001ed4682742423c253b5938bcf59fa681 --- /dev/null +++ b/video/recce/app.py @@ -0,0 +1,472 @@ +"""RECCE deepfake video detection service. + +Wraps the RECCE (CVPR 2022) reconstruction-classification face forgery +detection model with a FastAPI endpoint. Uses an Xception encoder with +guided attention and graph reasoning, trained via DeepfakeBench on +FaceForensics++ (c40). + +The checkpoint originates from DeepfakeBench v1.0.1 which wraps the +original RECCE architecture under a ``model.`` prefix and trains with +2-class (real/fake) output instead of the original repo's 1-class +sigmoid. Weights are loaded by stripping the ``model.`` prefix and +instantiating ``Recce(num_classes=2)``. + +Reference: Cao et al., "End-to-End Reconstruction-Classification +Learning for Face Forgery Detection", CVPR 2022. +""" + +import base64 +import gc +import logging +import os +import platform +import sys +import tempfile +import threading +import time +from typing import Any, Dict, List, Optional, Tuple + +import cv2 +import numpy as np +import torch +import torch.nn.functional as F +import uvicorn +from facenet_pytorch import MTCNN +from fastapi import FastAPI, HTTPException +from PIL import Image +from pydantic import BaseModel, ConfigDict, Field + +# ── Model code import ────────────────────────────────────────────────────── +# The original RECCE model code is vendored under model_code/. We add it +# to sys.path so that ``from model.network import Recce`` resolves. +_MODEL_CODE_DIR = os.path.join(os.path.dirname(__file__), "model_code") +if _MODEL_CODE_DIR not in sys.path: + sys.path.insert(0, _MODEL_CODE_DIR) + +# Patch timm's xception to skip pretrained-weight download. We load our +# own checkpoint, so downloading ImageNet weights wastes bandwidth and +# fails in air-gapped Docker containers. +import timm.models # noqa: E402 + +_original_xception = timm.models.xception + + +def _xception_no_pretrained(**kwargs): + """Force pretrained=False to avoid downloading default weights.""" + kwargs["pretrained"] = False + return _original_xception(**kwargs) + + +timm.models.xception = _xception_no_pretrained + +from model.network import Recce # noqa: E402 + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +MODEL_PORT = int(os.environ.get("MODEL_PORT", 7007)) +PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true" +MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", 1800)) +WEIGHTS_PATH = os.environ.get("WEIGHTS_PATH", "/app/weights/recce_checkpoint.pth") + +# RECCE uses 299x299 face crops (from config/Recce.yml and inference.py) +IMAGE_SIZE = (299, 299) +# Number of frames to sample from each video +NUM_FRAMES = 32 +# Face crop margin factor +MARGIN_FACTOR = 0.5 + + +def _get_device() -> torch.device: + """Select optimal device: CUDA > MPS (Apple) > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + return torch.device("cpu") + + +# ── Global state ─────────────────────────────────────────────────────────── + +_model: Optional[Recce] = None +_face_detector: Optional[MTCNN] = None +_device: Optional[torch.device] = None +_load_lock = threading.Lock() + + +def _load_models() -> None: + """Load RECCE detector and MTCNN face detector (thread-safe).""" + global _model, _face_detector, _device + + if _model is not None: + return + + with _load_lock: + # Double-check after acquiring lock + if _model is not None: + return + + _device = _get_device() + if _device.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + if _device.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) + else: + logger.warning( + "Device: %s (no CUDA available" " -- check nvidia-container-toolkit)", + _device, + ) + logger.info("Loading RECCE model on %s ...", _device) + + # ── Face detector (MTCNN) ────────────────────────────────────── + _face_detector = MTCNN( + keep_all=True, + device=_device, + post_process=False, + ) + + # ── RECCE classifier ─────────────────────────────────────────── + # DeepfakeBench trains RECCE with 2-class output (real, fake) + net = Recce(num_classes=2) + + if not os.path.exists(WEIGHTS_PATH): + raise FileNotFoundError(f"RECCE weights not found at {WEIGHTS_PATH}") + + checkpoint = torch.load(WEIGHTS_PATH, map_location="cpu", weights_only=False) + + # DeepfakeBench wraps the RECCE model under a 'model.' prefix. + # Strip it so the keys match the original Recce class. + if any(k.startswith("model.") for k in checkpoint.keys()): + state_dict = { + k[len("model.") :]: v + for k, v in checkpoint.items() + if k.startswith("model.") + } + else: + state_dict = checkpoint + + net.load_state_dict(state_dict) + net = net.to(_device) + net.eval() + + _model = net + logger.info("RECCE model loaded successfully.") + + +def _is_model_loaded() -> bool: + """Return True if both the classifier and face detector are loaded.""" + return _model is not None and _face_detector is not None + + +# ── FastAPI app ──────────────────────────────────────────────────────────── + +app = FastAPI( + title="RECCE Detection Service", + description=( + "End-to-End Reconstruction-Classification Learning " + "for Face Forgery Detection (Xception backbone, CVPR 2022)" + ), + version="1.0.0", +) + + +class PredictRequest(BaseModel): + """Incoming prediction request.""" + + video_data: str # Base64-encoded video bytes + threshold: float = 0.5 + + +class PredictResponse(BaseModel): + """Outgoing prediction result.""" + + model_config = ConfigDict(populate_by_name=True) + + model: str = "recce_detection" + probability: float + prediction: int + class_name: str = Field(..., alias="class") + inference_time: float + metadata: Dict[str, Any] + + +@app.on_event("startup") +async def startup_event(): + """Optionally preload model at startup.""" + if PRELOAD_MODEL: + _load_models() + + +@app.get("/") +def root(): + """Service info endpoint.""" + return { + "service": "recce_detection", + "port": MODEL_PORT, + "model_loaded": _is_model_loaded(), + "device": str(_device) if _device else "unknown", + } + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and _device is not None and _device.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +def health(): + """Health check endpoint.""" + return { + "status": "healthy", + "model": "recce_detection", + "device": str(_device) if _device else "cpu", + "model_loaded": _is_model_loaded(), + "weights_exist": os.path.exists(WEIGHTS_PATH), + **_gpu_health_info(), + } + + +# ── Video / face utilities ───────────────────────────────────────────────── + + +def _extract_frames(video_path: str, num_frames: int = NUM_FRAMES) -> List[np.ndarray]: + """Uniformly sample *num_frames* RGB frames from a video file. + + Args: + video_path: Path to the video on disk. + num_frames: Number of frames to extract. + + Returns: + List of RGB uint8 numpy arrays (H, W, 3). + """ + cap = cv2.VideoCapture(video_path) + total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + + if total <= 0: + cap.release() + return [] + + indices = np.linspace(0, total - 1, num_frames, endpoint=True, dtype=int) + frames: List[np.ndarray] = [] + + for idx in indices: + cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx)) + ret, frame = cap.read() + if ret: + frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + + cap.release() + return frames + + +def _crop_face( + img: np.ndarray, + bbox: Tuple[float, float, float, float], + margin: float = MARGIN_FACTOR, +) -> np.ndarray: + """Crop a face region from an image with a relative margin. + + Args: + img: RGB image array (H, W, 3). + bbox: (x0, y0, x1, y1) face bounding box. + margin: Fraction of bbox dimension to add as padding. + + Returns: + Cropped face region as a numpy array. + """ + h_img, w_img = img.shape[:2] + x0, y0, x1, y1 = bbox + w = x1 - x0 + h = y1 - y0 + + x0_new = max(0, int(x0 - w * margin / 2)) + x1_new = min(w_img, int(x1 + w * margin / 2) + 1) + y0_new = max(0, int(y0 - h * margin / 2)) + y1_new = min(h_img, int(y1 + h * margin / 2) + 1) + + return img[y0_new:y1_new, x0_new:x1_new] + + +def _detect_and_crop_faces(frame: np.ndarray) -> List[np.ndarray]: + """Detect faces in a single frame and return cropped+resized chips. + + Uses MTCNN for face detection, then crops with margin and resizes + to IMAGE_SIZE (299x299). + + Args: + frame: RGB image array (H, W, 3). + + Returns: + List of face crops resized to IMAGE_SIZE, as uint8 arrays. + """ + assert _face_detector is not None + + pil_img = Image.fromarray(frame) + boxes, _ = _face_detector.detect(pil_img) + + if boxes is None or len(boxes) == 0: + return [] + + crops: List[np.ndarray] = [] + for box in boxes: + x0, y0, x1, y1 = box.tolist() + face = _crop_face(frame, (x0, y0, x1, y1)) + if face.size == 0: + continue + resized = cv2.resize(face, IMAGE_SIZE) + crops.append(resized) + + return crops + + +def _preprocess_face(face_crop: np.ndarray) -> torch.Tensor: + """Apply RECCE-specific preprocessing to a face crop. + + RECCE uses Normalize(mean=[0.5]*3, std=[0.5]*3) which maps + [0, 255] uint8 to [-1, 1] float32. This matches the + albumentations pipeline in the original inference.py. + + Args: + face_crop: RGB uint8 array of shape (299, 299, 3). + + Returns: + Tensor of shape (3, 299, 299) in range [-1, 1]. + """ + # uint8 [0, 255] -> float32 [0, 1] -> normalized [-1, 1] + tensor = torch.tensor(face_crop).permute(2, 0, 1).float().div(255.0) + tensor = (tensor - 0.5) / 0.5 + return tensor + + +# ── Prediction endpoint ──────────────────────────────────────────────────── + + +@app.post("/predict", response_model=PredictResponse) +async def predict(request: PredictRequest): + """Run RECCE face-forgery detection on a base64-encoded video. + + Pipeline: + 1. Decode video and write to temp file. + 2. Extract uniformly-sampled frames. + 3. Detect and crop faces per frame (MTCNN). + 4. Preprocess each face crop (resize 299x299, normalize to [-1,1]). + 5. Classify each face crop with RECCE (2-class softmax). + 6. For each frame, take the max fake probability across faces. + 7. Average the per-frame max probabilities. + + If no faces are detected in any frame the service returns + probability=0.5 (undetermined) rather than raising an error. + """ + if not _is_model_loaded(): + _load_models() + + start_time = time.time() + + # ── Decode video ─────────────────────────────────────────────── + with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp: + try: + video_bytes = base64.b64decode(request.video_data) + tmp.write(video_bytes) + tmp_path = tmp.name + except Exception as e: + raise HTTPException(status_code=400, detail=f"Failed to decode video: {e}") + + try: + # ── Extract frames ───────────────────────────────────────── + frames = _extract_frames(tmp_path, NUM_FRAMES) + if not frames: + raise HTTPException( + status_code=400, + detail="Could not extract frames from video.", + ) + + # ── Detect faces and classify ────────────────────────────── + per_frame_max: List[float] = [] + total_faces = 0 + + for frame in frames: + crops = _detect_and_crop_faces(frame) + if not crops: + continue + + # Preprocess and build batch tensor + tensors = [_preprocess_face(c) for c in crops] + batch_tensor = torch.stack(tensors).to(_device) + + with torch.no_grad(): + logits = _model(batch_tensor) + # RECCE forward() calls squeeze() which drops the + # batch dim when batch=1, producing shape (2,) + # instead of (1, 2). Always ensure 2D. + if logits.dim() == 1: + logits = logits.unsqueeze(0) + probs = F.softmax(logits, dim=1)[:, 1] # fake prob + + frame_max = probs.max().cpu().item() + per_frame_max.append(frame_max) + total_faces += len(crops) + + # ── Aggregate ────────────────────────────────────────────── + if per_frame_max: + probability = float(np.mean(per_frame_max)) + else: + # No faces detected in any frame -- undetermined + probability = 0.5 + + prediction = 1 if probability >= request.threshold else 0 + class_name = "fake" if prediction == 1 else "real" + + return PredictResponse( + probability=probability, + prediction=prediction, + class_name=class_name, + inference_time=time.time() - start_time, + metadata={ + "frames_sampled": len(frames), + "frames_with_faces": len(per_frame_max), + "total_faces_detected": total_faces, + "device": str(_device), + }, + ) + + except HTTPException: + raise + except Exception as e: + logger.exception("Error during RECCE prediction") + raise HTTPException(status_code=500, detail=str(e)) + finally: + if os.path.exists(tmp_path): + os.remove(tmp_path) + gc.collect() + + +if __name__ == "__main__": + uvicorn.run(app, host="0.0.0.0", port=MODEL_PORT) diff --git a/video/recce/requirements.txt b/video/recce/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..f632c804b4cd0d6f82dde683d120aa6bd51c169e --- /dev/null +++ b/video/recce/requirements.txt @@ -0,0 +1,11 @@ +fastapi +uvicorn +pydantic +python-multipart +torch>=1.8.0 +torchvision>=0.9.0 +timm +facenet-pytorch +opencv-python-headless +numpy<2.0.0 +Pillow diff --git a/video/sbi/Dockerfile b/video/sbi/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..e92bd90622129a0f43116ed6bc470a078df16f2f --- /dev/null +++ b/video/sbi/Dockerfile @@ -0,0 +1,52 @@ +FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +WORKDIR /app + +# Install Python 3.10 and system dependencies for OpenCV +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3 python3-pip \ + libgl1 \ + libglib2.0-0 \ + libsm6 \ + libxext6 \ + libxrender-dev \ + && rm -rf /var/lib/apt/lists/* + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +# Install PyTorch with CUDA 12.1 +RUN pip install --no-cache-dir \ + torch==2.5.1 torchvision==0.20.1 \ + --index-url https://download.pytorch.org/whl/cu121 + +# Copy requirements and install (facenet-pytorch needs --no-deps due to torch<2.3 pin) +COPY requirements.txt . +RUN pip install --no-cache-dir --no-deps facenet-pytorch && \ + pip install --no-cache-dir -r requirements.txt + +# Create logs and weights directories +RUN mkdir -p logs weights +COPY weights/ /app/weights/ + +# Copy application code +COPY app.py . + +# Environment variables +ENV MODEL_PORT=7002 +ENV PRELOAD_MODEL=false +ENV MODEL_TIMEOUT=1800 +ENV WEIGHTS_PATH=/app/weights/FFc23.tar + +# Expose port +EXPOSE 7002 + +# Drop root privileges +RUN adduser --disabled-password --gecos '' appuser && \ + chown -R appuser:appuser /app/logs /app/weights +USER appuser + +# Run the service +CMD ["python", "app.py"] diff --git a/video/sbi/app.py b/video/sbi/app.py new file mode 100644 index 0000000000000000000000000000000000000000..c6752930d604b9f670c340bb5cac0462d87e4587 --- /dev/null +++ b/video/sbi/app.py @@ -0,0 +1,430 @@ +"""SBI (Self-Blended Images) deepfake video detection service. + +Wraps the SBI (CVPR 2022) face forgery detection model with a +FastAPI endpoint. Uses EfficientNet-B4 trained on self-blended image +augmentation for robust cross-dataset face-swap detection. + +Reference: Shiohara & Yamasaki, "Detecting Deepfakes with +Self-Blended Images", CVPR 2022. +""" + +import base64 +import gc +import logging +import os +import platform +import sys +import tempfile +import threading +import time +from typing import Any, Dict, List, Optional, Tuple + +import cv2 +import numpy as np +import torch +import torch.nn.functional as F +import uvicorn +from efficientnet_pytorch import EfficientNet +from facenet_pytorch import MTCNN +from fastapi import FastAPI, HTTPException +from PIL import Image +from pydantic import BaseModel, ConfigDict, Field +from torch import nn + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +MODEL_PORT = int(os.environ.get("MODEL_PORT", 7002)) +PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true" +MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", 1800)) +WEIGHTS_PATH = os.environ.get("WEIGHTS_PATH", "/app/weights/FFc23.tar") + +# SBI uses 380x380 face crops (from configs/sbi/base.json) +IMAGE_SIZE = (380, 380) +# Number of frames to sample from each video +NUM_FRAMES = 32 +# Face crop margin factor (matches SBI preprocess.py test-phase behaviour) +MARGIN_FACTOR = 0.5 + + +class Detector(nn.Module): + """EfficientNet-B4 binary classifier (real vs fake). + + Mirrors the inference-time Detector from the SBI repository + (src/inference/model.py). The training-time SAM optimizer and + training_step are intentionally omitted. + """ + + def __init__(self): + super().__init__() + self.net = EfficientNet.from_pretrained( + "efficientnet-b4", advprop=True, num_classes=2 + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass returning 2-class logits.""" + return self.net(x) + + +def _get_device() -> torch.device: + """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU.""" + override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower() + if override == "cpu": + return torch.device("cpu") + if override == "cuda" and torch.cuda.is_available(): + return torch.device("cuda") + if ( + override == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if ( + platform.system() == "Darwin" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return torch.device("mps") + if torch.cuda.is_available(): + return torch.device("cuda") + return torch.device("cpu") + + +# ── Global state ──────────────────────────────────────────────────────────── + +_model: Optional[Detector] = None +_face_detector: Optional[MTCNN] = None +_device: Optional[torch.device] = None +_load_lock = threading.Lock() + + +def _load_models() -> None: + """Load SBI detector and MTCNN face detector (thread-safe).""" + global _model, _face_detector, _device + + if _model is not None: + return + + with _load_lock: + # Double-check after acquiring lock + if _model is not None: + return + + _device = _get_device() + if _device.type == "cuda": + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + if _device.type == "cuda": + logger.info( + "Device: cuda (%s, %.1f GB VRAM)", + torch.cuda.get_device_name(0), + torch.cuda.get_device_properties(0).total_memory / 1024**3, + ) + else: + logger.warning( + "Device: %s (no CUDA available -- check nvidia-container-toolkit)", + _device, + ) + logger.info("Loading SBI model on %s ...", _device) + + # ── Face detector (MTCNN) ─────────────────────────────────────── + _face_detector = MTCNN( + keep_all=True, + device=_device, + post_process=False, + ) + + # ── SBI classifier ────────────────────────────────────────────── + detector = Detector() + + if not os.path.exists(WEIGHTS_PATH): + raise FileNotFoundError(f"SBI weights not found at {WEIGHTS_PATH}") + + checkpoint = torch.load(WEIGHTS_PATH, map_location="cpu", weights_only=False) + state_dict = checkpoint.get("model", checkpoint) + detector.load_state_dict(state_dict) + detector = detector.to(_device) + detector.eval() + + _model = detector + logger.info("SBI model loaded successfully.") + + +def _is_model_loaded() -> bool: + """Return True if both the classifier and face detector are loaded.""" + return _model is not None and _face_detector is not None + + +# ── FastAPI app ───────────────────────────────────────────────────────────── + +app = FastAPI( + title="SBI Detection Service", + description=( + "Detecting Deepfakes with Self-Blended Images " "(EfficientNet-B4, CVPR 2022)" + ), + version="1.0.0", +) + + +class PredictRequest(BaseModel): + """Incoming prediction request.""" + + video_data: str # Base64-encoded video bytes + threshold: float = 0.5 + + +class PredictResponse(BaseModel): + """Outgoing prediction result.""" + + model_config = ConfigDict(populate_by_name=True) + + model: str = "sbi_detection" + probability: float + prediction: int + class_name: str = Field(..., alias="class") + inference_time: float + metadata: Dict[str, Any] + + +@app.on_event("startup") +async def startup_event(): + """Optionally preload model at startup.""" + if PRELOAD_MODEL: + _load_models() + + +@app.get("/") +def root(): + """Service info endpoint.""" + return { + "service": "sbi_detection", + "port": MODEL_PORT, + "model_loaded": _is_model_loaded(), + "device": str(_device) if _device else "unknown", + } + + +def _gpu_health_info() -> dict: + """Return GPU metrics for the health endpoint.""" + if torch.cuda.is_available() and _device is not None and _device.type == "cuda": + return { + "gpu_name": torch.cuda.get_device_name(0), + "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2), + "vram_total_mb": round( + torch.cuda.get_device_properties(0).total_memory / 1024**2 + ), + } + return {} + + +@app.get("/health") +def health(): + """Health check endpoint.""" + return { + "status": "healthy", + "model": "sbi_detection", + "device": str(_device) if _device else "cpu", + "model_loaded": _is_model_loaded(), + "weights_exist": os.path.exists(WEIGHTS_PATH), + **_gpu_health_info(), + } + + +# ── Video / face utilities ────────────────────────────────────────────────── + + +def _extract_frames(video_path: str, num_frames: int = NUM_FRAMES) -> List[np.ndarray]: + """Uniformly sample *num_frames* RGB frames from a video file. + + Args: + video_path: Path to the video on disk. + num_frames: Number of frames to extract. + + Returns: + List of RGB uint8 numpy arrays (H, W, 3). + """ + cap = cv2.VideoCapture(video_path) + total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + + if total <= 0: + cap.release() + return [] + + indices = np.linspace(0, total - 1, num_frames, endpoint=True, dtype=int) + frames: List[np.ndarray] = [] + + for idx in indices: + cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx)) + ret, frame = cap.read() + if ret: + frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + + cap.release() + return frames + + +def _crop_face( + img: np.ndarray, + bbox: Tuple[float, float, float, float], + margin: float = MARGIN_FACTOR, +) -> np.ndarray: + """Crop a face region from an image with a relative margin. + + Mirrors the test-phase crop logic from SBI preprocess.py where + margin multipliers are 0.5 in each direction. + + Args: + img: RGB image array (H, W, 3). + bbox: (x0, y0, x1, y1) face bounding box. + margin: Fraction of bbox dimension to add as padding. + + Returns: + Cropped face region as a numpy array. + """ + h_img, w_img = img.shape[:2] + x0, y0, x1, y1 = bbox + w = x1 - x0 + h = y1 - y0 + + x0_new = max(0, int(x0 - w * margin / 2)) + x1_new = min(w_img, int(x1 + w * margin / 2) + 1) + y0_new = max(0, int(y0 - h * margin / 2)) + y1_new = min(h_img, int(y1 + h * margin / 2) + 1) + + return img[y0_new:y1_new, x0_new:x1_new] + + +def _detect_and_crop_faces( + frame: np.ndarray, +) -> List[np.ndarray]: + """Detect faces in a single frame and return cropped+resized chips. + + Uses MTCNN for face detection, then crops with margin and resizes + to IMAGE_SIZE (380x380). + + Args: + frame: RGB image array (H, W, 3). + + Returns: + List of face crops resized to IMAGE_SIZE, as uint8 arrays. + """ + assert _face_detector is not None + + pil_img = Image.fromarray(frame) + boxes, _ = _face_detector.detect(pil_img) + + if boxes is None or len(boxes) == 0: + return [] + + crops: List[np.ndarray] = [] + for box in boxes: + x0, y0, x1, y1 = box.tolist() + face = _crop_face(frame, (x0, y0, x1, y1)) + if face.size == 0: + continue + resized = cv2.resize(face, IMAGE_SIZE) + crops.append(resized) + + return crops + + +# ── Prediction endpoint ───────────────────────────────────────────────────── + + +@app.post("/predict", response_model=PredictResponse) +async def predict(request: PredictRequest): + """Run SBI face-forgery detection on a base64-encoded video. + + Pipeline: + 1. Decode video and write to temp file. + 2. Extract uniformly-sampled frames. + 3. Detect and crop faces per frame (MTCNN). + 4. Classify each face crop with EfficientNet-B4. + 5. For each frame, take the max probability across faces. + 6. Average the per-frame max probabilities. + + If no faces are detected in any frame the service returns + probability=0.5 (undetermined) rather than raising an error. + """ + if not _is_model_loaded(): + _load_models() + + start_time = time.time() + + # ── Decode video ──────────────────────────────────────────────── + with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp: + try: + video_bytes = base64.b64decode(request.video_data) + tmp.write(video_bytes) + tmp_path = tmp.name + except Exception as e: + raise HTTPException(status_code=400, detail=f"Failed to decode video: {e}") + + try: + # ── Extract frames ────────────────────────────────────────── + frames = _extract_frames(tmp_path, NUM_FRAMES) + if not frames: + raise HTTPException( + status_code=400, + detail="Could not extract frames from video.", + ) + + # ── Detect faces and classify ─────────────────────────────── + per_frame_max: List[float] = [] + total_faces = 0 + + for frame in frames: + crops = _detect_and_crop_faces(frame) + if not crops: + continue + + # Build tensor: (N, C, H, W) float32 in [0, 1] + batch = np.stack(crops, axis=0) # (N, H, W, C) uint8 + batch_tensor = ( + torch.tensor(batch).permute(0, 3, 1, 2).float().div(255.0).to(_device) + ) + + with torch.no_grad(): + logits = _model(batch_tensor) # (N, 2) + probs = F.softmax(logits, dim=1)[:, 1] # fake prob + + frame_max = probs.max().cpu().item() + per_frame_max.append(frame_max) + total_faces += len(crops) + + # ── Aggregate ─────────────────────────────────────────────── + if per_frame_max: + probability = float(np.mean(per_frame_max)) + else: + # No faces detected in any frame — undetermined + probability = 0.5 + + prediction = 1 if probability >= request.threshold else 0 + class_name = "fake" if prediction == 1 else "real" + + return PredictResponse( + probability=probability, + prediction=prediction, + class_name=class_name, + inference_time=time.time() - start_time, + metadata={ + "frames_sampled": len(frames), + "frames_with_faces": len(per_frame_max), + "total_faces_detected": total_faces, + "device": str(_device), + }, + ) + + except HTTPException: + raise + except Exception as e: + logger.exception("Error during SBI prediction") + raise HTTPException(status_code=500, detail=str(e)) + finally: + if os.path.exists(tmp_path): + os.remove(tmp_path) + gc.collect() + + +if __name__ == "__main__": + uvicorn.run(app, host="0.0.0.0", port=MODEL_PORT) diff --git a/video/sbi/requirements.txt b/video/sbi/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..569d05d0e01da28635f9c28405f3041756d164bb --- /dev/null +++ b/video/sbi/requirements.txt @@ -0,0 +1,11 @@ +fastapi +uvicorn +pydantic +python-multipart +torch>=1.8.0 +torchvision>=0.9.0 +efficientnet_pytorch +facenet-pytorch +opencv-python-headless +numpy<2.0.0 +Pillow