Spaces:
Sleeping
Sleeping
| """DeepShield adapter for MDMF's DINOv2 embeddings and MMD² score.""" | |
| import json | |
| import math | |
| import pickle | |
| import sys | |
| import time | |
| from pathlib import Path | |
| HERE = Path(__file__).resolve().parent | |
| ROOT = HERE.parents[1] | |
| sys.path[:0] = [str(ROOT), str(HERE / "upstream" / "src")] | |
| import torch | |
| from PIL import Image | |
| from torchvision import transforms | |
| from transformers import AutoModel | |
| from support.recent_common import checkpoint, detection_result, device_for, parser, timed_score, validate_inputs, write_results | |
| from support.recent_data import samples | |
| from test_MDMF import MMD_MIL_Detector, pool_tokens_to_patch_size, precompute_ref_stats, score_mmd2_batch | |
| def backbone_for(device): | |
| location = HERE / "checkpoint" / "backbones" / "dinov2-large" | |
| if not (location / "config.json").is_file(): | |
| raise FileNotFoundError(f"DINOv2 backbone missing: {location}; run support/setup_recent_weights.py MDMF") | |
| return AutoModel.from_pretrained(location).to(device).eval() | |
| TRANSFORM = transforms.Compose([ | |
| transforms.Resize((224, 224)), transforms.ToTensor(), | |
| transforms.Normalize([.485, .456, .406], [.229, .224, .225]), | |
| ]) | |
| def embedding_batch(backbone, device, paths, patch_size=32): | |
| images = [] | |
| for path in paths: | |
| with Image.open(path) as image: | |
| images.append(TRANSFORM(image.convert("RGB"))) | |
| tensor = torch.stack(images).to(device) | |
| tokens = backbone(tensor).last_hidden_state[:, 1:, :].float() | |
| return pool_tokens_to_patch_size(tokens, patch_size, 224).cpu().numpy() | |
| def embedding(backbone, device, path, patch_size=32): | |
| return embedding_batch(backbone, device, [path], patch_size)[0] | |
| def embedding_bank(backbone, device, entries, split, target, batch_size=4): | |
| vectors = [] | |
| for offset in range(0, len(entries), batch_size): | |
| vectors.extend(embedding_batch(backbone, device, | |
| [path for path, _ in entries[offset:offset + batch_size]])) | |
| data = {"embeddings": vectors, "labels": [label for _, label in entries], | |
| "image_paths": [str(path) for path, _ in entries], "patch_size": 32, | |
| "split": split, "embed_dim": 1024} | |
| target.parent.mkdir(parents=True, exist_ok=True) | |
| with target.open("wb") as handle: | |
| pickle.dump(data, handle, protocol=4) | |
| return data | |
| def model_for(path, device): | |
| state = torch.load(path, map_location="cpu", weights_only=False) | |
| config = state.get("config", {}) | |
| model = MMD_MIL_Detector(config.get("embed_dim", 1024), config.get("hidden_dim", 256), | |
| config.get("dropout", .3), config.get("init_sigma", 1.)).to(device) | |
| model.load_state_dict(state["model_state_dict"]) | |
| return model.eval(), int(config.get("patch_size", 32)) | |
| def reference_for(model, device, path): | |
| if not path.is_file(): | |
| raise FileNotFoundError(f"MDMF real training reference bank missing: {path}; run detectors/MDMF/prepare.py") | |
| with path.open("rb") as handle: | |
| bank = pickle.load(handle) | |
| if bank.get("split") != "train" or any(bank["labels"]): | |
| raise ValueError("MDMF reference bank must contain training-split real images only") | |
| if not bank["embeddings"]: | |
| raise ValueError("MDMF reference bank is empty") | |
| with torch.inference_mode(): | |
| logits = torch.stack([model.compute_logits(torch.from_numpy(x).float().unsqueeze(0).to(device)).squeeze(0) | |
| for x in bank["embeddings"]]) | |
| stats = precompute_ref_stats(logits, model.get_sigma()) | |
| return logits, stats | |
| def scorer(backbone, model, device, patch_size, reference, calibration): | |
| ref, (xx_mean, norm2) = reference | |
| threshold, scale = calibration["threshold"], calibration["scale"] | |
| def score_batch(paths): | |
| vectors = torch.from_numpy(embedding_batch(backbone, device, paths, patch_size)).float().to(device) | |
| logits = model.compute_logits(vectors) | |
| raws = score_mmd2_batch(ref, norm2, xx_mean, logits, model.get_sigma()) | |
| fakes = torch.sigmoid((raws - threshold) / scale) | |
| return list(zip(raws.cpu().tolist(), fakes.cpu().tolist())) | |
| def score(path): | |
| return score_batch([path])[0] | |
| score.batch = score_batch | |
| return score | |
| def calibrate(backbone, model, device, patch_size, reference, entries, path): | |
| # The released MDMF checkpoint contains no binary decision threshold. | |
| neutral = {"threshold": 0., "scale": 1.} | |
| score = scorer(backbone, model, device, patch_size, reference, neutral) | |
| real = [score(p)[0] for p, label in entries if label == 0] | |
| fake = [score(p)[0] for p, label in entries if label == 1] | |
| if not real or not fake: | |
| raise ValueError("MDMF calibration needs both real and fake validation images") | |
| import statistics | |
| a, b = statistics.median(real), statistics.median(fake) | |
| if b <= a: | |
| raise ValueError("MDMF validation scores do not separate in the expected direction; inspect data and checkpoint") | |
| result = {"threshold": (a + b) / 2, "scale": max((b - a) / 4, 1e-6), | |
| "source": "validation split", "real_count": len(real), "fake_count": len(fake)} | |
| path.write_text(json.dumps(result, indent=2) + "\n") | |
| return result | |
| def main(task): | |
| args = parser("MDMF", task).parse_args() | |
| validate_inputs(args, task) | |
| device = device_for(args.device) | |
| name = args.model or args.name if task == "detect" else args.name | |
| directory = HERE / "checkpoint" / name / "weights" | |
| backbone = backbone_for(device) | |
| if task == "train": | |
| train = samples(args.data_root, args.split_file, "train", args.data_keys) | |
| val = samples(args.data_root, args.split_file, "val", args.data_keys) | |
| real = [(p, y) for p, y in train if y == 0] | |
| if not real or not any(y == 1 for _, y in train): | |
| raise ValueError("MDMF training requires real and fake training images") | |
| directory.mkdir(parents=True, exist_ok=True) | |
| train_bank = embedding_bank(backbone, device, train, "train", directory / "train.pkl", args.batch_size) | |
| embedding_bank(backbone, device, real, "train", directory / "reference.pkl", args.batch_size) | |
| model = MMD_MIL_Detector().to(device) | |
| optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) | |
| real_vec = [torch.from_numpy(x).float() for x, y in zip(train_bank["embeddings"], train_bank["labels"]) if y == 0] | |
| fake_vec = [torch.from_numpy(x).float() for x, y in zip(train_bank["embeddings"], train_bank["labels"]) if y == 1] | |
| for epoch in range(args.epochs): | |
| model.train() | |
| half_batch = max(1, args.batch_size // 2) | |
| steps = math.ceil(max(len(real_vec), len(fake_vec)) / half_batch) | |
| if args.max_steps: | |
| steps = min(steps, args.max_steps) | |
| for step in range(steps): | |
| a = torch.stack([real_vec[(step * half_batch + i) % len(real_vec)] | |
| for i in range(half_batch)]).to(device) | |
| b = torch.stack([fake_vec[(step * half_batch + i) % len(fake_vec)] | |
| for i in range(half_batch)]).to(device) | |
| optimizer.zero_grad(set_to_none=True) | |
| loss, mmd, _ = model(a, b) | |
| loss.backward() | |
| optimizer.step() | |
| print(f"epoch={epoch + 1} step={step + 1} loss={loss.item():.5f} mmd2={mmd.item():.5f}") | |
| torch.save({"model_state_dict": model.state_dict(), "config": {"embed_dim": 1024, "hidden_dim": 256, | |
| "dropout": .3, "init_sigma": 1., "patch_size": 32}}, directory / "best.pt") | |
| model.eval() | |
| reference = reference_for(model, device, directory / "reference.pkl") | |
| calibrate(backbone, model, device, 32, reference, val, directory / "calibration.json") | |
| return | |
| model, patch_size = model_for(checkpoint("MDMF", name), device) | |
| reference = reference_for(model, device, directory / "reference.pkl") | |
| calibration_path = directory / "calibration.json" | |
| if not calibration_path.is_file(): | |
| raise FileNotFoundError(f"MDMF calibration missing: {calibration_path}; run detectors/MDMF/prepare.py") | |
| score = scorer(backbone, model, device, patch_size, reference, json.loads(calibration_path.read_text())) | |
| if task == "detect": | |
| image = Path(args.image) | |
| if not image.is_file(): | |
| raise FileNotFoundError(image) | |
| raw, fake, elapsed = timed_score(score, image) | |
| detection_result("MDMF", image, raw, fake, elapsed, args.output) | |
| else: | |
| rows = [] | |
| entries = samples(args.data_root, args.split_file, "test", args.data_keys) | |
| for offset in range(0, len(entries), args.batch_size): | |
| chunk = entries[offset:offset + args.batch_size] | |
| start = time.perf_counter() | |
| scores = score.batch([path for path, _ in chunk]) | |
| elapsed = (time.perf_counter() - start) / len(chunk) | |
| for (path, label), (raw, fake) in zip(chunk, scores): | |
| rows.append({"path": str(path), "label": label, "raw_score": raw, | |
| "fake_score": fake, "prediction": int(fake >= .5), "elapsed_time": elapsed}) | |
| write_results("MDMF", name, args.data_keys, rows) | |