| 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 | |