🚀 Update: new UI design, classification plot, and attention improvements

app/Inference.py

@@ -1,29 +1,242 @@
+# # app/inference.py
+# import os
+# import io
+# import json
+# import torch
+# import torch.nn.functional as F
+# from PIL import Image
+# from transformers import AutoImageProcessor, AutoModelForImageClassification
+
+# # ─────────────────────────────────────────────
+# # CONFIG
+# # ─────────────────────────────────────────────
+# # Hugging Face repo for the trained checkpoint
+# MODEL_REPO = "Arew99/dinov2-costum"
+
+# # optional: path to local label map (bundled in your repo)
+# ID2NAME_PATH = os.path.join(os.path.dirname(__file__), "id2name.json")
+
+# # detect device automatically
+# DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# print(f"🧠 Using device: {DEVICE}")
+
+# # global cache (so model loads only once)
+# _model = None
+# _processor = None
+# _id2name = None
+
+
+# # ─────────────────────────────────────────────
+# # HELPER: load label map
+# # ─────────────────────────────────────────────
+# def _load_id2name():
+#     if os.path.exists(ID2NAME_PATH):
+#         with open(ID2NAME_PATH, "r") as f:
+#             data = json.load(f)
+#         # ensure integer keys
+#         return {int(k): v for k, v in data.items()}
+#     print("⚠️  id2name.json not found — using placeholder labels.")
+#     return {i: f"Class {i}" for i in range(101)}  # fallback
+
+
+# # ─────────────────────────────────────────────
+# # INIT: load model & processor once
+# # ─────────────────────────────────────────────
+# def load_classification_model():
+#     global _model, _processor, _id2name
+
+#     if _model is not None:
+#         return _model, _processor, _id2name
+
+#     print(f"🔁 Loading model from Hugging Face repo: {MODEL_REPO}")
+#     _processor = AutoImageProcessor.from_pretrained(MODEL_REPO)
+#     _model = AutoModelForImageClassification.from_pretrained(
+#         MODEL_REPO,
+#         ignore_mismatched_sizes=True,
+#     ).to(DEVICE)
+#     _model.eval()
+#     _id2name = _load_id2name()
+
+#     print(f"✅ Model loaded and ready on {DEVICE}")
+#     return _model, _processor, _id2name
+
+
+# # ─────────────────────────────────────────────
+# # INFERENCE: classify raw image bytes
+# # ─────────────────────────────────────────────
+# def classify_bytes(image_bytes: bytes):
+#     model, processor, id2name = load_classification_model()
+
+#     image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+#     inputs = processor(images=image, return_tensors="pt").to(DEVICE)
+
+#     with torch.no_grad():
+#         outputs = model(**inputs)
+#         probs = F.softmax(outputs.logits, dim=-1)
+
+#     topk = torch.topk(probs, k=5)
+#     indices = topk.indices[0].tolist()
+#     values = topk.values[0].tolist()
+
+#     results = []
+#     for rank, (idx, prob) in enumerate(zip(indices, values), 1):
+#         label = id2name.get(int(idx), f"Class {idx}")
+#         results.append({
+#             "rank": rank,
+#             "id": int(idx),
+#             "label": label,
+#             "score": float(prob),
+#         })
+
+#     # concise summary for API
+#     return {
+#         "top1": results[0],
+#         "top5": results,
+#     }
+
+# app/inference.py
+import os
+import io
 import json
+import base64
 import torch
-from transformers import AutoProcessor, Dinov2ForImageClassification
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
 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}%")
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+
+# ─────────────────────────────────────────────
+# CONFIG
+# ─────────────────────────────────────────────
+MODEL_REPO = "Arew99/dinov2-costum"  # your Hugging Face repo
+ID2NAME_PATH = os.path.join(os.path.dirname(__file__), "id2name.json")
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"🧠 Using device: {DEVICE}")
+
+_model = None
+_processor = None
+_id2name = None
+
+
+# ─────────────────────────────────────────────
+# HELPER — load id2name mapping
+# ─────────────────────────────────────────────
+def _load_id2name():
+    if os.path.exists(ID2NAME_PATH):
+        with open(ID2NAME_PATH, "r") as f:
+            data = json.load(f)
+        return {int(k): v for k, v in data.items()}
+    print("⚠️ id2name.json not found — using placeholder labels.")
+    return {i: f"Class {i}" for i in range(101)}
+
+
+# ─────────────────────────────────────────────
+# LOAD MODEL (cached globally)
+# ─────────────────────────────────────────────
+def load_classification_model():
+    global _model, _processor, _id2name
+
+    if _model is not None:
+        return _model, _processor, _id2name
+
+    print(f"🔁 Loading model from Hugging Face repo: {MODEL_REPO}")
+    _processor = AutoImageProcessor.from_pretrained(MODEL_REPO)
+    _model = AutoModelForImageClassification.from_pretrained(
+        MODEL_REPO,
+        ignore_mismatched_sizes=True,
+    ).to(DEVICE)
+    _model.eval()
+    _id2name = _load_id2name()
+
+    print(f"✅ Model loaded and ready on {DEVICE}")
+    return _model, _processor, _id2name
+
+
+# ─────────────────────────────────────────────
+# CLASSIFY IMAGE BYTES
+# ─────────────────────────────────────────────
+def classify_bytes(image_bytes: bytes):
+    model, processor, id2name = load_classification_model()
+
+    # Load and preprocess image
+    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    inputs = processor(images=image, return_tensors="pt").to(DEVICE)
+
+    # Forward pass
+    with torch.no_grad():
+        outputs = model(**inputs)
+        probs = F.softmax(outputs.logits, dim=-1)
+
+    # Top-5 predictions
+    topk = torch.topk(probs, k=5)
+    indices = topk.indices[0].tolist()
+    values = topk.values[0].tolist()
+
+    results = []
+    for rank, (idx, prob) in enumerate(zip(indices, values), 1):
+        label = id2name.get(int(idx), f"Class {idx}")
+        results.append({
+            "rank": rank,
+            "id": int(idx),
+            "label": label,
+            "score": float(prob),
+        })
+
+    # ───────────────────────────────
+    # MATPLOTLIB TOP-3 PLOT
+    # ───────────────────────────────
+    top3 = results[:3]
+    labels = [p["label"] for p in top3]
+    probs_top3 = [p["score"] * 100 for p in top3]
+
+    plt.style.use("seaborn-v0_8-whitegrid")
+    fig, ax = plt.subplots(1, 2, figsize=(9, 4))
+
+    # Left: input image
+    ax[0].imshow(image)
+    ax[0].axis("off")
+    ax[0].set_title("Input Image", fontsize=12, weight="bold")
+
+    # Right: horizontal bar chart
+    bars = ax[1].barh(labels[::-1], probs_top3[::-1],
+                      color=["#C44E52", "#55A868", "#4C72B0"],
+                      edgecolor="none", height=0.6)
+    ax[1].set_xlim(0, 100)
+    ax[1].set_xlabel("Probability (%)", fontsize=11)
+    ax[1].set_title("Top-3 Predicted Species", fontsize=12, weight="bold")
+
+    for bar, prob in zip(bars, probs_top3[::-1]):
+        ax[1].text(prob + 1, bar.get_y() + bar.get_height()/2,
+                   f"{prob:.1f}%", va="center", fontsize=10, weight="bold")
+
+    plt.tight_layout()
+
+    # Encode plot as base64
+    buf = io.BytesIO()
+    plt.savefig(buf, format="png", bbox_inches="tight")
+    plt.close(fig)
+    buf.seek(0)
+    plot_b64 = base64.b64encode(buf.read()).decode("utf-8")
+    buf.close()
+
+    # ───────────────────────────────
+    # FINAL OUTPUT
+    # ───────────────────────────────
+    return {
+        "top1": results[0],
+        "top5": results,
+        "plot": f"data:image/png;base64,{plot_b64}"
+    }
+
+
+# ─────────────────────────────────────────────
+# LOCAL TEST
+# ──────────────────────────────���──────────────
+if __name__ == "__main__":
+    test_img = "sample3.jpg"
+    with open(test_img, "rb") as f:
+        img_bytes = f.read()
+    out = classify_bytes(img_bytes)
+    print(json.dumps(out["top5"], indent=2))
+    print("\nPlot base64 length:", len(out["plot"]))

app/MyInference.py

@@ -0,0 +1,73 @@
+import torch
+from PIL import Image
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+import torch.nn.functional as F
+import json
+import matplotlib.pyplot as plt
+
+# ===== Load label names =====
+with open("id2name.json", "r") as f:
+    id2name = json.load(f)
+
+# ===== Paths =====
+model_dir = "ckpt_merged_large"
+image_path = "sample3.jpg"
+
+# ===== Auto-detect device =====
+device = "cpu"
+print(f"Using device: {device}")
+
+# ===== Load processor & model =====
+processor = AutoImageProcessor.from_pretrained(model_dir)
+model = AutoModelForImageClassification.from_pretrained(model_dir)
+model.to(device)
+model.eval()
+
+# ===== Preprocess image =====
+image = Image.open(image_path).convert("RGB")
+inputs = processor(images=image, return_tensors="pt").to(device)
+
+# ===== Inference =====
+with torch.no_grad():
+    outputs = model(**inputs)
+    probs = F.softmax(outputs.logits, dim=-1)
+
+# ===== Top-5 predictions =====
+topk = torch.topk(probs, k=5)
+indices = topk.indices[0].tolist()
+values = topk.values[0].tolist()
+
+print("\nTop-5 Predictions:")
+for rank, (idx, prob) in enumerate(zip(indices, values), 1):
+    label = id2name[str(idx)]
+    print(f"{rank}. {label:<30} ({prob*100:.2f}%)")
+
+# ===== Prepare Top-3 for plotting =====
+top3_labels = [id2name[str(indices[i])] for i in range(3)]
+top3_probs = [values[i] * 100 for i in range(3)]
+
+# ===== Styled plot =====
+plt.style.use("seaborn-v0_8-whitegrid")
+fig, ax = plt.subplots(1, 2, figsize=(10, 4))
+
+# -- Left: input image --
+ax[0].imshow(image)
+ax[0].axis("off")
+ax[0].set_title("Input Image", fontsize=13, weight="bold")
+
+# -- Right: bar chart --
+bars = ax[1].barh(top3_labels[::-1], top3_probs[::-1],
+                  color=["#4C72B0", "#55A868", "#C44E52"],
+                  edgecolor="none", height=0.6)
+
+ax[1].set_xlim(0, 100)
+ax[1].set_xlabel("Probability (%)", fontsize=12)
+ax[1].set_title("Top-3 Predicted Species", fontsize=13, weight="bold")
+
+# Add percentage labels next to bars
+for bar, prob in zip(bars, top3_probs[::-1]):
+    ax[1].text(prob + 1, bar.get_y() + bar.get_height() / 2,
+               f"{prob:.1f}%", va="center", fontsize=11, weight="bold", color="#333")
+
+plt.tight_layout()
+plt.show()

app/main.py

@@ -1,10 +1,11 @@
 # app/main.py
-import os
-from fastapi import FastAPI, File, UploadFile
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import HTMLResponse
+from fastapi import FastAPI, UploadFile, File
 from fastapi.staticfiles import StaticFiles
+from fastapi.responses import HTMLResponse, JSONResponse
 from app.model import load_model, predict_from_bytes
+from app.inference import load_classification_model, classify_bytes
+import json, os
+
 
 
 # ──────────────────────────────────────────────
@@ -12,13 +13,13 @@ from app.model import load_model, predict_from_bytes
 # ──────────────────────────────────────────────
 app = FastAPI(title="NEMO Tools")
 
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+# app.add_middleware(
+#     CORSMiddleware,
+#     allow_origins=["*"],
+#     allow_credentials=True,
+#     allow_methods=["*"],
+#     allow_headers=["*"],
+# )
 
 # ──────────────────────────────────────────────
 # Static Frontend
@@ -42,6 +43,17 @@ print("🚀 Loading DINOv2 custom model...")
 model_device_tuple = load_model()
 print("✅ Model loaded and ready for inference!")
 
+# warm-up on startup
+load_classification_model()
+
+# --- Load classification model & labels once at startup ---
+MAP_PATH = os.path.join(os.path.dirname(__file__), "id2name.json")
+with open(MAP_PATH, "r") as f:
+    ID2NAME = json.load(f)
+
+cls_model = load_model()
+print("✅ Classification model loaded and ready for inference!")
+
 # ──────────────────────────────────────────────
 # API Endpoints
 # ──────────────────────────────────────────────
@@ -52,6 +64,11 @@ async def generate_attention(file: UploadFile = File(...)):
     result = predict_from_bytes(model_device_tuple, image_bytes)
     return result
 
+@app.post("/classify")
+async def classify(file: UploadFile = File(...)):
+    image_bytes = await file.read()
+    return classify_bytes(image_bytes)
+
 @app.get("/api")
 def api_root():
     return {"message": "NEMO Tools backend running."}

app/model.py

@@ -25,7 +25,8 @@ CKPT_PATH = hf_hub_download(
 )
 
 PATCH_SIZE = 14
-IMAGE_SIZE = (1000,1000)
+IMAGE_SIZE = (800,800)
+
 
 
 # -------------------------------------------------------
@@ -46,7 +47,6 @@ 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]:
@@ -81,10 +81,6 @@ def preprocess_image(image_bytes):
     img = img[:, :w, :h].unsqueeze(0)
     return img, (w, h)
 
-
-# -------------------------------------------------------
-# Prediction logic (generate attention map)
-# -------------------------------------------------------
 def predict_from_bytes(model_device_tuple, image_bytes):
     model, device = model_device_tuple
     img, (w, h) = preprocess_image(image_bytes)
@@ -96,6 +92,7 @@ def predict_from_bytes(model_device_tuple, image_bytes):
     attentions = model.get_last_self_attention(x)
     nh = attentions.shape[1]  # number of heads
 
+    # Reshape attention maps
     attentions = attentions[0, :, 0, 1:].reshape(nh, -1)
     attentions = attentions.reshape(nh, w_featmap, h_featmap)
     attentions = nn.functional.interpolate(
@@ -104,22 +101,34 @@ def predict_from_bytes(model_device_tuple, image_bytes):
         mode="nearest"
     )[0].cpu().numpy()
 
-    # Mean attention map
-    mean_attention = np.mean(attentions, axis=0)
+    # --- Normalize and visualize ---
+    all_heads_base64 = []
+    original_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
 
-    # Normalize to [0,1] for visualization
-    mean_attention_norm = (mean_attention - mean_attention.min()) / (
-        mean_attention.max() - mean_attention.min() + 1e-8
-    )
+    for i in range(nh):
+        head_attn = attentions[i]
+        head_norm = (head_attn - head_attn.min()) / (head_attn.max() - head_attn.min() + 1e-8)
+        heatmap = (cm.viridis(head_norm)[:, :, :3] * 255).astype(np.uint8)
+        heatmap_img = Image.fromarray(heatmap).resize(original_image.size, Image.BILINEAR)
 
-    # Apply colormap (viridis) and resize to match input image size
-    heatmap = (cm.viridis(mean_attention_norm)[:, :, :3] * 255).astype(np.uint8)
-    heatmap_img = Image.fromarray(heatmap).resize(IMAGE_SIZE)
+        buf = BytesIO()
+        heatmap_img.save(buf, format="PNG")
+        buf.seek(0)
+        head_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+        all_heads_base64.append(head_b64)
+
+    # --- Mean attention map ---
+    mean_attention = np.mean(attentions, axis=0)
+    mean_norm = (mean_attention - mean_attention.min()) / (mean_attention.max() - mean_attention.min() + 1e-8)
+    heatmap = (cm.viridis(mean_norm)[:, :, :3] * 255).astype(np.uint8)
+    mean_img = Image.fromarray(heatmap).resize(original_image.size, Image.BILINEAR)
 
-    # Convert to base64
     buf = BytesIO()
-    heatmap_img.save(buf, format="PNG")
+    mean_img.save(buf, format="PNG")
     buf.seek(0)
-    img_base64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+    mean_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
 
-    return {"attention_map": img_base64}
+    return {
+        "mean_attention_map": mean_b64,
+        "head_attention_maps": all_heads_base64,
+    }

app/static/Correctindex.html

@@ -0,0 +1,345 @@
+<!doctype html>
+<html lang="en">
+  <head>
+    <meta charset="utf-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1" />
+    <title>NEMO Tools</title>
+
+    <!-- TailwindCSS CDN -->
+    <script src="https://cdn.tailwindcss.com"></script>
+    <style>
+      body {
+        background-image: url('static/background.jpg');
+        background-size: cover;
+        background-position: center;
+        background-attachment: fixed;
+        background-repeat: no-repeat;
+        color: #f9fafb;
+      }
+      body::before {
+        content: "";
+        position: fixed;
+        top: 0;
+        left: 0;
+        width: 100%;
+        height: 100%;
+        background: rgba(0, 10, 20, 0.3);
+        z-index: -1;
+      }
+      .card {
+        background: rgba(255, 255, 255, 0.12);
+        backdrop-filter: blur(10px);
+        border-radius: 20px;
+        box-shadow: 0 4px 20px rgba(0, 0, 0, 0.2);
+        color: #f1f1f1;
+      }
+      .nav-link {
+        color: #e0e0e0;
+      }
+      .nav-link.active {
+        color: #60a5fa;
+        border-bottom: 2px solid #60a5fa;
+      }
+    </style>
+  </head>
+
+  <body class="bg-gray-100 min-h-screen">
+    <header class="bg-white shadow-sm">
+      <div class="max-w-6xl mx-auto px-4 py-4 flex items-center justify-between">
+        <div class="flex items-center gap-3">
+          <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>
+          </div>
+        </div>
+
+        <nav class="flex gap-3">
+          <button id="tab-research" class="tab-btn text-gray-500 hover:text-indigo-600 px-3 py-1 rounded-md text-sm font-medium" onclick="showTab('research')">
+            Research
+          </button>
+          <button id="tab-people" class="tab-btn text-gray-500 hover:text-indigo-600 px-3 py-1 rounded-md text-sm font-medium" onclick="showTab('people')">
+            People
+          </button>
+          <button id="tab-tools" class="tab-btn text-indigo-600 bg-indigo-50 px-3 py-1 rounded-md text-sm font-medium" onclick="showTab('tools')">
+            Tools
+          </button>
+        </nav>
+      </div>
+    </header>
+
+    <main class="max-w-6xl mx-auto px-4 py-8">
+      <section id="page-research" class="hidden">
+        <h2 class="text-2xl font-semibold text-gray-800 mb-4">Research</h2>
+        <p class="text-gray-500 mb-4">We can list publications, datasets, and experiment notes here later.</p>
+        <div class="bg-white rounded-xl shadow p-6 text-gray-400 text-sm">(placeholder) Add your papers, abstracts, or GitHub repos here.</div>
+      </section>
+
+      <section id="page-people" class="hidden">
+        <h2 class="text-2xl font-semibold text-gray-800 mb-4">People</h2>
+        <p class="text-gray-500 mb-4">We can add your name, collaborators, and links to profiles here.</p>
+        <div class="grid gap-4 md:grid-cols-3">
+          <div class="bg-white rounded-xl shadow p-5">
+            <h3 class="font-semibold text-gray-700">You</h3>
+            <p class="text-gray-400 text-sm">Project owner</p>
+          </div>
+          <div class="bg-white rounded-xl shadow p-5">
+            <h3 class="font-semibold text-gray-700">To add</h3>
+            <p class="text-gray-400 text-sm">Collaborators / advisors</p>
+          </div>
+          <div class="bg-white rounded-xl shadow p-5">
+            <h3 class="font-semibold text-gray-700">Contact</h3>
+            <p class="text-gray-400 text-sm">Add email / GitHub here</p>
+          </div>
+        </div>
+      </section>
+
+      <section id="page-tools">
+        <div class="bg-white shadow-lg rounded-2xl p-8 w-full">
+          <h2 class="text-2xl font-bold text-indigo-600 mb-6 flex items-center gap-2"><span>🧰 Tools</span></h2>
+          <div class="flex gap-3 mb-6 border-b pb-2">
+            <button id="sub-attention" class="subtab-btn text-indigo-600 font-medium border-b-2 border-indigo-600 pb-1" onclick="showSubTool('attention')">🧠 Attention Maps</button>
+            <button id="sub-classification" class="subtab-btn text-gray-500 hover:text-indigo-600 pb-1" onclick="showSubTool('classification')">🔍 Run Classification</button>
+          </div>
+
+          <!-- Attention Tool -->
+          <div id="tool-attention">
+            <div class="flex flex-col items-center gap-4 mb-6 justify-center">
+              <input id="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
+                      file:rounded-full file:border-0
+                      file:text-sm file:font-semibold
+                      file:bg-indigo-50 file:text-indigo-600
+                      hover:file:bg-indigo-100"
+                onchange="onAttentionImageSelected()" />
+            </div>
+
+            <div id="attention-extra" class="hidden flex flex-col items-center gap-6">
+              <button id="runButton" onclick="runActiveTool()" class="px-8 py-3 bg-indigo-600 text-white text-lg font-semibold rounded-full shadow-md hover:bg-indigo-700 transition">
+                ▶️ Run Attention
+              </button>
+              <div class="flex flex-col md:flex-row justify-center items-start gap-6">
+                <div class="flex flex-col items-center">
+                  <h4 class="text-gray-600 mb-2 font-medium">Original Image</h4>
+                  <img id="original" class="rounded-lg shadow-md max-w-xs hidden" />
+                </div>
+                <div class="flex flex-col items-center">
+                  <h4 class="text-gray-600 mb-2 font-medium">Mean Attention Map</h4>
+                  <img id="output" class="rounded-lg shadow-md max-w-xs hidden" />
+                </div>
+              </div>
+              <div id="headsContainer" class="hidden mt-8">
+                <h4 class="text-gray-600 mb-3 font-medium text-center">All Attention Heads</h4>
+                <div id="headsGrid" class="flex flex-wrap justify-center gap-3"></div>
+              </div>
+              <p id="status" class="text-center text-gray-500 mt-2 text-sm"></p>
+            </div>
+          </div>
+
+          <!-- Classification Tool -->
+          <div id="tool-classification" class="hidden">
+            <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
+                          file:rounded-full file:border-0
+                          file:text-sm file:font-semibold
+                          file:bg-indigo-50 file:text-indigo-600
+                          hover:file:bg-indigo-100" />
+              <button id="cls-run" onclick="runClassification()" style="display:none;"
+                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>
+            <div id="cls-result" class="text-center text-gray-700 mt-4 text-lg font-medium"></div>
+          </div>
+        </div>
+      </section>
+    </main>
+
+    <script>
+      function showTab(name) {
+        const tabs = ["research", "people", "tools"];
+        tabs.forEach(t => {
+          document.getElementById("page-" + t).classList.add("hidden");
+          document.getElementById("tab-" + t).classList.remove("bg-indigo-50", "text-indigo-600");
+          document.getElementById("tab-" + t).classList.add("text-gray-500");
+        });
+        document.getElementById("page-" + name).classList.remove("hidden");
+        document.getElementById("tab-" + name).classList.add("bg-indigo-50", "text-indigo-600");
+        document.getElementById("tab-" + name).classList.remove("text-gray-500");
+      }
+
+      let activeTool = "attention";
+      const runButton = () => document.getElementById("runButton");
+
+      function showSubTool(name) {
+        const subs = ["attention", "classification"];
+        subs.forEach(s => {
+          document.getElementById("tool-" + s).classList.add("hidden");
+          document.getElementById("sub-" + s).classList.remove("text-indigo-600", "font-medium", "border-b-2", "border-indigo-600");
+          document.getElementById("sub-" + s).classList.add("text-gray-500");
+        });
+        document.getElementById("tool-" + name).classList.remove("hidden");
+        document.getElementById("sub-" + name).classList.add("text-indigo-600", "font-medium", "border-b-2", "border-indigo-600");
+        document.getElementById("sub-" + name).classList.remove("text-gray-500");
+        activeTool = name;
+      }
+
+      async function runActiveTool() {
+        const btn = runButton();
+        btn.disabled = true;
+        btn.textContent = "⌛️ Running...";
+        btn.classList.add("opacity-70", "cursor-not-allowed");
+        try {
+          if (activeTool === "attention") await sendAttention();
+          else await runClassification();
+        } catch (err) {
+          console.error(err);
+          alert("❌ Error while running model: " + err.message);
+        }
+        btn.disabled = false;
+        btn.classList.remove("opacity-70", "cursor-not-allowed");
+        btn.textContent = activeTool === "attention" ? "▶️ Run Attention" : "▶️ Run Classification";
+      }
+
+      function onAttentionImageSelected() {
+        const fileInput = document.getElementById("file");
+        const extra = document.getElementById("attention-extra");
+        const original = document.getElementById("original");
+        if (fileInput.files.length > 0) {
+          extra.classList.remove("hidden");
+          const reader = new FileReader();
+          reader.onload = e => {
+            original.src = e.target.result;
+            original.classList.remove("hidden");
+          };
+          reader.readAsDataURL(fileInput.files[0]);
+        } else {
+          extra.classList.add("hidden");
+          original.classList.add("hidden");
+        }
+      }
+
+      async function sendAttention() {
+        const fileInput = document.getElementById("file");
+        const file = fileInput.files[0];
+        const output = document.getElementById("output");
+        const status = document.getElementById("status");
+        const headsContainer = document.getElementById("headsContainer");
+        const headsGrid = document.getElementById("headsGrid");
+        if (!file) {
+          alert("Please choose an image first!");
+          return;
+        }
+        output.classList.add("hidden");
+        headsContainer.classList.add("hidden");
+        headsGrid.innerHTML = "";
+        status.textContent = "⌛️ Model is running...";
+        const fd = new FormData();
+        fd.append("file", file);
+        try {
+          const res = await fetch("/attention", { method: "POST", body: fd });
+          if (!res.ok) throw new Error(`Server error: ${res.status}`);
+          const json = await res.json();
+          output.src = "data:image/png;base64," + json.mean_attention_map;
+          output.classList.remove("hidden");
+          if (json.head_attention_maps && json.head_attention_maps.length > 0) {
+            json.head_attention_maps.forEach((headB64, i) => {
+              const img = document.createElement("img");
+              img.src = "data:image/png;base64," + headB64;
+              img.alt = `Head ${i + 1}`;
+              img.className = "rounded-md shadow-sm";
+              img.style.width = "120px";
+              img.style.transition = "transform 0.2s";
+              img.onmouseenter = () => (img.style.transform = "scale(1.1)");
+              img.onmouseleave = () => (img.style.transform = "scale(1)");
+              headsGrid.appendChild(img);
+            });
+            headsContainer.classList.remove("hidden");
+          }
+          status.textContent = "✅ Done!";
+        } catch (err) {
+          console.error(err);
+          status.textContent = "❌ Error: " + err.message;
+        }
+      }
+
+      async function runClassification() {
+        const fileInput = document.getElementById("cls-file");
+        const file = fileInput.files[0];
+        const result = document.getElementById("cls-result");
+        const btn = document.getElementById("cls-run");
+
+        if (!file) return alert("Please choose an image to classify!");
+
+        // Disable button + show "Running..."
+        btn.disabled = true;
+        btn.textContent = "⏳ Running...";
+        btn.classList.add("opacity-70", "cursor-not-allowed");
+
+        result.textContent = "⌛️ Model is running...";
+
+        const fd = new FormData();
+        fd.append("file", file);
+
+        try {
+          const res = await fetch("/classify", { method: "POST", body: fd });
+          if (!res.ok) throw new Error(`Server error: ${res.status}`);
+          const json = await res.json();
+
+          if (json.top5) {
+            result.innerHTML = `
+              <h3 class="font-semibold text-indigo-600 mb-2">Top-5 Predictions</h3>
+              ${json.top5.map(p => `
+                <div class="flex justify-between border-b py-1">
+                  <span>${p.rank}. ${p.label}</span>
+                  <span class="text-gray-500">${(p.score * 100).toFixed(2)}%</span>
+                </div>
+              `).join("")}
+            `;
+          } else {
+            result.textContent = "No predictions returned.";
+          }
+
+          if (json.plot) {
+            const plotImg = document.createElement("img");
+            plotImg.src = json.plot;
+            plotImg.alt = "Top-3 Predicted Species";
+            plotImg.style.display = "block";
+            plotImg.style.margin = "20px auto";
+            plotImg.style.maxWidth = "800px";
+            result.appendChild(plotImg);
+          }
+        } catch (err) {
+          console.error(err);
+          result.textContent = "❌ Error: " + err.message;
+        } finally {
+          // Restore button state
+          btn.disabled = false;
+          btn.classList.remove("opacity-70", "cursor-not-allowed");
+          btn.textContent = "▶️ Run Classification";
+        }
+      }
+
+
+      // Show classification button only after image is selected
+      document.addEventListener("DOMContentLoaded", () => {
+        const fileInput = document.getElementById("cls-file");
+        const runBtn = document.getElementById("cls-run");
+        fileInput.addEventListener("change", () => {
+          if (fileInput.files && fileInput.files.length > 0) {
+            runBtn.style.display = "inline-block";
+          } else {
+            runBtn.style.display = "none";
+          }
+        });
+      });
+
+      showTab("tools");
+      showSubTool("attention");
+    </script>
+  </body>
+</html>

app/static/index.html

@@ -1,364 +1,266 @@
 <!doctype html>
 <html lang="en">
-  <head>
-    <meta charset="utf-8" />
-    <meta name="viewport" content="width=device-width, initial-scale=1" />
-    <title>NEMO Tools</title>
-
-    <!-- TailwindCSS CDN -->
-    <script src="https://cdn.tailwindcss.com"></script>
-    <style>
-      body {
-        background-image: url('/background.jpg');
-        background-size: cover;
-        background-position: center;
-        background-attachment: fixed;
-        background-repeat: no-repeat;
-        color: #f9fafb; /* soft white text for better contrast */
-      }
-    
-      /* Add slight overlay to make text pop */
-      body::before {
-        content: "";
-        position: fixed;
-        top: 0;
-        left: 0;
-        width: 100%;
-        height: 100%;
-        background: rgba(0, 10, 20, 0.3); /* deep ocean tint */
-        z-index: -1;
-      }
-    
-      .card {
-        background: rgba(255, 255, 255, 0.12);
-        backdrop-filter: blur(10px);
-        border-radius: 20px;
-        box-shadow: 0 4px 20px rgba(0, 0, 0, 0.2);
-        color: #f1f1f1;
-      }
-    
-      .nav-link {
-        color: #e0e0e0;
-      }
-    
-      .nav-link.active {
-        color: #60a5fa;
-        border-bottom: 2px solid #60a5fa;
-      }
-    </style>
-    
-  </head>
-
-  <body class="bg-gray-100 min-h-screen">
-    <!-- Header -->
-    <header class="bg-white shadow-sm">
-      <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="/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>
-          </div>
+<head>
+  <meta charset="utf-8" />
+  <meta name="viewport" content="width=device-width, initial-scale=1" />
+  <title>NEMO Tools</title>
+
+  <!-- TailwindCSS -->
+  <script src="https://cdn.tailwindcss.com"></script>
+  <style>
+    body {
+      background-image: url('/static/background.jpg');
+      background-size: cover;
+      background-position: center;
+      background-attachment: fixed;
+      background-repeat: no-repeat;
+      color: #f9fafb;
+    }
+    body::before {
+      content: "";
+      position: fixed;
+      top: 0; left: 0;
+      width: 100%; height: 100%;
+      background: rgba(0, 10, 20, 0.3);
+      z-index: -1;
+    }
+  </style>
+</head>
+
+<body class="bg-gray-100 min-h-screen">
+
+  <!-- Header -->
+  <header class="bg-white shadow-sm">
+    <div class="max-w-6xl mx-auto px-4 py-4 flex items-center justify-between">
+      <div class="flex items-center gap-3">
+        <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>
         </div>
-
-        <!-- Top navigation -->
-        <nav class="flex gap-3">
-          <button id="tab-research" class="tab-btn text-gray-500 hover:text-indigo-600 px-3 py-1 rounded-md text-sm font-medium" onclick="showTab('research')">
-            Research
-          </button>
-          <button id="tab-people" class="tab-btn text-gray-500 hover:text-indigo-600 px-3 py-1 rounded-md text-sm font-medium" onclick="showTab('people')">
-            People
-          </button>
-          <button id="tab-tools" class="tab-btn text-indigo-600 bg-indigo-50 px-3 py-1 rounded-md text-sm font-medium" onclick="showTab('tools')">
-            Tools
-          </button>
-        </nav>
       </div>
-    </header>
-
-    <!-- Main content -->
-    <main class="max-w-6xl mx-auto px-4 py-8">
-      <!-- Research tab -->
-      <section id="page-research" class="hidden">
-        <h2 class="text-2xl font-semibold text-gray-800 mb-4">Research</h2>
-        <p class="text-gray-500 mb-4">
-          We can list publications, datasets, and experiment notes here later.
-        </p>
-        <div class="bg-white rounded-xl shadow p-6 text-gray-400 text-sm">
-          (placeholder) Add your papers, abstracts, or GitHub repos here.
-        </div>
-      </section>
-
-      <!-- People tab -->
-      <section id="page-people" class="hidden">
-        <h2 class="text-2xl font-semibold text-gray-800 mb-4">People</h2>
-        <p class="text-gray-500 mb-4">
-          We can add your name, collaborators, and links to profiles here.
-        </p>
-        <div class="grid gap-4 md:grid-cols-3">
-          <div class="bg-white rounded-xl shadow p-5">
-            <h3 class="font-semibold text-gray-700">You</h3>
-            <p class="text-gray-400 text-sm">Project owner</p>
-          </div>
-          <div class="bg-white rounded-xl shadow p-5">
-            <h3 class="font-semibold text-gray-700">To add</h3>
-            <p class="text-gray-400 text-sm">Collaborators / advisors</p>
-          </div>
-          <div class="bg-white rounded-xl shadow p-5">
-            <h3 class="font-semibold text-gray-700">Contact</h3>
-            <p class="text-gray-400 text-sm">Add email / GitHub here</p>
-          </div>
+      <nav class="flex gap-3">
+        <button id="tab-research" class="tab-btn text-gray-500 hover:text-indigo-600 px-3 py-1 rounded-md text-sm font-medium" onclick="showTab('research')">Research</button>
+        <button id="tab-people" class="tab-btn text-gray-500 hover:text-indigo-600 px-3 py-1 rounded-md text-sm font-medium" onclick="showTab('people')">People</button>
+        <button id="tab-tools" class="tab-btn text-indigo-600 bg-indigo-50 px-3 py-1 rounded-md text-sm font-medium" onclick="showTab('tools')">Tools</button>
+      </nav>
+    </div>
+  </header>
+
+  <!-- Main Content -->
+  <main class="max-w-6xl mx-auto px-4 py-8">
+    <section id="page-tools">
+      <div class="bg-white shadow-lg rounded-2xl p-8 w-full">
+        <h2 class="text-2xl font-bold text-indigo-600 mb-6 flex items-center gap-2">🧰 Tools</h2>
+
+        <!-- Subtabs -->
+        <div class="flex gap-3 mb-6 border-b pb-2">
+          <button id="sub-attention" class="subtab-btn text-indigo-600 font-medium border-b-2 border-indigo-600 pb-1" onclick="showSubTool('attention')">🧠 Mean Attention Map</button>
+          <button id="sub-classification" class="subtab-btn text-gray-500 hover:text-indigo-600 pb-1" onclick="showSubTool('classification')">🔍 Run Classification</button>
         </div>
-      </section>
 
-      <!-- Tools tab -->
-      <section id="page-tools">
-        <div class="bg-white shadow-lg rounded-2xl p-8 w-full">
-          <h2 class="text-2xl font-bold text-indigo-600 mb-6 flex items-center gap-2">
-            <span>🧰 Tools</span>
-          </h2>
-
-          <!-- Sub-tab buttons -->
-          <div class="flex gap-3 mb-6 border-b pb-2">
-            <button id="sub-attention" class="subtab-btn text-indigo-600 font-medium border-b-2 border-indigo-600 pb-1"
-                    onclick="showSubTool('attention')">
-              🧠 Mean Attention Map
-            </button>
-            <button id="sub-classification" class="subtab-btn text-gray-500 hover:text-indigo-600 pb-1"
-                    onclick="showSubTool('classification')">
-              🔍 Run Classification
-            </button>
-          </div>
-
-
-          <!-- Attention Map tool -->
-          <div id="tool-attention">
-            <!-- 1) always visible file input -->
-            <div class="flex flex-col items-center gap-4 mb-6 justify-center">
-              <input id="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
-                      file:rounded-full file:border-0
-                      file:text-sm file:font-semibold
-                      file:bg-indigo-50 file:text-indigo-600
-                      hover:file:bg-indigo-100"
-                onchange="onAttentionImageSelected()" />
+        <!-- 🧠 Attention Tool -->
+        <div id="tool-attention">
+          <div class="flex flex-col items-center gap-4 mb-6 justify-center">
+            <div class="w-full max-w-md mx-auto text-center border-2 border-dashed border-gray-300 rounded-xl p-6 bg-gray-50 hover:bg-gray-100 transition">
+              <p class="text-gray-700 font-semibold mb-2">Upload Image</p>
+              <p class="text-gray-400 text-sm mb-4">Supported formats: JPG, PNG</p>
+              <label for="file" class="inline-block bg-indigo-600 hover:bg-indigo-700 text-white font-semibold py-2 px-6 rounded-full cursor-pointer shadow-md transition">Choose File</label>
+              <input id="file" type="file" accept="image/*" class="hidden" onchange="onAttentionFileSelected()" />
+              <p id="attention-filename" class="text-gray-500 mt-3 text-sm"></p>
             </div>
+          </div>
 
-            <!-- 2) hidden until user selects an image -->
-            <div id="attention-extra" class="hidden flex flex-col items-center gap-6">
-              <!-- run button -->
-              <button id="runButton"
-                onclick="runActiveTool()"
-                class="px-8 py-3 bg-indigo-600 text-white text-lg font-semibold rounded-full shadow-md hover:bg-indigo-700 transition">
-                ▶️ Run Attention
-              </button>
+          <div id="attention-extra" class="hidden flex flex-col items-center gap-6">
+            <button id="runButton" onclick="runActiveTool()" class="px-8 py-3 bg-indigo-600 text-white text-lg font-semibold rounded-full shadow-md hover:bg-indigo-700 transition">▶️ Run Attention</button>
 
-              <!-- image area -->
-              <div class="flex flex-col md:flex-row justify-center items-start gap-6">
-                <div class="flex flex-col items-center">
-                  <h4 class="text-gray-600 mb-2 font-medium">Original Image</h4>
-                  <img id="original" class="rounded-lg shadow-md max-w-xs hidden" />
-                </div>
-                <div class="flex flex-col items-center">
-                  <h4 class="text-gray-600 mb-2 font-medium">Mean Attention Map</h4>
-                  <!-- start empty, we'll fill after run -->
-                  <img id="output" class="rounded-lg shadow-md max-w-xs hidden" />
-                </div>
+            <div class="flex flex-col md:flex-row justify-center items-start gap-6">
+              <div class="flex flex-col items-center">
+                <h4 class="text-gray-600 mb-2 font-medium">Original Image</h4>
+                <img id="original" class="rounded-lg shadow-md max-w-xs hidden" />
+              </div>
+              <div class="flex flex-col items-center">
+                <h4 class="text-gray-600 mb-2 font-medium">Mean Attention Map</h4>
+                <img id="output" class="rounded-lg shadow-md max-w-xs hidden" />
               </div>
-
-              <p id="status" class="text-center text-gray-500 mt-2 text-sm"></p>
             </div>
-          </div>
 
-          <!-- Classification tool -->
-          <!-- 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">
-              <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
-                          file:rounded-full file:border-0
-                          file:text-sm file:font-semibold
-                          file:bg-indigo-50 file:text-indigo-600
-                          hover:file:bg-indigo-100" />
-              
-              <!-- ✅ Add this button -->
-              <button id="cls-run-btn"
-                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 id="headsContainer" class="hidden mt-8">
+              <h4 class="text-gray-600 mb-3 font-medium text-center">All Attention Heads</h4>
+              <div id="headsGrid" class="flex flex-wrap justify-center gap-3"></div>
             </div>
-
-            <div id="cls-result" class="text-center text-gray-700 mt-4 text-lg font-medium"></div>
+            <p id="status" class="text-center text-gray-500 mt-2 text-sm"></p>
           </div>
-
-          
-          
-            <!-- <div id="cls-result" class="text-center text-gray-700 mt-4 text-lg font-medium"></div> -->
-          </div>          
         </div>
-      </section>
-    </main>
-
-    <!-- JS -->
-    <script>
-      // --- top navigation ---
-      function showTab(name) {
-        const tabs = ["research", "people", "tools"];
-        tabs.forEach(t => {
-          document.getElementById("page-" + t).classList.add("hidden");
-          document.getElementById("tab-" + t).classList.remove("bg-indigo-50", "text-indigo-600");
-          document.getElementById("tab-" + t).classList.add("text-gray-500");
-        });
-        document.getElementById("page-" + name).classList.remove("hidden");
-        document.getElementById("tab-" + name).classList.add("bg-indigo-50", "text-indigo-600");
-        document.getElementById("tab-" + name).classList.remove("text-gray-500");
-      }
 
-      // --- tools logic ---
-      let activeTool = "attention";
-      const runButton = () => document.getElementById("runButton");
-
-      function showSubTool(name) {
-        const subs = ["attention", "classification"];
-        subs.forEach(s => {
-          document.getElementById("tool-" + s).classList.add("hidden");
-          document.getElementById("sub-" + s).classList.remove("text-indigo-600", "font-medium", "border-b-2", "border-indigo-600");
-          document.getElementById("sub-" + s).classList.add("text-gray-500");
-        });
-        document.getElementById("tool-" + name).classList.remove("hidden");
-        document.getElementById("sub-" + name).classList.add("text-indigo-600", "font-medium", "border-b-2", "border-indigo-600");
-        document.getElementById("sub-" + name).classList.remove("text-gray-500");
+        <!-- 🔍 Classification Tool -->
+        <div id="tool-classification" class="hidden flex flex-col items-center">
+          <div class="w-full max-w-md mx-auto text-center border-2 border-dashed border-gray-300 rounded-xl p-6 bg-gray-50 hover:bg-gray-100 transition">
+            <p class="text-gray-700 font-semibold mb-2">Upload Image</p>
+            <p class="text-gray-400 text-sm mb-4">Supported formats: JPG, PNG</p>
+            <label for="cls-file" class="inline-block bg-green-600 hover:bg-green-700 text-white font-semibold py-2 px-6 rounded-full cursor-pointer shadow-md transition">Choose File</label>
+            <input id="cls-file" type="file" accept="image/*" class="hidden" onchange="onClsFileSelected()" />
+            <p id="cls-filename" class="text-gray-500 mt-3 text-sm"></p>
+          </div>
 
-        activeTool = name;
+          <button id="cls-run" onclick="runClassification()" style="display:none;" class="mt-5 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>
 
-        const btn = runButton();
-        if (name === "attention") {
-          btn.textContent = "▶️ Run Attention";
-          btn.className = "px-8 py-3 bg-indigo-600 text-white text-lg font-semibold rounded-full shadow-md hover:bg-indigo-700 transition";
-        } else {
-          btn.textContent = "▶️ Run Classification";
-          btn.className = "px-8 py-3 bg-green-600 text-white text-lg font-semibold rounded-full shadow-md hover:bg-green-700 transition";
-        }
+          <div id="cls-result" class="text-center text-gray-700 mt-6 text-lg font-medium"></div>
+        </div>
+      </div>
+    </section>
+  </main>
+
+  <!-- Scripts -->
+  <script>
+    let activeTool = "attention";
+
+    function showSubTool(name) {
+      const subs = ["attention", "classification"];
+      subs.forEach(s => {
+        document.getElementById("tool-" + s).classList.add("hidden");
+        document.getElementById("sub-" + s).classList.remove("text-indigo-600", "font-medium", "border-b-2", "border-indigo-600");
+        document.getElementById("sub-" + s).classList.add("text-gray-500");
+      });
+      document.getElementById("tool-" + name).classList.remove("hidden");
+      document.getElementById("sub-" + name).classList.add("text-indigo-600", "font-medium", "border-b-2", "border-indigo-600");
+      document.getElementById("sub-" + name).classList.remove("text-gray-500");
+      activeTool = name;
+    }
+
+    // 🧠 Attention Tool
+    function onAttentionFileSelected() {
+      const fileInput = document.getElementById("file");
+      const extra = document.getElementById("attention-extra");
+      const original = document.getElementById("original");
+      const nameEl = document.getElementById("attention-filename");
+
+      if (fileInput.files.length > 0) {
+        extra.classList.remove("hidden");
+        const file = fileInput.files[0];
+        nameEl.textContent = file.name;
+
+        const reader = new FileReader();
+        reader.onload = e => {
+          original.src = e.target.result;
+          original.classList.remove("hidden");
+        };
+        reader.readAsDataURL(file);
+      } else {
+        extra.classList.add("hidden");
+        original.classList.add("hidden");
+        nameEl.textContent = "";
       }
-
-      async function runActiveTool() {
-        const btn = runButton();
-        btn.disabled = true;
-        btn.textContent = "⌛️ Running...";
-        btn.classList.add("opacity-70", "cursor-not-allowed");
-
-        try {
-          if (activeTool === "attention") await sendAttention();
-          else await runClassification();
-        } catch (err) {
-          console.error(err);
-          alert("❌ Error while running model: " + err.message);
+    }
+
+    async function sendAttention() {
+      const file = document.getElementById("file").files[0];
+      const output = document.getElementById("output");
+      const status = document.getElementById("status");
+      const headsContainer = document.getElementById("headsContainer");
+      const headsGrid = document.getElementById("headsGrid");
+      const btn = document.getElementById("runButton");
+
+      if (!file) return alert("Please choose an image first!");
+      output.classList.add("hidden");
+      headsContainer.classList.add("hidden");
+      headsGrid.innerHTML = "";
+      // status.textContent = "⌛️ Model is running...";
+
+      btn.disabled = true;
+      btn.textContent = "⏳ Running...";
+      btn.classList.add("opacity-70", "cursor-not-allowed");
+
+      const fd = new FormData();
+      fd.append("file", file);
+
+      try {
+        const res = await fetch("/attention", { method: "POST", body: fd });
+        const json = await res.json();
+        output.src = "data:image/png;base64," + json.mean_attention_map;
+        output.classList.remove("hidden");
+        if (json.head_attention_maps) {
+          json.head_attention_maps.forEach((b64, i) => {
+            const img = document.createElement("img");
+            img.src = "data:image/png;base64," + b64;
+            img.className = "rounded-md shadow-sm w-[120px]";
+            headsGrid.appendChild(img);
+          });
+          headsContainer.classList.remove("hidden");
         }
-
+        status.textContent = "✅ Done!";
+      } catch (err) {
+        status.textContent = "❌ Error: " + err.message;
+      } finally {
         btn.disabled = false;
+        btn.textContent = "▶️ Run Attention";
         btn.classList.remove("opacity-70", "cursor-not-allowed");
-        btn.textContent = activeTool === "attention" ? "▶️ Run Attention" : "▶️ Run Classification";
       }
-
-      // --- Attention Map Tool ---
-      // called when the user picks a file in the Attention tool
-      function onAttentionImageSelected() {
-        const fileInput = document.getElementById("file");
-        const extra = document.getElementById("attention-extra");
-        const original = document.getElementById("original");
-
-        if (fileInput.files.length > 0) {
-          // show the extra area (run button + image boxes)
-          extra.classList.remove("hidden");
-
-          // preview original immediately
-          const reader = new FileReader();
-          reader.onload = e => {
-            original.src = e.target.result;
-            original.classList.remove("hidden");
-          };
-          reader.readAsDataURL(fileInput.files[0]);
-        } else {
-          extra.classList.add("hidden");
-          original.classList.add("hidden");
-        }
+    }
+
+    // 🔍 Classification Tool
+    function onClsFileSelected() {
+      const fileInput = document.getElementById("cls-file");
+      const fileName = document.getElementById("cls-filename");
+      const runBtn = document.getElementById("cls-run");
+      if (fileInput.files.length > 0) {
+        fileName.textContent = fileInput.files[0].name;
+        runBtn.style.display = "block";
+      } else {
+        fileName.textContent = "";
+        runBtn.style.display = "none";
       }
-
-      // unified run button will call this when attention is active
-      async function sendAttention() {
-        const fileInput = document.getElementById("file");
-        const file = fileInput.files[0];
-        const output = document.getElementById("output");
-        const status = document.getElementById("status");
-
-        if (!file) {
-          alert("Please choose an image first!");
-          return;
-        }
-
-        // show loading
-        status.textContent = "⌛️ Model is running...";
-        output.classList.add("hidden");
-
-        const fd = new FormData();
-        fd.append("file", file);
-
-        try {
-          const res = await fetch("/attention", { method: "POST", body: fd });
-          if (!res.ok) throw new Error(`Server error: ${res.status}`);
-          const json = await res.json();
-
-          // show result
-          output.src = "data:image/png;base64," + json.attention_map;
-          output.classList.remove("hidden");
-          status.textContent = "✅ Done!";
-        } catch (err) {
-          console.error(err);
-          status.textContent = "❌ Error: " + err.message;
+    }
+
+    async function runClassification() {
+      const file = document.getElementById("cls-file").files[0];
+      const result = document.getElementById("cls-result");
+      const btn = document.getElementById("cls-run");
+
+      if (!file) return alert("Please choose an image to classify!");
+
+      btn.disabled = true;
+      btn.textContent = "⏳ Running...";
+      btn.classList.add("opacity-70", "cursor-not-allowed");
+      // result.textContent = "⌛️ Model is running...";
+
+      const fd = new FormData();
+      fd.append("file", file);
+
+      try {
+        const res = await fetch("/classify", { method: "POST", body: fd });
+        const json = await res.json();
+        if (json.top5) {
+          result.innerHTML = `
+            <h3 class="font-semibold text-indigo-600 mb-2">Top-5 Predictions</h3>
+            ${json.top5.map(p => `
+              <div class="flex justify-between border-b py-1">
+                <span>${p.rank}. ${p.label}</span>
+                <span class="text-gray-500">${(p.score * 100).toFixed(2)}%</span>
+              </div>`).join("")}`;
+        } else result.textContent = "No predictions returned.";
+
+        if (json.plot) {
+          const plotImg = document.createElement("img");
+          plotImg.src = json.plot;
+          plotImg.className = "block mx-auto mt-6 max-w-2xl";
+          result.appendChild(plotImg);
         }
+      } catch (err) {
+        result.textContent = "❌ Error: " + err.message;
+      } finally {
+        btn.disabled = false;
+        btn.textContent = "▶️ Run Classification";
+        btn.classList.remove("opacity-70", "cursor-not-allowed");
       }
+    }
 
+    async function runActiveTool() {
+      if (activeTool === "attention") await sendAttention();
+      else await runClassification();
+    }
 
-      // --- Classification Tool ---
-      async function runClassification() {
-        const fileInput = document.getElementById("cls-file");
-        const file = fileInput.files[0];
-        const result = document.getElementById("cls-result");
-
-        if (!file) return alert("Please choose an image to classify!");
-
-        result.textContent = "⌛️ Model is running...";
-
-        const fd = new FormData();
-        fd.append("file", file);
-
-        try {
-          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();
-
-          // ✅ 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;
-        }
-      }
-
-      // Initialize default tab
-      showTab("tools");
-      showSubTool("attention");
-    </script>
-  </body>
+    showSubTool("attention");
+  </script>
+</body>
 </html>

app/test.py

@@ -0,0 +1,10 @@
+from safetensors.torch import load_file
+from pprint import pprint
+import os
+
+path = os.path.expanduser("~/.cache/huggingface/hub/models--Arew99--dinov2-costum/snapshots/055a10af249d426a5b9a6ac07550f011e5739bbf/model.safetensors")
+print(f"Loading checkpoint: {path}")
+sd = load_file(path)
+print(f"✅ Loaded {len(sd)} tensors")
+print("Sample keys:")
+pprint(list(sd.keys())[:30])