Update app.py

app.py

@@ -8,32 +8,35 @@ import io
 import base64
 from torchvision import transforms
 import torch.nn.functional as F
+import warnings
+warnings.filterwarnings("ignore")
+
+# Global variable to store model
+model = None
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Load the pretrained model
-@gr.utils.cache
 def load_model():
     """Load the pretrained brain segmentation model"""
-    try:
-        model = torch.hub.load(
-            'mateuszbuda/brain-segmentation-pytorch',
-            'unet',
-            in_channels=3,
-            out_channels=1,
-            init_features=32,
-            pretrained=True,
-            force_reload=False
-        )
-        model.eval()
-        return model
-    except Exception as e:
-        print(f"Error loading model: {e}")
-        return None
-
-# Initialize model
-model = load_model()
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-if model:
-    model = model.to(device)
+    global model
+    if model is None:
+        try:
+            print("Loading brain segmentation model...")
+            model = torch.hub.load(
+                'mateuszbuda/brain-segmentation-pytorch',
+                'unet',
+                in_channels=3,
+                out_channels=1,
+                init_features=32,
+                pretrained=True,
+                force_reload=False
+            )
+            model.eval()
+            model = model.to(device)
+            print("Model loaded successfully!")
+        except Exception as e:
+            print(f"Error loading model: {e}")
+            model = None
+    return model
 
 def preprocess_image(image):
     """Preprocess the input image for the model"""
@@ -45,7 +48,12 @@ def preprocess_image(image):
         image = image.convert('RGB')
 
     # Resize to 256x256 (model's expected input size)
-    image = image.resize((256, 256), Image.Resampling.LANCZOS)
+    # Use LANCZOS if available, otherwise use BILINEAR
+    try:
+        image = image.resize((256, 256), Image.Resampling.LANCZOS)
+    except AttributeError:
+        # For older PIL versions
+        image = image.resize((256, 256), Image.LANCZOS)
 
     # Convert to tensor and normalize
     transform = transforms.Compose([
@@ -73,20 +81,24 @@ def create_overlay_visualization(original_img, mask, alpha=0.6):
 
 def predict_tumor(image):
     """Main prediction function"""
-    if model is None:
-        return None, "❌ Model failed to load. Please try again."
+    # Load model if not loaded
+    current_model = load_model()
+    
+    if current_model is None:
+        return None, "❌ Model failed to load. Please try again later."
 
     if image is None:
         return None, "⚠️ Please upload an image first."
 
     try:
+        print("Processing image...")
         # Preprocess the image
         input_tensor, original_img = preprocess_image(image)
         input_tensor = input_tensor.to(device)
 
         # Make prediction
         with torch.no_grad():
-            prediction = model(input_tensor)
+            prediction = current_model(input_tensor)
             # Apply sigmoid to get probability map
             prediction = torch.sigmoid(prediction)
             # Convert to numpy
@@ -109,24 +121,25 @@ def predict_tumor(image):
 
         # 4. Side-by-side comparison
         fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+        fig.suptitle('Brain Tumor Segmentation Results', fontsize=16, fontweight='bold')
 
         axes[0].imshow(original_array)
-        axes[0].set_title('Original Image', fontsize=14, fontweight='bold')
+        axes[0].set_title('Original Image', fontsize=12, fontweight='bold')
         axes[0].axis('off')
 
         axes[1].imshow(mask_colored)
-        axes[1].set_title('Tumor Segmentation', fontsize=14, fontweight='bold')
+        axes[1].set_title('Tumor Segmentation', fontsize=12, fontweight='bold')
         axes[1].axis('off')
 
         axes[2].imshow(overlay)
-        axes[2].set_title('Overlay (Red = Tumor)', fontsize=14, fontweight='bold')
+        axes[2].set_title('Overlay (Red = Tumor)', fontsize=12, fontweight='bold')
         axes[2].axis('off')
 
         plt.tight_layout()
 
         # Save plot to bytes
         buf = io.BytesIO()
-        plt.savefig(buf, format='png', dpi=150, bbox_inches='tight')
+        plt.savefig(buf, format='png', dpi=150, bbox_inches='tight', facecolor='white')
         buf.seek(0)
         plt.close()
 
@@ -140,38 +153,41 @@ def predict_tumor(image):
 
         # Create analysis report
         analysis_text = f"""
-        ## 🧠 Brain Tumor Segmentation Analysis
-
-        **📊 Tumor Statistics:**
-        - Total pixels analyzed: {total_pixels:,}
-        - Tumor pixels detected: {tumor_pixels:,}
-        - Tumor area percentage: {tumor_percentage:.2f}%
-
-        **🎯 Model Performance:**
-        - Model: U-Net with attention mechanism
-        - Input resolution: 256×256 pixels
-        - Detection threshold: {threshold}
-
-        **⚠️ Medical Disclaimer:**
-        This is an AI tool for research purposes only.
-        Always consult qualified medical professionals for diagnosis.
+## 🧠 Brain Tumor Segmentation Analysis
+
+**📊 Tumor Statistics:**
+- Total pixels analyzed: {total_pixels:,}
+- Tumor pixels detected: {tumor_pixels:,}
+- Tumor area percentage: {tumor_percentage:.2f}%
+
+**🎯 Model Information:**
+- Model: Pre-trained U-Net for brain segmentation
+- Input resolution: 256×256 pixels  
+- Detection threshold: {threshold}
+- Device: {device.type.upper()}
+
+**⚠️ Medical Disclaimer:**
+This is an AI tool for research and educational purposes only. 
+Always consult qualified medical professionals for diagnosis.
         """
 
+        print("Processing completed successfully!")
         return result_image, analysis_text
 
     except Exception as e:
         error_msg = f"❌ Error during prediction: {str(e)}"
+        print(error_msg)
         return None, error_msg
 
 def clear_all():
     """Clear all inputs and outputs"""
-    return None, None, ""
+    return None, None, "Upload an image and click 'Analyze Image' to see results."
 
 # Custom CSS for better styling
 css = """
-#main-container {
-    max-width: 1200px;
-    margin: 0 auto;
+.gradio-container {
+    max-width: 1200px !important;
+    margin: auto !important;
 }
 #title {
     text-align: center;
@@ -181,21 +197,21 @@ css = """
     border-radius: 10px;
     margin-bottom: 20px;
 }
-#upload-box {
-    border: 2px dashed #ccc;
-    border-radius: 10px;
-    padding: 20px;
-    text-align: center;
-    margin: 10px 0;
-}
 .output-image {
     border-radius: 10px;
     box-shadow: 0 4px 8px rgba(0,0,0,0.1);
 }
+button {
+    border-radius: 8px;
+    font-weight: 500;
+}
+.progress-bar {
+    background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
+}
 """
 
 # Create Gradio interface
-with gr.Blocks(css=css, title="Brain Tumor Segmentation") as app:
+with gr.Blocks(css=css, title="🧠 Brain Tumor Segmentation AI", theme=gr.themes.Soft()) as app:
 
     # Header
     gr.HTML("""
@@ -207,34 +223,34 @@ with gr.Blocks(css=css, title="Brain Tumor Segmentation") as app:
 
     with gr.Row():
         with gr.Column(scale=1):
-            gr.HTML("<h3>📤 Input Image</h3>")
+            gr.Markdown("### 📤 Input Image")
 
             # Image input with camera option
             image_input = gr.Image(
                 label="Upload Brain MRI Scan",
                 type="pil",
-                sources=["upload", "webcam"],  # Allow both upload and camera
+                sources=["upload", "webcam"],
                 height=300
             )
 
             with gr.Row():
-                predict_btn = gr.Button("🔍 Analyze Image", variant="primary", size="lg")
-                clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
+                predict_btn = gr.Button("🔍 Analyze Image", variant="primary", scale=2)
+                clear_btn = gr.Button("🗑️ Clear", variant="secondary", scale=1)
 
             gr.HTML("""
-            <div style="margin-top: 20px; padding: 15px; background-color: #f0f8ff; border-radius: 8px;">
+            <div style="margin-top: 20px; padding: 15px; background-color: #f0f8ff; border-radius: 8px; border-left: 4px solid #667eea;">
                 <h4>📋 Instructions:</h4>
-                <ul>
+                <ul style="margin: 10px 0; padding-left: 20px;">
                     <li>Upload a brain MRI scan image</li>
                     <li>Supported formats: PNG, JPG, JPEG</li>
                     <li>For best results, use clear, high-contrast MRI images</li>
-                    <li>You can also use the camera to capture an image from your device</li>
+                    <li>Camera option available for mobile devices</li>
                 </ul>
             </div>
             """)
 
         with gr.Column(scale=2):
-            gr.HTML("<h3>📊 Segmentation Results</h3>")
+            gr.Markdown("### 📊 Segmentation Results")
 
             # Output image
             output_image = gr.Image(
@@ -246,26 +262,31 @@ with gr.Blocks(css=css, title="Brain Tumor Segmentation") as app:
 
             # Analysis text
             analysis_output = gr.Markdown(
-                label="Analysis Report",
-                value="Upload an image and click 'Analyze Image' to see results."
+                value="Upload an image and click 'Analyze Image' to see results.",
+                elem_id="analysis"
             )
 
     # Add footer with information
     gr.HTML("""
-    <div style="margin-top: 30px; padding: 20px; background-color: #f9f9f9; border-radius: 10px;">
-        <h4>🔬 About This Tool</h4>
-        <p><strong>Model:</strong> Pre-trained U-Net architecture optimized for brain tumor segmentation</p>
-        <p><strong>Technology:</strong> PyTorch, Deep Learning, Computer Vision</p>
-        <p><strong>Dataset:</strong> Trained on medical MRI brain scans</p>
-
-        <h4>⚠️ Important Medical Disclaimer</h4>
-        <p style="color: #d73027; font-weight: bold;">
-        This AI tool is for research and educational purposes only. It should NOT be used for medical diagnosis.
-        Always consult qualified healthcare professionals for medical advice and diagnosis.
-        </p>
-
-        <p style="text-align: center; margin-top: 20px; color: #666;">
-        Made with ❤️ using Gradio • Powered by PyTorch • Hosted on 🤗 Hugging Face Spaces
+    <div style="margin-top: 30px; padding: 20px; background-color: #f9f9f9; border-radius: 10px; border: 1px solid #e1e4e8;">
+        <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 20px;">
+            <div>
+                <h4 style="color: #667eea; margin-bottom: 10px;">🔬 About This Tool</h4>
+                <p><strong>Model:</strong> Pre-trained U-Net for brain segmentation</p>
+                <p><strong>Technology:</strong> PyTorch + Deep Learning</p>
+                <p><strong>Purpose:</strong> Research & Educational Use</p>
+            </div>
+            <div>
+                <h4 style="color: #d73027; margin-bottom: 10px;">⚠️ Medical Disclaimer</h4>
+                <p style="color: #d73027; font-weight: 500;">
+                This AI tool is for research and educational purposes only.<br>
+                <strong>NOT for medical diagnosis.</strong> Always consult healthcare professionals.
+                </p>
+            </div>
+        </div>
+        <hr style="margin: 20px 0; border: none; border-top: 1px solid #e1e4e8;">
+        <p style="text-align: center; color: #666; margin: 10px 0;">
+            Made with ❤️ using Gradio • Powered by PyTorch • Hosted on 🤗 Hugging Face Spaces
         </p>
     </div>
     """)
@@ -274,26 +295,22 @@ with gr.Blocks(css=css, title="Brain Tumor Segmentation") as app:
     predict_btn.click(
         fn=predict_tumor,
         inputs=[image_input],
-        outputs=[output_image, analysis_output]
+        outputs=[output_image, analysis_output],
+        show_progress=True
     )
 
     clear_btn.click(
         fn=clear_all,
+        inputs=[],
         outputs=[image_input, output_image, analysis_output]
     )
 
-    # Auto-predict when image is uploaded
-    image_input.change(
-        fn=predict_tumor,
-        inputs=[image_input],
-        outputs=[output_image, analysis_output]
-    )
-
 # Launch the app
 if __name__ == "__main__":
+    print("Starting Brain Tumor Segmentation App...")
     app.launch(
-        share=True,
         server_name="0.0.0.0",
         server_port=7860,
-        show_error=True
+        show_error=True,
+        share=False
     )