Update app.py

app.py

@@ -2,57 +2,130 @@ import io
 from PIL import Image
 import gradio as gr
 
-def compress_image(file, target_value, unit):
-    if file is None:
-        return None, "Please upload an image"
+ALLOWED_EXTENSIONS = {"png", "jpg", "jpeg", "webp"}
+
+def get_file_size_label(size_bytes):
+    if size_bytes < 1024 * 1024:
+        return f"{size_bytes / 1024:.1f} KB"
+    return f"{size_bytes / (1024 * 1024):.2f} MB"
 
-    img = Image.open(file)
-    original_bytes_io = io.BytesIO()
-    img.save(original_bytes_io, format="PNG")
-    original_size = original_bytes_io.tell()
 
-    target_bytes = target_value * 1024 if unit == "KB" else target_value * 1024 * 1024
+def compress_to_target(img, fmt, target_bytes):
+    buf = io.BytesIO()
 
-    def compress(img, target_bytes):
-        for q in range(95, 74, -2):
+    # ---- PNG handling ----
+    if fmt == "PNG":
+        for level in range(1, 10):
             buf = io.BytesIO()
-            img.convert("RGB").save(buf, format="JPEG", quality=q, optimize=True)
+            img.save(buf, format="PNG", optimize=True, compress_level=level)
             if buf.tell() <= target_bytes:
                 buf.seek(0)
-                return buf
+                return buf, "PNG"
+
+        # fallback to WEBP lossless
+        buf = io.BytesIO()
+        img.save(buf, format="WEBP", lossless=True, quality=100)
+        if buf.tell() <= target_bytes:
+            buf.seek(0)
+            return buf, "WEBP"
+
+        for q in range(95, 79, -5):
+            buf = io.BytesIO()
+            img.save(buf, format="WEBP", quality=q, method=6)
+            if buf.tell() <= target_bytes:
+                buf.seek(0)
+                return buf, "WEBP"
+
         buf.seek(0)
-        return buf
+        return buf, "WEBP"
+
+    # ---- JPEG / WEBP ----
+    save_fmt = "JPEG" if fmt in ("JPEG", "JPG") else "WEBP"
 
-    if target_bytes >= original_size:
-        return img, f"Already under target size ({original_size/1024:.1f} KB)"
+    for q in range(95, 74, -2):
+        buf = io.BytesIO()
+        if save_fmt == "JPEG":
+            img.save(buf, format="JPEG", quality=q, optimize=True,
+                     subsampling=0 if q >= 85 else 2)
+        else:
+            img.save(buf, format="WEBP", quality=q, method=6)
 
-    compressed_buf = compress(img, target_bytes)
-    compressed_img = Image.open(compressed_buf)
+        if buf.tell() <= target_bytes:
+            buf.seek(0)
+            return buf, save_fmt
 
-    saved = original_size - compressed_buf.getbuffer().nbytes
-    pct = (saved / original_size) * 100
+    buf.seek(0)
+    return buf, save_fmt
 
-    msg = f"Original: {original_size/1024:.1f} KB | Compressed: {compressed_buf.getbuffer().nbytes/1024:.1f} KB | Saved: {pct:.1f}%"
 
-    return compressed_img, msg
+def process_image(file, target_value, unit):
+    if file is None:
+        return None, "❌ Please upload an image"
+
+    try:
+        img = Image.open(file)
+        fmt = img.format or "JPEG"
 
+        # normalize image mode
+        if img.mode in ("RGBA", "LA", "P"):
+            if fmt == "JPEG":
+                background = Image.new("RGB", img.size, (255, 255, 255))
+                if img.mode == "P":
+                    img = img.convert("RGBA")
+                background.paste(img, mask=img.split()[-1] if img.mode in ("RGBA", "LA") else None)
+                img = background
+        elif img.mode not in ("RGB", "L"):
+            img = img.convert("RGB")
 
-with gr.Blocks() as demo:
-    gr.Markdown("# 🗜️ ImagePress (Hugging Face Version)")
+        # original size
+        original_buf = io.BytesIO()
+        img.save(original_buf, format="PNG")
+        original_size = original_buf.tell()
+
+        target_bytes = int(target_value * 1024) if unit == "KB" else int(target_value * 1024 * 1024)
+
+        # already small
+        if target_bytes >= original_size:
+            return img, f"✅ Already optimized: {get_file_size_label(original_size)}"
+
+        # compress
+        compressed_buf, out_fmt = compress_to_target(img, fmt, target_bytes)
+        compressed_size = compressed_buf.getbuffer().nbytes
+
+        saved = max(0, original_size - compressed_size)
+        pct = (saved / original_size) * 100
+
+        result_msg = (
+            f"📦 Original: {get_file_size_label(original_size)}\n"
+            f"🗜️ Compressed: {get_file_size_label(compressed_size)}\n"
+            f"💾 Saved: {pct:.1f}%"
+        )
+
+        return Image.open(compressed_buf), result_msg
+
+    except Exception as e:
+        return None, f"❌ Error: {str(e)}"
+
+
+# ---- UI ----
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🗜️ ImagePress — Smart Compressor")
+    gr.Markdown("Compress images to target size **without visible quality loss**")
 
     with gr.Row():
         input_img = gr.Image(type="filepath", label="Upload Image")
-        output_img = gr.Image(label="Compressed Image")
+        output_img = gr.Image(label="Compressed Output")
 
-    target = gr.Number(value=200, label="Target Size")
-    unit = gr.Radio(["KB", "MB"], value="KB", label="Unit")
+    with gr.Row():
+        target = gr.Number(value=200, label="Target Size")
+        unit = gr.Radio(["KB", "MB"], value="KB", label="Unit")
 
-    btn = gr.Button("⚡ Compress")
+    btn = gr.Button("⚡ Compress Image")
 
-    result = gr.Textbox()
+    result = gr.Textbox(label="Result")
 
     btn.click(
-        compress_image,
+        fn=process_image,
         inputs=[input_img, target, unit],
         outputs=[output_img, result]
     )