| import torch |
| import gradio as gr |
| from torchvision import transforms |
| from PIL import Image |
| import numpy as np |
| from model import model |
| import tempfile |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| transform = transforms.Compose([ |
| transforms.Resize((32, 32)), |
| transforms.ToTensor() |
| ]) |
|
|
| resize_output = transforms.Resize((512, 512)) |
|
|
| def interpolate_vectors(v1, v2, num_steps): |
| return [v1 * (1 - alpha) + v2 * alpha for alpha in np.linspace(0, 1, num_steps)] |
|
|
| def to_pil(img_tensor): |
| img = img_tensor.squeeze(0).permute(1, 2, 0).cpu().numpy() |
| img = (img * 255).clip(0, 255).astype(np.uint8) |
| return Image.fromarray(img) |
|
|
| def interpolate_images_gif(img1, img2, num_interpolations=10, duration=100): |
| img1 = Image.fromarray(img1).convert('RGB') |
| img2 = Image.fromarray(img2).convert('RGB') |
| img1_tensor = transform(img1).unsqueeze(0).to(device) |
| img2_tensor = transform(img2).unsqueeze(0).to(device) |
|
|
| with torch.no_grad(): |
| mu1, _ = model.encode(img1_tensor) |
| mu2, _ = model.encode(img2_tensor) |
| interpolated = interpolate_vectors(mu1, mu2, num_interpolations) |
| decoded_images = [] |
| for z in interpolated: |
| out = model.decode(z) |
| img = to_pil(out) |
| img_resized = resize_output(img) |
| decoded_images.append(img_resized) |
|
|
| tmp_file = tempfile.NamedTemporaryFile(suffix=".gif", delete=False) |
| decoded_images[0].save( |
| tmp_file.name, |
| save_all=True, |
| append_images=decoded_images[1:], |
| duration=duration, |
| loop=0 |
| ) |
| return tmp_file.name |
|
|
| def get_interface(): |
| with gr.Blocks() as iface: |
| gr.Markdown("## Latent Space Interpolation (GIF Output)") |
| with gr.Row(): |
| img1 = gr.Image(label="First Image", type="numpy") |
| img2 = gr.Image(label="Second Image", type="numpy") |
| slider_steps = gr.Slider(5, 30, value=10, step=1, label="Number of Interpolations") |
| slider_duration = gr.Slider(50, 500, value=100, step=10, label="Duration per Frame (ms)") |
| output_gif = gr.Image(label="Interpolation GIF") |
| run_button = gr.Button("Interpolate") |
|
|
| run_button.click( |
| fn=interpolate_images_gif, |
| inputs=[img1, img2, slider_steps, slider_duration], |
| outputs=output_gif |
| ) |
|
|
| examples = [ |
| ["example_images/image1.jpg", "example_images/image2.jpg", 10, 100], |
| ["example_images/image3.jpg", "example_images/image4.jpg", 15, 150], |
| ["example_images/image5.jpg", "example_images/image6.jpg", 20, 200], |
| ["example_images/image7.jpg", "example_images/image8.jpg", 25, 250], |
| ] |
|
|
| gr.Examples( |
| examples=examples, |
| inputs=[img1, img2, slider_steps, slider_duration], |
| outputs=output_gif, |
| fn=interpolate_images_gif, |
| cache_examples=False |
| ) |
| return iface |
|
|
