Update app.py

app.py

@@ -1,214 +1,217 @@
-import gradio as gr
-import torch
-import nibabel as nib
-import numpy as np
-from monai.networks.nets import SwinUNETR
-from monai.inferers import sliding_window_inference
-from monai.transforms import (
-    Compose, LoadImage, EnsureChannelFirst, Orientation,
-    Spacing, ScaleIntensityRange, CropForeground, Resize, EnsureType
-)
-import tempfile
-import matplotlib.pyplot as plt
-from io import BytesIO
-from PIL import Image
-import os
-
-# ============================================================================
-# MODEL INITIALIZATION
-# ============================================================================
-print(" Initializing model...")
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f" Using device: {device}")
-
-model = SwinUNETR(
-    in_channels=1,
-    out_channels=2,
-    patch_size=(2, 2, 2),
-    depths=(2, 2, 2, 2),
-    num_heads=(3, 6, 12, 24),
-    window_size=7,
-    feature_size=48,
-    norm_name="instance",
-    use_checkpoint=True,
-    spatial_dims=3,
-).to(device)
-
-# Load trained weights
-model_path = "best_metric_model.pth"
-if os.path.exists(model_path):
-    model.load_state_dict(torch.load(model_path, map_location=device))
-    print(" Model loaded successfully!")
-else:
-    print(" Model file not found! Please upload best_metric_model.pth")
-
-model.eval()
-
-# ============================================================================
-# PREPROCESSING PIPELINE
-# ============================================================================
-transforms = Compose([
-    LoadImage(image_only=True),
-    EnsureChannelFirst(),
-    Orientation(axcodes="RAS"),
-    Spacing(pixdim=(1.5, 1.5, 1.5), mode="bilinear"),
-    ScaleIntensityRange(a_min=-200, a_max=200, b_min=0.0, b_max=1.0, clip=True),
-    CropForeground(source_key="image"),
-    Resize(spatial_size=(128, 128, 128), mode="trilinear"),
-    EnsureType(dtype=torch.float32),
-])
-
-# ============================================================================
-# SEGMENTATION FUNCTION
-# ============================================================================
-def segment_liver(nifti_file, show_slice):
-    """
-    Segment liver from uploaded NIfTI file
-    
-    Args:
-        nifti_file: Uploaded .nii.gz file
-        show_slice: Which slice to visualize (0-127)
-    
-    Returns:
-        Visualization image and downloadable segmentation file
-    """
-    if nifti_file is None:
-        return None, None
-    
-    try:
-        print(f" Processing file: {nifti_file.name}")
-        
-        # Preprocess
-        volume = transforms(nifti_file.name)
-        volume_input = volume.unsqueeze(0).to(device)
-        
-        # Inference with sliding window
-        print(" Running inference...")
-        with torch.no_grad():
-            output = sliding_window_inference(
-                volume_input, 
-                roi_size=(128, 128, 64), 
-                sw_batch_size=2, 
-                predictor=model, 
-                overlap=0.75
-            )
-            prediction = torch.argmax(output, dim=1, keepdim=True)
-        
-        # Convert to numpy
-        vol_np = volume.cpu().numpy()
-        pred_np = prediction.cpu().numpy()[0, 0]
-        
-        # Normalize volume for display
-        vol_display = (vol_np - vol_np.min()) / (vol_np.max() - vol_np.min() + 1e-8)
-        
-        # Clamp slice index
-        max_slice = vol_np.shape[2] - 1
-        slice_idx = min(max(0, show_slice), max_slice)
-        
-        # Create visualization (3 views)
-        fig, axes = plt.subplots(1, 3, figsize=(18, 6))
-        
-        # CT scan
-        axes[0].imshow(vol_display[:, :, slice_idx].T, cmap='gray', origin='lower')
-        axes[0].set_title(f'CT Scan (Slice {slice_idx}/{max_slice})', fontsize=14)
-        axes[0].axis('off')
-        
-        # Segmentation only
-        axes[1].imshow(pred_np[:, :, slice_idx].T, cmap='Reds', origin='lower')
-        axes[1].set_title('Liver Segmentation', fontsize=14)
-        axes[1].axis('off')
-        
-        # Overlay
-        axes[2].imshow(vol_display[:, :, slice_idx].T, cmap='gray', origin='lower')
-        axes[2].imshow(pred_np[:, :, slice_idx].T, cmap='Greens', alpha=0.5, origin='lower')
-        axes[2].set_title('Overlay (Green = Liver)', fontsize=14)
-        axes[2].axis('off')
-        
-        plt.tight_layout()
-        
-        # Convert plot to image
-        buf = BytesIO()
-        plt.savefig(buf, format='png', dpi=100, bbox_inches='tight')
-        buf.seek(0)
-        result_image = Image.open(buf)
-        plt.close()
-        
-        # Save segmentation as NIfTI
-        output_file = tempfile.NamedTemporaryFile(delete=False, suffix='.nii.gz')
-        
-        # Load original NIfTI to preserve metadata
-        original_nii = nib.load(nifti_file.name)
-        
-        # Create segmentation NIfTI with original affine
-        pred_nifti = nib.Nifti1Image(pred_np.astype(np.uint8), affine=original_nii.affine)
-        nib.save(pred_nifti, output_file.name)
-        
-        print("Segmentation complete!")
-        
-        return result_image, output_file.name
-    
-    except Exception as e:
-        print(f" Error: {str(e)}")
-        return None, None
-
-# ============================================================================
-# GRADIO INTERFACE
-# ============================================================================
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        # AI-Powered Liver Segmentation
-        
-        Upload a liver CT scan in NIfTI format to get automatic liver segmentation using deep learning.
-        
-        **Model**: SwinUNETR (Vision Transformer) trained on Medical Segmentation Decathlon  
-        **Performance**: Mean Dice Score = **95.5% ± 2.0%** on test set  
-        **Input**: NIfTI files (.nii or .nii.gz)
-        """
-    )
-    
-    with gr.Row():
-        with gr.Column():
-            file_input = gr.File(
-                label=" Upload CT Scan (NIfTI format)", 
-                file_types=[".nii", ".nii.gz"]
-            )
-            slice_slider = gr.Slider(
-                minimum=0, 
-                maximum=127, 
-                value=64, 
-                step=1, 
-                label="🔍 Select Slice to Visualize"
-            )
-            segment_btn = gr.Button("🚀 Segment Liver", variant="primary", size="lg")
-        
-        with gr.Column():
-            output_image = gr.Image(label=" Segmentation Result")
-            output_file = gr.File(label=" Download Segmentation (NIfTI)")
-    
-    gr.Markdown(
-        """
-        ---
-        ### How to Use:
-        1. Upload a liver CT scan in NIfTI format (.nii.gz)
-        2. Adjust the slice slider to view different cross-sections
-        3. Click "Segment Liver" to run the AI model
-        4. Download the segmentation mask for further analysis
-        
-        ### ⚠️ Note:
-        - This is a research prototype, not for clinical use
-        - Processing time: ~10-30 seconds depending on GPU availability
-        - Best results with contrast-enhanced CT scans
-        """
-    )
-    
-    # Connect button to function
-    segment_btn.click(
-        fn=segment_liver,
-        inputs=[file_input, slice_slider],
-        outputs=[output_image, output_file]
-    )
-
-# Launch the app
-if __name__ == "__main__":
-    demo.launch()
+import gradio as gr
+import torch
+import nibabel as nib
+import numpy as np
+from monai.networks.nets import SwinUNETR
+from monai.inferers import sliding_window_inference
+from monai.transforms import (
+    Compose, LoadImage, EnsureChannelFirst, Orientation,
+    Spacing, ScaleIntensityRange, CropForeground, Resize, EnsureType
+)
+import tempfile
+import matplotlib.pyplot as plt
+from io import BytesIO
+from PIL import Image
+import os
+
+# ============================================================================
+# MODEL INITIALIZATION
+# ============================================================================
+print(" Initializing model...")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f" Using device: {device}")
+
+# Initialize SwinUNETR with EXACT same config as your training
+model = SwinUNETR(
+    in_channels=1,
+    out_channels=2,
+    depths=(2, 2, 2, 2),
+    num_heads=(3, 6, 12, 24),
+    feature_size=48,
+    norm_name="instance",
+    drop_rate=0.0,
+    attn_drop_rate=0.0,
+    dropout_path_rate=0.0,
+    normalize=True,
+    use_checkpoint=True,
+    spatial_dims=3,
+).to(device)
+
+# Load trained weights
+model_path = "best_metric_model.pth"
+if os.path.exists(model_path):
+    model.load_state_dict(torch.load(model_path, map_location=device, weights_only=False))
+    print(" Model loaded successfully!")
+else:
+    print(" Model file not found! Please upload best_metric_model.pth")
+
+model.eval()
+
+# ============================================================================
+# PREPROCESSING PIPELINE
+# ============================================================================
+transforms = Compose([
+    LoadImage(image_only=True),
+    EnsureChannelFirst(),
+    Orientation(axcodes="RAS"),
+    Spacing(pixdim=(1.5, 1.5, 1.5), mode="bilinear"),
+    ScaleIntensityRange(a_min=-200, a_max=200, b_min=0.0, b_max=1.0, clip=True),
+    CropForeground(source_key="image"),
+    Resize(spatial_size=(128, 128, 128), mode="trilinear"),
+    EnsureType(dtype=torch.float32),
+])
+
+# ============================================================================
+# SEGMENTATION FUNCTION
+# ============================================================================
+def segment_liver(nifti_file, show_slice):
+    """
+    Segment liver from uploaded NIfTI file
+    
+    Args:
+        nifti_file: Uploaded .nii.gz file
+        show_slice: Which slice to visualize (0-127)
+    
+    Returns:
+        Visualization image and downloadable segmentation file
+    """
+    if nifti_file is None:
+        return None, None
+    
+    try:
+        print(f" Processing file: {nifti_file.name}")
+        
+        # Preprocess
+        volume = transforms(nifti_file.name)
+        volume_input = volume.unsqueeze(0).to(device)
+        
+        # Inference with sliding window
+        print(" Running inference...")
+        with torch.no_grad():
+            output = sliding_window_inference(
+                volume_input, 
+                roi_size=(96, 96, 96),  # Smaller ROI for CPU
+                sw_batch_size=1,        # Batch size 1 for CPU
+                predictor=model, 
+                overlap=0.5
+            )
+            prediction = torch.argmax(output, dim=1, keepdim=True)
+        
+        # Convert to numpy
+        vol_np = volume.cpu().numpy()
+        pred_np = prediction.cpu().numpy()[0, 0]
+        
+        # Normalize volume for display
+        vol_display = (vol_np - vol_np.min()) / (vol_np.max() - vol_np.min() + 1e-8)
+        
+        # Clamp slice index
+        max_slice = vol_np.shape[2] - 1
+        slice_idx = min(max(0, show_slice), max_slice)
+        
+        # Create visualization (3 views)
+        fig, axes = plt.subplots(1, 3, figsize=(18, 6))
+        
+        # CT scan
+        axes[0].imshow(vol_display[:, :, slice_idx].T, cmap='gray', origin='lower')
+        axes[0].set_title(f'CT Scan (Slice {slice_idx}/{max_slice})', fontsize=14)
+        axes[0].axis('off')
+        
+        # Segmentation only
+        axes[1].imshow(pred_np[:, :, slice_idx].T, cmap='Reds', origin='lower')
+        axes[1].set_title('Liver Segmentation', fontsize=14)
+        axes[1].axis('off')
+        
+        # Overlay
+        axes[2].imshow(vol_display[:, :, slice_idx].T, cmap='gray', origin='lower')
+        axes[2].imshow(pred_np[:, :, slice_idx].T, cmap='Greens', alpha=0.5, origin='lower')
+        axes[2].set_title('Overlay (Green = Liver)', fontsize=14)
+        axes[2].axis('off')
+        
+        plt.tight_layout()
+        
+        # Convert plot to image
+        buf = BytesIO()
+        plt.savefig(buf, format='png', dpi=100, bbox_inches='tight')
+        buf.seek(0)
+        result_image = Image.open(buf)
+        plt.close()
+        
+        # Save segmentation as NIfTI
+        output_file = tempfile.NamedTemporaryFile(delete=False, suffix='.nii.gz')
+        
+        # Create segmentation NIfTI
+        pred_nifti = nib.Nifti1Image(pred_np.astype(np.uint8), affine=np.eye(4))
+        nib.save(pred_nifti, output_file.name)
+        
+        print(" Segmentation complete!")
+        
+        return result_image, output_file.name
+    
+    except Exception as e:
+        print(f"❌ Error: {str(e)}")
+        import traceback
+        traceback.print_exc()
+        return None, None
+
+# ============================================================================
+# GRADIO INTERFACE
+# ============================================================================
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # 🫀 AI-Powered Liver Segmentation
+        
+        Upload a liver CT scan in NIfTI format to get automatic liver segmentation using deep learning.
+        
+        **Model**: SwinUNETR (Vision Transformer) trained on Medical Segmentation Decathlon  
+        **Performance**: Mean Dice Score = **95.5% ± 2.0%** on test set  
+        **Input**: NIfTI files (.nii or .nii.gz)
+        """
+    )
+    
+    with gr.Row():
+        with gr.Column():
+            file_input = gr.File(
+                label=" Upload CT Scan (NIfTI format)", 
+                file_types=[".nii", ".nii.gz"]
+            )
+            slice_slider = gr.Slider(
+                minimum=0, 
+                maximum=127, 
+                value=64, 
+                step=1, 
+                label=" Select Slice to Visualize"
+            )
+            segment_btn = gr.Button(" Segment Liver", variant="primary", size="lg")
+        
+        with gr.Column():
+            output_image = gr.Image(label=" Segmentation Result")
+            output_file = gr.File(label=" Download Segmentation (NIfTI)")
+    
+    gr.Markdown(
+        """
+        ---
+        ###  How to Use:
+        1. Upload a liver CT scan in NIfTI format (.nii.gz)
+        2. Adjust the slice slider to view different cross-sections
+        3. Click "Segment Liver" to run the AI model
+        4. Download the segmentation mask for further analysis
+        
+        ### ⚠️ Note:
+        - This is a research prototype, not for clinical use
+        - Processing time: ~30-60 seconds on CPU
+        - Best results with contrast-enhanced CT scans
+        """
+    )
+    
+    # Connect button to function
+    segment_btn.click(
+        fn=segment_liver,
+        inputs=[file_input, slice_slider],
+        outputs=[output_image, output_file]
+    )
+
+# Launch the app
+if __name__ == "__main__":
+    demo.launch()
+