Spaces:
Sleeping
Sleeping
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]) | |
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) | |
if __name__ == "__main__": | |
demo.launch() | |