Update app.py

app.py

@@ -18,17 +18,31 @@ live_preds = None
 
 def rotate_btn_fn(img_path, xt, yt, zt, add_bone_cmap=False):
     global current_img
-    
-    angles = (xt, yt, zt)
-    out_img_path = f'data/cached_outputs/{os.path.basename(img_path)[:-4]}_{angles}.png'
-    out_img = skimage.io.imread(out_img_path)
-    if not add_bone_cmap:
+
+    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}")
+
+        out_img_path = f'data/cached_outputs/{os.path.basename(img_path)[:-4]}_{angles}.png'
+        if not os.path.exists(out_img_path):
+            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
+
+        cmap = plt.get_cmap('bone')
+        out_img = cmap(out_img)
+        out_img = (out_img[..., :3] * 255).astype(np.uint8)
+        current_img = out_img
         return out_img
-    cmap = plt.get_cmap('bone')
-    out_img = cmap(out_img)
-    out_img = (out_img[..., :3] * 255).astype(np.uint8)
-    current_img = out_img
-    return out_img
+
+    except Exception as e:
+        print(f"Error in rotate_btn_fn: {e}")
+        return None
     
 css_style = "./style.css"
 callback = gr.CSVLogger()
@@ -40,7 +54,7 @@ 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', source='upload', interactive=True, elem_classes='imgs')
+            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')
         with gr.Row():
             with gr.Column(scale=0.25):
@@ -70,8 +84,8 @@ with gr.Blocks(css=css_style, title="RadRotator") as app:
 try:
     app.close()
     gr.close_all()
-except:
-    pass
+except Exception as e:
+    print(f"Error closing app: {e}")
 
 demo = app.launch(
     max_threads=4,