Update app.py

app.py

@@ -266,181 +266,191 @@ def generate_attention_heatmap(attention_maps):
     return heatmap
 
 def analyze_image(image, ground_truth, filename):
-    """Main analysis function - FIXED VERSION matching the working notebook"""
+    """
+    Robust replacement for the original analyze_image.
+    - Fixes broadcasting issues between 2D masks and 3-channel images.
+    - Converts attention heatmap (BGR from OpenCV) to RGB for correct plotting.
+    - Ensures masks are strict binary uint8 arrays.
+    - Returns (PIL.Image result_plot, markdown_text).
+    """
     if model is None:
         return None, "Model not loaded. Please restart the application."
-    
+
     if image is None:
         return None, "Please select an image first."
-    
+
     try:
-        print("="*50)
+        print("=" * 50)
         print("DEBUG: Starting analysis...")
         print(f"Input image mode: {image.mode}")
         print(f"Input image size: {image.size}")
-        
-        # Preprocess - exactly like the working notebook
+
+        # Preprocess - keeps same behavior as notebook
         input_tensor = preprocess_for_model(image).to(device)
         print(f"Input tensor shape: {input_tensor.shape}")
         print(f"Input tensor min/max: {input_tensor.min():.4f}/{input_tensor.max():.4f}")
-        
+
         # Get prediction and attention maps
         with torch.no_grad():
             print("Getting model output...")
             model_output, attention_maps = model(input_tensor)
-            
+
+            # model_output shape expected: [1, 1, 256, 256]
             print(f"Model output shape: {model_output.shape}")
             print(f"Model output min/max BEFORE sigmoid: {model_output.min():.4f}/{model_output.max():.4f}")
-            
-            # Apply sigmoid and threshold - EXACTLY like the working notebook
-            pred_mask = torch.sigmoid(model_output)
-            print(f"After sigmoid min/max: {pred_mask.min():.4f}/{pred_mask.max():.4f}")
-            
-            # Apply threshold to get binary mask
-            binary_mask = (pred_mask > 0.5).float()
-            print(f"Binary mask sum (number of 1s): {binary_mask.sum()}")
-            
-            # Convert to numpy - following notebook approach
-            pred_mask_np = binary_mask.cpu().squeeze().numpy()
+
+            pred_prob = torch.sigmoid(model_output)  # probabilities in [0,1]
+            print(f"After sigmoid min/max: {pred_prob.min():.4f}/{pred_prob.max():.4f}")
+
+            # DEFAULT THRESHOLD: 0.5 (same as your notebook). Change if debugging low-confidence.
+            pred_mask = (pred_prob > 0.5).float()
+            print(f"Binary mask sum (number of 1s): {pred_mask.sum():.4f}")
+
+            # Convert prediction to numpy
+            pred_mask_np = pred_mask.cpu().squeeze().numpy()     # shape: (H, W)
             print(f"Numpy binary mask shape: {pred_mask_np.shape}")
             print(f"Numpy binary mask unique values: {np.unique(pred_mask_np)}")
             print(f"Numpy binary mask sum: {np.sum(pred_mask_np)}")
-        
-        # Create visualization mask like in the notebook
-        # The notebook uses: inv_pred_mask_np = np.where(pred_mask_np == 1, 0, 255)
-        # This inverts the mask for better visualization
-        inv_pred_mask_np = np.where(pred_mask_np == 1, 0, 255)
-        
-        # Generate attention heatmap
+
+        # Create attention heatmap (the helper resizes & returns a 3-channel BGR heatmap)
         print("Generating attention heatmap...")
-        att_heatmap = generate_attention_heatmap(attention_maps)
-        print(f"Attention heatmap shape: {att_heatmap.shape}")
-        
-        # Prepare original image array
-        original_np = np.array(image.resize((256, 256)))
-        
-        # Create tumor-only image (like in notebook)
-        tumor_only = np.where(pred_mask_np == 1, original_np, 255)
-        
-        # Create visualization
+        att_heatmap = generate_attention_heatmap(attention_maps)  # likely BGR (cv2)
+        print(f"Raw attention heatmap shape: {att_heatmap.shape}")
+
+        # Convert heatmap to RGB (OpenCV returns BGR)
+        if att_heatmap is not None and att_heatmap.size != 0:
+            try:
+                att_heatmap = cv2.cvtColor(att_heatmap, cv2.COLOR_BGR2RGB)
+            except Exception:
+                # if conversion fails, proceed with what we have
+                pass
+
+        # Prepare original image arrays:
+        original_gray = np.array(image.convert('L').resize((256, 256))).astype(np.uint8)   # 2D
+        original_rgb = np.array(image.convert('RGB').resize((256, 256))).astype(np.uint8) # 3D
+
+        # Ensure pred_mask_np is strict binary 0/1 uint8
+        pred_mask_bin = (pred_mask_np > 0.5).astype(np.uint8)  # shape: (256,256), dtype: uint8
+
+        # Inverted predicted mask for visualization (white background, tumor black)
+        inv_pred_mask_np = np.where(pred_mask_bin == 1, 0, 255).astype(np.uint8)
+
+        # Tumor-only images:
+        tumor_only_gray = np.where(pred_mask_bin == 1, original_gray, 255).astype(np.uint8)
+        tumor_only_rgb = original_rgb.copy()
+        tumor_only_rgb[pred_mask_bin == 0] = 255
+
+        # Begin plotting (match existing layout: 2x4 with GT or 2x3 without)
         if ground_truth is not None:
             fig, axes = plt.subplots(2, 4, figsize=(16, 8))
         else:
             fig, axes = plt.subplots(2, 3, figsize=(15, 8))
-            
+
         fig.suptitle('Brain Tumor Segmentation Analysis', fontsize=16, weight='bold')
-        
-        # Row 1: Original, Attention, Predicted Mask, Tumor Only
-        axes[0,0].imshow(original_np, cmap='gray')
-        axes[0,0].set_title('Original Image', fontsize=12, weight='bold')
-        axes[0,0].axis('off')
-        
-        # Attention heatmap overlay
-        axes[0,1].imshow(original_np, cmap='gray')
-        axes[0,1].imshow(att_heatmap, alpha=0.4)
-        axes[0,1].set_title('Attention Heatmap', fontsize=12, weight='bold')
-        axes[0,1].axis('off')
-        
+
+        # Row 1: Original, Attention, Predicted Mask, Tumor Only (if GT exists show 4th)
+        axes[0, 0].imshow(original_gray, cmap='gray')
+        axes[0, 0].set_title('Original Image', fontsize=12, weight='bold')
+        axes[0, 0].axis('off')
+
+        # Attention overlay on RGB original (blend)
+        axes[0, 1].imshow(original_rgb)
+        if att_heatmap is not None and att_heatmap.size != 0:
+            axes[0, 1].imshow(att_heatmap, alpha=0.4)
+        axes[0, 1].set_title('Attention Heatmap', fontsize=12, weight='bold')
+        axes[0, 1].axis('off')
+
         # Predicted mask (inverted for visualization)
-        axes[0,2].imshow(inv_pred_mask_np, cmap='gray')
-        axes[0,2].set_title('Predicted Mask', fontsize=12, weight='bold')
-        axes[0,2].axis('off')
-        
+        axes[0, 2].imshow(inv_pred_mask_np, cmap='gray')
+        axes[0, 2].set_title('Predicted Mask', fontsize=12, weight='bold')
+        axes[0, 2].axis('off')
+
         if ground_truth is not None:
+            axes[0, 3].imshow(tumor_only_rgb)
+            axes[0, 3].set_title('Tumor Only', fontsize=12, weight='bold')
+            axes[0, 3].axis('off')
+
             # Ground truth processing - convert to binary like notebook
-            gt_array = np.array(ground_truth.resize((256, 256)))
-            # Apply same preprocessing as notebook
             val_test_transform = transforms.Compose([
-                transforms.Resize((256,256)),
+                transforms.Resize((256, 256)),
                 transforms.ToTensor()
             ])
             mask_np = val_test_transform(ground_truth).cpu().squeeze().numpy()
-            
-            print(f"Ground truth array shape: {gt_array.shape}")
-            print(f"Ground truth unique values: {np.unique(gt_array)}")
-            
-            # Tumor only image
-            axes[0,3].imshow(tumor_only, cmap='gray')
-            axes[0,3].set_title('Tumor Only', fontsize=12, weight='bold')
-            axes[0,3].axis('off')
-            
-            # Row 2: Ground truth, overlay comparison, metrics
-            axes[1,0].imshow(mask_np, cmap='gray')
-            axes[1,0].set_title('Ground Truth Mask', fontsize=12, weight='bold')
-            axes[1,0].axis('off')
-            
-            # Overlay comparison - following notebook style
-            overlay = np.array(image.convert('RGB').resize((256, 256)))
-            overlay[pred_mask_np == 1] = [0, 255, 0]  # Green for prediction
-            overlay[mask_np > 0.5] = [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 and Dice exactly like notebook
-            intersection = np.logical_and(pred_mask_np, mask_np).sum()
-            union = np.logical_or(pred_mask_np, mask_np).sum()
+            mask_bin = (mask_np > 0.5).astype(np.uint8)
+
+            print(f"Ground truth array shape: {np.array(ground_truth.resize((256,256))).shape}")
+            print(f"Ground truth unique values: {np.unique(np.array(ground_truth.resize((256,256))))}")
+
+            # Row 2: Ground truth, overlay comparison, metrics, segmented tumor
+            axes[1, 0].imshow(mask_bin, cmap='gray')
+            axes[1, 0].set_title('Ground Truth Mask', fontsize=12, weight='bold')
+            axes[1, 0].axis('off')
+
+            overlay = original_rgb.copy()
+            overlay[pred_mask_bin == 1] = [0, 255, 0]   # predicted green
+            overlay[mask_bin == 1] = [255, 0, 0]       # ground truth red
+            axes[1, 1].imshow(overlay)
+            axes[1, 1].set_title('Prediction (Green) vs GT (Red)', fontsize=12, weight='bold')
+            axes[1, 1].axis('off')
+
+            # Metrics calculation (IoU and Dice)
+            intersection = np.logical_and(pred_mask_bin, mask_bin).sum()
+            union = np.logical_or(pred_mask_bin, mask_bin).sum()
             iou = intersection / (union + 1e-7)
-            
-            # Dice score
-            dice = (2 * intersection) / (pred_mask_np.sum() + mask_np.sum() + 1e-7)
-            
+            dice = (2 * intersection) / (pred_mask_bin.sum() + mask_bin.sum() + 1e-7)
+
             print(f"Final IoU: {iou:.4f}")
             print(f"Final Dice: {dice:.4f}")
             print(f"Intersection: {intersection}")
             print(f"Union: {union}")
-            print(f"Pred pixels: {np.sum(pred_mask_np)}")
-            print(f"GT pixels: {np.sum(mask_np > 0.5)}")
-            
-            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')
-            axes[1,2].set_xlim(0, 1)
-            axes[1,2].set_ylim(0, 1)
-            axes[1,2].axis('off')
-            axes[1,2].set_title('Metrics', fontsize=12, weight='bold')
-            
-            # Additional tumor statistics
-            axes[1,3].imshow(tumor_only, cmap='gray')
-            axes[1,3].set_title('Segmented Tumor', fontsize=12, weight='bold')
-            axes[1,3].axis('off')
-            
+            print(f"Pred pixels: {np.sum(pred_mask_bin)}")
+            print(f"GT pixels: {np.sum(mask_bin)}")
+
+            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')
+            axes[1, 2].set_xlim(0, 1)
+            axes[1, 2].set_ylim(0, 1)
+            axes[1, 2].axis('off')
+            axes[1, 2].set_title('Metrics', fontsize=12, weight='bold')
+
+            axes[1, 3].imshow(tumor_only_gray, cmap='gray')
+            axes[1, 3].set_title('Segmented Tumor', fontsize=12, weight='bold')
+            axes[1, 3].axis('off')
+
         else:
             # No ground truth case
-            axes[1,0].imshow(inv_pred_mask_np, cmap='gray')
-            axes[1,0].set_title('Predicted Mask', fontsize=12, weight='bold')
-            axes[1,0].axis('off')
-            
-            # Tumor only
-            axes[1,1].imshow(tumor_only, cmap='gray')
-            axes[1,1].set_title('Tumor Only', fontsize=12, weight='bold')
-            axes[1,1].axis('off')
-            
-            # Overlay
-            overlay = np.array(image.convert('RGB').resize((256, 256)))
-            overlay[pred_mask_np == 1] = [255, 0, 0]
-            axes[1,2].imshow(overlay)
-            axes[1,2].set_title('Prediction Overlay', fontsize=12, weight='bold')
-            axes[1,2].axis('off')
+            axes[1, 0].imshow(inv_pred_mask_np, cmap='gray')
+            axes[1, 0].set_title('Predicted Mask', fontsize=12, weight='bold')
+            axes[1, 0].axis('off')
+
+            axes[1, 1].imshow(tumor_only_gray, cmap='gray')
+            axes[1, 1].set_title('Tumor Only', fontsize=12, weight='bold')
+            axes[1, 1].axis('off')
+
+            overlay = original_rgb.copy()
+            overlay[pred_mask_bin == 1] = [255, 0, 0]  # red for prediction overlay
+            axes[1, 2].imshow(overlay)
+            axes[1, 2].set_title('Prediction Overlay', fontsize=12, weight='bold')
+            axes[1, 2].axis('off')
 
         plt.tight_layout()
-        
-        # Save plot
+
+        # Save plot to buffer and return as PIL image
         buf = io.BytesIO()
         plt.savefig(buf, format='png', dpi=150, bbox_inches='tight', facecolor='white')
         buf.seek(0)
         plt.close()
-        
-        result_image = Image.open(buf)
-        
-        # Generate analysis text
-        tumor_pixels = np.sum(pred_mask_np)
-        total_pixels = pred_mask_np.size
-        tumor_percentage = (tumor_pixels / total_pixels) * 100
-        
+        result_image = Image.open(buf).convert("RGB")
+
+        # Analysis text: tumor area
+        tumor_pixels = int(np.sum(pred_mask_bin))
+        total_pixels = int(pred_mask_bin.size)
+        tumor_percentage = (tumor_pixels / total_pixels) * 100 if total_pixels > 0 else 0.0
+
         print(f"Final tumor pixels: {tumor_pixels}")
         print(f"Final tumor percentage: {tumor_percentage:.2f}%")
-        print("="*50)
-        
+        print("=" * 50)
+
         analysis_text = f"""
 # Analysis Results
 
@@ -454,20 +464,20 @@ def analyze_image(image, ground_truth, filename):
 - Attention Visualization: Generated
 - Post-processing: Applied
 """
-        
+
         if ground_truth is not None:
             analysis_text += f"""
 **Performance Metrics:**
 - IoU Score: {iou:.4f}
 - Dice Score: {dice:.4f}
 """
-        
+
         return result_image, analysis_text
-        
+
     except Exception as e:
         import traceback
         error_msg = f"Analysis failed: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
-        print(error_msg)  # For debugging
+        print(error_msg)
         return None, error_msg