151 lines
5.2 KiB
Python
151 lines
5.2 KiB
Python
import math, cv2, random, torch, torchvision, json, os
|
|
import numpy as np
|
|
import nodes, folder_paths # 기본노드, 파일로드
|
|
|
|
# class name:
|
|
# def __init__(self):
|
|
# pass
|
|
|
|
# @classmethod
|
|
# def INPUT_TYPES(s):
|
|
# return {
|
|
# "required": {},
|
|
# "optional": {},
|
|
# }
|
|
# RETURN_TYPES = ()
|
|
# RETURN_NAMES = ()
|
|
|
|
# FUNCTION = "run"
|
|
|
|
# CATEGORY = "abyz22"
|
|
|
|
# def run(self,*args,**kwargs):
|
|
# return None
|
|
|
|
|
|
def find_white_bbox(mask):
|
|
mask = (mask > 0).astype(np.uint8)
|
|
# 행과 열의 합을 계산하여 흰색 영역을 찾습니다.
|
|
row_sum = np.sum(mask, axis=1) # 768
|
|
col_sum = np.sum(mask, axis=0) # 512
|
|
|
|
# 흰색 영역의 시작과 끝 인덱스를 찾습니다.
|
|
for i, x in enumerate(row_sum):
|
|
if x > 0:
|
|
break
|
|
for ii, x in enumerate(col_sum):
|
|
if x > 0:
|
|
break
|
|
y1, x1 = i, ii
|
|
|
|
for i, x in enumerate(row_sum[::-1]):
|
|
if x > 0:
|
|
break
|
|
for ii, x in enumerate(col_sum[::-1]):
|
|
if x > 0:
|
|
break
|
|
y2, x2 = row_sum.shape[0] - i, col_sum.shape[0] - ii
|
|
return y1, x1, y2, x2
|
|
|
|
|
|
class abyz22_smallhead:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",), # image= 1,768,512,3
|
|
"segs": ("IMAGE",), # mask= 1,768,512,3
|
|
"ratio": ("FLOAT", {"default": 0.8, "min": 0.01, "max": 1.5, "step": 0.01, "round": 0.001, "dispaly": "slider"}),
|
|
"empty_space": (["noise", "black"],),
|
|
"model_detector": ("BOOLEAN", {"default": True, "label_on": "face", "label_off": "head"}),
|
|
"kernel_size": ("INT", {"default": 10, "min": 0, "max": 100, "step": 1, "dispaly": "slider"}),
|
|
"sigma": ("FLOAT", {"default": 10.0, "min": 0.1, "max": 100.0, "step": 0.1, "round": 0.01, "dispaly": "slider"}),
|
|
"mask_dilation": ("INT", {"default": 10, "min": 0, "max": 128, "step": 1, "dispaly": "slider"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"IMAGE",
|
|
"MASK",
|
|
)
|
|
RETURN_NAMES = (
|
|
"images",
|
|
"masks",
|
|
)
|
|
|
|
FUNCTION = "run"
|
|
|
|
CATEGORY = "abyz22"
|
|
|
|
def run(sefl, *args, **kwargs): # image= 1,768,512,3 mask= n,1,h,w
|
|
images, mask = None, None
|
|
if len(kwargs["segs"]) == 1:
|
|
kwargs["segs"] = kwargs["segs"][0]
|
|
for i in range(kwargs["image"].shape[0]):
|
|
# y1, x1, y2, x2 = find_white_bbox(kwargs["segs"].permute(0, 3, 1, 2)[i][0].numpy())
|
|
y1, x1, y2, x2 = find_white_bbox(kwargs["segs"][i].permute(2, 0, 1)[0].numpy())
|
|
w = x2 - x1
|
|
h = y2 - y1
|
|
|
|
image = kwargs["image"][i]
|
|
face = kwargs["image"][i].clone()
|
|
|
|
face_cropped = face * kwargs["segs"][i]
|
|
if kwargs["empty_space"] == "noise":
|
|
image[y1:y2, x1:x2, :] = torch.rand_like(image[y1:y2, x1:x2, :])
|
|
face_noised = torch.rand_like(face) * -1 * (kwargs["segs"] - 1)[i]
|
|
elif kwargs["empty_space"] == "black":
|
|
image[y1:y2, x1:x2, :] = torch.zeros_like(image[y1:y2, x1:x2, :])
|
|
face_noised = torch.zeros_like(face) * -1 * (kwargs["segs"] - 1)[i]
|
|
face_cropped = face_cropped + face_noised
|
|
|
|
face_cropped = face_cropped[y1:y2, x1:x2, :].permute(2, 0, 1) # 768,512,3 -> 3,768,512
|
|
face_resized = torchvision.transforms.Resize((int(h * kwargs["ratio"]), int(w * kwargs["ratio"])))(face_cropped).permute(
|
|
1, 2, 0
|
|
) # 3,768,512 -> 768,512,3
|
|
|
|
h, w, c = face_resized.shape
|
|
|
|
image[y2 - h : y2, int((x1 + x2) // 2 - w // 2) : int((x1 + x2) // 2 - w // 2) + w, :] = face_resized
|
|
|
|
# mask 출력
|
|
mask = torch.zeros_like(kwargs["segs"][i, :, :, 0])
|
|
|
|
mask_face = kwargs["segs"][i][y1:y2, x1:x2, :].permute(2, 0, 1)
|
|
mask_face = torchvision.transforms.Resize((int(h * kwargs["ratio"]), int(w * kwargs["ratio"])))(mask_face).permute(1, 2, 0)
|
|
|
|
h, w, c = mask_face.shape
|
|
mask[y2 - h : y2, int((x1 + x2) // 2 - w // 2) : int((x1 + x2) // 2 - w // 2) + w] = mask_face[:, :, 0]
|
|
mask = -1 * mask + 1
|
|
mask[:y1] = 0
|
|
mask[y2:] = 0
|
|
mask[:, :x1] = 0
|
|
mask[:, x2:] = 0
|
|
|
|
kernel = np.ones((kwargs["mask_dilation"], kwargs["mask_dilation"]), np.uint8)
|
|
mask = cv2.dilate(mask.numpy(), kernel, iterations=1)
|
|
mask = torch.tensor(mask)
|
|
mask = mask.unsqueeze(0)
|
|
|
|
if kwargs["model_detector"]:
|
|
mask[:, :y1, x1:x2] = 1
|
|
|
|
if i == 0:
|
|
images = image.unsqueeze(0)
|
|
masks = mask
|
|
else:
|
|
images = torch.cat([images, image.unsqueeze(0)])
|
|
masks = torch.cat([masks, mask])
|
|
|
|
# mask= n,1,h,w
|
|
masks = torchvision.transforms.GaussianBlur(kernel_size=kwargs["kernel_size"] * 2 + 1, sigma=kwargs["sigma"])(masks)
|
|
if masks.shape[0] ==1:
|
|
masks=masks.unsqueeze(0)
|
|
return (
|
|
images,
|
|
masks,
|
|
) # image shoulde be n,768,512,3
|