app.py

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()