Fix DINOv2 processor and update requirements

app/main.py

@@ -1,16 +1,31 @@
 # 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 transformers import AutoImageProcessor, Dinov2ForImageClassification
+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
 
-app = FastAPI()
+# -------------------------------------------------
+# paths
+# -------------------------------------------------
+BASE_DIR = Path(__file__).parent
+STATIC_DIR = BASE_DIR / "static"
+INDEX_HTML = STATIC_DIR / "index.html"
+MAP_PATH = BASE_DIR / "id2name.json"
+
+app = FastAPI(title="NEMO Tools")
 
-# Allow frontend to call backend
+# CORS so the JS can call us
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -19,44 +34,80 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# --- Load model and mapping on startup ---
+# serve /static/*
+app.mount("/static", StaticFiles(directory=str(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"
 
-print("🚀 Loading model and label mapping...")
+# model: your fine-tuned one
 model = Dinov2ForImageClassification.from_pretrained(
     MODEL_ID,
     num_labels=101,
-    ignore_mismatched_sizes=True
+    ignore_mismatched_sizes=True,
 )
-processor = AutoImageProcessor.from_pretrained("facebook/dinov2-large")
-
 model.eval()
 
-# Load id2name.json
-MAP_PATH = os.path.join(os.path.dirname(__file__), "id2name.json")
-with open(MAP_PATH, "r") as f:
-    id2name = json.load(f)
+# 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")
 
-@app.get("/")
-def root():
-    return {"message": "Welcome to NEMOtools API"}
 
+# -------------------------------------------------
+# endpoints
+# -------------------------------------------------
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
-    """Perform top-5 inference on an uploaded image."""
-    image = Image.open(file.file).convert("RGB")
-    inputs = processor(images=image, return_tensors="pt")
+    # 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.squeeze(0)
+        logits = model(**inputs).logits[0]   # shape [101]
         probs, idxs = softmax(logits, dim=0).topk(5)
 
-    results = [
-        {"label": id2name[str(i)], "confidence": float(p)}
-        for p, i in zip(probs, idxs)
-    ]
+    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}")}
+
+
+@app.get("/api")
+def api_root():
+    return {"message": "NEMO Tools backend is running."}
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -8,7 +8,7 @@ pillow
 numpy
 
 # Hugging Face bits
-transformers
+transformers>=4.42.0
 huggingface-hub
 peft