Update app.py

app.py

@@ -19,32 +19,32 @@ class DynamicMLPAutoencoder(nn.Module):
     """
     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)
@@ -58,21 +58,26 @@ class DynamicMLPAutoencoder(nn.Module):
 # 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)
+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, resolution):
+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)
@@ -98,13 +103,38 @@ def create_loss_plot(losses):
     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,
-    resolution,
+    res_mode,
+    custom_width,
+    custom_height,
     latent_dim,
     hidden_dim,
     num_hidden_layers,
@@ -117,20 +147,26 @@ def run_autoencoder_sandbox(
     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="Preprocessing image...")
-    img_tensor, original_resized = preprocess_image(input_image, resolution)
-    
+    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 * resolution * resolution
+    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, resolution)
+    noisy_preview = postprocess_tensor(noisy_img_tensor, width, height)
 
     # 3. Model Initialization
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -140,7 +176,7 @@ def run_autoencoder_sandbox(
         latent_dim=latent_dim,
         num_hidden_layers=num_hidden_layers
     ).to(device)
-    
+
     flat_input = flat_input.to(device)
     flat_clean = flat_clean.to(device)
 
@@ -149,48 +185,49 @@ def run_autoencoder_sandbox(
         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) 
+    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, resolution, resolution)
-            reconstructed_pil = postprocess_tensor(reconstructed_tensor, resolution)
-            
+            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,
@@ -205,6 +242,16 @@ def run_autoencoder_sandbox(
 # 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(
         """
@@ -212,87 +259,104 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Autoencoder Sandbox") as demo:
         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."
+                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)", 
+                    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", 
+                    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", 
+                    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=2.0, 
-                    value=0.0, 
-                    step=0.05, 
-                    label="Noise Injection (Denoising Mode)", 
+                    minimum=0.0,
+                    maximum=2.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", 
+                    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, 
+                    minimum=0.0001,
+                    maximum=0.05,
+                    value=0.005,
+                    step=0.0005,
                     label="Learning Rate"
                 )
                 optimizer_name = gr.Radio(
-                    choices=["Adam", "SGD"], 
-                    value="Adam", 
+                    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(
@@ -300,14 +364,16 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Autoencoder Sandbox") as demo:
                     )
                 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,
-            resolution,
+            res_mode,
+            custom_width,
+            custom_height,
             latent_dim,
             hidden_dim,
             num_hidden_layers,