diff --git a/image_c_postprocessor.py b/image_c_postprocessor.py index 54bce4b..bc169ec 100644 --- a/image_c_postprocessor.py +++ b/image_c_postprocessor.py @@ -140,6 +140,95 @@ class ImageCPostprocessor: return (tensors, ) +class PrivateImageMask: + def __init__(self): + pass + # def apply_mask_with_opencv(self, input_image, mask_image, mask_gray_threshold): + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "input_image": ("IMAGE", ), + "mask_image": ("IMAGE", ), + "mask_gray_threshold": ("INT", {"default": 127}) + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "apply_postprocess" + CATEGORY = "Tools" + + def apply_mask_with_opencv(self, input_image, mask_image, mask_gray_threshold): + """ + 使用OpenCV和NumPy对输入的PIL图像应用掩膜,并返回处理后的PIL图像。 + + Args: + input_image (PIL.Image): 被遮罩的PIL图像。 + mask_image (PIL.Image): 用作遮罩的PIL图像。 + + Returns: + PIL.Image: 处理后的PIL图像,掩膜区域变为透明。 + """ + # 将PIL图像转换为OpenCV图像格式 + input_cv = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGBA2BGRA) + mask_cv = cv2.cvtColor(np.array(mask_image), cv2.COLOR_RGB2GRAY) + + # 确保掩膜是单通道的 + if len(mask_cv.shape) == 3: + mask_cv = cv2.cvtColor(mask_cv, cv2.COLOR_BGR2GRAY) + + # 将掩膜二值化,阈值可以根据需要调整 + _, binary_mask = cv2.threshold(mask_cv, mask_gray_threshold, 255, cv2.THRESH_BINARY) + + # 创建一个新的只有alpha通道的图像,使用掩膜作为alpha通道 + alpha_channel = cv2.bitwise_not(binary_mask) + b, g, r, a = cv2.split(input_cv) + bgra = [b, g, r, alpha_channel] + + # 合并通道生成最终图像 + result_cv = cv2.merge(bgra) + + # 将OpenCV图像转换回PIL图像 + result_pil = Image.fromarray(cv2.cvtColor(result_cv, cv2.COLOR_BGRA2RGBA)) + + return result_pil + + def apply_postprocess(self, image, mask_image, mask_gray_threshold): + tensors = [] + if len(image) > 1: + for img in image: + + pil_image = None + # PIL Image + pil_image = tensor2pil(img) + + # Apply Fliter + new_img = self.apply_mask_with_opencv(pil_image, mask_image, mask_gray_threshold) + + # Output image + out_image = (pil2tensor(new_img) if pil_image else img) + + tensors.append(out_image) + + tensors = torch.cat(tensors, dim=0) + + else: + pil_image = None + img = image + # PIL Image + + pil_image = tensor2pil(img) + + # Apply Fliter + new_img = self.apply_mask_with_opencv(pil_image, mask_image, mask_gray_threshold) + + # Output image + out_image = (pil2tensor(new_img) if pil_image else img) + + tensors = out_image + + return (tensors, ) # Set the web directory, any .js file in that directory will be loaded by the frontend as a frontend extension # WEB_DIRECTORY = "./somejs" @@ -148,10 +237,12 @@ class ImageCPostprocessor: # A dictionary that contains all nodes you want to export with their names # NOTE: names should be globally unique NODE_CLASS_MAPPINGS = { - "ImageCPostprocessor": ImageCPostprocessor + "ImageCPostprocessor": ImageCPostprocessor, + "PrivateImageMask": PrivateImageMask } # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { - "ImageCPostprocessor": "Private ImageCPostprocessor" + "ImageCPostprocessor": "Private ImageCPostprocessor", + "PrivateImageMask": "Private Image Mask" }