Update app.py

app.py

@@ -7,7 +7,7 @@ import numpy as np
 import gradio.components as gc
 import gradio as gr
 
-
+ 
 def pixart(
     i,
     block_size=4,
@@ -16,6 +16,8 @@ def pixart(
     local_contrast_blur_radius=51,  # has to be odd
     upscale=True,
     seed=None,
+    output_scaling=1,
+    dither_amount=15
 ):
     w, h = i.size
     dw = w//block_size
@@ -56,17 +58,35 @@ def pixart(
         most_sat_color = (hsv[label_grid == cluster] @
                           np.array(hsv_weights)).argmax()
         return hsv[label_grid == cluster][most_sat_color]
+    
     cluster_colors = np.array([
         pick_representative_pixel(c)
         for c in range(centers.shape[0])])
+    
+
 
-    # assign each pixel the color of its cluster
-    ki = cluster_colors[label_grid]
+    if dither_amount == 0:
+        # assign each pixel the color of its cluster
+        ki = cluster_colors[label_grid]
+    else:
+        # add noise to the colors before selecting the nearest color, this acts as a dithering effect
+        noised_colors = hsv32 + np.random.normal(0, dither_amount, hsv.shape)
+        noised_colors = np.clip(noised_colors, 0, 255)
+        flattened = noised_colors.reshape(-1, 3)
+        # use the dot product to find the closest cluster (could also try euclidean distance)
+        closest_clusters = np.argmax(flattened @ centers.T,axis=1)
+        closest_clusters_eucledian = np.argmin(np.linalg.norm(centers - flattened[:, None], axis=-1), axis=1)
+        label_grid = closest_clusters_eucledian.reshape(hsv32.shape[:2])
+        ki = cluster_colors[label_grid]
 
     rgb = cv2.cvtColor(ki.astype(np.uint8), cv2.COLOR_HSV2RGB)
     i = Image.fromarray(rgb)
     if upscale:
         i = i.resize((w, h), Image.Resampling.NEAREST)
+
+    if output_scaling != 1:
+        i = i.resize(
+            (w*output_scaling, h*output_scaling), Image.Resampling.NEAREST)
     return i, seed
 
 
@@ -77,6 +97,8 @@ def query(
     hsv_weights,  # ='0,0,1'
     local_contrast_blur_radius,  # =51 has to be odd
     seed,  # =42,
+    output_scaling,
+    dither_amount
 ):
     bs = float(block_size)
     w, h = i.size
@@ -95,12 +117,13 @@ def query(
         local_contrast_blur_radius=local_contrast_blur_radius,
         upscale=True,
         seed=int(seed) if seed != '' else None,
+        output_scaling=output_scaling,
+        dither_amount=dither_amount
     )
-    if n_clusters<=256:
-        pxart=pxart.convert('P', palette=Image.Palette.ADAPTIVE, colors=n_clusters)
-        
+    if n_clusters <= 256:
+        pxart = pxart.convert('P', palette=Image.Palette.ADAPTIVE, colors=n_clusters)
+    pxart.save('temp.bmp')
     return pxart, usedseed
-        
 
 
 # %%
@@ -112,7 +135,7 @@ block_size = gc.Textbox(
     label='Block Size ',
     placeholder="e.g. 8 for 8 pixels. 0.01 for 1% of min(w,h) (<1 for percentages, >= 1 for pixels)")
 palette_size = gc.Slider(
-    1, 2048, 32, step=1, label='Palette Size (Number of Colors)')
+    1, 1024, 32, step=1, label='Palette Size (Number of Colors)')
 hsv_weights = gc.Textbox(
     "0,0,1",
     label='HSV Weights. Weights of the channels when selecting a "representative pixel"/centroid from a cluster of pixels',
@@ -126,14 +149,19 @@ seed = gc.Textbox(
     placeholder='e.g. 42')
 
 outimage = gc.Image(
-    #shape=(224, 224), 
+    # shape=(224, 224), 
     label="Output", type='pil')
 seedout = gc.Textbox(label='used seed')
 
+output_scaling = gc.Slider(
+    0, 16, 1, step=1, label='Output scaling factor')
+
+dither_amount = gc.Slider(
+    0, 255, 0, step=1, label='Dithering amount')
 
 gr.Interface(
     query,
-    [searchimage, block_size, palette_size, hsv_weights,  lcbr, seed],
+    [searchimage, block_size, palette_size, hsv_weights,  lcbr, seed, output_scaling, dither_amount],
     [outimage, seedout],
     title="kmeans-Pixartifier",
     description=f"Turns images into pixel art using kmeans clustering",