commit 1fdd5c50bfe2e00164f2026871715e201a2eb684 Author: jerrylongyan Date: Wed Jan 8 16:38:41 2025 +0800 first commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..edee190 --- /dev/null +++ b/.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 \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..5ee0cfa --- /dev/null +++ b/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) + + diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..dbc2e24 --- /dev/null +++ b/__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'] \ No newline at end of file diff --git a/example.png b/example.png new file mode 100644 index 0000000..3a9c99e Binary files /dev/null and b/example.png differ diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..d938241 --- /dev/null +++ b/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,) \ No newline at end of file diff --git a/requirement.txt b/requirement.txt new file mode 100644 index 0000000..0bb4ad9 --- /dev/null +++ b/requirement.txt @@ -0,0 +1,3 @@ +numpy>=1.21.0 +opencv-python>=4.5.0 +torch>=2.0.0 \ No newline at end of file