Merge pull request #41 from licyk/pixeloe

Add PixelOE and OutlineExpansion
main
Kohaku-Blueleaf 2025-03-17 13:15:38 +08:00 committed by GitHub
commit ac80e1e941
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 242 additions and 1 deletions

View File

@ -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

View File

@ -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

37
hakuimg/pixeloe.py Normal file
View File

@ -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

62
install.py Normal file
View File

@ -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)

9
requirements.txt Normal file
View File

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

View File

@ -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,