first commit

.gitignore

@@ -0,0 +1,148 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the executable, but when you build from the source
+#  directory, the .spec file may not have been created.
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+doc/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration over git, it will be ignored.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type inference
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# ComfyUI specific
+comfyui/
+comfyui.egg-info/
+*.ckpt
+*.safetensors
+*.pt
+*.bin
+*.json
+*.yaml
+*.yml
+*.ckpt.meta
+*.safetensors.meta
+*.pt.meta
+*.bin.meta
+*.json.meta
+*.yaml.meta
+*.yml.meta

README.md

@@ -0,0 +1,13 @@
+## Description
+
+Some nodes for processing masks, currently including nodes that fill in the concave parts of existing masks with convex hulls.
+
+## Mask Nodes examples
+
+MaskToConvexMask is responsible for filling in all concave areas of an existing mask with a convex hull.
+
+MaskToBottomHalfConvexMask is responsible for filling in the concave areas of the lower half of an existing mask with a convex hull.
+
+![alt text](example.png)
+
+

__init__.py

@@ -0,0 +1,12 @@
+from . import nodes
+
+NODE_CLASS_MAPPINGS = {
+    "MaskToBottonHalfConvexMask": nodes.MaskToBottonHalfConvexMask,
+    "MaskToConvexMask": nodes.MaskToConvexMask,
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "MaskToBottonHalfConvexMask": "Mask To Botton Half Convex Mask",
+    "MaskToConvexMask": "Mask To Convex Mask",
+}
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

nodes.py

@@ -0,0 +1,98 @@
+import torch
+import numpy as np
+import torch
+import cv2
+
+class MaskToBottonHalfConvexMask:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "mask": ("MASK",),
+            },
+        }
+
+    RETURN_TYPES = ("MASK",)
+    CATEGORY = "MaskToBottonHalfConvexMask"
+    FUNCTION = "generate_convex_mask"
+
+    def generate_convex_mask(self, mask):
+        """
+        生成一个凸形遮罩，填充输入遮罩下半部分的凹区域。
+        
+        参数：
+            mask (torch.Tensor): 输入遮罩，尺寸为 (batch_size, height, width)。
+        
+        返回：
+            torch.Tensor: 生成的凸形遮罩，尺寸与输入遮罩相同。
+        """
+        # 将 PyTorch 张量转换为 NumPy 数组
+        mask_np = mask.squeeze(0).numpy()  # 去掉 batch_size 维度
+
+        height, width = mask_np.shape
+        bottom_half = mask_np[height // 2:, :]
+        contours, _ = cv2.findContours(
+            bottom_half.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+        )
+
+        if len(contours) == 0:
+            return (mask,)
+
+        all_points = np.vstack(contours)
+        # 计算凸包
+        hull = cv2.convexHull(all_points)
+        # 创建一个空白图像用于绘制凸包
+        convex_mask = np.zeros_like(bottom_half, dtype=np.float32)
+        # 填充凸包区域
+        cv2.fillPoly(convex_mask, [hull], 1.0)
+
+        new_mask = mask_np.copy()
+        new_mask[height // 2:, :] = convex_mask
+
+        new_mask_tensor = torch.from_numpy(new_mask).unsqueeze(0)  # 添加 batch_size 维度
+
+        return (new_mask_tensor,)
+
+
+class MaskToConvexMask:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "mask": ("MASK",),  # 输入的 MASK
+            },
+        }
+    
+    RETURN_TYPES = ("MASK",)
+    CATEGORY = "MaskToConvexMask"
+    FUNCTION = "generate_convex_mask"
+
+    def generate_convex_mask(self, mask):
+        """
+        将输入的 MASK 的凹区域填充为凸区域。
+        
+        参数：
+            mask (torch.Tensor): 输入的 MASK，形状为 (batch_size, height, width)。
+        
+        返回：
+            torch.Tensor: 新的 MASK，形状为 (batch_size, height, width)。
+        """
+        mask_np = mask.squeeze(0).numpy()  # 去掉 batch_size 维度
+
+        contours, _ = cv2.findContours(
+            mask_np.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+        )
+
+        if len(contours) == 0:
+            return (mask,)
+
+        all_points = np.vstack(contours)
+        # 计算凸包
+        hull = cv2.convexHull(all_points)
+        # 创建一个空白图像用于绘制凸包
+        convex_mask = np.zeros_like(mask_np, dtype=np.float32)
+        # 填充凸包区域
+        cv2.fillPoly(convex_mask, [hull], 1.0)
+        convex_mask_tensor = torch.from_numpy(convex_mask).unsqueeze(0)  # 添加 batch_size 维度
+
+        return (convex_mask_tensor,)

requirement.txt

@@ -0,0 +1,3 @@
+numpy>=1.21.0
+opencv-python>=4.5.0
+torch>=2.0.0