Update app.py

app.py

@@ -6,20 +6,19 @@ import cv2
 from PIL import Image
 import matplotlib.pyplot as plt
 import io
-from torchvision import transforms
 import torchvision.transforms.functional as TF
-import urllib.request
+from torchvision import transforms
 import os
-import kagglehub
 import random
-from pathlib import Path
-import seaborn as sns
+import urllib.request
+import zipfile
+import kagglehub
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = None
-dataset_path = None
+DATASET_PATH = "brain_tumor_dataset"
 
-# Define your Attention U-Net architecture (from your training code)
+# Your model classes (from previous code)
 class DoubleConv(nn.Module):
     def __init__(self, in_channels, out_channels):
         super(DoubleConv, self).__init__()
@@ -61,7 +60,7 @@ class AttentionBlock(nn.Module):
         x1 = self.W_x(x)
         psi = self.relu(g1 + x1)
         psi = self.psi(psi)
-        return x * psi, psi  # Return attention coefficients for visualization
+        return x * psi, psi  # Return both attended features and attention map
 
 class AttentionUNET(nn.Module):
     def __init__(self, in_channels=1, out_channels=1, features=[32, 64, 128, 256]):
@@ -72,7 +71,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
@@ -80,7 +79,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))
@@ -88,9 +87,9 @@ class AttentionUNET(nn.Module):
 
         self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
 
-    def forward(self, x, return_attention=False):
+    def forward(self, x):
         skip_connections = []
-        attention_maps = []
+        attention_maps = []  # To collect attention maps
 
         for down in self.downs:
             x = down(x)
@@ -107,30 +106,13 @@ class AttentionUNET(nn.Module):
             if x.shape != skip_connection.shape:
                 x = TF.resize(x, size=skip_connection.shape[2:])
 
-            skip_connection, attention_coeff = self.attentions[idx // 2](skip_connection, x)
-            if return_attention:
-                attention_maps.append(attention_coeff)
+            attended, attn_map = self.attentions[idx // 2](x, skip_connection)  # Get attention map
+            attention_maps.append(attn_map)
             
-            concat_skip = torch.cat((skip_connection, x), dim=1)
+            concat_skip = torch.cat((attended, x), dim=1)
             x = self.ups[idx+1](concat_skip)
 
-        output = self.final_conv(x)
-        
-        if return_attention:
-            return output, attention_maps
-        return output
-
-def download_dataset():
-    """Download Brain Tumor Segmentation dataset from Kaggle"""
-    global dataset_path
-    try:
-        print("📥 Downloading Brain Tumor Segmentation dataset...")
-        dataset_path = kagglehub.dataset_download('nikhilroxtomar/brain-tumor-segmentation')
-        print(f"✅ Dataset downloaded to: {dataset_path}")
-        return dataset_path
-    except Exception as e:
-        print(f"❌ Failed to download dataset: {e}")
-        return None
+        return self.final_conv(x), attention_maps
 
 def download_model():
     """Download your trained model from HuggingFace"""
@@ -138,758 +120,307 @@ def download_model():
     model_path = "best_attention_model.pth.tar"
     
     if not os.path.exists(model_path):
-        print("📥 Downloading trained model...")
+        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 None
-    else:
-        print("✅ Model already exists!")
-    
     return model_path
 
-def load_attention_model():
-    """Load trained Attention U-Net model"""
+def load_model():
+    """Load your trained Attention U-Net model"""
     global model
     if model is None:
         try:
-            print("🔄 Loading Attention U-Net model...")
+            print("🔄 Loading your trained Attention U-Net model...")
             
+            # Download model if needed
             model_path = download_model()
             if model_path is None:
                 return None
             
+            # Initialize your model architecture
             model = AttentionUNET(in_channels=1, out_channels=1).to(device)
+            
+            # Load your trained weights
             checkpoint = torch.load(model_path, map_location=device, weights_only=True)
             model.load_state_dict(checkpoint["state_dict"])
             model.eval()
             
-            print("✅ Attention U-Net model loaded successfully!")
+            print("✅ Your Attention U-Net model loaded successfully!")
         except Exception as e:
-            print(f"❌ Error loading model: {e}")
+            print(f"❌ Error loading your model: {e}")
             model = None
     return model
 
-def get_random_sample_from_dataset():
-    """Get a random sample image and ground truth mask from the dataset"""
-    global dataset_path
+def preprocess_image(image):
+    """Preprocessing like your Colab code"""
+    # Convert to grayscale
+    if image.mode != 'L':
+        image = image.convert('L')
     
-    if dataset_path is None:
-        dataset_path = download_dataset()
-        if dataset_path is None:
-            return None, None
+    # Use your exact transform
+    val_test_transform = transforms.Compose([
+        transforms.Resize((256,256)),
+        transforms.ToTensor()
+    ])
     
-    try:
-        images_path = Path(dataset_path) / "images"
-        masks_path = Path(dataset_path) / "masks"
-        
-        if not images_path.exists() or not masks_path.exists():
-            print("❌ Dataset structure not found")
-            return None, None
-        
-        # Get all image files
-        image_files = list(images_path.glob("*.jpg")) + list(images_path.glob("*.png")) + list(images_path.glob("*.tif"))
-        
-        if not image_files:
-            print("❌ No image files found in dataset")
-            return None, None
-        
-        # Select random image
-        random_image_file = random.choice(image_files)
-        image_name = random_image_file.stem
-        
-        # Find corresponding mask
-        possible_mask_extensions = ['.jpg', '.png', '.tif', '.gif']
-        mask_file = None
-        
-        for ext in possible_mask_extensions:
-            potential_mask = masks_path / f"{image_name}{ext}"
-            if potential_mask.exists():
-                mask_file = potential_mask
-                break
-        
-        if mask_file is None:
-            print(f"❌ No corresponding mask found for {image_name}")
-            return None, None
-        
-        # Load image and mask
-        image = Image.open(random_image_file).convert('L')
-        mask = Image.open(mask_file).convert('L')
-        
-        print(f"✅ Loaded random sample: {image_name}")
-        return image, mask
-        
-    except Exception as e:
-        print(f"❌ Error loading random sample: {e}")
-        return None, None
+    return val_test_transform(image).unsqueeze(0)  # Add batch dimension
 
-def test_time_augmentation(model, image_tensor):
-    """Apply Test-Time Augmentation (TTA) for robust predictions"""
-    augmentations = [
-        lambda x: x,  # Original
-        lambda x: torch.flip(x, dims=[3]),  # Horizontal flip
-        lambda x: torch.flip(x, dims=[2]),  # Vertical flip
-        lambda x: torch.flip(x, dims=[2, 3]),  # Both flips
-        lambda x: torch.rot90(x, k=1, dims=[2, 3]),  # 90° rotation
-        lambda x: torch.rot90(x, k=3, dims=[2, 3]),  # 270° rotation
-    ]
+def post_process_mask(pred_mask_np):
+    """Post-processing with morphological operations (Novelty 1)"""
+    # Binarize
+    binary_mask = (pred_mask_np > 0.5).astype(np.uint8)
     
-    reverse_augmentations = [
-        lambda x: x,  # Original
-        lambda x: torch.flip(x, dims=[3]),  # Reverse horizontal flip
-        lambda x: torch.flip(x, dims=[2]),  # Reverse vertical flip
-        lambda x: torch.flip(x, dims=[2, 3]),  # Reverse both flips
-        lambda x: torch.rot90(x, k=3, dims=[2, 3]),  # Reverse 90° rotation
-        lambda x: torch.rot90(x, k=1, dims=[2, 3]),  # Reverse 270° rotation
-    ]
+    # Morphological opening to remove small noise
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5,5))
+    binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_OPEN, kernel)
     
+    # Morphological closing to fill gaps
+    binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel)
+    
+    return binary_mask
+
+def test_time_augmentation(input_tensor, model):
+    """Test-Time Augmentation (Novelty 2)"""
     predictions = []
     
-    with torch.no_grad():
-        for aug, rev_aug in zip(augmentations, reverse_augmentations):
-            # Apply augmentation
-            aug_input = aug(image_tensor)
-            
-            # Get prediction
-            pred = torch.sigmoid(model(aug_input))
-            
-            # Reverse augmentation on prediction
-            pred = rev_aug(pred)
-            
-            predictions.append(pred)
+    # Original
+    pred, _ = model(input_tensor)
+    predictions.append(torch.sigmoid(pred))
     
-    # Average all predictions
-    tta_prediction = torch.mean(torch.stack(predictions), dim=0)
+    # Horizontal flip
+    hflip = TF.hflip(input_tensor)
+    pred_h, _ = model(hflip)
+    pred_h = TF.hflip(pred_h)
+    predictions.append(torch.sigmoid(pred_h))
     
-    return tta_prediction
-
-def generate_attention_heatmaps(model, image_tensor):
-    """Generate attention heatmaps for interpretability"""
-    with torch.no_grad():
-        pred, attention_maps = model(image_tensor, return_attention=True)
-        
-        # Convert attention maps to numpy for visualization
-        heatmaps = []
-        for i, att_map in enumerate(attention_maps):
-            # Resize attention map to match input size
-            att_map_resized = TF.resize(att_map, (256, 256))
-            att_np = att_map_resized.cpu().squeeze().numpy()
-            heatmaps.append(att_np)
-        
-        return heatmaps
-
-def preprocess_image(image):
-    """Preprocessing exactly like training code"""
-    if image.mode != 'L':
-        image = image.convert('L')
+    # Vertical flip
+    vflip = TF.vflip(input_tensor)
+    pred_v, _ = model(vflip)
+    pred_v = TF.vflip(pred_v)
+    predictions.append(torch.sigmoid(pred_v))
     
-    val_test_transform = transforms.Compose([
-        transforms.Resize((256, 256)),
-        transforms.ToTensor()
-    ])
+    # Average predictions
+    avg_pred = torch.mean(torch.stack(predictions), dim=0)
+    return avg_pred.squeeze().cpu().numpy()
+
+def generate_attention_heatmap(attention_maps, size=(256,256)):
+    """Generate attention heatmap visualization (Novelty 3)"""
+    # Average attention maps from all levels
+    avg_attn = torch.mean(torch.cat([TF.resize(m, size) for m in attention_maps]), dim=0)
+    attn_np = avg_attn.squeeze().cpu().numpy()
     
-    return val_test_transform(image).unsqueeze(0)
+    # Normalize and apply colormap
+    attn_norm = (attn_np - attn_np.min()) / (attn_np.max() - attn_np.min() + 1e-8)
+    heatmap = plt.cm.hot(attn_norm)[:,:,:3] * 255
+    return heatmap.astype(np.uint8)
 
-def calculate_metrics(pred_mask, ground_truth_mask):
-    """Calculate Dice and IoU metrics"""
-    pred_binary = (pred_mask > 0.5).float()
-    gt_binary = (ground_truth_mask > 0.5).float()
+def download_dataset():
+    """Download and extract the dataset if not present"""
+    if not os.path.exists(DATASET_PATH):
+        print("📥 Downloading brain tumor dataset...")
+        try:
+            path = kagglehub.dataset_download('nikhilroxtomar/brain-tumor-segmentation')
+            print(f"Dataset downloaded to: {path}")
+            
+            # Extract if zipped
+            for file in os.listdir(path):
+                if file.endswith('.zip'):
+                    with zipfile.ZipFile(os.path.join(path, file), 'r') as zip_ref:
+                        zip_ref.extractall(DATASET_PATH)
+                    print("✅ Dataset extracted!")
+                    return True
+        except Exception as e:
+            print(f"❌ Failed to download dataset: {e}")
+            return False
+    print("✅ Dataset already exists!")
+    return True
+
+def get_random_sample():
+    """Get random image and mask from dataset"""
+    if not os.path.exists(DATASET_PATH):
+        if not download_dataset():
+            return None, None
+    
+    images_path = os.path.join(DATASET_PATH, "images")
+    masks_path = os.path.join(DATASET_PATH, "masks")
+    
+    image_files = [f for f in os.listdir(images_path) if f.endswith(('.png', '.jpg'))]
+    
+    if not image_files:
+        return None, None
     
-    # Dice coefficient
-    intersection = torch.sum(pred_binary * gt_binary)
-    dice = (2.0 * intersection) / (torch.sum(pred_binary) + torch.sum(gt_binary) + 1e-8)
+    random_file = random.choice(image_files)
     
-    # IoU
-    union = torch.sum(pred_binary) + torch.sum(gt_binary) - intersection
-    iou = intersection / (union + 1e-8)
+    img_path = os.path.join(images_path, random_file)
+    mask_path = os.path.join(masks_path, random_file)
     
-    return dice.item(), iou.item()
+    if not os.path.exists(mask_path):
+        return None, None
+    
+    return Image.open(img_path), Image.open(mask_path)
 
-def predict_with_enhancements(image, ground_truth=None, use_tta=True, show_attention=True):
-    """Enhanced prediction with TTA and attention visualization"""
-    current_model = load_attention_model()
+def predict_tumor(image, use_tta=True, show_attention=True, is_dataset_sample=False, ground_truth=None):
+    current_model = load_your_attention_model()
     
     if current_model is None:
-        return None, "❌ Failed to load 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 enhanced Attention U-Net...")
+        print("Processing image...")
         
         input_tensor = preprocess_image(image).to(device)
         
-        # Standard prediction
-        with torch.no_grad():
-            standard_pred = torch.sigmoid(current_model(input_tensor))
-        
-        # Test-Time Augmentation
+        # Use TTA if enabled
         if use_tta:
-            tta_pred = test_time_augmentation(current_model, input_tensor)
-            final_pred = tta_pred
-        else:
-            final_pred = standard_pred
-        
-        # Generate attention heatmaps
-        attention_heatmaps = []
-        if show_attention:
-            attention_heatmaps = generate_attention_heatmaps(current_model, input_tensor)
-        
-        # Convert predictions to binary
-        pred_mask_binary = (final_pred > 0.5).float()
-        pred_mask_np = pred_mask_binary.cpu().squeeze().numpy()
-        standard_mask_np = (standard_pred > 0.5).float().cpu().squeeze().numpy()
-        
-        # Prepare images for visualization
-        original_np = np.array(image.convert('L').resize((256, 256)))
-        
-        # Create comprehensive visualization
-        if ground_truth is not None:
-            # With ground truth comparison
-            gt_np = np.array(ground_truth.convert('L').resize((256, 256)))
-            gt_binary = (gt_np > 127).astype(np.float32)  # Threshold ground truth
-            
-            # Calculate metrics
-            gt_tensor = torch.tensor(gt_binary).unsqueeze(0).unsqueeze(0).to(device)
-            dice_score, iou_score = calculate_metrics(final_pred, gt_tensor)
-            
-            # Create figure with ground truth comparison
-            n_cols = 6 if show_attention and attention_heatmaps else 5
-            fig, axes = plt.subplots(2, n_cols, figsize=(4*n_cols, 8))
-            fig.suptitle('🧠 Enhanced Attention U-Net Analysis with Ground Truth Comparison', fontsize=16, weight='bold')
-            
-            # Top row - Standard analysis
-            axes[0, 0].imshow(original_np, cmap='gray')
-            axes[0, 0].set_title('Original Image', fontsize=12, weight='bold')
-            axes[0, 0].axis('off')
-            
-            axes[0, 1].imshow(standard_mask_np * 255, cmap='hot')
-            axes[0, 1].set_title('Standard Prediction', fontsize=12, weight='bold')
-            axes[0, 1].axis('off')
-            
-            axes[0, 2].imshow(pred_mask_np * 255, cmap='hot')
-            axes[0, 2].set_title(f'{"TTA Enhanced" if use_tta else "Final Prediction"}', fontsize=12, weight='bold')
-            axes[0, 2].axis('off')
-            
-            axes[0, 3].imshow(gt_binary * 255, cmap='hot')
-            axes[0, 3].set_title('Ground Truth', fontsize=12, weight='bold')
-            axes[0, 3].axis('off')
-            
-            # Overlay comparison
-            overlay = original_np.copy()
-            overlay = np.stack([overlay, overlay, overlay], axis=-1)
-            overlay[pred_mask_np > 0.5] = [255, 0, 0]  # Red for prediction
-            overlay[gt_binary > 0.5] = [0, 255, 0]  # Green for ground truth
-            overlap = (pred_mask_np > 0.5) & (gt_binary > 0.5)
-            overlay[overlap] = [255, 255, 0]  # Yellow for overlap
-            
-            axes[0, 4].imshow(overlay.astype(np.uint8))
-            axes[0, 4].set_title('Overlay (Red:Pred, Green:GT, Yellow:Match)', fontsize=10, weight='bold')
-            axes[0, 4].axis('off')
-            
-            if show_attention and attention_heatmaps:
-                # Show combined attention
-                combined_attention = np.mean(attention_heatmaps, axis=0)
-                axes[0, 5].imshow(combined_attention, cmap='jet', alpha=0.7)
-                axes[0, 5].imshow(original_np, cmap='gray', alpha=0.3)
-                axes[0, 5].set_title('Attention Heatmap', fontsize=12, weight='bold')
-                axes[0, 5].axis('off')
-            
-            # Bottom row - Individual attention maps or detailed analysis
-            if show_attention and attention_heatmaps:
-                for i, heatmap in enumerate(attention_heatmaps[:n_cols]):
-                    axes[1, i].imshow(heatmap, cmap='jet', alpha=0.7)
-                    axes[1, i].imshow(original_np, cmap='gray', alpha=0.3)
-                    axes[1, i].set_title(f'Attention Gate {i+1}', fontsize=10, weight='bold')
-                    axes[1, i].axis('off')
-            else:
-                # Show tumor extraction and analysis
-                tumor_only = np.where(pred_mask_np == 1, original_np, 255)
-                inv_mask = np.where(pred_mask_np == 1, 0, 255)
-                
-                axes[1, 0].imshow(tumor_only, cmap='gray')
-                axes[1, 0].set_title('Tumor Extraction', fontsize=12, weight='bold')
-                axes[1, 0].axis('off')
-                
-                axes[1, 1].imshow(inv_mask, cmap='gray')
-                axes[1, 1].set_title('Inverted Mask', fontsize=12, weight='bold')
-                axes[1, 1].axis('off')
-                
-                # Difference map
-                diff_map = np.abs(pred_mask_np - gt_binary)
-                axes[1, 2].imshow(diff_map, cmap='Reds')
-                axes[1, 2].set_title('Difference Map', fontsize=12, weight='bold')
-                axes[1, 2].axis('off')
-                
-                # Clear remaining axes
-                for j in range(3, n_cols):
-                    axes[1, j].axis('off')
+            pred_np = test_time_augmentation(input_tensor, current_model)
         else:
-            # Without ground truth
-            n_cols = 5 if show_attention and attention_heatmaps else 4
-            fig, axes = plt.subplots(2, n_cols, figsize=(4*n_cols, 8))
-            fig.suptitle('🧠 Enhanced Attention U-Net Analysis', fontsize=16, weight='bold')
-            
-            # Top row
-            images = [original_np, standard_mask_np * 255, pred_mask_np * 255]
-            titles = ["Original Image", "Standard Prediction", f'{"TTA Enhanced" if use_tta else "Final Prediction"}']
-            cmaps = ['gray', 'hot', 'hot']
-            
-            for i in range(3):
-                axes[0, i].imshow(images[i], cmap=cmaps[i])
-                axes[0, i].set_title(titles[i], fontsize=12, weight='bold')
-                axes[0, i].axis('off')
-            
-            # Tumor extraction
-            tumor_only = np.where(pred_mask_np == 1, original_np, 255)
-            axes[0, 3].imshow(tumor_only, cmap='gray')
-            axes[0, 3].set_title('Tumor Extraction', fontsize=12, weight='bold')
-            axes[0, 3].axis('off')
-            
-            if show_attention and attention_heatmaps:
-                combined_attention = np.mean(attention_heatmaps, axis=0)
-                axes[0, 4].imshow(combined_attention, cmap='jet', alpha=0.7)
-                axes[0, 4].imshow(original_np, cmap='gray', alpha=0.3)
-                axes[0, 4].set_title('Combined Attention', fontsize=12, weight='bold')
-                axes[0, 4].axis('off')
-            
-            # Bottom row - Individual attention maps
-            if show_attention and attention_heatmaps:
-                for i, heatmap in enumerate(attention_heatmaps[:n_cols]):
-                    axes[1, i].imshow(heatmap, cmap='jet', alpha=0.7)
-                    axes[1, i].imshow(original_np, cmap='gray', alpha=0.3)
-                    axes[1, i].set_title(f'Attention Gate {i+1}', fontsize=10, weight='bold')
-                    axes[1, i].axis('off')
-            else:
-                # Clear bottom row
-                for j in range(n_cols):
-                    axes[1, j].axis('off')
+            pred, attn_maps = current_model(input_tensor)
+            pred_np = torch.sigmoid(pred).squeeze().cpu().numpy()
+            attn_maps = attn_maps if show_attention else None
+        
+        # Post-processing
+        binary_mask = post_process_mask(pred_np)
+        
+        # Generate attention heatmap if enabled
+        attention_heatmap = None
+        if show_attention and attn_maps:
+            attention_heatmap = generate_attention_heatmap(attn_maps)
+        
+        # Create visualization
+        fig, axes = plt.subplots(1, 3 + int(show_attention) + int(is_dataset_sample and ground_truth is not None), figsize=(20, 5))
+        fig.suptitle('Brain Tumor Segmentation Results', fontsize=16)
+        
+        # Original image
+        original_np = np.array(image.resize((256, 256)))
+        axes[0].imshow(original_np, cmap='gray')
+        axes[0].set_title('Original Image')
+        axes[0].axis('off')
+        
+        # Predicted mask
+        axes[1].imshow(binary_mask * 255, cmap='gray')
+        axes[1].set_title('Predicted Mask')
+        axes[1].axis('off')
+        
+        # Overlay
+        overlay = cv2.cvtColor(original_np, cv2.COLOR_GRAY2RGB) if len(original_np.shape) == 2 else original_np
+        overlay[binary_mask == 1] = [255, 0, 0]  # Red for tumor
+        overlay = cv2.addWeighted(original_np, 0.7, overlay, 0.3, 0)
+        axes[2].imshow(overlay)
+        axes[2].set_title('Overlay')
+        axes[2].axis('off')
+        
+        col = 3
+        if show_attention and attention_heatmap is not None:
+            axes[col].imshow(attention_heatmap)
+            axes[col].set_title('Attention Heatmap')
+            axes[col].axis('off')
+            col += 1
+        
+        if is_dataset_sample and ground_truth is not None:
+            gt_np = np.array(ground_truth.resize((256, 256)))
+            axes[col].imshow(gt_np, cmap='gray')
+            axes[col].set_title('Ground Truth')
+            axes[col].axis('off')
         
         plt.tight_layout()
         
-        # Save result
         buf = io.BytesIO()
-        plt.savefig(buf, format='png', dpi=150, bbox_inches='tight', facecolor='white')
+        plt.savefig(buf, format='png', dpi=150, bbox_inches='tight')
         buf.seek(0)
         plt.close()
         
         result_image = Image.open(buf)
         
-        # Calculate statistics
-        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
         
-        max_confidence = torch.max(final_pred).item()
-        mean_confidence = torch.mean(final_pred).item()
-        
-        # Enhanced analysis text
         analysis_text = f"""
-## 🧠 Enhanced Attention U-Net Analysis Results
-
-### 📊 Detection Summary
-- **Status**: {'🔴 TUMOR DETECTED' if tumor_pixels > 50 else '🟢 NO SIGNIFICANT TUMOR'}
-- **Tumor Coverage**: {tumor_percentage:.2f}% of brain region
-- **Tumor Pixels**: {tumor_pixels:,} pixels
-- **Max Confidence**: {max_confidence:.4f}
-- **Mean Confidence**: {mean_confidence:.4f}
+### Segmentation Statistics
+- Tumor Area Percentage: {tumor_percentage:.2f}%
+- Tumor Pixels: {tumor_pixels}
+- Total Pixels: {total_pixels}
+- TTA Used: {'Yes' if use_tta else 'No'}
+- Attention Visualization: {'Yes' if show_attention else 'No'}
 """
-
-        if ground_truth is not None:
+        
+        if is_dataset_sample and ground_truth is not None:
+            gt_np = np.array(ground_truth.resize((256, 256)))
+            intersection = np.logical_and(binary_mask, gt_np > 0).sum()
+            union = np.logical_or(binary_mask, gt_np > 0).sum()
+            iou = intersection / (union + 1e-8)
+            dice = (2 * intersection) / (binary_mask.sum() + (gt_np > 0).sum() + 1e-8)
+            
             analysis_text += f"""
-### 🎯 Ground Truth Comparison
-- **Dice Score**: {dice_score:.4f} {'✅ Excellent' if dice_score > 0.8 else '⚠️ Good' if dice_score > 0.6 else '❌ Poor'}
-- **IoU Score**: {iou_score:.4f} {'✅ Excellent' if iou_score > 0.7 else '⚠️ Good' if iou_score > 0.5 else '❌ Poor'}
-- **Model Accuracy**: {'High precision match' if dice_score > 0.8 else 'Reasonable match' if dice_score > 0.6 else 'Needs improvement'}
-"""
-
-        analysis_text += f"""
-### 🚀 Enhancement Features
-- **Test-Time Augmentation**: {'✅ Applied (6 augmentations averaged)' if use_tta else '❌ Disabled'}
-- **Attention Visualization**: {'✅ Generated attention heatmaps' if show_attention else '❌ Disabled'}
-- **Boundary Enhancement**: {'✅ TTA improves edge detection' if use_tta else '⚠️ Standard prediction only'}
-- **Interpretability**: {'✅ Attention gates show focus areas' if show_attention else '❌ Black box mode'}
-
-### 🔬 Model Architecture
-- **Base Model**: Attention U-Net with skip connections
-- **Training Performance**: Dice: 0.8420, IoU: 0.7297, Accuracy: 98.90%
-- **Attention Gates**: 4 levels with soft attention mechanism
-- **Features Channels**: [32, 64, 128, 256] progression
-- **Device**: {device.type.upper()}
-
-### 📈 Enhanced Processing Pipeline
-- **Preprocessing**: Resize(256×256) + Normalization
-- **Augmentations**: Flips (H,V), Rotations (90°,270°), Combined
-- **Attention Fusion**: Multi-scale attention coefficient extraction
-- **Post-processing**: Ensemble averaging + Binary thresholding (0.5)
-
-### ⚠️ Medical Disclaimer
-This enhanced AI model is for **research and educational purposes only**. 
-Results include advanced features for better accuracy and interpretability.
-Always consult medical professionals for clinical applications.
-
-### 🏆 Research Contributions
-✅ **Attention Gates**: Enhanced boundary detection through selective feature passing
-✅ **Test-Time Augmentation**: Robust predictions via ensemble averaging  
-✅ **Interpretability**: Attention heatmaps for clinical trust and validation
-✅ **Efficiency**: No retraining required, minimal computational overhead
+### Comparison with Ground Truth
+- IoU Score: {iou:.4f}
+- Dice Score: {dice:.4f}
 """
         
-        print(f"✅ Enhanced analysis completed! Tumor coverage: {tumor_percentage:.2f}%")
         return result_image, analysis_text
         
     except Exception as e:
-        error_msg = f"❌ Error during enhanced analysis: {str(e)}"
-        print(error_msg)
-        return None, error_msg
+        return None, f"Error: {str(e)}"
 
-def load_random_sample():
-    """Load a random sample from the dataset"""
-    image, mask = get_random_sample_from_dataset()
+def test_random_sample():
+    image, mask = get_random_sample()
     if image is None:
-        return None, None, "❌ Failed to load random sample from dataset"
-    return image, mask, "✅ Random sample loaded from dataset"
+        return None, "Failed to load dataset sample. Please download dataset first."
+    return predict_tumor(image, use_tta=True, show_attention=True, is_dataset_sample=True, ground_truth=mask)
 
-def clear_all():
-    return None, None, None, "Upload a brain MRI image or load a random sample to test the enhanced model"
-
-# Enhanced professional CSS
+# Custom CSS for professional, minimalist look
 css = """
-.gradio-container {
-    max-width: 1600px !important;
-    margin: auto !important;
-    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
-}
-
-#title {
-    text-align: center;
-    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-    color: white;
-    padding: 40px;
-    border-radius: 20px;
-    margin-bottom: 30px;
-    box-shadow: 0 12px 24px rgba(102, 126, 234, 0.4);
-}
-
-.feature-box {
-    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
-    border-radius: 15px;
-    padding: 25px;
-    margin: 15px 0;
-    color: white;
-    box-shadow: 0 8px 16px rgba(240, 147, 251, 0.3);
-}
-
-.metric-card {
-    background: linear-gradient(135deg, #4facfe 0%, #00f2fe 100%);
-    border-radius: 12px;
-    padding: 20px;
-    text-align: center;
-    margin: 10px;
-    box-shadow: 0 6px 12px rgba(79, 172, 254, 0.3);
-}
-
-.enhancement-badge {
-    display: inline-block;
-    background: linear-gradient(45deg, #fa709a 0%, #fee140 100%);
-    color: white;
-    padding: 8px 16px;
-    border-radius: 25px;
-    margin: 5px;
-    font-weight: bold;
-    box-shadow: 0 4px 8px rgba(250, 112, 154, 0.3);
-}
+body, .gradio-container { font-family: 'Arial', sans-serif; color: #333; }
+h1, h2, h3, h4 { color: #2c3e50; font-weight: 500; }
+.button { background-color: #3498db; color: white; border: none; border-radius: 4px; padding: 10px 20px; font-size: 16px; cursor: pointer; }
+.button:hover { background-color: #2980b9; }
+.card { border: 1px solid #e0e0e0; border-radius: 8px; padding: 20px; background: white; box-shadow: 0 2px 4px rgba(0,0,0,0.1); }
 """
 
-# Create enhanced Gradio interface
-with gr.Blocks(css=css, title="🧠 Enhanced Brain Tumor Segmentation", theme=gr.themes.Soft()) as app:
-    
-    gr.HTML("""
-    <div id="title">
-        <h1>🧠 Enhanced Attention U-Net Brain Tumor Segmentation</h1>
-        <p style="font-size: 20px; margin-top: 20px; font-weight: 300;">
-            🚀 Advanced Medical AI with Test-Time Augmentation & Attention Visualization
-        </p>
-        <p style="font-size: 16px; margin-top: 15px; opacity: 0.9;">
-            📊 Performance: Dice 0.8420 • IoU 0.7297 • Accuracy 98.90% | 
-            🔬 Research-Grade Interpretability & Robustness
-        </p>
-    </div>
-    """)
+with gr.Blocks(css=css, title="Brain Tumor Segmentation") as app:
+    gr.Markdown("# Brain Tumor Segmentation Using Attention U-Net")
     
     with gr.Row():
         with gr.Column(scale=1):
-            gr.Markdown("### 📤 Input & Controls")
-            
-            with gr.Tab("📸 Upload Image"):
-                image_input = gr.Image(
-                    label="Brain MRI Scan",
-                    type="pil",
-                    sources=["upload", "webcam"],
-                    height=300
-                )
-            
-            with gr.Tab("🎲 Random Sample"):
-                random_image = gr.Image(
-                    label="Sample Image",
-                    type="pil",
-                    height=300,
-                    interactive=False
-                )
-                random_ground_truth = gr.Image(
-                    label="Ground Truth Mask",
-                    type="pil",
-                    height=300,
-                    interactive=False
-                )
-                load_sample_btn = gr.Button("🎲 Load Random Sample", variant="secondary", size="lg")
-                sample_status = gr.Textbox(label="Sample Status", interactive=False)
-            
-            gr.Markdown("### ⚙️ Enhancement Options")
-            
-            use_tta = gr.Checkbox(
-                label="🔄 Test-Time Augmentation",
-                value=True,
-                info="Apply multiple augmentations for robust predictions"
-            )
-            
-            show_attention = gr.Checkbox(
-                label="🔥 Attention Visualization",
-                value=True,
-                info="Generate attention heatmaps for interpretability"
-            )
+            gr.Markdown("### Input")
+            image_input = gr.Image(label="Upload Image", type="pil")
             
             with gr.Row():
-                analyze_btn = gr.Button(
-                    "🧠 Analyze with Enhanced Model", 
-                    variant="primary", 
-                    scale=3, 
-                    size="lg"
-                )
-                clear_btn = gr.Button("🗑️ Clear All", variant="secondary", scale=1)
-            
-            gr.HTML("""
-            <div class="feature-box">
-                <h4 style="margin-bottom: 15px;">🎯 Research Innovations</h4>
-                <div class="enhancement-badge">Attention Gates</div>
-                <div class="enhancement-badge">Test-Time Augmentation</div>
-                <div class="enhancement-badge">Interpretability</div>
-                <div class="enhancement-badge">Ground Truth Comparison</div>
-                <p style="margin-top: 15px; font-size: 14px; opacity: 0.9;">
-                    Advanced medical AI combining accuracy, robustness, and clinical interpretability
-                </p>
-            </div>
-            """)
+                predict_btn = gr.Button("Predict")
+                random_btn = gr.Button("Test Random Sample")
+                download_btn = gr.Button("Download Dataset")
         
         with gr.Column(scale=2):
-            gr.Markdown("### 📊 Enhanced Analysis Results")
-            
-            output_image = gr.Image(
-                label="Comprehensive Analysis Visualization",
-                type="pil",
-                height=600
-            )
-            
-            with gr.Accordion("📈 Detailed Analysis Report", open=True):
-                analysis_output = gr.Markdown(
-                    value="Upload a brain MRI image or load a random sample to test the enhanced Attention U-Net model.",
-                    elem_id="analysis"
-                )
-
-    # Performance metrics section
-    gr.HTML("""
-    <div style="margin-top: 40px;">
-        <h3 style="text-align: center; color: #4a5568; margin-bottom: 25px;">📊 Model Performance & Research Contributions</h3>
-        <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(300px, 1fr)); gap: 20px; margin-bottom: 30px;">
-            
-            <div class="metric-card">
-                <h4 style="color: white; margin-bottom: 10px;">🎯 Segmentation Accuracy</h4>
-                <div style="font-size: 24px; font-weight: bold; margin: 10px 0;">98.90%</div>
-                <p style="font-size: 14px; opacity: 0.9;">Training accuracy on brain tumor dataset</p>
-            </div>
-            
-            <div class="metric-card">
-                <h4 style="color: white; margin-bottom: 10px;">📐 Dice Score</h4>
-                <div style="font-size: 24px; font-weight: bold; margin: 10px 0;">0.8420</div>
-                <p style="font-size: 14px; opacity: 0.9;">Overlap similarity coefficient</p>
-            </div>
-            
-            <div class="metric-card">
-                <h4 style="color: white; margin-bottom: 10px;">🔲 IoU Score</h4>
-                <div style="font-size: 24px; font-weight: bold; margin: 10px 0;">0.7297</div>
-                <p style="font-size: 14px; opacity: 0.9;">Intersection over Union metric</p>
-            </div>
-            
-            <div class="metric-card">
-                <h4 style="color: white; margin-bottom: 10px;">⚡ Enhancement Features</h4>
-                <div style="font-size: 20px; font-weight: bold; margin: 10px 0;">TTA + Attention</div>
-                <p style="font-size: 14px; opacity: 0.9;">Advanced robustness & interpretability</p>
-            </div>
-            
-        </div>
-    </div>
-    """)
-
-    # Research contributions section
-    gr.HTML("""
-    <div style="margin-top: 30px; padding: 30px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 20px; color: white;">
-        <h3 style="text-align: center; margin-bottom: 25px; color: white;">🚀 Novel Research Contributions</h3>
-        
-        <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 30px; margin-bottom: 20px;">
-            
-            <div>
-                <h4 style="margin-bottom: 15px; color: #ffd700;">🔍 1. Enhanced Boundary Detection</h4>
-                <ul style="line-height: 1.8; margin-left: 20px;">
-                    <li><strong>Problem:</strong> Traditional U-Net passes noisy features through skip connections</li>
-                    <li><strong>Solution:</strong> Attention gates filter irrelevant encoder features</li>
-                    <li><strong>Impact:</strong> Cleaner boundaries, reduced false positives</li>
-                </ul>
-            </div>
-            
-            <div>
-                <h4 style="margin-bottom: 15px; color: #ffd700;">🔄 2. Test-Time Augmentation</h4>
-                <ul style="line-height: 1.8; margin-left: 20px;">
-                    <li><strong>Problem:</strong> Medical datasets are small, MRI scans vary across centers</li>
-                    <li><strong>Solution:</strong> Multiple augmentations averaged for robust predictions</li>
-                    <li><strong>Impact:</strong> Improved robustness without retraining</li>
-                </ul>
-            </div>
-            
-            <div>
-                <h4 style="margin-bottom: 15px; color: #ffd700;">🔥 3. Attention Visualization</h4>
-                <ul style="line-height: 1.8; margin-left: 20px;">
-                    <li><strong>Problem:</strong> Deep networks are "black boxes" for clinicians</li>
-                    <li><strong>Solution:</strong> Extract attention coefficients as interpretable heatmaps</li>
-                    <li><strong>Impact:</strong> Build clinical trust through transparency</li>
-                </ul>
-            </div>
-            
-            <div>
-                <h4 style="margin-bottom: 15px; color: #ffd700;">⚡ 4. Efficient Implementation</h4>
-                <ul style="line-height: 1.8; margin-left: 20px;">
-                    <li><strong>Problem:</strong> Complex architectures are hard to deploy</li>
-                    <li><strong>Solution:</strong> Low-overhead enhancements within existing backbone</li>
-                    <li><strong>Impact:</strong> Practical for real-world medical workflows</li>
-                </ul>
-            </div>
-            
-        </div>
-        
-        <div style="text-align: center; padding-top: 20px; border-top: 2px solid rgba(255,255,255,0.3);">
-            <p style="font-size: 16px; font-weight: 600; margin-bottom: 10px;">
-                🎯 Research Gap Addressed: Accuracy + Robustness + Interpretability
-            </p>
-            <p style="font-size: 14px; opacity: 0.9;">
-                This combination tackles three major challenges in medical AI with minimal architectural changes
-            </p>
-        </div>
-    </div>
-    """)
-
-    # Dataset and disclaimer section
-    gr.HTML("""
-    <div style="margin-top: 30px; padding: 25px; background-color: #f7fafc; border-radius: 15px; border-left: 5px solid #667eea;">
-        <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 30px;">
-            
-            <div>
-                <h4 style="color: #667eea; margin-bottom: 15px;">📚 Dataset Information</h4>
-                <p><strong>Source:</strong> Brain Tumor Segmentation (Kaggle)</p>
-                <p><strong>Author:</strong> nikhilroxtomar</p>
-                <p><strong>Structure:</strong> Images + Ground Truth Masks</p>
-                <p><strong>Format:</strong> Grayscale MRI scans</p>
-                <p><strong>Use Case:</strong> Medical image segmentation research</p>
-                <p><strong>Ground Truth:</strong> Available for metric calculation</p>
-            </div>
-            
-            <div>
-                <h4 style="color: #dc2626; margin-bottom: 15px;">⚠️ Medical Disclaimer</h4>
-                <p style="color: #dc2626; font-weight: 600; line-height: 1.5;">
-                    This enhanced AI system is designed for <strong>research and educational purposes only</strong>.<br><br>
-                    
-                    While the model includes advanced features like attention visualization and test-time augmentation 
-                    for improved accuracy and interpretability, all results must be validated by qualified medical professionals.<br><br>
-                    
-                    <strong>Not approved for clinical diagnosis or medical decision making.</strong>
-                </p>
-            </div>
-            
-        </div>
-        
-        <hr style="margin: 25px 0; border: none; border-top: 2px solid #e2e8f0;">
-        
-        <p style="text-align: center; color: #4a5568; margin: 15px 0; font-weight: 600;">
-            🔬 Research-Grade Medical AI • Enhanced Interpretability • Robust Predictions • Ground Truth Validation
-        </p>
-    </div>
-    """)
+            gr.Markdown("### Output")
+            output_image = gr.Image(label="Result")
+            analysis_output = gr.Textbox(label="Analysis", lines=10)
     
     # Event handlers
-    def analyze_with_ground_truth(image, gt_mask, use_tta, show_attention):
-        """Wrapper function to handle ground truth comparison"""
-        return predict_with_enhancements(image, gt_mask, use_tta, show_attention)
-    
-    def analyze_uploaded_image(image, use_tta, show_attention):
-        """Wrapper function for uploaded images without ground truth"""
-        return predict_with_enhancements(image, None, use_tta, show_attention)
-    
-    # Button event handlers
-    analyze_btn.click(
-        fn=lambda img, rand_img, rand_gt, tta, attention: (
-            analyze_with_ground_truth(rand_img, rand_gt, tta, attention) 
-            if rand_img is not None 
-            else analyze_uploaded_image(img, tta, attention)
-        ),
-        inputs=[image_input, random_image, random_ground_truth, use_tta, show_attention],
-        outputs=[output_image, analysis_output],
-        show_progress=True
+    predict_btn.click(
+        fn=predict_tumor,
+        inputs=[image_input],
+        outputs=[output_image, analysis_output]
     )
     
-    load_sample_btn.click(
-        fn=load_random_sample,
+    random_btn.click(
+        fn=test_random_sample,
         inputs=[],
-        outputs=[random_image, random_ground_truth, sample_status],
-        show_progress=True
+        outputs=[output_image, analysis_output]
     )
     
-    clear_btn.click(
-        fn=clear_all,
+    download_btn.click(
+        fn=download_dataset,
         inputs=[],
-        outputs=[image_input, random_image, random_ground_truth, analysis_output]
+        outputs=gr.Textbox(value="Dataset download status...")
     )
 
-    # Auto-load dataset on startup
-    gr.HTML("""
-    <script>
-    document.addEventListener('DOMContentLoaded', function() {
-        console.log('Enhanced Brain Tumor Segmentation App Loaded');
-        console.log('Features: TTA + Attention Visualization + Ground Truth Comparison');
-    });
-    </script>
-    """)
-
 if __name__ == "__main__":
-    print("🚀 Starting Enhanced Brain Tumor Segmentation System...")
-    print("📊 Model Performance: Dice 0.8420, IoU 0.7297, Accuracy 98.90%")
-    print("🔬 Research Features: Attention Gates + TTA + Interpretability")
-    print("📥 Auto-downloading dataset and model...")
-    
-    # Initialize dataset download
-    print("📚 Initializing dataset...")
-    try:
-        dataset_path = download_dataset()
-        if dataset_path:
-            print(f"✅ Dataset ready at: {dataset_path}")
-        else:
-            print("⚠️ Dataset download failed, random samples unavailable")
-    except Exception as e:
-        print(f"⚠️ Dataset initialization error: {e}")
-    
-    app.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True,
-        share=False
-    )
+    app.launch()