csabhay/ImageCompression

1import gradio as gr2import numpy as np3from sklearn.cluster import MiniBatchKMeans4from sklearn.utils import shuffle5from PIL import Image6import os7import tempfile8from io import BytesIO9 10def compress_kmeans(image: np.ndarray, k: int, random_state: 42) -> np.ndarray:11    """12    Faster K‑Means compression using a random pixel sample for centroid fitting.13    """14    pixels = image.reshape(-1, 3).astype(np.float32)15 16    sample_size = min(20000, len(pixels))17    sample = shuffle(pixels, random_state=random_state)[:sample_size]18 19    model = MiniBatchKMeans(20        n_clusters=k,21        batch_size=min(1024, sample_size),22        n_init='auto',23        max_iter=50,24        random_state=random_state,25        verbose=0,26    )27    model.fit(sample)28 29    labels = model.predict(pixels)30    centres = model.cluster_centers_.astype(np.uint8)31    compressed_pixels = centres[labels]32 33    return compressed_pixels.reshape(image.shape)34 35def process(filepath, k):36    if filepath is None:37        return None, "Please upload an image.", None38 39    # Load original40    orig = np.array(Image.open(filepath).convert('RGB'))41 42    # Compress43    comp = compress_kmeans(orig, k, random_state=42)44 45    def get_size(img):46        with BytesIO() as buf:47            Image.fromarray(img).save(buf, format='PNG')48            return len(buf.getvalue())49 50    orig_size = get_size(orig)51    comp_size = get_size(comp)52    ratio = orig_size / comp_size53    saved = (1 - comp_size/orig_size) * 10054 55    # PSNR56    mse = np.mean((orig.astype(float) - comp.astype(float)) ** 2)57    psnr = 20 * np.log10(255.0 / np.sqrt(mse)) if mse > 0 else float('inf')58 59    stats = (f"**Original:** {orig_size/1024:.1f} KB  \n"60             f"**Compressed:** {comp_size/1024:.1f} KB  \n"61             f"**Compression ratio:** {ratio:.1f}x  \n"62             f"**Space saved:** {saved:.1f}%  \n"63             f"**PSNR:** {psnr:.1f} dB")64 65    temp = tempfile.mkdtemp()66    out_path = os.path.join(temp, 'compressed.png')67    Image.fromarray(comp).save(out_path)68 69    return comp, stats, out_path70 71iface = gr.Interface(72    fn=process,73    inputs=[74        gr.Image(type='filepath', label='Upload an Image'),75        gr.Slider(minimum=2, maximum=64, step=2, value=16, label='Number of colours (k)')76    ],77    outputs=[78        gr.Image(label='Compressed Image'),79        gr.Markdown(label='Compression Statistics'),80        gr.File(label='Download Compressed Image')81    ],82    title='Image Compression with K‑Means',83    description='Reduce the number of colours in an image using K-Means clustering.'84)85 86iface.launch()