203 lines
4.2 KiB
Python
203 lines
4.2 KiB
Python
import numpy as np
|
|
import cv2
|
|
import math
|
|
import torch
|
|
|
|
|
|
def from_torch_image(image):
|
|
image = image.cpu().numpy() * 255.0
|
|
image = np.clip(image, 0, 255).astype(np.uint8)
|
|
return image
|
|
|
|
|
|
def to_torch_image(image):
|
|
image = image.astype(dtype=np.float32)
|
|
image /= 255.0
|
|
image = torch.from_numpy(image)
|
|
return image
|
|
|
|
|
|
def smooth_step_plain(x):
|
|
|
|
if x < -1:
|
|
return -1
|
|
elif x <= 1:
|
|
return math.sin(x * np.pi / 2.0)
|
|
else:
|
|
return 1
|
|
|
|
|
|
def smooth_step_np(x):
|
|
truths = np.logical_and(-1 < x, x < 1).astype(np.float32)
|
|
x1 = np.clip(x, -1, 1)
|
|
x2 = np.sin(x * np.pi / 2.0)
|
|
ret = (truths * x2) + ((1 - truths) * x1)
|
|
return ret
|
|
|
|
|
|
def smooth_step_stretch(x, a, b):
|
|
|
|
if b < a:
|
|
tmp = b
|
|
b = a
|
|
a = tmp
|
|
|
|
if a < 0:
|
|
a = 0
|
|
|
|
if b > 1:
|
|
b = 1
|
|
|
|
if a == b:
|
|
a = 0
|
|
b = 1
|
|
|
|
return smooth_step_np((2 * (x - a) / (b - a)) - 1)
|
|
|
|
|
|
def get_light_layer(image,
|
|
ref_r=255,
|
|
ref_g=255,
|
|
ref_b=255,
|
|
do_scale=True,
|
|
scale_a=0.0,
|
|
scale_b=1.0):
|
|
|
|
sqmax = 3 * 255 * 255
|
|
scalemax = math.sqrt(sqmax)
|
|
|
|
b = image[:, :, 0].astype(dtype=np.float32)
|
|
g = image[:, :, 1].astype(dtype=np.float32)
|
|
r = image[:, :, 2].astype(dtype=np.float32)
|
|
|
|
b2 = b * b
|
|
g2 = g * g
|
|
r2 = r * r
|
|
|
|
d2 = np.zeros(b2.shape, dtype=np.float32)
|
|
d2 += sqmax - b2 - g2 - r2
|
|
d = np.sqrt(d2)
|
|
|
|
ref_r2 = ref_r * ref_r
|
|
ref_g2 = ref_g * ref_g
|
|
ref_b2 = ref_b * ref_b
|
|
ref_d2 = sqmax - ref_r2 - ref_g2 - ref_b2
|
|
|
|
ref_d = math.sqrt(ref_d2)
|
|
|
|
dot = (b * ref_b) + (g * ref_g) + (r * ref_r) + (d * ref_d)
|
|
dot /= sqmax
|
|
|
|
if do_scale:
|
|
dot = smooth_step_stretch(x=dot, a=scale_a, b=scale_b)
|
|
|
|
dot *= 255
|
|
dot = dot.astype(np.uint8)
|
|
|
|
return dot
|
|
|
|
|
|
class main_light_layer():
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", ),
|
|
"ref_r": ("INT", {
|
|
"default": 255,
|
|
"min": 0,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"ref_g": ("INT", {
|
|
"default": 255,
|
|
"min": 0,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"ref_b": ("INT", {
|
|
"default": 255,
|
|
"min": 0,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"do_scale": (["enable", "disable"], ),
|
|
"thresh_low": ("FLOAT", {
|
|
"default": 0.6,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"thresh_high": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
},
|
|
}
|
|
|
|
FUNCTION = "run"
|
|
RETURN_TYPES = ("MASK", )
|
|
CATEGORY = "HackNode"
|
|
|
|
def run(
|
|
self,
|
|
image,
|
|
ref_r,
|
|
ref_g,
|
|
ref_b,
|
|
do_scale,
|
|
thresh_low,
|
|
thresh_high,
|
|
):
|
|
|
|
do_scale = (do_scale == "enable")
|
|
print('do_scale', do_scale)
|
|
|
|
image = from_torch_image(image)
|
|
print('image.shape', image.shape)
|
|
|
|
batch_size = image.shape[0]
|
|
print('batch_size', batch_size)
|
|
|
|
mask = []
|
|
|
|
for i in range(batch_size):
|
|
|
|
tmp_img = image[i]
|
|
print('tmp_img.shape', tmp_img.shape)
|
|
|
|
tmp_mask = get_light_layer(
|
|
tmp_img,
|
|
ref_b,
|
|
ref_g,
|
|
ref_r,
|
|
do_scale,
|
|
scale_a=thresh_low,
|
|
scale_b=thresh_high,
|
|
)
|
|
print('tmp_mask.shape', tmp_mask.shape)
|
|
|
|
mask.append(tmp_mask)
|
|
|
|
mask = np.array(mask)
|
|
|
|
mask = to_torch_image(mask)
|
|
print(mask.shape)
|
|
|
|
return (mask, )
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
'main_light_layer': main_light_layer,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
'main_light_layer': 'main_light_layer',
|
|
}
|