Update app.py

app.py

@@ -81,6 +81,10 @@ def get_fused_image(img, pred_mask, view, alpha=0.8):
 # Define the inference function
 def inference(learn, reorder, resample, org_img, input_img, org_size):
     """Perform segmentation using the loaded model."""
+    # Ensure input_img is a torch.Tensor
+    if not isinstance(input_img, torch.Tensor):
+        raise ValueError(f"Expected input_img to be a torch.Tensor, but got {type(input_img)}")
+    
     # Perform the segmentation
     with torch.no_grad():
         pred = learn.predict(input_img)
@@ -122,6 +126,10 @@ def gradio_image_segmentation(fileobj, learn, reorder, resample, save_dir, view)
         # Infer org_size from org_img
         org_size = org_img.shape[1:]  # Assuming org_img has a shape attribute
 
+    # Ensure input_img is a torch.Tensor
+    if not isinstance(input_img, torch.Tensor):
+        raise ValueError(f"Expected input_img to be a torch.Tensor, but got {type(input_img)}")
+    
     mask_data = inference(learn, reorder=reorder, resample=resample,
                           org_img=org_img, input_img=input_img,
                           org_size=org_size)