File size: 3,338 Bytes
df66507 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | """Single-image inference using the unmodified official DoUNet."""
import argparse
import json
import sys
from pathlib import Path
def main():
p = argparse.ArgumentParser()
source = p.add_mutually_exclusive_group(required=True)
source.add_argument('--repo-id')
source.add_argument('--ckpt', type=Path)
p.add_argument('--revision', default='main', help='Prefer the uploaded commit SHA')
p.add_argument('--image', required=True, type=Path)
p.add_argument('--device', default='auto')
p.add_argument('--seed', type=int, default=42)
p.add_argument('--disable-amp', action='store_true')
p.add_argument('--trust-checkpoint', action='store_true', help='Allow pickle loading for a trusted official checkpoint only')
a = p.parse_args()
import torch
from PIL import Image
from huggingface_hub import hf_hub_download
sys.path.insert(0, str(Path(__file__).resolve().parent / 'upstream'))
from models.dou import DoUNet, load_clip_to_cpu
from utils.options import TrainCfg
from dataset.dataset import translate_duplicate
from torchvision import transforms
torch.manual_seed(a.seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(a.seed)
device = torch.device(('cuda' if torch.cuda.is_available() else 'cpu') if a.device == 'auto' else a.device)
path = a.ckpt or Path(hf_hub_download(a.repo_id, 'best.pth', revision=a.revision))
cfg = TrainCfg(model_name='ViT-L/14', classnames=('real', 'generated'), learn_logit_scale=False, amp=not a.disable_amp)
backbone = load_clip_to_cpu(cfg).eval()
for param in backbone.parameters():
param.requires_grad_(False)
model = DoUNet(cfg, backbone).to(device)
ckpt = torch.load(path, map_location='cpu', weights_only=False)
if not isinstance(ckpt, dict) or not isinstance(ckpt.get('model'), dict):
raise ValueError("Expected a checkpoint dictionary containing ckpt['model']")
# Upstream uses strict=False silently. Reject missing/unexpected tensors here
# to avoid presenting an incompletely loaded detector as a valid result.
report = model.load_state_dict(ckpt['model'], strict=False)
if report.missing_keys or report.unexpected_keys:
raise RuntimeError(f'Checkpoint mismatch: {report}. Resolve against the official release; do not ignore it.')
model.eval()
# Obtain the exact official evaluation transform without scanning a dataset.
transform = transforms.Compose([
transforms.Lambda(
lambda img: translate_duplicate(img, 224)),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(
[0.48145466, 0.4578275, 0.40821073], [0.26862954, 0.26130258, 0.27577711]),
])
with Image.open(a.image) as img:
x = transform(img.convert('RGB')).unsqueeze(0).to(device)
with torch.inference_mode(), torch.amp.autocast('cuda', enabled=cfg.amp and device.type == 'cuda'):
probs = model(x).softmax(dim=-1)[0].float().cpu()
print(json.dumps({'label': ['real', 'generated'][probs.argmax().item()],
'real_score': probs[0].item(), 'generated_score': probs[1].item(),
'seed': a.seed, 'note': 'Softmax scores are not calibrated probabilities of authenticity.'}))
if __name__ == '__main__':
main()
|