fix: file size mismatch, download filename change

app.py

@@ -2,9 +2,10 @@ import gradio as gr
 import numpy as np
 from sklearn.cluster import KMeans
 from PIL import Image
+import os
+from io import BytesIO
 
 def compress_kmeans(image: np.ndarray, k: int) -> np.ndarray:
-    """Reduce image palette to k colours using K-Means vector quantisation."""
     pixels = image.reshape(-1, 3)
     model = KMeans(n_clusters=k, n_init="auto", random_state=42)
     model.fit(pixels)
@@ -12,57 +13,79 @@ def compress_kmeans(image: np.ndarray, k: int) -> np.ndarray:
     labels = model.predict(pixels)
     return centres[labels].reshape(image.shape)
 
-def compute_psnr_mse(original: np.ndarray, compressed: np.ndarray):
-    """Return (MSE, PSNR) between two images."""
-    mse = np.mean((original.astype(np.float64) - compressed.astype(np.float64)) ** 2)
+def compress_and_report(filepath, k):
+    orig_filename = os.path.splitext(os.path.basename(filepath))[0]
+    orig_pil = Image.open(filepath).convert("RGB")
+    orig_np = np.array(orig_pil)
+
+    comp_np = compress_kmeans(orig_np, k)
+    comp_pil = Image.fromarray(comp_np)
+
+    buf_orig = BytesIO()
+    orig_pil.save(buf_orig, format="PNG")
+    orig_size = len(buf_orig.getvalue())
+
+    buf_comp = BytesIO()
+    comp_pil.save(buf_comp, format="PNG")
+    comp_size = len(buf_comp.getvalue())
+
+    mse = np.mean((orig_np.astype(np.float64) - comp_np.astype(np.float64)) ** 2)
     psnr = float("inf") if mse == 0 else 20.0 * np.log10(255.0 / np.sqrt(mse))
-    return mse, psnr
 
-def compute_metrics(image: np.ndarray, k: int):
-    """Estimate compression ratio and file sizes (theoretical)."""
-    h, w, c = image.shape
-    original_bits = h * w * c * 8
-    compressed_bits = (k * 24) + (h * w * int(np.ceil(np.log2(k))))
-    ratio = original_bits / compressed_bits
-    return {
-        "ratio": ratio,
-        "original_kb": original_bits / (8 * 1024),
-        "compressed_kb": compressed_bits / (8 * 1024),
-        "saved_percent": (1.0 - (1.0 / ratio)) * 100.0,
-    }
+    ratio = orig_size / comp_size
+    saved_pct = (1.0 - (comp_size / orig_size)) * 100.0
+    orig_kb = orig_size / 1024
+    comp_kb = comp_size / 1024
 
-def compress_and_report(image, k):
-    """Gradio entry point: compress, compute metrics, return image + report."""
-    if isinstance(image, np.ndarray):
-        orig = image[:, :, :3] if image.shape[-1] == 4 else image
+    if psnr > 35:
+        quality_badge = '<span style="color:#16a34a;">Excellent</span>'
+    elif psnr > 28:
+        quality_badge = '<span style="color:#eab308;">Good</span>'
     else:
-        orig = np.array(image.convert("RGB"))
+        quality_badge = '<span style="color:#dc2626;">Acceptable</span>'
 
-    compressed = compress_kmeans(orig, k)
-    mse, psnr = compute_psnr_mse(orig, compressed)
-    metrics = compute_metrics(orig, k)
+    report_html = f"""
+    <div style="background: #f8f9fa; padding: 16px; border-radius: 8px; border-left: 4px solid #2563eb; font-family: sans-serif;">
+        <h4 style="margin: 0 0 12px 0; color: #1e293b;">Compression Report</h4>
+        <table style="width: 100%; border-collapse: collapse; font-size: 14px;">
+            <tr style="border-bottom: 1px solid #e2e8f0;">
+                <td style="padding: 6px 0; font-weight: 600; color: #475569;">PSNR</td>
+                <td style="padding: 6px 0; color: #0f172a;">{psnr:.2f} dB</td>
+                <td style="padding: 6px 0;">{quality_badge}</td>
+            </tr>
+            <tr style="border-bottom: 1px solid #e2e8f0;">
+                <td style="padding: 6px 0; font-weight: 600; color: #475569;">MSE</td>
+                <td style="padding: 6px 0; color: #0f172a;" colspan="2">{mse:.2f}</td>
+            </tr>
+            <tr style="border-bottom: 1px solid #e2e8f0;">
+                <td style="padding: 6px 0; font-weight: 600; color: #475569;">Compression Ratio</td>
+                <td style="padding: 6px 0; color: #0f172a;" colspan="2">{ratio:.2f}x</td>
+            </tr>
+            <tr style="border-bottom: 1px solid #e2e8f0;">
+                <td style="padding: 6px 0; font-weight: 600; color: #475569;">Original Size</td>
+                <td style="padding: 6px 0; color: #0f172a;" colspan="2">{orig_kb:.2f} KB</td>
+            </tr>
+            <tr>
+                <td style="padding: 6px 0; font-weight: 600; color: #475569;">Compressed Size</td>
+                <td style="padding: 6px 0; color: #0f172a;">{comp_kb:.2f} KB</td>
+                <td style="padding: 6px 0; color: #2563eb; font-weight: 600;">↓ {saved_pct:.1f}%</td>
+            </tr>
+        </table>
+    </div>
+    """
 
-    report = (
-        f"COMPRESSION REPORT\n"
-        f"{'-' * 40}\n"
-        f"PSNR               : {psnr:.2f} dB\n"
-        f"MSE                : {mse:.2f}\n"
-        f"Compression Ratio  : {metrics['ratio']:.2f}x\n"
-        f"Original Size (KB) : {metrics['original_kb']:.2f}\n"
-        f"Compressed Size (KB): {metrics['compressed_kb']:.2f}\n"
-        f"Space Saved (%)    : {metrics['saved_percent']:.1f}"
-    )
-    return compressed, report
+    download_name = f"{orig_filename}_compressed_{int(comp_kb)}KB.png"
+    return gr.update(value=comp_np, download_name=download_name), report_html
 
 iface = gr.Interface(
     fn=compress_and_report,
     inputs=[
-        gr.Image(type="pil", label="Upload Image"),
+        gr.Image(type="filepath", label="Upload Image"),
         gr.Slider(minimum=2, maximum=64, step=2, value=16, label="Number of Colors (k)"),
     ],
     outputs=[
         gr.Image(type="numpy", label="Compressed Image"),
-        gr.Textbox(label="Metrics Report", lines=10),
+        gr.Markdown(label="Metrics Report"),
     ],
     title="Image Compression using K-Means",
     description="Upload an image and adjust the colour palette size to see lossy compression in action.",