Update app.py

app.py

@@ -1,36 +1,52 @@
+import os
+import tempfile
+from io import BytesIO
+
 import gradio as gr
 import torch
 import nibabel as nib
 import numpy as np
+import matplotlib.pyplot as plt
+from PIL import Image
+
 from monai.networks.nets import SwinUNETR
 from monai.inferers import sliding_window_inference
 from monai.transforms import (
-    Compose, LoadImaged, EnsureChannelFirstd, Orientationd, 
-    Spacingd, ScaleIntensityRanged, CropForegroundd, Resized, EnsureTyped
+    Compose,
+    LoadImaged,
+    EnsureChannelFirstd,
+    Orientationd,
+    Spacingd,
+    ScaleIntensityRanged,
+    CropForegroundd,
+    Resized,
+    EnsureTyped,
 )
-import matplotlib.pyplot as plt
-from io import BytesIO
-from PIL import Image
-import os
-import tempfile
 
 print("Loading model...")
-device = torch.device("cpu") # Keep CPU for stability
+# HF Spaces: assume CPU, GPU not guaranteed
+device = torch.device("cpu")
 print(f"Device: {device}")
 
-# Model config
+# --------- MODEL DEFINITION (must match training) ----------
 model = SwinUNETR(
-    img_size=(128, 128, 128),
     in_channels=1,
     out_channels=2,
+    patch_size=2,
+    depths=(2, 2, 2, 2),
+    num_heads=(3, 6, 12, 24),
+    window_size=7,
     feature_size=48,
+    norm_name="instance",
+    use_checkpoint=False,
     spatial_dims=3,
 ).to(device)
 
 model_path = "best_metric_model.pth"
 if os.path.exists(model_path):
     try:
-        model.load_state_dict(torch.load(model_path, map_location=device, weights_only=False))
+        state = torch.load(model_path, map_location=device)
+        model.load_state_dict(state)
         print("Model loaded!")
     except Exception as e:
         print(f"Error loading model: {e}")
@@ -39,110 +55,148 @@ else:
 
 model.eval()
 
-# --- CORRECTED TRANSFORMS ---
-# Using LoadImaged (dictionary version) to match the rest of the pipeline
-test_transforms = Compose([
-    LoadImaged(keys=["image"]),  # FIXED: Dictionary transform
-    EnsureChannelFirstd(keys=["image"]),
-    Orientationd(keys=["image"], axcodes="RAS"),
-    Spacingd(keys=["image"], pixdim=(1.5, 1.5, 1.5), mode="bilinear"),
-    ScaleIntensityRanged(keys=["image"], a_min=-200, a_max=200, b_min=0.0, b_max=1.0, clip=True),
-    CropForegroundd(keys=["image"], source_key="image", allow_smaller=False),
-    Resized(keys=["image"], spatial_size=(128, 128, 128)),
-    EnsureTyped(keys=["image"], dtype=torch.float32),
-])
+# --------- PREPROCESSING (mirror training pipeline) ----------
+test_transforms = Compose(
+    [
+        LoadImaged(keys=["image"]),
+        EnsureChannelFirstd(keys=["image"]),
+        Orientationd(keys=["image"], axcodes="RAS"),
+        Spacingd(keys=["image"], pixdim=(1.5, 1.5, 1.0), mode="bilinear"),
+        ScaleIntensityRanged(
+            keys=["image"],
+            a_min=-200,
+            a_max=200,
+            b_min=0.0,
+            b_max=1.0,
+            clip=True,
+        ),
+        CropForegroundd(keys=["image"], source_key="image", allow_smaller=False),
+        Resized(keys=["image"], spatial_size=(128, 128, 64)),
+        EnsureTyped(keys=["image"], dtype=torch.float32),
+    ]
+)
+
 
+def _get_path_from_gradio_file(file_obj):
+    """
+    Gradio / HF can pass:
+      - dict with "name"
+      - tempfile-like object with .name
+      - plain string path (local)
+    """
+    if file_obj is None:
+        return None
+
+    if isinstance(file_obj, dict):
+        return file_obj.get("name")
+    if hasattr(file_obj, "name"):
+        return file_obj.name
+    if isinstance(file_obj, str):
+        return file_obj
+    raise ValueError(f"Unsupported file object type: {type(file_obj)}")
+
+
+# --------- INFERENCE FUNCTION ----------
 def segment_liver(file_obj, slice_num=64):
     try:
         if file_obj is None:
             return None, None
-            
-        # Handle Gradio file object
-        file_path = file_obj.name if hasattr(file_obj, "name") else file_obj
+
+        file_path = _get_path_from_gradio_file(file_obj)
         print(f"Processing: {file_path}")
-        
-        # Load and preprocess
+
+        if file_path is None or not os.path.exists(file_path):
+            raise FileNotFoundError("Uploaded file path not found")
+
+        # Manual extension validation (since we removed file_types)
+        if not (file_path.endswith(".nii") or file_path.endswith(".nii.gz")):
+            raise ValueError("Invalid file type. Please upload a .nii or .nii.gz NIfTI file.")
+
+        # Preprocess
         data_dict = {"image": file_path}
         data_dict = test_transforms(data_dict)
-        volume = data_dict["image"].unsqueeze(0).to(device)
-        
+        volume = data_dict["image"].unsqueeze(0).to(device)  # [1,1,H,W,D]
         print(f"Input shape: {volume.shape}")
-        
+
         # Inference
         with torch.no_grad():
             outputs = sliding_window_inference(
-                volume, (96, 96, 96), 1, model, overlap=0.25
+                volume,
+                roi_size=(96, 96, 96),
+                sw_batch_size=1,
+                predictor=model,
+                overlap=0.25,
             )
-            pred = torch.argmax(outputs, dim=1).float()
-        
-        # Visualization
+            pred = torch.argmax(outputs, dim=1).float()  # [1,H,W,D]
+
         vol_np = volume[0, 0].cpu().numpy()
-        pred_np = pred[0, 0].cpu().numpy()
-        
-        # Normalize volume
+        pred_np = pred[0].cpu().numpy()
+
+        # Normalize CT for display
         vol_display = (vol_np - vol_np.min()) / (vol_np.max() - vol_np.min() + 1e-8)
-        
+
         # Slice selection
         z_dim = vol_np.shape[2]
-        slice_idx = int(slice_num)
-        if slice_idx >= z_dim: slice_idx = z_dim // 2
-        
+        idx = int(slice_num)
+        if idx < 0:
+            idx = 0
+        if idx >= z_dim:
+            idx = z_dim // 2
+
         # Plot
         fig, axes = plt.subplots(1, 3, figsize=(15, 5))
-        
-        # CT
-        axes[0].imshow(vol_display[:, :, slice_idx], cmap='gray')
-        axes[0].set_title('CT Scan')
-        axes[0].axis('off')
-        
-        # Prediction
-        axes[1].imshow(pred_np[:, :, slice_idx], cmap='Reds', vmin=0, vmax=1)
-        axes[1].set_title('Liver Prediction')
-        axes[1].axis('off')
-        
-        # Overlay
-        axes[2].imshow(vol_display[:, :, slice_idx], cmap='gray')
-        axes[2].imshow(pred_np[:, :, slice_idx], cmap='Greens', alpha=0.5, vmin=0, vmax=1)
-        axes[2].set_title('Overlay')
-        axes[2].axis('off')
-        
+
+        axes[0].imshow(vol_display[:, :, idx], cmap="gray")
+        axes[0].set_title("CT Scan")
+        axes[0].axis("off")
+
+        axes[1].imshow(pred_np[:, :, idx], cmap="Reds", vmin=0, vmax=1)
+        axes[1].set_title("Liver Prediction")
+        axes[1].axis("off")
+
+        axes[2].imshow(vol_display[:, :, idx], cmap="gray")
+        axes[2].imshow(pred_np[:, :, idx], cmap="Greens", alpha=0.5, vmin=0, vmax=1)
+        axes[2].set_title("Overlay")
+        axes[2].axis("off")
+
         plt.tight_layout()
-        
-        # Save image
+
+        # Convert figure to numpy image for Gradio
         buf = BytesIO()
-        plt.savefig(buf, format='png', bbox_inches='tight')
+        fig.savefig(buf, format="png", bbox_inches="tight")
         buf.seek(0)
-        img = Image.open(buf)
-        plt.close()
-        
-        # Save prediction as NIfTI
+        img = np.array(Image.open(buf))
+        plt.close(fig)
+
+        # Save prediction as NIfTI for download
         pred_nii = nib.Nifti1Image(pred_np.astype(np.uint8), np.eye(4))
-        out_path = tempfile.mktemp(suffix='.nii.gz')
+        out_path = tempfile.mktemp(suffix=".nii.gz")
         nib.save(pred_nii, out_path)
-        
+
         print("Success!")
         return img, out_path
-        
+
     except Exception as e:
-        print(f"Error: {e}")
+        print(f"Error in segment_liver: {e}")
         import traceback
         traceback.print_exc()
         return None, None
 
-# Interface
+
+# --------- GRADIO INTERFACE ----------
 iface = gr.Interface(
     fn=segment_liver,
     inputs=[
-        gr.File(label="Upload NIfTI"),
-        gr.Slider(0, 127, value=64, label="Slice")
+        gr.File(label="Upload NIfTI volume (.nii or .nii.gz)"),
+        gr.Slider(0, 127, value=64, label="Slice index"),
     ],
     outputs=[
-        gr.Image(label="Result"),
-        gr.File(label="Download Mask")
+        gr.Image(label="Result", type="numpy"),
+        gr.File(label="Download Mask (.nii.gz)"),
     ],
-    title="Liver Segmentation (95.5% Dice)",
-    description="Upload file from TestVolumes folder"
+    title="Liver Segmentation (SwinUNETR, MONAI)",
+    description="Upload a 3D liver CT volume (.nii or .nii.gz). The app runs a SwinUNETR model trained on MSD Task03 Liver.",
 )
 
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+    iface.launch()