Mask.py

 
import scipy.ndimage
import torch

import numpy as np 
# from PIL import Image, ImageDraw
from PIL import Image, ImageOps 

from comfy.cli_args import args
import cv2,os
from nodes import MAX_RESOLUTION, SaveImage, common_ksampler
import folder_paths,random

# Tensor to PIL
def tensor2pil(image):
    return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))

# Convert PIL to Tensor
def pil2tensor(image):
    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)


def add_masks(mask1, mask2):
    mask1 = mask1.cpu()
    mask2 = mask2.cpu()
    cv2_mask1 = np.array(mask1) * 255
    cv2_mask2 = np.array(mask2) * 255

    if cv2_mask1.shape == cv2_mask2.shape:
        cv2_mask = cv2.add(cv2_mask1, cv2_mask2)
        return torch.clamp(torch.from_numpy(cv2_mask) / 255.0, min=0, max=1)
    else:
        return mask1
    

def grow(mask, expand, tapered_corners):
    c = 0 if tapered_corners else 1
    kernel = np.array([[c, 1, c],
                            [1, 1, 1],
                            [c, 1, c]])
    mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1]))
    out = []
    for m in mask:
        output = m.numpy()
        for _ in range(abs(expand)):
            if expand < 0:
                output = scipy.ndimage.grey_erosion(output, footprint=kernel)
            else:
                output = scipy.ndimage.grey_dilation(output, footprint=kernel)
        output = torch.from_numpy(output)
        out.append(output)
    return torch.stack(out, dim=0)

def combine(destination, source, x, y):
    output = destination.reshape((-1, destination.shape[-2], destination.shape[-1])).clone()
    source = source.reshape((-1, source.shape[-2], source.shape[-1]))

    left, top = (x, y,)
    right, bottom = (min(left + source.shape[-1], destination.shape[-1]), min(top + source.shape[-2], destination.shape[-2]))
    visible_width, visible_height = (right - left, bottom - top,)

    source_portion = source[:, :visible_height, :visible_width]
    destination_portion = destination[:, top:bottom, left:right]
 
    #operation == "subtract":
    output[:, top:bottom, left:right] = destination_portion - source_portion
        
    output = torch.clamp(output, 0.0, 1.0)

    return output


class PreviewMask_(SaveImage):
    def __init__(self):
        self.output_dir = folder_paths.get_temp_directory()
        self.type = "temp"
        self.prefix_append =''.join(random.choice("abcdehijklmnopqrstupvxyzfg") for x in range(5))
        self.compress_level = 4

    @classmethod
    def INPUT_TYPES(s):
        return {
                "required": {
                                "mask": ("MASK",), 
                            }
            }
    
    RETURN_TYPES = ()
    OUTPUT_NODE = True
    FUNCTION = "run"
    CATEGORY = "♾️Mixlab/Mask"

    # 运行的函数
    def run(self, mask ):
        img=tensor2pil(mask)
        img=img.convert('RGB')
        img=pil2tensor(img) 
        return self.save_images(img, 'temp_', None,  None)


class OutlineMask:

    @classmethod
    def INPUT_TYPES(s):
        return {
                "required": {
                                "mask": ("MASK",),
                                "outline_width":("INT", {"default": 10,"min": 1, "max": MAX_RESOLUTION, "step": 1}),
                                "tapered_corners": ("BOOLEAN", {"default": True}),
                            }
            }
    
    RETURN_TYPES = ('MASK',)

    FUNCTION = "run"

    CATEGORY = "♾️Mixlab/Mask"

    # 运行的函数
    def run(self, mask, outline_width, tapered_corners):

        m1=grow(mask,outline_width,tapered_corners)
        m2=grow(mask,-outline_width,tapered_corners)

        m3=combine(m1,m2,0,0)

        return (m3,)
    

class MaskListReplace:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "masks": ("MASK",),
            "mask_replace": ("MASK",),
            "start_index":("INT", {"default": 0, "min": 0, "step": 1}),
            "end_index":("INT", {"default": 0, "min": 0, "step": 1}),
            "invert": ("BOOLEAN", {"default": False}),
        }
        }

    RETURN_TYPES = ("MASK",)
    FUNCTION = "run"
    CATEGORY = "♾️Mixlab/Video"

    INPUT_IS_LIST = True
    OUTPUT_IS_LIST = (True,)

    def run(self, masks,mask_replace,start_index,end_index,invert):
        mask_replace=mask_replace[0]
        start_index=start_index[0]
        end_index=end_index[0]
        invert=invert[0]
        
        new_masks=[]
        for i in range(len(masks)):
            if i>=start_index and i<=end_index:
                if invert:
                    new_masks.append(masks[i])
                else:
                    new_masks.append(mask_replace)
            else:
                if invert:
                    new_masks.append(mask_replace)
                else:
                    new_masks.append(masks[i])

        return (new_masks,)


class MaskListMerge:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "masks": ("MASK",),
        }
        }

    RETURN_TYPES = ("MASK",)
    FUNCTION = "run"
    CATEGORY = "♾️Mixlab/Mask"

    INPUT_IS_LIST = True
    OUTPUT_IS_LIST = (False,)

    def run(self, masks):
        mask=masks[0]
        if isinstance(masks, list):
            for m in masks:
                # print(m.shape)
                mask = add_masks(mask, m)
        return (mask,)


class FeatheredMask:
    @classmethod
    def INPUT_TYPES(s):
        return {
                "required": {
                                "mask": ("MASK",),
                                "start_offset":("INT", {"default": 1, 
                                                        "min": -150, 
                                                        "max": 150, 
                                                        "step": 1,
                                                        "display": "slider"}),
                                "feathering_weight":("FLOAT", {"default": 0.1,
                                                                "min": 0.0,
                                                                "max": 1,
                                                                "step": 0.1,
                                                                "display": "slider"})
                            }
            }
    
    RETURN_TYPES = ('MASK',)

    FUNCTION = "run"

    CATEGORY = "♾️Mixlab/Mask"

    OUTPUT_IS_LIST = (True,)
  
    # 运行的函数
    def run(self,mask,start_offset, feathering_weight):
        # print(mask.shape,mask.size())
        
        num,_,_=mask.size()

        masks=[]

        for i in range(num):
            mm=mask[i]
            image=tensor2pil(mm)

            # Open the image using PIL
            image = image.convert("L")
            if start_offset>0:
                image=ImageOps.invert(image)

            # Convert the image to a numpy array
            image_np = np.array(image)

            # Use Canny edge detection to get black contours
            edges = cv2.Canny(image_np, 30, 150)

            for i in range(0,abs(start_offset)):
                # int(100*feathering_weight)
                a=int(abs(start_offset)*0.1*i)
                # Dilate the black contours to make them wider
                kernel = np.ones((a, a), np.uint8)

                dilated_edges = cv2.dilate(edges, kernel, iterations=1)
                # dilated_edges = cv2.erode(edges, kernel, iterations=1)
                # Smooth the dilated edges using Gaussian blur
                smoothed_edges = cv2.GaussianBlur(dilated_edges, (5, 5), 0)

                # Adjust the feathering weight
                feathering_weight = max(0, min(feathering_weight, 1))

                # Blend the smoothed edges with the original image to achieve feathering effect
                image_np = cv2.addWeighted(image_np, 1, smoothed_edges, feathering_weight, feathering_weight)

            # Convert the result back to PIL image
            result_image = Image.fromarray(np.uint8(image_np))
            result_image=result_image.convert("L")

            if start_offset>0:
                result_image=ImageOps.invert(result_image)
            
            result_image=result_image.convert("L")
            mt=pil2tensor(result_image)
            masks.append(mt)
         
            # print( mt.size())
        return (masks,)