From d1aaa4a0975b28e71aaaad7a1eeb427c5f03f51c Mon Sep 17 00:00:00 2001 From: "Dr.Lt.Data" Date: Sat, 23 Dec 2023 11:43:38 +0900 Subject: [PATCH] feat: Gaussian Blur Mask, Dilate Mask (SEGS), Gaussian Blur Mask (SEGS) refactor: feather/gaussian blur fix: Mask To SEGS - invalid conversion when mask is donut shape fix: SEGS Paste/SEGS Preview - proper alpha handling --- README.md | 21 ++++++---- __init__.py | 6 +++ modules/impact/config.py | 2 +- modules/impact/core.py | 8 +++- modules/impact/impact_pack.py | 2 +- modules/impact/segs_nodes.py | 76 ++++++++++++++++++++++++++++++++++- modules/impact/utils.py | 16 +++++--- 7 files changed, 112 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index c312960..4a5d778 100644 --- a/README.md +++ b/README.md @@ -86,15 +86,20 @@ This custom node helps to conveniently enhance images through Detector, Detailer * SEGS Filter (range) - This node retrieves only SEGs from SEGS that have a size and position within a certain range. * SEGSConcat - Concatenate segs1 and segs2. If source shape of segs1 and segs2 are different from segs2 will be ignored. * Picker (SEGS) - Among the input SEGS, you can select a specific SEG through a dialog. If no SEG is selected, it outputs an empty SEGS. Increasing the batch_size of SEGSDetailer can be used for the purpose of selecting from the candidates. - * Set Default Image For SEGS - Set a default image for SEGS. SEGS with images set this way do not need to have a fallback image set. When override is set to false, the original image is preserved. - * DecomposeSEGS - Decompose SEGS to allow for detailed manipulation. - * AssembleSEGS - Reassemble the decomposed SEGS. - * From SEG_ELT - Extract detailed information from SEG_ELT. - * Edit SEG_ELT - Modify some of the information in SEG_ELT. - * Dilate SEG_ELT - Dilate the mask of SEG_ELT. + * Set Default Image For SEGS - Set a default image for SEGS. SEGS with images set this way do not need to have a fallback image set. When override is set to false, the original image is preserved. + * Dilate Mask (SEGS) - Dilate/Erosion Mask in SEGS + * Gaussian Blur Mask (SEGS) - Apply Gaussian Blur to Mask in SEGS + * SEGS_ELT Manipulation - experimental nodes + * DecomposeSEGS - Decompose SEGS to allow for detailed manipulation. + * AssembleSEGS - Reassemble the decomposed SEGS. + * From SEG_ELT - Extract detailed information from SEG_ELT. + * Edit SEG_ELT - Modify some of the information in SEG_ELT. + * Dilate SEG_ELT - Dilate the mask of SEG_ELT. -* Dilate Mask - Dilate Mask. - * Support erosion for negative value. +* Mask Manipulation + * Dilate Mask - Dilate Mask. + * Support erosion for negative value. + * Gaussian Blur Mask - Apply Gaussian Blur to Mask. You can utilize this for mask feathering. * Pipe nodes * ToDetailerPipe, FromDetailerPipe - These nodes are used to bundle multiple inputs used in the detailer, such as models and vae, ..., into a single DETAILER_PIPE or extract the elements that are bundled in the DETAILER_PIPE. diff --git a/__init__.py b/__init__.py index 4dff234..33983d1 100644 --- a/__init__.py +++ b/__init__.py @@ -213,6 +213,9 @@ NODE_CLASS_MAPPINGS = { "ImpactEdit_SEG_ELT": Edit_SEG_ELT, "ImpactDilate_Mask_SEG_ELT": Dilate_SEG_ELT, "ImpactDilateMask": DilateMask, + "ImpactGaussianBlurMask": GaussianBlurMask, + "ImpactDilateMaskInSEGS": DilateMaskInSEGS, + "ImpactGaussianBlurMaskInSEGS": GaussianBlurMaskInSEGS, "ImpactScaleBy_BBOX_SEG_ELT": SEG_ELT_BBOX_ScaleBy, "BboxDetectorCombined_v2": BboxDetectorCombined, @@ -375,6 +378,9 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ImpactDilate_Mask_SEG_ELT": "Dilate Mask (SEG_ELT)", "ImpactScaleBy_BBOX_SEG_ELT": "ScaleBy BBOX (SEG_ELT)", "ImpactDilateMask": "Dilate Mask", + "ImpactGaussianBlurMask": "Gaussian Blur Mask", + "ImpactDilateMaskInSEGS": "Dilate Mask (SEGS)", + "ImpactGaussianBlurMaskInSEGS": "Gaussian Blur Mask (SEGS)", "PreviewBridge": "Preview Bridge (Image)", "PreviewBridgeLatent": "Preview Bridge (Latent)", diff --git a/modules/impact/config.py b/modules/impact/config.py index b0781f8..0571e71 100644 --- a/modules/impact/config.py +++ b/modules/impact/config.py @@ -2,7 +2,7 @@ import configparser import os -version = "V4.48.12" +version = "V4.49" dependency_version = 19 diff --git a/modules/impact/core.py b/modules/impact/core.py index d16f055..c750d33 100644 --- a/modules/impact/core.py +++ b/modules/impact/core.py @@ -971,8 +971,12 @@ def mask_to_segs(mask, combined, crop_factor, bbox_fill, drop_size=1, label='A', else: mask_i_uint8 = (mask_i * 255.0).astype(np.uint8) - contours, _ = cv2.findContours(mask_i_uint8, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) - for contour in contours: + contours, ctree = cv2.findContours(mask_i_uint8, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) + for j, contour in enumerate(contours): + hierarchy = ctree[0][j] + if hierarchy[3] != -1: + continue + separated_mask = np.zeros_like(mask_i_uint8) cv2.drawContours(separated_mask, [contour], 0, 255, -1) separated_mask = np.array(separated_mask / 255.0).astype(np.float32) diff --git a/modules/impact/impact_pack.py b/modules/impact/impact_pack.py index 9829119..3bc01c9 100644 --- a/modules/impact/impact_pack.py +++ b/modules/impact/impact_pack.py @@ -207,7 +207,7 @@ class DetailerForEach: else crop_ndarray4(image.numpy(), seg.crop_region) cropped_image = to_tensor(cropped_image) mask = to_tensor(seg.cropped_mask) - mask = tensor_feather_mask(mask, feather) + mask = tensor_gaussian_blur_mask(mask, feather) is_mask_all_zeros = (seg.cropped_mask == 0).all().item() if is_mask_all_zeros: diff --git a/modules/impact/segs_nodes.py b/modules/impact/segs_nodes.py index e7b4967..c851750 100644 --- a/modules/impact/segs_nodes.py +++ b/modules/impact/segs_nodes.py @@ -238,7 +238,7 @@ class SEGSPaste: ref_tensor = ref_image_opt[i].unsqueeze(0) ref_image = crop_image(ref_tensor, seg.crop_region) if ref_image is not None: - mask = tensor_feather_mask(seg.cropped_mask, feather, alpha/255) + mask = tensor_gaussian_blur_mask(seg.cropped_mask, feather) * (alpha/255) x, y, *_ = seg.crop_region tensor_paste(image_i, ref_image, (x, y), mask) @@ -305,7 +305,7 @@ class SEGSPreview: cropped_image = to_pil(cropped_image) if alpha_mode: - mask_array = seg.cropped_mask.astype(np.uint8) * 255 + mask_array = (seg.cropped_mask * 255).astype(np.uint8) mask_image = Image.fromarray(mask_array, mode='L').resize(cropped_image.size) cropped_image.putalpha(mask_image) @@ -754,6 +754,78 @@ class DilateMask: return (torch.from_numpy(mask), ) +class GaussianBlurMask: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "mask": ("MASK", ), + "kernel_size": ("INT", {"default": 10, "min": 0, "max": 100, "step": 1}), + "sigma": ("FLOAT", {"default": 10.0, "min": 0.1, "max": 100.0, "step": 0.1}), + }} + + RETURN_TYPES = ("MASK", ) + + FUNCTION = "doit" + + CATEGORY = "ImpactPack/Util" + + def doit(self, mask, kernel_size, sigma): + mask = torch.unsqueeze(mask, dim=-1) + mask = utils.tensor_gaussian_blur_mask(mask, kernel_size, sigma) + mask = torch.squeeze(mask, dim=-1) + return (mask, ) + + +class DilateMaskInSEGS: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "segs": ("SEGS", ), + "dilation": ("INT", {"default": 10, "min": -512, "max": 512, "step": 1}), + }} + + RETURN_TYPES = ("SEGS", ) + + FUNCTION = "doit" + + CATEGORY = "ImpactPack/Util" + + def doit(self, segs, dilation): + new_segs = [] + for seg in segs[1]: + mask = core.dilate_mask(seg.cropped_mask, dilation) + seg = SEG(seg.cropped_image, mask, seg.confidence, seg.crop_region, seg.bbox, seg.label, seg.control_net_wrapper) + new_segs.append(seg) + + return ((segs[0], new_segs), ) + + +class GaussianBlurMaskInSEGS: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "segs": ("SEGS", ), + "kernel_size": ("INT", {"default": 10, "min": 0, "max": 100, "step": 1}), + "sigma": ("FLOAT", {"default": 10.0, "min": 0.1, "max": 100.0, "step": 0.1}), + }} + + RETURN_TYPES = ("SEGS", ) + + FUNCTION = "doit" + + CATEGORY = "ImpactPack/Util" + + def doit(self, segs, kernel_size, sigma): + new_segs = [] + for seg in segs[1]: + mask = utils.tensor_gaussian_blur_mask(seg.cropped_mask, kernel_size, sigma) + mask = torch.squeeze(mask, dim=-1).squeeze(0).numpy() + seg = SEG(seg.cropped_image, mask, seg.confidence, seg.crop_region, seg.bbox, seg.label, seg.control_net_wrapper) + new_segs.append(seg) + + return ((segs[0], new_segs), ) + + class Dilate_SEG_ELT: @classmethod def INPUT_TYPES(s): diff --git a/modules/impact/utils.py b/modules/impact/utils.py index 59bbef7..c52538e 100644 --- a/modules/impact/utils.py +++ b/modules/impact/utils.py @@ -332,7 +332,7 @@ def _gaussian_kernel(kernel_size, sigma): return kernel / kernel.sum() -def tensor_feather_mask(mask, thickness, base_alpha=1.0): +def tensor_gaussian_blur_mask(mask, kernel_size, sigma=10.0): """Return NHWC torch.Tenser from ndim == 2 or 4 `np.ndarray` or `torch.Tensor`""" if isinstance(mask, np.ndarray): mask = torch.from_numpy(mask) @@ -341,14 +341,20 @@ def tensor_feather_mask(mask, thickness, base_alpha=1.0): mask = mask[None, ..., None] _tensor_check_mask(mask) - if thickness <= 0: + if kernel_size <= 0: return mask - # Create a feathered mask by applying a Gaussian blur to the mask - mask = mask[:, None, ..., 0] + prev_device = mask.device + device = comfy.model_management.get_torch_device() + mask.to(device) - blurred_mask = torchvision.transforms.GaussianBlur(kernel_size=thickness*2+1, sigma=10.0)(mask) + # apply gaussian blur + mask = mask[:, None, ..., 0] + blurred_mask = torchvision.transforms.GaussianBlur(kernel_size=kernel_size*2+1, sigma=sigma)(mask) blurred_mask = blurred_mask[:, 0, ..., None] + + blurred_mask.to(prev_device) + return blurred_mask