Update app.py

app.py

@@ -10,14 +10,22 @@ import torchvision.transforms as transforms
 import torchvision.transforms.functional as TF
 import random
 import os
-import zipfile
 import urllib.request
 import kagglehub
+from glob import glob
 
+# Global variables - loaded once at startup
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = None
+dataset_images = []
+dataset_masks = []
+dataset_loaded = False
 
-# Your Attention U-Net classes (from your code)
+print("="*50)
+print("BRAIN TUMOR SEGMENTATION APPLICATION")
+print("="*50)
+
+# Your Attention U-Net classes (unchanged)
 class DoubleConv(nn.Module):
     def __init__(self, in_channels, out_channels):
         super(DoubleConv, self).__init__()
@@ -59,7 +67,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 map as well
+        return x * psi, psi
 
 class AttentionUNET(nn.Module):
     def __init__(self, in_channels=1, out_channels=1, features=[32, 64, 128, 256]):
@@ -70,15 +78,12 @@ class AttentionUNET(nn.Module):
         self.attentions = nn.ModuleList()
         self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
 
-        # Down part
         for feature in features:
             self.downs.append(DoubleConv(in_channels, feature))
             in_channels = feature
 
-        # Bottleneck
         self.bottleneck = DoubleConv(features[-1], features[-1]*2)
         
-        # 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,7 +93,7 @@ class AttentionUNET(nn.Module):
 
     def forward(self, x):
         skip_connections = []
-        attention_maps = []  # To store attention maps
+        attention_maps = []
 
         for down in self.downs:
             x = down(x)
@@ -105,222 +110,265 @@ class AttentionUNET(nn.Module):
             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)  # Get attention map
-            attention_maps.append(att_map)  # Store attention map
+            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)
             x = self.ups[idx+1](concat_skip)
 
         return self.final_conv(x), attention_maps
 
-def download_model():
-    """Download your trained model from HuggingFace"""
+def download_and_load_model():
+    """Download and load model once at startup"""
+    global model
+    print("Loading Attention U-Net model...")
+    
     model_url = "https://huggingface.co/spaces/ArchCoder/the-op-segmenter/resolve/main/best_attention_model.pth.tar"
     model_path = "best_attention_model.pth.tar"
     
+    # Download model if needed
     if not os.path.exists(model_path):
-        print("📥 Downloading your trained model...")
+        print("Downloading model weights...")
         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
-    return model_path
-
-def load_your_attention_model():
-    """Load YOUR trained Attention U-Net model"""
-    global model
-    if model is None:
-        try:
-            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("✅ Your Attention U-Net model loaded successfully!")
         except Exception as e:
-            print(f"❌ Error loading your model: {e}")
-            model = None
-    return model
-
-def download_dataset():
-    """Download and extract the dataset using kagglehub"""
-    dataset_path = kagglehub.dataset_download('nikhilroxtomar/brain-tumor-segmentation')
+            print(f"Failed to download model: {e}")
+            return False
     
-    # 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)
+    # Load model
+    try:
+        model = AttentionUNET(in_channels=1, out_channels=1).to(device)
+        checkpoint = torch.load(model_path, map_location=device, weights_only=True)
+        model.load_state_dict(checkpoint["state_dict"])
+        model.eval()
+        print("✓ Model loaded successfully!")
+        return True
+    except Exception as e:
+        print(f"Failed to load model: {e}")
+        return False
+
+def download_and_load_dataset():
+    """Download and load entire dataset once at startup"""
+    global dataset_images, dataset_masks, dataset_loaded
     
-    images_path = os.path.join(extracted_path, 'images')
-    masks_path = os.path.join(extracted_path, 'masks')
+    if dataset_loaded:
+        return True
+        
+    print("Loading brain tumor dataset...")
     
-    return images_path, masks_path
+    try:
+        # Download dataset using kagglehub - returns directory path
+        dataset_path = kagglehub.dataset_download('nikhilroxtomar/brain-tumor-segmentation')
+        print(f"Dataset downloaded to: {dataset_path}")
+        
+        # Find images and masks directories
+        images_dir = os.path.join(dataset_path, 'images')
+        masks_dir = os.path.join(dataset_path, 'masks')
+        
+        # If direct path doesn't exist, search subdirectories
+        if not os.path.exists(images_dir):
+            # Search for images and masks directories
+            for root, dirs, files in os.walk(dataset_path):
+                if 'images' in dirs:
+                    images_dir = os.path.join(root, 'images')
+                if 'masks' in dirs:
+                    masks_dir = os.path.join(root, 'masks')
+        
+        if not os.path.exists(images_dir) or not os.path.exists(masks_dir):
+            print("Could not find images/masks directories. Searching all files...")
+            # Fallback: find all image files
+            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:
+            # Load image and mask file paths
+            dataset_images = glob(os.path.join(images_dir, "*.*"))
+            dataset_masks = glob(os.path.join(masks_dir, "*.*"))
+        
+        dataset_images = sorted(dataset_images)
+        dataset_masks = sorted(dataset_masks)
+        
+        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 load_random_sample():
-    """Load a random image and mask from the dataset"""
-    images_path, masks_path = download_dataset()
+def get_random_sample():
+    """Get a random image and corresponding mask from dataset"""
+    if not dataset_loaded:
+        return None, None, "Dataset not loaded"
     
-    image_files = [f for f in os.listdir(images_path) if f.endswith(('.png', '.jpg'))]
-    if not image_files:
+    if not dataset_images:
         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)
+    # Get random index
+    idx = random.randint(0, len(dataset_images) - 1)
+    img_path = dataset_images[idx]
     
-    image = Image.open(img_path).convert("L")
-    mask = Image.open(mask_path).convert("L") if os.path.exists(mask_path) else None
+    # Find corresponding mask
+    img_name = os.path.basename(img_path)
+    mask_path = None
+    for mask in dataset_masks:
+        if os.path.basename(mask) == img_name:
+            mask_path = mask
+            break
     
-    return image, mask, random_file
+    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}"
 
-def preprocess_for_your_model(image):
-    """Preprocessing exactly like your Colab code"""
+def preprocess_for_model(image):
+    """Preprocessing for your model"""
     if image.mode != 'L':
         image = image.convert('L')
     
-    val_test_transform = transforms.Compose([
+    transform = transforms.Compose([
         transforms.Resize((256,256)),
         transforms.ToTensor()
     ])
     
-    return val_test_transform(image).unsqueeze(0)
+    return transform(image).unsqueeze(0)
 
 def apply_tta(model, input_tensor):
-    """Test-Time Augmentation: Apply augmentations and average predictions"""
+    """Test-Time Augmentation"""
     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
+        lambda x: TF.vflip(x),  # Vertical flip
     ]
     
     predictions = []
-    for aug in augmentations:
+    for i, aug in enumerate(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, _ = model(aug_input)
+        pred = torch.sigmoid(pred)
+        
+        # Reverse augmentation
+        if i == 1:  # Reverse hflip
             pred = TF.hflip(pred)
-        elif aug == augmentations[4]:  # Reverse vflip
+        elif i == 2:  # Reverse vflip
             pred = TF.vflip(pred)
+            
         predictions.append(pred)
     
-    # Average predictions
-    avg_pred = torch.mean(torch.stack(predictions), dim=0)
-    return avg_pred
+    return torch.mean(torch.stack(predictions), dim=0)
 
 def generate_attention_heatmap(attention_maps):
-    """Generate combined attention heatmap"""
+    """Generate attention heatmap"""
     if not attention_maps:
-        return np.zeros((256, 256))
+        return np.zeros((256, 256, 3))
     
-    # Average attention maps from different levels
+    # Combine attention maps
     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, ground_truth=None, filename=None):
-    current_model = load_your_attention_model()
+def analyze_image(image, ground_truth, filename):
+    """Main analysis function"""
+    if model is None:
+        return None, "Model not loaded. Please restart the application."
     
-    if current_model is None:
-        return None, "Failed to load your trained model."
-
     if image is None:
-        return None, "Please upload or load an image first."
+        return None, "Please select an image first."
     
     try:
         # Preprocess
-        input_tensor = preprocess_for_your_model(image).to(device)
+        input_tensor = preprocess_for_model(image).to(device)
         
         # Apply TTA
-        avg_pred = apply_tta(current_model, input_tensor)
+        with torch.no_grad():
+            avg_pred = apply_tta(model, input_tensor)
+            _, attention_maps = model(input_tensor)
         
         # Get binary mask
-        binary_mask = (avg_pred > 0.5).float().squeeze().cpu().numpy()
+        binary_mask = (avg_pred > 0.5).squeeze().cpu().numpy()
         
         # Post-processing
-        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5,5))
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3,3))
         binary_mask = cv2.morphologyEx(binary_mask.astype(np.uint8), cv2.MORPH_OPEN, kernel)
         binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel)
         
-        # Extract attention maps
-        _, attention_maps = current_model(input_tensor)
+        # Generate attention heatmap
         att_heatmap = generate_attention_heatmap(attention_maps)
         
         # Create visualization
-        fig, axes = plt.subplots(2, 3, figsize=(18, 12))
-        fig.suptitle('Brain Tumor Segmentation Analysis', fontsize=20)
+        if ground_truth is not None:
+            fig, axes = plt.subplots(2, 3, figsize=(15, 10))
+        else:
+            fig, axes = plt.subplots(2, 2, figsize=(12, 10))
+            
+        fig.suptitle('Brain Tumor Segmentation Analysis', fontsize=16, weight='bold')
         
-        # Original
+        # Original image
         axes[0,0].imshow(image, cmap='gray')
-        axes[0,0].set_title('Original Image')
+        axes[0,0].set_title('Original Image', fontsize=12, weight='bold')
         axes[0,0].axis('off')
         
-        # 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')
+        # Attention heatmap
+        axes[0,1].imshow(image, cmap='gray')
+        axes[0,1].imshow(att_heatmap, alpha=0.4)
+        axes[0,1].set_title('Attention Heatmap', fontsize=12, weight='bold')
         axes[0,1].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)
+        # Predicted mask
         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[0,2].imshow(binary_mask, cmap='gray')
+            axes[0,2].set_title('Predicted Mask', fontsize=12, weight='bold')
+            axes[0,2].axis('off')
+            
+            # Ground truth
+            gt_array = np.array(ground_truth.resize((256, 256)))
+            axes[1,0].imshow(gt_array, cmap='gray')
+            axes[1,0].set_title('Ground Truth Mask', fontsize=12, weight='bold')
             axes[1,0].axis('off')
             
-            # Comparison Overlay
-            overlay = np.array(image.convert('RGB'))
+            # Overlay comparison
+            overlay = np.array(image.convert('RGB').resize((256, 256)))
             overlay[binary_mask > 0] = [0, 255, 0]  # Green for prediction
-            overlay[gt_np > 0] = [255, 0, 0]  # Red for ground truth
+            overlay[gt_array > 128] = [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].set_title('Prediction (Green) vs GT (Red)', fontsize=12, weight='bold')
             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
+            # Calculate IoU
+            pred_binary = binary_mask > 0
+            gt_binary = gt_array > 128
+            intersection = np.sum(pred_binary & gt_binary)
+            union = np.sum(pred_binary | gt_binary)
             iou = intersection / (union + 1e-8)
             
-            axes[1,2].text(0.1, 0.5, f'IoU Score: {iou:.4f}', fontsize=20)
+            # Dice score
+            dice = (2 * intersection) / (np.sum(pred_binary) + np.sum(gt_binary) + 1e-8)
+            
+            axes[1,2].text(0.1, 0.6, f'IoU: {iou:.4f}', fontsize=16, weight='bold')
+            axes[1,2].text(0.1, 0.4, f'Dice: {dice:.4f}', fontsize=16, weight='bold')
+            axes[1,2].set_xlim(0, 1)
+            axes[1,2].set_ylim(0, 1)
             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].imshow(binary_mask, cmap='gray')
+            axes[1,0].set_title('Predicted Mask', fontsize=12, weight='bold')
             axes[1,0].axis('off')
             
+            # Overlay
+            overlay = np.array(image.convert('RGB').resize((256, 256)))
+            overlay[binary_mask > 0] = [255, 0, 0]
+            axes[1,1].imshow(overlay)
+            axes[1,1].set_title('Prediction Overlay', fontsize=12, weight='bold')
             axes[1,1].axis('off')
-            axes[1,2].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)
@@ -328,128 +376,176 @@ def predict_tumor(image, ground_truth=None, filename=None):
         
         result_image = Image.open(buf)
         
-        # Statistics
+        # Generate analysis text
         tumor_pixels = np.sum(binary_mask)
         total_pixels = binary_mask.size
         tumor_percentage = (tumor_pixels / total_pixels) * 100
         
         analysis_text = f"""
-## Brain Tumor Segmentation Results
+# Analysis Results
+
+**File:** {filename if filename else 'Uploaded Image'}
 
-### Detection Summary
-- Tumor Percentage: {tumor_percentage:.2f}%
-- Tumor Pixels: {tumor_pixels}
-- File: {filename if filename else 'Uploaded Image'}
+**Tumor Detection:**
+- Tumor Area: {tumor_percentage:.2f}%
+- Tumor Pixels: {tumor_pixels:,}
 
-### Model Information
-- Your Attention U-Net Model
+**Model Features:**
 - Test-Time Augmentation: Applied
-- Attention Visualization: Included
-        """
+- Attention Visualization: Generated
+- Post-processing: Morphological cleanup
+"""
         
         if ground_truth is not None:
-            analysis_text += f"\n- IoU with Ground Truth: {iou:.4f}"
+            analysis_text += f"""
+**Performance Metrics:**
+- IoU Score: {iou:.4f}
+- Dice Score: {dice:.4f}
+"""
         
         return result_image, analysis_text
         
     except Exception as e:
-        return None, f"Error: {str(e)}"
+        return None, f"Analysis failed: {str(e)}"
+
+# Initialize model and 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!")
 
-def clear_all():
-    return None, None, None, "Upload or load an image for analysis"
+print("Application ready!")
 
-# Professional CSS (white, clean, professional)
+# Professional CSS
 css = """
 .gradio-container {
-    max-width: 1400px !important;
+    max-width: 1600px !important;
     margin: auto !important;
-    background-color: white !important;
-    font-family: 'Arial', sans-serif !important;
+    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
 }
-h1, h2, h3, h4 {
-    color: #333333 !important;
+.gr-button {
+    border-radius: 6px !important;
+    font-weight: 500 !important;
 }
-button {
-    background-color: #f0f0f0 !important;
-    color: #333333 !important;
-    border: 1px solid #dddddd !important;
-    border-radius: 4px !important;
+.gr-button-primary {
+    background: #2563eb !important;
+    border-color: #2563eb !important;
 }
-button.primary {
-    background-color: #007bff !important;
-    color: white !important;
+.gr-button-secondary {
+    background: #6b7280 !important;
+    border-color: #6b7280 !important;
 }
-.output-image {
-    border: 1px solid #dddddd !important;
-    border-radius: 4px !important;
+h1, h2, h3 {
+    color: #1f2937 !important;
 }
-.markdown {
-    line-height: 1.6 !important;
-    color: #555555 !important;
+.gr-form {
+    border: 1px solid #e5e7eb !important;
+    border-radius: 8px !important;
 }
 """
 
-# Create professional Gradio interface
-with gr.Blocks(css=css, title="Brain Tumor Segmentation Application") as app:
+# Create Gradio interface
+with gr.Blocks(css=css, title="Brain Tumor Segmentation Analysis") as app:
     
     gr.Markdown("""
     # Brain Tumor Segmentation Using Attention U-Net
-    A professional tool for medical image analysis
+    
+    **Advanced Medical Image Analysis Tool**
+    
+    Features: Test-Time Augmentation, Attention Visualization, Dataset Integration
     """)
     
+    # Status display
+    with gr.Row():
+        with gr.Column():
+            status_text = f"Model Status: {'✓ Loaded' if model_loaded else '✗ Failed'} | Dataset Status: {'✓ Loaded' if dataset_loaded_success else '✗ Failed'}"
+            if dataset_loaded_success:
+                status_text += f" | Images: {len(dataset_images)} | Masks: {len(dataset_masks)}"
+            gr.Markdown(f"**{status_text}**")
+    
     with gr.Row():
         with gr.Column(scale=1):
             gr.Markdown("### Input Selection")
             
-            image_input = gr.Image(
-                label="Upload Brain MRI",
+            # Image display
+            image_display = gr.Image(
+                label="Selected Image",
                 type="pil",
-                sources=["upload", "webcam"],
                 height=300
             )
             
-            load_random_btn = gr.Button("Load Random Sample from Dataset", variant="primary")
-            
+            # Control buttons
             with gr.Row():
-                analyze_btn = gr.Button("Analyze Image", variant="primary", scale=2)
-                clear_btn = gr.Button("Clear", scale=1)
+                load_sample_btn = gr.Button("Load Random Sample", variant="primary", scale=1)
+                upload_btn = gr.UploadButton("Upload Image", file_types=["image"], scale=1)
+            
+            analyze_btn = gr.Button("Analyze Image", variant="primary", size="lg")
+            
+            # Dataset info
+            gr.Markdown(f"""
+            **Dataset Information:**
+            - Total Images: {len(dataset_images) if dataset_loaded_success else 'N/A'}
+            - Total Masks: {len(dataset_masks) if dataset_loaded_success else 'N/A'}
+            - Source: nikhilroxtomar/brain-tumor-segmentation
+            """)
         
         with gr.Column(scale=2):
             gr.Markdown("### Analysis Results")
             
-            output_image = gr.Image(
-                label="Segmentation Results",
+            result_display = gr.Image(
+                label="Segmentation Analysis",
                 type="pil",
-                height=400
+                height=500
             )
             
-            analysis_output = gr.Markdown(
-                value="Select an input method to begin analysis."
+            analysis_text = gr.Markdown(
+                value="Load an image and click 'Analyze Image' to begin."
             )
     
-    # Hidden state for ground truth and filename
-    ground_truth_state = gr.State()
-    filename_state = gr.State()
-
+    # Hidden states
+    current_ground_truth = gr.State()
+    current_filename = gr.State()
+    
     # Event handlers
-    analyze_btn.click(
-        fn=predict_tumor,
-        inputs=[image_input, ground_truth_state, filename_state],
-        outputs=[output_image, analysis_output]
+    def handle_sample_load():
+        image, mask, filename = get_random_sample()
+        return image, mask, filename
+    
+    def handle_upload(file):
+        if file is not None:
+            image = Image.open(file.name).convert("L")
+            return image, None, os.path.basename(file.name)
+        return None, None, ""
+    
+    load_sample_btn.click(
+        fn=handle_sample_load,
+        outputs=[image_display, current_ground_truth, current_filename]
     )
     
-    load_random_btn.click(
-        fn=load_random_sample,
-        inputs=[],
-        outputs=[image_input, ground_truth_state, filename_state, analysis_output]
+    upload_btn.upload(
+        fn=handle_upload,
+        inputs=[upload_btn],
+        outputs=[image_display, current_ground_truth, current_filename]
     )
     
-    clear_btn.click(
-        fn=clear_all,
-        inputs=[],
-        outputs=[image_input, output_image, ground_truth_state, analysis_output]
+    analyze_btn.click(
+        fn=analyze_image,
+        inputs=[image_display, current_ground_truth, current_filename],
+        outputs=[result_display, analysis_text]
     )
 
 if __name__ == "__main__":
-    print("Starting Brain Tumor Segmentation Application...")
-    app.launch()
+    print("\n" + "="*50)
+    print("LAUNCHING BRAIN TUMOR SEGMENTATION APPLICATION")
+    print("="*50)
+    
+    app.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True,
+        share=False
+    )