commit YoloV8Detect node

2024-05-01 22:19:50 +08:00
parent 08018e2f6e
commit f493767c0e
6 changed files with 143 additions and 4 deletions
@@ -70,8 +70,9 @@ When this error has occurred, please check the network environment.
 ## Update
 <font size="4">**If the dependency package error after updating, please reinstall the relevant dependency packages. </font><br />    

+* 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.

+### <a id="table1">YoloV8Detect</a>
+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.
+

 ### <a id="table1">Shadow</a> & Highlight Mask
 Generate masks for the dark and bright parts of the image.
@@ -70,8 +70,9 @@ git clone https://github.com/chflame163/ComfyUI_LayerStyle.git
 ## 更新说明
 <font size="4">**如果本插件更新后出现依赖包错误，请重新安装相关依赖包。

-* 添加 [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: 遮罩边缘向外扩张范围。数值越大，向外修复的范围越大。


+### <a id="table1">YoloV8Detect</a>
+使用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。
+
+
 ### <a id="table1">Shadow</a> & Highlight Mask
 生成图像暗部和亮部的遮罩。
 ![image](image/shadow_and_highlight_mask_example.png)
@@ -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"
+}
@@ -23,4 +23,5 @@ omegaconf
 tqdm
 transformers
 kornia
-image-reward
+image-reward
+ultralytics