Update app.py

app.py

@@ -21,19 +21,18 @@ import numpy as np
 logging.basicConfig(level=logging.INFO)
 log = logging.getLogger(__name__)
 
-# Load model
+# Load model - IMPORTANT: Don't call .eval() here for CAM to work
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-MODEL = xrv.models.get_model("densenet121-res224-all").to(DEVICE).eval()
+MODEL = xrv.models.get_model("densenet121-res224-all").to(DEVICE)
+# Note: We'll switch between train/eval modes as needed
 LABELS = MODEL.pathologies
 
-# Correct transform for TorchXRayVision (grayscale, normalized to [-1024, 1024])
+# Initialize CAM extractor with correct input shape for grayscale
+cam_extractor = SmoothGradCAMpp(MODEL, input_shape=(1, 224, 224))
+
 def preprocess_xray(pil_img: Image.Image) -> torch.Tensor:
     """
     Preprocess PIL image for TorchXRayVision model
-    TorchXRayVision expects:
-    - Single channel grayscale image
-    - Values normalized to [-1024, 1024] range
-    - Resolution of 224x224
     """
     # Convert to grayscale if needed
     if pil_img.mode != "L":
@@ -43,11 +42,10 @@ def preprocess_xray(pil_img: Image.Image) -> torch.Tensor:
     img_array = np.array(pil_img, dtype=np.float32)
     
     # Normalize to [-1024, 1024] range (TorchXRayVision standard)
-    # Assume input is 8-bit (0-255), scale to [-1024, 1024]
     img_array = xrv.datasets.normalize(img_array, 255)
     
-    # Add channel dimension and resize
-    img_array = img_array[None, ...]  # Add channel dimension
+    # Add channel dimension
+    img_array = img_array[None, ...]
     
     # Use TorchXRayVision transforms
     transform = transforms.Compose([
@@ -57,25 +55,26 @@ def preprocess_xray(pil_img: Image.Image) -> torch.Tensor:
     
     img_array = transform(img_array)
     
-    # Convert to tensor
+    # Convert to tensor with gradient enabled
     img_tensor = torch.from_numpy(img_array).unsqueeze(0).to(DEVICE)
+    img_tensor.requires_grad_(True)  # Enable gradients for CAM
     
     return img_tensor
 
-# Initialize CAM extractor with correct input shape for grayscale
-cam_extractor = SmoothGradCAMpp(MODEL, input_shape=(1, 224, 224))  # Single channel
-
 def analyse_xray(img: Image.Image):
     if img is None: 
         return "Please upload an image.", None
 
     try:
-        # Preprocess image for TorchXRayVision
+        # Preprocess image
         x = preprocess_xray(img)
         
-        with torch.no_grad():
-            logits = MODEL(x)
-            probs = torch.sigmoid(logits)[0] * 100  # Convert to percentages
+        # Set model to train mode temporarily for gradient computation
+        MODEL.train()
+        
+        # Forward pass with gradient tracking
+        logits = MODEL(x)
+        probs = torch.sigmoid(logits)[0] * 100  # Convert to percentages
         
         # Get top 5 predictions
         topk = torch.topk(probs, 5)
@@ -86,13 +85,16 @@ def analyse_xray(img: Image.Image):
         # Generate activation map
         activation_map = cam_extractor(target, logits)[0]
         
-        # Overlay heatmap on original image
         # Convert single channel to 3-channel for overlay
-        input_for_overlay = x.squeeze(0).cpu()
-        input_for_overlay = input_for_overlay.repeat(3, 1, 1)  # Repeat single channel 3 times
+        input_for_overlay = x.squeeze(0).cpu().detach()
+        input_for_overlay = input_for_overlay.repeat(3, 1, 1)
         
+        # Generate heatmap
         heatmap = cam_extractor.overlay(input_for_overlay, activation_map)
         
+        # Set model back to eval mode
+        MODEL.eval()
+        
         # Build HTML summary table
         table_rows = ""
         for i in range(len(topk.indices)):
@@ -148,7 +150,6 @@ MEDICAL_ADVICE = {
 def get_medical_advice(condition: str) -> str:
     return MEDICAL_ADVICE.get(condition, "Consult with a radiologist or pulmonologist for proper interpretation.")
 
-# PDF report analysis (simplified - focusing on the main issue)
 def analyse_report(file):
     if file is None: 
         return "Please upload a PDF file."
@@ -252,24 +253,6 @@ with gr.Blocks(title="🩻 RadiologyScan AI", theme=gr.themes.Soft()) as demo:
                 outputs=[pdf_input, report_output]
             )
 
-    gr.Markdown("""
-    ### 📖 How to Use
-    1. **X-ray Analysis**: Upload a chest X-ray image (JPEG, PNG) and click "Analyze X-ray"
-    2. **Report Analysis**: Upload a medical report PDF and click "Analyze Report"
-    
-    ### 🔬 Technical Details
-    - Uses TorchXRayVision pre-trained DenseNet-121 model
-    - Trained on multiple chest X-ray datasets
-    - Provides attention heatmaps for interpretability
-    - Supports 18 different pathological conditions
-    
-    ### ⚠️ Limitations
-    - For educational use only
-    - Not a substitute for professional medical diagnosis
-    - Results may vary based on image quality
-    - Always consult healthcare professionals
-    """)
-
 if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",