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, width, height): """Resizes (to width x height), converts to tensor, and flattens an image. Note: PIL's `.resize()` takes a (width, height) tuple, which is exactly the ordering we want to preserve here so non-square (n x m) resolutions work correctly. """ img_resized = pil_img.resize((width, height), 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, width, height): """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 def resolve_resolution(res_mode, custom_width, custom_height, input_image): """ Determines the final (width, height) to train on based on the selected mode: - "Use image's native resolution": pulls straight from the uploaded image - Preset choices like "64x64", "128x128", "64x256": parsed directly - "Custom": uses the custom_width / custom_height slider values """ if res_mode == "Use image's native resolution": if input_image is None: raise gr.Error("Please upload an image first so its native resolution can be used!") width, height = input_image.size # PIL gives (width, height) return width, height if res_mode == "Custom (set width/height below)": return int(custom_width), int(custom_height) # Preset like "64x64" or "64x256" try: w_str, h_str = res_mode.lower().split("x") return int(w_str), int(h_str) except Exception: raise gr.Error(f"Could not parse resolution preset: {res_mode}") # ===================================================================== # 3. INTERACTIVE TRAINING FUNCTION (GRADIO GENERATOR) # ===================================================================== def run_autoencoder_sandbox( input_image, res_mode, custom_width, custom_height, 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!") # 0. Resolve target width/height (supports n x n and n x m, or native size) width, height = resolve_resolution(res_mode, custom_width, custom_height, input_image) if width < 4 or height < 4: raise gr.Error("Width and height must each be at least 4 pixels.") # 1. Preprocessing progress(0, desc=f"Preprocessing image to {width}x{height}...") img_tensor, original_resized = preprocess_image(input_image, width, height) # Flatten dimensions using .reshape() instead of .view() to avoid layout conflicts input_dim = 3 * width * height flat_clean = img_tensor.reshape(1, -1) # 2. Setup Noise (for Denoising mode) noisy_img_tensor = add_gaussian_noise(img_tensor, noise_factor) flat_input = noisy_img_tensor.reshape(1, -1) # Pre-render the noisy input for user preview noisy_preview = postprocess_tensor(noisy_img_tensor, width, height) # 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 using .reshape() reconstructed_tensor = reconstructed.reshape(1, 3, height, width) reconstructed_pil = postprocess_tensor(reconstructed_tensor, width, height) # 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"- **Resolution:** `{width} x {height}` ({input_dim} input values)\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 # ===================================================================== RESOLUTION_PRESETS = [ "Use image's native resolution", "32x32", "64x64", "128x128", "64x256", "256x64", "Custom (set width/height below)", ] 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"): res_mode = gr.Dropdown( choices=RESOLUTION_PRESETS, value="64x64", label="Training Resolution", info="Pick a preset (n x n or n x m), use the image's native size, or set a custom width/height below." ) with gr.Row(): custom_width = gr.Slider( minimum=8, maximum=512, value=64, step=8, label="Custom Width", info="Only used when 'Custom' is selected above." ) custom_height = gr.Slider( minimum=8, maximum=512, value=64, step=8, label="Custom Height", info="Only used when 'Custom' is selected above." ) 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=5.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 # Added show_progress="hidden" to suppress output-level flickering/loading animations train_btn.click( fn=run_autoencoder_sandbox, inputs=[ input_img, res_mode, custom_width, custom_height, 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 ], show_progress="hidden" ) if __name__ == "__main__": demo.queue().launch()