Recreate clean HF Space using HF Dataset for reference images

Dockerfile

@@ -1,31 +1,22 @@
 FROM python:3.11-slim
 
 WORKDIR /app
-
-# Make local modules importable
 ENV PYTHONPATH=/app
 
-# Install system dependencies
 RUN apt-get update && apt-get install -y \
     build-essential \
     && rm -rf /var/lib/apt/lists/*
 
-# Copy requirements first for better caching
 COPY requirements.txt .
 
-# Install Python dependencies
 RUN pip install --no-cache-dir --upgrade pip && \
-    pip install --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/cpu && \
+    pip install --no-cache-dir torch torchvision \
+        --index-url https://download.pytorch.org/whl/cpu && \
     pip install --no-cache-dir -r requirements.txt
 
-# Copy ALL application files
 COPY *.py ./
-
-# Copy model weights
 COPY resnet50_multitask_mold.pth ./
 
-# Expose HF Spaces port
 EXPOSE 7860
 
-# Run the app
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -2,14 +2,17 @@ from fastapi import FastAPI, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image
 import torch, io
-from pathlib import Path
 from torchvision import transforms
 
 from model import MultiTaskResNet50
-from decision import final_decision                #
-from advanced_decision import *
+from decision import final_decision
+from advanced_decision import (
+    mc_uncertainty,
+    patch_consistency,
+    final_decision_v2
+)
 from gradcam import GradCAM
-from dino import *
+from dino import load_dino, build_embeddings, similarity
 
 app = FastAPI(title="Mold Detection API v2")
 
@@ -23,24 +26,49 @@ app.add_middleware(
 device = "cuda" if torch.cuda.is_available() else "cpu"
 mold_idx = 4
 
-# Load model
+# ------------------
+# Load main model
+# ------------------
 model = MultiTaskResNet50().to(device)
-model.load_state_dict(torch.load("resnet50_multitask_mold.pth", map_location=device))
+model.load_state_dict(
+    torch.load("resnet50_multitask_mold.pth", map_location=device)
+)
 model.eval()
 
+# ------------------
 # Transforms
+# ------------------
 transform = transforms.Compose([
-    transforms.Resize((224,224)),
+    transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
+    transforms.Normalize(
+        [0.485, 0.456, 0.406],
+        [0.229, 0.224, 0.225]
+    )
 ])
 
+# ------------------
 # Grad-CAM
+# ------------------
 gradcam = GradCAM(model, model.backbone.layer4[-1].conv3)
 
-# DINO
-dino = load_dino(device)
-mold_embs = build_embeddings(dino, transform, "mold_reference_images", device)
+# ------------------
+# DINO (lazy loaded)
+# ------------------
+dino = None
+mold_embs = None
+
+
+def ensure_dino():
+    global dino, mold_embs
+    if dino is None:
+        dino = load_dino(device)
+        mold_embs = build_embeddings(dino, transform, device)
+
+
+# ------------------
+# API endpoints
+# ------------------
 
 @app.post("/predict/v1")
 async def predict_v1(file: UploadFile):
@@ -48,22 +76,29 @@ async def predict_v1(file: UploadFile):
     img_t = transform(img).to(device)
     return final_decision(model, img_t)
 
+
 @app.post("/predict/v2")
 async def predict_v2(file: UploadFile):
+    ensure_dino()
+
     img = Image.open(io.BytesIO(await file.read())).convert("RGB")
     img_t = transform(img).to(device)
 
     with torch.no_grad():
         out = model(img_t.unsqueeze(0))
-        cp = torch.softmax(out["class"],1)[0]
-        bp = torch.softmax(out["bio"],1)[0]
+        cp = torch.softmax(out["class"], 1)[0]
+        bp = torch.softmax(out["bio"], 1)[0]
 
     mold_p = cp[mold_idx].item()
-    bio_p  = bp[1].item()
+    bio_p = bp[1].item()
 
     mean_p, std_p = mc_uncertainty(model, img_t, mold_idx)
-    patch_ratio = patch_consistency(model, img, transform, mold_idx, device)
-    dino_sim = similarity(dino, mold_embs, img, transform, device)
+    patch_ratio = patch_consistency(
+        model, img, transform, mold_idx, device
+    )
+    dino_sim = similarity(
+        dino, mold_embs, img, transform, device
+    )
 
     decision = final_decision_v2(
         mold_p, bio_p, std_p, patch_ratio, dino_sim
@@ -72,16 +107,17 @@ async def predict_v2(file: UploadFile):
     return {
         "decision": decision,
         "model_outputs": {
-            "mold_probability": round(mold_p,3),
-            "biological_probability": round(bio_p,3)
+            "mold_probability": round(mold_p, 3),
+            "biological_probability": round(bio_p, 3),
         },
         "confidence_checks": {
-            "uncertainty": round(std_p,3),
-            "patch_ratio": round(patch_ratio,3),
-            "dino_similarity": round(dino_sim,3)
-        }
+            "uncertainty": round(std_p, 3),
+            "patch_ratio": round(patch_ratio, 3),
+            "dino_similarity": round(dino_sim, 3),
+        },
     }
 
+
 @app.post("/explain/gradcam")
 async def explain_gradcam(file: UploadFile):
     img = Image.open(io.BytesIO(await file.read())).convert("RGB")

dino.py

@@ -1,28 +1,47 @@
-import os
-import numpy as np
 import torch
-import torch.hub
+import numpy as np
 from PIL import Image
+from datasets import load_dataset
 from sklearn.metrics.pairwise import cosine_similarity
 
+
 def load_dino(device):
-    model = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14")
+    model = torch.hub.load(
+        "facebookresearch/dinov2",
+        "dinov2_vits14"
+    )
     model.eval().to(device)
     return model
 
-def build_embeddings(dino, transform, image_dir, device):
+
+def build_embeddings(dino, transform, device):
+    dataset = load_dataset(
+        "AdarshDS/mold-reference-images",
+        split="train"
+    )
+
     embs = []
-    for f in os.listdir(image_dir):
-        if f.lower().endswith((".jpg",".png",".jpeg")):
-            img = Image.open(os.path.join(image_dir,f)).convert("RGB")
-            t = transform(img).unsqueeze(0).to(device)
-            with torch.no_grad():
-                e = dino(t)
-            embs.append(e.squeeze().cpu().numpy())
+
+    for sample in dataset:
+        img: Image.Image = sample["image"].convert("RGB")
+        t = transform(img).unsqueeze(0).to(device)
+
+        with torch.no_grad():
+            e = dino(t)
+
+        embs.append(e.squeeze().cpu().numpy())
+
+    if not embs:
+        raise RuntimeError(
+            "No reference images found in HF dataset"
+        )
+
     return np.vstack(embs)
 
+
 def similarity(dino, mold_embs, image, transform, device):
     t = transform(image).unsqueeze(0).to(device)
     with torch.no_grad():
         e = dino(t).cpu().numpy()
+
     return float(cosine_similarity(e, mold_embs).max())

requirements.txt

@@ -6,6 +6,8 @@ pillow
 numpy<2
 python-multipart
 scikit-learn
+scikit-learn
+datasets