diff --git a/README.md b/README.md index abf2a7f..13006ef 100644 --- a/README.md +++ b/README.md @@ -7,8 +7,8 @@ * SegmDetector - Detect segmentation from input image. * BboxDetector - Detect bbox from input image. -### ![example](https://user-images.githubusercontent.com/128333288/228896955-d84e164e-af87-4c4d-a377-662a6363381f.png) -Detected mask example. +### ![example](https://user-images.githubusercontent.com/128333288/229093882-239c4a2a-0c94-4b77-b0fa-8ac07eda17b9.png) +Head detection example. # Installation diff --git a/comfyui-impact-pack.py b/comfyui-impact-pack.py index 7089d4f..b0c9e32 100644 --- a/comfyui-impact-pack.py +++ b/comfyui-impact-pack.py @@ -14,18 +14,23 @@ def packages_mim(): return [r.decode().split('==')[0] for r in subprocess.check_output([sys.executable, '-m', 'mim', 'list']).split()] # INSTALL -print("### Check dependencies") -if "openmim" not in packages_pip(): - subprocess.check_call([sys.executable, '-m', 'pip', '-U', 'install', 'openmim']) +print("Loading: ComfyUI-Impact-Pack") +print("### ComfyUI-Impact-Pack: Check dependencies") +installed_pip = packages_pip() -if "mmcv-full" not in packages_mim(): +if "openmim" not in installed_pip: + subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-U', 'openmim']) + +installed_mim = packages_mim() + +if "mmcv-full" not in installed_mim: subprocess.check_call([sys.executable, '-m', 'mim', 'install', 'mmcv-full==1.7.0']) -if "mmdet" not in packages_mim(): +if "mmdet" not in installed_mim: subprocess.check_call([sys.executable, '-m', 'mim', 'install', 'mmdet==2.28.2']) # Download model -print("### Check basic models") +print("### ComfyUI-Impact-Pack: Check basic models") if os.path.realpath("..").endswith("custom_nodes"): # For user @@ -66,7 +71,6 @@ from PIL import Image import model_management def load_mmdet(model_path): - print(model_management.vram_state) model_config = os.path.splitext(model_path)[0] + ".py" model = init_detector(model_config, model_path, device="cpu") return model @@ -78,9 +82,7 @@ def create_segmasks(results): segms = results[2] segmasks = [] for i in range(len(segms)): - cv2_mask = segms[i].astype(np.uint8) * 255 - mask = Image.fromarray(cv2_mask) - segmasks.append(mask) + segmasks.append(segms[i].astype(np.float32)) return segmasks def combine_masks(masks): @@ -90,8 +92,35 @@ def combine_masks(masks): cv2_mask = np.array(masks[i]) combined_cv2_mask = cv2.bitwise_or(combined_cv2_mask, cv2_mask) - combined_mask = Image.fromarray(combined_cv2_mask) - return combined_mask + # combined_mask = Image.fromarray(combined_cv2_mask) + # return combined_mask + mask = torch.from_numpy(combined_cv2_mask) + return mask + +def bitwise_and_masks(mask1, mask2): + cv2_mask1 = np.array(mask1) + cv2_mask2 = np.array(mask2) + cv2_mask = cv2.bitwise_and(cv2_mask1, cv2_mask2) + mask = torch.from_numpy(cv2_mask) + return mask + +def dilate_masks(masks, dilation_factor, iter=1): + if dilation_factor == 0: + return masks + dilated_masks = [] + kernel = np.ones((dilation_factor,dilation_factor), np.uint8) + for i in range(len(masks)): + cv2_mask = masks[i] + dilated_mask = cv2.dilate(cv2_mask, kernel, iter) + dilated_masks.append(dilated_mask) + return dilated_masks + +def subtract_masks(mask1, mask2): + cv2_mask1 = np.array(mask1) * 255 + cv2_mask2 = np.array(mask2) * 255 + cv2_mask = cv2.subtract(cv2_mask1, cv2_mask2) + mask = torch.from_numpy(cv2_mask) / 255.0 + return mask def inference_segm(model, image, conf_threshold): image = image.numpy()[0] * 255 @@ -192,21 +221,22 @@ class SegmDetector: "segm_model": ("SEGM_MODEL", ), "image": ("IMAGE", ), "threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "dilation": ("INT", {"default": 0, "min": 0, "max": 255, "step": 1}), } } - RETURN_TYPES = ("IMAGE",) + RETURN_TYPES = ("MASK",) FUNCTION = "doit" CATEGORY = "ImpactPack" - def doit(self, segm_model, image, threshold): + def doit(self, segm_model, image, threshold, dilation): mmdet_results = inference_segm(segm_model, image, threshold) segmasks = create_segmasks(mmdet_results) + if dilation > 0: + segmasks = dilate_masks(segmasks, dilation) mask = combine_masks(segmasks) - - image = pil2tensor(mask) - return (image,) + return (mask,) class BboxDetector(SegmDetector): @classmethod @@ -216,21 +246,62 @@ class BboxDetector(SegmDetector): "bbox_model": ("BBOX_MODEL", ), "image": ("IMAGE", ), "threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "dilation": ("INT", {"default": 4, "min": 0, "max": 255, "step": 1}), } } - def doit(self, bbox_model, image, threshold): + def doit(self, bbox_model, image, threshold, dilation): mmdet_results = inference_bbox(bbox_model, image, threshold) segmasks = create_segmasks(mmdet_results) + if dilation > 0: + segmasks = dilate_masks(segmasks, dilation) mask = combine_masks(segmasks) + return (mask,) - image = pil2tensor(mask) - return (image,) +class BitwiseAndMask: + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "mask1": ("MASK", ), + "mask2": ("MASK", ), + } + } + + RETURN_TYPES = ("MASK",) + FUNCTION = "doit" + + CATEGORY = "ImpactPack" + + def doit(self, mask1, mask2): + mask = bitwise_and_masks(mask1, mask2) + return (mask,) + +class SubtractMask: + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "mask1": ("MASK", ), + "mask2": ("MASK", ), + } + } + + RETURN_TYPES = ("MASK",) + FUNCTION = "doit" + + CATEGORY = "ImpactPack" + + def doit(self, mask1, mask2): + mask = subtract_masks(mask1, mask2) + return (mask,) NODE_CLASS_MAPPINGS = { "MMDetLoader": MMDetLoader, "BboxDetector": BboxDetector, "SegmDetector": SegmDetector, -} \ No newline at end of file + "BitwiseAndMask": BitwiseAndMask, + "SubtractMask": SubtractMask, +} diff --git a/misc/example.json b/misc/example.json index 7d5c745..d56027f 100644 --- a/misc/example.json +++ b/misc/example.json @@ -1,26 +1,26 @@ { - "last_node_id": 21, - "last_link_id": 23, + "last_node_id": 34, + "last_link_id": 47, "nodes": [ { "id": 4, "type": "PreviewImage", "pos": [ - 1008, - 319 + 495, + 145 + ], + "size": [ + 440, + 206 ], - "size": { - "0": 339, - "1": 380 - }, "flags": {}, - "order": 5, + "order": 12, "mode": 0, "inputs": [ { "name": "images", "type": "IMAGE", - "link": 20 + "link": 31 } ], "title": "PreviewImage - Faces", @@ -28,176 +28,26 @@ "Node name for S&R": "PreviewImage" } }, - { - "id": 9, - "type": "PreviewImage", - "pos": [ - 97, - 321 - ], - "size": { - "0": 366, - "1": 380 - }, - "flags": {}, - "order": 6, - "mode": 0, - "inputs": [ - { - "name": "images", - "type": "IMAGE", - "link": 23 - } - ], - "title": "PreviewImage - Persons", - "properties": { - "Node name for S&R": "PreviewImage" - } - }, - { - "id": 17, - "type": "BboxDetector", - "pos": [ - 1017, - 136 - ], - "size": { - "0": 315, - "1": 78 - }, - "flags": {}, - "order": 3, - "mode": 0, - "inputs": [ - { - "name": "bbox_model", - "type": "BBOX_MODEL", - "link": 18 - }, - { - "name": "image", - "type": "IMAGE", - "link": 21 - } - ], - "outputs": [ - { - "name": "IMAGE", - "type": "IMAGE", - "links": [ - 20 - ], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "BboxDetector" - }, - "widgets_values": [ - 0.5 - ] - }, - { - "id": 1, - "type": "LoadImage", - "pos": [ - 510, - 314 - ], - "size": { - "0": 448, - "1": 387 - }, - "flags": {}, - "order": 0, - "mode": 0, - "outputs": [ - { - "name": "IMAGE", - "type": "IMAGE", - "links": [ - 21, - 22 - ], - "slot_index": 0 - }, - { - "name": "MASK", - "type": "MASK", - "links": null - } - ], - "properties": { - "Node name for S&R": "LoadImage" - }, - "widgets_values": [ - "snh.jpg", - "image" - ] - }, - { - "id": 21, - "type": "SegmDetector", - "pos": [ - 122, - 810 - ], - "size": { - "0": 315, - "1": 78 - }, - "flags": {}, - "order": 4, - "mode": 0, - "inputs": [ - { - "name": "segm_model", - "type": "SEGM_MODEL", - "link": 19 - }, - { - "name": "image", - "type": "IMAGE", - "link": 22 - } - ], - "outputs": [ - { - "name": "IMAGE", - "type": "IMAGE", - "links": [ - 23 - ], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "SegmDetector" - }, - "widgets_values": [ - 0.5 - ] - }, { "id": 18, "type": "MMDetLoader", "pos": [ - 511, - 135 + 961, + 149 + ], + "size": [ + 317, + 87 ], - "size": { - "0": 437, - "1": 87 - }, "flags": {}, - "order": 1, + "order": 0, "mode": 0, "outputs": [ { "name": "BBOX_MODEL", "type": "BBOX_MODEL", "links": [ - 18 + 24 ], "slot_index": 0 }, @@ -215,31 +65,227 @@ ] }, { - "id": 20, - "type": "MMDetLoader", + "id": 25, + "type": "VAEDecode", "pos": [ - 523, - 801 + 294, + 299 + ], + "size": [ + 177, + 46 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 29 + }, + { + "name": "vae", + "type": "VAE", + "link": 30 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 31 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEDecode" + } + }, + { + "id": 23, + "type": "VAEEncodeForInpaint", + "pos": [ + 295, + 147 + ], + "size": [ + 177, + 99 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "pixels", + "type": "IMAGE", + "link": 27 + }, + { + "name": "vae", + "type": "VAE", + "link": 28 + }, + { + "name": "mask", + "type": "MASK", + "link": 26 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 29 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEEncodeForInpaint" + } + }, + { + "id": 30, + "type": "VAEEncodeForInpaint", + "pos": [ + 265, + 985 ], "size": { - "0": 437, - "1": 87 + "0": 210, + "1": 66 }, "flags": {}, - "order": 2, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "pixels", + "type": "IMAGE", + "link": 36 + }, + { + "name": "vae", + "type": "VAE", + "link": 35 + }, + { + "name": "mask", + "type": "MASK", + "link": 34 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 37 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEEncodeForInpaint" + } + }, + { + "id": 28, + "type": "VAEDecode", + "pos": [ + 265, + 1104 + ], + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 37 + }, + { + "name": "vae", + "type": "VAE", + "link": 38 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 39 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEDecode" + } + }, + { + "id": 9, + "type": "PreviewImage", + "pos": [ + 495, + 967 + ], + "size": [ + 439, + 193 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 39 + } + ], + "title": "PreviewImage - Persons", + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 26, + "type": "MMDetLoader", + "pos": [ + 961, + 1079 + ], + "size": [ + 387, + 78 + ], + "flags": {}, + "order": 1, "mode": 0, "outputs": [ { "name": "BBOX_MODEL", "type": "BBOX_MODEL", - "links": [], + "links": null, "slot_index": 0 }, { "name": "SEGM_MODEL", "type": "SEGM_MODEL", "links": [ - 19 + 32 ], "slot_index": 1 } @@ -250,59 +296,552 @@ "widgets_values": [ "segm/mmdet_dd-person_mask2former.pth" ] + }, + { + "id": 27, + "type": "SegmDetector", + "pos": [ + 963, + 922 + ], + "size": [ + 386, + 102 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "segm_model", + "type": "SEGM_MODEL", + "link": 32 + }, + { + "name": "image", + "type": "IMAGE", + "link": 33 + } + ], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [ + 34, + 40 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "SegmDetector" + }, + "widgets_values": [ + 0.5, + 0 + ] + }, + { + "id": 32, + "type": "VAEEncodeForInpaint", + "pos": [ + 1399.0002441406255, + 495.19991455078116 + ], + "size": { + "0": 210, + "1": 66 + }, + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "pixels", + "type": "IMAGE", + "link": 43 + }, + { + "name": "vae", + "type": "VAE", + "link": 44 + }, + { + "name": "mask", + "type": "MASK", + "link": 42 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 45 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEEncodeForInpaint" + } + }, + { + "id": 31, + "type": "BitwiseAndMask", + "pos": [ + 1399.0002441406255, + 619.1999145507809 + ], + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "mask1", + "type": "MASK", + "link": 40 + }, + { + "name": "mask2", + "type": "MASK", + "link": 41 + } + ], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [ + 42 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "BitwiseAndMask" + } + }, + { + "id": 33, + "type": "VAEDecode", + "pos": [ + 1399.0002441406255, + 722.1999145507809 + ], + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 45 + }, + { + "name": "vae", + "type": "VAE", + "link": 46 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 47 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEDecode" + } + }, + { + "id": 34, + "type": "PreviewImage", + "pos": [ + 1641.0002441406255, + 494.19991455078116 + ], + "size": [ + 421, + 271 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 47 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 22, + "type": "BboxDetector", + "pos": [ + 963, + 276 + ], + "size": { + "0": 315, + "1": 102 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "bbox_model", + "type": "BBOX_MODEL", + "link": 24 + }, + { + "name": "image", + "type": "IMAGE", + "link": 25 + } + ], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [ + 26, + 41 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "BboxDetector" + }, + "widgets_values": [ + 0.5, + 60 + ] + }, + { + "id": 24, + "type": "VAELoader", + "pos": [ + 221, + 617 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 3, + "mode": 0, + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 28, + 30, + 35, + 38, + 44, + 46 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAELoader" + }, + "widgets_values": [ + "vae-ft-mse-840000-ema-pruned.safetensors" + ] + }, + { + "id": 1, + "type": "LoadImage", + "pos": [ + 669, + 498 + ], + "size": [ + 441, + 293 + ], + "flags": {}, + "order": 2, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 25, + 27, + 33, + 36, + 43 + ], + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "joyouriz.png", + "image" + ] } ], "links": [ [ - 18, + 24, 18, 0, - 17, + 22, 0, "BBOX_MODEL" ], [ - 19, - 20, + 25, 1, - 21, 0, - "SEGM_MODEL" + 22, + 1, + "IMAGE" ], [ - 20, - 17, + 26, + 22, + 0, + 23, + 2, + "MASK" + ], + [ + 27, + 1, + 0, + 23, + 0, + "IMAGE" + ], + [ + 28, + 24, + 0, + 23, + 1, + "VAE" + ], + [ + 29, + 23, + 0, + 25, + 0, + "LATENT" + ], + [ + 30, + 24, + 0, + 25, + 1, + "VAE" + ], + [ + 31, + 25, 0, 4, 0, "IMAGE" ], [ - 21, + 32, + 26, + 1, + 27, + 0, + "SEGM_MODEL" + ], + [ + 33, 1, 0, - 17, + 27, 1, "IMAGE" ], [ - 22, + 34, + 27, + 0, + 30, + 2, + "MASK" + ], + [ + 35, + 24, + 0, + 30, + 1, + "VAE" + ], + [ + 36, 1, 0, - 21, - 1, + 30, + 0, "IMAGE" ], [ - 23, - 21, + 37, + 30, + 0, + 28, + 0, + "LATENT" + ], + [ + 38, + 24, + 0, + 28, + 1, + "VAE" + ], + [ + 39, + 28, 0, 9, 0, "IMAGE" + ], + [ + 40, + 27, + 0, + 31, + 0, + "MASK" + ], + [ + 41, + 22, + 0, + 31, + 1, + "MASK" + ], + [ + 42, + 31, + 0, + 32, + 2, + "MASK" + ], + [ + 43, + 1, + 0, + 32, + 0, + "IMAGE" + ], + [ + 44, + 24, + 0, + 32, + 1, + "VAE" + ], + [ + 45, + 32, + 0, + 33, + 0, + "LATENT" + ], + [ + 46, + 24, + 0, + 33, + 1, + "VAE" + ], + [ + 47, + 33, + 0, + 34, + 0, + "IMAGE" ] ], - "groups": [], + "groups": [ + { + "title": "Face Detection", + "bounding": [ + 210, + 15, + 1130, + 404 + ], + "color": "#3f789e" + }, + { + "title": "Person Detection", + "bounding": [ + 220, + 844, + 1156, + 367 + ], + "color": "#3f789e" + }, + { + "title": "Head Mask", + "bounding": [ + 1350, + 374, + 780, + 466 + ], + "color": "#3f789e" + } + ], "config": {}, "extra": {}, "version": 0.4 diff --git a/misc/example.png b/misc/example.png index f8bf3ca..ee9c0a2 100644 Binary files a/misc/example.png and b/misc/example.png differ