Update app.py

app.py

@@ -1,4 +1,3 @@
-# full_app_with_heatmap.py
 import gradio as gr
 import torch
 import torch.nn as nn
@@ -11,22 +10,11 @@ from torchvision import transforms
 import torchvision.transforms.functional as TF
 import urllib.request
 import os
-import random
-from glob import glob
-import kagglehub  # if you use dataset download in the app; remove if not needed
 
-# -------------------------
-# Setup / Globals
-# -------------------------
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = None
-dataset_images = []
-dataset_masks = []
-dataset_loaded = False
 
-# -------------------------
-# Model classes (Attention U-Net)
-# -------------------------
+# ---- model classes (kept equivalent to your working code) ----
 class DoubleConv(nn.Module):
     def __init__(self, in_channels, out_channels):
         super(DoubleConv, self).__init__()
@@ -38,11 +26,9 @@ class DoubleConv(nn.Module):
             nn.BatchNorm2d(out_channels),
             nn.ReLU(inplace=True),
         )
-
     def forward(self, x):
         return self.conv(x)
 
-
 class AttentionBlock(nn.Module):
     def __init__(self, F_g, F_l, F_int):
         super(AttentionBlock, self).__init__()
@@ -60,14 +46,12 @@ class AttentionBlock(nn.Module):
             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 attended skip, attention map
-
+        return x * psi  # matches the old working code
 
 class AttentionUNET(nn.Module):
     def __init__(self, in_channels=1, out_channels=1, features=[32, 64, 128, 256]):
@@ -78,12 +62,15 @@ class AttentionUNET(nn.Module):
         self.attentions = nn.ModuleList()
         self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
 
+        # down
         for feature in features:
             self.downs.append(DoubleConv(in_channels, feature))
             in_channels = feature
 
+        # bottleneck
         self.bottleneck = DoubleConv(features[-1], features[-1]*2)
 
+        # up
         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))
@@ -93,291 +80,283 @@ class AttentionUNET(nn.Module):
 
     def forward(self, x):
         skip_connections = []
-        attention_maps = []
-
         for down in self.downs:
             x = down(x)
             skip_connections.append(x)
             x = self.pool(x)
-
         x = self.bottleneck(x)
-        skip_connections = skip_connections[::-1]
+        skip_connections = skip_connections[::-1]  # reverse
 
         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:])
-
-            attended_skip, att_map = self.attentions[idx // 2](x, skip_connection)
-            attention_maps.append(att_map)
-            concat_skip = torch.cat((attended_skip, x), dim=1)
+            # attention applied exactly as in your working code
+            skip_connection = self.attentions[idx // 2](skip_connection, x)
+            concat_skip = torch.cat((skip_connection, x), dim=1)
             x = self.ups[idx+1](concat_skip)
+        return self.final_conv(x)
 
-        return self.final_conv(x), attention_maps
-
-# -------------------------
-# Model download / load
-# -------------------------
-def download_and_load_model():
-    global model
-    print("Loading Attention U-Net model...")
-
+# ---- model download/load helpers (same as yours) ----
+def download_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"
-
     if not os.path.exists(model_path):
-        print("Downloading model weights...")
+        print("📥 Downloading your trained model...")
         try:
             urllib.request.urlretrieve(model_url, model_path)
+            print("✅ Model downloaded successfully!")
         except Exception as e:
-            print(f"Failed to download model: {e}")
-            return False
+            print(f"❌ Failed to download model: {e}")
+            return None
+    else:
+        print("✅ Model already exists!")
+    return model_path
 
-    try:
-        model = AttentionUNET(in_channels=1, out_channels=1).to(device)
-        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
+def load_your_attention_model():
+    global model
+    if model is None:
         try:
+            print("🔄 Loading your trained Attention U-Net model...")
+            model_path = download_model()
+            if model_path is None:
+                return None
+            model = AttentionUNET(in_channels=1, out_channels=1).to(device)
+            checkpoint = torch.load(model_path, map_location=device)
+            # checkpoint expected to have "state_dict" key as in your working code
+            if "state_dict" in checkpoint:
+                sd = checkpoint["state_dict"]
+            else:
+                sd = checkpoint
             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():
-    global dataset_images, dataset_masks, dataset_loaded
-    if dataset_loaded:
-        return True
-    try:
-        print("Loading brain tumor dataset (kagglehub)...")
-        dataset_path = kagglehub.dataset_download('nikhilroxtomar/brain-tumor-segmentation')
-        images_dir = os.path.join(dataset_path, 'images')
-        masks_dir = os.path.join(dataset_path, 'masks')
-        if not os.path.exists(images_dir) or not os.path.exists(masks_dir):
-            # fallback search
-            all_files = glob(os.path.join(dataset_path, "**/*.png"), 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:
-            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():
-    if not dataset_loaded:
-        return None, None, "Dataset not loaded"
-    if not dataset_images:
-        return None, None, "No images found"
-    idx = random.randint(0, len(dataset_images)-1)
-    img_path = dataset_images[idx]
-    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
-        return image, mask, img_name
-    except Exception as e:
-        return None, None, f"Error loading sample: {e}"
+            model.eval()
+            print("✅ Your Attention U-Net model loaded successfully!")
+        except Exception as e:
+            print(f"❌ Error loading your model: {e}")
+            model = None
+    return model
 
-# -------------------------
-# Preprocessing & Heatmap utils
-# -------------------------
-def preprocess_for_model(image):
+# ---- preprocessing (same as your Colab code) ----
+def preprocess_for_your_model(image):
     if image.mode != 'L':
         image = image.convert('L')
-    transform = transforms.Compose([
-        transforms.Resize((256, 256)),
+    val_test_transform = transforms.Compose([
+        transforms.Resize((256,256)),
         transforms.ToTensor()
     ])
-    return transform(image).unsqueeze(0)
-
-def generate_attention_heatmap(attention_maps):
-    if not attention_maps:
-        return np.zeros((256, 256, 3), dtype=np.uint8)
-    resized_maps = []
-    target_size = (256, 256)
-    for att_map in attention_maps:
-        att_np = att_map.squeeze().cpu().numpy()
-        att_resized = cv2.resize(att_np, target_size)
-        resized_maps.append(att_resized)
-    combined_att = np.mean(resized_maps, axis=0)
-    combined_att = (combined_att - combined_att.min()) / (combined_att.max() - combined_att.min() + 1e-8)
-    heatmap = cv2.applyColorMap((combined_att * 255).astype(np.uint8), cv2.COLORMAP_JET)
-    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."
+    return val_test_transform(image).unsqueeze(0)  # Add batch dimension
+
+# ---- main predict function (modified to add separate heatmap, no change to 1-4) ----
+def predict_tumor(image):
+    """
+    Keeps the exact old 4-panel outputs the same, and adds a 5th panel with
+    the probability heatmap. The heatmap is computed from the sigmoid(probabilities)
+    and does not change any tensors used for predictions.
+    """
+    current_model = load_your_attention_model()
+    if current_model is None:
+        return None, "❌ Failed to load your trained model."
     if image is None:
-        return None, "Please select an image first."
-
-    # 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)
+        return None, "⚠️ Please upload an image first."
 
-    # Preprocess for model
-    prep = transforms.Compose([transforms.Resize((256, 256)), transforms.ToTensor()])
-    input_tensor = prep(img_gray).unsqueeze(0).to(device)
-
-    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
+    try:
+        print("🧠 Processing with YOUR trained Attention U-Net...")
+
+        # Preprocess exactly like your Colab
+        input_tensor = preprocess_for_your_model(image).to(device)  # [1,1,256,256]
+
+        # Forward and prediction (identical to your working code)
+        with torch.no_grad():
+            logits = current_model(input_tensor)  # model returns logits tensor
+            pred_prob = torch.sigmoid(logits)     # keep prob map for heatmap
+            pred_mask = (pred_prob > 0.5).float() # binary mask (same threshold as old code)
+
+        # Convert to numpy like old code
+        pred_mask_np = pred_mask.cpu().squeeze().numpy()          # shape (256,256)
+        original_np = np.array(image.convert('L').resize((256, 256)))
+        inv_pred_mask_np = np.where(pred_mask_np == 1, 0, 255)
+        tumor_only = np.where(pred_mask_np == 1, original_np, 255)
+
+        # -------------------------
+        # Create heatmap (NO CHANGES to pred_mask or any prediction tensors)
+        # -------------------------
+        # Use the probability map (float) as the basis
+        pred_prob_np = pred_prob.cpu().squeeze().numpy()  # float in [0,1]
+        # ensure same shape 256x256
+        if pred_prob_np.shape != (256, 256):
+            pred_prob_resized = cv2.resize(pred_prob_np, (256, 256))
         else:
-            logits = out
-            attention_maps = []
-
-        pred_prob = torch.sigmoid(logits)
-        pred_mask = (pred_prob > threshold).float()
+            pred_prob_resized = pred_prob_np.copy()
+
+        # Normalize to 0-1 and convert to uint8 for colormap
+        prob_norm = (pred_prob_resized - pred_prob_resized.min()) / (pred_prob_resized.max() - pred_prob_resized.min() + 1e-8)
+        prob_uint8 = (prob_norm * 255).astype(np.uint8)
+        prob_heatmap_bgr = cv2.applyColorMap(prob_uint8, cv2.COLORMAP_JET)  # OpenCV BGR
+        # Convert BGR -> RGB for matplotlib/PIL visualization
+        prob_heatmap_rgb = cv2.cvtColor(prob_heatmap_bgr, cv2.COLOR_BGR2RGB)
+
+        # -------------------------
+        # Build the 5-panel figure
+        # Panels (left->right): Original | Pred segmentation (pred*255) | Inverted mask | Tumor only | Heatmap
+        # Panels 1-4 are produced exactly the same as your old code
+        # -------------------------
+        fig, axes = plt.subplots(1, 5, figsize=(24, 5))
+        fig.suptitle('🧠 Your Attention U-Net Results (with Heatmap)', fontsize=16, fontweight='bold')
+
+        # 1 Original (gray)
+        axes[0].imshow(original_np, cmap='gray')
+        axes[0].set_title("Original Image", fontsize=12, fontweight='bold')
+        axes[0].axis('off')
+
+        # 2 Tumor Segmentation (pred*255) — identical to old code's second panel
+        axes[1].imshow(pred_mask_np * 255, cmap='hot')
+        axes[1].set_title("Tumor Segmentation (pred × 255)", fontsize=12, fontweight='bold')
+        axes[1].axis('off')
+
+        # 3 Inverted mask — identical
+        axes[2].imshow(inv_pred_mask_np, cmap='gray')
+        axes[2].set_title("Inverted Mask (visual)", fontsize=12, fontweight='bold')
+        axes[2].axis('off')
+
+        # 4 Tumor only (grayscale crop) — identical
+        axes[3].imshow(tumor_only, cmap='gray')
+        axes[3].set_title("Tumor Only", fontsize=12, fontweight='bold')
+        axes[3].axis('off')
+
+        # 5 Heatmap (RGB)
+        axes[4].imshow(prob_heatmap_rgb)
+        axes[4].set_title("Probability Heatmap (sigmoid)", fontsize=12, fontweight='bold')
+        axes[4].axis('off')
+
+        plt.tight_layout()
+
+        # Save plot
+        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)
+
+        # Calculate statistics (like your Colab code)
+        tumor_pixels = int(np.sum(pred_mask_np))
+        total_pixels = int(pred_mask_np.size)
+        tumor_percentage = (tumor_pixels / total_pixels) * 100 if total_pixels > 0 else 0.0
+
+        # Confidence metrics (from the probability tensor)
+        max_confidence = float(pred_prob.max().item())
+        mean_confidence = float(pred_prob.mean().item())
+
+        analysis_text = f"""
+## 🧠 Your Attention U-Net Analysis Results
+### 📊 Detection Summary:
+- **Status**: {'🔴 TUMOR DETECTED' if tumor_pixels > 50 else '🟢 NO SIGNIFICANT TUMOR'}
+- **Tumor Area**: {tumor_percentage:.2f}% of image
+- **Tumor Pixels**: {tumor_pixels:,} pixels
+- **Max Confidence**: {max_confidence:.4f}
+- **Mean Confidence**: {mean_confidence:.4f}
+
+### 🔬 Model Info:
+- **Architecture**: YOUR Attention U-Net
+- **Input**: Grayscale (single channel), resized to 256×256
+- **Threshold**: 0.5 (sigmoid > 0.5)
+- **Device**: {device.type.upper()}
+
+### ⚠️ Disclaimer:
+This is for research/education only. Validate results with medical professionals.
+"""
+        print(f"✅ Your model analysis completed! Tumor area: {tumor_percentage:.2f}%")
+        return result_image, analysis_text
 
-    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)
+    except Exception as e:
+        error_msg = f"❌ Error with your model: {str(e)}"
+        print(error_msg)
+        return None, error_msg
 
-    # 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 Area: {tumor_pct:.2f}%
-- Tumor Pixels: {tumor_pixels:,}
-- Max confidence: {float(pred_prob.max()):.4f}
-- Threshold used: {threshold}
-"""
+def clear_all():
+    return None, "Upload a brain MRI image to test YOUR trained Attention U-Net model"
 
-    return result_img, analysis_text
-
-# -------------------------
-# 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!")
-
-# -------------------------
-# Gradio UI
-# -------------------------
+# ---- Gradio UI (kept as you had it, but wired to the new predict function) ----
 css = """
-.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; }
+.gradio-container {
+    max-width: 1400px !important;
+    margin: auto !important;
+}
+#title {
+    text-align: center;
+    background: linear-gradient(135deg, #8B5CF6 0%, #7C3AED 100%);
+    color: white;
+    padding: 30px;
+    border-radius: 15px;
+    margin-bottom: 25px;
+    box-shadow: 0 8px 16px rgba(139, 92, 246, 0.3);
+}
 """
 
-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.Blocks(css=css, title="🧠 Your Attention U-Net Model", theme=gr.themes.Soft()) as app:
+    gr.HTML("""
+    <div id="title">
+        <h1>🧠 YOUR Attention U-Net Model</h1>
+        <p style="font-size: 18px; margin-top: 15px;">
+            Using Your Own Trained Model • Dice: 0.8420 • IoU: 0.7297
+        </p>
+        <p style="font-size: 14px; margin-top: 10px; opacity: 0.9;">
+            Loaded from: ArchCoder/the-op-segmenter HuggingFace Space
+        </p>
+    </div>
+    """)
+
     with gr.Row():
         with gr.Column(scale=1):
-            image_display = gr.Image(label="Selected Image", type="pil", height=300)
+            gr.Markdown("### 📤 Upload Brain MRI")
+            image_input = gr.Image(
+                label="Brain MRI Scan",
+                type="pil",
+                sources=["upload", "webcam"],
+                height=350
+            )
+
             with gr.Row():
-                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")
-            gr.Markdown(f"**Model Status:** {'✓ Loaded' if model_loaded else '✗ Failed'}  \n**Dataset:** {'✓ Loaded' if dataset_loaded_success else '✗ Failed'}")
+                analyze_btn = gr.Button("🔍 Analyze with YOUR Model", variant="primary", scale=2, size="lg")
+                clear_btn = gr.Button("🗑️ Clear", variant="secondary", scale=1)
+
+            gr.HTML("""
+            <div style="margin-top: 20px; padding: 20px; background: linear-gradient(135deg, #F3E8FF 0%, #EDE9FE 100%); border-radius: 10px; border-left: 4px solid #8B5CF6;">
+                <h4 style="color: #8B5CF6; margin-bottom: 15px;">🏆 Your Model Features:</h4>
+                <ul style="margin: 10px 0; padding-left: 20px; line-height: 1.6;">
+                    <li><strong>Personal Model:</strong> Your own trained Attention U-Net</li>
+                    <li><strong>Proven Performance:</strong> 84.2% Dice Score, 72.97% IoU</li>
+                    <li><strong>Attention Gates:</strong> Advanced feature selection</li>
+                    <li><strong>Clean Output:</strong> Binary segmentation masks</li>
+                    <li><strong>5-Panel View:</strong> Original, Segmentation, Inverted, Tumor-only, Heatmap</li>
+                </ul>
+            </div>
+            """)
         with gr.Column(scale=2):
-            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()
-
-    def handle_sample_load():
-        image, mask, filename = get_random_sample()
-        return image, mask, filename
-
-    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=results_with_heatmap, inputs=[image_display, current_ground_truth, current_filename], outputs=[result_display, analysis_text])
+            gr.Markdown("### 📊 Your Model Results")
+            output_image = gr.Image(
+                label="Your Attention U-Net Analysis",
+                type="pil",
+                height=500
+            )
+            analysis_output = gr.Markdown(
+                value="Upload a brain MRI image to test YOUR trained Attention U-Net model.",
+                elem_id="analysis"
+            )
+
+    analyze_btn.click(
+        fn=predict_tumor,
+        inputs=[image_input],
+        outputs=[output_image, analysis_output],
+        show_progress=True
+    )
+
+    clear_btn.click(
+        fn=clear_all,
+        inputs=[],
+        outputs=[image_input, analysis_output]
+    )
 
 if __name__ == "__main__":
+    print("🚀 Starting YOUR Attention U-Net Model System...")
     app.launch(server_name="0.0.0.0", server_port=7860, show_error=True, share=False)