Update app.py

app.py

@@ -10,10 +10,19 @@ from torchvision import transforms
 import torchvision.transforms.functional as TF
 import urllib.request
 import os
+import random
+import kagglehub
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = None
 
+# Download dataset
+dataset_path = kagglehub.dataset_download('nikhilroxtomar/brain-tumor-segmentation')
+image_path = os.path.join(dataset_path, 'images')
+mask_path = os.path.join(dataset_path, 'masks')
+test_imgs = sorted([f for f in os.listdir(image_path) if f.endswith('.jpg') or f.endswith('.png')])
+test_masks = sorted([f for f in os.listdir(mask_path) if f.endswith('.jpg') or f.endswith('.png')])
+
 # Define your Attention U-Net architecture (from your training code)
 class DoubleConv(nn.Module):
     def __init__(self, in_channels, out_channels):
@@ -65,12 +74,10 @@ class AttentionUNET(nn.Module):
         self.downs = nn.ModuleList()
         self.attentions = nn.ModuleList()
         self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
-        
         # Down part of UNET
         for feature in features:
             self.downs.append(DoubleConv(in_channels, feature))
             in_channels = feature
-        
         # Bottleneck
         self.bottleneck = DoubleConv(features[-1], features[-1]*2)
         
@@ -79,19 +86,15 @@ class AttentionUNET(nn.Module):
             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))
             self.ups.append(DoubleConv(feature*2, feature))
-        
         self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
-    
     def forward(self, x):
         skip_connections = []
         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] #reverse list
-        
         for idx in range(0, len(self.ups), 2):  #do up and double_conv
             x = self.ups[idx](x)
             skip_connection = skip_connections[idx//2]
@@ -100,7 +103,6 @@ class AttentionUNET(nn.Module):
             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)
 
 def download_model():
@@ -109,15 +111,15 @@ def download_model():
     model_path = "best_attention_model.pth.tar"
     
     if not os.path.exists(model_path):
-        print("📥 Downloading your trained model...")
+        print("Downloading your trained model...")
         try:
             urllib.request.urlretrieve(model_url, model_path)
-            print("✅ Model downloaded successfully!")
+            print("Model downloaded successfully!")
         except Exception as e:
-            print(f"❌ Failed to download model: {e}")
+            print(f"Failed to download model: {e}")
             return None
     else:
-        print("✅ Model already exists!")
+        print("Model already exists!")
     
     return model_path
 
@@ -126,7 +128,7 @@ def load_your_attention_model():
     global model
     if model is None:
         try:
-            print("🔄 Loading your trained Attention U-Net model...")
+            print("Loading your trained Attention U-Net model...")
             
             # Download model if needed
             model_path = download_model()
@@ -141,9 +143,9 @@ def load_your_attention_model():
             model.load_state_dict(checkpoint["state_dict"])
             model.eval()
             
-            print("✅ Your Attention U-Net model loaded successfully!")
+            print("Your Attention U-Net model loaded successfully!")
         except Exception as e:
-            print(f"❌ Error loading your model: {e}")
+            print(f"Error loading your model: {e}")
             model = None
     return model
 
@@ -161,75 +163,63 @@ def preprocess_for_your_model(image):
     
     return val_test_transform(image).unsqueeze(0)  # Add batch dimension
 
-def create_heatmap_visualization(pred_mask_continuous, original_image):
-    """Create heatmap visualization from continuous prediction values"""
-    # pred_mask_continuous should be the raw sigmoid output (0-1 values)
-    heatmap_np = pred_mask_continuous.cpu().squeeze().numpy()
-    
-    # Normalize to 0-255 for better visualization
-    heatmap_normalized = (heatmap_np * 255).astype(np.uint8)
-    
-    # Apply colormap (using 'hot' colormap like in medical imaging)
-    heatmap_colored = cv2.applyColorMap(heatmap_normalized, cv2.COLORMAP_HOT)
-    heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)
-    
-    # Convert original image to RGB for overlay
-    if len(original_image.shape) == 2:  # Grayscale
-        original_rgb = cv2.cvtColor(original_image.astype(np.uint8), cv2.COLOR_GRAY2RGB)
-    else:
-        original_rgb = original_image.astype(np.uint8)
-    
-    # Create overlay (blend original image with heatmap)
-    alpha = 0.6  # Transparency factor
-    overlay = cv2.addWeighted(original_rgb, 1-alpha, heatmap_colored, alpha, 0)
-    
-    return overlay
-
-def predict_tumor(image):
+def predict_tumor(image, mask=None):
     current_model = load_your_attention_model()
     
     if current_model is None:
-        return None, "❌ Failed to load your trained model."
+        return None, "Failed to load your trained model."
     if image is None:
-        return None, "⚠️ Please upload an image first."
+        return None, "Please upload an image first."
     
     try:
-        print("🧠 Processing with YOUR trained Attention U-Net...")
+        print("Processing with YOUR trained Attention U-Net...")
         
         # Use the exact preprocessing from your Colab code
         input_tensor = preprocess_for_your_model(image).to(device)
         
         # Predict using your model (exactly like your Colab code)
         with torch.no_grad():
-            pred_mask_continuous = torch.sigmoid(current_model(input_tensor))  # Keep continuous values for heatmap
-            pred_mask_binary = (pred_mask_continuous > 0.5).float()  # Binary mask for original visualizations
+            pred_mask = torch.sigmoid(current_model(input_tensor))
+            pred_mask_binary = (pred_mask > 0.5).float()
         
-        # Convert to numpy (like your Colab code) - KEEPING ORIGINAL LOGIC
+        # Convert to numpy (like your Colab code)
         pred_mask_np = pred_mask_binary.cpu().squeeze().numpy()
+        prob_mask_np = pred_mask.cpu().squeeze().numpy()  # Probability for heatmap
         original_np = np.array(image.convert('L').resize((256, 256)))
         
-        # Create inverted mask for visualization (like your Colab code) - UNCHANGED
+        # Create inverted mask for visualization (like your Colab code)
         inv_pred_mask_np = np.where(pred_mask_np == 1, 0, 255)
         
-        # Create tumor-only image (like your Colab code) - UNCHANGED
+        # Create tumor-only image (like your Colab code)
         tumor_only = np.where(pred_mask_np == 1, original_np, 255)
         
-        # NEW: Create heatmap visualization
-        heatmap_overlay = create_heatmap_visualization(pred_mask_continuous, original_np)
+        # Handle ground truth if provided
+        mask_np = None
+        dice_score = None
+        iou_score = None
+        if mask is not None:
+            mask_transform = transforms.Compose([
+                transforms.Resize((256,256)),
+                transforms.ToTensor()
+            ])
+            mask_tensor = mask_transform(mask).squeeze().numpy()
+            mask_np = (mask_tensor > 0.5).astype(float)
+            
+            intersection = np.logical_and(pred_mask_np, mask_np).sum()
+            union = np.logical_or(pred_mask_np, mask_np).sum()
+            iou_score = intersection / (union + 1e-7)
+            dice_score = (2 * intersection) / (pred_mask_np.sum() + mask_np.sum() + 1e-7)
         
-        # Create visualization with 5 panels (original 4 + heatmap)
+        # Create visualization (5-panel layout)
         fig, axes = plt.subplots(1, 5, figsize=(25, 5))
-        fig.suptitle('🧠 Your Attention U-Net Results with Heatmap', fontsize=16, fontweight='bold')
+        fig.suptitle('Your Attention U-Net Results', fontsize=16, fontweight='bold')
         
-        titles = ["Original Image", "Predicted Mask", "Inverted Mask", "Tumor Only", "Prediction Heatmap"]
-        images = [original_np, pred_mask_np * 255, inv_pred_mask_np, tumor_only, heatmap_overlay]
-        cmaps = ['gray', 'hot', 'gray', 'gray', None]  # None for RGB heatmap
+        titles = ["Original Image", "Ground Truth", "Predicted Mask", "Tumor Only", "Heatmap"]
+        images = [original_np, mask_np if mask_np is not None else np.zeros_like(original_np), inv_pred_mask_np, tumor_only, prob_mask_np]
+        cmaps = ['gray', 'gray', 'gray', 'gray', 'hot']
         
         for i, ax in enumerate(axes):
-            if cmaps[i] is not None:
-                ax.imshow(images[i], cmap=cmaps[i])
-            else:
-                ax.imshow(images[i])  # RGB image
+            ax.imshow(images[i], cmap=cmaps[i])
             ax.set_title(titles[i], fontsize=12, fontweight='bold')
             ax.axis('off')
         
@@ -243,68 +233,73 @@ def predict_tumor(image):
         
         result_image = Image.open(buf)
         
-        # Calculate statistics (like your Colab code) - UNCHANGED
+        # Calculate statistics (like your Colab code)
         tumor_pixels = np.sum(pred_mask_np)
         total_pixels = pred_mask_np.size
         tumor_percentage = (tumor_pixels / total_pixels) * 100
         
         # Calculate confidence metrics
-        max_confidence = torch.max(pred_mask_continuous).item()
-        mean_confidence = torch.mean(pred_mask_continuous).item()
+        max_confidence = torch.max(pred_mask).item()
+        mean_confidence = torch.mean(pred_mask).item()
         
         analysis_text = f"""
-## 🧠 Your Attention U-Net Analysis Results
-### 📊 Detection Summary:
-- **Status**: {'🔴 TUMOR DETECTED' if tumor_pixels > 50 else '🟢 NO SIGNIFICANT TUMOR'}
+## 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 brain region
 - **Tumor Pixels**: {tumor_pixels:,} pixels
 - **Max Confidence**: {max_confidence:.4f}
 - **Mean Confidence**: {mean_confidence:.4f}
-
-### 🔥 New Heatmap Features:
-- **Continuous Predictions**: Shows confidence levels (0-1)
-- **Color Coding**: Red/Yellow = High confidence, Blue/Black = Low confidence
-- **Overlay Visualization**: Heatmap overlaid on original image
-- **Enhanced Analysis**: Better understanding of model uncertainty
-
-### 🔬 Your Model Information:
+"""
+        if dice_score is not None and iou_score is not None:
+            analysis_text += f"""
+- **Dice Score**: {dice_score:.4f}
+- **IoU Score**: {iou_score:.4f}
+"""
+        analysis_text += """
+### Model Information:
 - **Architecture**: YOUR trained Attention U-Net
 - **Training Performance**: Dice: 0.8420, IoU: 0.7297 
 - **Input**: Grayscale (single channel)
-- **Output**: Binary segmentation mask + Continuous heatmap
+- **Output**: Binary segmentation mask
 - **Device**: {device.type.upper()}
-
-### 🎯 Model Performance:
+### Model Performance:
 - **Training Accuracy**: 98.90%
 - **Best Dice Score**: 0.8420
 - **Best IoU Score**: 0.7297
 - **Training Dataset**: Brain tumor segmentation dataset
-
-### 📈 Processing Details:
+### Processing Details:
 - **Preprocessing**: Resize(256×256) + ToTensor (your exact method)
 - **Threshold**: 0.5 (sigmoid > 0.5)
 - **Architecture**: Attention gates + Skip connections
 - **Features**: [32, 64, 128, 256] channels
-- **Heatmap**: Continuous sigmoid output with hot colormap
-
-### ⚠️ Medical Disclaimer:
+### Medical Disclaimer:
 This is YOUR trained AI model for **research and educational purposes only**. 
 Results should be validated by medical professionals. Not for clinical diagnosis.
-
-### 🏆 Model Quality:
-✅ This is your own trained model with proven {tumor_percentage:.2f}% detection capability!
+### Model Quality:
+This is your own trained model with {tumor_percentage:.2f}% detection capability!
         """
         
-        print(f"✅ Your model analysis completed! Tumor area: {tumor_percentage:.2f}%")
+        print(f"Your model analysis completed! Tumor area: {tumor_percentage:.2f}%")
         return result_image, analysis_text
         
     except Exception as e:
-        error_msg = f"❌ Error with your model: {str(e)}"
+        error_msg = f"Error with your model: {str(e)}"
         print(error_msg)
         return None, error_msg
 
+def load_random_sample():
+    if not test_imgs:
+        return None, None, "Dataset not available."
+    rand_idx = random.randint(0, len(test_imgs) - 1)
+    img_path = os.path.join(image_path, test_imgs[rand_idx])
+    msk_path = os.path.join(mask_path, test_masks[rand_idx])  # Assuming paired by index
+    image = Image.open(img_path).convert('L')
+    mask = Image.open(msk_path).convert('L')
+    return image, mask, "Loaded random sample from dataset."
+
 def clear_all():
-    return None, None, "Upload a brain MRI image to test YOUR trained Attention U-Net model with heatmap visualization"
+    return None, None, "Upload a brain MRI image to test YOUR trained Attention U-Net model", None
 
 # Enhanced CSS for your model
 css = """
@@ -324,13 +319,13 @@ css = """
 """
 
 # Create Gradio interface for your model
-with gr.Blocks(css=css, title="🧠 Your Attention U-Net Model with Heatmap", theme=gr.themes.Soft()) as app:
+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 with Heatmap</h1>
+        <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 • Now with Heatmap Visualization
+            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
@@ -338,9 +333,11 @@ with gr.Blocks(css=css, title="🧠 Your Attention U-Net Model with Heatmap", th
     </div>
     """)
     
+    mask_state = gr.State(None)
+    
     with gr.Row():
         with gr.Column(scale=1):
-            gr.Markdown("### 📤 Upload Brain MRI")
+            gr.Markdown("### Upload Brain MRI")
             
             image_input = gr.Image(
                 label="Brain MRI Scan",
@@ -350,51 +347,49 @@ with gr.Blocks(css=css, title="🧠 Your Attention U-Net Model with Heatmap", th
             )
             
             with gr.Row():
-                analyze_btn = gr.Button("🔍 Analyze with YOUR Model", variant="primary", scale=2, size="lg")
-                clear_btn = gr.Button("🗑️ Clear", variant="secondary", scale=1)
+                analyze_btn = gr.Button("Analyze with YOUR Model", variant="primary", scale=1, size="lg")
+                random_btn = gr.Button("Load Random Sample", variant="secondary", scale=1, 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>
+                <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>NEW: Heatmap:</strong> Continuous confidence visualization</li>
-                    <li><strong>5-Panel View:</strong> Complete analysis with heatmap</li>
+                    <li><strong>5-Panel View:</strong> Complete analysis like your Colab</li>
                 </ul>
             </div>
             """)
         
         with gr.Column(scale=2):
-            gr.Markdown("### 📊 Your Model Results with Heatmap")
+            gr.Markdown("### Your Model Results")
             
             output_image = gr.Image(
-                label="Your Attention U-Net Analysis with Heatmap",
+                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 with heatmap visualization.",
+                value="Upload a brain MRI image to test YOUR trained Attention U-Net model.",
                 elem_id="analysis"
             )
-    
-    # Footer highlighting your model with heatmap features
+    # Footer highlighting your model
     gr.HTML("""
     <div style="margin-top: 30px; padding: 25px; background-color: #F8FAFC; border-radius: 15px; border: 2px solid #8B5CF6;">
         <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 30px;">
             <div>
-                <h4 style="color: #8B5CF6; margin-bottom: 15px;">🏆 Your Personal AI Model</h4>
+                <h4 style="color: #8B5CF6; margin-bottom: 15px;">Your Personal AI Model</h4>
                 <p><strong>Architecture:</strong> Attention U-Net with skip connections</p>
                 <p><strong>Performance:</strong> Dice: 0.8420, IoU: 0.7297, Accuracy: 98.90%</p>
                 <p><strong>Training:</strong> Your own dataset-specific training</p>
                 <p><strong>Features:</strong> [32, 64, 128, 256] channel progression</p>
-                <p><strong>NEW:</strong> Continuous heatmap visualization for confidence</p>
             </div>
             <div>
-                <h4 style="color: #DC2626; margin-bottom: 15px;">⚠️ Your Model Disclaimer</h4>
+                <h4 style="color: #DC2626; margin-bottom: 15px;">Your Model Disclaimer</h4>
                 <p style="color: #DC2626; font-weight: 600; line-height: 1.4;">
                     This is YOUR personally trained AI model for <strong>research purposes only</strong>.<br>
                     Results reflect your model's training performance.<br>
@@ -404,7 +399,7 @@ with gr.Blocks(css=css, title="🧠 Your Attention U-Net Model with Heatmap", th
         </div>
         <hr style="margin: 20px 0; border: none; border-top: 2px solid #E5E7EB;">
         <p style="text-align: center; color: #6B7280; margin: 10px 0; font-weight: 600;">
-            🚀 Your Personal Attention U-Net • Downloaded from HuggingFace • Research-Grade Performance • Now with Heatmap! 🔥
+            Your Personal Attention U-Net • Downloaded from HuggingFace • Research-Grade Performance
         </p>
     </div>
     """)
@@ -412,27 +407,32 @@ with gr.Blocks(css=css, title="🧠 Your Attention U-Net Model with Heatmap", th
     # Event handlers
     analyze_btn.click(
         fn=predict_tumor,
-        inputs=[image_input],
+        inputs=[image_input, mask_state],
         outputs=[output_image, analysis_output],
         show_progress=True
     )
     
+    random_btn.click(
+        fn=load_random_sample,
+        inputs=[],
+        outputs=[image_input, mask_state, analysis_output]
+    )
+    
     clear_btn.click(
         fn=clear_all,
         inputs=[],
-        outputs=[image_input, output_image, analysis_output]
+        outputs=[image_input, output_image, analysis_output, mask_state]
     )
 
 if __name__ == "__main__":
-    print("🚀 Starting YOUR Attention U-Net Model System with Heatmap...")
-    print("🏆 Using your personally trained model")
-    print("📥 Auto-downloading from HuggingFace...")
-    print("🎯 Expected performance: Dice 0.8420, IoU 0.7297")
-    print("🔥 NEW: Heatmap visualization added!")
+    print("Starting YOUR Attention U-Net Model System...")
+    print("Using your personally trained model")
+    print("Auto-downloading from HuggingFace...")
+    print("Expected performance: Dice 0.8420, IoU 0.7297")
     
     app.launch(
         server_name="0.0.0.0",
         server_port=7860,
         show_error=True,
         share=False
-    )
+    )