Update app.py

app.py

@@ -94,12 +94,22 @@ def gradio_image_segmentation(fileobj, learn, reorder, resample, save_dir, view)
     save_path = save_dir / save_fn
     
     # Ensure only_tensor is set to False to get all values
-    org_img, input_img, org_size = med_img_reader(img_path_str, 
-                                                  reorder=reorder,
-                                                  resample=resample,
-                                                  only_tensor=False,
-                                                  dtype=torch.Tensor)
-    
+    try:
+        org_img, input_img, org_size = med_img_reader(img_path_str, 
+                                                      reorder=reorder,
+                                                      resample=resample,
+                                                      only_tensor=False,
+                                                      dtype=torch.Tensor)
+    except ValueError:
+        # Handle the case where med_img_reader returns only two values
+        org_img, input_img = med_img_reader(img_path_str, 
+                                           reorder=reorder,
+                                           resample=resample,
+                                           only_tensor=False,
+                                           dtype=torch.Tensor)
+        # Infer org_size from org_img
+        org_size = org_img.shape[1:]  # Assuming org_img has a shape attribute
+
     mask_data = inference(learn, reorder=reorder, resample=resample,
                           org_img=org_img, input_img=input_img,
                           org_size=org_size).data