bug fix

app.py

@@ -41,7 +41,7 @@ def show_image(img, title):
     plt.title(title)
 
 # Visualize a Single Image
-def visualize_single_image(image_path, image_name, model, device):
+def visualize_single_image(image_path):
     img = load_image(image_path, transform).to(device)
     
     # Run inference
@@ -56,46 +56,38 @@ def visualize_single_image(image_path, image_name, model, device):
     # MAE reconstruction pasted with visible patches
     im_paste = img * (1 - mask) + predicted_img * mask
     
-   # make the plt figure larger
-    plt.figure(figsize=(12, 4))
+#    # make the plt figure larger
+#     plt.figure(figsize=(12, 4))
 
-    plt.subplot(1, 4, 1)
-    show_image(img[0], "original")
+#     plt.subplot(1, 4, 1)
+#     show_image(img[0], "original")
 
-    plt.subplot(1, 4, 2)
-    show_image(im_masked[0], "masked")
+#     plt.subplot(1, 4, 2)
+#     show_image(im_masked[0], "masked")
 
-    plt.subplot(1, 4, 3)
-    show_image(predicted_img[0], "reconstruction")
+#     plt.subplot(1, 4, 3)
+#     show_image(predicted_img[0], "reconstruction")
 
-    plt.subplot(1, 4, 4)
-    show_image(im_paste[0], "reconstruction + visible")
+#     plt.subplot(1, 4, 4)
+#     show_image(im_paste[0], "reconstruction + visible")
 
-    plt.tight_layout()
+#     plt.tight_layout()
 
-    return plt
-
-# Example Usage
-image_path = 'images/dog.jpg'  # Replace with the actual path to your image
-# take the string after the last '/' as the image name
-image_name = image_path.split('/')[-1].split('.')[0]
-visualize_single_image(image_path, image_name, model, device)
+    return img[0].numpy(), im_masked[0].numpy(), predicted_img[0].numpy(), im_paste[0].numpy()
 
 inputs_image = [
     gr.components.Image(type="filepath", label="Input Image"),
 ]
 
 outputs_image = [
-    gr.components.Image(type="plot", label="Output Image"),
+    gr.components.Image(type="numpy", label="Output Image"),
 ]
 
 gr.Interface(
     fn=visualize_single_image,
     inputs=inputs_image,
     outputs=outputs_image,
+    examples=path,
     title="MAE-ViT Image Reconstruction",
     description="This is a demo of the MAE-ViT model for image reconstruction.",
-    allow_flagging=False,
-    allow_screenshot=False,
-    allow_remote_access=False,
 ).launch()