diff --git a/README.MD b/README.MD
index b50219b..a82c1f0 100644
--- a/README.MD
+++ b/README.MD
@@ -70,8 +70,9 @@ When this error has occurred, please check the network environment.
## Update
**If the dependency package error after updating, please reinstall the relevant dependency packages.
+* Commit [YoloV8Detect](#YoloV8Detect) node.
* Commit [QWenImage2Prompt](#QWenImage2Prompt) node, this node is repackage of the [ComfyUI_VLM_nodes](https://github.com/gokayfem/ComfyUI_VLM_nodes)'s ```UForm-Gen2 Qwen Node```, thanks to the original author.
-* Commit [BooleanOperator](#BooleanOperator), [NumberCalculator](#NumberCalculator), [TextBox](#TextBox), [Integer](#Integer), [Float](#Float), [Boolean](#Boolean) nodes. These nodes can perform mathematical and logical operations.
+* Commit [BooleanOperator](#BooleanOperator), [NumberCalculator](#NumberCalculator), [TextBox](#TextBox), [Integer](#Integer), [Float](#Float), [Boolean](#Boolean)nodes. These nodes can perform mathematical and logical operations.
* Commit [ExtendCanvasV2](#ExtendCanvasV2) node,support color value input.
* Commit [AutoBrightness](#AutoBrightness) node,it can automatically adjust the brightness of image.
* [CreateGradientMask](#CreateGradientMask) node add ```center``` option.
@@ -1141,6 +1142,22 @@ On the basis of PersonMaskUltra, the following changes have been made:
* detail_erode: Mask the erosion range inward from the edge. the larger the value, the larger the range of inward repair.
* detail_dilate: The edge of the mask expands outward. the larger the value, the wider the range of outward repair.
+### YoloV8Detect
+Use the YoloV8 model to detect faces, hand box areas, or character segmentation. Supports the output of the selected number of channels.
+Download the model files from [GoogleDrive](https://drive.google.com/drive/folders/1I5TISO2G1ArSkKJu1O9b4Uvj3DVgn5d2) or [BaiduNetdisk](https://pan.baidu.com/s/1ImoJrzL1zDgaCqaSzrNEtw?pwd=5xgk) to ```ComfyUI/models/yolo``` folder.
+
+
+
+Node Options:
+
+* yolo_model: Yolo model selection. the model with ```seg``` name can output segmented masks, otherwise they can only output box masks.
+* mask_merge: Select the merged mask. ```all``` is to merge all mask outputs. The selected number is how many masks to output, sorted by recognition confidence to merge the output.
+
+Outputs:
+* mask: The output mask.
+* yolo_plot_image: Preview of yolo recognition results.
+* yolo_masks: For all masks identified by yolo, each individual mask is output as a mask.
+
### Shadow & Highlight Mask
Generate masks for the dark and bright parts of the image.
diff --git a/README_CN.MD b/README_CN.MD
index a0cb55c..3c8a402 100644
--- a/README_CN.MD
+++ b/README_CN.MD
@@ -70,8 +70,9 @@ git clone https://github.com/chflame163/ComfyUI_LayerStyle.git
## 更新说明
**如果本插件更新后出现依赖包错误,请重新安装相关依赖包。
-* 添加 [QWenImage2Prompt](#QWenImage2Prompt)节点, 用本地模型反推提示词。(需要下载模型到models文件夹)。这个节点是[ComfyUI_VLM_nodes](https://github.com/gokayfem/ComfyUI_VLM_nodes)中的```UForm-Gen2 Qwen Node```节点的重新封装,感谢原作者。
-* 添加 [BooleanOperator](#BooleanOperator), [NumberCalculator](#NumberCalculator), [TextBox](#TextBox), [Integer](#Integer), [Float](#Float), [Boolean](#Boolean) 节点。这些节点可进行数学和逻辑运算。
+* 添加 [YoloV8Detect](#YoloV8Detect) 节点。
+* 添加 [QWenImage2Prompt](#QWenImage2Prompt)节点, 用本地模型反推提示词。(需要下载模型到models文件夹)
+* 添加 [BooleanOperator](#BooleanOperator), [NumberCalculator](#NumberCalculator), [TextBox](#TextBox), [Integer](#Integer), [Float](#Float), [Boolean](#Boolean)节点。这些节点可进行数学和逻辑运算。
* 添加 [ExtendCanvasV2](#ExtendCanvasV2) 节点,支持color值输入。
* 添加 [AutoBrightness](#AutoBrightness) 节点,可自动调整图片亮度。
* [CreateGradientMask](#CreateGradientMask) 节点增加 ```center``` 选项。
@@ -1135,6 +1136,23 @@ PersonMaskUltra的V2升级版,增加了VITMatte边缘处理方法。(注意:
* detail_dilate: 遮罩边缘向外扩张范围。数值越大,向外修复的范围越大。
+### YoloV8Detect
+使用YoloV8模型检测人脸、手部box区域,或者人物分割。支持输出所选择数量的通道。
+请在 [GoogleDrive](https://drive.google.com/drive/folders/1I5TISO2G1ArSkKJu1O9b4Uvj3DVgn5d2) 或者 [百度网盘](https://pan.baidu.com/s/1ImoJrzL1zDgaCqaSzrNEtw?pwd=5xgk) 下载模型文件并放到 ```ComfyUI/models/yolo``` 文件夹。
+
+
+
+节点选项说明:
+
+* yolo_model: yolo模型选择。带有```seg```名字的模型可以输出分割的mask, 否则只能输出box区域的遮罩。
+* mask_merge: 选择合并的遮罩。```all```是合并全部遮罩输出。选数值是输出多少个遮罩,按识别置信度排序合并输出。
+
+输出:
+* mask: 输出的遮罩。
+* yolo_plot_image: yolo识别结果预览图。
+* yolo_masks: yolo识别出来的所有遮罩,每个单独的遮罩输出为一个mask。
+
+
### Shadow & Highlight Mask
生成图像暗部和亮部的遮罩。

diff --git a/image/yolov8_detect_example.png b/image/yolov8_detect_example.png
new file mode 100644
index 0000000..a1f3050
Binary files /dev/null and b/image/yolov8_detect_example.png differ
diff --git a/image/yolov8_detect_node.png b/image/yolov8_detect_node.png
new file mode 100644
index 0000000..ef3aaaa
Binary files /dev/null and b/image/yolov8_detect_node.png differ
diff --git a/py/yolov8_detect.py b/py/yolov8_detect.py
new file mode 100644
index 0000000..aad45e4
--- /dev/null
+++ b/py/yolov8_detect.py
@@ -0,0 +1,103 @@
+import copy
+import os.path
+
+from .imagefunc import *
+
+NODE_NAME = 'YoloV8Detect'
+
+model_path = os.path.join(folder_paths.models_dir, 'yolo')
+
+class YoloV8Detect:
+
+ def __init__(self):
+ pass
+
+ @classmethod
+ def INPUT_TYPES(self):
+ __file_list = glob.glob(model_path + '/*.pt')
+ # __file_list.extend(glob.glob(model_path + '/*.safetensors'))
+ FILES_DICT = {}
+ for i in range(len(__file_list)):
+ _, __filename = os.path.split(__file_list[i])
+ FILES_DICT[__filename] = __file_list[i]
+ FILE_LIST = list(FILES_DICT.keys())
+
+ mask_merge = ["all", "1", "2", "3", "4", "5", "6", "7", "8", "9"]
+ return {
+ "required": {
+ "image": ("IMAGE", ),
+ "yolo_model": (FILE_LIST,),
+ "mask_merge": (mask_merge,),
+ },
+ "optional": {
+ }
+ }
+
+ RETURN_TYPES = ("MASK", "IMAGE", "MASK" )
+ RETURN_NAMES = ("mask", "yolo_plot_image", "yolo_masks")
+ FUNCTION = 'yolo_detect'
+ CATEGORY = '😺dzNodes/LayerMask'
+
+ def yolo_detect(self, image,
+ yolo_model, mask_merge
+ ):
+
+ ret_masks = []
+ ret_yolo_plot_images = []
+ ret_yolo_masks = []
+
+ from ultralytics import YOLO
+ yolo_model = YOLO(os.path.join(model_path, yolo_model))
+
+ for i in image:
+ i = torch.unsqueeze(i, 0)
+ _image = tensor2pil(i)
+ results = yolo_model(_image, retina_masks=True)
+ for result in results:
+ yolo_plot_image = cv2.cvtColor(result.plot(), cv2.COLOR_BGR2RGB)
+ ret_yolo_plot_images.append(pil2tensor(Image.fromarray(yolo_plot_image)))
+ # have mask
+ if result.masks is not None and len(result.masks) > 0:
+ masks = []
+ masks_data = result.masks.data
+ for index, mask in enumerate(masks_data):
+ _mask = mask.cpu().numpy() * 255
+ _mask = np2pil(_mask).convert("L")
+ ret_yolo_masks.append(image2mask(_mask))
+ # no mask, if have box, draw box
+ elif result.boxes is not None and len(result.boxes.xyxy) > 0:
+ white_image = Image.new('L', _image.size, "white")
+ for box in result.boxes:
+ x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
+ x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+ _mask = Image.new('L', _image.size, "black")
+ _mask.paste(white_image.crop((x1, y1, x2, y2)), (x1, y1))
+ ret_yolo_masks.append(image2mask(_mask))
+ # no mask and box, add a black mask
+ else:
+ ret_yolo_masks.append(torch.zeros((1, _image.size[1], _image.size[0]), dtype=torch.float32))
+ # ret_yolo_masks.append(image2mask(Image.new('L', _image.size, "black")))
+ log(f"{NODE_NAME} mask or box not detected.")
+
+ # merge mask
+ _mask = ret_yolo_masks[0]
+ if mask_merge == "all":
+ for i in range(len(ret_yolo_masks) - 1):
+ _mask = add_mask(_mask, ret_yolo_masks[i + 1])
+ else:
+ for i in range(min(len(ret_yolo_masks), int(mask_merge)) - 1):
+ _mask = add_mask(_mask, ret_yolo_masks[i + 1])
+ ret_masks.append(_mask)
+
+ log(f"{NODE_NAME} Processed {len(ret_masks)} image(s).", message_type='finish')
+ return (torch.cat(ret_masks, dim=0),
+ torch.cat(ret_yolo_plot_images, dim=0),
+ torch.cat(ret_yolo_masks, dim=0),)
+
+NODE_CLASS_MAPPINGS = {
+ "LayerMask: YoloV8Detect": YoloV8Detect
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+ "LayerMask: YoloV8Detect": "LayerMask: YoloV8 Detect"
+}
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index fc294c5..ba084c1 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -23,4 +23,5 @@ omegaconf
tqdm
transformers
kornia
-image-reward
\ No newline at end of file
+image-reward
+ultralytics
\ No newline at end of file