Fix Docker CMD and UI run button

Dockerfile

@@ -9,7 +9,7 @@ COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
 # Copy your FastAPI app and dinov2 module
-COPY app ./app
+COPY app /code/app
 
 # Set environment variable for module discovery
 ENV PYTHONPATH=/code
@@ -17,7 +17,5 @@ ENV PYTHONPATH=/code
 # Expose the default Hugging Face Spaces port
 EXPOSE 7860
 
-# Run the FastAPI app
+# ✅ Run FastAPI app
 CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]
-
-

app/main.py

@@ -1,32 +1,55 @@
-from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import JSONResponse
-from fastapi.staticfiles import StaticFiles
-import uvicorn
-from .model import load_model, predict_from_bytes
+# app/main.py
 import os
+import json
+import torch
+from fastapi import FastAPI, File, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
+from transformers import AutoProcessor, Dinov2ForImageClassification
+from torch.nn.functional import softmax
+from PIL import Image
+
+app = FastAPI()
+
+# Allow frontend to call backend
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
-app = FastAPI(title="DINOv2 Attention Map Demo")
+# --- Load model and mapping on startup ---
+print("🚀 Loading model and label mapping...")
+MODEL_ID = "Arew99/dinov2-costum"
 
-MODEL = None
+processor = AutoProcessor.from_pretrained(MODEL_ID)
+model = Dinov2ForImageClassification.from_pretrained(MODEL_ID)
+model.eval()
 
-@app.on_event("startup")
-async def startup_event():
-    global MODEL
-    MODEL = load_model()
+# 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)
+
+print(f"✓ Loaded {len(id2name)} labels from id2name.json")
+
+@app.get("/")
+def root():
+    return {"message": "Welcome to NEMOtools API"}
 
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
-    contents = await file.read()
-    result = predict_from_bytes(MODEL, contents)
-    return JSONResponse(result)
-
-@app.get("/health")
-async def health():
-    return {"status": "ok"}
+    """Perform top-5 inference on an uploaded image."""
+    image = Image.open(file.file).convert("RGB")
+    inputs = processor(images=image, return_tensors="pt")
 
-# 👇 Mount static *after* routes
-static_dir = os.path.join(os.path.dirname(__file__), "static")
-app.mount("/", StaticFiles(directory=static_dir, html=True), name="static")
+    with torch.no_grad():
+        logits = model(**inputs).logits.squeeze(0)
+        probs, idxs = softmax(logits, dim=0).topk(5)
 
-if __name__ == "__main__":
-    uvicorn.run("app.main:app", host="0.0.0.0", port=7860)
+    results = [
+        {"label": id2name[str(i)], "confidence": float(p)}
+        for p, i in zip(probs, idxs)
+    ]
+    return {"predictions": results}

app/static/index.html

@@ -174,7 +174,7 @@
 
           <!-- Classification tool -->
           <div id="tool-classification" class="hidden">
-            <div class="flex flex-col md:flex-row md:items-center gap-4 mb-6 justify-center">
+            <div class="flex flex-col items-center gap-4 mb-6 justify-center">
               <input id="cls-file" type="file" accept="image/*"
                     class="block w-full md:w-auto text-sm text-gray-600
                           file:mr-4 file:py-2 file:px-4
@@ -182,10 +182,22 @@
                           file:text-sm file:font-semibold
                           file:bg-indigo-50 file:text-indigo-600
                           hover:file:bg-indigo-100" />
+          
+              <!-- 👇 THIS BUTTON TRIGGERS INFERENCE -->
+              <button id="runClsBtn"
+                onclick="runClassification()"
+                class="px-8 py-3 bg-green-600 text-white text-lg font-semibold rounded-full shadow-md hover:bg-green-700 transition">
+                ▶️ Run Classification
+              </button>
             </div>
-
+          
+            <!-- Where results will appear -->
             <div id="cls-result" class="text-center text-gray-700 mt-4 text-lg font-medium"></div>
           </div>
+          
+          
+            <!-- <div id="cls-result" class="text-center text-gray-700 mt-4 text-lg font-medium"></div> -->
+          </div>          
         </div>
       </section>
     </main>
@@ -324,10 +336,19 @@
         fd.append("file", file);
 
         try {
-          const res = await fetch("/classify", { method: "POST", body: fd });
+          const res = await fetch("/predict", { method: "POST", body: fd }); // ✅ must match FastAPI route
           if (!res.ok) throw new Error(`Server error: ${res.status}`);
           const json = await res.json();
-          result.textContent = "✅ Predicted class: " + json.label;
+
+          // ✅ display top-5 predictions if available
+          if (json.predictions) {
+            result.innerHTML = "<h3 class='font-semibold text-indigo-600 mb-2'>Top-5 Predictions:</h3>" +
+              json.predictions.map(p =>
+                `<div>${p.label} — ${(p.confidence * 100).toFixed(2)}%</div>`
+              ).join("");
+          } else {
+            result.textContent = "✅ Predicted class: " + json.label;
+          }
         } catch (err) {
           console.error(err);
           result.textContent = "❌ Error: " + err.message;