Spaces:
Sleeping
Sleeping
File size: 2,017 Bytes
ecf1dca | 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 53 54 55 56 57 58 59 60 61 62 | import torch
import torch.nn as nn
import gradio as gr
import numpy as np
from torchvision.transforms import ToPILImage
# === Generator Model Definition ===
class Generator(nn.Module):
def __init__(self, latent_dim, label_dim, img_shape):
super().__init__()
input_dim = latent_dim + label_dim
self.model = nn.Sequential(
nn.Linear(input_dim, 128),
nn.ReLU(True),
nn.Linear(128, 256),
nn.BatchNorm1d(256),
nn.ReLU(True),
nn.Linear(256, 512),
nn.BatchNorm1d(512),
nn.ReLU(True),
nn.Linear(512, int(torch.prod(torch.tensor(img_shape)))),
nn.Tanh()
)
self.img_shape = img_shape
def forward(self, noise, labels):
x = torch.cat((noise, labels), dim=1)
img = self.model(x)
return img.view(img.size(0), *self.img_shape)
# === Load Model ===
latent_dim = 100
label_dim = 10
img_shape = (1, 28, 28)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
generator = Generator(latent_dim, label_dim, img_shape).to(device)
generator.load_state_dict(torch.load("models/mnist_generator.pth", map_location=device))
generator.eval()
# === Generate Function ===
def generate_images(digit: str):
digit = int(digit)
z = torch.randn(5, latent_dim, device=device)
labels = torch.eye(label_dim, device=device)[[digit] * 5]
with torch.no_grad():
gen_imgs = generator(z, labels).cpu()
gen_imgs = (gen_imgs + 1) / 2 # Normalize to [0, 1]
images = [ToPILImage()(img.squeeze(0)) for img in gen_imgs]
return images
# === Gradio Interface ===
iface = gr.Interface(
fn=generate_images,
inputs=gr.Dropdown(choices=[str(i) for i in range(10)], label="Pick a digit"),
outputs=[gr.Image(type="pil") for _ in range(5)],
title="MNIST Digit Generator",
description="Select a digit from 0–9 to generate 5 synthetic handwritten digits using a GAN trained on MNIST."
)
iface.launch()
|