| """Extract CLIP features for AF TEST videos (separate from train extraction). |
| |
| Test-set verifier features are needed ONLY at inference time for the |
| verifier-as-context experiment (Phase 1). The verifier itself was trained |
| without test labels (no leakage); this step just runs the trained verifier on |
| test data so the policy model can see per-second forgery scores as context. |
| """ |
| import argparse |
| import os |
| import sys |
| import time |
|
|
| import torch |
| import torch.nn.functional as F |
| from decord import VideoReader, cpu |
| from PIL import Image |
| from transformers import CLIPModel, CLIPProcessor |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from src.open_r1.data_loader import TEST_GENERATORS, build_examples |
| from src.open_r1.forgery_head import frame_labels_from_segments |
|
|
| ANNOT = "/mnt/local-fast/zhangt/annot/annot" |
| VROOT = "/mnt/local-fast/zhangt/video" |
| CACHE = "/mnt/local-fast/zhangt/forensics_verifier_clip_l14" |
| MODEL_ID = "openai/clip-vit-large-patch14" |
| SAMPLE_FPS = 1.0 |
|
|
|
|
| def decode_at_1fps(video_path: str, duration: float): |
| vr = VideoReader(video_path, ctx=cpu(0)) |
| fps_video = vr.get_avg_fps() |
| n_secs = max(1, int(duration)) |
| idxs = [min(int(s * fps_video), len(vr) - 1) for s in range(n_secs)] |
| frames = vr.get_batch(idxs).asnumpy() |
| return [Image.fromarray(f) for f in frames], n_secs |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--rank", type=int, default=0) |
| ap.add_argument("--world_size", type=int, default=1) |
| ap.add_argument("--device", type=int, default=0) |
| ap.add_argument("--batch", type=int, default=32) |
| args = ap.parse_args() |
|
|
| device = f"cuda:{args.device}" |
| print(f"[rank {args.rank}/{args.world_size}] device={device}", flush=True) |
|
|
| print("loading CLIP ...", flush=True) |
| t0 = time.time() |
| clip = CLIPModel.from_pretrained(MODEL_ID, torch_dtype=torch.float32).to(device).eval() |
| proc = CLIPProcessor.from_pretrained(MODEL_ID) |
| print(f" loaded in {time.time()-t0:.1f}s", flush=True) |
|
|
| split = "test" |
| examples = build_examples( |
| annot_dir=ANNOT, video_root=VROOT, generators=TEST_GENERATORS, |
| split_prefix=split, preprocessed_data_path=None, require_video_exists=True, |
| ) |
| examples = [ex for i, ex in enumerate(examples) if i % args.world_size == args.rank] |
| print(f"[{split}] rank={args.rank} processing {len(examples)} videos", flush=True) |
|
|
| t_start = time.time() |
| done = skipped = failed = 0 |
| for ex in examples: |
| sample_id = os.path.splitext(os.path.basename(ex["video_path"]))[0] |
| out_dir = os.path.join(CACHE, split, ex["generator"], sample_id) |
| feat_path = os.path.join(out_dir, "clip_feats.pt") |
| if os.path.exists(feat_path): |
| skipped += 1 |
| continue |
| try: |
| pil_imgs, n_secs = decode_at_1fps(ex["video_path"], ex["durations"]) |
| except Exception: |
| failed += 1 |
| continue |
| feats_all = [] |
| for i in range(0, len(pil_imgs), args.batch): |
| chunk = pil_imgs[i:i + args.batch] |
| with torch.no_grad(): |
| inputs = proc(images=chunk, return_tensors="pt").to(device) |
| f = clip.get_image_features(**inputs) |
| f = F.normalize(f, dim=-1) |
| feats_all.append(f.cpu()) |
| feats = torch.cat(feats_all, dim=0).float() |
| labels = frame_labels_from_segments( |
| ex["solution"], n_secs, fps_to_groups=SAMPLE_FPS |
| ).float() |
| os.makedirs(out_dir, exist_ok=True) |
| torch.save(feats, feat_path) |
| torch.save(labels, os.path.join(out_dir, "clip_labels.pt")) |
| done += 1 |
| if done % 30 == 0: |
| elapsed = time.time() - t_start |
| rate = done / max(1e-6, elapsed) |
| remaining = (len(examples) - done - skipped) / max(1e-6, rate) |
| print(f" rank={args.rank} done={done} skipped={skipped} failed={failed} " |
| f"rate={rate:.2f}/s eta={remaining/60:.1f}min", flush=True) |
|
|
| print(f"rank={args.rank} DONE done={done} skipped={skipped} failed={failed}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
