TheArtOfficial-ComfyUI-MaskMorph/nodes.py

import math
from typing import Dict

import torch
import torch.nn.functional as F


class DirectionalMaskExpand:
    @classmethod
    def INPUT_TYPES(cls) -> Dict:
        return {
            "required": {
                "mask": ("MASK", {}),
                "angle_deg": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 360.0, "step": 1.0}),
                "pixels": ("INT", {"default": 0, "min": 0, "max": 2048, "step": 1}),
            }
        }

    RETURN_TYPES = ("MASK",)
    FUNCTION = "expand"
    CATEGORY = "mask"

    @staticmethod
    def _shift(mask: torch.Tensor, dy: int, dx: int) -> torch.Tensor:
        """Shift mask by integer offsets without wrapping."""
        shifted = torch.zeros_like(mask)
        _, h, w = mask.shape

        if dy >= 0:
            src_y0, src_y1 = 0, h - dy
            dst_y0, dst_y1 = dy, h
        else:
            src_y0, src_y1 = -dy, h
            dst_y0, dst_y1 = 0, h + dy

        if dx >= 0:
            src_x0, src_x1 = 0, w - dx
            dst_x0, dst_x1 = dx, w
        else:
            src_x0, src_x1 = -dx, w
            dst_x0, dst_x1 = 0, w + dx

        if src_y1 <= src_y0 or src_x1 <= src_x0:
            return shifted

        shifted[:, dst_y0:dst_y1, dst_x0:dst_x1] = mask[:, src_y0:src_y1, src_x0:src_x1]
        return shifted

    def expand(
        self,
        mask: torch.Tensor,
        angle_deg: float,
        pixels: int,
    ):
        if mask.dim() == 2:
            mask = mask.unsqueeze(0)

        expanded = mask

        # Directional expansion along a vector.
        if pixels > 0:
            angle_rad = math.radians(angle_deg % 360.0)
            dir_x = math.sin(angle_rad)
            dir_y = -math.cos(angle_rad)  # 0 degrees = up

            offsets = []
            for dist in range(1, pixels + 1):
                ox = int(round(dir_x * dist))
                oy = int(round(dir_y * dist))
                offsets.append((oy, ox))

            offsets = list(dict.fromkeys(offsets))  # deduplicate

            for oy, ox in offsets:
                if oy == 0 and ox == 0:
                    continue
                shifted = self._shift(mask, oy, ox)
                expanded = torch.maximum(expanded, shifted)

        # Apply tapering warps after expansion (replace, not union).
        return (expanded.clamp(0.0, 1.0), )

    @classmethod
    def IS_CHANGED(
        cls,
        angle_deg: float,
        pixels: int,
        **kwargs,
    ):
        # Deterministic value prevents re-runs when the node inputs are unchanged.
        return hash(
            (
                round(angle_deg, 3),
                int(pixels),
            )
        )


class MaskPinchWiden:
    @classmethod
    def INPUT_TYPES(cls) -> Dict:
        return {
            "required": {
                "mask": ("MASK", {}),
                "taper_amount": ("INT", {"default": 0, "min": 0, "max": 2048, "step": 1}),
                "vertical_pinch": ("BOOLEAN", {"default": False}),
                "horizontal_pinch": ("BOOLEAN", {"default": False}),
                "vertical_widen": ("BOOLEAN", {"default": False}),
                "horizontal_widen": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ("MASK",)
    FUNCTION = "taper"
    CATEGORY = "mask"

    @staticmethod
    def _bbox(mask: torch.Tensor):
        b, h, w = mask.shape
        coords = (mask > 0).nonzero()
        if coords.numel() == 0:
            return 0, 0, 0, 0
        _, ys, xs = coords[:,0], coords[:,1], coords[:,2]
        y0 = int(ys.min())
        y1 = int(ys.max()) + 1
        x0 = int(xs.min())
        x1 = int(xs.max()) + 1
        return y0, y1, x0, x1

    @staticmethod
    def _pinch_widen(mask: torch.Tensor, amount: int, axis: str, mode: str) -> torch.Tensor:
        """Morph inside the mask bbox using inverse-mask erosion/dilation, keeping outer silhouette."""
        if amount <= 0:
            return mask

        b, h, w = mask.shape
        y0, y1, x0, x1 = MaskPinchWiden._bbox(mask)
        if y1 <= y0 or x1 <= x0:
            return mask

        sub = mask[:, y0:y1, x0:x1]
        inv = 1.0 - sub

        if axis == "vertical":
            kernel = torch.ones((1, 1, 1, amount * 2 + 1), device=mask.device, dtype=mask.dtype)
            inv = F.pad(inv.unsqueeze(1), (amount, amount, 0, 0), value=1.0)
        else:
            kernel = torch.ones((1, 1, amount * 2 + 1, 1), device=mask.device, dtype=mask.dtype)
            inv = F.pad(inv.unsqueeze(1), (0, 0, amount, amount), value=1.0)

        conv = F.conv2d(inv, kernel, padding=0)
        if mode == "pinch":
            inv_new = (conv >= kernel.numel()).float()
        else:
            inv_new = (conv > 0).float()
        sub_new = (1.0 - inv_new).squeeze(1)

        out = mask.clone()
        out[:, y0:y1, x0:x1] = sub_new
        return out

    def taper(
        self,
        mask: torch.Tensor,
        taper_amount: int,
        vertical_pinch: bool = False,
        horizontal_pinch: bool = False,
        vertical_widen: bool = False,
        horizontal_widen: bool = False,
    ):
        if mask.dim() == 2:
            mask = mask.unsqueeze(0)
        out = mask
        amt = max(0, int(taper_amount))
        if amt > 0:
            if vertical_pinch:
                out = self._pinch_widen(out, amt, axis="vertical", mode="pinch")
            if vertical_widen:
                out = self._pinch_widen(out, amt, axis="vertical", mode="widen")
            if horizontal_pinch:
                out = self._pinch_widen(out, amt, axis="horizontal", mode="pinch")
            if horizontal_widen:
                out = self._pinch_widen(out, amt, axis="horizontal", mode="widen")
        return (out.clamp(0.0, 1.0), )

    @classmethod
    def IS_CHANGED(
        cls,
        taper_amount: int,
        vertical_pinch: bool = False,
        horizontal_pinch: bool = False,
        vertical_widen: bool = False,
        horizontal_widen: bool = False,
        **kwargs,
    ):
        return hash(
            (
                int(taper_amount),
                bool(vertical_pinch),
                bool(horizontal_pinch),
                bool(vertical_widen),
                bool(horizontal_widen),
            )
        )


NODE_CLASS_MAPPINGS = {
    "DirectionalMaskExpand": DirectionalMaskExpand,
    "MaskPinchWiden": MaskPinchWiden,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "DirectionalMaskExpand": "Directional Mask Expand",
    "MaskPinchWiden": "Mask Pinch/Widen",
}