Update app.py

app.py

@@ -226,7 +226,7 @@ def get_random_sample():
         return None, None, f"Error loading sample: {e}"
 
 def preprocess_for_model(image):
-    """Preprocessing for your model"""
+    """Preprocessing for your model - matches the working notebook"""
     if image.mode != 'L':
         image = image.convert('L')
     
@@ -238,7 +238,7 @@ def preprocess_for_model(image):
     return transform(image).unsqueeze(0)
 
 def generate_attention_heatmap(attention_maps):
-    """Generate attention heatmap - Fixed version"""
+    """Generate attention heatmap"""
     if not attention_maps:
         return np.zeros((256, 256, 3))
     
@@ -266,7 +266,7 @@ def generate_attention_heatmap(attention_maps):
     return heatmap
 
 def analyze_image(image, ground_truth, filename):
-    """Main analysis function - DEBUG VERSION"""
+    """Main analysis function - FIXED VERSION matching the working notebook"""
     if model is None:
         return None, "Model not loaded. Please restart the application."
     
@@ -279,7 +279,7 @@ def analyze_image(image, ground_truth, filename):
         print(f"Input image mode: {image.mode}")
         print(f"Input image size: {image.size}")
         
-        # Preprocess
+        # Preprocess - exactly like the working 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}")
@@ -292,128 +292,135 @@ def analyze_image(image, ground_truth, filename):
             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
+            # 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}")
             
-            # Check values before thresholding
-            unique_vals = torch.unique(pred_mask)
-            print(f"Unique values in prediction: {unique_vals[:10]}")  # Show first 10 unique values
-            
-            # Apply threshold
+            # Apply threshold to get binary mask
             binary_mask = (pred_mask > 0.5).float()
-            print(f"Binary mask shape: {binary_mask.shape}")
             print(f"Binary mask sum (number of 1s): {binary_mask.sum()}")
             
-            # Convert to numpy
-            binary_mask_np = binary_mask.squeeze().cpu().numpy()
-            print(f"Numpy binary mask shape: {binary_mask_np.shape}")
-            print(f"Numpy binary mask unique values: {np.unique(binary_mask_np)}")
-            print(f"Numpy binary mask sum: {np.sum(binary_mask_np)}")
-        
-        # Try different thresholds if 0.5 doesn't work
-        if np.sum(binary_mask_np) == 0:
-            print("No pixels detected with threshold 0.5, trying lower thresholds...")
-            for thresh in [0.3, 0.2, 0.1, 0.05]:
-                test_mask = (pred_mask > thresh).float().squeeze().cpu().numpy()
-                pixel_count = np.sum(test_mask)
-                print(f"Threshold {thresh}: {pixel_count} pixels")
-                if pixel_count > 0:
-                    print(f"Using threshold {thresh} instead of 0.5")
-                    binary_mask_np = test_mask
-                    break
-        
-        # Post-processing (morphological operations)
-        print("Applying morphological operations...")
-        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3,3))
-        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)
-        print(f"After morphological ops sum: {np.sum(binary_mask_np)}")
+            # Convert to numpy - following notebook approach
+            pred_mask_np = binary_mask.cpu().squeeze().numpy()
+            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
         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
         if ground_truth is not None:
-            fig, axes = plt.subplots(2, 3, figsize=(15, 10))
+            fig, axes = plt.subplots(2, 4, figsize=(16, 8))
         else:
-            fig, axes = plt.subplots(2, 2, figsize=(12, 10))
+            fig, axes = plt.subplots(2, 3, figsize=(15, 8))
             
         fig.suptitle('Brain Tumor Segmentation Analysis', fontsize=16, weight='bold')
         
-        # Original image
-        axes[0,0].imshow(image, cmap='gray')
+        # 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
-        axes[0,1].imshow(image, cmap='gray')
+        # 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')
         
-        # Predicted mask
+        # 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')
+        
         if ground_truth is not None:
-            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
+            # 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.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)}")
             
-            # Normalize ground truth to binary (0 or 1)
-            gt_binary = (gt_array > 128).astype(np.uint8)
-            print(f"GT binary sum: {np.sum(gt_binary)}")
+            # 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')
             
-            axes[1,0].imshow(gt_binary, cmap='gray')
+            # 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
+            # Overlay comparison - following notebook style
             overlay = np.array(image.convert('RGB').resize((256, 256)))
-            overlay[binary_mask_np > 0] = [0, 255, 0]  # Green for prediction
-            overlay[gt_binary > 0] = [255, 0, 0]       # Red for ground truth
+            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
-            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)
+            # 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()
+            iou = intersection / (union + 1e-7)
             
             # Dice score
-            dice = (2 * intersection) / (np.sum(pred_binary) + np.sum(gt_binary_bool) + 1e-8)
+            dice = (2 * intersection) / (pred_mask_np.sum() + mask_np.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_binary)}")
-            print(f"GT pixels: {np.sum(gt_binary_bool)}")
+            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')
+            
         else:
-            axes[1,0].imshow(binary_mask_np, cmap='gray')
+            # 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[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')
+            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')
 
         plt.tight_layout()
         
@@ -426,8 +433,8 @@ def analyze_image(image, ground_truth, filename):
         result_image = Image.open(buf)
         
         # Generate analysis text
-        tumor_pixels = np.sum(binary_mask_np)
-        total_pixels = binary_mask_np.size
+        tumor_pixels = np.sum(pred_mask_np)
+        total_pixels = pred_mask_np.size
         tumor_percentage = (tumor_pixels / total_pixels) * 100
         
         print(f"Final tumor pixels: {tumor_pixels}")
@@ -445,7 +452,7 @@ def analyze_image(image, ground_truth, filename):
 
 **Model Features:**
 - Attention Visualization: Generated
-- Post-processing: Morphological cleanup
+- Post-processing: Applied
 """
         
         if ground_truth is not None: