Update app.py

app.py

@@ -6,15 +6,18 @@ import cv2
 from PIL import Image
 import matplotlib.pyplot as plt
 import io
-from torchvision import transforms
+import torchvision.transforms as transforms
 import torchvision.transforms.functional as TF
-import urllib.request
+import random
 import os
+import zipfile
+import urllib.request
+import kagglehub
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = None
 
-# Define your Attention U-Net architecture (from your training code)
+# Your Attention U-Net classes (from your code)
 class DoubleConv(nn.Module):
     def __init__(self, in_channels, out_channels):
         super(DoubleConv, self).__init__()
@@ -56,7 +59,7 @@ class AttentionBlock(nn.Module):
         x1 = self.W_x(x)
         psi = self.relu(g1 + x1)
         psi = self.psi(psi)
-        return x * psi
+        return x * psi, psi  # Return attention map as well
 
 class AttentionUNET(nn.Module):
     def __init__(self, in_channels=1, out_channels=1, features=[32, 64, 128, 256]):
@@ -67,7 +70,7 @@ class AttentionUNET(nn.Module):
         self.attentions = nn.ModuleList()
         self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
 
-        # Down part of UNET
+        # Down part
         for feature in features:
             self.downs.append(DoubleConv(in_channels, feature))
             in_channels = feature
@@ -75,7 +78,7 @@ class AttentionUNET(nn.Module):
         # Bottleneck
         self.bottleneck = DoubleConv(features[-1], features[-1]*2)
         
-        # Up part of UNET
+        # Up part
         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))
@@ -85,6 +88,7 @@ class AttentionUNET(nn.Module):
 
     def forward(self, x):
         skip_connections = []
+        attention_maps = []  # To store attention maps
 
         for down in self.downs:
             x = down(x)
@@ -92,20 +96,21 @@ class AttentionUNET(nn.Module):
             x = self.pool(x)
 
         x = self.bottleneck(x)
-        skip_connections = skip_connections[::-1] #reverse list
+        skip_connections = skip_connections[::-1]
 
-        for idx in range(0, len(self.ups), 2):  #do up and double_conv
+        for idx in range(0, len(self.ups), 2):
             x = self.ups[idx](x)
             skip_connection = skip_connections[idx//2]
 
             if x.shape != skip_connection.shape:
                 x = TF.resize(x, size=skip_connection.shape[2:])
 
-            skip_connection = self.attentions[idx // 2](skip_connection, x)
-            concat_skip = torch.cat((skip_connection, x), dim=1)
+            attended_skip, att_map = self.attentions[idx // 2](x, skip_connection)  # Get attention map
+            attention_maps.append(att_map)  # Store attention map
+            concat_skip = torch.cat((attended_skip, x), dim=1)
             x = self.ups[idx+1](concat_skip)
 
-        return self.final_conv(x)
+        return self.final_conv(x), attention_maps
 
 def download_model():
     """Download your trained model from HuggingFace"""
@@ -120,9 +125,6 @@ def download_model():
         except Exception as e:
             print(f"❌ Failed to download model: {e}")
             return None
-    else:
-        print("✅ Model already exists!")
-    
     return model_path
 
 def load_your_attention_model():
@@ -151,66 +153,174 @@ def load_your_attention_model():
             model = None
     return model
 
+def download_dataset():
+    """Download and extract the dataset using kagglehub"""
+    dataset_path = kagglehub.dataset_download('nikhilroxtomar/brain-tumor-segmentation')
+    
+    # Extract if it's a zip
+    extracted_path = "brain_tumor_dataset"
+    if not os.path.exists(extracted_path):
+        with zipfile.ZipFile(dataset_path, 'r') as zip_ref:
+            zip_ref.extractall(extracted_path)
+    
+    images_path = os.path.join(extracted_path, 'images')
+    masks_path = os.path.join(extracted_path, 'masks')
+    
+    return images_path, masks_path
+
+def load_random_sample():
+    """Load a random image and mask from the dataset"""
+    images_path, masks_path = download_dataset()
+    
+    image_files = [f for f in os.listdir(images_path) if f.endswith(('.png', '.jpg'))]
+    if not image_files:
+        return None, None, "No images found in dataset"
+    
+    random_file = random.choice(image_files)
+    img_path = os.path.join(images_path, random_file)
+    mask_path = os.path.join(masks_path, random_file)
+    
+    image = Image.open(img_path).convert("L")
+    mask = Image.open(mask_path).convert("L") if os.path.exists(mask_path) else None
+    
+    return image, mask, random_file
+
 def preprocess_for_your_model(image):
     """Preprocessing exactly like your Colab code"""
-    # Convert to grayscale (like your Colab code)
     if image.mode != 'L':
         image = image.convert('L')
     
-    # Use the exact same transform as your Colab code
     val_test_transform = transforms.Compose([
         transforms.Resize((256,256)),
         transforms.ToTensor()
     ])
     
-    return val_test_transform(image).unsqueeze(0)  # Add batch dimension
+    return val_test_transform(image).unsqueeze(0)
+
+def apply_tta(model, input_tensor):
+    """Test-Time Augmentation: Apply augmentations and average predictions"""
+    augmentations = [
+        lambda x: x,  # Original
+        lambda x: TF.rotate(x, 90),  # 90 deg rotation
+        lambda x: TF.rotate(x, -90),  # -90 deg rotation
+        lambda x: TF.hflip(x),  # Horizontal flip
+        lambda x: TF.vflip(x)   # Vertical flip
+    ]
+    
+    predictions = []
+    for aug in augmentations:
+        aug_input = aug(input_tensor)
+        pred = torch.sigmoid(model(aug_input)[0])  # Get prediction
+        # Reverse the augmentation for averaging
+        if aug == augmentations[1]:  # Reverse 90 deg
+            pred = TF.rotate(pred, -90)
+        elif aug == augmentations[2]:  # Reverse -90 deg
+            pred = TF.rotate(pred, 90)
+        elif aug == augmentations[3]:  # Reverse hflip
+            pred = TF.hflip(pred)
+        elif aug == augmentations[4]:  # Reverse vflip
+            pred = TF.vflip(pred)
+        predictions.append(pred)
+    
+    # Average predictions
+    avg_pred = torch.mean(torch.stack(predictions), dim=0)
+    return avg_pred
+
+def generate_attention_heatmap(attention_maps):
+    """Generate combined attention heatmap"""
+    if not attention_maps:
+        return np.zeros((256, 256))
+    
+    # Average attention maps from different levels
+    combined_att = torch.mean(torch.stack(attention_maps), dim=0).squeeze().cpu().numpy()
+    combined_att = cv2.resize(combined_att, (256, 256))
+    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
 
-def predict_tumor(image):
+def predict_tumor(image, ground_truth=None, filename=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 or load an image first."
     
     try:
-        print("🧠 Processing with YOUR trained Attention U-Net...")
-        
-        # Use the exact preprocessing from your Colab code
+        # Preprocess
         input_tensor = preprocess_for_your_model(image).to(device)
         
-        # Predict using your model (exactly like your Colab code)
-        with torch.no_grad():
-            pred_mask = torch.sigmoid(current_model(input_tensor))
-            pred_mask_binary = (pred_mask > 0.5).float()
+        # Apply TTA
+        avg_pred = apply_tta(current_model, input_tensor)
+        
+        # Get binary mask
+        binary_mask = (avg_pred > 0.5).float().squeeze().cpu().numpy()
         
-        # Convert to numpy (like your Colab code)
-        pred_mask_np = pred_mask_binary.cpu().squeeze().numpy()
-        original_np = np.array(image.convert('L').resize((256, 256)))
+        # Post-processing
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5,5))
+        binary_mask = cv2.morphologyEx(binary_mask.astype(np.uint8), cv2.MORPH_OPEN, kernel)
+        binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel)
         
-        # Create inverted mask for visualization (like your Colab code)
-        inv_pred_mask_np = np.where(pred_mask_np == 1, 0, 255)
+        # Extract attention maps
+        _, attention_maps = current_model(input_tensor)
+        att_heatmap = generate_attention_heatmap(attention_maps)
         
-        # Create tumor-only image (like your Colab code)
-        tumor_only = np.where(pred_mask_np == 1, original_np, 255)
+        # Create visualization
+        fig, axes = plt.subplots(2, 3, figsize=(18, 12))
+        fig.suptitle('Brain Tumor Segmentation Analysis', fontsize=20)
         
-        # Create visualization (matching your Colab 4-panel layout)
-        fig, axes = plt.subplots(1, 4, figsize=(20, 5))
-        fig.suptitle('🧠 Your Attention U-Net Results', fontsize=16, fontweight='bold')
+        # Original
+        axes[0,0].imshow(image, cmap='gray')
+        axes[0,0].set_title('Original Image')
+        axes[0,0].axis('off')
         
-        titles = ["Original Image", "Tumor Segmentation", "Inverted Mask", "Tumor Only"]
-        images = [original_np, pred_mask_np * 255, inv_pred_mask_np, tumor_only]
-        cmaps = ['gray', 'hot', 'gray', 'gray']
+        # Attention Heatmap
+        axes[0,1].imshow(np.array(image), cmap='gray')
+        axes[0,1].imshow(att_heatmap, alpha=0.5)
+        axes[0,1].set_title('Attention Heatmap')
+        axes[0,1].axis('off')
         
-        for i, ax in enumerate(axes):
-            ax.imshow(images[i], cmap=cmaps[i])
-            ax.set_title(titles[i], fontsize=12, fontweight='bold')
-            ax.axis('off')
+        # Predicted Mask
+        axes[0,2].imshow(binary_mask, cmap='gray')
+        axes[0,2].set_title('Predicted Mask')
+        axes[0,2].axis('off')
         
+        # Ground Truth (if available)
+        if ground_truth is not None:
+            gt_np = np.array(ground_truth.resize((256, 256)))
+            axes[1,0].imshow(gt_np, cmap='gray')
+            axes[1,0].set_title('Ground Truth Mask')
+            axes[1,0].axis('off')
+            
+            # Comparison Overlay
+            overlay = np.array(image.convert('RGB'))
+            overlay[binary_mask > 0] = [0, 255, 0]  # Green for prediction
+            overlay[gt_np > 0] = [255, 0, 0]  # Red for ground truth
+            axes[1,1].imshow(overlay)
+            axes[1,1].set_title('Prediction (Green) vs GT (Red)')
+            axes[1,1].axis('off')
+            
+            # IoU Calculation
+            intersection = np.sum(binary_mask * (gt_np > 0))
+            union = np.sum(binary_mask) + np.sum(gt_np > 0) - intersection
+            iou = intersection / (union + 1e-8)
+            
+            axes[1,2].text(0.1, 0.5, f'IoU Score: {iou:.4f}', fontsize=20)
+            axes[1,2].axis('off')
+        else:
+            # Overlay for prediction only
+            overlay = np.array(image.convert('RGB'))
+            overlay[binary_mask > 0] = [255, 0, 0]
+            axes[1,0].imshow(overlay)
+            axes[1,0].set_title('Prediction Overlay')
+            axes[1,0].axis('off')
+            
+            axes[1,1].axis('off')
+            axes[1,2].axis('off')
+
         plt.tight_layout()
         
-        # Save result
         buf = io.BytesIO()
         plt.savefig(buf, format='png', dpi=150, bbox_inches='tight', facecolor='white')
         buf.seek(0)
@@ -218,187 +328,128 @@ def predict_tumor(image):
         
         result_image = Image.open(buf)
         
-        # Calculate statistics (like your Colab code)
-        tumor_pixels = np.sum(pred_mask_np)
-        total_pixels = pred_mask_np.size
+        # Statistics
+        tumor_pixels = np.sum(binary_mask)
+        total_pixels = binary_mask.size
         tumor_percentage = (tumor_pixels / total_pixels) * 100
         
-        # Calculate confidence metrics
-        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'}
-- **Tumor Area**: {tumor_percentage:.2f}% of brain region
-- **Tumor Pixels**: {tumor_pixels:,} pixels
-- **Max Confidence**: {max_confidence:.4f}
-- **Mean Confidence**: {mean_confidence:.4f}
-
-### 🔬 Your Model Information:
-- **Architecture**: YOUR trained Attention U-Net
-- **Training Performance**: Dice: 0.8420, IoU: 0.7297 
-- **Input**: Grayscale (single channel)
-- **Output**: Binary segmentation mask
-- **Device**: {device.type.upper()}
-
-### 🎯 Model Performance:
-- **Training Accuracy**: 98.90%
-- **Best Dice Score**: 0.8420
-- **Best IoU Score**: 0.7297
-- **Training Dataset**: Brain tumor segmentation dataset
-
-### 📈 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
-
-### ⚠️ 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!
+## Brain Tumor Segmentation Results
+
+### Detection Summary
+- Tumor Percentage: {tumor_percentage:.2f}%
+- Tumor Pixels: {tumor_pixels}
+- File: {filename if filename else 'Uploaded Image'}
+
+### Model Information
+- Your Attention U-Net Model
+- Test-Time Augmentation: Applied
+- Attention Visualization: Included
         """
         
-        print(f"✅ Your model analysis completed! Tumor area: {tumor_percentage:.2f}%")
+        if ground_truth is not None:
+            analysis_text += f"\n- IoU with Ground Truth: {iou:.4f}"
+        
         return result_image, analysis_text
         
     except Exception as e:
-        error_msg = f"❌ Error with your model: {str(e)}"
-        print(error_msg)
-        return None, error_msg
+        return None, f"Error: {str(e)}"
 
 def clear_all():
-    return None, None, "Upload a brain MRI image to test YOUR trained Attention U-Net model"
+    return None, None, None, "Upload or load an image for analysis"
 
-# Enhanced CSS for your model
+# Professional CSS (white, clean, professional)
 css = """
 .gradio-container {
     max-width: 1400px !important;
     margin: auto !important;
+    background-color: white !important;
+    font-family: 'Arial', sans-serif !important;
+}
+h1, h2, h3, h4 {
+    color: #333333 !important;
+}
+button {
+    background-color: #f0f0f0 !important;
+    color: #333333 !important;
+    border: 1px solid #dddddd !important;
+    border-radius: 4px !important;
+}
+button.primary {
+    background-color: #007bff !important;
+    color: white !important;
+}
+.output-image {
+    border: 1px solid #dddddd !important;
+    border-radius: 4px !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);
+.markdown {
+    line-height: 1.6 !important;
+    color: #555555 !important;
 }
 """
 
-# Create Gradio interface for your model
-with gr.Blocks(css=css, title="🧠 Your Attention U-Net Model", theme=gr.themes.Soft()) as app:
+# Create professional Gradio interface
+with gr.Blocks(css=css, title="Brain Tumor Segmentation Application") 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>
+    gr.Markdown("""
+    # Brain Tumor Segmentation Using Attention U-Net
+    A professional tool for medical image analysis
     """)
     
     with gr.Row():
         with gr.Column(scale=1):
-            gr.Markdown("### 📤 Upload Brain MRI")
+            gr.Markdown("### Input Selection")
             
             image_input = gr.Image(
-                label="Brain MRI Scan",
+                label="Upload Brain MRI",
                 type="pil",
                 sources=["upload", "webcam"],
-                height=350
+                height=300
             )
             
-            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)
+            load_random_btn = gr.Button("Load Random Sample from Dataset", variant="primary")
             
-            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>4-Panel View:</strong> Complete analysis like your Colab</li>
-                </ul>
-            </div>
-            """)
+            with gr.Row():
+                analyze_btn = gr.Button("Analyze Image", variant="primary", scale=2)
+                clear_btn = gr.Button("Clear", scale=1)
         
         with gr.Column(scale=2):
-            gr.Markdown("### 📊 Your Model Results")
+            gr.Markdown("### Analysis Results")
             
             output_image = gr.Image(
-                label="Your Attention U-Net Analysis",
+                label="Segmentation Results",
                 type="pil",
-                height=500
+                height=400
             )
             
             analysis_output = gr.Markdown(
-                value="Upload a brain MRI image to test YOUR trained Attention U-Net model.",
-                elem_id="analysis"
+                value="Select an input method to begin analysis."
             )
-
-    # 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>
-                <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>
-            </div>
-            <div>
-                <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>
-                    Always validate with medical professionals for any clinical application.
-                </p>
-            </div>
-        </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
-        </p>
-    </div>
-    """)
     
+    # Hidden state for ground truth and filename
+    ground_truth_state = gr.State()
+    filename_state = gr.State()
+
     # Event handlers
     analyze_btn.click(
         fn=predict_tumor,
-        inputs=[image_input],
-        outputs=[output_image, analysis_output],
-        show_progress=True
+        inputs=[image_input, ground_truth_state, filename_state],
+        outputs=[output_image, analysis_output]
+    )
+    
+    load_random_btn.click(
+        fn=load_random_sample,
+        inputs=[],
+        outputs=[image_input, ground_truth_state, filename_state, analysis_output]
     )
     
     clear_btn.click(
         fn=clear_all,
         inputs=[],
-        outputs=[image_input, output_image, analysis_output]
+        outputs=[image_input, output_image, ground_truth_state, analysis_output]
     )
 
 if __name__ == "__main__":
-    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
-    )
+    print("Starting Brain Tumor Segmentation Application...")
+    app.launch()