Update app.py

app.py

@@ -16,21 +16,28 @@ from io_utils import LoadImageD
 current_img = None
 live_preds = None
 
-def rotate_btn_fn(img_path, xt, yt, zt, add_bone_cmap=False):
+def rotate_btn_fn(img, xt, yt, zt, add_bone_cmap=False):
     global current_img
 
     try:
         angles = (xt, yt, zt)
         print(f"Rotating with angles: {angles}")
 
-        if not os.path.exists(img_path):
-            raise FileNotFoundError(f"Input image path does not exist: {img_path}")
+        if isinstance(img, np.ndarray):
+            input_img_path = "uploaded_image.png"
+            skimage.io.imsave(input_img_path, img)
+        elif isinstance(img, str) and os.path.exists(img):
+            input_img_path = img
+        else:
+            raise ValueError("Invalid input image")
 
-        out_img_path = f'data/cached_outputs/{os.path.basename(img_path)[:-4]}_{angles}.png'
-        if not os.path.exists(out_img_path):
+        out_img_path = f'data/cached_outputs/{os.path.basename(input_img_path)[:-4]}_{angles}.png'
+
+        if os.path.exists(out_img_path):
+            out_img = skimage.io.imread(out_img_path)
+        else:
             raise FileNotFoundError(f"Output image path does not exist: {out_img_path}")
 
-        out_img = skimage.io.imread(out_img_path)
         if not add_bone_cmap:
             return out_img
 
@@ -54,8 +61,8 @@ with gr.Blocks(css=css_style, title="RadRotator") as app:
     
     with gr.TabItem("Demo"):
         with gr.Row():
-            input_img = gr.Image(type='filepath', label='Input image', interactive=True, elem_classes='imgs')
-            output_img = gr.Image(type='pil', label='Output image', interactive=False, elem_classes='imgs')
+            input_img = gr.Image(type='numpy', label='Input image', interactive=True, elem_classes='imgs')
+            output_img = gr.Image(type='numpy', label='Output image', interactive=False, elem_classes='imgs')
         with gr.Row():
             with gr.Column(scale=0.25):
                 pass