Update app filew

app.py

@@ -1,117 +1,115 @@
-# from transformers import ViTFeatureExtractor, ViTMAEForPreTraining
-# import numpy as np
-# import torch
-# from PIL import Image
-#
-# import gradio as gr
-#
-# 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])
-#
-#     # out_masked, out_rec, out_rec_vis = img, img, img
-#
-#     return [(out_masked, 'masked'), (out_rec, 'reconstruction'), (out_rec_vis, 'reconstruction + visible')]
-#     # return [(img, '1')]
-#
-#
-# 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], height='auto', container=True)
-#
-#         gallery = gr.Gallery(columns=3,
-#                              rows=1,
-#                              height='800px',
-#                              object_fit='none')
-#
-#     btn.click(reconstruct, img, gallery)
-#
-# if __name__ == "__main__":
-#     demo.launch()
-
-# This demo needs to be run from the repo folder.
-# python demo/fake_gan/run.py
-import random
+from transformers import ViTFeatureExtractor, ViTMAEForPreTraining
+import numpy as np
+import torch
+from PIL import Image
 
 import gradio as gr
 
+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 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
+
+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')]
+    return [(img, 'label 1')]
 
 
 with gr.Blocks() as demo:
     with gr.Column(variant="panel"):
         with gr.Row():
-            text = gr.Textbox(
+            img = gr.Image(
                 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")
+        # gallery = gr.Gallery(
+        #     label="Generated images", show_label=False, elem_id="gallery"
+        # , columns=[3], rows=[1], height='auto', container=True)
+
+        gallery = gr.Gallery(columns=1,
+                             rows=1,
+                             height='800px',
+                             object_fit='none')
 
-    btn.click(fake_gan, None, gallery)
+    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()
+