Enhance AI detector initialization and improve TruFor result handling

app.py

@@ -1,23 +1,54 @@
+import os
+
 import gradio as gr
 import torch
 from PIL import Image
 from typing import Dict, Optional, Tuple
-from transformers import AutoImageProcessor, SiglipForImageClassification
+from torchvision import transforms
+from timm import create_model
+from huggingface_hub import hf_hub_download
+from huggingface_hub.errors import GatedRepoError, HfHubHTTPError
 
 from trufor_runner import TruForEngine, TruForResult, TruForUnavailableError
 
-MODEL_ID = "Ateeqq/ai-vs-human-image-detector"
+IMG_SIZE = 380
+LABEL_MAPPING = {0: "human", 1: "ai"}
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Use GPU when available so large batches stay responsive.
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+transform: Optional[transforms.Compose]
+model: Optional[torch.nn.Module]
+MODEL_STATUS: str
 
 try:
-    processor = AutoImageProcessor.from_pretrained(MODEL_ID)
-    model = SiglipForImageClassification.from_pretrained(MODEL_ID)
-    model.to(device)
+    token = os.getenv("HF_TOKEN")
+    model_path = hf_hub_download(repo_id="Dafilab/ai-image-detector", filename="pytorch_model.pth", token=token)
+    transform = transforms.Compose([
+        transforms.Resize(IMG_SIZE + 20),
+        transforms.CenterCrop(IMG_SIZE),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    model = create_model("efficientnet_b4", pretrained=False, num_classes=len(LABEL_MAPPING))
+    model.load_state_dict(torch.load(model_path, map_location=DEVICE))
+    model.to(DEVICE)
     model.eval()
+    MODEL_STATUS = "AI detector ready."
+except GatedRepoError:
+    transform = None
+    model = None
+    MODEL_STATUS = (
+        "AI detector requires approved Hugging Face access. Configure HF_TOKEN with a permitted token."
+    )
+except (HfHubHTTPError, OSError) as exc:
+    transform = None
+    model = None
+    MODEL_STATUS = f"AI detector unavailable: {exc}"
 except Exception as exc:  # pragma: no cover - surface loading issues early.
-    raise RuntimeError(f"Failed to load model from {MODEL_ID}") from exc
+    transform = None
+    model = None
+    MODEL_STATUS = f"AI detector failed to initialize: {exc}"
+
+AI_INTRO_SUMMARY = MODEL_STATUS if model is None else "Upload an image to view the prediction."
 
 try:
     TRUFOR_ENGINE: Optional[TruForEngine] = TruForEngine(device="cpu")
@@ -28,64 +59,72 @@ except TruForUnavailableError as exc:
 
 
 def analyze_ai_vs_human(image: Image.Image) -> Tuple[Dict[str, float], str]:
-    """Run the Hugging Face detector and return confidences with a readable summary."""
+    """Run the EfficientNet-based detector and return confidences with a readable summary."""
+    empty_scores = {label: 0.0 for label in LABEL_MAPPING.values()}
+
+    if model is None or transform is None:
+        return empty_scores, MODEL_STATUS
+
     if image is None:
-        empty_scores = {label: 0.0 for label in model.config.id2label.values()}
         return empty_scores, "No image provided."
 
     image = image.convert("RGB")
-    inputs = processor(images=image, return_tensors="pt").to(device)
+    inputs = transform(image).unsqueeze(0).to(DEVICE)
 
     with torch.no_grad():
-        logits = model(**inputs).logits
-
-    probabilities = torch.softmax(logits, dim=-1)[0]
-    scores = {
-        model.config.id2label[idx]: float(probabilities[idx])
-        for idx in range(probabilities.size(0))
-    }
+        logits = model(inputs)
 
-    top_idx = int(probabilities.argmax().item())
-    top_label = model.config.id2label[top_idx]
-    top_score = scores[top_label]
-    summary = f"**Predicted Label:** {top_label}  \
-**Confidence:** {top_score:.4f}"
+    probabilities = torch.softmax(logits, dim=1)[0]
+    ordered_scores = sorted(
+        ((LABEL_MAPPING[idx], float(probabilities[idx])) for idx in LABEL_MAPPING),
+        key=lambda item: item[1],
+        reverse=True,
+    )
+    scores = dict(ordered_scores)
+
+    top_label, top_score = ordered_scores[0]
+    second_label, second_score = ordered_scores[1]
+    summary = (
+        f"**Predicted Label:** {top_label}  \
+**Confidence:** {top_score:.2%}\n"
+        f"`{top_label}: {top_score:.2%} | {second_label}: {second_score:.2%}`"
+    )
 
     return scores, summary
 
 
-def analyze_trufor(image: Image.Image) -> Tuple[str, Optional[Image.Image], Optional[Image.Image]]:
+def analyze_trufor(image: Image.Image) -> Tuple[str, Optional[Image.Image]]:
     """Run TruFor inference when available, otherwise return diagnostics."""
     if TRUFOR_ENGINE is None:
-        return TRUFOR_STATUS, None, None
+        return TRUFOR_STATUS, None
 
     if image is None:
-        return "Upload an image to run TruFor.", None, None
+        return "Upload an image to run TruFor.", None
 
     try:
         result: TruForResult = TRUFOR_ENGINE.infer(image)
     except TruForUnavailableError as exc:
-        return str(exc), None, None
+        return str(exc), None
 
-    summary_lines = []
-    if result.score is not None:
-        summary_lines.append(f"**Tamper Score:** {result.score:.4f}")
-    extras_dict = result.raw_scores.copy()
-    if result.score is not None:
-        extras_dict.pop("tamper_score", None)
-    if extras_dict:
-        extras = "  ".join(f"{key}: {value:.4f}" for key, value in extras_dict.items())
-        summary_lines.append(f"`{extras}`")
-    if not summary_lines:
-        summary_lines.append("TruFor returned no scores for this image.")
+    # Determine if image is altered based on tamper score threshold
+    if result.score is None:
+        return "TruFor returned no prediction for this image.", result.map_overlay
+    
+    # Threshold for altered vs not altered (adjust as needed)
+    threshold = 0.5
+    is_altered = result.score >= threshold
+    prediction = "Altered" if is_altered else "Not Altered"
+    confidence = result.score if is_altered else (1.0 - result.score)
+    
+    summary = f"**Prediction:** {prediction}\n**Confidence:** {confidence:.2%}"
 
-    return "\n".join(summary_lines), result.map_overlay, result.confidence_overlay
+    return summary, result.map_overlay
 
 
-def analyze_image(image: Image.Image) -> Tuple[Dict[str, float], str, str, Optional[Image.Image], Optional[Image.Image]]:
+def analyze_image(image: Image.Image) -> Tuple[Dict[str, float], str, str, Optional[Image.Image]]:
     ai_scores, ai_summary = analyze_ai_vs_human(image)
-    trufor_summary, tamper_overlay, conf_overlay = analyze_trufor(image)
-    return ai_scores, ai_summary, trufor_summary, tamper_overlay, conf_overlay
+    trufor_summary, tamper_overlay = analyze_trufor(image)
+    return ai_scores, ai_summary, trufor_summary, tamper_overlay
 
 
 with gr.Blocks() as demo:
@@ -93,7 +132,7 @@ with gr.Blocks() as demo:
         """# Image Authenticity Workbench\nUpload an image to compare the AI-vs-human classifier with the TruFor forgery detector."""
     )
 
-    status_box = gr.Markdown(f"`{TRUFOR_STATUS}`")
+    status_box = gr.Markdown(f"`TruFor: {TRUFOR_STATUS}`\n`AI Detector: {MODEL_STATUS}`")
 
     image_input = gr.Image(label="Input Image", type="pil")
     analyze_button = gr.Button("Analyze", variant="primary", size="sm")
@@ -101,18 +140,16 @@ with gr.Blocks() as demo:
     with gr.Tabs():
         with gr.TabItem("AI vs Human"):
             ai_label_output = gr.Label(label="Prediction", num_top_classes=2)
-            ai_summary_output = gr.Markdown("Upload an image to view the prediction.")
+            ai_summary_output = gr.Markdown(AI_INTRO_SUMMARY)
         with gr.TabItem("TruFor Forgery Detection"):
             trufor_summary_output = gr.Markdown("Configure TruFor assets to enable tamper analysis.")
-            tamper_overlay_output = gr.Image(label="Tamper Heatmap", type="pil", interactive=False)
-            conf_overlay_output = gr.Image(label="Confidence Heatmap", type="pil", interactive=False)
+            tamper_overlay_output = gr.Image(label="Altered Regions Map", type="pil", interactive=False)
 
     output_components = [
         ai_label_output,
         ai_summary_output,
         trufor_summary_output,
         tamper_overlay_output,
-        conf_overlay_output,
     ]
 
     analyze_button.click(

trufor_runner.py

@@ -23,8 +23,6 @@ class TruForUnavailableError(RuntimeError):
 class TruForResult:
     score: Optional[float]
     map_overlay: Optional[Image.Image]
-    confidence_overlay: Optional[Image.Image]
-    raw_scores: Dict[str, float]
 
 
 class TruForEngine:
@@ -151,35 +149,15 @@ class TruForEngine:
     def _infer_native(self, image: Image.Image) -> TruForResult:
         outputs = self.native_model.predict(image)
 
-        overlays: Dict[str, Optional[Image.Image]] = {"map": None, "conf": None}
+        map_overlay = None
         try:
-            overlays["map"] = self._apply_heatmap(image, outputs.tamper_map)
+            map_overlay = self._apply_heatmap(image, outputs.tamper_map)
         except Exception as exc:  # pragma: no cover - visualisation fallback
             LOGGER.debug("Failed to build tamper heatmap: %s", exc)
 
-        if outputs.confidence_map is not None:
-            try:
-                overlays["conf"] = self._apply_heatmap(image, outputs.confidence_map)
-            except Exception as exc:  # pragma: no cover
-                LOGGER.debug("Failed to build confidence heatmap: %s", exc)
-
-        raw_scores: Dict[str, float] = {
-            "tamper_mean": float(np.mean(outputs.tamper_map)),
-            "tamper_max": float(np.max(outputs.tamper_map)),
-        }
-
-        if outputs.confidence_map is not None:
-            raw_scores["confidence_mean"] = float(np.mean(outputs.confidence_map))
-            raw_scores["confidence_max"] = float(np.max(outputs.confidence_map))
-
-        if outputs.detection_score is not None:
-            raw_scores["tamper_score"] = float(outputs.detection_score)
-
         return TruForResult(
             score=outputs.detection_score,
-            map_overlay=overlays["map"],
-            confidence_overlay=overlays["conf"],
-            raw_scores=raw_scores,
+            map_overlay=map_overlay,
         )
 
     def _infer_docker(self, image: Image.Image) -> TruForResult:
@@ -257,35 +235,17 @@ class TruForEngine:
 
         data = np.load(npz_files[0], allow_pickle=False)
         tamper_map = data.get("map")
-        conf_map = data.get("conf")
         score = float(data["score"]) if "score" in data.files else None
 
-        overlays: Dict[str, Optional[Image.Image]] = {"map": None, "conf": None}
+        map_overlay = None
         try:
-            overlays["map"] = self._apply_heatmap(image, tamper_map) if tamper_map is not None else None
+            map_overlay = self._apply_heatmap(image, tamper_map) if tamper_map is not None else None
         except Exception as exc:  # pragma: no cover
             LOGGER.debug("Failed to build tamper heatmap: %s", exc)
 
-        try:
-            overlays["conf"] = self._apply_heatmap(image, conf_map) if conf_map is not None else None
-        except Exception as exc:  # pragma: no cover
-            LOGGER.debug("Failed to build confidence heatmap: %s", exc)
-
-        raw_scores: Dict[str, float] = {}
-        if score is not None:
-            raw_scores["tamper_score"] = score
-        if tamper_map is not None:
-            raw_scores["tamper_mean"] = float(np.mean(tamper_map))
-            raw_scores["tamper_max"] = float(np.max(tamper_map))
-        if conf_map is not None:
-            raw_scores["confidence_mean"] = float(np.mean(conf_map))
-            raw_scores["confidence_max"] = float(np.max(conf_map))
-
         return TruForResult(
             score=score,
-            map_overlay=overlays["map"],
-            confidence_overlay=overlays["conf"],
-            raw_scores=raw_scores,
+            map_overlay=map_overlay,
         )
 
     @staticmethod