add pixeloe and outline expansion

hakuimg/image_preprocess.py

@@ -0,0 +1,28 @@
+import torch
+import numpy as np
+from PIL import Image
+
+
+def tensor_to_pil(image: torch.Tensor):
+    return Image.fromarray(
+        np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
+    )
+
+
+def pil_to_tensor(image: Image):
+    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
+
+
+def image_preprocess(img: torch.Tensor, device: str):
+    use_channel_last = False
+    if img.ndim == 3:
+        img = img.unsqueeze(0)
+    if img.size(3) <= 4:
+        img = img.permute(0, 3, 1, 2)
+        use_channel_last = True
+    if img.size(1) == 4:
+        img = img[:, :3]
+    org_device = img.device
+    if device != "default":
+        img = img.to(device)
+    return img, use_channel_last, org_device

hakuimg/outline_expansion.py

@@ -0,0 +1,19 @@
+from PIL import Image
+from .image_preprocess import image_preprocess, pil_to_tensor, tensor_to_pil
+from pixeloe.torch.outline import outline_expansion
+
+
+def run(
+    img: Image,
+    pixel_size: int,
+    thickness: int,
+    device: str,
+):
+    img = pil_to_tensor(img)
+    img, use_channel_last, org_device = image_preprocess(img, device)
+    oe_image, _ = outline_expansion(img, thickness, thickness, pixel_size)
+    oe_image = oe_image.to(org_device)
+    if use_channel_last:
+        oe_image = oe_image.permute(0, 2, 3, 1)
+    oe_image = tensor_to_pil(oe_image)
+    return oe_image

hakuimg/pixeloe.py

@@ -0,0 +1,37 @@
+from PIL import Image
+from pixeloe.torch.pixelize import pixelize
+from .image_preprocess import image_preprocess, pil_to_tensor, tensor_to_pil
+
+
+def run(
+    img: Image,
+    pixel_size: int,
+    thickness: int,
+    num_colors: int,
+    mode: str,
+    quant_mode: str,
+    dither_mode: str,
+    device: str,
+    color_quant: bool,
+    no_post_upscale: bool,
+):
+    img = pil_to_tensor(img)
+    img, use_channel_last, org_device = image_preprocess(img, device)
+    result, _, _ = pixelize(
+        img,
+        pixel_size,
+        thickness,
+        mode,
+        do_color_match=True,
+        do_quant=color_quant,
+        num_colors=num_colors,
+        quant_mode=quant_mode,
+        dither_mode=dither_mode,
+        no_post_upscale=no_post_upscale,
+        return_intermediate=True,
+    )
+    result = result.to(org_device)
+    if use_channel_last:
+        result = result.permute(0, 2, 3, 1)
+    result = tensor_to_pil(result)
+    return result

install.py

@@ -0,0 +1,62 @@
+import launch
+import pkg_resources
+from pathlib import Path
+from typing import Tuple, Optional
+
+
+repo_root = Path(__file__).parent
+main_req_file = repo_root / "requirements.txt"
+
+
+def comparable_version(version: str) -> Tuple:
+    return tuple(version.split("."))
+
+
+def get_installed_version(package: str) -> Optional[str]:
+    try:
+        return pkg_resources.get_distribution(package).version
+    except Exception:
+        return None
+
+
+def extract_base_package(package_string: str) -> str:
+    base_package = package_string.split("@git")[0]
+    return base_package
+
+
+def install_requirements(req_file):
+    with open(req_file) as file:
+        for package in file:
+            try:
+                package = package.strip()
+                if "==" in package:
+                    package_name, package_version = package.split("==")
+                    installed_version = get_installed_version(package_name)
+                    if installed_version != package_version:
+                        launch.run_pip(
+                            f"install -U {package}",
+                            f"a1111-sd-webui-haku-img requirement: changing {package_name} version from {installed_version} to {package_version}",
+                        )
+                elif ">=" in package:
+                    package_name, package_version = package.split(">=")
+                    installed_version = get_installed_version(package_name)
+                    if not installed_version or comparable_version(
+                        installed_version
+                    ) < comparable_version(package_version):
+                        launch.run_pip(
+                            f"install -U {package}",
+                            f"a1111-sd-webui-haku-img requirement: changing {package_name} version from {installed_version} to {package_version}",
+                        )
+                elif not launch.is_installed(extract_base_package(package)):
+                    launch.run_pip(
+                        f"install {package}",
+                        f"a1111-sd-webui-haku-img requirement: {package}",
+                    )
+            except Exception as e:
+                print(e)
+                print(
+                    f"Warning: Failed to install {package}"
+                )
+
+
+install_requirements(main_req_file)

requirements.txt

@@ -0,0 +1,9 @@
+torch
+torchvision
+pillow
+numpy
+opencv-python
+scipy
+matplotlib
+blendmodes
+pixeloe>=0.1.4

scripts/main.py

@@ -18,6 +18,8 @@ from hakuimg import (
     color,
     sketch,
     pixel,
+    pixeloe,
+    outline_expansion,
     neon,
     curve,
     chromatic,
@@ -286,6 +288,60 @@ def add_tab():
                                     pixel_btn = gr.Button("refresh", variant="primary")
                                     pixel_rst_btn = gr.Button("reset")
 
+                            with gr.TabItem("PixelOE", elem_id="haku_PixelOE"):
+                                pixeloe_pixel_size = gr.Slider(
+                                    1, 32, 4, step=1, label="pixel size"
+                                )
+                                pixeloe_thickness = gr.Slider(
+                                    0, 6, 2, step=1, label="thickness"
+                                )
+                                pixeloe_num_colors = gr.Slider(
+                                    2, 256, 256, step=1, label="num colors"
+                                )
+                                pixeloe_mode = gr.Radio(
+                                    ["contrast", "k_centroid", "lanczos", "nearest", "bilinear"],
+                                    value="contrast",
+                                    label="mode",
+                                )
+                                pixeloe_quant_mode = gr.Radio(
+                                    ["kmeans", "weighted-kmeans", "repeat-kmeans"],
+                                    value="kmeans",
+                                    label="quant mode",
+                                )
+                                pixeloe_dither_mode = gr.Radio(
+                                    ["ordered", "error_diffusion", "none"],
+                                    value="ordered",
+                                    label="dither mode",
+                                )
+                                pixeloe_device = gr.Radio(
+                                    ["default", "cpu", "cuda", "mps"],
+                                    value="default",
+                                    label="device",
+                                )
+                                pixeloe_color_quant = gr.Checkbox(label="color quant", value=False)
+                                pixeloe_no_post_upscale = gr.Checkbox(label="no post upscale", value=False)
+
+                                with gr.Row():
+                                    pixeloe_btn = gr.Button("refresh", variant="primary")
+                                    pixeloe_rst_btn = gr.Button("reset")
+
+                            with gr.TabItem("OutlineExpansion", elem_id="haku_OutlineExpansion"):
+                                outline_expansion_pixel_size = gr.Slider(
+                                    1, 32, 4, step=1, label="pixel size"
+                                )
+                                outline_expansion_thickness = gr.Slider(
+                                    1, 6, 3, step=1, label="thickness"
+                                )
+                                outline_expansion_device = gr.Radio(
+                                    ["default", "cpu", "cuda", "mps"],
+                                    value="default",
+                                    label="device",
+                                )
+
+                                with gr.Row():
+                                    outline_expansion_btn = gr.Button("refresh", variant="primary")
+                                    outline_expansion_rst_btn = gr.Button("reset")
+
                             with gr.TabItem("Glow", elem_id="haku_Glow"):
                                 neon_mode = gr.Radio(
                                     ["BS", "BMBL"], value="BS", label="Glow mode"
@@ -534,6 +590,36 @@ def add_tab():
         pixel_btn.click(pixel.run, all_p_input, image_out)
         pixel_rst_btn.click(lambda: [16, 8, 0, 5, "kmeans"], None, all_p_set)
 
+        # pixeloe
+        all_pixeloe_set = [
+            pixeloe_pixel_size,
+            pixeloe_thickness,
+            pixeloe_num_colors,
+            pixeloe_mode,
+            pixeloe_quant_mode,
+            pixeloe_dither_mode,
+            pixeloe_device,
+            pixeloe_color_quant,
+            pixeloe_no_post_upscale
+        ]
+        all_pixeloe_input = [image_eff] + all_pixeloe_set
+        for component in all_pixeloe_set:
+            _release_if_possible(component, pixeloe.run, all_pixeloe_input, image_out)
+        pixeloe_btn.click(pixeloe.run, all_pixeloe_input, image_out)
+        pixeloe_rst_btn.click(lambda: [4, 2, 256, "constrast", "kmeans", "kmeans", "default", False, False], None, all_pixeloe_set)
+
+        # outline expansion
+        all_outline_expansion_set = [
+            outline_expansion_pixel_size,
+            outline_expansion_thickness,
+            outline_expansion_device
+        ]
+        all_outline_expansion_input = [image_eff] + all_outline_expansion_set
+        for component in all_outline_expansion_set:
+            _release_if_possible(component, outline_expansion.run, all_outline_expansion_input, image_out)
+        outline_expansion_btn.click(outline_expansion.run, all_outline_expansion_input, image_out)
+        outline_expansion_rst_btn.click(lambda: [4, 3, "default"], None, all_outline_expansion_set)
+
         # neon
         all_neon_set = [
             neon_blur,
@@ -659,4 +745,4 @@ def on_ui_settings():
 
 
 script_callbacks.on_ui_tabs(add_tab)
-script_callbacks.on_ui_settings(on_ui_settings)
+script_callbacks.on_ui_settings(on_ui_settings)