Create app.py

app.py

@@ -0,0 +1,327 @@
+import gradio as gr
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+import io
+import time
+
+# =====================================================================
+# 1. MODEL DEFINITION
+# =====================================================================
+
+class DynamicMLPAutoencoder(nn.Module):
+    """
+    A customizable Multi-Layer Perceptron Autoencoder.
+    Allows dynamic configuration of input size, hidden layers, and latent bottleneck.
+    """
+    def __init__(self, input_dim, hidden_dim, latent_dim, num_hidden_layers=1):
+        super().__init__()
+        
+        # Build Encoder
+        encoder_layers = []
+        current_dim = input_dim
+        
+        # Add progressive downscaling hidden layers
+        for i in range(num_hidden_layers):
+            next_dim = max(hidden_dim // (2 ** i), latent_dim * 2)
+            encoder_layers.append(nn.Linear(current_dim, next_dim))
+            encoder_layers.append(nn.ReLU())
+            current_dim = next_dim
+            
+        encoder_layers.append(nn.Linear(current_dim, latent_dim))
+        self.encoder = nn.Sequential(*encoder_layers)
+        
+        # Build Decoder
+        decoder_layers = []
+        current_dim = latent_dim
+        
+        # Add progressive upscaling hidden layers matching the encoder's reverse path
+        for i in reversed(range(num_hidden_layers)):
+            next_dim = max(hidden_dim // (2 ** i), latent_dim * 2)
+            decoder_layers.append(nn.Linear(current_dim, next_dim))
+            decoder_layers.append(nn.ReLU())
+            current_dim = next_dim
+            
+        decoder_layers.append(nn.Linear(current_dim, input_dim))
+        decoder_layers.append(nn.Sigmoid())  # Clamp output pixels between [0, 1]
+        self.decoder = nn.Sequential(*decoder_layers)
+
+    def forward(self, x):
+        latent = self.encoder(x)
+        reconstruction = self.decoder(latent)
+        return reconstruction, latent
+
+# =====================================================================
+# 2. HELPER UTILITIES
+# =====================================================================
+
+def preprocess_image(pil_img, resolution):
+    """Resizes, converts to tensor, and flattens an image."""
+    img_resized = pil_img.resize((resolution, resolution), Image.Resampling.LANCZOS)
+    img_np = np.array(img_resized).astype(np.float32) / 255.0
+    
+    # Handle Grayscale / RGBA conversions
+    if len(img_np.shape) == 2:  # Grayscale to RGB
+        img_np = np.stack([img_np] * 3, axis=-1)
+    elif img_np.shape[2] == 4:  # RGBA to RGB
+        img_np = img_np[:, :, :3]
+        
+    img_tensor = torch.tensor(img_np).permute(2, 0, 1).unsqueeze(0)  # Shape: [1, 3, H, W]
+    return img_tensor, img_resized
+
+def postprocess_tensor(tensor, resolution):
+    """Converts flattened/raw image tensors back to PIL Images."""
+    img_np = tensor.squeeze(0).permute(1, 2, 0).detach().cpu().numpy()
+    img_np = np.clip(img_np * 255.0, 0, 255).astype(np.uint8)
+    return Image.fromarray(img_np)
+
+def add_gaussian_noise(tensor, noise_factor):
+    """Adds zero-mean Gaussian noise to the image tensor."""
+    if noise_factor <= 0.0:
+        return tensor
+    noise = torch.randn_like(tensor) * noise_factor
+    noisy_tensor = torch.clamp(tensor + noise, 0.0, 1.0)
+    return noisy_tensor
+
+def create_loss_plot(losses):
+    """Generates a matplotlib line plot for training loss history."""
+    fig, ax = plt.subplots(figsize=(6, 3))
+    ax.plot(losses, color='#4F46E5', linewidth=2, label="Reconstruction Loss")
+    ax.set_title("Training Loss Curve", fontsize=11, fontweight='bold', pad=10)
+    ax.set_xlabel("Epoch", fontsize=9)
+    ax.set_ylabel("Loss (MSE)", fontsize=9)
+    ax.grid(True, linestyle='--', alpha=0.5)
+    ax.legend(loc="upper right")
+    fig.tight_layout()
+    return fig
+
+# =====================================================================
+# 3. INTERACTIVE TRAINING FUNCTION (GRADIO GENERATOR)
+# =====================================================================
+
+def run_autoencoder_sandbox(
+    input_image,
+    resolution,
+    latent_dim,
+    hidden_dim,
+    num_hidden_layers,
+    noise_factor,
+    epochs,
+    learning_rate,
+    optimizer_name,
+    progress=gr.Progress()
+):
+    if input_image is None:
+        raise gr.Error("Please upload or select an image first!")
+
+    # 1. Preprocessing
+    progress(0, desc="Preprocessing image...")
+    img_tensor, original_resized = preprocess_image(input_image, resolution)
+    
+    # Flatten dimensions
+    input_dim = 3 * resolution * resolution
+    flat_clean = img_tensor.view(1, -1)
+    
+    # 2. Setup Noise (for Denoising mode)
+    noisy_img_tensor = add_gaussian_noise(img_tensor, noise_factor)
+    flat_input = noisy_img_tensor.view(1, -1)
+    
+    # Pre-render the noisy input for user preview
+    noisy_preview = postprocess_tensor(noisy_img_tensor, resolution)
+
+    # 3. Model Initialization
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = DynamicMLPAutoencoder(
+        input_dim=input_dim,
+        hidden_dim=hidden_dim,
+        latent_dim=latent_dim,
+        num_hidden_layers=num_hidden_layers
+    ).to(device)
+    
+    flat_input = flat_input.to(device)
+    flat_clean = flat_clean.to(device)
+
+    # 4. Optimizer & Loss Config
+    if optimizer_name == "Adam":
+        optimizer = optim.Adam(model.parameters(), lr=learning_rate)
+    else:
+        optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9)
+        
+    criterion = nn.MSELoss()
+    losses = []
+    
+    # Determine UI update frequency to preserve execution speed
+    update_every = max(1, epochs // 40) 
+
+    # 5. Training Loop
+    for epoch in range(epochs):
+        model.train()
+        optimizer.zero_grad()
+        
+        # Forward Pass
+        reconstructed, latent = model(flat_input)
+        
+        # Calculate loss (against clean image always, to support denoising autoencoder structure)
+        loss = criterion(reconstructed, flat_clean)
+        
+        # Backward Pass & Step
+        loss.backward()
+        optimizer.step()
+        
+        losses.append(loss.item())
+        
+        # Emit live training updates
+        if epoch % update_every == 0 or epoch == epochs - 1:
+            progress((epoch + 1) / epochs, desc=f"Epoch {epoch+1}/{epochs} | Loss: {loss.item():.5f}")
+            
+            # Format reconstructed image back to normal shape
+            reconstructed_tensor = reconstructed.view(1, 3, resolution, resolution)
+            reconstructed_pil = postprocess_tensor(reconstructed_tensor, resolution)
+            
+            # Generate loss graph
+            loss_plot = create_loss_plot(losses)
+            
+            status_text = (
+                f"### Training Diagnostics\n"
+                f"- **Current Epoch:** {epoch + 1} / {epochs}\n"
+                f"- **Current Loss (MSE):** `{loss.item():.6f}`\n"
+                f"- **Compression Ratio:** `{input_dim} inputs ➜ {latent_dim} bottleneck` (Compacted by **{input_dim / latent_dim:.1f}x**)"
+            )
+            
+            # Yield components iteratively for active visual rendering
+            yield (
+                noisy_preview,
+                reconstructed_pil,
+                loss_plot,
+                status_text
+            )
+            plt.close(loss_plot)  # Cleanup plots to prevent memory overflow
+            time.sleep(0.01)
+
+# =====================================================================
+# 4. GRADIO APP INTERFACE LAYOUT
+# =====================================================================
+
+with gr.Blocks(theme=gr.themes.Soft(), title="Autoencoder Sandbox") as demo:
+    gr.Markdown(
+        """
+        # 🧠 Autoencoder Image Bottleneck Sandbox
+        Explore how deep neural networks compress, reconstruct, and denoise raw images. By training an autoencoder on *just* this single image, you can observe how limiting the **latent bottleneck** restricts the reconstruction capability—forcing the network to blur details or filter noise!
+        """
+    )
+    
+    with gr.Row():
+        # --- LEFT SIDEBAR: CONTROLS & ARCHITECTURE ---
+        with gr.Column(scale=1):
+            gr.Markdown("### 🛠️ Step 1: Input & Parameters")
+            input_img = gr.Image(type="pil", label="Upload Source Image", value=None)
+            
+            with gr.Tab("Network Architecture"):
+                resolution = gr.Dropdown(
+                    choices=[32, 64, 128], 
+                    value=64, 
+                    label="Image Training Resolution", 
+                    info="Higher resolutions take longer to train. 64x64 is recommended."
+                )
+                latent_dim = gr.Slider(
+                    minimum=1, 
+                    maximum=256, 
+                    value=16, 
+                    step=1, 
+                    label="Latent Dimension (Bottleneck)", 
+                    info="Lower values yield abstract, blurry representations ('not 1:1')."
+                )
+                hidden_dim = gr.Slider(
+                    minimum=16, 
+                    maximum=512, 
+                    value=128, 
+                    step=16, 
+                    label="Hidden Dimension", 
+                    info="Size of the intermediate neural network layers."
+                )
+                num_hidden_layers = gr.Slider(
+                    minimum=1, 
+                    maximum=3, 
+                    value=1, 
+                    step=1, 
+                    label="Number of Hidden Layers", 
+                    info="Deepens the feature extraction depth."
+                )
+            
+            with gr.Tab("Training Configuration"):
+                noise_factor = gr.Slider(
+                    minimum=0.0, 
+                    maximum=1.0, 
+                    value=0.0, 
+                    step=0.05, 
+                    label="Noise Injection (Denoising Mode)", 
+                    info="Introduces random noise to the training inputs. Model learns to clean it!"
+                )
+                epochs = gr.Slider(
+                    minimum=20, 
+                    maximum=1500, 
+                    value=400, 
+                    step=20, 
+                    label="Training Epochs", 
+                    info="Total gradient updates."
+                )
+                learning_rate = gr.Slider(
+                    minimum=0.0001, 
+                    maximum=0.05, 
+                    value=0.005, 
+                    step=0.0005, 
+                    label="Learning Rate"
+                )
+                optimizer_name = gr.Radio(
+                    choices=["Adam", "SGD"], 
+                    value="Adam", 
+                    label="Optimizer Mode"
+                )
+                
+            train_btn = gr.Button("🚀 Begin Real-Time Training", variant="primary")
+
+        # --- RIGHT SIDEBAR: OUTPUTS & METRICS ---
+        with gr.Column(scale=2):
+            gr.Markdown("### 📊 Step 2: Live Training Output")
+            
+            with gr.Row():
+                with gr.Column():
+                    noisy_output_img = gr.Image(label="Model Input (Noisy / Resized)", interactive=False)
+                with gr.Column():
+                    reconstruction_output_img = gr.Image(label="Live Model Reconstruction", interactive=False)
+            
+            with gr.Row():
+                with gr.Column(scale=1):
+                    diagnostics = gr.Markdown(
+                        "### Training Diagnostics\n*Click **Begin Real-Time Training** to launch the network optimization loop.*"
+                    )
+                with gr.Column(scale=1.2):
+                    loss_curve_plot = gr.Plot(label="Live Loss Curve")
+                    
+    # Link action to the training function
+    train_btn.click(
+        fn=run_autoencoder_sandbox,
+        inputs=[
+            input_img,
+            resolution,
+            latent_dim,
+            hidden_dim,
+            num_hidden_layers,
+            noise_factor,
+            epochs,
+            learning_rate,
+            optimizer_name
+        ],
+        outputs=[
+            noisy_output_img,
+            reconstruction_output_img,
+            loss_curve_plot,
+            diagnostics
+        ]
+    )
+
+if __name__ == "__main__":
+    demo.queue().launch()