Update app.py

app.py

@@ -1,12 +1,11 @@
 # Archive paper: https://arxiv.org/abs/2404.13000
 
 import os
-
 import gradio as gr
 import matplotlib.pyplot as plt
 import numpy as np
 import skimage
-import monai as mn
+from skimage import io
 import torch
 
 from io_utils import LoadImageD
@@ -18,7 +17,6 @@ live_preds = None
 
 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}")
@@ -33,10 +31,13 @@ def rotate_btn_fn(img, xt, yt, zt, add_bone_cmap=False):
 
         out_img_path = f'data/cached_outputs/{os.path.basename(input_img_path)[:-4]}_{angles}.png'
 
+        # Assuming you have precomputed outputs for the example images
         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}")
+            # Perform your rotation here if the precomputed image doesn't exist
+            # For now, let's just return the input image for demonstration purposes
+            out_img = img
 
         if not add_bone_cmap:
             return out_img
@@ -50,7 +51,7 @@ def rotate_btn_fn(img, xt, yt, zt, add_bone_cmap=False):
     except Exception as e:
         print(f"Error in rotate_btn_fn: {e}")
         return None
-    
+
 css_style = "./style.css"
 callback = gr.CSVLogger()