Use MiniBatchKmeans for faster compression using random pixel sample for centroid fitting

app.py

@@ -1,19 +1,36 @@
 import gradio as gr
 import numpy as np
-from sklearn.cluster import KMeans
+from sklearn.cluster import MiniBatchKMeans
+from sklearn.utils import shuffle
 from PIL import Image
 import os
 import tempfile
 from io import BytesIO
 
-def compress_image(image, k):
-    pixels = image.reshape(-1, 3)
-    model = KMeans(n_clusters=k, n_init=10, random_state=42)
-    model.fit(pixels)
-    centres = model.cluster_centers_.astype(np.uint8)
+def compress_kmeans(image: np.ndarray, k: int, random_state: 42) -> np.ndarray:
+    """
+    Faster K‑Means compression using a random pixel sample for centroid fitting.
+    """
+    pixels = image.reshape(-1, 3).astype(np.float32)
+
+    sample_size = min(20000, len(pixels))
+    sample = shuffle(pixels, random_state=random_state)[:sample_size]
+
+    model = MiniBatchKMeans(
+        n_clusters=k,
+        batch_size=min(1024, sample_size),
+        n_init='auto',
+        max_iter=50,
+        random_state=random_state,
+        verbose=0,
+    )
+    model.fit(sample)
+
     labels = model.predict(pixels)
-    compressed = centres[labels].reshape(image.shape)
-    return compressed
+    centres = model.cluster_centers_.astype(np.uint8)
+    compressed_pixels = centres[labels]
+
+    return compressed_pixels.reshape(image.shape)
 
 def process(filepath, k):
     if filepath is None:
@@ -23,7 +40,7 @@ def process(filepath, k):
     orig = np.array(Image.open(filepath).convert('RGB'))
 
     # Compress
-    comp = compress_image(orig, k)
+    comp = compress_kmeans(orig, k, random_state=42)
 
     def get_size(img):
         with BytesIO() as buf: