From 680e7d00d2080fd090d6246a660c57acf50df1a3 Mon Sep 17 00:00:00 2001
From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com>
Date: Fri, 20 Jan 2023 01:28:23 +0800
Subject: [PATCH] Add Pixelize

---
 .gitignore           |   1 +
 hakuimg/pixel.py     | 183 +++++++++++++++++++++++++++++++++++++++++++
 javascript/swicth.js |   6 ++
 javascript/utils.js  |   6 --
 scripts/main.py      |  27 ++++++-
 5 files changed, 213 insertions(+), 10 deletions(-)
 create mode 100644 .gitignore
 create mode 100644 hakuimg/pixel.py

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..c18dd8d
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+__pycache__/
diff --git a/hakuimg/pixel.py b/hakuimg/pixel.py
new file mode 100644
index 0000000..4720ed3
--- /dev/null
+++ b/hakuimg/pixel.py
@@ -0,0 +1,183 @@
+from typing import Any
+from numpy.typing import NDArray
+
+import cv2
+from PIL import Image
+import numpy as np
+
+
+INFLATE_FILTER = [
+    None,
+    np.array(
+        [[0, 1, 0],
+        [1, 1, 1],
+        [0, 1, 0]], np.uint8
+    ),
+    np.array(
+        [[1, 1, 1],
+        [1, 1, 1],
+        [1, 1, 1]], np.uint8
+    ),
+    np.array(
+        [[0, 0, 1, 0, 0],
+        [0, 1, 1, 1, 0],
+        [1, 1, 1, 1, 1],
+        [0, 1, 1, 1, 0],
+        [0, 0, 1, 0, 0]], np.uint8
+    ),
+    np.array(
+        [[1, 1, 1, 1, 1],
+        [1, 1, 1, 1, 1],
+        [1, 1, 1, 1, 1],
+        [1, 1, 1, 1, 1],
+        [1, 1, 1, 1, 1]], np.uint8
+    ),
+    np.ones((7, 7), np.uint8),
+    np.ones((9, 9), np.uint8),
+    np.ones((11, 11), np.uint8),
+    np.ones((13, 13), np.uint8),
+    np.ones((15, 15), np.uint8),
+    np.ones((17, 17), np.uint8)
+]
+
+
+def read_img_as_array(
+    img
+) -> NDArray[Any]:
+    '''Convert image to RGBA and read to ndarray'''
+    img = Image.fromarray(img)
+    img = img.convert('RGBA')
+    img_arr = np.asarray(img)
+    return img_arr
+
+
+def preprocess(
+    img: NDArray[Any],
+    blur: int = 0, 
+    erode: int = 0,
+) -> NDArray[Any]:
+    '''
+    Process for 
+    * outline inflation (erode)
+    * smoothing (blur)
+    * saturation
+    * contrast
+    '''
+    # outline process
+    if erode:
+        img = cv2.erode(
+            img, INFLATE_FILTER[erode], 
+            iterations = 1,
+        )
+    
+    # blur process
+    if blur:
+        img = cv2.bilateralFilter(
+            img, 15, blur*20, 20
+        )
+    img = img.astype(np.float32)
+    return img
+
+
+def pixelize(
+    img: NDArray[Any],
+    k: int, c: int,
+    d_w: int, d_h: int,
+    o_w: int, o_h: int,
+    precise: int,
+) -> tuple[NDArray[Any], NDArray[Any]]:
+    '''
+    Use down scale and up scale to make pixel image.
+    
+    And use k-means to confine the num of colors.
+    '''
+    img = cv2.resize(
+        img, (d_w, d_h), 
+        interpolation = cv2.INTER_NEAREST
+    )
+    
+    # reshape to 1-dim array(for every color) for k-means
+    # use k-means to abstract the colors to use
+    img_cp = img.reshape(-1, c)
+    criteria = (
+        cv2.TERM_CRITERIA_EPS+cv2.TERM_CRITERIA_MAX_ITER, 
+        precise*5, 0.01
+    )
+    _, label, center = cv2.kmeans(
+        img_cp, k, None,
+        criteria, 1, cv2.KMEANS_PP_CENTERS
+    )
+    result = center[label.flatten()]
+    result = result.reshape((img.shape))
+    
+    result = cv2.resize(
+        result, 
+        (o_w, o_h), 
+        interpolation=cv2.INTER_NEAREST
+    )
+    return result.astype(np.uint8), center.astype(np.uint8)
+
+
+def run(
+    src: Image.Image,
+    k: int = 3,
+    scale: int = 2,
+    blur: int = 0, 
+    erode: int = 0,
+    precise: int = 10,
+) -> tuple[Image.Image, list[list[str|float]]]:
+    #print('Start process.')
+    #print('Read raw image... ', end='', flush=True)
+    img = read_img_as_array(src)
+    
+    #convert color space
+    alpha_channel = img[:, :, 3]
+    img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)
+    h, w, c = img.shape
+    d_h = h // scale
+    d_w = w // scale
+    o_h = d_h * scale
+    o_w = d_w * scale
+    #print('done!')
+    
+    #print('Image preprocess... ', end='', flush=True)
+    # preprocess(erode, blur, saturation, contrast)
+    img = preprocess(
+        img,
+        blur, erode
+    )
+    #print('done!')
+    
+    #print('Pixelize... ', end='', flush=True)
+    # pixelize(using k-means)
+    result, colors = pixelize(
+        img, k, c,
+        d_w, d_h,
+        o_w, o_h,
+        precise
+    )
+    #print('done!')
+    
+    #print('Process output image... ', end='', flush=True)
+    # add alpha channel
+    a = cv2.resize(
+        alpha_channel, (d_w, d_h), 
+        interpolation = cv2.INTER_NEAREST
+    )
+    a = cv2.resize(
+        a, (o_w, o_h), 
+        interpolation = cv2.INTER_NEAREST
+    )
+    a[a!=0]=255
+    if 0 not in a:
+        a[0, 0] = 0
+    r, g, b = cv2.split(result)
+    result = cv2.merge((r, g, b, a))
+    
+    # for saving to png
+    result = cv2.cvtColor(
+        result, cv2.COLOR_RGBA2BGRA
+    )
+    #print('done!')
+    
+    return Image.fromarray(result)
\ No newline at end of file
diff --git a/javascript/swicth.js b/javascript/swicth.js
index 8175cae..bd61741 100644
--- a/javascript/swicth.js
+++ b/javascript/swicth.js
@@ -64,4 +64,10 @@ function switch_to_haku_blend(){
 }
 function switch_to_haku_eff(){
   return switch_to_inner_tab('haku_eff')
+}
+
+
+function switch_to_inpaint_upload(){
+  switch_to_img2img_tab(4);
+  return args_to_array(arguments);
 }
\ No newline at end of file
diff --git a/javascript/utils.js b/javascript/utils.js
index 9a2ed2b..d903f3f 100644
--- a/javascript/utils.js
+++ b/javascript/utils.js
@@ -18,10 +18,4 @@ function change_img_height(id, height){
 function get_change_height(id){
   id = "div[id*='" + id + "']"
   return (height)=>{change_img_height(id, height)}
-}
-
-
-function switch_to_inpaint_upload(){
-  switch_to_img2img_tab(4);
-  return args_to_array(arguments);
 }
\ No newline at end of file
diff --git a/scripts/main.py b/scripts/main.py
index 9caa2b1..8e3f88c 100644
--- a/scripts/main.py
+++ b/scripts/main.py
@@ -9,10 +9,13 @@ from PIL import Image, ImageFilter, ImageEnhance, ImageColor
 import cv2
 import numpy as np
 
-from hakuimg import blend
-from hakuimg import color
-from hakuimg import blur
-from hakuimg import sketch
+from hakuimg import(
+    blend,
+    blur,
+    color,
+    sketch,
+    pixel
+)
 
 
 '''
@@ -120,6 +123,13 @@ def add_tab():
                                 sk_color = gr.Radio(['gray', 'rgb'], value='gray', label='color mode')
                                 sk_scale = gr.Checkbox(False, label='use scale')
                                 sketch_btn = gr.Button("refresh", variant="primary")
+                            
+                            with gr.TabItem('Pixelize', elem_id='haku_Pixelize'):
+                                p_colors = gr.Slider(2, 128, 16, step=1, label='colors')
+                                p_dot_size = gr.Slider(2, 32, 8, step=1, label='dot size')
+                                p_outline = gr.Slider(0, 10, 5, step=1, label='outline inflating')
+                                p_smooth = gr.Slider(0, 10, 0, step=1, label='Smoothing')
+                                pixel_btn = gr.Button("refresh", variant="primary")
                     
                     with gr.TabItem('Other'):
                         with gr.Tabs(elem_id='function list'):
@@ -178,6 +188,15 @@ def add_tab():
             component.change(sketch.run, all_sk_input, image_out)
         sketch_btn.click(sketch.run, all_sk_input, image_out)
         
+        #pixelize
+        all_p_set = [
+            p_colors, p_dot_size, p_smooth, p_outline
+        ]
+        all_p_input = [image_eff] + all_p_set
+        for component in all_p_set:
+            component.change(pixel.run, all_p_input, image_out)
+        pixel_btn.click(pixel.run, all_p_input, image_out)
+        
         #send
         for btns, btn3, img_src in all_btns:
             for btn, img in zip(btns, all_layers):