Uploaded files

app.py

@@ -0,0 +1,84 @@
+import torch
+import torch.nn as nn
+import streamlit as st
+import torchvision.utils as vutils
+import matplotlib.pyplot as plt
+
+class Generator(nn.Module):
+    def __init__(self, channels_noise, channels_img, features_g):
+        super(Generator, self).__init__()
+        self.net = nn.Sequential(
+            # Input: N x channels_noise x 1 x 1
+            self._block(channels_noise, features_g * 16, 4, 1, 0),  # img: 4x4
+            self._block(features_g * 16, features_g * 8, 4, 2, 1),  # img: 8x8
+            self._block(features_g * 8, features_g * 4, 4, 2, 1),  # img: 16x16
+            self._block(features_g * 4, features_g * 2, 4, 2, 1),  # img: 32x32
+            nn.ConvTranspose2d(
+                features_g * 2, channels_img, kernel_size=4, stride=2, padding=1
+            ),
+            # Output: N x channels_img x 64 x 64
+            nn.Tanh(),
+        )
+
+    def _block(self, in_channels, out_channels, kernel_size, stride, padding):
+        return nn.Sequential(
+            nn.ConvTranspose2d(
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride,
+                padding,
+                bias=False,
+            ),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+    
+
+# Load the trained model
+@st.cache_resource
+def load_model(model_path="gan_final.pth", noise_dim=100, device="cpu"):
+    checkpoint = torch.load(model_path, map_location=device)
+    
+    # Recreate generator model
+    gen = Generator(channels_noise=noise_dim, channels_img=3, features_g=64).to(device)
+    gen.load_state_dict(checkpoint["generator"])
+    gen.eval()
+    
+    return gen
+
+# Function to generate images
+def generate_images(generator, num_images=1, noise_dim=100, device="cpu"):
+    noise = torch.randn(num_images, noise_dim, 1, 1, device=device)
+    with torch.no_grad():
+        fake_images = generator(noise).cpu()
+    
+    # Denormalize from [-1,1] to [0,1]
+    fake_images = (fake_images * 0.5) + 0.5  
+
+    return fake_images
+
+# Streamlit UI
+st.title("GAN Image Generator 🎨")
+st.write("Generate images using a trained GAN model.")
+
+# Load the model
+device = "cuda" if torch.cuda.is_available() else "cpu"
+generator = load_model(device=device)
+
+# User input for number of images
+num_images = st.slider("Select number of images", 1, 8, 4)
+
+# Generate button
+if st.button("Generate Images"):
+    st.write("🖌️ Generating images...")
+    fake_images = generate_images(generator, num_images=num_images, device=device)
+
+    # Display images
+    fig, ax = plt.subplots(figsize=(num_images, num_images))
+    ax.axis("off")
+    ax.imshow(vutils.make_grid(fake_images, padding=2, normalize=False).permute(1, 2, 0))
+    st.pyplot(fig)

gan_final.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:585d62c44b3618d62e4f46f6c1abadc5fb852685093f615347c6bd6e4b08b93d
+size 61734830

requirements.txt

@@ -0,0 +1,4 @@
+streamlit
+torch
+torchvision
+matplotlib