bug fix

app.py

@@ -36,10 +36,6 @@ def show_image(img, title):
     img = rearrange(img, "c h w -> h w c")
     img = (img.cpu().detach().numpy() + 1) / 2  # Normalize to [0, 1]
 
-    plt.imshow(img)
-    plt.axis('off')
-    plt.title(title)
-
 # Visualize a Single Image
 def visualize_single_image(image_path):
     img = load_image(image_path, transform).to(device)
@@ -56,8 +52,8 @@ def visualize_single_image(image_path):
     # MAE reconstruction pasted with visible patches
     im_paste = img * (1 - mask) + predicted_img * mask
 
-    # convert to numpy
-    img = img[0].cpu().detach().numpy()
+    img = rearrange(img[0], "c h w -> h w c")
+    img = (img.cpu().detach().numpy() + 1) / 2  # Normalize to [0, 1]
     
 #    # make the plt figure larger
 #     plt.figure(figsize=(12, 4))