Update app file

app.py

@@ -18,38 +18,32 @@ def prep_image(image):
 
 
 def reconstruct(img):
-    # image = Image.fromarray(img)
-    # pixel_values = feature_extractor(image, return_tensors='pt').pixel_values
-    #
-    # outputs = model(pixel_values)
-    # y = model.unpatchify(outputs.logits)
-    # y = torch.einsum('nchw->nhwc', y).detach().cpu()
-    #
-    # # visualize the mask
-    # mask = outputs.mask.detach()
-    # mask = mask.unsqueeze(-1).repeat(1, 1, model.config.patch_size ** 2 * 3)  # (N, H*W, p*p*3)
-    # mask = model.unpatchify(mask)  # 1 is removing, 0 is keeping
-    # mask = torch.einsum('nchw->nhwc', mask).detach().cpu()
-    #
-    # x = torch.einsum('nchw->nhwc', pixel_values).detach().cpu()
-    #
-    # # masked image
-    # im_masked = x * (1 - mask)
-    #
-    # # MAE reconstruction pasted with visible patches
-    # im_paste = x * (1 - mask) + y * mask
-    #
-    # out_masked = prep_image(im_masked[0])
-    # out_rec = prep_image(y[0])
-    # out_rec_vis = prep_image(im_paste[0])
-    #
-    # # out_masked, out_rec, out_rec_vis = img, img, img
-    #
-    # return [(out_masked, 'masked'), (out_rec, 'reconstruction'), (out_rec_vis, 'reconstruction + visible')]
-    url = "https://images.unsplash.com/photo-1507003211169-0a1dd7228f2d?ixlib=rb-1.2.1&ixid=MnwxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8&auto=format&fit=crop&w=387&q=80"
-    im = Image.open(requests.get(url, stream=True).raw)
-    im = np.array(im)
-    return [(im, 'label 1')]
+    image = Image.fromarray(img)
+    pixel_values = feature_extractor(image, return_tensors='pt').pixel_values
+
+    outputs = model(pixel_values)
+    y = model.unpatchify(outputs.logits)
+    y = torch.einsum('nchw->nhwc', y).detach().cpu()
+
+    # visualize the mask
+    mask = outputs.mask.detach()
+    mask = mask.unsqueeze(-1).repeat(1, 1, model.config.patch_size ** 2 * 3)  # (N, H*W, p*p*3)
+    mask = model.unpatchify(mask)  # 1 is removing, 0 is keeping
+    mask = torch.einsum('nchw->nhwc', mask).detach().cpu()
+
+    x = torch.einsum('nchw->nhwc', pixel_values).detach().cpu()
+
+    # masked image
+    im_masked = x * (1 - mask)
+
+    # MAE reconstruction pasted with visible patches
+    im_paste = x * (1 - mask) + y * mask
+
+    out_masked = prep_image(im_masked[0])
+    out_rec = prep_image(y[0])
+    out_rec_vis = prep_image(im_paste[0])
+
+    return [(out_masked, 'masked'), (out_rec, 'reconstruction'), (out_rec_vis, 'reconstruction + visible')]
 
 
 with gr.Blocks() as demo:
@@ -65,7 +59,7 @@ with gr.Blocks() as demo:
         #     label="Generated images", show_label=False, elem_id="gallery"
         # , columns=[3], rows=[1], height='auto', container=True)
 
-        gallery = gr.Gallery(columns=1,
+        gallery = gr.Gallery(columns=3,
                              rows=1,
                              height='800px',
                              object_fit='none')