Added Grad-CAM

app.py

@@ -8,6 +8,11 @@ from dotenv import load_dotenv
 from model_loader import load_alexnet_model, preprocess_image
 from flask_cors import CORS
 
+from io import BytesIO
+import base64
+import numpy as np
+
+
 load_dotenv(override=True)
 
 # HF sets PORT dynamically. Fall back to 7860 locally.
@@ -55,6 +60,96 @@ def load_image(file_stream_or_path):
         return Image.open(file_stream_or_path).convert("RGB")
     return Image.open(file_stream_or_path).convert("RGB")
 
+def generate_gradcam(img_pil: Image.Image, target_idx: int) -> str:
+    """
+    Returns a data URL (PNG) of the Grad-CAM overlay for the target class.
+    """
+    model.eval()
+    orig_w, orig_h = img_pil.size
+
+    # Last conv of standard AlexNet
+    target_layer = model.features[10]
+
+    activations = []
+    gradients = []
+
+    def fwd_hook(_, __, out):
+        # Save activations (detached) and attach a tensor hook to capture gradients
+        activations.append(out.detach())
+        out.register_hook(lambda g: gradients.append(g.detach().clone()))
+
+    handle = target_layer.register_forward_hook(fwd_hook)
+    try:
+        # Forward
+        input_tensor = preprocess_image(img_pil).to(DEVICE)
+        output = model(input_tensor)  # [1, C]
+
+        # Backward on the selected class
+        if target_idx < 0 or target_idx >= output.shape[1]:
+            raise ValueError(f"target_idx {target_idx} out of range for output dim {output.shape[1]}")
+        score = output[0, target_idx]
+        model.zero_grad(set_to_none=True)
+        score.backward()
+
+        # Ensure hooks fired
+        if not activations or not gradients:
+            raise RuntimeError("Grad-CAM hooks did not capture activations/gradients")
+
+        A  = activations[-1]         # [1, C, H, W]
+        dA = gradients[-1]           # [1, C, H, W]
+
+        # Weights: global-average-pool the gradients
+        weights = dA.mean(dim=(2, 3), keepdim=True)  # [1, C, 1, 1]
+        cam = (weights * A).sum(dim=1, keepdim=False)  # [1, H, W]
+        cam = torch.relu(cam)[0]  # [H, W]
+
+        # Normalize to [0,1]
+        cam -= cam.min()
+        if cam.max() > 0:
+            cam /= cam.max()
+
+        # Resize CAM to original image size
+        cam_np = cam.detach().cpu().numpy()
+        cam_img = Image.fromarray((cam_np * 255).astype(np.uint8), mode="L")
+        cam_img = cam_img.resize((orig_w, orig_h), resample=Image.BILINEAR)
+
+        # Red alpha overlay
+        heat_rgba = Image.new("RGBA", (orig_w, orig_h), (255, 0, 0, 0))
+        heat_rgba.putalpha(cam_img)
+        base = img_pil.convert("RGBA")
+        overlayed = Image.alpha_composite(base, heat_rgba)
+
+        # Encode to data URL
+        buff = BytesIO()
+        overlayed.save(buff, format="PNG")
+        b64 = base64.b64encode(buff.getvalue()).decode("utf-8")
+        return f"data:image/png;base64,{b64}"
+    finally:
+        handle.remove()   # <-- remove the actual handle you registered
+
+
+def run_inference_with_gradcam(img: Image.Image) -> Dict[str, Any]:
+    """Run softmax inference and also compute Grad-CAM for the predicted class."""
+    # Regular inference (no grad) for probabilities
+    input_tensor = preprocess_image(img).to(DEVICE)
+    with torch.no_grad():
+        output = model(input_tensor)
+        probabilities = F.softmax(output[0], dim=0).detach().cpu()
+
+    pred_prob, pred_idx = torch.max(probabilities, dim=0)
+    predicted_class = classes[int(pred_idx)]
+
+    # Grad-CAM for predicted index
+    gradcam_data_url = generate_gradcam(img, int(pred_idx))
+
+    return {
+        "class": predicted_class,
+        "confidence": float(pred_prob),
+        "probabilities": {cls: float(prob) for cls, prob in zip(classes, probabilities.tolist())},
+        "gradcam": gradcam_data_url,
+    }
+
+
 def run_inference(img: Image.Image) -> Dict[str, Any]:
     input_tensor = preprocess_image(img).to(DEVICE)
     with torch.no_grad():
@@ -78,7 +173,7 @@ def predict_alexnet() -> Any:
         return jsonify({"error": "Empty file."}), 400
     try:
         img = load_image(file.stream)
-        result = run_inference(img)
+        result = run_inference_with_gradcam(img)  # << changed
         return jsonify(result)
     except Exception as e:
         return jsonify({"error": f"Failed to process image: {e}"}), 400
@@ -103,7 +198,7 @@ def predict_preset() -> Any:
 
     try:
         img = load_image(path)
-        result = run_inference(img)
+        result = run_inference_with_gradcam(img)  # << changed
         result.update({"preset": key, "path": path})
         return jsonify(result)
     except Exception as e:

requirements.txt

@@ -2,4 +2,5 @@ flask>=3.0.0
 pillow>=10.0.0
 gunicorn>=21.2.0
 python-dotenv>=1.0.0
-Flask-Cors>=4.0.0
+Flask-Cors>=4.0.0
+numpy>=1.24.0

static/index.html

@@ -155,6 +155,12 @@
             <div class="probabilities-title">All Class Probabilities</div>
             <div id="probabilitiesList"></div>
         </div>
+
+        <div class="gradcam-container" id="gradcamContainer" style="display:none; margin:16px 0 20px;">
+            <div class="probabilities-title" style="margin-bottom:10px;">Grad-CAM (Predicted Class)</div>
+            <img id="gradcamImage" class="preview-image" alt="Grad-CAM visualization" style="max-width:480px; width:100%; border-radius:10px; box-shadow:0 4px 20px rgba(0,0,0,0.08);" />
+        </div>
+
     </div>
 
     <div class="error-message" id="errorMessage"></div>
@@ -180,6 +186,10 @@
     const loadingSpinner = document.getElementById('loadingSpinner');
     const presetGrid = document.getElementById('presetGrid');
 
+    const gradcamContainer = document.getElementById('gradcamContainer');
+    const gradcamImage = document.getElementById('gradcamImage');
+
+
     let currentFile = null;
     let currentPreset = null; // 'TP' | 'TN' | 'FN' | 'FP' | null
 
@@ -344,35 +354,49 @@
         }
     });
 
-    function displayResults(result) {
-        predictedClass.textContent = result.class;
-        confidenceScore.textContent = `${(result.confidence * 100).toFixed(2)}% Confidence`;
-        const sortedProbs = Object.entries(result.probabilities).sort(([,a],[,b])=>b-a).slice(0,10);
-        probabilitiesList.innerHTML = '';
-        sortedProbs.forEach(([className, prob], index) => {
-            const probPercent = (prob * 100).toFixed(2);
-            const isTop = index === 0;
-            const div = document.createElement('div');
-            div.className = 'probability-item';
-            div.innerHTML = `
-                <div class="probability-label">
-                    <span class="class-name" style="${isTop?'font-weight:700;color:#667eea;':''}">${className}</span>
-                    <span class="class-prob" style="${isTop?'font-weight:700;color:#667eea;':''}">${probPercent}%</span>
-                </div>
-                <div class="probability-bar-bg">
-                    <div class="probability-bar" style="width:0%;" data-width="${probPercent}"></div>
-                </div>
-            `;
-            probabilitiesList.appendChild(div);
-        });
-        resultsSection.classList.add('active');
-        setTimeout(() => {
-            probabilitiesList.querySelectorAll('.probability-bar').forEach(bar => {
-                bar.style.width = bar.getAttribute('data-width') + '%';
-            });
-        }, 100);
+function displayResults(result) {
+    predictedClass.textContent = result.class;
+    confidenceScore.textContent = `${(result.confidence * 100).toFixed(2)}% Confidence`;
+
+    // --- NEW: Grad-CAM rendering ---
+    if (result.gradcam) {
+        gradcamImage.src = result.gradcam;
+        gradcamContainer.style.display = 'block';
+    } else {
+        gradcamContainer.style.display = 'none';
+        gradcamImage.removeAttribute('src');
     }
 
+    const sortedProbs = Object.entries(result.probabilities)
+        .sort(([, a], [, b]) => b - a)
+        .slice(0, 10);
+
+    probabilitiesList.innerHTML = '';
+    sortedProbs.forEach(([className, prob], index) => {
+        const probPercent = (prob * 100).toFixed(2);
+        const isTop = index === 0;
+        const div = document.createElement('div');
+        div.className = 'probability-item';
+        div.innerHTML = `
+            <div class="probability-label">
+                <span class="class-name" style="${isTop ? 'font-weight:700;color:#667eea;' : ''}">${className}</span>
+                <span class="class-prob" style="${isTop ? 'font-weight:700;color:#667eea;' : ''}">${probPercent}%</span>
+            </div>
+            <div class="probability-bar-bg">
+                <div class="probability-bar" style="width:0%;" data-width="${probPercent}"></div>
+            </div>
+        `;
+        probabilitiesList.appendChild(div);
+    });
+
+    resultsSection.classList.add('active');
+    setTimeout(() => {
+        probabilitiesList.querySelectorAll('.probability-bar').forEach(bar => {
+            bar.style.width = bar.getAttribute('data-width') + '%';
+        });
+    }, 100);
+}
+
     function showLoading() { loadingSpinner.classList.add('active'); classifyBtn.disabled = true; }
     function hideLoading() { loadingSpinner.classList.remove('active'); classifyBtn.disabled = false; }
     function hideResults() { resultsSection.classList.remove('active'); }