377 lines
11 KiB
Python
377 lines
11 KiB
Python
from PIL import Image
|
|
import numpy as np
|
|
import torch
|
|
import os
|
|
import folder_paths
|
|
from .ldivider.ld_utils import save_psd, load_masks, divide_folder, load_seg_model
|
|
from .ldivider.ld_convertor import pil2cv, cv2pil, df2bgra
|
|
from .ldivider.ld_processor import get_base, get_normal_layer, get_composite_layer, get_seg_base
|
|
from .ldivider.ld_segment import get_mask_generator, get_masks, show_anns
|
|
from pytoshop.enums import BlendMode
|
|
import requests
|
|
|
|
comfy_path = os.path.dirname(folder_paths.__file__)
|
|
|
|
layer_divider_path = f'{comfy_path}/custom_nodes/ComfyUI-LayerDivider'
|
|
|
|
output_dir = f"{layer_divider_path}/output"
|
|
input_dir = f"{layer_divider_path}/input"
|
|
model_dir = f"{layer_divider_path}/segment_model"
|
|
|
|
if not os.path.exists(f'{output_dir}'):
|
|
os.makedirs(f'{output_dir}')
|
|
|
|
import uuid
|
|
|
|
import cv2
|
|
|
|
|
|
def HWC3(x):
|
|
assert x.dtype == np.uint8
|
|
if x.ndim == 2:
|
|
x = x[:, :, None]
|
|
assert x.ndim == 3
|
|
H, W, C = x.shape
|
|
assert C == 1 or C == 3 or C == 4
|
|
if C == 3:
|
|
return x
|
|
if C == 1:
|
|
return np.concatenate([x, x, x], axis=2)
|
|
if C == 4:
|
|
color = x[:, :, 0:3].astype(np.float32)
|
|
alpha = x[:, :, 3:4].astype(np.float32) / 255.0
|
|
y = color * alpha + 255.0 * (1.0 - alpha)
|
|
y = y.clip(0, 255).astype(np.uint8)
|
|
return y
|
|
|
|
|
|
def to_comfy_img(np_img):
|
|
out_imgs = []
|
|
|
|
out_imgs.append(HWC3(np_img))
|
|
|
|
out_imgs = np.stack(out_imgs)
|
|
|
|
out_imgs = torch.from_numpy(out_imgs.astype(np.float32) / 255.)
|
|
|
|
return out_imgs
|
|
|
|
|
|
def to_comfy_imgs(np_imgs):
|
|
out_imgs = []
|
|
|
|
for np_img in np_imgs:
|
|
out_imgs.append(HWC3(np_img))
|
|
|
|
out_imgs = np.stack(out_imgs)
|
|
|
|
out_imgs = torch.from_numpy(out_imgs.astype(np.float32) / 255.)
|
|
|
|
return out_imgs
|
|
|
|
|
|
def generate_layers(input_image, cv_image, df, layer_mode, divide_mode):
|
|
base_image = to_comfy_img(df2bgra(df))
|
|
comfy_image = to_comfy_img(cv_image)
|
|
|
|
if layer_mode == "composite":
|
|
base_layer_list, shadow_layer_list, bright_layer_list, addition_layer_list, subtract_layer_list = (
|
|
get_composite_layer(input_image, df))
|
|
|
|
filename = save_psd(
|
|
input_image,
|
|
[base_layer_list, bright_layer_list, shadow_layer_list, subtract_layer_list, addition_layer_list],
|
|
["base", "screen", "multiply", "subtract", "addition"],
|
|
[BlendMode.normal, BlendMode.screen, BlendMode.multiply, BlendMode.subtract, BlendMode.linear_dodge],
|
|
output_dir,
|
|
layer_mode,
|
|
divide_mode
|
|
)
|
|
|
|
# base_layer_list = [cv2pil(layer) for layer in base_layer_list]
|
|
|
|
divide_folder(filename, input_dir, layer_mode)
|
|
|
|
base_layer_list = to_comfy_imgs(base_layer_list)
|
|
bright_layer_list = to_comfy_imgs(bright_layer_list)
|
|
shadow_layer_list = to_comfy_imgs(shadow_layer_list)
|
|
|
|
return (comfy_image, base_image, base_layer_list,
|
|
bright_layer_list, shadow_layer_list, filename)
|
|
|
|
elif layer_mode == "normal":
|
|
base_layer_list, bright_layer_list, shadow_layer_list = get_normal_layer(input_image, df)
|
|
filename = save_psd(
|
|
input_image,
|
|
[base_layer_list, bright_layer_list, shadow_layer_list],
|
|
["base", "bright", "shadow"],
|
|
[BlendMode.normal, BlendMode.normal, BlendMode.normal],
|
|
output_dir,
|
|
layer_mode,
|
|
divide_mode
|
|
)
|
|
|
|
divide_folder(filename, input_dir, layer_mode)
|
|
|
|
return (comfy_image, base_image, to_comfy_imgs(base_layer_list), to_comfy_imgs(bright_layer_list),
|
|
to_comfy_imgs(shadow_layer_list), filename)
|
|
else:
|
|
return None
|
|
|
|
|
|
class LayerDividerColorBase:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image1": ("IMAGE",),
|
|
"loops": ("INT", {
|
|
"default": 1,
|
|
"min": 1,
|
|
"max": 20,
|
|
"step": 1,
|
|
"display": "slider"
|
|
}),
|
|
"init_cluster": ("INT", {
|
|
"default": 10,
|
|
"min": 1,
|
|
"max": 50,
|
|
"step": 1,
|
|
"display": "slider"
|
|
}),
|
|
"ciede_threshold": ("INT", {
|
|
"default": 5,
|
|
"min": 1,
|
|
"max": 50,
|
|
"step": 1,
|
|
"display": "slider"
|
|
}),
|
|
"blur_size": ("INT", {
|
|
"default": 5,
|
|
"min": 1,
|
|
"max": 20,
|
|
"step": 1,
|
|
"display": "slider"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LD_INPUT_IMAGE", "LD_DF", "LD_DIVIDE_MODE")
|
|
RETURN_NAMES = ("input_image", "df", "divide_mode")
|
|
|
|
FUNCTION = "execute"
|
|
|
|
# OUTPUT_NODE = False
|
|
|
|
CATEGORY = "LayerDivider"
|
|
|
|
def execute(self,
|
|
image1,
|
|
loops, init_cluster, ciede_threshold, blur_size):
|
|
|
|
# Disable bg remove for now
|
|
|
|
split_bg = False
|
|
h_split = -1
|
|
v_split = -1
|
|
n_cluster = -1
|
|
alpha = -1
|
|
th_rate = 0
|
|
|
|
img_batch_np = image1.cpu().detach().numpy().__mul__(255.).astype(np.uint8)
|
|
|
|
input_image = Image.fromarray(img_batch_np[0])
|
|
|
|
image = pil2cv(input_image)
|
|
|
|
self.input_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGBA)
|
|
|
|
df = get_base(self.input_image, loops, init_cluster, ciede_threshold, blur_size, h_split, v_split, n_cluster,
|
|
alpha, th_rate, split_bg, False)
|
|
|
|
return self.input_image, df, "color_base"
|
|
|
|
|
|
class LayerDividerLoadMaskGenerator:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pred_iou_thresh": ("FLOAT", {
|
|
"default": 0.8,
|
|
"min": 0,
|
|
"max": 1,
|
|
"step": 0.01,
|
|
"display": "slider"
|
|
}),
|
|
"stability_score_thresh": ("FLOAT", {
|
|
"default": 0.8,
|
|
"min": 0,
|
|
"max": 1,
|
|
"step": 0.01,
|
|
"display": "slider"
|
|
}),
|
|
"min_mask_region_area": ("INT", {
|
|
"default": 100,
|
|
"min": 1,
|
|
"max": 1000,
|
|
"step": 1,
|
|
"display": "slider"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK_GENERATOR",)
|
|
RETURN_NAMES = ("mask_generator",)
|
|
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "LayerDivider"
|
|
|
|
def execute(self, pred_iou_thresh, stability_score_thresh, min_mask_region_area):
|
|
if not os.path.exists(model_dir):
|
|
os.makedirs(model_dir)
|
|
|
|
load_seg_model(model_dir)
|
|
|
|
mask_generator = get_mask_generator(pred_iou_thresh, stability_score_thresh, min_mask_region_area, model_dir)
|
|
|
|
return (mask_generator,)
|
|
|
|
|
|
class LayerDividerSegmentMask:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image1": ("IMAGE",),
|
|
"mask_generator": ("MASK_GENERATOR",),
|
|
"area_th": ("INT", {
|
|
"default": 20000,
|
|
"min": 1,
|
|
"max": 100000,
|
|
"step": 100,
|
|
"display": "slider"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LD_INPUT_IMAGE", "LD_DF", "LD_DIVIDE_MODE", "IMAGE")
|
|
RETURN_NAMES = ("input_image", "df", "divide_mode", "masks_preview")
|
|
|
|
FUNCTION = "execute"
|
|
|
|
# OUTPUT_NODE = False
|
|
|
|
CATEGORY = "LayerDivider"
|
|
|
|
def execute(self, image1, mask_generator, area_th):
|
|
img_batch_np = image1.cpu().detach().numpy().__mul__(255.).astype(np.uint8)
|
|
|
|
input_image = Image.fromarray(img_batch_np[0])
|
|
|
|
masks = get_masks(pil2cv(input_image), mask_generator)
|
|
|
|
masked_image = show_anns(input_image, masks, output_dir)
|
|
|
|
masked_image = to_comfy_img(np.array(masked_image))
|
|
|
|
input_image.putalpha(255)
|
|
|
|
input_image = Image.fromarray(img_batch_np[0])
|
|
|
|
image = pil2cv(input_image)
|
|
|
|
self.input_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGBA)
|
|
|
|
masks = load_masks(output_dir)
|
|
|
|
df = get_seg_base(self.input_image, masks, area_th)
|
|
|
|
return self.input_image, df, "seg_mask", masked_image
|
|
|
|
|
|
class LayerDividerDivideLayer:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"input_image": ("LD_INPUT_IMAGE",),
|
|
"df": ("LD_DF",),
|
|
"divide_mode": ("LD_DIVIDE_MODE",),
|
|
"layer_mode": (["composite", "normal"],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "STRING")
|
|
RETURN_NAMES = ("base_image", "base", "bright", "shadow", "filepath")
|
|
|
|
FUNCTION = "execute"
|
|
|
|
# OUTPUT_NODE = False
|
|
|
|
CATEGORY = "LayerDivider"
|
|
|
|
def execute(self, input_image, df, divide_mode, layer_mode):
|
|
if layer_mode == "composite":
|
|
base_layer_list, shadow_layer_list, bright_layer_list, addition_layer_list, subtract_layer_list = get_composite_layer(
|
|
input_image, df)
|
|
filename = save_psd(
|
|
input_image,
|
|
[base_layer_list, bright_layer_list, shadow_layer_list, subtract_layer_list, addition_layer_list],
|
|
["base", "screen", "multiply", "subtract", "addition"],
|
|
[BlendMode.normal, BlendMode.screen, BlendMode.multiply, BlendMode.subtract, BlendMode.linear_dodge],
|
|
output_dir,
|
|
layer_mode,
|
|
divide_mode
|
|
)
|
|
|
|
elif layer_mode == "normal":
|
|
base_layer_list, bright_layer_list, shadow_layer_list = get_normal_layer(input_image, df)
|
|
filename = save_psd(
|
|
input_image,
|
|
[base_layer_list, bright_layer_list, shadow_layer_list],
|
|
["base", "bright", "shadow"],
|
|
[BlendMode.normal, BlendMode.normal, BlendMode.normal],
|
|
output_dir,
|
|
layer_mode,
|
|
divide_mode
|
|
)
|
|
|
|
print("filename:" + filename)
|
|
|
|
divide_folder(filename, input_dir, layer_mode)
|
|
|
|
return (to_comfy_img(input_image),
|
|
to_comfy_imgs(base_layer_list),
|
|
to_comfy_imgs(bright_layer_list),
|
|
to_comfy_imgs(shadow_layer_list),
|
|
filename)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LayerDivider - Color Base": LayerDividerColorBase,
|
|
"LayerDivider - Load SAM Mask Generator": LayerDividerLoadMaskGenerator,
|
|
"LayerDivider - Segment Mask": LayerDividerSegmentMask,
|
|
"LayerDivider - Divide Layer": LayerDividerDivideLayer
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LayerDivider - Color Base": LayerDividerColorBase,
|
|
"LayerDivider - Load SAM Mask Generator": LayerDividerLoadMaskGenerator,
|
|
"LayerDivider - Segment Mask": LayerDividerSegmentMask,
|
|
"LayerDivider - Divide Layer": LayerDividerDivideLayer
|
|
}
|