Spaces:
Build error
Build error
| import streamlit as st | |
| import torch | |
| from torchvision.utils import save_image | |
| import matplotlib.pyplot as plt | |
| from PIL import Image | |
| import io | |
| # Importa il modello del generatore | |
| class EfficientGenerator(torch.nn.Module): | |
| def __init__(self, latent_dim): | |
| super(EfficientGenerator, self).__init__() | |
| self.init_size = 64 // 16 | |
| self.l1 = torch.nn.Sequential(torch.nn.Linear(latent_dim, 128 * self.init_size ** 2)) | |
| self.conv_blocks = torch.nn.Sequential( | |
| torch.nn.BatchNorm2d(128), | |
| torch.nn.Upsample(scale_factor=2), | |
| torch.nn.Conv2d(128, 128, 3, stride=1, padding=1), | |
| torch.nn.BatchNorm2d(128, 0.8), | |
| torch.nn.ReLU(inplace=True), | |
| torch.nn.Upsample(scale_factor=2), | |
| torch.nn.Conv2d(128, 64, 3, stride=1, padding=1), | |
| torch.nn.BatchNorm2d(64, 0.8), | |
| torch.nn.ReLU(inplace=True), | |
| torch.nn.Upsample(scale_factor=2), | |
| torch.nn.Conv2d(64, 32, 3, stride=1, padding=1), | |
| torch.nn.BatchNorm2d(32, 0.8), | |
| torch.nn.ReLU(inplace=True), | |
| torch.nn.Upsample(scale_factor=2), | |
| torch.nn.Conv2d(32, 3, 3, stride=1, padding=1), | |
| torch.nn.Tanh(), | |
| ) | |
| def forward(self, z): | |
| out = self.l1(z) | |
| out = out.view(out.shape[0], 128, self.init_size, self.init_size) | |
| img = self.conv_blocks(out) | |
| return img | |
| # Carica il modello salvato | |
| def load_model(model_path): | |
| latent_dim = 32 | |
| model = EfficientGenerator(latent_dim) | |
| checkpoint = torch.load(model_path, map_location=torch.device('cpu')) | |
| model.load_state_dict(checkpoint['generator_state_dict']) | |
| model.eval() | |
| return model | |
| # Funzione per generare un'immagine | |
| def generate_image(model): | |
| latent_dim = 32 | |
| z = torch.randn(1, latent_dim) | |
| with torch.no_grad(): | |
| gen_img = model(z).detach().cpu() | |
| img = gen_img.squeeze(0).permute(1, 2, 0) | |
| img = (img + 1) / 2 # Denormalizza | |
| return img | |
| # Interfaccia Streamlit | |
| st.title("Generazione di Immagini GAN") | |
| st.write("Carica un modello GAN addestrato per generare immagini di gatti.") | |
| uploaded_model = st.file_uploader("Carica il modello", type=["pth"]) | |
| if uploaded_model is not None: | |
| # Carica il modello dal file caricato | |
| with open("temp_model.pth", "wb") as f: | |
| f.write(uploaded_model.getbuffer()) | |
| model = load_model("temp_model.pth") | |
| st.write("Modello caricato correttamente!") | |
| if st.button("Genera Immagine"): | |
| img = generate_image(model) | |
| fig, ax = plt.subplots() | |
| ax.imshow(img) | |
| ax.axis("off") | |
| st.pyplot(fig) | |
| # Salva l'immagine generata | |
| buffer = io.BytesIO() | |
| save_image(img.permute(2, 0, 1), buffer, format="PNG", normalize=True) | |
| st.download_button( | |
| label="Scarica Immagine", | |
| data=buffer.getvalue(), | |
| file_name="generated_image.png", | |
| mime="image/png" | |
| ) |