Update app.py

app.py

@@ -1,94 +1,106 @@
-#!/usr/bin/env python
-# ──────────────────────────────────────────────────────────────
-#  BubbleAI Image-Safety Detector – Hugging Face Space
-#  -------------------------------------------------------------
-#  Gradio app that classifies uploaded images as “Safe” or
-#  “Unsafe” using a fine-tuned ResNet-50.
-#
-#  Coder: Amir Mehdi Memari  (2025-08-06)
-#  Description:
-#     • Loads the checkpoint `resnet_safety_classifier.pth`
-#       (must be in the same repo directory; tracked with Git-LFS).
-#     • Applies standard ImageNet preprocessing.
-#     • Returns class probabilities via a simple Gradio UI.
-#  Usage:
-#     The HF Space builder executes `python app.py` automatically.
-# ──────────────────────────────────────────────────────────────
-
-from __future__ import annotations
-import pathlib, typing as t
-
-import torch
-import torchvision
-from torchvision import transforms
-from PIL import Image
-import gradio as gr
-
-
-# ── 1. Locate checkpoint ──────────────────────────────────────
-REPO_DIR    = pathlib.Path(__file__).parent
-CKPT_PATH   = REPO_DIR / "resnet_safety_classifier.pth"
-DEVICE      = "cuda" if torch.cuda.is_available() else "cpu"
-
-
-# ── 2. Define the model architecture exactly as trained ───────
-class SafetyResNet(torch.nn.Module):
-    """
-    ResNet-50 backbone with a 2-unit classifier head.
-    """
-    def __init__(self) -> None:
-        super().__init__()
-        base = torchvision.models.resnet50(weights=None)
-        base.fc = torch.nn.Linear(base.fc.in_features, 2)   # Safe / Unsafe
-        self.model = base
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.model(x)
-
-
-# ── 3. Instantiate & load weights ─────────────────────────────
-model = SafetyResNet().to(DEVICE)
-state = torch.load(CKPT_PATH, map_location=DEVICE)
-model.load_state_dict(state, strict=True)
-model.eval()                                           # inference mode
-CLASSES = ["Safe", "Unsafe"]
-
-
-# ── 4. Pre-processing pipeline (ImageNet stats) ───────────────
-preprocess = transforms.Compose([
-    transforms.Resize(256),
-    transforms.CenterCrop(224),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                         std =[0.229, 0.224, 0.225]),
-])
-
-
-# ── 5. Inference helper ───────────────────────────────────────
-@torch.inference_mode()
-def predict(img: Image.Image) -> t.Dict[str, float]:
-    """
-    Returns {class_name: probability} for a single PIL image.
-    """
-    tensor = preprocess(img).unsqueeze(0).to(DEVICE)
-    probs  = torch.softmax(model(tensor)[0], dim=0).cpu().tolist()
-    return {CLASSES[i]: float(probs[i]) for i in range(2)}
-
-
-# ── 6. Build Gradio Interface ─────────────────────────────────
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil", label="Upload an image"),
-    outputs=gr.Label(num_top_classes=2, label="Prediction"),
-    title="BubbleAI Image-Safety Detector",
-    description=(
-        "This demo classifies images as **Safe** or **Unsafe** (NSFW) "
-        "using a fine-tuned ResNet-50. Probabilities are shown for both "
-        "classes. Model weights © 2025 Amir Mehdi Memari."
-    ),
-    cache_examples=False,           # disable weighty example cache rebuilds
-)
-
-# ── 7. Launch (HF Spaces auto-calls this in production) ───────
-if __name__ == "__main__":
-    demo.launch()
+#!/usr/bin/env python
+# ──────────────────────────────────────────────────────────────
+#  BubbleAI Image-Safety Detector – fixed version
+#  -------------------------------------------------------------
+#  Loads checkpoint `resnet_safety_classifier.pth` whose keys are
+#  saved under feature_extractor.* and classifier.*; serves a
+#  Gradio UI that predicts “Safe” vs “Unsafe”.
+#
+#  Coder: Amir Mehdi Memari   (2025-08-06)
+# ──────────────────────────────────────────────────────────────
+
+from __future__ import annotations
+import pathlib
+import typing as t
+
+import torch
+import torchvision
+from torchvision import transforms
+from PIL import Image
+import gradio as gr
+
+
+# ── 1. Paths & device ─────────────────────────────────────────
+REPO_DIR  = pathlib.Path(__file__).parent
+CKPT_PATH = REPO_DIR / "resnet_safety_classifier.pth"
+DEVICE    = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+# ── 2. Architecture that matches checkpoint keys ──────────────
+class SafetyResNet(torch.nn.Module):
+    """
+    ResNet-50 backbone (conv1 ▸ layer4)  ➜  global avg-pool
+    ➜  MLP (2048→512→2).  Keys align with:
+        • feature_extractor.*
+        • classifier.*
+    """
+    def __init__(self) -> None:
+        super().__init__()
+
+        base = torchvision.models.resnet50(weights=None)
+
+        # keep stem + 4 stages (0-7)  [conv1, bn1, relu, maxpool, layer1-4]
+        self.feature_extractor = torch.nn.Sequential(*list(base.children())[:8])
+
+        self.pool = torch.nn.AdaptiveAvgPool2d((1, 1))
+        self.classifier = torch.nn.Sequential(
+            torch.nn.Flatten(),                   # (B, 2048, 1, 1) → (B, 2048)
+            torch.nn.Linear(2048, 512, bias=True),
+            torch.nn.ReLU(inplace=True),
+            torch.nn.Dropout(p=0.30),
+            torch.nn.Linear(512, 2, bias=True)
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.feature_extractor(x)   # (B, 2048, H/32, W/32)
+        x = self.pool(x)                # (B, 2048, 1, 1)
+        x = self.classifier(x)          # (B, 2)
+        return x
+
+
+# ── 3. Instantiate & load weights ─────────────────────────────
+model = SafetyResNet().to(DEVICE)
+state = torch.load(CKPT_PATH, map_location=DEVICE)
+model.load_state_dict(state, strict=True)
+model.eval()
+
+CLASSES = ["Safe", "Unsafe"]
+
+
+# ── 4. Pre-processing pipeline (ImageNet stats) ───────────────
+preprocess = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std =[0.229, 0.224, 0.225]),
+])
+
+
+# ── 5. Inference helper ───────────────────────────────────────
+@torch.inference_mode()
+def predict(img: Image.Image) -> t.Dict[str, float]:
+    """
+    Returns {class_name: probability} for a single PIL image.
+    """
+    tensor = preprocess(img).unsqueeze(0).to(DEVICE)
+    probs  = torch.softmax(model(tensor)[0], dim=0).cpu().tolist()
+    return {CLASSES[i]: float(probs[i]) for i in range(2)}
+
+
+# ── 6. Gradio interface ───────────────────────────────────────
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Upload an image"),
+    outputs=gr.Label(num_top_classes=2, label="Prediction"),
+    title="BubbleAI Image-Safety Detector",
+    description=(
+        "Drop an image to check whether it's **Safe** or **Unsafe** "
+        "using BubbleAI’s ResNet-50 classifier."
+    ),
+    cache_examples=False,
+)
+
+# ── 7. Launch ────────────────────────────────────────────────
+if __name__ == "__main__":
+    demo.launch()