Droping Metrics

app.py

@@ -13,11 +13,8 @@ try:
 except Exception:
     kc = None
 
-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 = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 CKPT = "checkpoints/checkpoint_epoch_017_20250810_193435.pt"
@@ -57,49 +54,34 @@ def pad_to_multiple(img_np, m=16):
     ph, pw = math.ceil(h/m)*m, math.ceil(w/m)*m
     return cv2.copyMakeBorder(img_np,0,ph-h,0,pw-w,cv2.BORDER_CONSTANT,value=(0,0,0)), (h,w)
 
-def compute_metrics(pred, gt):
-    p = pred.astype(np.float32)/255.; g = gt.astype(np.float32)/255.
-    mae  = float(np.mean(np.abs(p-g)))
-    psnr = float(psnr_metric(g, p, data_range=1.0))
-    try:
-        ssim = float(ssim_metric(g, p, channel_axis=2, data_range=1.0, win_size=7))
-    except TypeError:
-        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)
-
 def to_grayscale(image):
     if image is None:
         return None
     return image.convert("L").convert("RGB")
 
-def infer(image: Image.Image, want_metrics: bool):
+def infer(image: Image.Image):
     if image is None:
-        return None, None, None, None, None
+        return None, None
     if model is None:
-        return None, None, None, None, "<div>Checkpoint not found.</div>"
+        return None, "<div>Checkpoint not found.</div>"
 
     pil = image.convert("RGB")
     rgb = np.array(pil)
-    
+
     proc, (oh, ow) = pad_to_multiple(rgb, 16); back = (ow, oh)
 
     L = to_L(proc)
     with torch.no_grad():
         ab = model(L)
     out = lab_to_rgb(L, ab)
-
     out = out[:back[1], :back[0]]
 
-    mae = psnr = ssim = None
-    if want_metrics:
-        mae, psnr, ssim = compute_metrics(out, np.array(pil))
-
     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_html = f"""
     <div style="margin:auto; border-radius:14px; overflow:hidden;">
         <img-comparison-slider>
@@ -109,7 +91,7 @@ def infer(image: Image.Image, want_metrics: bool):
     </div>
     """
 
-    return out, mae, psnr, ssim, compare_html
+    return out, compare_html
 
 def save_for_download(image_array):
     """Saves a NumPy array to a temporary file and returns the path."""
@@ -159,7 +141,6 @@ with gr.Blocks(theme=THEME, title="Image Colorizer", head=HEAD) as demo:
                 sources=["upload", "webcam", "clipboard"]
             )
             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")
@@ -174,33 +155,25 @@ with gr.Blocks(theme=THEME, title="Image Colorizer", head=HEAD) as demo:
 
         with gr.Column(scale=7):
             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)
-
-    def _go(image, want_metrics):
-        out_image, mae, psnr, ssim, cmp_html = infer(image, want_metrics)
-        if not want_metrics:
-            mae = psnr = ssim = None
-        
+
+    def _go(image):
+        out_image, cmp_html = infer(image)
         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
+        return out_image, cmp_html, download_button_update
 
     run.click(
         _go,
-        inputs=[img_in, show_m],
-        outputs=[result_state, out_html, mae_box, psnr_box, ssim_box, download_btn]
+        inputs=[img_in],
+        outputs=[result_state, out_html, download_btn]
     )
 
     def _clear():
-        return None, None, PLACEHOLDER_HTML, None, None, None, gr.update(visible=False)
+        return None, None, PLACEHOLDER_HTML, gr.update(visible=False)
     
     clr.click(
         _clear,
         inputs=None,
-        outputs=[img_in, result_state, out_html, mae_box, psnr_box, ssim_box, download_btn]
+        outputs=[img_in, result_state, out_html, download_btn]
     )
     
     download_btn.click(
@@ -213,4 +186,4 @@ if __name__ == "__main__":
     try:
         demo.launch(show_api=False)
     except TypeError:
-        demo.launch()
+        demo.launch()