Update app.py

app.py

@@ -5,7 +5,7 @@ import numpy as np
 from monai.networks.nets import SwinUNETR
 from monai.inferers import sliding_window_inference
 from monai.transforms import (
-    Compose, LoadImage, EnsureChannelFirstd, Orientationd, 
+    Compose, LoadImaged, EnsureChannelFirstd, Orientationd, 
     Spacingd, ScaleIntensityRanged, CropForegroundd, Resized, EnsureTyped
 )
 import matplotlib.pyplot as plt
@@ -15,7 +15,7 @@ import os
 import tempfile
 
 print("Loading model...")
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device("cpu") # Keep CPU for stability
 print(f"Device: {device}")
 
 # Model config
@@ -27,33 +27,38 @@ model = SwinUNETR(
     spatial_dims=3,
 ).to(device)
 
-try:
-    model.load_state_dict(torch.load("best_metric_model.pth", map_location=device, weights_only=False))
-    print("Model loaded!")
-except Exception as e:
-    print(f"Error loading model: {e}")
+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))
+        print("Model loaded!")
+    except Exception as e:
+        print(f"Error loading model: {e}")
+else:
+    print("WARNING: Model file not found")
 
 model.eval()
 
-# CORRECTED TRANSFORMS
-# Note: Using Resized instead of Resize
+# --- CORRECTED TRANSFORMS ---
+# Using LoadImaged (dictionary version) to match the rest of the pipeline
 test_transforms = Compose([
-    LoadImage(image_only=True),
+    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), # Fix deprecated warning
-    Resized(keys=["image"], spatial_size=(128, 128, 128)),  # Fix: Resized instead of Resize
+    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),
 ])
 
-def segment_liver(file_path, slice_num=64):
-    """Simplified segmentation function"""
+def segment_liver(file_obj, slice_num=64):
     try:
-        if file_path is None:
+        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
         print(f"Processing: {file_path}")
         
         # Load and preprocess
@@ -66,7 +71,7 @@ def segment_liver(file_path, slice_num=64):
         # Inference
         with torch.no_grad():
             outputs = sliding_window_inference(
-                volume, (96, 96, 96), 1, model, overlap=0.5
+                volume, (96, 96, 96), 1, model, overlap=0.25
             )
             pred = torch.argmax(outputs, dim=1).float()
         
@@ -77,7 +82,10 @@ def segment_liver(file_path, slice_num=64):
         # Normalize volume
         vol_display = (vol_np - vol_np.min()) / (vol_np.max() - vol_np.min() + 1e-8)
         
-        slice_idx = min(int(slice_num), vol_np.shape[2] - 1)
+        # Slice selection
+        z_dim = vol_np.shape[2]
+        slice_idx = int(slice_num)
+        if slice_idx >= z_dim: slice_idx = z_dim // 2
         
         # Plot
         fig, axes = plt.subplots(1, 3, figsize=(15, 5))
@@ -119,23 +127,13 @@ def segment_liver(file_path, slice_num=64):
         print(f"Error: {e}")
         import traceback
         traceback.print_exc()
-        
-        fig, ax = plt.subplots(figsize=(10, 6))
-        ax.text(0.5, 0.5, f'Error: {str(e)}', 
-                ha='center', va='center', fontsize=12, color='red')
-        ax.axis('off')
-        buf = BytesIO()
-        plt.savefig(buf, format='png')
-        buf.seek(0)
-        img = Image.open(buf)
-        plt.close()
-        return img, None
+        return None, None
 
-# Simple Gradio interface
+# Interface
 iface = gr.Interface(
     fn=segment_liver,
     inputs=[
-        gr.File(label="Upload NIfTI"), # Removed file_types constraint
+        gr.File(label="Upload NIfTI"),
         gr.Slider(0, 127, value=64, label="Slice")
     ],
     outputs=[
@@ -143,9 +141,8 @@ iface = gr.Interface(
         gr.File(label="Download Mask")
     ],
     title="Liver Segmentation (95.5% Dice)",
-    description="Upload file from TestVolumes folder (like liver_104.nii.gz)"
+    description="Upload file from TestVolumes folder"
 )
 
-
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch(server_name="0.0.0.0", server_port=7860)