Update app.py

app.py

@@ -265,13 +265,14 @@ def generate_attention_heatmap(attention_maps):
     
     return heatmap
 
-def analyze_image(image, ground_truth, filename):
+def analyze_image(image, ground_truth, filename, debug=True):
     """
-    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).
+    Replacement analyze_image that:
+    - Accepts model returning either logits or (logits, attention_maps)
+    - Prints detailed stats and shapes
+    - Produces prob heatmap (no threshold) for debugging
+    - Fixes broadcasting/color issues for visualization
+    - Returns (PIL.Image, markdown_text)
     """
     if model is None:
         return None, "Model not loaded. Please restart the application."
@@ -285,153 +286,146 @@ def analyze_image(image, ground_truth, filename):
         print(f"Input image mode: {image.mode}")
         print(f"Input image size: {image.size}")
 
-        # Preprocess - keeps same behavior as notebook
-        input_tensor = preprocess_for_model(image).to(device)
+        # Preprocess - same as your notebook/app
+        input_tensor = preprocess_for_model(image).to(device)  # shape [1,1,256,256]
         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
+        # Forward pass and robust unpacking (support both return styles)
         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}")
-
-            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 attention heatmap (the helper resizes & returns a 3-channel BGR heatmap)
-        print("Generating attention heatmap...")
-        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)
+            out = model(input_tensor)
+            # If model returned tuple/list: (logits, attention_maps)
+            if isinstance(out, (list, tuple)) and len(out) == 2:
+                logits, attention_maps = out
+            else:
+                # assume out is logits tensor and no attention maps were returned
+                logits = out
+                attention_maps = []
+
+            # Ensure logits is a tensor
+            if not torch.is_tensor(logits):
+                raise RuntimeError("Model output is not a tensor. Check model forward() return type.")
+
+            print(f"Model output (logits) shape: {logits.shape}")
+            print(f"Model output min/max BEFORE sigmoid: {logits.min():.4f}/{logits.max():.4f}")
+
+            # Probabilities (sigmoid)
+            pred_prob = torch.sigmoid(logits)
+            print(f"Pred prob min/max: {pred_prob.min():.4f}/{pred_prob.max():.4f}")
+
+            # Convert to numpy for visualization; keep a float prob map for the heatmap
+            pred_prob_np = pred_prob.cpu().squeeze().numpy()  # shape (H, W)
+            pred_mask_bin = (pred_prob_np > 0.5).astype(np.uint8)  # default threshold 0.5
+
+            print(f"Binary mask (0.5 threshold) sum: {pred_mask_bin.sum()}")
+
+        # Debug: print attention maps shapes and stats
+        if debug:
+            print("Attention maps info:")
+            for i, att in enumerate(attention_maps):
+                try:
+                    att_np = att.squeeze().cpu().numpy()
+                    print(f"  att[{i}] shape: {att_np.shape} min/max: {att_np.min():.4f}/{att_np.max():.4f}")
+                except Exception as ex:
+                    print(f"  att[{i}] inspect failed: {ex}")
+
+        # Build prob heatmap (no threshold) for debugging
+        try:
+            prob_resized = cv2.resize(pred_prob_np, (256, 256)) if pred_prob_np.shape != (256, 256) else pred_prob_np
+            prob_norm = (prob_resized - prob_resized.min()) / (prob_resized.max() - prob_resized.min() + 1e-8)
+            prob_heatmap_bgr = cv2.applyColorMap((prob_norm * 255).astype(np.uint8), cv2.COLORMAP_JET)
+            prob_heatmap = cv2.cvtColor(prob_heatmap_bgr, cv2.COLOR_BGR2RGB)
+        except Exception:
+            prob_heatmap = np.zeros((256, 256, 3), dtype=np.uint8)
+
+        # Generate attention heatmap (reuse your function), convert BGR->RGB
+        att_heatmap = generate_attention_heatmap(attention_maps)
         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
+        # Prepare images (gray and rgb)
+        original_gray = np.array(image.convert('L').resize((256, 256))).astype(np.uint8)
+        original_rgb  = np.array(image.convert('RGB').resize((256, 256))).astype(np.uint8)
 
-        # 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)
+        # Ensure binary mask dtype/shape consistency
+        pred_mask_bin = (pred_mask_bin > 0).astype(np.uint8)
         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  = original_rgb.copy()
         tumor_only_rgb[pred_mask_bin == 0] = 255
 
-        # Begin plotting (match existing layout: 2x4 with GT or 2x3 without)
+        # Decide grid: show prob heatmap next to attention so you can compare
         if ground_truth is not None:
-            fig, axes = plt.subplots(2, 4, figsize=(16, 8))
+            fig, axes = plt.subplots(3, 4, figsize=(16, 12))  # add an extra row for debug heatmap
         else:
-            fig, axes = plt.subplots(2, 3, figsize=(15, 8))
-
-        fig.suptitle('Brain Tumor Segmentation Analysis', fontsize=16, weight='bold')
+            fig, axes = plt.subplots(3, 3, figsize=(15, 12))
 
-        # 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')
+        fig.suptitle('Brain Tumor Segmentation Analysis (debug)', fontsize=18, weight='bold')
 
-        # Attention overlay on RGB original (blend)
-        axes[0, 1].imshow(original_rgb)
+        # Row 1
+        axes[0,0].imshow(original_gray, cmap='gray'); axes[0,0].set_title('Original'); axes[0,0].axis('off')
+        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,1].imshow(att_heatmap, alpha=0.45)
+        axes[0,1].set_title('Attention Heatmap (overlay)'); axes[0,1].axis('off')
+        axes[0,2].imshow(inv_pred_mask_np, cmap='gray'); axes[0,2].set_title('Pred Mask (inv)'); 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 (RGB)'); axes[0,3].axis('off')
 
+        # Row 2
         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
-            val_test_transform = transforms.Compose([
-                transforms.Resize((256, 256)),
-                transforms.ToTensor()
-            ])
+            # show GT and overlay and metrics
+            val_test_transform = transforms.Compose([transforms.Resize((256,256)), transforms.ToTensor()])
             mask_np = val_test_transform(ground_truth).cpu().squeeze().numpy()
             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')
-
+            axes[1,0].imshow(mask_bin, cmap='gray'); axes[1,0].set_title('Ground Truth Mask'); 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')
+            overlay[pred_mask_bin == 1] = [0,255,0]
+            overlay[mask_bin == 1] = [255,0,0]
+            axes[1,1].imshow(overlay); axes[1,1].set_title('Prediction (G) vs GT (R)'); 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 = (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_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')
+            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].axis('off'); axes[1,2].set_title('Metrics')
 
+            axes[1,3].imshow(tumor_only_gray, cmap='gray'); axes[1,3].set_title('Segmented Tumor'); 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')
-
-            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')
+            # No GT: second row shows predicted mask, tumor only and overlay
+            axes[1,0].imshow(inv_pred_mask_np, cmap='gray'); axes[1,0].set_title('Predicted Mask'); axes[1,0].axis('off')
+            axes[1,1].imshow(tumor_only_gray, cmap='gray'); axes[1,1].set_title('Tumor Only'); axes[1,1].axis('off')
+            overlay = original_rgb.copy(); overlay[pred_mask_bin==1] = [255,0,0]
+            axes[1,2].imshow(overlay); axes[1,2].set_title('Prediction Overlay'); axes[1,2].axis('off')
+
+        # Row 3 (debug): probability heatmap + (optional) raw att channel thumbnails
+        axes[2,0].imshow(original_rgb); axes[2,0].imshow(prob_heatmap, alpha=0.5); axes[2,0].set_title('Prob Heatmap (overlay)'); axes[2,0].axis('off')
+        # show the plain probability heatmap
+        axes[2,1].imshow(prob_heatmap); axes[2,1].set_title('Prob Heatmap (plain)'); axes[2,1].axis('off')
+
+        # if we have attention maps, show up to two scaled maps for quick check
+        if len(attention_maps) >= 1:
+            try:
+                att0 = attention_maps[0].squeeze().cpu().numpy()
+                att0 = cv2.resize((att0 - att0.min())/(att0.max()-att0.min()+1e-8), (256,256))
+                axes[2,2].imshow(att0, cmap='viridis'); axes[2,2].set_title('Att map 0 (rescaled)'); axes[2,2].axis('off')
+            except Exception:
+                axes[2,2].axis('off')
+        else:
+            axes[2,2].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')
+        # hide any unused axes (robust)
+        for ax_row in axes.reshape(-1):
+            if not hasattr(ax_row, 'has_data') or ax_row.images == []:
+                ax_row.axis('off')
 
         plt.tight_layout()
 
@@ -442,15 +436,11 @@ def analyze_image(image, ground_truth, filename):
         plt.close()
         result_image = Image.open(buf).convert("RGB")
 
-        # Analysis text: tumor area
-        tumor_pixels = int(np.sum(pred_mask_bin))
+        # Numeric analysis text
+        tumor_pixels = int(pred_mask_bin.sum())
         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)
-
         analysis_text = f"""
 # Analysis Results
 
@@ -462,7 +452,7 @@ def analyze_image(image, ground_truth, filename):
 
 **Model Features:**
 - Attention Visualization: Generated
-- Post-processing: Applied
+- Probability Heatmap: Generated
 """
 
         if ground_truth is not None:
@@ -472,6 +462,10 @@ def analyze_image(image, ground_truth, filename):
 - Dice Score: {dice:.4f}
 """
 
+        # Extra helpful hint when predictions are all zero
+        if debug and pred_prob_np.max() < 0.5:
+            analysis_text += "\n\n**Debug hint:** model probabilities are low (max < 0.5). Try lowering threshold (e.g. 0.3) or inspect model weights/loading."
+
         return result_image, analysis_text
 
     except Exception as e:
@@ -480,7 +474,7 @@ def analyze_image(image, ground_truth, filename):
         print(error_msg)
         return None, error_msg
 
-        
+
 # Initialize model and dataset at startup
 print("Initializing application components...")
 model_loaded = download_and_load_model()