Update app.py

app.py

@@ -12,23 +12,14 @@ try:
 except Exception as e:
     print(f"Error loading model: {e}")
 image_processor = ViTImageProcessor.from_pretrained(model_name)
+
 # NIH Chest X-ray predefined conditions
 labels = [
-    "Atelectasis",
-    "Cardiomegaly",
-    "Effusion",
-    "Infiltration",
-    "Mass",
-    "Nodule",
-    "Pneumonia",
-    "Pneumothorax",
-    "Consolidation",
-    "Edema",
-    "Emphysema",
-    "Fibrosis",
-    "Pleural Thickening",
-    "Hernia"
+    "Atelectasis", "Cardiomegaly", "Effusion", "Infiltration", "Mass", "Nodule",
+    "Pneumonia", "Pneumothorax", "Consolidation", "Edema", "Emphysema", 
+    "Fibrosis", "Pleural Thickening", "Hernia"
 ]
+
 # Function to apply Grad-CAM visualization
 def generate_grad_cam(image, target_layer):
     try:
@@ -37,18 +28,28 @@ def generate_grad_cam(image, target_layer):
         
         # Forward pass to get logits
         input_tensor = inputs["pixel_values"]
+        input_tensor.requires_grad = True  # Enable gradient tracking
         outputs = model(input_tensor)
         logits = outputs.logits
         
-        # Calculate Grad-CAM
-        cam_weights = torch.mean(torch.relu(logits), dim=(2, 3))
-        cam_map = (torch.unsqueeze(cam_weights, 1) * 
-                   torch.sigmoid(outputs.pooler_output)).sum(dim=1).squeeze()
+        # Get the predicted class and calculate gradients
+        predicted_class = logits.argmax(-1)
+        class_score = logits[0, predicted_class]
+        class_score.backward()
+        
+        # Get gradients and weights from the target layer
+        gradients = model.get_input_embeddings().weight.grad
+        pooled_gradients = torch.mean(gradients, dim=[0, 2, 3])
+        
+        # Apply Grad-CAM calculation (modify this part as per the model architecture)
+        cam = torch.mean(pooled_gradients * inputs["pixel_values"], dim=1).squeeze()
+        cam = torch.clamp(cam, min=0).numpy()  # Ensure non-negative values
         
-        return cam_map.numpy(), logits.argmax(-1)
+        return cam, predicted_class.item()
     except Exception as e:
         print(f"Error generating Grad-CAM: {e}")
         return None
+
 # Function to predict classes and visualize Grad-CAM
 def predict_and_explain(image):
     try:
@@ -63,13 +64,25 @@ def predict_and_explain(image):
         predicted_class = logits.argmax(-1).item()
         cam_map, _ = generate_grad_cam(image, "pooler_output")
         
+        # Convert cam_map to a visualizable format (heatmap)
+        plt.imshow(cam_map, cmap='jet', alpha=0.5)
+        plt.axis('off')
+        plt.title(f"Grad-CAM for {labels[predicted_class]}")
+        plt.colorbar()
+        plt.savefig("grad_cam_output.png")
+        plt.close()
+        
+        # Load the saved image to return it
+        grad_cam_image = Image.open("grad_cam_output.png")
+        
         return {
             "predicted class": labels[predicted_class],
-            "Grad-CAM map": cam_map,
+            "Grad-CAM map": grad_cam_image,
         }
     except Exception as e:
         print(f"Error predicting and explaining: {e}")
         return None
+
 # Create a Gradio interface
 iface = gr.Interface(
     fn=predict_and_explain, 
@@ -79,4 +92,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()