Auto image size, working attention map

app/Inference.py

@@ -0,0 +1,29 @@
+import json
+import torch
+from transformers import AutoProcessor, Dinov2ForImageClassification
+from PIL import Image
+from torch.nn.functional import softmax
+
+# --- Load mapping ---
+with open("id2name.json", "r") as f:
+    id2name = json.load(f)
+
+# --- Load model ---
+model_name = "Arew99/dinov2-costum"
+processor = AutoProcessor.from_pretrained(model_name)
+model = Dinov2ForImageClassification.from_pretrained(model_name)
+model.eval()
+
+# --- Load image (example) ---
+image = Image.open("sample_fish.jpg").convert("RGB")
+inputs = processor(images=image, return_tensors="pt")
+
+# --- Inference ---
+with torch.no_grad():
+    logits = model(**inputs).logits.squeeze(0)
+    probs, idxs = softmax(logits, dim=0).topk(5)
+
+print("\nTop-5 predictions:")
+for p, i in zip(probs.tolist(), idxs.tolist()):
+    label = id2name[str(i)]
+    print(f"{label:30s}  {p*100:.2f}%")

app/main.py

@@ -1,31 +1,17 @@
 # app/main.py
 import os
-import json
-from pathlib import Path
-
-import torch
 from fastapi import FastAPI, File, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import HTMLResponse
 from fastapi.staticfiles import StaticFiles
-from transformers import (
-    Dinov2ForImageClassification,
-    Dinov2ImageProcessor,   # <-- needs the newer transformers
-)
-from torch.nn.functional import softmax
-from PIL import Image
+from app.model import load_model, predict_from_bytes
 
-# -------------------------------------------------
-# paths
-# -------------------------------------------------
-BASE_DIR = Path(__file__).parent
-STATIC_DIR = BASE_DIR / "static"
-INDEX_HTML = STATIC_DIR / "index.html"
-MAP_PATH = BASE_DIR / "id2name.json"
 
+# ──────────────────────────────────────────────
+# FastAPI setup
+# ──────────────────────────────────────────────
 app = FastAPI(title="NEMO Tools")
 
-# CORS so the JS can call us
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -34,80 +20,43 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# serve /static/*
-app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
+# ──────────────────────────────────────────────
+# Static Frontend
+# ──────────────────────────────────────────────
+BASE_DIR = os.path.dirname(__file__)
+STATIC_DIR = os.path.join(BASE_DIR, "static")
+INDEX_HTML = os.path.join(STATIC_DIR, "index.html")
 
+app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
 
 @app.get("/", response_class=HTMLResponse)
 def serve_frontend():
-    return INDEX_HTML.read_text(encoding="utf-8")
-
-
-# -------------------------------------------------
-# load model + processor + labels ONCE
-# -------------------------------------------------
-print("🚀 Loading model and label mapping...")
-
-MODEL_ID = "Arew99/dinov2-costum"
-
-# model: your fine-tuned one
-model = Dinov2ForImageClassification.from_pretrained(
-    MODEL_ID,
-    num_labels=101,
-    ignore_mismatched_sizes=True,
-)
-model.eval()
-
-# processor: from the ORIGINAL dino repo (not your custom one)
-processor = Dinov2ImageProcessor.from_pretrained("facebook/dinov2-large")
-
-# labels
-with MAP_PATH.open("r") as f:
-    id2name = json.load(f)
-print(f"✓ Loaded {len(id2name)} labels from id2name.json")
-
-
-# -------------------------------------------------
-# endpoints
-# -------------------------------------------------
-@app.post("/predict")
-async def predict(file: UploadFile = File(...)):
-    # this is your “top-5 for an image” endpoint
-    img = Image.open(file.file).convert("RGB")
-
-    # Dinov2ImageProcessor wants a list → [img]
-    inputs = processor(images=[img], return_tensors="pt")
-
-    with torch.no_grad():
-        logits = model(**inputs).logits[0]   # shape [101]
-        probs, idxs = softmax(logits, dim=0).topk(5)
-
-    results = []
-    for p, i in zip(probs.tolist(), idxs.tolist()):
-        label = id2name.get(str(i), f"Class {i}")
-        results.append({"label": label, "confidence": p})
-
-    return {"predictions": results}
-
-
-@app.post("/classify")
-async def classify(file: UploadFile = File(...)):
-    img = Image.open(file.file).convert("RGB")
-    inputs = processor(images=[img], return_tensors="pt")
-
-    with torch.no_grad():
-        logits = model(**inputs).logits[0]
-        pred = int(logits.argmax().item())
-
-    return {"label": id2name.get(str(pred), f"Class {pred}")}
-
+    """Serve the web interface."""
+    with open(INDEX_HTML, "r", encoding="utf-8") as f:
+        return f.read()
+
+# ──────────────────────────────────────────────
+# Model Initialization
+# ──────────────────────────────────────────────
+print("🚀 Loading DINOv2 custom model...")
+model_device_tuple = load_model()
+print("✅ Model loaded and ready for inference!")
+
+# ──────────────────────────────────────────────
+# API Endpoints
+# ──────────────────────────────────────────────
+@app.post("/attention")
+async def generate_attention(file: UploadFile = File(...)):
+    """Generate and return mean attention map for uploaded image."""
+    image_bytes = await file.read()
+    result = predict_from_bytes(model_device_tuple, image_bytes)
+    return result
 
 @app.get("/api")
 def api_root():
-    return {"message": "NEMO Tools backend is running."}
-
+    return {"message": "NEMO Tools backend running."}
 
+# ──────────────────────────────────────────────
 if __name__ == "__main__":
     import uvicorn
-
     uvicorn.run(app, host="0.0.0.0", port=7860)

app/model.py

@@ -25,7 +25,7 @@ CKPT_PATH = hf_hub_download(
 )
 
 PATCH_SIZE = 14
-IMAGE_SIZE = (1024, 720)
+IMAGE_SIZE = (1000,1000)
 
 
 # -------------------------------------------------------
@@ -46,6 +46,7 @@ def load_model():
     # Load weights
     state_dict = load_file(CKPT_PATH)
     keys_list = list(state_dict.keys())
+    print(f"Loaded {len(state_dict.keys())} weights from {CKPT_PATH}")
 
     # Handle "model." prefix if present
     if keys_list and "model." in keys_list[0]:

app/static/Dockerfile

@@ -1,22 +0,0 @@
-# Use a lightweight Python image
-FROM python:3.10-slim
-
-# Set working directory
-WORKDIR /code
-
-# Copy requirements and install dependencies
-COPY requirements.txt .
-RUN pip install --no-cache-dir -r requirements.txt
-
-# Copy your FastAPI app and dinov2 module
-COPY app ./app
-COPY dinov2 ./dinov2
-
-# Set environment variable for module discovery
-ENV PYTHONPATH=/code
-
-# Expose the default Hugging Face Spaces port
-EXPOSE 7860
-
-# Run the FastAPI app
-CMD ["python", "-m", "app.main"]

app/static/index.html

@@ -55,7 +55,7 @@
       <div class="max-w-6xl mx-auto px-4 py-4 flex items-center justify-between">
         <!-- Logo and title -->
         <div class="flex items-center gap-3">
-          <img src="assets/logo.png" alt="NEMO logo" class="h-10 w-10 rounded-full shadow-sm" />
+          <img src="/static/assets/logo.png" alt="NEMO logo" class="h-10 w-10 rounded-full shadow-sm" />
           <div>
             <h1 class="text-lg font-bold text-indigo-600">NEMO tools</h1>
             <p class="text-xs text-gray-400">DINOv2 visualisation sandbox</p>
@@ -308,7 +308,7 @@
         fd.append("file", file);
 
         try {
-          const res = await fetch("/predict", { method: "POST", body: fd });
+          const res = await fetch("/attention", { method: "POST", body: fd });
           if (!res.ok) throw new Error(`Server error: ${res.status}`);
           const json = await res.json();
 
@@ -337,7 +337,7 @@
         fd.append("file", file);
 
         try {
-          const res = await fetch("/predict", { method: "POST", body: fd }); // ✅ must match FastAPI route
+          const res = await fetch("/attention", { method: "POST", body: fd }); // ✅ must match FastAPI route
           if (!res.ok) throw new Error(`Server error: ${res.status}`);
           const json = await res.json();
 

requirements.txt

@@ -6,6 +6,7 @@ torch
 torchvision
 pillow
 numpy
+matplotlib
 
 # Hugging Face bits
 transformers>=4.42.0