Update app.py

app.py

@@ -2,6 +2,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import gradio as gr
+from gradio_imageslider import ImageSlider
 from PIL import Image
 import torchvision.transforms.functional as TF
 
@@ -43,61 +44,138 @@ class FastEDSR(nn.Module):
         return base_upscaled + details
 
 # --- 2. INITIALIZATION ---
-device = torch.device('cpu') # Hugging Face Free Tier runs on CPU
+device = torch.device('cpu') 
 model = FastEDSR(scale_factor=2, num_blocks=8, channels=64)
 
-# Load the weights (Update this string if your file is named differently in the HF root)
+# Load the weights
 model_path = "FastEDSR_x2_31dB.pth" 
 model.load_state_dict(torch.load(model_path, map_location=device))
 model.eval()
 
-# --- 3. INFERENCE FUNCTION ---
-def upscale_image(img):
-    if img is None:
-        return None
+def calc_psnr(pred, target):
+    mse = torch.mean((pred - target) ** 2)
+    if mse == 0:
+        return 100.0
+    return 10 * torch.log10(1.0 / mse).item()
+
+# --- 3. INFERENCE FUNCTIONS ---
+
+def standard_upscale(img):
+    if img is None: return None, ""
         
-    # Enforce constraints to prevent CPU OOM timeouts
-    # Max input 1024px -> Max output 2048px (2K)
     max_input_dim = 1024 
     w, h = img.size
     
     if w > max_input_dim or h > max_input_dim:
         scale = max_input_dim / max(w, h)
-        new_w, new_h = int(w * scale), int(h * scale)
-        img = img.resize((new_w, new_h), Image.BICUBIC)
+        w, h = int(w * scale), int(h * scale)
+        img = img.resize((w, h), Image.BICUBIC)
 
-    # Preprocess
     img = img.convert('RGB')
     input_tensor = TF.to_tensor(img).unsqueeze(0).to(device)
 
-    # Forward Pass
     with torch.no_grad():
         output_tensor = model(input_tensor)
 
-    # Postprocess
     output_tensor = output_tensor.squeeze(0).clamp(0, 1)
     output_img = TF.to_pil_image(output_tensor)
     
-    return output_img
+    new_w, new_h = output_img.size
+    
+    details = (
+        f"### Resolution Details\n"
+        f"**Before:** {w} x {h} ({w * h:,} pixels)\n\n"
+        f"**After:** {new_w} x {new_h} ({new_w * new_h:,} pixels)"
+    )
+    
+    return output_img, details
+
+def benchmark_upscale(hr_img):
+    if hr_img is None: return "", None, None
+    
+    hr_img = hr_img.convert('RGB')
+    w, h = hr_img.size
+    
+    # Enforce even dimensions so 2x scaling mathematically matches
+    w = w - (w % 2)
+    h = h - (h % 2)
+    hr_img = hr_img.crop((0, 0, w, h))
+    
+    max_input_dim = 2048 # HR can be 2048 because LR will be 1024
+    if w > max_input_dim or h > max_input_dim:
+        scale = max_input_dim / max(w, h)
+        w, h = int(w * scale), int(h * scale)
+        # Ensure even dimensions again after resize
+        w = w - (w % 2)
+        h = h - (h % 2)
+        hr_img = hr_img.resize((w, h), Image.BICUBIC)
+
+    # Create the simulated Low-Res image
+    lr_w, lr_h = w // 2, h // 2
+    lr_img = hr_img.resize((lr_w, lr_h), Image.BICUBIC)
+    
+    # Run Inference
+    lr_tensor = TF.to_tensor(lr_img).unsqueeze(0).to(device)
+    hr_tensor = TF.to_tensor(hr_img).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        pred_tensor = model(lr_tensor).clamp(0, 1)
+
+    # Calculate PSNR
+    psnr = calc_psnr(pred_tensor, hr_tensor)
+    pred_img = TF.to_pil_image(pred_tensor.squeeze(0))
+    
+    # Resize LR using NEAREST so it looks accurately pixelated in the slider comparison
+    lr_slider_img = lr_img.resize((w, h), Image.NEAREST)
+    
+    details = (
+        f"### Benchmark Results\n"
+        f"**PSNR:** {psnr:.2f} dB\n\n"
+        f"**Low-Res Input:** {lr_w} x {lr_h} ({lr_w * lr_h:,} pixels)\n\n"
+        f"**Model Output & Ground Truth:** {w} x {h} ({w * h:,} pixels)"
+    )
+    
+    return details, (lr_slider_img, pred_img), (hr_img, pred_img)
 
 # --- 4. GRADIO UI ---
 with gr.Blocks(theme=gr.themes.Soft()) as app:
     gr.Markdown(
         """
         # ⚡ FastEDSR 2x Image Upscaler
-        Upload an image to enhance and upscale it by 2x.  
-        *Note: To ensure stability on CPU infrastructure, input images larger than 1024px are proportionally downscaled before processing to guarantee a maximum 2K output.*
+        Upload an image to enhance and upscale it by 2x. 
         """
     )
     
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(type="pil", label="Low Resolution Input")
-            upscale_btn = gr.Button("Upscale Image", variant="primary")
-        with gr.Column():
-            output_image = gr.Image(type="pil", label="2x High Resolution Output")
-            
-    upscale_btn.click(fn=upscale_image, inputs=input_image, outputs=output_image)
+    with gr.Tabs():
+        # TAB 1: STANDARD
+        with gr.TabItem("⚡ Standard Upscaling"):
+            gr.Markdown("Directly upscale any low-resolution image.")
+            with gr.Row():
+                with gr.Column():
+                    std_input = gr.Image(type="pil", label="Low Resolution Input")
+                    std_btn = gr.Button("Upscale Image", variant="primary")
+                with gr.Column():
+                    std_output = gr.Image(type="pil", label="2x High Resolution Output")
+                    std_details = gr.Markdown()
+                    
+            std_btn.click(fn=standard_upscale, inputs=std_input, outputs=[std_output, std_details])
+
+        # TAB 2: BENCHMARK
+        with gr.TabItem("📊 Benchmark Mode"):
+            gr.Markdown("Upload a high-quality image. The app will compress it, upscale it, and measure the PSNR quality against the original.")
+            with gr.Row():
+                with gr.Column():
+                    bm_input = gr.Image(type="pil", label="Ground Truth (High Res) Image")
+                    bm_btn = gr.Button("Run Benchmark", variant="primary")
+                    bm_details = gr.Markdown()
+                with gr.Column():
+                    gr.Markdown("### Low-Res vs. Model Prediction")
+                    slider_lr_pred = ImageSlider(label="Left: Pixelated Low-Res | Right: FastEDSR")
+                    
+                    gr.Markdown("### Ground Truth vs. Model Prediction")
+                    slider_hr_pred = ImageSlider(label="Left: Original HR | Right: FastEDSR")
+                    
+            bm_btn.click(fn=benchmark_upscale, inputs=bm_input, outputs=[bm_details, slider_lr_pred, slider_hr_pred])
 
 if __name__ == "__main__":
     app.launch()