fix: file size mismatch, download filename change

app.py

@@ -3,6 +3,7 @@ import numpy as np
 from sklearn.cluster import KMeans
 from PIL import Image
 import os
+import traceback
 from io import BytesIO
 
 def compress_kmeans(image: np.ndarray, k: int) -> np.ndarray:
@@ -14,69 +15,93 @@ def compress_kmeans(image: np.ndarray, k: int) -> np.ndarray:
     return centres[labels].reshape(image.shape)
 
 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)
+    if filepath is None:
+        return None, "<div style='padding: 20px; color: #64748b;'>Upload an image to view the report.</div>"
 
-    comp_np = compress_kmeans(orig_np, k)
-    comp_pil = Image.fromarray(comp_np)
+    try:
+        # Safely extract original filename
+        if isinstance(filepath, str):
+            base = os.path.basename(filepath)
+        else:
+            base = os.path.basename(str(filepath))
+        orig_filename = os.path.splitext(base)[0]
 
-    buf_orig = BytesIO()
-    orig_pil.save(buf_orig, format="PNG")
-    orig_size = len(buf_orig.getvalue())
+        orig_pil = Image.open(filepath).convert("RGB")
+        orig_np = np.array(orig_pil)
 
-    buf_comp = BytesIO()
-    comp_pil.save(buf_comp, format="PNG")
-    comp_size = len(buf_comp.getvalue())
+        comp_np = compress_kmeans(orig_np, k)
+        comp_pil = Image.fromarray(comp_np)
 
-    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))
+        # Measure actual PNG byte sizes
+        buf_orig = BytesIO()
+        orig_pil.save(buf_orig, format="PNG")
+        orig_size = len(buf_orig.getvalue())
 
-    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
+        buf_comp = BytesIO()
+        comp_pil.save(buf_comp, format="PNG")
+        comp_size = len(buf_comp.getvalue())
 
-    if psnr > 35:
-        quality_badge = '<span style="color:#16a34a;">Excellent</span>'
-    elif psnr > 28:
-        quality_badge = '<span style="color:#eab308;">Good</span>'
-    else:
-        quality_badge = '<span style="color:#dc2626;">Acceptable</span>'
+        # Metrics
+        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))
+        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
 
-    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>
-    """
+        # Quality badge
+        if psnr > 35:
+            badge = '<span style="color:#16a34a;">Excellent</span>'
+        elif psnr > 28:
+            badge = '<span style="color:#eab308;">Good</span>'
+        else:
+            badge = '<span style="color:#dc2626;">Acceptable</span>'
 
-    download_name = f"{orig_filename}_compressed_{int(comp_kb)}KB.png"
-    return gr.update(value=comp_np, download_name=download_name), report_html
+        # HTML Report
+        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;">{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>
+        """
 
+        download_name = f"{orig_filename}_compressed_{int(comp_kb)}KB.png"
+        return gr.update(value=comp_np, download_name=download_name), report_html
+
+    except Exception as e:
+        error_trace = traceback.format_exc()
+        error_html = f"""
+        <div style="background: #fee2e2; padding: 16px; border-radius: 8px; border-left: 4px solid #dc2626; font-family: monospace; font-size: 13px;">
+            <strong style="color: #991b1b;">Runtime Error</strong>
+            <pre style="white-space: pre-wrap; word-wrap: break-word; margin-top: 8px;">{error_trace}</pre>
+        </div>
+        """
+        return None, error_html
+
+# Gradio Interface
 iface = gr.Interface(
     fn=compress_and_report,
     inputs=[