Update app.py

app.py

@@ -23,39 +23,41 @@ from monai.transforms import (
     EnsureTyped,
 )
 
-print("Loading model...")
-# HF Spaces: assume CPU, GPU not guaranteed
+print("Starting app...")
+
+# ----------------- DEVICE -----------------
 device = torch.device("cpu")
-print(f"Device: {device}")
+print(f"Using device: {device}")
 
-# --------- MODEL DEFINITION (must match training) ----------
+# ----------------- MODEL -----------------
+# NOTE: SwinUNETR in current MONAI versions does NOT take `patch_size` or `window_size`.
+# Use img_size consistent with your pre-processing (Resized to 128x128x64).
 model = SwinUNETR(
+    img_size=(128, 128, 64),
     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):
+ckpt_path = "best_metric_model.pth"
+if os.path.exists(ckpt_path):
     try:
-        state = torch.load(model_path, map_location=device)
+        state = torch.load(ckpt_path, map_location=device)
         model.load_state_dict(state)
-        print("Model loaded!")
+        print("Model loaded successfully.")
     except Exception as e:
-        print(f"Error loading model: {e}")
+        print(f"ERROR loading model weights: {e}")
 else:
-    print("WARNING: Model file not found")
+    print(f"WARNING: checkpoint '{ckpt_path}' not found in Space.")
 
 model.eval()
 
-# --------- PREPROCESSING (mirror training pipeline) ----------
+# ----------------- TRANSFORMS -----------------
 test_transforms = Compose(
     [
         LoadImaged(keys=["image"]),
@@ -79,44 +81,75 @@ test_transforms = Compose(
 
 def _get_path_from_gradio_file(file_obj):
     """
-    Gradio / HF can pass:
-      - dict with "name"
-      - tempfile-like object with .name
-      - plain string path (local)
+    Convert the Gradio file object into a real path on disk.
+    Handles dicts, tempfiles, and plain string paths.
     """
     if file_obj is None:
         return None
 
+    # Case 1: dict (HF Spaces often passes this)
     if isinstance(file_obj, dict):
-        return file_obj.get("name")
+        if "path" in file_obj and file_obj["path"]:
+            return file_obj["path"]
+        if "name" in file_obj and file_obj["name"]:
+            return file_obj["name"]
+        # If we only have raw bytes, write them to a temp file
+        if "data" in file_obj and file_obj["data"] is not None:
+            suffix = ".nii.gz"
+            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
+            tmp.write(file_obj["data"])
+            tmp.flush()
+            tmp.close()
+            return tmp.name
+
+    # Case 2: tempfile-like with .name
     if hasattr(file_obj, "name"):
         return file_obj.name
+
+    # Case 3: already a string path (local testing)
     if isinstance(file_obj, str):
         return file_obj
+
     raise ValueError(f"Unsupported file object type: {type(file_obj)}")
 
 
-# --------- INFERENCE FUNCTION ----------
+def _error_image(msg: str):
+    """
+    Create a simple image with an error message so the UI
+    never looks 'empty' when something goes wrong.
+    """
+    fig, ax = plt.subplots(figsize=(8, 3))
+    ax.text(0.5, 0.5, msg, ha="center", va="center", color="red", fontsize=12)
+    ax.axis("off")
+    buf = BytesIO()
+    fig.savefig(buf, format="png", bbox_inches="tight")
+    buf.seek(0)
+    img = np.array(Image.open(buf))
+    plt.close(fig)
+    return img
+
+
+# ----------------- INFERENCE -----------------
 def segment_liver(file_obj, slice_num=64):
     try:
         if file_obj is None:
-            return None, None
+            return _error_image("No file uploaded."), None
 
         file_path = _get_path_from_gradio_file(file_obj)
-        print(f"Processing: {file_path}")
+        print(f"[segment_liver] file_path = {file_path}")
 
         if file_path is None or not os.path.exists(file_path):
-            raise FileNotFoundError("Uploaded file path not found")
+            raise FileNotFoundError("Uploaded file path not found on server.")
 
-        # Manual extension validation (since we removed file_types)
+        # Manual extension check
         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)  # [1,1,H,W,D]
-        print(f"Input shape: {volume.shape}")
+        volume = data_dict["image"].unsqueeze(0).to(device)  # [1, 1, H, W, D]
+        print(f"[segment_liver] preprocessed volume shape: {volume.shape}")
 
         # Inference
         with torch.no_grad():
@@ -127,7 +160,7 @@ def segment_liver(file_obj, slice_num=64):
                 predictor=model,
                 overlap=0.25,
             )
-            pred = torch.argmax(outputs, dim=1).float()  # [1,H,W,D]
+            pred = torch.argmax(outputs, dim=1).float()  # [1, H, W, D]
 
         vol_np = volume[0, 0].cpu().numpy()
         pred_np = pred[0].cpu().numpy()
@@ -135,23 +168,21 @@ def segment_liver(file_obj, slice_num=64):
         # Normalize CT for display
         vol_display = (vol_np - vol_np.min()) / (vol_np.max() - vol_np.min() + 1e-8)
 
-        # Slice selection
+        # Handle slice index safely
         z_dim = vol_np.shape[2]
         idx = int(slice_num)
-        if idx < 0:
-            idx = 0
-        if idx >= z_dim:
+        if idx < 0 or idx >= z_dim:
             idx = z_dim // 2
 
-        # Plot
+        # Plot CT / mask / overlay
         fig, axes = plt.subplots(1, 3, figsize=(15, 5))
 
         axes[0].imshow(vol_display[:, :, idx], cmap="gray")
-        axes[0].set_title("CT Scan")
+        axes[0].set_title("CT Slice")
         axes[0].axis("off")
 
         axes[1].imshow(pred_np[:, :, idx], cmap="Reds", vmin=0, vmax=1)
-        axes[1].set_title("Liver Prediction")
+        axes[1].set_title("Predicted Liver Mask")
         axes[1].axis("off")
 
         axes[2].imshow(vol_display[:, :, idx], cmap="gray")
@@ -161,29 +192,30 @@ def segment_liver(file_obj, slice_num=64):
 
         plt.tight_layout()
 
-        # Convert figure to numpy image for Gradio
+        # Convert figure to numpy image
         buf = BytesIO()
         fig.savefig(buf, format="png", bbox_inches="tight")
         buf.seek(0)
         img = np.array(Image.open(buf))
         plt.close(fig)
 
-        # Save prediction as NIfTI for download
+        # Save prediction mask as NIfTI for download
         pred_nii = nib.Nifti1Image(pred_np.astype(np.uint8), np.eye(4))
         out_path = tempfile.mktemp(suffix=".nii.gz")
         nib.save(pred_nii, out_path)
 
-        print("Success!")
+        print("[segment_liver] success.")
         return img, out_path
 
     except Exception as e:
-        print(f"Error in segment_liver: {e}")
         import traceback
+
+        print("[segment_liver] ERROR:", e)
         traceback.print_exc()
-        return None, None
+        return _error_image(f"Error: {e}"), None
 
 
-# --------- GRADIO INTERFACE ----------
+# ----------------- GRADIO UI -----------------
 iface = gr.Interface(
     fn=segment_liver,
     inputs=[
@@ -199,4 +231,5 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
+    # On HF Spaces: iface.launch(server_name=\"0.0.0.0\", server_port=7860)
     iface.launch()