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. + +![image](image/yolov8_detect_example.png) + +Node Options: +![image](image/yolov8_detect_node.png) +* 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``` 文件夹。 + +![image](image/yolov8_detect_example.png) + +节点选项说明: +![image](image/yolov8_detect_node.png) +* yolo_model: yolo模型选择。带有```seg```名字的模型可以输出分割的mask, 否则只能输出box区域的遮罩。 +* mask_merge: 选择合并的遮罩。```all```是合并全部遮罩输出。选数值是输出多少个遮罩,按识别置信度排序合并输出。 + +输出: +* mask: 输出的遮罩。 +* yolo_plot_image: yolo识别结果预览图。 +* yolo_masks: yolo识别出来的所有遮罩,每个单独的遮罩输出为一个mask。 + + ### Shadow & Highlight Mask 生成图像暗部和亮部的遮罩。 ![image](image/shadow_and_highlight_mask_example.png) 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