File size: 1,823 Bytes
2c480a0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
import torch
import gradio as gr
from torchvision import transforms
from PIL import Image
import numpy as np
from model import model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
resize_input = transforms.Resize((32, 32))
to_tensor = transforms.ToTensor()
def reconstruct_image(image):
image = Image.fromarray(image).convert('RGB')
image_32 = resize_input(image)
image_tensor = to_tensor(image_32).unsqueeze(0).to(device)
with torch.no_grad():
mu, _ = model.encode(image_tensor)
recon = model.decode(mu)
recon_np = recon.squeeze(0).permute(1, 2, 0).cpu().numpy()
recon_img = Image.fromarray((recon_np * 255).astype(np.uint8)).resize((512, 512))
orig_resized = image_32.resize((512, 512))
return orig_resized, recon_img
def get_interface():
with gr.Blocks() as iface:
gr.Markdown("## Encoding & Reconstruction")
with gr.Row():
input_image = gr.Image(label="Input (Downsampled to 32x32)", type="numpy")
output_image = gr.Image(label="Reconstructed", type="pil")
run_button = gr.Button("Run Reconstruction")
run_button.click(fn=reconstruct_image, inputs=input_image, outputs=[input_image, output_image])
examples = [
["example_images/image1.jpg"],
["example_images/image2.jpg"],
["example_images/image3.jpg"],
["example_images/image10.jpg"],
["example_images/image4.jpg"],
["example_images/image5.jpg"],
["example_images/image6.jpg"],
["example_images/image7.jpg"],
["example_images/image8.jpg"],
]
gr.Examples(
examples=examples,
inputs=[input_image],
)
return iface
|