From 9165dc5188858c79ce9e2a1b3dff988237a9823f Mon Sep 17 00:00:00 2001 From: "Dr.Lt.Data" Date: Wed, 11 Dec 2024 21:52:20 +0900 Subject: [PATCH] v1.0 The Impact Subpack is a custom node that exists independently from the Impact Pack. --- LICENSE => LICENSE.txt | 0 README.md | 15 ++- __init__.py | 37 +++++++ install.py | 11 +- {impact => modules}/subcore.py | 74 +++++++++---- {impact => modules}/subpack_nodes.py | 17 ++- modules/utils.py | 151 +++++++++++++++++++++++++++ pyproject.toml | 15 +++ requirements.txt | 2 + 9 files changed, 279 insertions(+), 43 deletions(-) rename LICENSE => LICENSE.txt (100%) create mode 100644 __init__.py rename {impact => modules}/subcore.py (75%) rename {impact => modules}/subpack_nodes.py (57%) create mode 100644 modules/utils.py create mode 100644 pyproject.toml diff --git a/LICENSE b/LICENSE.txt similarity index 100% rename from LICENSE rename to LICENSE.txt diff --git a/README.md b/README.md index 4145c72..a54d9b6 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,18 @@ # ComfyUI-Impact-Subpack -This extension serves as a complement to the Impact Pack, offering features that are not deemed suitable for inclusion by default in the ComfyUI Impact Pack. +This node pack provides nodes that complement the [ComfyUI Impact Pack](https://github.com/ltdrdata/ComfyUI-Impact-Pack), such as the UltralyticsDetectorProvider. -The nodes in this repository cannot be used standalone and depend on [ComfyUI-Impact-Pack](https://github.com/ltdrdata/ComfyUI-Impact-Pack). ## Nodes -* UltralyticsDetectorProvider - This node provides an object detection detector based on Ultralystics. - * By using this Detector Provider, you can replace the existing mmdet-based detector. +* `UltralyticsDetectorProvider` - Loads the Ultralystics model to provide SEGM_DETECTOR, BBOX_DETECTOR. +- Unlike `MMDetDetectorProvider`, for segm models, `BBOX_DETECTOR` is also provided. +- The various models available in UltralyticsDetectorProvider can be downloaded through **ComfyUI-Manager**. + + +## Ultralytics models +* When using ultralytics models, save them separately in `models/ultralytics/bbox` and `models/ultralytics/segm` depending on the type of model. Many models can be downloaded by searching for `ultralytics` in the Model Manager of ComfyUI-Manager. +* huggingface.co/Bingsu/[adetailer](https://huggingface.co/Bingsu/adetailer/tree/main) - You can download face, people detection models, and clothing detection models. +* ultralytics/[assets](https://github.com/ultralytics/assets/releases/) - You can download various types of detection models other than faces or people. +* civitai/[adetailer](https://civitai.com/search/models?sortBy=models_v5&query=adetailer) - You can download various types detection models....Many models are associated with NSFW content. ## Credits diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e048f1d --- /dev/null +++ b/__init__.py @@ -0,0 +1,37 @@ +""" +@author: Dr.Lt.Data +@title: Impact Subpack +@nickname: Impact Subpack +@description: This extension provides UltralyticsDetectorProvider node +""" + +import importlib + +version_code = [1, 0] +version_str = f"V{version_code[0]}.{version_code[1]}" + (f'.{version_code[2]}' if len(version_code) > 2 else '') +print(f"### Loading: ComfyUI-Impact-Subpack ({version_str})") + +node_list = [ + "subpack_nodes", +] + +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +for module_name in node_list: + imported_module = importlib.import_module(".modules.{}".format(module_name), __name__) + + NODE_CLASS_MAPPINGS = {**NODE_CLASS_MAPPINGS, **imported_module.NODE_CLASS_MAPPINGS} + NODE_DISPLAY_NAME_MAPPINGS = {**NODE_DISPLAY_NAME_MAPPINGS, **imported_module.NODE_DISPLAY_NAME_MAPPINGS} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] + +try: + import cm_global + cm_global.register_extension('ComfyUI-Impact-Subpack', + {'version': version_code, + 'name': 'Impact Subpack', + 'nodes': set(NODE_CLASS_MAPPINGS.keys()), + 'description': 'This extension provides UltralyticsDetectorProvider node.', }) +except: + pass diff --git a/install.py b/install.py index 5a55af7..0164347 100644 --- a/install.py +++ b/install.py @@ -6,20 +6,21 @@ subpack_path = os.path.join(os.path.dirname(__file__)) comfy_path = os.environ.get('COMFYUI_PATH') if comfy_path is None: - print(f"\n[bold yellow]WARN: The `COMFYUI_PATH` environment variable is not set. Assuming `{os.path.dirname(__file__)}/../../../` as the ComfyUI path.[/bold yellow]", file=sys.stderr) - comfy_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..')) - + comfy_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')) + print(f"\nWARN: The `COMFYUI_PATH` environment variable is not set. Assuming `{comfy_path}` as the ComfyUI path.", file=sys.stderr) + sys.path.append(comfy_path) model_path = os.environ.get('COMFYUI_MODEL_PATH') if model_path is None: - print(f"\n[bold yellow]WARN: The `COMFYUI_MODEL_PATH` environment variable is not set. Assuming `{model_path}` as the ComfyUI path.[/bold yellow]", file=sys.stderr) model_path = os.path.abspath(os.path.join(comfy_path, 'models')) + print(f"\nWARN: The `COMFYUI_MODEL_PATH` environment variable is not set. Assuming `{model_path}` as the ComfyUI path.", file=sys.stderr) + ultralytics_bbox_path = os.path.join(model_path, "ultralytics", "bbox") ultralytics_segm_path = os.path.join(model_path, "ultralytics", "segm") -if not os.path.exists(os.path.join(subpack_path, '..', '..', 'skip_download_model')): +if not os.path.exists(os.path.join(subpack_path, '..', 'skip_download_model')): if not os.path.exists(ultralytics_bbox_path): os.makedirs(ultralytics_bbox_path) diff --git a/impact/subcore.py b/modules/subcore.py similarity index 75% rename from impact/subcore.py rename to modules/subcore.py index e37609f..125766d 100644 --- a/impact/subcore.py +++ b/modules/subcore.py @@ -1,14 +1,42 @@ from pathlib import Path from PIL import Image -import impact.core as core import cv2 import numpy as np from torchvision.transforms.functional import to_pil_image import torch +from collections import namedtuple +from . import utils + orig_torch_load = torch.load + +SEG = namedtuple("SEG", + ['cropped_image', 'cropped_mask', 'confidence', 'crop_region', 'bbox', 'label', 'control_net_wrapper'], + defaults=[None]) + + +class NO_BBOX_DETECTOR: + pass + + +class NO_SEGM_DETECTOR: + pass + + +def create_segmasks(results): + bboxs = results[1] + segms = results[2] + confidence = results[3] + + results = [] + for i in range(len(segms)): + item = (bboxs[i], segms[i].astype(np.float32), confidence[i]) + results.append(item) + return results + + try: from ultralytics import YOLO except Exception as e: @@ -134,11 +162,11 @@ class UltraBBoxDetector: def detect(self, image, threshold, dilation, crop_factor, drop_size=1, detailer_hook=None): drop_size = max(drop_size, 1) - detected_results = inference_bbox(self.bbox_model, core.tensor2pil(image), threshold) - segmasks = core.create_segmasks(detected_results) + detected_results = inference_bbox(self.bbox_model, utils.tensor2pil(image), threshold) + segmasks = create_segmasks(detected_results) if dilation > 0: - segmasks = core.dilate_masks(segmasks, dilation) + segmasks = utils.dilate_masks(segmasks, dilation) items = [] h = image.shape[1] @@ -151,17 +179,17 @@ class UltraBBoxDetector: y1, x1, y2, x2 = item_bbox if x2 - x1 > drop_size and y2 - y1 > drop_size: # minimum dimension must be (2,2) to avoid squeeze issue - crop_region = core.make_crop_region(w, h, item_bbox, crop_factor) + crop_region = utils.make_crop_region(w, h, item_bbox, crop_factor) if detailer_hook is not None: crop_region = detailer_hook.post_crop_region(w, h, item_bbox, crop_region) - cropped_image = core.crop_image(image, crop_region) - cropped_mask = core.crop_ndarray2(item_mask, crop_region) + cropped_image = utils.crop_image(image, crop_region) + cropped_mask = utils.crop_ndarray2(item_mask, crop_region) confidence = x[2] # bbox_size = (item_bbox[2]-item_bbox[0],item_bbox[3]-item_bbox[1]) # (w,h) - item = core.SEG(cropped_image, cropped_mask, confidence, crop_region, item_bbox, label, None) + item = SEG(cropped_image, cropped_mask, confidence, crop_region, item_bbox, label, None) items.append(item) @@ -174,12 +202,12 @@ class UltraBBoxDetector: return segs def detect_combined(self, image, threshold, dilation): - detected_results = inference_bbox(self.bbox_model, core.tensor2pil(image), threshold) - segmasks = core.create_segmasks(detected_results) + detected_results = inference_bbox(self.bbox_model, utils.tensor2pil(image), threshold) + segmasks = create_segmasks(detected_results) if dilation > 0: - segmasks = core.dilate_masks(segmasks, dilation) + segmasks = utils.dilate_masks(segmasks, dilation) - return core.combine_masks(segmasks) + return utils.combine_masks(segmasks) def setAux(self, x): pass @@ -193,11 +221,11 @@ class UltraSegmDetector: def detect(self, image, threshold, dilation, crop_factor, drop_size=1, detailer_hook=None): drop_size = max(drop_size, 1) - detected_results = inference_segm(self.bbox_model, core.tensor2pil(image), threshold) - segmasks = core.create_segmasks(detected_results) + detected_results = inference_segm(self.bbox_model, utils.tensor2pil(image), threshold) + segmasks = create_segmasks(detected_results) if dilation > 0: - segmasks = core.dilate_masks(segmasks, dilation) + segmasks = utils.dilate_masks(segmasks, dilation) items = [] h = image.shape[1] @@ -210,17 +238,17 @@ class UltraSegmDetector: y1, x1, y2, x2 = item_bbox if x2 - x1 > drop_size and y2 - y1 > drop_size: # minimum dimension must be (2,2) to avoid squeeze issue - crop_region = core.make_crop_region(w, h, item_bbox, crop_factor) + crop_region = utils.make_crop_region(w, h, item_bbox, crop_factor) if detailer_hook is not None: crop_region = detailer_hook.post_crop_region(w, h, item_bbox, crop_region) - cropped_image = core.crop_image(image, crop_region) - cropped_mask = core.crop_ndarray2(item_mask, crop_region) + cropped_image = utils.crop_image(image, crop_region) + cropped_mask = utils.crop_ndarray2(item_mask, crop_region) confidence = x[2] # bbox_size = (item_bbox[2]-item_bbox[0],item_bbox[3]-item_bbox[1]) # (w,h) - item = core.SEG(cropped_image, cropped_mask, confidence, crop_region, item_bbox, label, None) + item = SEG(cropped_image, cropped_mask, confidence, crop_region, item_bbox, label, None) items.append(item) @@ -233,12 +261,12 @@ class UltraSegmDetector: return segs def detect_combined(self, image, threshold, dilation): - detected_results = inference_segm(self.bbox_model, core.tensor2pil(image), threshold) - segmasks = core.create_segmasks(detected_results) + detected_results = inference_segm(self.bbox_model, utils.tensor2pil(image), threshold) + segmasks = create_segmasks(detected_results) if dilation > 0: - segmasks = core.dilate_masks(segmasks, dilation) + segmasks = utils.dilate_masks(segmasks, dilation) - return core.combine_masks(segmasks) + return utils.combine_masks(segmasks) def setAux(self, x): pass \ No newline at end of file diff --git a/impact/subpack_nodes.py b/modules/subpack_nodes.py similarity index 57% rename from impact/subpack_nodes.py rename to modules/subpack_nodes.py index 98902a8..ff80287 100644 --- a/impact/subpack_nodes.py +++ b/modules/subpack_nodes.py @@ -1,17 +1,12 @@ import os import folder_paths -import impact.core as core -import impact.subcore as subcore -from impact.utils import add_folder_path_and_extensions - -version_code = 24 - -print(f"### Loading: ComfyUI-Impact-Pack (Subpack: V0.8)") +from . import subcore +from . import utils model_path = folder_paths.models_dir -add_folder_path_and_extensions("ultralytics_bbox", [os.path.join(model_path, "ultralytics", "bbox")], folder_paths.supported_pt_extensions) -add_folder_path_and_extensions("ultralytics_segm", [os.path.join(model_path, "ultralytics", "segm")], folder_paths.supported_pt_extensions) -add_folder_path_and_extensions("ultralytics", [os.path.join(model_path, "ultralytics")], folder_paths.supported_pt_extensions) +utils.add_folder_path_and_extensions("ultralytics_bbox", [os.path.join(model_path, "ultralytics", "bbox")], folder_paths.supported_pt_extensions) +utils.add_folder_path_and_extensions("ultralytics_segm", [os.path.join(model_path, "ultralytics", "segm")], folder_paths.supported_pt_extensions) +utils.add_folder_path_and_extensions("ultralytics", [os.path.join(model_path, "ultralytics")], folder_paths.supported_pt_extensions) class UltralyticsDetectorProvider: @@ -30,7 +25,7 @@ class UltralyticsDetectorProvider: model = subcore.load_yolo(model_path) if model_name.startswith("bbox"): - return subcore.UltraBBoxDetector(model), core.NO_SEGM_DETECTOR() + return subcore.UltraBBoxDetector(model), subcore.NO_SEGM_DETECTOR() else: return subcore.UltraBBoxDetector(model), subcore.UltraSegmDetector(model) diff --git a/modules/utils.py b/modules/utils.py new file mode 100644 index 0000000..f6cca0a --- /dev/null +++ b/modules/utils.py @@ -0,0 +1,151 @@ +import folder_paths +from PIL import Image +import numpy as np +import cv2 + + +def add_folder_path_and_extensions(folder_name, full_folder_paths, extensions): + # Iterate over the list of full folder paths + for full_folder_path in full_folder_paths: + # Use the provided function to add each model folder path + folder_paths.add_model_folder_path(folder_name, full_folder_path) + + # Now handle the extensions. If the folder name already exists, update the extensions + if folder_name in folder_paths.folder_names_and_paths: + # Unpack the current paths and extensions + current_paths, current_extensions = folder_paths.folder_names_and_paths[folder_name] + # Update the extensions set with the new extensions + updated_extensions = current_extensions | extensions + # Reassign the updated tuple back to the dictionary + folder_paths.folder_names_and_paths[folder_name] = (current_paths, updated_extensions) + else: + # If the folder name was not present, add_model_folder_path would have added it with the last path + # Now we just need to update the set of extensions as it would be an empty set + # Also ensure that all paths are included (since add_model_folder_path adds only one path at a time) + folder_paths.folder_names_and_paths[folder_name] = (full_folder_paths, extensions) + + +def normalize_region(limit, startp, size): + if startp < 0: + new_endp = min(limit, size) + new_startp = 0 + elif startp + size > limit: + new_startp = max(0, limit - size) + new_endp = limit + else: + new_startp = startp + new_endp = min(limit, startp+size) + + return int(new_startp), int(new_endp) + + +def _tensor_check_image(image): + if image.ndim != 4: + raise ValueError(f"Expected NHWC tensor, but found {image.ndim} dimensions") + if image.shape[-1] not in (1, 3, 4): + raise ValueError(f"Expected 1, 3 or 4 channels for image, but found {image.shape[-1]} channels") + return + + +def tensor2pil(image): + _tensor_check_image(image) + return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(0), 0, 255).astype(np.uint8)) + + +def dilate_masks(segmasks, dilation_factor, iter=1): + if dilation_factor == 0: + return segmasks + + dilated_masks = [] + kernel = np.ones((abs(dilation_factor), abs(dilation_factor)), np.uint8) + + for i in range(len(segmasks)): + cv2_mask = segmasks[i][1] + + if dilation_factor > 0: + dilated_mask = cv2.dilate(cv2_mask, kernel, iter) + else: + dilated_mask = cv2.erode(cv2_mask, kernel, iter) + + item = (segmasks[i][0], dilated_mask, segmasks[i][2]) + dilated_masks.append(item) + + return dilated_masks + + +def combine_masks(masks): + if len(masks) == 0: + return None + else: + initial_cv2_mask = np.array(masks[0][1]) + combined_cv2_mask = initial_cv2_mask + + for i in range(1, len(masks)): + cv2_mask = np.array(masks[i][1]) + + if combined_cv2_mask.shape == cv2_mask.shape: + combined_cv2_mask = cv2.bitwise_or(combined_cv2_mask, cv2_mask) + else: + # do nothing - incompatible mask + pass + + mask = torch.from_numpy(combined_cv2_mask) + return mask + + +def make_crop_region(w, h, bbox, crop_factor, crop_min_size=None): + x1 = bbox[0] + y1 = bbox[1] + x2 = bbox[2] + y2 = bbox[3] + + bbox_w = x2 - x1 + bbox_h = y2 - y1 + + crop_w = bbox_w * crop_factor + crop_h = bbox_h * crop_factor + + if crop_min_size is not None: + crop_w = max(crop_min_size, crop_w) + crop_h = max(crop_min_size, crop_h) + + kernel_x = x1 + bbox_w / 2 + kernel_y = y1 + bbox_h / 2 + + new_x1 = int(kernel_x - crop_w / 2) + new_y1 = int(kernel_y - crop_h / 2) + + # make sure position in (w,h) + new_x1, new_x2 = normalize_region(w, new_x1, crop_w) + new_y1, new_y2 = normalize_region(h, new_y1, crop_h) + + return [new_x1, new_y1, new_x2, new_y2] + + +def crop_ndarray2(npimg, crop_region): + x1 = crop_region[0] + y1 = crop_region[1] + x2 = crop_region[2] + y2 = crop_region[3] + + cropped = npimg[y1:y2, x1:x2] + + return cropped + + +def crop_ndarray4(npimg, crop_region): + x1 = crop_region[0] + y1 = crop_region[1] + x2 = crop_region[2] + y2 = crop_region[3] + + cropped = npimg[:, y1:y2, x1:x2, :] + + return cropped + + +crop_tensor4 = crop_ndarray4 + + +def crop_image(image, crop_region): + return crop_tensor4(image, crop_region) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..42f1c83 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,15 @@ +[project] +name = "comfyui-impact-subpack" +description = "This node pack provides nodes that complement the Impact Pack, such as the UltralyticsDetectorProvider." +version = "1.0" +license = { file = "LICENSE.txt" } +dependencies = ["matplotlib", "ultralytics==8.3.40"] + +[project.urls] +Repository = "https://github.com/ltdrdata/ComfyUI-Impact-Subpack" +# Used by Comfy Registry https://comfyregistry.org + +[tool.comfy] +PublisherId = "drltdata" +DisplayName = "ComfyUI Impact Subpack" +Icon = "" diff --git a/requirements.txt b/requirements.txt index 3d19f3a..5f7489e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,4 @@ matplotlib ultralytics==8.3.40 +numpy<2 +opencv-python-headless \ No newline at end of file