Update app.py

app.py

@@ -42,7 +42,6 @@ class DoubleConv(nn.Module):
         return self.conv(x)
 
 
-# Also, make sure your AttentionBlock.forward() returns the attention map:
 class AttentionBlock(nn.Module):
     def __init__(self, F_g, F_l, F_int):
         super(AttentionBlock, self).__init__()
@@ -238,30 +237,6 @@ def preprocess_for_model(image):
     
     return transform(image).unsqueeze(0)
 
-def apply_tta(model, input_tensor):
-    """Test-Time Augmentation"""
-    augmentations = [
-        lambda x: x,  # Original
-        lambda x: TF.hflip(x),  # Horizontal flip
-        lambda x: TF.vflip(x),  # Vertical flip
-    ]
-    
-    predictions = []
-    for i, aug in enumerate(augmentations):
-        aug_input = aug(input_tensor)
-        pred, _ = model(aug_input)
-        pred = torch.sigmoid(pred)
-        
-        # Reverse augmentation
-        if i == 1:  # Reverse hflip
-            pred = TF.hflip(pred)
-        elif i == 2:  # Reverse vflip
-            pred = TF.vflip(pred)
-            
-        predictions.append(pred)
-    
-    return torch.mean(torch.stack(predictions), dim=0)
-
 def generate_attention_heatmap(attention_maps):
     """Generate attention heatmap - Fixed version"""
     if not attention_maps:
@@ -290,9 +265,8 @@ def generate_attention_heatmap(attention_maps):
     
     return heatmap
 
-
 def analyze_image(image, ground_truth, filename):
-    """Main analysis function"""
+    """Main analysis function - FIXED VERSION"""
     if model is None:
         return None, "Model not loaded. Please restart the application."
     
@@ -303,18 +277,22 @@ def analyze_image(image, ground_truth, filename):
         # Preprocess
         input_tensor = preprocess_for_model(image).to(device)
         
-        # Apply TTA
+        # Get prediction and attention maps
         with torch.no_grad():
-            avg_pred = apply_tta(model, input_tensor)
-            _, attention_maps = model(input_tensor)
-        
-        # Get binary mask
-        binary_mask = (avg_pred > 0.5).squeeze().cpu().numpy()
+            # Get model output (prediction + attention maps)
+            model_output, attention_maps = model(input_tensor)
+            
+            # Apply sigmoid and threshold to get binary mask
+            pred_mask = torch.sigmoid(model_output)
+            binary_mask = (pred_mask > 0.5).float()
+            
+            # Convert to numpy for further processing
+            binary_mask_np = binary_mask.squeeze().cpu().numpy()
         
-        # Post-processing
+        # Post-processing (morphological operations)
         kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3,3))
-        binary_mask = cv2.morphologyEx(binary_mask.astype(np.uint8), cv2.MORPH_OPEN, kernel)
-        binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel)
+        binary_mask_np = cv2.morphologyEx(binary_mask_np.astype(np.uint8), cv2.MORPH_OPEN, kernel)
+        binary_mask_np = cv2.morphologyEx(binary_mask_np, cv2.MORPH_CLOSE, kernel)
         
         # Generate attention heatmap
         att_heatmap = generate_attention_heatmap(attention_maps)
@@ -340,33 +318,36 @@ def analyze_image(image, ground_truth, filename):
         
         # Predicted mask
         if ground_truth is not None:
-            axes[0,2].imshow(binary_mask, cmap='gray')
+            axes[0,2].imshow(binary_mask_np, cmap='gray')
             axes[0,2].set_title('Predicted Mask', fontsize=12, weight='bold')
             axes[0,2].axis('off')
             
             # Ground truth
             gt_array = np.array(ground_truth.resize((256, 256)))
-            axes[1,0].imshow(gt_array, cmap='gray')
+            # Normalize ground truth to binary (0 or 1)
+            gt_binary = (gt_array > 128).astype(np.uint8)
+            
+            axes[1,0].imshow(gt_binary, cmap='gray')
             axes[1,0].set_title('Ground Truth Mask', fontsize=12, weight='bold')
             axes[1,0].axis('off')
             
             # Overlay comparison
             overlay = np.array(image.convert('RGB').resize((256, 256)))
-            overlay[binary_mask > 0] = [0, 255, 0]  # Green for prediction
-            overlay[gt_array > 128] = [255, 0, 0]   # Red for ground truth
+            overlay[binary_mask_np > 0] = [0, 255, 0]  # Green for prediction
+            overlay[gt_binary > 0] = [255, 0, 0]       # Red for ground truth
             axes[1,1].imshow(overlay)
             axes[1,1].set_title('Prediction (Green) vs GT (Red)', fontsize=12, weight='bold')
             axes[1,1].axis('off')
             
-            # Calculate IoU
-            pred_binary = binary_mask > 0
-            gt_binary = gt_array > 128
-            intersection = np.sum(pred_binary & gt_binary)
-            union = np.sum(pred_binary | gt_binary)
+            # Calculate IoU and Dice
+            pred_binary = binary_mask_np > 0
+            gt_binary_bool = gt_binary > 0
+            intersection = np.sum(pred_binary & gt_binary_bool)
+            union = np.sum(pred_binary | gt_binary_bool)
             iou = intersection / (union + 1e-8)
             
             # Dice score
-            dice = (2 * intersection) / (np.sum(pred_binary) + np.sum(gt_binary) + 1e-8)
+            dice = (2 * intersection) / (np.sum(pred_binary) + np.sum(gt_binary_bool) + 1e-8)
             
             axes[1,2].text(0.1, 0.6, f'IoU: {iou:.4f}', fontsize=16, weight='bold')
             axes[1,2].text(0.1, 0.4, f'Dice: {dice:.4f}', fontsize=16, weight='bold')
@@ -374,13 +355,13 @@ def analyze_image(image, ground_truth, filename):
             axes[1,2].set_ylim(0, 1)
             axes[1,2].axis('off')
         else:
-            axes[1,0].imshow(binary_mask, cmap='gray')
+            axes[1,0].imshow(binary_mask_np, cmap='gray')
             axes[1,0].set_title('Predicted Mask', fontsize=12, weight='bold')
             axes[1,0].axis('off')
             
             # Overlay
             overlay = np.array(image.convert('RGB').resize((256, 256)))
-            overlay[binary_mask > 0] = [255, 0, 0]
+            overlay[binary_mask_np > 0] = [255, 0, 0]
             axes[1,1].imshow(overlay)
             axes[1,1].set_title('Prediction Overlay', fontsize=12, weight='bold')
             axes[1,1].axis('off')
@@ -396,8 +377,8 @@ def analyze_image(image, ground_truth, filename):
         result_image = Image.open(buf)
         
         # Generate analysis text
-        tumor_pixels = np.sum(binary_mask)
-        total_pixels = binary_mask.size
+        tumor_pixels = np.sum(binary_mask_np)
+        total_pixels = binary_mask_np.size
         tumor_percentage = (tumor_pixels / total_pixels) * 100
         
         analysis_text = f"""
@@ -410,7 +391,6 @@ def analyze_image(image, ground_truth, filename):
 - Tumor Pixels: {tumor_pixels:,}
 
 **Model Features:**
-- Test-Time Augmentation: Applied
 - Attention Visualization: Generated
 - Post-processing: Morphological cleanup
 """
@@ -425,7 +405,10 @@ def analyze_image(image, ground_truth, filename):
         return result_image, analysis_text
         
     except Exception as e:
-        return None, f"Analysis failed: {str(e)}"
+        import traceback
+        error_msg = f"Analysis failed: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
+        print(error_msg)  # For debugging
+        return None, error_msg
 
 # Initialize model and dataset at startup
 print("Initializing application components...")
@@ -475,7 +458,7 @@ with gr.Blocks(css=css, title="Brain Tumor Segmentation Analysis") as app:
     
     **Advanced Medical Image Analysis Tool**
     
-    Features: Test-Time Augmentation, Attention Visualization, Dataset Integration
+    Features: Attention Visualization, Dataset Integration, Morphological Post-processing
     """)
     
     # Status display
@@ -567,4 +550,4 @@ if __name__ == "__main__":
         server_port=7860,
         show_error=True,
         share=False
-    )
+    )