Update predict.py

predict.py

@@ -15,11 +15,8 @@ import torch.nn as nn
 import torch.nn.functional as F
 import torchvision.transforms as T
 
-import requests  # NEW: for Baseten VLM calls
+import requests
 
-# =========================
-# GLOBAL PATHS (YOU SET)
-# =========================
 
 TRUFOR_TRAIN_TEST_DIR = "TruFor_train_test"
 TRUFOR_CFG_PATH = "TruFor_train_test/lib/config/trufor_ph3.yaml"
@@ -28,19 +25,15 @@ TRUFOR_CKPT_PATH = "weights/trufor.pth.tar"
 UFD_FC_WEIGHTS_PATH = "fc_weights.pth"
 UFD_CLIP_NAME = "ViT-L/14"
 
-# NEW: EfficientNet metric+classifier checkpoint
 EFFNET_CKPT_PATH = "best_metric_cls_effnet.pt"
 
-# Weights for fusion
 W_TRUFOR = 0.5
 W_UFD = 0.4
-W_EFFNET = 0.1  # NEW
+W_EFFNET = 0.1  
 
 IMG_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"}
 
-# =========================
-# NEW: Baseten VLM (your model)
-# =========================
+
 BASETEN_VLM_MODEL_ID = "zq8pe88w"
 BASETEN_VLM_URL = f"https://model-{BASETEN_VLM_MODEL_ID}.api.baseten.co/development/predict"
 
@@ -83,23 +76,16 @@ def get_vlm_reasoning_from_baseten(pil: Image.Image, authenticity_score: float)
     r.raise_for_status()
     out = r.json()
 
-    # Tolerate different response shapes
     if isinstance(out, dict):
-        # Common keys you might return from the Truss model
         for k in ("output", "text", "result", "prediction", "vlm_reasoning"):
             v = out.get(k)
             if isinstance(v, str) and v.strip():
                 return v.strip()
-        # If your Truss returns {"data": "..."} or similar, you'll see it here
         return json.dumps(out, ensure_ascii=False)
 
-    # If it’s a raw string/list/etc.
     return str(out).strip()
 
 
-# =========================
-# UFD CLIPModel
-# =========================
 import clip  # openai/CLIP
 
 CHANNELS = {
@@ -162,10 +148,7 @@ class UniversalFakeDetectDetector:
         return float(torch.sigmoid(logit).item())
 
 
-# =========================
-# NEW: EfficientNet Metric+Classifier
-# =========================
-# Requires: pip install timm
+
 import timm
 
 
@@ -189,9 +172,9 @@ class EffNetMetricClassifier(nn.Module):
 
     def forward(self, x):
         feat = self.backbone(x)
-        z = self.proj(feat)                 # unnormalized projected features
-        emb = F.normalize(z, p=2, dim=1)    # embeddings (not used here, but kept for completeness)
-        logits = self.classifier(z)         # 2-class logits
+        z = self.proj(feat)                 
+        emb = F.normalize(z, p=2, dim=1)    
+        logits = self.classifier(z)         
         return emb, logits
 
 
@@ -225,7 +208,6 @@ class EffNetDetector:
         self.model.to(self.device)
         self.model.eval()
 
-        # Match validation preprocessing from training
         self.transform = T.Compose([
             T.Resize(int(img_size * 1.15)),
             T.CenterCrop(img_size),
@@ -238,19 +220,15 @@ class EffNetDetector:
         x = self.transform(pil.convert("RGB")).unsqueeze(0).to(self.device)
         _, logits = self.model(x)
 
-        # If logits is [B,2], use softmax prob of class 1 (AI)
         if logits.shape[-1] == 2:
             p1 = torch.softmax(logits, dim=1)[0, 1]
             return float(p1.item())
 
-        # fallback (if someone trained 1-logit head): sigmoid
         logit = logits.view(-1)[0]
         return float(torch.sigmoid(logit).item())
 
 
-# =========================
-# TruFor
-# =========================
+
 def _add_trufor_to_syspath():
     if not os.path.isdir(TRUFOR_TRAIN_TEST_DIR):
         raise FileNotFoundError(f"TRUFOR_TRAIN_TEST_DIR not found: {TRUFOR_TRAIN_TEST_DIR}")
@@ -348,9 +326,6 @@ class TruForDetector:
         return TruForOutputs(score=score, loc_prob=loc_prob, conf_prob=conf_prob)
 
 
-# =========================
-# Mask saving + fusion
-# =========================
 def list_images(input_dir: str) -> List[str]:
     paths = []
     for root, _, files in os.walk(input_dir):
@@ -393,14 +368,6 @@ def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--input_dir", required=True)
     ap.add_argument("--output_file", required=True)
-    ap.add_argument("--threshold", type=float, default=0.5)
-    ap.add_argument("--only_flagged", action="store_true")
-
-    ap.add_argument("--mask_dir", default="", help="If set, save TruFor loc/conf maps as PNGs into this folder.")
-    ap.add_argument("--save_conf", action="store_true", help="If set, also save TruFor confidence maps.")
-    ap.add_argument("--print_scores", action="store_true", help="If set, print TruFor/UFD/EffNet/Fused scores per image.")
-
-    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
     args = ap.parse_args()
 
     for p, name in [
@@ -413,10 +380,7 @@ def main():
         if not p or not os.path.exists(p):
             raise FileNotFoundError(f"{name} missing or not found: {p}")
 
-    if args.mask_dir:
-        os.makedirs(args.mask_dir, exist_ok=True)
-
-    device = torch.device(args.device)
+    device = torch.device("cuda")
 
     trufor = TruForDetector(device=device)
     ufd = UniversalFakeDetectDetector(device=device)
@@ -425,32 +389,17 @@ def main():
     preds: List[Dict[str, Any]] = []
     for img_path in list_images(args.input_dir):
         img_name = os.path.basename(img_path)
-        stem = os.path.splitext(img_name)[0]
         pil = Image.open(img_path)
 
         tru = trufor.predict(pil)
         ufd_prob = ufd.predict_prob(pil)
-        eff_prob = effnet.predict_prob(pil)  # NEW
+        eff_prob = effnet.predict_prob(pil)
 
         fused = fuse_scores(tru.score, ufd_prob, eff_prob)
 
-        if args.print_scores:
-            print(
-                f"{img_name}\tTruFor={tru.score:.4f}\tUFD={ufd_prob:.4f}\tEffNet={eff_prob:.4f}\tFused={fused:.4f}",
-                flush=True,
-            )
-
-        if args.mask_dir:
-            loc_path = os.path.join(args.mask_dir, f"{stem}_trufor_loc.png")
-            save_prob_map_png(tru.loc_prob, loc_path)
-
-            if args.save_conf:
-                conf_path = os.path.join(args.mask_dir, f"{stem}_trufor_conf.png")
-                save_prob_map_png(tru.conf_prob, conf_path)
-
-        # NEW: VLM reasoning from Baseten
         if fused < 0.5:
             vlm_reasoning = "It looks natural."
+            continue
         else:
             try:
                 vlm_reasoning = get_vlm_reasoning_from_baseten(pil, fused)
@@ -462,15 +411,9 @@ def main():
             "authenticity_score": float(fused),
             "manipulation_type": manipulation_type_from_maps(tru, ufd_prob, fused),
             "vlm_reasoning": vlm_reasoning,
-            "debug": {
-                "trufor_score": float(tru.score),
-                "ufd_score": float(ufd_prob),
-                "effnet_score": float(eff_prob),
-            },
         }
 
-        if (not args.only_flagged) or (fused >= args.threshold):
-            preds.append(rec)
+        preds.append(rec)
 
     with open(args.output_file, "w", encoding="utf-8") as f:
         json.dump(preds, f, indent=2)