UI Updates

Simplified UI

app.py

@@ -1,10 +1,13 @@
-# app.py — Gradio-native metrics, clean UI, CUDA/CPU only
-
-import os, math, cv2, base64
-import torch, numpy as np, gradio as gr
+import os
+import math
+import cv2
+import base64
+import torch
+import numpy as np
+import gradio as gr
 from PIL import Image
+import tempfile
 
-# Optional (fine if missing)
 try:
     import kornia.color as kc
 except Exception:
@@ -12,31 +15,24 @@ except Exception:
 
 from skimage.metrics import peak_signal_noise_ratio as psnr_metric
 from skimage.metrics import structural_similarity as ssim_metric
+from model import ViTUNetColorizer
+
 
-# ---------------- Device & Model (no MPS) ----------------
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-from model import ViTUNetColorizer
 CKPT = "checkpoints/checkpoint_epoch_015_20250808_154437.pt"
 model = None
 if os.path.exists(CKPT):
+    print(f"Loading model from: {CKPT}")
     model = ViTUNetColorizer(vit_model_name="vit_tiny_patch16_224").to(device)
     state = torch.load(CKPT, map_location=device)
     sd = state.get("generator_state_dict", state)
     model.load_state_dict(sd)
     model.eval()
-
-# ---------------- Utils ----------------
-def is_grayscale(img: Image.Image) -> bool:
-    a = np.array(img)
-    if a.ndim == 2: return True
-    if a.ndim == 3 and a.shape[2] == 1: return True
-    if a.ndim == 3 and a.shape[2] == 3:
-        return np.allclose(a[...,0], a[...,1]) and np.allclose(a[...,1], a[...,2])
-    return False
+else:
+    print(f"Warning: Checkpoint not found at {CKPT}. The app will not be able to colorize images.")
 
 def to_L(rgb_np: np.ndarray):
-    # ViTUNetColorizer expects L in [0,1]
     if kc is None:
         gray = cv2.cvtColor(rgb_np, cv2.COLOR_RGB2GRAY).astype(np.float32)
         L = gray / 100.0
@@ -71,18 +67,19 @@ def compute_metrics(pred, gt):
         ssim = float(ssim_metric(g, p, multichannel=True, data_range=1.0, win_size=7))
     return round(mae,4), round(psnr,2), round(ssim,4)
 
-# ---------------- Inference ----------------
-def infer(image: Image.Image, want_metrics: bool, show_L: bool):
+def to_grayscale(image):
     if image is None:
-        return None, None, None, None, None, "", ""
+        return None
+    return image.convert("L").convert("RGB")
+
+def infer(image: Image.Image, want_metrics: bool):
+    if image is None:
+        return None, None, None, None, None
     if model is None:
-        return None, None, None, None, None, "", "<div>Checkpoint not found in /checkpoints.</div>"
+        return None, None, None, None, "<div>Checkpoint not found.</div>"
 
     pil = image.convert("RGB")
     rgb = np.array(pil)
-    w,h = pil.size
-    was_color = not is_grayscale(pil)
-
     
     proc, (oh, ow) = pad_to_multiple(rgb, 16); back = (ow, oh)
 
@@ -93,46 +90,36 @@ def infer(image: Image.Image, want_metrics: bool, show_L: bool):
 
     out = out[:back[1], :back[0]]
 
-    # Metrics (Gradio-native numbers)
     mae = psnr = ssim = None
     if want_metrics:
         mae, psnr, ssim = compute_metrics(out, np.array(pil))
 
-    # Optional L preview
-    extra_html = ""
-    if show_L:
-        L01 = np.clip(L[0,0].detach().cpu().numpy(),0,1)
-        L_vis = (L01*255).astype(np.uint8)
-        L_vis = cv2.cvtColor(L_vis, cv2.COLOR_GRAY2RGB)
-        _, buf = cv2.imencode(".png", cv2.cvtColor(L_vis, cv2.COLOR_RGB2BGR))
-        L_b64 = "data:image/png;base64," + base64.b64encode(buf).decode()
-        extra_html += f"<div><b>L-channel</b><br/><img style='max-height:140px;border-radius:12px' src='{L_b64}'/></div>"
-
-    # Subtle notice only if needed
-    if was_color:
-        extra_html += "<div style='opacity:.8;margin-top:8px'>We used a grayscale version of your image for colorization.</div>"
-
-    # Compare slider (HTML only; easy to remove if you want 100% Gradio)
-    _, bo = cv2.imencode(".jpg", cv2.cvtColor(np.array(pil), cv2.COLOR_RGB2BGR))
-    _, bc = cv2.imencode(".jpg", cv2.cvtColor(out,         cv2.COLOR_RGB2BGR))
+    gray_pil = pil.convert("L").convert("RGB")
+    _, bo = cv2.imencode(".jpg", cv2.cvtColor(np.array(gray_pil), cv2.COLOR_RGB2BGR))
+    _, bc = cv2.imencode(".jpg", cv2.cvtColor(out, cv2.COLOR_RGB2BGR))
     so = "data:image/jpeg;base64," + base64.b64encode(bo).decode()
     sc = "data:image/jpeg;base64," + base64.b64encode(bc).decode()
-    compare = f"""
-    <div style="position:relative;max-width:500px;margin:auto;border-radius:14px;overflow:hidden;box-shadow:0 8px 20px rgba(0,0,0,.2)">
-    <img src="{so}" style="width:100%;display:block"/>
-    <div id="cmpTop" style="position:absolute;top:0;left:0;height:100%;width:50%;overflow:hidden">
-        <img src="{sc}" style="width:100%;display:block"/>
-    </div>
-    <input id="cmpRange" type="range" min="0" max="100" value="50"
-            oninput="document.getElementById('cmpTop').style.width=this.value+'%';"
-            style="position:absolute;left:0;right:0;bottom:8px;width:60%;margin:auto"/>
+ 
+    compare_html = f"""
+    <div style="margin:auto; border-radius:14px; overflow:hidden;">
+        <img-comparison-slider>
+            <img slot="first" src="{so}" />
+            <img slot="second" src="{sc}" />
+        </img-comparison-slider>
     </div>
     """
 
+    return out, mae, psnr, ssim, compare_html
 
-    return Image.fromarray(np.array(pil)), Image.fromarray(out), mae, psnr, ssim, compare, extra_html
+def save_for_download(image_array):
+    """Saves a NumPy array to a temporary file and returns the path."""
+    if image_array is not None:
+        pil_img = Image.fromarray(image_array)
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
+            pil_img.save(temp_file.name)
+            return temp_file.name
+    return None
 
-# ---------------- Theme (fallback-safe) ----------------
 def make_theme():
     try:
         from gradio.themes.utils import colors, fonts, sizes
@@ -146,9 +133,21 @@ def make_theme():
 
 THEME = make_theme()
 
-# ---------------- UI ----------------
-with gr.Blocks(theme=THEME, title="Image Colorizer") as demo:
+PLACEHOLDER_HTML = """
+<div style='display:flex; justify-content:center; align-items:center; height:480px; border: 2px dashed #4B5563; border-radius:12px; color:#4B5563; font-family:sans-serif;'>
+    <span>Result will be shown here</span>
+</div>
+"""
+
+HEAD = """
+<script type="module" src="https://unpkg.com/img-comparison-slider@8/dist/index.js"></script>
+<link rel="stylesheet" href="https://unpkg.com/img-comparison-slider@8/dist/themes/default.css" />
+"""
+
+with gr.Blocks(theme=THEME, title="Image Colorizer", head=HEAD) as demo:
     gr.Markdown("# 🎨 Image Colorizer")
+    
+    result_state = gr.State()
 
     with gr.Row():
         with gr.Column(scale=5):
@@ -159,12 +158,12 @@ with gr.Blocks(theme=THEME, title="Image Colorizer") as demo:
                 height=320,
                 sources=["upload", "webcam", "clipboard"]
             )
-            with gr.Row():
-                show_L = gr.Checkbox(label="Show L-channel", value=False)
+            img_in.upload(fn=to_grayscale, inputs=img_in, outputs=img_in)
             show_m = gr.Checkbox(label="Show metrics", value=True)
             with gr.Row():
                 run = gr.Button("Colorize")
                 clr = gr.Button("Clear")
+                download_btn = gr.DownloadButton("Download Result", visible=False)
 
             examples = gr.Examples(
                 examples=[os.path.join("examples", f) for f in os.listdir("examples")] if os.path.exists("examples") else [],
@@ -174,39 +173,44 @@ with gr.Blocks(theme=THEME, title="Image Colorizer") as demo:
             )
 
         with gr.Column(scale=7):
-            with gr.Row():
-                orig = gr.Image(label="Original", interactive=False, height=300, show_download_button=True)
-                out  = gr.Image(label="Result",   interactive=False, height=300, show_download_button=True)
-
-            # Pure Gradio metric fields
+            out_html = gr.HTML(label="Result", value=PLACEHOLDER_HTML)
             with gr.Row():
                 mae_box  = gr.Number(label="MAE",       interactive=False, precision=4)
                 psnr_box = gr.Number(label="PSNR (dB)", interactive=False, precision=2)
                 ssim_box = gr.Number(label="SSIM",      interactive=False, precision=4)
 
-            gr.Markdown("**Compare**")
-            compare = gr.HTML()
-            extras  = gr.HTML()
-
-    def _go(image, want_metrics, sizing_mode, show_L):
-        o, c, mae, psnr, ssim, cmp_html, extra = infer(image, want_metrics, show_L)
+    def _go(image, want_metrics):
+        out_image, mae, psnr, ssim, cmp_html = infer(image, want_metrics)
         if not want_metrics:
             mae = psnr = ssim = None
-        return o, c, mae, psnr, ssim, cmp_html, extra
+        
+        download_button_update = gr.update(visible=True) if out_image is not None else gr.update(visible=False)
+        
+        return out_image, cmp_html, mae, psnr, ssim, download_button_update
 
     run.click(
         _go,
-        inputs=[img_in, show_m, show_L],
-        outputs=[orig, out, mae_box, psnr_box, ssim_box, compare, extras]
+        inputs=[img_in, show_m],
+        outputs=[result_state, out_html, mae_box, psnr_box, ssim_box, download_btn]
     )
 
     def _clear():
-        return None, None, None, None, None, "", ""
-    clr.click(_clear, inputs=None, outputs=[orig, out, mae_box, psnr_box, ssim_box, compare, extras])
+        return None, PLACEHOLDER_HTML, None, None, None, gr.update(visible=False)
+    
+    clr.click(
+        _clear,
+        inputs=None,
+        outputs=[result_state, out_html, mae_box, psnr_box, ssim_box, download_btn]
+    )
+    
+    download_btn.click(
+        save_for_download,
+        inputs=[result_state],
+        outputs=[download_btn]
+    )
 
 if __name__ == "__main__":
-    # No queue, no API panel
     try:
         demo.launch(show_api=False)
     except TypeError:
-        demo.launch()
+        demo.launch()