Create app.py

app.py

@@ -0,0 +1,69 @@
+import gradio as gr 
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+from transformers import ViTMAEForPreTraining, ViTFeatureExtractor
+from PIL import Image
+import uuid
+
+feature_extractor = ViTFeatureExtractor.from_pretrained("facebook/vit-mae-base")
+model = ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base")
+
+imagenet_mean = np.array(feature_extractor.image_mean)
+imagenet_std = np.array(feature_extractor.image_std)
+
+def show_image(image, title=''):
+    # image is [H, W, 3]
+    assert image.shape[2] == 3
+    unique_id = str(uuid.uuid4())
+    plt.imshow(torch.clip((image * imagenet_std + imagenet_mean) * 255, 0, 255).int())
+    plt.axis('off')
+    plt.savefig(f"{unique_id}.png", bbox_inches='tight', pad_inches=0)
+
+    return f"{unique_id}.png"
+
+def visualize(image):
+    pixel_values = feature_extractor(image, return_tensors="pt").pixel_values
+    # forward pass
+    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)
+
+    # masked image
+    im_masked = x * (1 - mask)
+
+    # MAE reconstruction pasted with visible patches
+    im_paste = x * (1 - mask) + y * mask
+
+    gallery_labels = ["Original Image", "Masked Image", "Reconstruction", "Reconstruction with Patches"]
+    gallery_out = [show_image(out) for out in [x[0], im_masked[0], y[0], im_paste[0]]]
+
+    return [(k,v) for k,v in zip(gallery_out, gallery_labels)]
+
+  
+
+
+with gr.Blocks() as demo:
+  gr.Markdown("## ViTMAE Demo")
+  gr.Markdown("ViTMAE is an architecture that combine masked autoencoder and Vision Transformer (ViT) for self-supervised pre-training.")
+  gr.Markdown("By pre-training a ViT to reconstruct pixel values for masked patches, one can get results after fine-tuning that outperform supervised pre-training.")
+
+  with gr.Row():
+
+    input_img = gr.Image()
+    output = gr.Gallery()
+  
+  input_img.change(visualize, inputs=input_img, outputs=output)
+  
+  gr.Examples([["./cat.png"]], inputs=input_img, outputs=output, fn=visualize)
+
+demo.launch(debug=True)
+