Update app file

.gitignore

@@ -1 +1,2 @@
-.idea/
+.idea/
+flagged/

app.py

@@ -1,7 +1,64 @@
+from transformers import ViTFeatureExtractor, ViTMAEForPreTraining
+import numpy as np
+import torch
+from PIL import Image
+
 import gradio as gr
 
-def greet(name):
-    return "Hello " + name + "!!"
+feature_extractor = ViTFeatureExtractor.from_pretrained('andrewbo29/vit-mae-base-formula1')
+model = ViTMAEForPreTraining.from_pretrained('andrewbo29/vit-mae-base-formula1')
+
+imagenet_mean = np.array(feature_extractor.image_mean)
+imagenet_std = np.array(feature_extractor.image_std)
+
+
+def prep_image(image):
+    return torch.clip((image * imagenet_std + imagenet_mean) * 255, 0, 255).int().cpu().numpy()
+
+
+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])
+
+    return [(out_masked, 'masked'), (out_rec, 'reconstruction'), (out_rec_vis, 'reconstruction + visible')]
+
+
+with gr.Blocks() as demo:
+    with gr.Column(variant="panel"):
+        with gr.Row():
+            img = gr.Image(
+                label="Enter your prompt",
+                container=False,
+            )
+            btn = gr.Button("Generate image", scale=0)
+
+        gallery = gr.Gallery(
+            label="Generated images", show_label=False, elem_id="gallery"
+        , columns=[3], rows=[1], object_fit="contain", height='auto', container=True)
+
+    btn.click(reconstruct, img, gallery)
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,2 @@
+torch
+transformers