Update app.py

app.py

@@ -1,3 +1,4 @@
+# full_app_with_heatmap.py
 import gradio as gr
 import torch
 import torch.nn as nn
@@ -6,26 +7,26 @@ import cv2
 from PIL import Image
 import matplotlib.pyplot as plt
 import io
-import torchvision.transforms as transforms
+from torchvision import transforms
 import torchvision.transforms.functional as TF
-import random
-import os
 import urllib.request
-import kagglehub
+import os
+import random
 from glob import glob
+import kagglehub  # if you use dataset download in the app; remove if not needed
 
-# Global variables - loaded once at startup
+# -------------------------
+# Setup / Globals
+# -------------------------
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = None
 dataset_images = []
 dataset_masks = []
 dataset_loaded = False
 
-print("="*50)
-print("BRAIN TUMOR SEGMENTATION APPLICATION")
-print("="*50)
-
-# Your Attention U-Net classes (unchanged)
+# -------------------------
+# Model classes (Attention U-Net)
+# -------------------------
 class DoubleConv(nn.Module):
     def __init__(self, in_channels, out_channels):
         super(DoubleConv, self).__init__()
@@ -49,26 +50,24 @@ class AttentionBlock(nn.Module):
             nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True),
             nn.BatchNorm2d(F_int)
         )
-        
         self.W_x = nn.Sequential(
             nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=True),
             nn.BatchNorm2d(F_int)
         )
-        
         self.psi = nn.Sequential(
             nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=True),
             nn.BatchNorm2d(1),
             nn.Sigmoid()
         )
-        
         self.relu = nn.ReLU(inplace=True)
-        
+
     def forward(self, g, x):
         g1 = self.W_g(g)
         x1 = self.W_x(x)
         psi = self.relu(g1 + x1)
         psi = self.psi(psi)
-        return x * psi, psi  # Return both attended features AND attention map
+        return x * psi, psi  # return attended skip, attention map
+
 
 class AttentionUNET(nn.Module):
     def __init__(self, in_channels=1, out_channels=1, features=[32, 64, 128, 256]):
@@ -84,7 +83,7 @@ class AttentionUNET(nn.Module):
             in_channels = feature
 
         self.bottleneck = DoubleConv(features[-1], features[-1]*2)
-        
+
         for feature in reversed(features):
             self.ups.append(nn.ConvTranspose2d(feature*2, feature, kernel_size=2, stride=2))
             self.attentions.append(AttentionBlock(F_g=feature, F_l=feature, F_int=feature // 2))
@@ -106,7 +105,7 @@ class AttentionUNET(nn.Module):
 
         for idx in range(0, len(self.ups), 2):
             x = self.ups[idx](x)
-            skip_connection = skip_connections[idx//2]
+            skip_connection = skip_connections[idx // 2]
 
             if x.shape != skip_connection.shape:
                 x = TF.resize(x, size=skip_connection.shape[2:])
@@ -118,15 +117,16 @@ class AttentionUNET(nn.Module):
 
         return self.final_conv(x), attention_maps
 
+# -------------------------
+# Model download / load
+# -------------------------
 def download_and_load_model():
-    """Download and load model once at startup"""
     global model
     print("Loading Attention U-Net model...")
-    
+
     model_url = "https://huggingface.co/spaces/ArchCoder/the-op-segmenter/resolve/main/best_attention_model.pth.tar"
     model_path = "best_attention_model.pth.tar"
-    
-    # Download model if needed
+
     if not os.path.exists(model_path):
         print("Downloading model weights...")
         try:
@@ -134,90 +134,70 @@ def download_and_load_model():
         except Exception as e:
             print(f"Failed to download model: {e}")
             return False
-    
-    # Load model
+
     try:
         model = AttentionUNET(in_channels=1, out_channels=1).to(device)
-        checkpoint = torch.load(model_path, map_location=device, weights_only=True)
-        model.load_state_dict(checkpoint["state_dict"])
+        checkpoint = torch.load(model_path, map_location=device)
+        # checkpoint format expected to have "state_dict"
+        if "state_dict" in checkpoint:
+            sd = checkpoint["state_dict"]
+        else:
+            sd = checkpoint
+        # Try exact load; if mismatch, try strict=False and warn
+        try:
+            model.load_state_dict(sd)
+        except Exception as ex:
+            print(f"Warning: strict load failed: {ex}. Trying strict=False...")
+            model.load_state_dict(sd, strict=False)
         model.eval()
         print("✓ Model loaded successfully!")
         return True
     except Exception as e:
         print(f"Failed to load model: {e}")
+        model = None
         return False
 
+# -------------------------
+# Dataset utilities (optional)
+# -------------------------
 def download_and_load_dataset():
-    """Download and load entire dataset once at startup"""
     global dataset_images, dataset_masks, dataset_loaded
-    
     if dataset_loaded:
         return True
-        
-    print("Loading brain tumor dataset...")
-    
     try:
-        # Download dataset using kagglehub - returns directory path
+        print("Loading brain tumor dataset (kagglehub)...")
         dataset_path = kagglehub.dataset_download('nikhilroxtomar/brain-tumor-segmentation')
-        print(f"Dataset downloaded to: {dataset_path}")
-        
-        # Find images and masks directories
         images_dir = os.path.join(dataset_path, 'images')
         masks_dir = os.path.join(dataset_path, 'masks')
-        
-        # If direct path doesn't exist, search subdirectories
-        if not os.path.exists(images_dir):
-            # Search for images and masks directories
-            for root, dirs, files in os.walk(dataset_path):
-                if 'images' in dirs:
-                    images_dir = os.path.join(root, 'images')
-                if 'masks' in dirs:
-                    masks_dir = os.path.join(root, 'masks')
-        
         if not os.path.exists(images_dir) or not os.path.exists(masks_dir):
-            print("Could not find images/masks directories. Searching all files...")
-            # Fallback: find all image files
+            # fallback search
             all_files = glob(os.path.join(dataset_path, "**/*.png"), recursive=True) + \
-                       glob(os.path.join(dataset_path, "**/*.jpg"), recursive=True)
-            
+                        glob(os.path.join(dataset_path, "**/*.jpg"), recursive=True)
             dataset_images = [f for f in all_files if '/images/' in f or 'image' in f.lower()]
             dataset_masks = [f for f in all_files if '/masks/' in f or 'mask' in f.lower()]
         else:
-            # Load image and mask file paths
-            dataset_images = glob(os.path.join(images_dir, "*.*"))
-            dataset_masks = glob(os.path.join(masks_dir, "*.*"))
-        
-        dataset_images = sorted(dataset_images)
-        dataset_masks = sorted(dataset_masks)
-        
+            dataset_images = sorted(glob(os.path.join(images_dir, "*.*")))
+            dataset_masks = sorted(glob(os.path.join(masks_dir, "*.*")))
         print(f"✓ Found {len(dataset_images)} images and {len(dataset_masks)} masks")
         dataset_loaded = True
         return True
-        
     except Exception as e:
         print(f"Failed to load dataset: {e}")
         return False
 
 def get_random_sample():
-    """Get a random image and corresponding mask from dataset"""
     if not dataset_loaded:
         return None, None, "Dataset not loaded"
-    
     if not dataset_images:
-        return None, None, "No images found in dataset"
-    
-    # Get random index
-    idx = random.randint(0, len(dataset_images) - 1)
+        return None, None, "No images found"
+    idx = random.randint(0, len(dataset_images)-1)
     img_path = dataset_images[idx]
-    
-    # Find corresponding mask
     img_name = os.path.basename(img_path)
     mask_path = None
     for mask in dataset_masks:
         if os.path.basename(mask) == img_name:
             mask_path = mask
             break
-    
     try:
         image = Image.open(img_path).convert("L")
         mask = Image.open(mask_path).convert("L") if mask_path else None
@@ -225,394 +205,179 @@ def get_random_sample():
     except Exception as e:
         return None, None, f"Error loading sample: {e}"
 
+# -------------------------
+# Preprocessing & Heatmap utils
+# -------------------------
 def preprocess_for_model(image):
-    """Preprocessing for your model - matches the working notebook"""
     if image.mode != 'L':
         image = image.convert('L')
-    
     transform = transforms.Compose([
-        transforms.Resize((256,256)),
+        transforms.Resize((256, 256)),
         transforms.ToTensor()
     ])
-    
     return transform(image).unsqueeze(0)
 
 def generate_attention_heatmap(attention_maps):
-    """Generate attention heatmap"""
     if not attention_maps:
-        return np.zeros((256, 256, 3))
-    
-    # Resize all attention maps to the same size (256x256) before combining
+        return np.zeros((256, 256, 3), dtype=np.uint8)
     resized_maps = []
     target_size = (256, 256)
-    
     for att_map in attention_maps:
-        # Convert to numpy and squeeze
         att_np = att_map.squeeze().cpu().numpy()
-        
-        # Resize to target size
         att_resized = cv2.resize(att_np, target_size)
         resized_maps.append(att_resized)
-    
-    # Now we can safely average the maps since they're all the same size
     combined_att = np.mean(resized_maps, axis=0)
-    
-    # Normalize to [0, 1]
     combined_att = (combined_att - combined_att.min()) / (combined_att.max() - combined_att.min() + 1e-8)
-    
-    # Apply colormap
     heatmap = cv2.applyColorMap((combined_att * 255).astype(np.uint8), cv2.COLORMAP_JET)
-    
-    return heatmap
-
-def analyze_image(image, ground_truth, filename, debug=True):
-    """
-    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)
-    """
+    return heatmap  # BGR (OpenCV)
+
+# -------------------------
+# Core: produce combined 1x5 image (preserve old 1-4 behavior)
+# -------------------------
+def results_with_heatmap(image, ground_truth=None, filename=None, threshold=0.5):
     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("DEBUG: Starting analysis...")
-        print(f"Input image mode: {image.mode}")
-        print(f"Input image size: {image.size}")
-
-        # 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}")
-
-        # Forward pass and robust unpacking (support both return styles)
-        with torch.no_grad():
-            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:
-                pass
-
-        # 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 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_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
-
-        # Decide grid: show prob heatmap next to attention so you can compare
-        if ground_truth is not None:
-            fig, axes = plt.subplots(3, 4, figsize=(16, 12))  # add an extra row for debug heatmap
-        else:
-            fig, axes = plt.subplots(3, 3, figsize=(15, 12))
-
-        fig.suptitle('Brain Tumor Segmentation Analysis (debug)', fontsize=18, weight='bold')
-
-        # 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.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:
-            # 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)
-
-            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]
-            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')
-
-            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)
-
-            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 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')
-
-        # 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')
+    # Keep preprocessing & prediction exactly like your working code
+    img_gray = image.convert('L') if image.mode != 'L' else image
+    original_np = np.array(img_gray.resize((256, 256))).astype(np.uint8)
 
-        plt.tight_layout()
+    # Preprocess for model
+    prep = transforms.Compose([transforms.Resize((256, 256)), transforms.ToTensor()])
+    input_tensor = prep(img_gray).unsqueeze(0).to(device)
 
-        # 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).convert("RGB")
-
-        # 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
-
-        analysis_text = f"""
+    with torch.no_grad():
+        out = model(input_tensor)
+        # support both: model -> logits OR (logits, att_maps)
+        if isinstance(out, (list, tuple)) and len(out) == 2:
+            logits, attention_maps = out
+        else:
+            logits = out
+            attention_maps = []
+
+        pred_prob = torch.sigmoid(logits)
+        pred_mask = (pred_prob > threshold).float()
+
+    pred_mask_np = pred_mask.cpu().squeeze().numpy()  # (256,256)
+    inv_pred_mask_np = np.where(pred_mask_np == 1, 0, 255).astype(np.uint8)
+    tumor_only = np.where(pred_mask_np == 1, original_np, 255).astype(np.uint8)
+
+    # ground truth handling (preserve old style)
+    if ground_truth is not None:
+        gt_gray = ground_truth.convert('L') if ground_truth.mode != 'L' else ground_truth
+        mask_np = prep(gt_gray).cpu().squeeze().numpy()
+        mask_vis = (mask_np > 0.5).astype(np.uint8)
+    else:
+        mask_vis = np.zeros_like(original_np)
+
+    # Try to build attention heatmap; fallback to probability heatmap
+    att_heat = generate_attention_heatmap(attention_maps)
+    if att_heat is None or att_heat.size == 0:
+        prob_np = pred_prob.cpu().squeeze().numpy()
+        prob_resized = cv2.resize(prob_np, (256, 256))
+        prob_norm = (prob_resized - prob_resized.min()) / (prob_resized.max() - prob_resized.min() + 1e-8)
+        att_heat_bgr = cv2.applyColorMap((prob_norm * 255).astype(np.uint8), cv2.COLORMAP_JET)
+        att_heat = att_heat_bgr
+
+    # convert BGR->RGB for display
+    try:
+        att_heat = cv2.cvtColor(att_heat, cv2.COLOR_BGR2RGB)
+    except Exception:
+        pass
+
+    # ensure dtype/shape
+    if att_heat.dtype != np.uint8:
+        att_heat = (att_heat * 255).astype(np.uint8) if att_heat.max() <= 1.0 else att_heat.astype(np.uint8)
+    if att_heat.ndim == 2:
+        att_heat = cv2.cvtColor(att_heat, cv2.COLOR_GRAY2RGB)
+
+    # Create 1x5 figure
+    fig, axes = plt.subplots(1, 5, figsize=(22, 5))
+    fig.suptitle('Results + Heatmap', fontsize=16, weight='bold')
+
+    axes[0].imshow(original_np, cmap='gray'); axes[0].set_title('Original Image'); axes[0].axis('off')
+    axes[1].imshow(mask_vis, cmap='gray'); axes[1].set_title('Ground Truth Mask' if ground_truth is not None else 'GT (none)'); axes[1].axis('off')
+    axes[2].imshow(inv_pred_mask_np, cmap='gray'); axes[2].set_title('Predicted Mask'); axes[2].axis('off')
+    axes[3].imshow(tumor_only, cmap='gray'); axes[3].set_title('Tumor Only'); axes[3].axis('off')
+    axes[4].imshow(att_heat); axes[4].set_title('Attention / Prob Heatmap'); axes[4].axis('off')
+
+    plt.tight_layout()
+
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png', dpi=150, bbox_inches='tight', facecolor='white')
+    buf.seek(0)
+    plt.close(fig)
+    result_img = Image.open(buf).convert("RGB")
+
+    tumor_pixels = int(np.sum(pred_mask_np))
+    total_pixels = int(pred_mask_np.size)
+    tumor_pct = (tumor_pixels / total_pixels) * 100 if total_pixels > 0 else 0.0
+
+    analysis_text = f"""
 # Analysis Results
-
 **File:** {filename if filename else 'Uploaded Image'}
-
-**Tumor Detection:**
-- Tumor Area: {tumor_percentage:.2f}%
+- Tumor Area: {tumor_pct:.2f}%
 - Tumor Pixels: {tumor_pixels:,}
-
-**Model Features:**
-- Attention Visualization: Generated
-- Probability Heatmap: Generated
-"""
-
-        if ground_truth is not None:
-            analysis_text += f"""
-**Performance Metrics:**
-- IoU Score: {iou:.4f}
-- Dice Score: {dice:.4f}
+- Max confidence: {float(pred_prob.max()):.4f}
+- Threshold used: {threshold}
 """
 
-        # 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_img, analysis_text
 
-        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)
-        return None, error_msg
-
-
-# Initialize model and dataset at startup
+# -------------------------
+# Initialize model & dataset at startup
+# -------------------------
 print("Initializing application components...")
 model_loaded = download_and_load_model()
 dataset_loaded_success = download_and_load_dataset()
-
 if not model_loaded:
     print("WARNING: Model failed to load!")
 if not dataset_loaded_success:
     print("WARNING: Dataset failed to load!")
-
 print("Application ready!")
 
-# Professional CSS
+# -------------------------
+# Gradio UI
+# -------------------------
 css = """
-.gradio-container {
-    max-width: 1600px !important;
-    margin: auto !important;
-    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
-}
-.gr-button {
-    border-radius: 6px !important;
-    font-weight: 500 !important;
-}
-.gr-button-primary {
-    background: #2563eb !important;
-    border-color: #2563eb !important;
-}
-.gr-button-secondary {
-    background: #6b7280 !important;
-    border-color: #6b7280 !important;
-}
-h1, h2, h3 {
-    color: #1f2937 !important;
-}
-.gr-form {
-    border: 1px solid #e5e7eb !important;
-    border-radius: 8px !important;
-}
+.gradio-container { max-width: 1400px !important; margin:auto !important; font-family: 'Segoe UI', Tahoma, Verdana; }
+.gr-button { border-radius: 6px !important; font-weight: 500 !important; }
 """
 
-# Create Gradio interface
-with gr.Blocks(css=css, title="Brain Tumor Segmentation Analysis") as app:
-    
-    gr.Markdown("""
-    # Brain Tumor Segmentation Using Attention U-Net
-    
-    **Advanced Medical Image Analysis Tool**
-    
-    Features: Attention Visualization, Dataset Integration, Morphological Post-processing
-    """)
-    
-    # Status display
-    with gr.Row():
-        with gr.Column():
-            status_text = f"Model Status: {'✓ Loaded' if model_loaded else '✗ Failed'} | Dataset Status: {'✓ Loaded' if dataset_loaded_success else '✗ Failed'}"
-            if dataset_loaded_success:
-                status_text += f" | Images: {len(dataset_images)} | Masks: {len(dataset_masks)}"
-            gr.Markdown(f"**{status_text}**")
-    
+with gr.Blocks(css=css, title="Brain Tumor Segmentation + Heatmap") as app:
+    gr.Markdown("# Brain Tumor Segmentation — Attention U-Net\nPreserves original 1–4 outputs; adds 5th: heatmap.")
     with gr.Row():
         with gr.Column(scale=1):
-            gr.Markdown("### Input Selection")
-            
-            # Image display
-            image_display = gr.Image(
-                label="Selected Image",
-                type="pil",
-                height=300
-            )
-            
-            # Control buttons
+            image_display = gr.Image(label="Selected Image", type="pil", height=300)
             with gr.Row():
-                load_sample_btn = gr.Button("Load Random Sample", variant="primary", scale=1)
-                upload_btn = gr.UploadButton("Upload Image", file_types=["image"], scale=1)
-            
+                load_sample_btn = gr.Button("Load Random Sample", variant="primary")
+                upload_btn = gr.UploadButton("Upload Image", file_types=["image"])
             analyze_btn = gr.Button("Analyze Image", variant="primary", size="lg")
-            
-            # Dataset info
-            gr.Markdown(f"""
-            **Dataset Information:**
-            - Total Images: {len(dataset_images) if dataset_loaded_success else 'N/A'}
-            - Total Masks: {len(dataset_masks) if dataset_loaded_success else 'N/A'}
-            - Source: nikhilroxtomar/brain-tumor-segmentation
-            """)
-        
+            gr.Markdown(f"**Model Status:** {'✓ Loaded' if model_loaded else '✗ Failed'}  \n**Dataset:** {'✓ Loaded' if dataset_loaded_success else '✗ Failed'}")
         with gr.Column(scale=2):
-            gr.Markdown("### Analysis Results")
-            
-            result_display = gr.Image(
-                label="Segmentation Analysis",
-                type="pil",
-                height=500
-            )
-            
-            analysis_text = gr.Markdown(
-                value="Load an image and click 'Analyze Image' to begin."
-            )
-    
-    # Hidden states
+            gr.Markdown("### Results (1x5 panel)")
+            result_display = gr.Image(label="Segmentation + Heatmap", type="pil", height=600)
+            analysis_text = gr.Markdown("Upload or load a sample and click Analyze.")
+
     current_ground_truth = gr.State()
     current_filename = gr.State()
-    
-    # Event handlers
+
     def handle_sample_load():
         image, mask, filename = get_random_sample()
         return image, mask, filename
-    
-    def handle_upload(file):
-        if file is not None:
-            image = Image.open(file.name).convert("L")
-            return image, None, os.path.basename(file.name)
+
+    def handle_upload(f):
+        if f is not None:
+            img = Image.open(f.name).convert("L")
+            return img, None, os.path.basename(f.name)
         return None, None, ""
-    
-    load_sample_btn.click(
-        fn=handle_sample_load,
-        outputs=[image_display, current_ground_truth, current_filename]
-    )
-    
-    upload_btn.upload(
-        fn=handle_upload,
-        inputs=[upload_btn],
-        outputs=[image_display, current_ground_truth, current_filename]
-    )
-    
-    analyze_btn.click(
-        fn=analyze_image,
-        inputs=[image_display, current_ground_truth, current_filename],
-        outputs=[result_display, analysis_text]
-    )
+
+    load_sample_btn.click(fn=handle_sample_load, outputs=[image_display, current_ground_truth, current_filename])
+    upload_btn.upload(fn=handle_upload, inputs=[upload_btn], outputs=[image_display, current_ground_truth, current_filename])
+
+    analyze_btn.click(fn=results_with_heatmap, inputs=[image_display, current_ground_truth, current_filename], outputs=[result_display, analysis_text])
 
 if __name__ == "__main__":
-    print("\n" + "="*50)
-    print("LAUNCHING BRAIN TUMOR SEGMENTATION APPLICATION")
-    print("="*50)
-    
-    app.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True,
-        share=False
-    )
+    app.launch(server_name="0.0.0.0", server_port=7860, show_error=True, share=False)