v1 version of app

app.py

@@ -1,25 +1,21 @@
-from transformers import ViTFeatureExtractor, ViTMAEForPreTraining
+from transformers import ViTMAEForPreTraining, ViTImageProcessor
 import numpy as np
 import torch
-from PIL import Image
-import requests
-
 import gradio as gr
 
-feature_extractor = ViTFeatureExtractor.from_pretrained('andrewbo29/vit-mae-base-formula1')
+image_processor = ViTImageProcessor.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)
+imagenet_mean = np.array(image_processor.image_mean)
+imagenet_std = np.array(image_processor.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
+def reconstruct(image):
+    pixel_values = image_processor.preprocess(image, return_tensors='pt').pixel_values
 
     outputs = model(pixel_values)
     y = model.unpatchify(outputs.logits)
@@ -39,75 +35,40 @@ def reconstruct(img):
     # MAE reconstruction pasted with visible patches
     im_paste = x * (1 - mask) + y * mask
 
+    out_orig = prep_image(x[0])
     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')]
+    return [(out_orig, 'original'),
+            (out_masked, 'masked'),
+            (out_rec, 'reconstruction'),
+            (out_rec_vis, 'reconstruction + visible')]
 
 
 with gr.Blocks() as demo:
-    with gr.Column(variant="panel"):
-        with gr.Row():
+    with gr.Column(variant='panel'):
+        with gr.Column():
             img = gr.Image(
-                label="Enter your prompt",
                 container=False,
+                type='pil'
+            )
+            btn = gr.Button(
+                'Apply F1 MAE',
+                 scale=0
             )
-            btn = gr.Button("Generate image", scale=0)
-
-        # gallery = gr.Gallery(
-        #     label="Generated images", show_label=False, elem_id="gallery"
-        # , columns=[3], rows=[1], height='auto', container=True)
 
-        gallery = gr.Gallery(columns=3,
-                             rows=1,
-                             height='800px',
-                             object_fit='none')
+            gallery = gr.Gallery(
+                columns=4,
+                rows=1,
+                height='300px',
+                object_fit='none'
+            )
 
     btn.click(reconstruct, img, gallery)
 
 if __name__ == "__main__":
     demo.launch()
 
-# import random
-#
-# import gradio as gr
-#
-#
-# def fake_gan():
-#     images = [
-#         (random.choice(
-#             [
-#                 "https://images.unsplash.com/photo-1507003211169-0a1dd7228f2d?ixlib=rb-1.2.1&ixid=MnwxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8&auto=format&fit=crop&w=387&q=80",
-#                 "https://images.unsplash.com/photo-1554151228-14d9def656e4?ixlib=rb-1.2.1&ixid=MnwxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8&auto=format&fit=crop&w=386&q=80",
-#                 "https://images.unsplash.com/photo-1542909168-82c3e7fdca5c?ixlib=rb-1.2.1&ixid=MnwxMjA3fDB8MHxzZWFyY2h8MXx8aHVtYW4lMjBmYWNlfGVufDB8fDB8fA%3D%3D&w=1000&q=80",
-#                 "https://images.unsplash.com/photo-1546456073-92b9f0a8d413?ixlib=rb-1.2.1&ixid=MnwxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8&auto=format&fit=crop&w=387&q=80",
-#                 "https://images.unsplash.com/photo-1601412436009-d964bd02edbc?ixlib=rb-1.2.1&ixid=MnwxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8&auto=format&fit=crop&w=464&q=80",
-#             ]
-#         ), f"label {i}" if i != 0 else "label" * 50)
-#         for i in range(3)
-#     ]
-#     return images
-#
-#
-# with gr.Blocks() as demo:
-#     with gr.Column(variant="panel"):
-#         with gr.Row():
-#             text = gr.Textbox(
-#                 label="Enter your prompt",
-#                 max_lines=1,
-#                 placeholder="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=[2], rows=[2], object_fit="contain", height="auto")
-#
-#     btn.click(fake_gan, None, gallery)
-#
-# if __name__ == "__main__":
-#     demo.launch()
 
 

requirements.txt

@@ -1,4 +1,2 @@
 torch
-transformers
-numpy
-pillow
+transformers