Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| import matplotlib.pyplot as plt | |
| import io | |
| from siren import SIREN | |
| from utils import ( | |
| get_image_coordinates, | |
| image_to_tensor, | |
| tensor_to_image, | |
| downsample_image, | |
| train_siren, | |
| compute_psnr, | |
| compute_mae, | |
| compute_ssim_simple, | |
| get_model_cache_path, | |
| save_model, | |
| load_model | |
| ) | |
| def super_resolve_image(input_image, scale_factor, training_steps, hidden_features, hidden_layers, use_cache=True, image_name="uploaded"): | |
| """Perform super-resolution using SIREN. | |
| Args: | |
| input_image: PIL Image (high-res ground truth) | |
| scale_factor: Upscaling factor (2, 4, or 8) | |
| training_steps: Number of training steps | |
| hidden_features: Number of hidden units | |
| hidden_layers: Number of hidden layers | |
| use_cache: Whether to use cached models | |
| image_name: Name for cache identification | |
| Returns: | |
| Tuple of images and metrics | |
| """ | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| print(f"Using device: {device}") | |
| # Get original (ground truth) dimensions | |
| gt_image = input_image | |
| W_gt, H_gt = gt_image.size | |
| # Downsample the image | |
| downsampled_image = downsample_image(gt_image, scale_factor) | |
| W_low, H_low = downsampled_image.size | |
| print(f"Ground truth size: {W_gt}x{H_gt}") | |
| print(f"Downsampled size: {W_low}x{H_low}") | |
| print(f"Target upscale: {scale_factor}x") | |
| # Convert downsampled image to tensor | |
| low_res_pixels = image_to_tensor(downsampled_image) | |
| low_res_coords = get_image_coordinates(H_low, W_low) | |
| # Check cache | |
| cache_path = get_model_cache_path( | |
| f"{image_name}_{W_gt}x{H_gt}", | |
| scale_factor, | |
| training_steps, | |
| hidden_features, | |
| hidden_layers | |
| ) | |
| # Create SIREN model | |
| model = SIREN( | |
| in_features=2, | |
| hidden_features=hidden_features, | |
| hidden_layers=hidden_layers, | |
| out_features=3, | |
| outermost_linear=True, | |
| first_omega_0=30, | |
| hidden_omega_0=30 | |
| ) | |
| # Try to load from cache | |
| losses = [] | |
| if use_cache: | |
| loaded_model = load_model(model, cache_path) | |
| if loaded_model is not None: | |
| model = loaded_model | |
| print("Using cached model!") | |
| # Generate dummy loss curve | |
| losses = [0.01] * training_steps | |
| # Train if not loaded from cache | |
| if not losses: | |
| print("Training SIREN model...") | |
| model, losses = train_siren( | |
| model=model, | |
| coords=low_res_coords, | |
| pixels=low_res_pixels, | |
| num_steps=training_steps, | |
| learning_rate=1e-4, | |
| device=device | |
| ) | |
| print("Training complete!") | |
| # Save to cache | |
| if use_cache: | |
| save_model(model, cache_path) | |
| # Generate super-resolved image at original resolution | |
| model.eval() | |
| with torch.no_grad(): | |
| high_res_coords = get_image_coordinates(H_gt, W_gt).to(device) | |
| super_resolved_pixels = model(high_res_coords) | |
| # Convert to image | |
| super_resolved_image = tensor_to_image(super_resolved_pixels, H_gt, W_gt) | |
| # Compute quality metrics | |
| gt_pixels = image_to_tensor(gt_image) | |
| psnr = compute_psnr(super_resolved_pixels.cpu(), gt_pixels) | |
| mae = compute_mae(super_resolved_pixels.cpu(), gt_pixels) | |
| ssim = compute_ssim_simple(super_resolved_pixels.cpu(), gt_pixels) | |
| print(f"\nQuality Metrics:") | |
| print(f" PSNR: {psnr:.2f} dB") | |
| print(f" SSIM: {ssim:.4f}") | |
| print(f" MAE: {mae:.4f}") | |
| # Create metrics display | |
| metrics_text = f""" | |
| 📊 Quality Metrics (vs Ground Truth): | |
| • PSNR: {psnr:.2f} dB (higher is better, >30 dB is good) | |
| • SSIM: {ssim:.4f} (closer to 1.0 is better) | |
| • MAE: {mae:.4f} (lower is better) | |
| Training completed in {training_steps} steps | |
| Final MSE Loss: {losses[-1]:.6f} | |
| """ | |
| # Create loss plot | |
| fig, ax = plt.subplots(figsize=(6, 3)) | |
| ax.plot(losses, linewidth=2, color='#2E86AB') | |
| ax.set_xlabel('Training Step', fontsize=10) | |
| ax.set_ylabel('MSE Loss', fontsize=10) | |
| ax.set_title('Training Loss Curve', fontsize=12, fontweight='bold') | |
| ax.grid(True, alpha=0.3, linestyle='--') | |
| ax.set_facecolor('#f8f9fa') | |
| # Convert plot to image | |
| buf = io.BytesIO() | |
| plt.savefig(buf, format='png', bbox_inches='tight', dpi=100, facecolor='white') | |
| buf.seek(0) | |
| loss_plot = Image.open(buf) | |
| plt.close() | |
| # Return individual images and metrics | |
| # Order: downsampled, loss_plot, super_resolved, gt, metrics (matches UI layout) | |
| return downsampled_image, loss_plot, super_resolved_image, gt_image, metrics_text | |
| # Create Gradio interface | |
| with gr.Blocks(title="SIREN Super-Resolution") as demo: | |
| gr.Markdown( | |
| """ | |
| # 🔥 SIREN Super-Resolution Demo | |
| Upload a high-resolution image, and watch **SIREN** (Sinusoidal Representation Networks) | |
| learn to super-resolve it from an artificially downsampled version. | |
| **How it works:** Your image is downsampled → SIREN learns the low-res → Generates high-res → Compare with original! | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("### 📤 Input") | |
| input_image = gr.Image( | |
| type="pil", | |
| label="Upload High-Resolution Image", | |
| height=300 | |
| ) | |
| scale_factor = gr.Radio( | |
| choices=[2, 4, 8], | |
| value=2, | |
| label="Downsampling Scale Factor", | |
| info="Higher scale = harder task" | |
| ) | |
| training_steps = gr.Dropdown( | |
| choices=[500, 1000, 1500, 2000, 3000, 4000, 5000], | |
| value=2000, | |
| label="Training Epochs/Steps", | |
| info="More steps = better quality but slower" | |
| ) | |
| use_cache = gr.Checkbox( | |
| value=True, | |
| label="Use Model Cache", | |
| info="Save/load trained models to avoid retraining" | |
| ) | |
| with gr.Accordion("⚙️ Advanced Settings", open=False): | |
| hidden_features = gr.Slider( | |
| minimum=128, | |
| maximum=512, | |
| value=256, | |
| step=64, | |
| label="Hidden Features", | |
| info="Network width" | |
| ) | |
| hidden_layers = gr.Slider( | |
| minimum=2, | |
| maximum=6, | |
| value=3, | |
| step=1, | |
| label="Hidden Layers", | |
| info="Network depth" | |
| ) | |
| run_btn = gr.Button("🚀 Run Super-Resolution", variant="primary", size="lg") | |
| with gr.Column(scale=2): | |
| gr.Markdown("### 📊 Results & Comparison") | |
| with gr.Tabs(): | |
| with gr.Tab("📉 Side-by-Side Comparison"): | |
| gr.Markdown("**Low-Resolution Input & Training**") | |
| with gr.Row(): | |
| output_downsampled = gr.Image( | |
| label="Downsampled (Input)", | |
| type="pil", | |
| height=300 | |
| ) | |
| output_loss_plot = gr.Image( | |
| label="Training Loss Curve", | |
| type="pil", | |
| height=300 | |
| ) | |
| gr.Markdown("**High-Resolution Comparison**") | |
| with gr.Row(): | |
| output_super_resolved = gr.Image( | |
| label="Super-Resolved (SIREN Prediction)", | |
| type="pil", | |
| height=300 | |
| ) | |
| output_ground_truth = gr.Image( | |
| label="Ground Truth (Original)", | |
| type="pil", | |
| height=300 | |
| ) | |
| with gr.Tab("📈 Quality Metrics"): | |
| metrics_display = gr.Textbox( | |
| label="Quality Analysis", | |
| lines=10, | |
| max_lines=15 | |
| ) | |
| # Examples | |
| gr.Markdown("### 📸 Try these examples:") | |
| # Wrapper function to handle examples with image names | |
| def super_resolve_with_name(input_image, scale_factor, training_steps, hidden_features, hidden_layers, use_cache): | |
| # Extract image name from the example path if it's from samples | |
| image_name = "uploaded" | |
| if hasattr(input_image, 'name') and input_image.name: | |
| image_name = input_image.name.split('/')[-1].split('.')[0] | |
| return super_resolve_image(input_image, scale_factor, training_steps, hidden_features, hidden_layers, use_cache, image_name) | |
| gr.Examples( | |
| examples=[ | |
| ["samples/cat.jpg", 2, 2000, 256, 3, True], | |
| ["samples/landscape.jpg", 4, 3000, 256, 3, True], | |
| ["samples/portrait.jpg", 2, 2000, 256, 3, True], | |
| ["samples/flower.jpg", 4, 3000, 256, 4, True], | |
| ], | |
| inputs=[input_image, scale_factor, training_steps, hidden_features, hidden_layers, use_cache], | |
| outputs=[output_downsampled, output_loss_plot, output_super_resolved, output_ground_truth, metrics_display], | |
| fn=super_resolve_with_name, | |
| cache_examples=False, | |
| ) | |
| gr.Markdown( | |
| """ | |
| ### 📚 About SIREN & Metrics | |
| **SIREN** uses sine activation functions for representing continuous signals with fine details. | |
| **Quality Metrics Explained:** | |
| - **PSNR** (Peak Signal-to-Noise Ratio): Measures reconstruction quality. >30 dB is good, >40 dB is excellent. | |
| - **SSIM** (Structural Similarity Index): Perceptual quality metric. 1.0 is perfect, >0.9 is very good. | |
| - **MAE** (Mean Absolute Error): Average pixel difference. Lower is better. | |
| **Tips for Better Results:** | |
| - Start with 2x scale for quick testing | |
| - Use 3000-5000 steps for 4x and 8x scaling | |
| - Enable model cache to avoid retraining identical settings | |
| - Higher scale factors need more training steps and network capacity | |
| **Reference:** [SIREN Paper](https://arxiv.org/abs/2006.09661) | | |
| [Tutorial](https://github.com/nipunbatra/pml-teaching/blob/master/notebooks/siren.ipynb) | |
| """ | |
| ) | |
| # Connect the button | |
| run_btn.click( | |
| fn=super_resolve_with_name, | |
| inputs=[input_image, scale_factor, training_steps, hidden_features, hidden_layers, use_cache], | |
| outputs=[output_downsampled, output_loss_plot, output_super_resolved, output_ground_truth, metrics_display] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |