Update app.py

app.py

@@ -78,6 +78,18 @@ def get_fused_image(img, pred_mask, view, alpha=0.8):
         rotated = cv2.flip(cv2.rotate(fused, cv2.ROTATE_90_COUNTERCLOCKWISE), 1)
         return rotated
 
+# Define the inference function
+def inference(learn, reorder, resample, org_img, input_img, org_size):
+    """Perform segmentation using the loaded model."""
+    # Perform the segmentation
+    with torch.no_grad():
+        pred = learn.predict(input_img)
+    
+    # Process the prediction if necessary
+    mask_data = pred[0]  # Assuming the first element of the prediction is the mask
+    
+    return mask_data
+
 # Function for Gradio image segmentation
 def gradio_image_segmentation(fileobj, learn, reorder, resample, save_dir, view):
     """Predict function using the learner and other resources."""
@@ -112,7 +124,7 @@ def gradio_image_segmentation(fileobj, learn, reorder, resample, save_dir, view)
 
     mask_data = inference(learn, reorder=reorder, resample=resample,
                           org_img=org_img, input_img=input_img,
-                          org_size=org_size).data
+                          org_size=org_size)
     
     if "".join(org_img.orientation) == "LSA":        
         mask_data = mask_data.permute(0,1,3,2)