Upload app.py with huggingface_hub

app.py

@@ -1,17 +1,11 @@
 import os
 import time
-import io
-import json
 import pickle
-import hashlib
-import random
 import warnings
-from pathlib import Path
-from typing import Dict, List, Optional, Tuple, Any
+from typing import Dict, Optional, Tuple
 
 import cv2
 import numpy as np
-import pandas as pd
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -20,6 +14,7 @@ import timm
 import albumentations as A
 from albumentations.pytorch import ToTensorV2
 from pytorch_grad_cam import GradCAMPlusPlus
+from pytorch_grad_cam.utils.image import show_cam_on_image
 import gradio as gr
 
 warnings.filterwarnings("ignore")
@@ -33,14 +28,7 @@ NUM_CLASSES = 5
 BACKBONE = "tf_efficientnetv2_m"
 FV_BACKBONE = "convnext_tiny.fb_in22k"
 
-GRADE_MAP = {
-    0: "No DR",
-    1: "Mild DR",
-    2: "Moderate DR",
-    3: "Severe DR",
-    4: "Proliferative DR",
-}
-
+GRADE_MAP = {0: "No DR", 1: "Mild DR", 2: "Moderate DR", 3: "Severe DR", 4: "Proliferative DR"}
 GRADE_COLORS = ["#2ecc71", "#f1c40f", "#e67e22", "#e74c3c", "#8e44ad"]
 SEVERITY_ICONS = ["🟢", "🟡", "🟠", "🔴", "🟣"]
 CLINICAL_ACTION = [
@@ -54,6 +42,9 @@ CLINICAL_ACTION = [
 IMAGENET_MEAN = [0.485, 0.456, 0.406]
 IMAGENET_STD = [0.229, 0.224, 0.225]
 
+STRONG_HEUR = {"red_dominance": 0.55, "disc_coverage": 0.25, "edge_density": 0.008}
+WEAK_HEUR   = {"red_dominance": 0.30, "disc_coverage": 0.12, "edge_density": 0.006}
+
 # =============================================================================
 # Model Definitions
 # =============================================================================
@@ -63,36 +54,34 @@ class GeM(nn.Module):
         super().__init__()
         self.p = nn.Parameter(torch.tensor(p))
         self.eps = eps
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return F.adaptive_avg_pool2d(x.clamp(min=self.eps).pow(self.p), 1).pow(1.0 / self.p)
 
+
 class CoralHead(nn.Module):
     def __init__(self, in_features: int, num_classes: int):
         super().__init__()
-        self.num_classes = num_classes
         self.linear = nn.Linear(in_features, 1, bias=False)
         self.bias = nn.Parameter(torch.zeros(num_classes - 1))
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        logit = self.linear(x)
-        return logit + self.bias
+        return self.linear(x) + self.bias
+
 
 class DRHead(nn.Module):
     def __init__(self, in_features: int, num_classes: int = 5, dropout: float = 0.3):
         super().__init__()
         self.net = nn.Sequential(
-            nn.Linear(in_features, 512),
-            nn.BatchNorm1d(512),
-            nn.SiLU(inplace=True),
-            nn.Dropout(dropout),
-            nn.Linear(512, 256),
-            nn.BatchNorm1d(256),
-            nn.SiLU(inplace=True),
-            nn.Dropout(dropout / 2),
+            nn.Linear(in_features, 512), nn.BatchNorm1d(512), nn.SiLU(inplace=True), nn.Dropout(dropout),
+            nn.Linear(512, 256), nn.BatchNorm1d(256), nn.SiLU(inplace=True), nn.Dropout(dropout / 2),
             nn.Linear(256, num_classes),
         )
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.net(x)
 
+
 class DRModel(nn.Module):
     def __init__(self, backbone: str = BACKBONE, num_classes: int = 5, pretrained: bool = False):
         super().__init__()
@@ -101,303 +90,369 @@ class DRModel(nn.Module):
         self.pool = GeM(p=3.0)
         self.head_cls = DRHead(in_feat, num_classes)
         self.head_coral = CoralHead(in_feat, num_classes)
+
     def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
-        feat = self.backbone(x)
-        feat = self.pool(feat).flatten(1)
+        feat = self.pool(self.backbone(x)).flatten(1)
         return {"logits": self.head_cls(feat), "coral": self.head_coral(feat)}
 
+
 class FundusValidator(nn.Module):
-    def __init__(self, pretrained: bool = False, backbone_name: str = FV_BACKBONE):
+    def __init__(self, backbone_name: str = FV_BACKBONE):
         super().__init__()
-        self.backbone = timm.create_model(backbone_name, pretrained=pretrained, num_classes=0, global_pool="avg")
-        for p in self.backbone.parameters(): p.requires_grad = False
+        self.backbone = timm.create_model(backbone_name, pretrained=False, num_classes=0, global_pool="avg")
+        for p in self.backbone.parameters():
+            p.requires_grad = False
         with torch.no_grad():
             feat_dim = self.backbone(torch.zeros(1, 3, 224, 224)).shape[1]
-        self.head = nn.Sequential(
-            nn.Linear(feat_dim, 256),
-            nn.GELU(),
-            nn.Dropout(0.3),
-            nn.Linear(256, 2),
-        )
-    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
-        with torch.no_grad(): return self.backbone(x)
+        self.head = nn.Sequential(nn.Linear(feat_dim, 256), nn.GELU(), nn.Dropout(0.3), nn.Linear(256, 2))
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        feats = self.extract_features(x)
+        with torch.no_grad():
+            feats = self.backbone(x)
         return self.head(feats)
 
+
 # =============================================================================
-# Preprocessing Utilities
+# Preprocessing
 # =============================================================================
-
 _clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
 
+
 def apply_clahe_lab(rgb: np.ndarray) -> np.ndarray:
     lab = cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB)
     lab[:, :, 0] = _clahe.apply(lab[:, :, 0])
     return cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
 
+
 def retinal_mask(rgb: np.ndarray) -> np.ndarray:
     gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
     _, m = cv2.threshold(gray, 15, 255, cv2.THRESH_BINARY)
-    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
-    m = cv2.morphologyEx(m, cv2.MORPH_CLOSE, kernel)
-    m = cv2.morphologyEx(m, cv2.MORPH_OPEN, kernel)
-    return m
+    k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
+    return cv2.morphologyEx(cv2.morphologyEx(m, cv2.MORPH_CLOSE, k), cv2.MORPH_OPEN, k)
+
 
 def crop_retinal_disc(rgb: np.ndarray, pad: int = 10) -> np.ndarray:
-    m = retinal_mask(rgb)
-    contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    if not contours: return rgb
+    contours, _ = cv2.findContours(retinal_mask(rgb), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours:
+        return rgb
     x, y, w, h = cv2.boundingRect(max(contours, key=cv2.contourArea))
     x, y = max(0, x - pad), max(0, y - pad)
     x2, y2 = min(rgb.shape[1], x + w + 2 * pad), min(rgb.shape[0], y + h + 2 * pad)
     return rgb[y:y2, x:x2]
 
-def ben_graham(rgb: np.ndarray, sigma_frac: float = 10.0, weight: float = 4.0, bias: int = 128) -> np.ndarray:
-    h, w = rgb.shape[:2]
-    sigma = max(1, int(min(h, w) / sigma_frac))
-    blurred = cv2.GaussianBlur(rgb, (0, 0), sigmaX=sigma, sigmaY=sigma)
-    out = cv2.addWeighted(rgb, weight, blurred, -weight, bias)
-    return np.clip(out, 0, 255).astype(np.uint8)
+
+def fundus_heuristics(rgb: np.ndarray) -> Dict[str, float]:
+    r, g, b = rgb[..., 0].astype(np.float32), rgb[..., 1].astype(np.float32), rgb[..., 2].astype(np.float32)
+    bright = (r > 15) | (g > 15) | (b > 15)
+    n_bright = int(bright.sum())
+    red_dom = float(((r > g) & (r > b))[bright].mean()) if n_bright > 100 else float(((r > g) & (r > b)).mean())
+    disc_cov = float(retinal_mask(rgb).mean() / 255.0)
+    gx = cv2.Sobel(cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY), cv2.CV_32F, 1, 0, ksize=3)
+    gy = cv2.Sobel(cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY), cv2.CV_32F, 0, 1, ksize=3)
+    mag = np.sqrt(gx**2 + gy**2)
+    edge_den = float(mag[bright].mean() / 255.0) if n_bright > 100 else float(mag.mean() / 255.0)
+    return {"red_dominance": red_dom, "disc_coverage": disc_cov, "edge_density": edge_den}
+
 
 # =============================================================================
-# Inference Helpers
+# Model Loading
 # =============================================================================
-
-def load_models():
-    # DR Model
-    model = DRModel(BACKBONE, NUM_CLASSES, pretrained=False).to(DEVICE)
+def load_dr_model():
+    m = DRModel(BACKBONE, NUM_CLASSES, pretrained=False).to(DEVICE)
     ckpt = torch.load("final_model.pt", map_location=DEVICE)
-    model.load_state_dict(ckpt["model_state"], strict=False)
-    model.eval()
+    state = ckpt.get("model_state", ckpt)
+    m.load_state_dict({k: v for k, v in state.items() if k in m.state_dict()}, strict=False)
+    return m.eval()
+
 
-    # Fundus Validator
-    fv_model = FundusValidator(pretrained=False, backbone_name=FV_BACKBONE).to(DEVICE)
-    fv_ckpt = torch.load("fundus_validator.pt", map_location=DEVICE)
-    fv_model.head.load_state_dict(fv_ckpt["head_state_dict"])
-    fv_model.eval()
+def load_fv_model():
+    m = FundusValidator(FV_BACKBONE).to(DEVICE)
+    ckpt = torch.load("fundus_validator.pt", map_location=DEVICE)
+    head_state = ckpt.get("head_state_dict", ckpt)
+    m.head.load_state_dict(head_state, strict=False)
+    return m.eval()
 
-    # Calibration
-    calib = None
+
+def load_calib():
     if os.path.exists("calibration.pkl"):
         with open("calibration.pkl", "rb") as f:
-            calib = pickle.load(f)
+            return pickle.load(f)
+    return None
 
-    return model, fv_model, calib
 
-model, fv_model, calib = load_models()
+print("Loading models…")
+model   = load_dr_model()
+fv_model = load_fv_model()
+calib   = load_calib()
+print(f"  DR model loaded | FV model loaded | calib={'yes' if calib else 'no'} | device={DEVICE}")
 
-inf_tfm = A.Compose([
-    A.Resize(IMG_SIZE, IMG_SIZE),
-    A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
-    ToTensorV2(),
-])
+# Transforms
+inf_tfm = A.Compose([A.Resize(IMG_SIZE, IMG_SIZE), A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), ToTensorV2()])
+fv_tfm  = A.Compose([A.Resize(224, 224),          A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), ToTensorV2()])
 
-fv_tfm = A.Compose([
-    A.Resize(224, 224),
-    A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
-    ToTensorV2(),
-])
+# Grad-CAM setup
+class _SingleLogit(nn.Module):
+    def __init__(self, inner): super().__init__(); self.inner = inner
+    def forward(self, x): return self.inner(x)["logits"]
 
-def fundus_heuristics(rgb: np.ndarray) -> Dict[str, float]:
-    r, g, b = rgb[..., 0].astype(np.float32), rgb[..., 1].astype(np.float32), rgb[..., 2].astype(np.float32)
-    bright = (r > 15) | (g > 15) | (b > 15)
-    n_bright = int(bright.sum())
-    red_dominance = float(((r > g) & (r > b))[bright].mean()) if n_bright > 100 else float(((r > g) & (r > b)).mean())
-    mask = retinal_mask(rgb)
-    disc_coverage = float(mask.mean() / 255.0)
-    gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
-    gx = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=3)
-    gy = cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=3)
-    mag = np.sqrt(gx**2 + gy**2)
-    edge_density = float(mag[bright].mean() / 255.0) if n_bright > 100 else float(mag.mean() / 255.0)
-    return {"red_dominance": red_dominance, "disc_coverage": disc_coverage, "edge_density": edge_density}
+_cam_model = _SingleLogit(model).to(DEVICE).eval()
+try:
+    _target_layers = [model.backbone.blocks[-1]]
+    cam_obj = GradCAMPlusPlus(model=_cam_model, target_layers=_target_layers)
+    GRADCAM_OK = True
+except Exception as e:
+    print(f"Grad-CAM init failed: {e}")
+    GRADCAM_OK = False
 
-STRONG_HEUR = {"red_dominance": 0.55, "disc_coverage": 0.25, "edge_density": 0.008}
-WEAK_HEUR = {"red_dominance": 0.30, "disc_coverage": 0.12, "edge_density": 0.006}
 
-def gate_fundus(pil_img: PILImage.Image) -> Tuple[bool, float, Dict[str, float], str]:
+# =============================================================================
+# Fundus Gating
+# =============================================================================
+def gate_fundus(pil_img: PILImage.Image) -> Tuple[bool, float, Dict, str]:
     rgb_raw = np.array(pil_img.convert("RGB"))
     try:
         rgb = apply_clahe_lab(crop_retinal_disc(rgb_raw))
-    except:
+    except Exception:
         rgb = rgb_raw
+
     h = fundus_heuristics(rgb)
     x = fv_tfm(image=rgb)["image"].unsqueeze(0).to(DEVICE)
     with torch.no_grad():
-        p_fundus = float(F.softmax(fv_model(x), dim=1).cpu().numpy()[0][1])
-    
+        p = float(F.softmax(fv_model(x), dim=1).cpu().numpy()[0][1])
+
     strong_ok = all(h[k] >= v for k, v in STRONG_HEUR.items())
-    weak_ok = all(h[k] >= v for k, v in WEAK_HEUR.items())
+    weak_ok   = all(h[k] >= v for k, v in WEAK_HEUR.items())
+
+    if strong_ok:                             return True, p, h, "Accepted (strong visual signatures)."
+    if p >= 0.75:                             return True, p, h, f"Accepted (validator: {p*100:.1f}%)."
+    if p >= 0.50 and weak_ok:                return True, p, h, "Accepted (validator + heuristics)."
+    if h["disc_coverage"] >= 0.25 and p >= 0.35: return True, p, h, "Accepted (clear retinal disc)."
+    if p >= 0.90:                             return True, p, h, "Accepted (high validator confidence)."
+    return False, p, h, "Not a fundus image. Please upload a colour retinal fundus photograph."
 
-    if strong_ok: return True, p_fundus, h, "Accepted (strong visual signatures)."
-    if p_fundus >= 0.75: return True, p_fundus, h, f"Accepted (validator: {p_fundus*100:.1f}%)."
-    if p_fundus >= 0.50 and weak_ok: return True, p_fundus, h, "Accepted (validator + heuristics)."
-    if h["disc_coverage"] >= 0.25 and p_fundus >= 0.35: return True, p_fundus, h, "Accepted (clear retinal disc)."
-    if p_fundus >= 0.90: return True, p_fundus, h, "Accepted (high validator confidence)."
-    return False, p_fundus, h, "Image rejected. Not a valid fundus photograph."
 
 # =============================================================================
-# Grad-CAM++ Integration
+# Grad-CAM Rendering
 # =============================================================================
-
-class _SingleLogit(nn.Module):
-    def __init__(self, inner):
-        super().__init__()
-        self.inner = inner
-    def forward(self, x):
-        return self.inner(x)["logits"]
-
-cam_model = _SingleLogit(model).to(DEVICE).eval()
-target_layers = [model.backbone.blocks[-1]] # Last block
-cam_obj = GradCAMPlusPlus(model=cam_model, target_layers=target_layers)
-
-def render_cam(rgb: np.ndarray, tensor: torch.Tensor, target_cls: int) -> np.ndarray:
+def render_cam(vis: np.ndarray, tensor: torch.Tensor) -> Optional[np.ndarray]:
+    if not GRADCAM_OK:
+        return None
     try:
-        grayscale_cam = cam_obj(input_tensor=tensor, targets=None)[0, :]
-        cam_image = show_cam_on_image(rgb.astype(np.float32) / 255, grayscale_cam, use_rgb=True)
-        return (cam_image * 255).astype(np.uint8)
-    except:
-        return rgb
+        gc = cam_obj(input_tensor=tensor, targets=None)[0, :]
+        overlay = show_cam_on_image(vis.astype(np.float32) / 255.0, gc, use_rgb=True)
+        return (overlay * 255).astype(np.uint8)
+    except Exception:
+        return None
 
-from pytorch_grad_cam.utils.image import show_cam_on_image
 
 # =============================================================================
-# Main Prediction Logic
+# Inference
 # =============================================================================
-
 def predict(uploaded):
-    start_time = time.time()
-    if uploaded is None: return None, None, None, "Please upload an image."
+    t0 = time.time()
+    if uploaded is None:
+        return None, None, "<p style='color:#666;padding:20px'>Upload a fundus image to begin.</p>"
 
     pil_img = PILImage.fromarray(uploaded.astype(np.uint8))
-    
-    # Gating
     accepted, p_fundus, heur, reason = gate_fundus(pil_img)
-    if not accepted:
-        return None, None, None, f"<div style='color:red; font-weight:bold;'>{reason}</div>"
 
-    # Preprocess
+    if not accepted:
+        rd  = heur.get("red_dominance", 0)
+        dc  = heur.get("disc_coverage", 0)
+        ed  = heur.get("edge_density", 0)
+        html = f"""
+        <div style="padding:20px;border-radius:10px;background:#fdecea;color:#8a1a14;
+                    border:1px solid #f5c6c2;font-family:sans-serif">
+          <h3 style="margin:0 0 8px">⚠️ Image Rejected</h3>
+          <p style="margin:0">{reason}</p>
+          <div style="margin-top:12px;font-size:12px;color:#555">
+            <strong>Validator score:</strong> {p_fundus*100:.1f}% (need ≥75%)<br>
+            <strong>Red dominance:</strong> {rd:.3f} (need ≥{WEAK_HEUR['red_dominance']:.2f})
+            {'✓' if rd >= WEAK_HEUR['red_dominance'] else '✗'}<br>
+            <strong>Disc coverage:</strong> {dc:.3f} (need ≥{WEAK_HEUR['disc_coverage']:.2f})
+            {'✓' if dc >= WEAK_HEUR['disc_coverage'] else '✗'}<br>
+            <strong>Edge density:</strong>  {ed:.4f} (need ≥{WEAK_HEUR['edge_density']:.3f})
+            {'✓' if ed >= WEAK_HEUR['edge_density'] else '✗'}
+          </div>
+        </div>"""
+        return None, None, html
+
+    # Pre-process
     rgb = np.array(pil_img.convert("RGB"))
-    rgb = apply_clahe_lab(crop_retinal_disc(rgb))
-    vis = cv2.resize(rgb, (320, 320))
+    try:
+        rgb = apply_clahe_lab(crop_retinal_disc(rgb))
+    except Exception:
+        pass
+    vis    = cv2.resize(rgb, (320, 320))
     tensor = inf_tfm(image=cv2.resize(rgb, (IMG_SIZE, IMG_SIZE)))["image"].unsqueeze(0).to(DEVICE)
 
-    # Inference (4-view TTA)
+    # 4-view TTA
     with torch.no_grad():
-        views = [tensor, torch.flip(tensor, dims=[-1]), torch.flip(tensor, dims=[-2]), torch.rot90(tensor, 2, dims=(-2, -1))]
-        logits_list, coral_list = [], []
+        views = [tensor,
+                 torch.flip(tensor, dims=[-1]),
+                 torch.flip(tensor, dims=[-2]),
+                 torch.rot90(tensor, 2, dims=(-2, -1))]
+        logits_all, coral_all = [], []
         for v in views:
             out = model(v)
-            logits_list.append(out["logits"].cpu().numpy())
-            coral_list.append(torch.sigmoid(out["coral"]).cpu().numpy())
-    
-    logits_mean = np.mean(logits_list, axis=0)
-    coral_mean = np.mean(coral_list, axis=0)
-
-    # Calibration
-    if calib:
-        probs = calib.apply(logits_mean)[0]
+            logits_all.append(out["logits"].cpu().numpy())
+            coral_all.append(torch.sigmoid(out["coral"]).cpu().numpy())
+
+    logits_mean = np.mean(logits_all, axis=0)  # (1, 5)
+    coral_mean  = np.mean(coral_all,  axis=0)  # (1, 4)
+
+    # Probabilities
+    if calib is not None:
+        try:
+            probs = calib.apply(logits_mean)[0]
+        except Exception:
+            ex = np.exp(logits_mean - logits_mean.max())
+            probs = (ex / ex.sum())[0]
     else:
         ex = np.exp(logits_mean - logits_mean.max())
         probs = (ex / ex.sum())[0]
 
-    # Fusion
+    # Fusion: CORAL ordinal + softmax
     argmax_pred = int(probs.argmax())
-    coral_pred = int((coral_mean > 0.5).sum())
-    pred = coral_pred if (coral_mean[0, coral_pred-1] if coral_pred > 0 else 1.0) > probs[argmax_pred] else argmax_pred
+    coral_pred  = int((coral_mean[0] > 0.5).sum())
+    coral_conf  = float(coral_mean[0, coral_pred - 1]) if coral_pred > 0 else 1.0
+    pred = coral_pred if coral_conf > float(probs[argmax_pred]) else argmax_pred
+    pred = max(0, min(4, pred))
     conf = float(probs[pred])
 
     # Grad-CAM
-    cam_img = render_cam(vis, tensor, pred)
-    
-    proc_time = time.time() - start_time
-    
-    # Result HTML
+    cam_img = render_cam(vis, tensor)
+
+    proc_time = time.time() - t0
+
+    # Build probability bars
+    bar_html = ""
+    for i, (lbl, col) in enumerate(zip(GRADE_MAP.values(), GRADE_COLORS)):
+        w = float(probs[i]) * 100
+        weight = "700" if i == pred else "400"
+        bar_html += f"""
+        <div style="margin-bottom:6px">
+          <div style="display:flex;justify-content:space-between;font-size:13px;font-weight:{weight}">
+            <span>{SEVERITY_ICONS[i]} {lbl}</span><span>{w:.1f}%</span>
+          </div>
+          <div style="height:6px;background:#eee;border-radius:3px;overflow:hidden">
+            <div style="width:{w}%;background:{col};height:100%"></div>
+          </div>
+        </div>"""
+
     html = f"""
-    <div style="background: white; padding: 20px; border-radius: 12px; box-shadow: 0 4px 12px rgba(0,0,0,0.1); font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;">
-        <div style="display: flex; align-items: center; gap: 15px;">
-            <span style="font-size: 48px;">{SEVERITY_ICONS[pred]}</span>
-            <div>
-                <h2 style="margin: 0; color: {GRADE_COLORS[pred]}; font-size: 24px;">{GRADE_MAP[pred]}</h2>
-                <p style="margin: 0; color: #666;">Grade {pred}/4 (ICDR Scale)</p>
-            </div>
+    <div style="background:#fff;padding:20px;border-radius:14px;
+                box-shadow:0 4px 16px rgba(0,0,0,.08);
+                font-family:'Segoe UI',Tahoma,sans-serif;line-height:1.5">
+
+      <!-- Grade header -->
+      <div style="display:flex;align-items:center;gap:14px;margin-bottom:16px">
+        <span style="font-size:52px">{SEVERITY_ICONS[pred]}</span>
+        <div>
+          <div style="font-size:22px;font-weight:700;color:{GRADE_COLORS[pred]}">{GRADE_MAP[pred]}</div>
+          <div style="font-size:13px;color:#888">Grade {pred} / 4 &nbsp;·&nbsp; ICDR Scale</div>
         </div>
-        <div style="margin-top: 15px;">
-            <div style="display: flex; justify-content: space-between; font-size: 14px; margin-bottom: 5px;">
-                <span>Confidence</span>
-                <span style="font-weight: bold;">{conf*100:.1f}%</span>
-            </div>
-            <div style="height: 8px; background: #eee; border-radius: 4px; overflow: hidden;">
-                <div style="width: {conf*100}%; background: {GRADE_COLORS[pred]}; height: 100%; transition: width 0.5s;"></div>
-            </div>
-        </div>
-        <div style="margin-top: 15px; padding: 12px; background: #f8f9fa; border-left: 4px solid {GRADE_COLORS[pred]}; border-radius: 4px;">
-            <p style="margin: 0; font-size: 14px; color: #333;">{CLINICAL_ACTION[pred]}</p>
-        </div>
-        <div style="margin-top: 15px; font-size: 12px; color: #888; border-top: 1px solid #eee; padding-top: 10px; display: flex; justify-content: space-between;">
-            <span>Processing time: {proc_time:.2f}s</span>
-            <span>Validator: {p_fundus*100:.1f}%</span>
-        </div>
-    </div>
-    """
+      </div>
+
+      <!-- Confidence bars -->
+      <div style="background:#f8f9fa;border-radius:8px;padding:12px;margin-bottom:14px">
+        <div style="font-size:13px;font-weight:600;color:#444;margin-bottom:8px">Grade Probabilities</div>
+        {bar_html}
+      </div>
+
+      <!-- Clinical action -->
+      <div style="padding:12px;background:#f0f4ff;border-left:4px solid {GRADE_COLORS[pred]};
+                  border-radius:6px;margin-bottom:14px">
+        <div style="font-size:12px;font-weight:600;color:#555;margin-bottom:4px">CLINICAL RECOMMENDATION</div>
+        <div style="font-size:14px;color:#222">{CLINICAL_ACTION[pred]}</div>
+      </div>
+
+      <!-- Metrics footer -->
+      <div style="display:flex;justify-content:space-between;font-size:11px;color:#aaa;
+                  border-top:1px solid #eee;padding-top:10px">
+        <span>⏱ Processing: <strong>{proc_time:.2f}s</strong></span>
+        <span>🔍 Validator: <strong>{p_fundus*100:.1f}%</strong></span>
+        <span>🎯 Confidence: <strong>{conf*100:.1f}%</strong></span>
+      </div>
+    </div>"""
+
     return vis, cam_img, html
 
+
 # =============================================================================
-# UI Definition
+# Gradio UI — Material Design 3
 # =============================================================================
-
-m3_theme = gr.themes.Soft(
+m3 = gr.themes.Soft(
     primary_hue="blue",
     secondary_hue="indigo",
     neutral_hue="slate",
-    font=["Outfit", "Inter", "system-ui"],
+    font=["Inter", "Outfit", "system-ui"],
 ).set(
     button_primary_background_fill="*primary_600",
     button_primary_background_fill_hover="*primary_700",
     block_title_text_weight="600",
     block_border_width="1px",
-    block_shadow="0 2px 4px rgba(0,0,0,0.05)",
+    block_shadow="0 2px 8px rgba(0,0,0,0.06)",
 )
 
-with gr.Blocks(theme=m3_theme, title="DR Grading AI") as demo:
+with gr.Blocks(theme=m3, title="DR Grading AI — Clinical Decision Support") as demo:
+
     gr.HTML("""
-    <div style="text-align: center; padding: 30px 20px; background: linear-gradient(135deg, #1A73E8, #0D47A1); color: white; border-radius: 16px; margin-bottom: 25px;">
-        <h1 style="margin: 0; font-size: 2.5em; font-weight: 700;">Diabetic Retinopathy Grading</h1>
-        <p style="margin: 10px 0 0; opacity: 0.9; font-size: 1.1em;">AI-Powered Clinical Diagnosis Support</p>
+    <div style="text-align:center;padding:28px 20px 22px;
+                background:linear-gradient(135deg,#1A73E8 0%,#0D47A1 100%);
+                color:#fff;border-radius:16px;margin-bottom:20px">
+      <div style="font-size:2.4em;font-weight:800;letter-spacing:-0.5px">
+        👁 Diabetic Retinopathy Grading
+      </div>
+      <div style="margin-top:8px;opacity:.88;font-size:1.05em">
+        EfficientNetV2-M · CORAL Ordinal · Grad-CAM++ · 4-View TTA
+      </div>
+      <div style="margin-top:6px;opacity:.7;font-size:.85em">
+        For research & educational use only — not a medical device
+      </div>
     </div>
     """)
 
-    with gr.Row():
-        with gr.Column(scale=1):
-            input_img = gr.Image(label="Upload Retina Fundus Image", type="numpy", height=400)
+    with gr.Row(equal_height=True):
+        # ── LEFT PANEL ──────────────────────────────────────────────────────
+        with gr.Column(scale=1, min_width=340):
+            inp = gr.Image(
+                label="Upload Retinal Fundus Photograph",
+                type="numpy",
+                height=380,
+                sources=["upload", "clipboard"],
+            )
             with gr.Row():
-                clear_btn = gr.Button("Clear", variant="secondary")
-                submit_btn = gr.Button("Analyze Image", variant="primary")
-            
+                btn_clear   = gr.Button("🗑 Clear",   variant="secondary", size="sm")
+                btn_analyze = gr.Button("🔬 Analyze", variant="primary",   size="lg")
+
             gr.Markdown("""
-            ### Instructions
-            1. Upload a clear, high-resolution color fundus photograph.
-            2. Click **Analyze Image**.
-            3. The AI will validate the image and provide a grading based on the ICDR scale.
-            
-            *Note: This tool is for research purposes only.*
+            **How to use**
+            1. Upload a colour fundus photograph (JPEG / PNG)
+            2. Click **Analyze** — the AI validates it first
+            3. Review the grade, heatmap, and recommendation
+
+            > *Grades 0–4 follow the International Clinical DR (ICDR) severity scale.*
             """)
 
-        with gr.Column(scale=1):
-            with gr.Group():
-                with gr.Row():
-                    out_vis = gr.Image(label="Enhanced View", interactive=False)
-                    out_cam = gr.Image(label="Pathology Heatmap", interactive=False)
-                out_html = gr.HTML(label="Diagnostic Report")
+        # ── RIGHT PANEL ─────────────────────────────────────────────────────
+        with gr.Column(scale=1, min_width=340):
+            with gr.Row():
+                out_vis = gr.Image(label="Enhanced Fundus", interactive=False, height=200)
+                out_cam = gr.Image(label="Pathology Heatmap (Grad-CAM++)", interactive=False, height=200)
+            out_report = gr.HTML()
 
-    submit_btn.click(
+    # ── Events ──────────────────────────────────────────────────────────────
+    btn_analyze.click(
         fn=predict,
-        inputs=input_img,
-        outputs=[out_vis, out_cam, out_html]
+        inputs=inp,
+        outputs=[out_vis, out_cam, out_report],
+    )
+    btn_clear.click(
+        fn=lambda: (None, None, None, ""),
+        inputs=[],
+        outputs=[inp, out_vis, out_cam, out_report],
     )
-    
-    clear_btn.click(lambda: (None, None, None, ""), None, [input_img, out_vis, out_cam, out_html])
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)