Update app.py

app.py

@@ -1,14 +1,13 @@
 import torch
 import torch.nn as nn
-import torchvision.utils as vutils
-import gradio as gr
 import numpy as np
-import matplotlib.pyplot as plt
-
+import gradio as gr
+from PIL import Image
+import os
 
-# Define Generator architecture - must match what you used during training
+# Define your Generator architecture - with ngf=128 to match your training parameters
 class Generator(nn.Module):
-    def __init__(self, ngpu=1, nz=100, ngf=64, nc=3):
+    def __init__(self, ngpu=1, nz=100, ngf=128, nc=3):
         super(Generator, self).__init__()
         self.ngpu = ngpu
         self.main = nn.Sequential(
@@ -37,90 +36,62 @@ class Generator(nn.Module):
     def forward(self, input):
         return self.main(input)
 
+# Load the model - Update path to point to the models folder
+device = torch.device("cpu")
+model_path = "models/netG_epoch_246.pth"
 
-# Load the generator
-def load_model(model_path="models/netG_best.pth"):
-    # Create the generator and load the saved weights
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    netG = Generator(ngpu=1, nz=100, ngf=64, nc=3).to(device)
-
-    try:
-        netG.load_state_dict(torch.load(model_path, map_location=device))
-        netG.eval()  # Set to evaluation mode
-        print(f"Model loaded successfully from {model_path}")
-        return netG, device
-    except Exception as e:
-        print(f"Error loading model: {e}")
-        return None, device
-
-
-# Generate images using the model
-def generate_images(num_images=16, seed=None, randomize=True):
-    # Load the model (do this once when needed)
-    global model, device
-    if 'model' not in globals():
-        model, device = load_model()
-        if model is None:
-            return np.zeros((299, 299, 3))
-
-    # Set random seed for reproducibility if provided
-    if seed is not None and not randomize:
-        torch.manual_seed(seed)
-        np.random.seed(seed)
+# Print file existence for debugging
+print(f"Checking if model file exists: {os.path.exists(model_path)}")
+print(f"Listing contents of models directory: {os.listdir('models') if os.path.exists('models') else 'models directory not found'}")
 
-    # Generate latent vectors
-    nz = 100  # Size of the latent vector
-    noise = torch.randn(num_images, nz, 1, 1, device=device)
+# Initialize the model with ngf=128 to match your training parameters
+model = Generator(ngf=128).to(device)
 
-    # Generate fake images
+# Try loading with error handling
+try:
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    print("Model loaded successfully!")
+except Exception as e:
+    print(f"Error loading model: {e}")
+    # Try alternative loading methods if the first fails
+    try:
+        model.load_state_dict(torch.load(model_path, map_location=device), strict=False)
+        print("Model loaded with strict=False")
+    except Exception as e2:
+        print(f"Error with alternative loading: {e2}")
+
+# Set model to evaluation mode
+model.eval()
+print(f"Model initialized: {model is not None}")
+
+def generate_images(random_seed=42):
+    """Generate images using the DCGAN model"""
+    # Set seed for reproducibility
+    torch.manual_seed(random_seed)
+    
+    # Generate random noise
+    noise = torch.randn(1, 100, 1, 1, device=device)
+    
+    # Generate fake image
     with torch.no_grad():
         fake_images = model(noise).detach().cpu()
-
-    # Convert to grid for display
-    grid = vutils.make_grid(fake_images, padding=2, normalize=True, nrow=int(np.sqrt(num_images)))
-
-    # Convert from tensor to numpy array for Gradio
-    grid_np = grid.numpy().transpose((1, 2, 0))
-
-    # Make sure values are in 0-1 range
-    grid_np = np.clip(grid_np, 0, 1)
-
-    return grid_np
-
+    
+    # Convert tensor to image
+    fake_img = fake_images * 0.5 + 0.5  # unnormalize
+    fake_img = fake_img.squeeze(0).permute(1, 2, 0).numpy()
+    fake_img = np.clip(fake_img * 255, 0, 255).astype(np.uint8)
+    return Image.fromarray(fake_img)
 
 # Create Gradio interface
-def create_gradio_app():
-    with gr.Blocks(title="Computer Mouse Generator") as app:
-        gr.Markdown("# Computer Mouse GAN Generator")
-        gr.Markdown("Generate computer mice using a Deep Convolutional GAN trained on ~2,500 augmented images")
-
-        with gr.Row():
-            with gr.Column():
-                num_images = gr.Slider(minimum=1, maximum=64, value=16, step=1, label="Number of Images")
-                seed = gr.Number(label="Random Seed", value=42, precision=0)
-                randomize = gr.Checkbox(label="Use Random Seeds (ignore seed value)", value=True)
-                generate_button = gr.Button("Generate Mice")
-
-            with gr.Column():
-                output_image = gr.Image(label="Generated Computer Mice")
-
-        generate_button.click(fn=generate_images, inputs=[num_images, seed, randomize], outputs=output_image)
-
-        gr.Markdown("## About")
-        gr.Markdown("""This model was trained using a PyTorch DCGAN implementation on a dataset of computer mouse images.
-
-The training process used data augmentation to expand a small dataset of 300+ original images into 2,500+ training samples through techniques like flipping, rotation, and brightness/contrast adjustments.
-
-The generator creates brand new, never-before-seen computer mice from random noise!""")
-
-    return app
-
-
-# Initialize global variables
-model = None
-device = None
-
-# Launch the app if the script is run directly
+demo = gr.Interface(
+    fn=generate_images,
+    inputs=gr.Slider(minimum=1, maximum=100, step=1, default=42, label="Random Seed"),
+    outputs=gr.Image(type="pil", label="Generated Computer Mouse"),
+    title="DCGAN Computer Mouse Generator",
+    description="Generate unique computer mouse designs using a DCGAN model trained on computer mice images using ngf=128 and ndf=128.",
+    examples=[[42], [23], [7], [99]]
+)
+
+# Launch the app
 if __name__ == "__main__":
-    app = create_gradio_app()
-    app.launch()
+    demo.launch()