Add spleen segmentation app with MONAI model

app.py

@@ -0,0 +1,169 @@
+import gradio as gr
+import spaces
+import torch
+import nibabel as nib
+import numpy as np
+from huggingface_hub import hf_hub_download
+from monai.transforms import (
+    Compose,
+    LoadImage,
+    EnsureChannelFirst,
+    ScaleIntensity,
+    Resize,
+    AsDiscrete,
+)
+from monai.networks.nets import UNet
+import tempfile
+import os
+
+# Load the model
+model = None
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def load_model():
+    global model
+    if model is None:
+        # Download model from HuggingFace
+        model_path = hf_hub_download(
+            repo_id="MONAI/example_spleen_segmentation",
+            filename="models/model.pt"
+        )
+        
+        # Initialize UNet architecture
+        model = UNet(
+            spatial_dims=3,
+            in_channels=1,
+            out_channels=2,
+            channels=(16, 32, 64, 128, 256),
+            strides=(2, 2, 2, 2),
+            num_res_units=2,
+        )
+        
+        # Load weights
+        checkpoint = torch.load(model_path, map_location=device)
+        model.load_state_dict(checkpoint)
+        model.to(device)
+        model.eval()
+    return model
+
+@spaces.GPU
+def segment_spleen(input_file):
+    """Segment spleen from CT NIfTI file"""
+    try:
+        # Load model
+        net = load_model()
+        
+        # Load NIfTI file
+        img = nib.load(input_file)
+        img_data = img.get_fdata()
+        
+        # Preprocessing
+        img_tensor = torch.from_numpy(img_data).float().unsqueeze(0).unsqueeze(0)
+        
+        # Normalize
+        img_tensor = (img_tensor - img_tensor.min()) / (img_tensor.max() - img_tensor.min())
+        
+        # Resize to model input size (96x96x96)
+        img_resized = torch.nn.functional.interpolate(
+            img_tensor,
+            size=(96, 96, 96),
+            mode="trilinear",
+            align_corners=True
+        )
+        
+        # Move to device and run inference
+        img_resized = img_resized.to(device)
+        
+        with torch.no_grad():
+            output = net(img_resized)
+            pred = torch.argmax(output, dim=1)
+        
+        # Resize back to original size
+        pred_resized = torch.nn.functional.interpolate(
+            pred.float().unsqueeze(0),
+            size=img_data.shape,
+            mode="nearest"
+        )
+        
+        pred_np = pred_resized.squeeze().cpu().numpy().astype(np.uint8)
+        
+        # Save segmentation as NIfTI
+        seg_img = nib.Nifti1Image(pred_np, img.affine, img.header)
+        output_path = tempfile.mktemp(suffix="_segmentation.nii.gz")
+        nib.save(seg_img, output_path)
+        
+        # Create visualization (middle slice)
+        mid_slice = img_data.shape[2] // 2
+        img_slice = img_data[:, :, mid_slice]
+        seg_slice = pred_np[:, :, mid_slice]
+        
+        # Normalize image for display
+        img_slice = (img_slice - img_slice.min()) / (img_slice.max() - img_slice.min()) * 255
+        
+        # Create overlay
+        overlay = np.stack([img_slice, img_slice, img_slice], axis=-1).astype(np.uint8)
+        overlay[seg_slice == 1] = [255, 0, 0]  # Red for spleen
+        
+        return overlay, output_path, "Segmentation completed successfully!"
+        
+    except Exception as e:
+        return None, None, f"Error: {str(e)}"
+
+# Create Gradio interface
+with gr.Blocks(title="Spleen Segmentation") as demo:
+    gr.Markdown("# 🏥 CT Spleen Segmentation")
+    gr.Markdown(
+        """Upload a CT scan in NIfTI format (.nii or .nii.gz) to segment the spleen using the 
+        [MONAI/example_spleen_segmentation](https://huggingface.co/MONAI/example_spleen_segmentation) model.
+        
+        **Model Info:**
+        - Architecture: UNet
+        - Input: 3D CT image (96×96×96)
+        - Output: Binary segmentation (spleen vs background)
+        - Mean Dice Score: 0.96
+        
+        **Instructions:**
+        1. Upload a NIfTI file (.nii or .nii.gz)
+        2. Click Submit
+        3. View the segmentation overlay and download the result
+        """
+    )
+    
+    with gr.Row():
+        with gr.Column():
+            input_file = gr.File(
+                label="Upload CT Scan (NIfTI format)",
+                file_types=[".nii", ".nii.gz"]
+            )
+            submit_btn = gr.Button("Segment Spleen", variant="primary")
+        
+        with gr.Column():
+            output_image = gr.Image(label="Segmentation Overlay (Middle Slice)", type="numpy")
+            output_file = gr.File(label="Download Segmentation")
+            status_text = gr.Textbox(label="Status")
+    
+    submit_btn.click(
+        fn=segment_spleen,
+        inputs=[input_file],
+        outputs=[output_image, output_file, status_text]
+    )
+    
+    gr.Markdown(
+        """### Requirements
+        - MONAI
+        - PyTorch
+        - nibabel
+        - numpy
+        - huggingface_hub
+        
+        ### Citation
+        If you use this model, please cite:
+        ```
+        Xia, Yingda, et al. "3D Semi-Supervised Learning with Uncertainty-Aware Multi-View Co-Training." 
+        arXiv preprint arXiv:1811.12506 (2018).
+        ```
+        """
+    )
+
+if __name__ == "__main__":
+    demo.launch()