Update app.py

app.py

@@ -9,6 +9,7 @@ from datetime import datetime
 import sqlite3
 import torch.nn as nn
 import torchvision.models as models
+import cv2
 
 app = Flask(__name__)
 
@@ -46,13 +47,11 @@ from densenet_withglam import get_model_with_attention
 # ✅ Instantiate the model
 model = get_model_with_attention('densenet169', num_classes=3)  # Will have GLAM
 model.load_state_dict(torch.load('densenet169_seed40_best.pt', map_location='cpu'))
- # Load your trained weights
 model.eval()
 
 # ✅ Class Names
 CLASS_NAMES = ["Advanced", "Early", "Normal"]
 
-
 # ✅ Transformation for input images
 transform = transforms.Compose([
     transforms.Resize(256),
@@ -61,6 +60,46 @@ transform = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                          std=[0.229, 0.224, 0.225]),
 ])
+
+# =========================
+# GRAD-CAM IMPLEMENTATION
+# =========================
+class GradCAM:
+    """Grad-CAM for the target layer."""
+    def __init__(self, model, target_layer_name):
+        self.model = model
+        self.target_layer_name = target_layer_name
+        self.activations = None
+        self.gradients = None
+        self._register_hooks()
+
+    def _register_hooks(self):
+        """Register forward and backward hooks."""
+        for name, module in self.model.named_modules():
+            if name == self.target_layer_name:
+                module.register_forward_hook(self._forward_hook)
+                module.register_full_backward_hook(self._backward_hook)
+
+    def _forward_hook(self, module, input, output):
+        """Save activations."""
+        self.activations = output
+
+    def _backward_hook(self, module, grad_in, grad_out):
+        """Save gradients."""
+        self.gradients = grad_out[0]
+
+    def generate(self, class_index):
+        """Generate the Grad-CAM."""
+        if self.activations is None or self.gradients is None:
+            raise ValueError("Must run forward and backward passes first.")
+        weights = self.gradients.mean(dim=(2, 3), keepdim=True)
+        cam = (weights * self.activations).sum(dim=1, keepdim=True)
+        cam = F.relu(cam)
+        cam = cam.squeeze().cpu().numpy()
+        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
+        return cam
+
+
 @app.route('/')
 def home():
     """Check that the API is working."""
@@ -75,9 +114,10 @@ def test_file():
     else:
         return "❌ Model file NOT found."
 
+
 @app.route('/predict', methods=['POST'])
 def predict():
-    """Perform prediction using PyTorch (3-class), save results, and save to SQLite database."""
+    """Perform prediction and save results (including Grad-CAM) to the database."""
     if 'file' not in request.files:
         return jsonify({'error': 'No file uploaded'}), 400
 
@@ -96,22 +136,40 @@ def predict():
         img = Image.open(uploaded_file_path).convert('RGB')
         input_tensor = transform(img).unsqueeze(0)
 
-        with torch.no_grad():
-            output = model(input_tensor)
-            probabilities = F.softmax(output, dim=1).cpu().numpy()[0]
+        # Grad-CAM setup
+        target_layer_name = "features.2.global_spatial_conv"
+        gradcam = GradCAM(model, target_layer_name)
+
+        # Forward pass
+        input_tensor.requires_grad = True
+        output = model(input_tensor)
 
-        # ✅ Get result
+        probabilities = F.softmax(output, dim=1).cpu().numpy()[0]
         class_index = np.argmax(probabilities)
         result = CLASS_NAMES[class_index]
         confidence = float(probabilities[class_index])
 
+        # Backward pass for Grad-CAM
+        model.zero_grad()
+        output[0, class_index].backward()
+        cam = gradcam.generate(class_index)
+
+        # ✅ Create overlay
+        original_img = np.asarray(img.resize((224, 224)))
+        heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_JET)
+        overlay = cv2.addWeighted(original_img, 0.6, heatmap, 0.4, 0)
+
+        gradcam_filename = f"gradcam_{timestamp}.png"
+        gradcam_file_path = os.path.join(OUTPUT_DIR, gradcam_filename)
+        cv2.imwrite(gradcam_file_path, overlay)
+
         # ✅ Save results to SQLite
         conn = sqlite3.connect(DB_PATH)
         cursor = conn.cursor()
         cursor.execute("""
             INSERT INTO results (image_filename, prediction, confidence, gradcam_filename, timestamp) 
             VALUES (?, ?, ?, ?, ?)
-        """, (uploaded_filename, result, confidence, '', datetime.now().isoformat()))
+        """, (uploaded_filename, result, confidence, gradcam_filename, datetime.now().isoformat()))
         conn.commit()
         conn.close()
 
@@ -121,13 +179,14 @@ def predict():
             'normal_probability': float(probabilities[0]),
             'early_glaucoma_probability': float(probabilities[1]),
             'advanced_glaucoma_probability': float(probabilities[2]),
-            'gradcam_image': '',  # Not used for now
+            'gradcam_image': gradcam_filename,
             'image_filename': uploaded_filename
         })
 
     except Exception as e:
         return jsonify({'error': str(e)}), 500
 
+
 @app.route('/results', methods=['GET'])
 def results():
     """List all results from the SQLite database."""
@@ -150,15 +209,17 @@ def results():
 
     return jsonify(results_list)
 
+
 @app.route('/gradcam/<filename>')
 def get_gradcam(filename):
-    """Serve the Grad-CAM overlay image (no-op for now)."""
+    """Serve the Grad-CAM overlay image."""
     filepath = os.path.join(OUTPUT_DIR, filename)
     if os.path.exists(filepath):
         return send_file(filepath, mimetype='image/png')
     else:
         return jsonify({'error': 'File not found'}), 404
 
+
 @app.route('/image/<filename>')
 def get_image(filename):
     """Serve the original uploaded image."""
@@ -168,5 +229,6 @@ def get_image(filename):
     else:
         return jsonify({'error': 'File not found'}), 404
 
+
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=7860)