Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
| """Faithful ReStraV reproduction (mirrors github.com/ChristianInterno/ReStraV). | |
| Matches the official pipeline exactly, except torchcodec (unavailable here) is replaced by | |
| decord with identical frame selection: | |
| * decode: 24 frames over a CENTER 2s window, repeat_last clamp (dinov2_features.decode_clip) | |
| * preprocess: float/255 + BICUBIC resize to 224, NO ImageNet normalization (dinov2_features.preprocess) | |
| * DINOv2 ViT-S/14 forward_features -> [CLS ; patches] flattened -> Z[T,D] (extract_dinov2_embeddings) | |
| * 21-d features: d[:7], theta[:6], [mu_d,min_d,max_d,var_d, mu_t,min_t,max_t,var_t] (features_from_Z) | |
| * label REAL=1, FAKE=0 (extract_training_features) | |
| * standardize with train mean/std; MLP 21->64->32->1, BCEWithLogits, Adam 1e-3, 10 epochs (train.py) | |
| * report ROC-AUC (and AP for fake=positive) | |
| Runs TWO protocols: | |
| IN-DISTRIBUTION : GenVideo-Val pool, balanced, random 50/50 stratified split (shared generators) | |
| CROSS-GENERATOR : train on d3_repro/extracted generators, test on disjoint GenVideo-Val generators | |
| """ | |
| import os, sys, glob, re, random | |
| os.environ.setdefault("HF_HOME", "/hddstore/ssalehi/SynthScope/.cache/huggingface") | |
| os.environ.setdefault("TORCH_HOME", "/hddstore/ssalehi/SynthScope/.cache/torch") | |
| import numpy as np, torch, torch.nn as nn, torch.nn.functional as F | |
| import torchvision.transforms.functional as TF | |
| from torchvision.transforms import InterpolationMode | |
| import decord | |
| from sklearn.model_selection import train_test_split | |
| from sklearn.metrics import roc_auc_score, average_precision_score | |
| EXTRACTED = "/hddstore/ssalehi/SynthScope/d3_repro/extracted" | |
| GV = "/hddstore/ssalehi/SynthScope/d3_repro/genvideo_val/GenVideo-Val" | |
| DEV = "cuda" | |
| N = int(os.environ.get("RESTRAV_REPRO_N", "500")) # per class, per pool | |
| T, WIN = 24, 2.0 | |
| dinov2 = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14").eval().to(DEV) | |
| def decode_clip(path): | |
| vr = decord.VideoReader(path, ctx=decord.cpu(0)) | |
| n = len(vr); fps = float(vr.get_avg_fps()) or 30.0; dur = n / fps | |
| center = dur / 2.0; start = max(0.0, center - WIN / 2) | |
| delta = WIN / max(T - 1, 1) | |
| idxs = [min(int(round((start + i * delta) * fps)), n - 1) for i in range(T)] # repeat_last clamp | |
| frames = vr.get_batch(idxs).asnumpy() # [T,H,W,3] uint8 | |
| return torch.from_numpy(frames).permute(0, 3, 1, 2) # [T,3,H,W] | |
| def preprocess(frames): | |
| frames = frames.float() / 255.0 | |
| return torch.stack([TF.resize(x, [224, 224], InterpolationMode.BICUBIC, antialias=True) for x in frames]) | |
| def feats_of(path): | |
| clip = preprocess(decode_clip(path)).to(DEV) | |
| f = dinov2.forward_features(clip) | |
| tokens = torch.cat([f["x_norm_clstoken"].unsqueeze(1), f["x_norm_patchtokens"]], dim=1) | |
| Z = tokens.flatten(1).unsqueeze(0) # [1,T,D] | |
| dz = Z[:, 1:] - Z[:, :-1]; d = dz.norm(dim=-1) | |
| cos = F.cosine_similarity(dz[:, :-1], dz[:, 1:], dim=-1) | |
| th = torch.rad2deg(torch.acos(cos.clamp(-1, 1))) | |
| def m4(x): return [x.mean(-1), x.amin(-1), x.amax(-1), x.var(-1, unbiased=False)] | |
| stats = torch.stack(m4(d) + m4(th), dim=1) # [1,8] (mu_d,min,max,var, mu_t,min,max,var) | |
| return torch.cat([d[:, :7], th[:, :6], stats], dim=1)[0].cpu().numpy() # [21] | |
| def build(items, tag): | |
| X, y = [], [] | |
| for i, (p, lab) in enumerate(items): | |
| try: | |
| f = feats_of(p) | |
| if np.isfinite(f).all(): X.append(f); y.append(lab) | |
| except Exception: | |
| pass | |
| if (i + 1) % 200 == 0: print(f" [{tag}] {i+1}/{len(items)}", flush=True) | |
| return np.array(X, np.float32), np.array(y, np.int64) | |
| def train_eval(Xtr, ytr, Xte, yte, tag): | |
| # standardize with TRAIN stats (official) | |
| mu = Xtr.mean(0, keepdims=True); sd = Xtr.std(0, keepdims=True) + 1e-8 | |
| Xtr = (Xtr - mu) / sd; Xte = (Xte - mu) / sd | |
| net = nn.Sequential(nn.Linear(21, 64), nn.ReLU(), nn.Linear(64, 32), nn.ReLU(), nn.Linear(32, 1)).to(DEV) | |
| opt = torch.optim.Adam(net.parameters(), lr=1e-3); crit = nn.BCEWithLogitsLoss() | |
| Xt = torch.tensor(Xtr, device=DEV); yt = torch.tensor(ytr, dtype=torch.float32, device=DEV).unsqueeze(1) | |
| for ep in range(10): | |
| perm = torch.randperm(len(Xt)) | |
| for i in range(0, len(Xt), 32): | |
| b = perm[i:i+32]; opt.zero_grad() | |
| loss = crit(net(Xt[b]), yt[b]); loss.backward(); opt.step() | |
| net.eval() | |
| with torch.no_grad(): | |
| p_real = torch.sigmoid(net(torch.tensor(Xte, device=DEV))).cpu().numpy().ravel() | |
| # label real=1; AUC on real. For fake-detection AP, positive=fake => score 1-p_real, label (1-yte) | |
| auc = roc_auc_score(yte, p_real) | |
| ap_fake = average_precision_score(1 - yte, 1 - p_real) | |
| print(f" {tag}: ROC-AUC={auc:.3f} AP(fake)={ap_fake:.3f} (n_test={len(yte)})", flush=True) | |
| def main(): | |
| random.seed(0) | |
| # GenVideo-Val pool (label real=1, fake=0) | |
| gvr = sorted(glob.glob(GV + "/Real/*.mp4")); gvf = sorted(glob.glob(GV + "/Fake/*/*.mp4")) | |
| random.Random(0).shuffle(gvr); random.Random(1).shuffle(gvf) | |
| print(f"Extracting GenVideo-Val features ({N}/class)...", flush=True) | |
| Xg, yg = build([(p, 1) for p in gvr[:N]] + [(p, 0) for p in gvf[:N]], "GV") | |
| print("\n=== IN-DISTRIBUTION (GenVideo-Val, random 50/50 split, shared generators) ===") | |
| Xtr, Xte, ytr, yte = train_test_split(Xg, yg, test_size=0.5, stratify=yg, random_state=42) | |
| train_eval(Xtr, ytr, Xte, yte, "in-distribution") | |
| # extracted pool (label real=1, fake=0) | |
| ex = sorted(glob.glob(EXTRACTED + "/*.mp4")) | |
| exr = [p for p in ex if re.search(r"_real_", os.path.basename(p))] | |
| exf = [p for p in ex if not re.search(r"_real_", os.path.basename(p))] | |
| random.Random(2).shuffle(exr); random.Random(3).shuffle(exf) | |
| print(f"\nExtracting train-pool (extracted) features ({N}/class)...", flush=True) | |
| Xe, ye = build([(p, 1) for p in exr[:N]] + [(p, 0) for p in exf[:N]], "EX") | |
| print("\n=== CROSS-GENERATOR (train=extracted generators, test=disjoint GenVideo-Val) ===") | |
| train_eval(Xe, ye, Xg, yg, "cross-generator") | |
| if __name__ == "__main__": | |
| main() | |