add multiple image plot

app.py

@@ -28,7 +28,6 @@ transform = v2.Compose([
 # Load and Preprocess the Image
 def load_image(image_path, transform):
     img = Image.open(image_path).convert('RGB')
-    # transform = Compose([ToTensor(), Normalize(0.5, 0.5), Resize((32, 32))])
     img = transform(img).unsqueeze(0)  # Add batch dimension
     return img
 
@@ -36,6 +35,10 @@ 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)
@@ -51,28 +54,28 @@ def visualize_single_image(image_path):
 
     # MAE reconstruction pasted with visible patches
     im_paste = img * (1 - mask) + predicted_img * mask
-
-    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))
+   # 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()
+    # convert the plt figure to a numpy array
+    plt.savefig("output.png")
 
-    return img # , im_masked[0].numpy(), predicted_img[0].numpy(), im_paste[0].numpy()
+    return np.array(plt.imread("output.png"))
 
 inputs_image = [
     gr.components.Image(type="filepath", label="Input Image"),