From 96e32784d097fa2e4a7dee81768a9ada624fe713 Mon Sep 17 00:00:00 2001 From: "Dr.Lt.Data" Date: Tue, 11 Apr 2023 18:39:37 +0900 Subject: [PATCH] ONNXLoader, ONNXDetector added --- README.md | 2 + comfyui-impact-pack.py | 171 +++++++++++++---- misc/onnx.json | 416 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 548 insertions(+), 41 deletions(-) create mode 100644 misc/onnx.json diff --git a/README.md b/README.md index 819eff8..40e19ab 100644 --- a/README.md +++ b/README.md @@ -5,6 +5,7 @@ # Features * MMDetLoader - Load MMDet model * SAMLoader - Load SAM model +* ONNXLoader - Load ONNX model * SegmDetectorCombined - Detect segmentation and return mask from input image. * BboxDetectorCombined - Detect bbox(bounding box) and return mask from input image. * SamDetectorCombined - Using the technology of SAM, extract the segment at the location indicated by the input SEGS on the input image, and output it as a unified mask. @@ -12,6 +13,7 @@ * SubtractMask - Perform subtract operations between 2 masks * SegmDetectorForEach - Detect segmentation and return SEGS from input image. * BboxDetectorForEach - Detect bbox(bounding box) and return SEGS from input image. +* ONNXDetectorForeach - Using the ONNX model, identify the bbox and retrieve the SEGS from the input image * DetailerForEach - Refine image rely on SEGS. * DetailerForEachDebug - Refine image rely on SEGS. Additionally, you can monitor cropped image and refined image of cropped image. * The 'DetailerForEach' and 'DetailerForEachDebug' now support an 'external_seed' that is obtained from the Seed node on the [WAS suite](https://github.com/WASasquatch/was-node-suite-comfyui) diff --git a/comfyui-impact-pack.py b/comfyui-impact-pack.py index ad8af5b..e2cffe3 100644 --- a/comfyui-impact-pack.py +++ b/comfyui-impact-pack.py @@ -23,6 +23,11 @@ def ensure_pip_packages(): except Exception: subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'scikit-image']) + try: + import onnxruntime + except Exception: + subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'onnxruntime']) + try: import pycocotools except Exception: @@ -68,6 +73,7 @@ model_path = folder_paths.models_dir bbox_path = os.path.join(model_path, "mmdets", "bbox") #segm_path = os.path.join(model_path, "mmdets", "segm") -- deprecated sam_path = os.path.join(model_path, "sams") +onnx_path = os.path.join(model_path, "onnx") if not os.path.exists(os.path.join(bbox_path, "mmdet_anime-face_yolov3.pth")): download_url("https://huggingface.co/dustysys/ddetailer/resolve/main/mmdet/bbox/mmdet_anime-face_yolov3.pth", bbox_path) @@ -78,6 +84,10 @@ if not os.path.exists(os.path.join(bbox_path, "mmdet_anime-face_yolov3.py")): if not os.path.exists(os.path.join(sam_path, "sam_vit_b_01ec64.pth")): download_url("https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth", sam_path) +if not os.path.exists(onnx_path): + print(f"### ComfyUI-Impact-Pack: onnx model directory created ({onnx_path})") + os.mkdir(onnx_path) + # ----- MAIN CODE -------------------------------------------------------------- # Core @@ -290,6 +300,7 @@ folder_paths.folder_names_and_paths["mmdets_bbox"] = ([os.path.join(model_path, folder_paths.folder_names_and_paths["mmdets_segm"] = ([os.path.join(model_path, "mmdets", "segm")], folder_paths.supported_pt_extensions) folder_paths.folder_names_and_paths["mmdets"] = ([os.path.join(model_path, "mmdets")], folder_paths.supported_pt_extensions) folder_paths.folder_names_and_paths["sams"] = ([os.path.join(model_path, "sams")], folder_paths.supported_pt_extensions) +folder_paths.folder_names_and_paths["onnx"] = ([os.path.join(model_path, "onnx")], set(['.onnx'])) class NO_BBOX_MODEL: @@ -300,47 +311,6 @@ class NO_SEGM_MODEL: ERROR = "" -class MMDetLoader: - @classmethod - def INPUT_TYPES(s): - bboxs = [ "bbox/"+x for x in folder_paths.get_filename_list("mmdets_bbox") ] - segms = [ "segm/"+x for x in folder_paths.get_filename_list("mmdets_segm") ] - return {"required": { "model_name": (bboxs + segms, )}} - RETURN_TYPES = ("BBOX_MODEL", "SEGM_MODEL") - FUNCTION = "load_mmdet" - - CATEGORY = "ImpactPack" - - def load_mmdet(self, model_name): - mmdet_path = folder_paths.get_full_path("mmdets", model_name) - model = load_mmdet(mmdet_path) - - if model_name.startswith("bbox"): - return (model, NO_SEGM_MODEL()) - else: - return (NO_BBOX_MODEL(), model) - - -from segment_anything import build_sam, SamPredictor -from segment_anything import sam_model_registry - - -class SAMLoader: - @classmethod - def INPUT_TYPES(s): - return {"required": { "model_name": (folder_paths.get_filename_list("sams"), )}} - RETURN_TYPES = ("SAM_MODEL", ) - FUNCTION = "load_model" - - CATEGORY = "ImpactPack" - - def load_model(selfself, model_name): - modelname = folder_paths.get_full_path("sams", model_name) - sam = sam_model_registry["vit_b"](checkpoint=modelname) - print(f"Loads SAM model: {modelname}") - return (sam, ) - - def normalize_region(limit, startp, size): if startp < 0: new_endp = size @@ -501,6 +471,122 @@ def composite_to(dest_latent, crop_region, src_latent): return orig_image[0] + +class MMDetLoader: + @classmethod + def INPUT_TYPES(s): + bboxs = ["bbox/"+x for x in folder_paths.get_filename_list("mmdets_bbox")] + segms = ["segm/"+x for x in folder_paths.get_filename_list("mmdets_segm")] + return {"required": {"model_name": (bboxs + segms, )}} + RETURN_TYPES = ("BBOX_MODEL", "SEGM_MODEL") + FUNCTION = "load_mmdet" + + CATEGORY = "ImpactPack" + + def load_mmdet(self, model_name): + mmdet_path = folder_paths.get_full_path("mmdets", model_name) + model = load_mmdet(mmdet_path) + + if model_name.startswith("bbox"): + return (model, NO_SEGM_MODEL()) + else: + return (NO_BBOX_MODEL(), model) + + +from segment_anything import build_sam, SamPredictor +from segment_anything import sam_model_registry +import onnxruntime + + +class SAMLoader: + @classmethod + def INPUT_TYPES(s): + return {"required": { "model_name": (folder_paths.get_filename_list("sams"), )}} + + RETURN_TYPES = ("SAM_MODEL", ) + FUNCTION = "load_model" + + CATEGORY = "ImpactPack" + + def load_model(self, model_name): + modelname = folder_paths.get_full_path("sams", model_name) + sam = sam_model_registry["vit_b"](checkpoint=modelname) + print(f"Loads SAM model: {modelname}") + return (sam, ) + + +class ONNXLoader: + @classmethod + def INPUT_TYPES(s): + return {"required": {"model_name": (folder_paths.get_filename_list("onnx"), )}} + + RETURN_TYPES = ("ONNX_MODEL", ) + FUNCTION = "load_model" + + CATEGORY = "ImpactPack" + + def load_model(self, model_name): + modelname = folder_paths.get_full_path("onnx", model_name) + print(f"Loads ONNX model: {modelname}") + return (modelname, ) + + +class ONNXDetectorForEach: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "onnx_model": ("ONNX_MODEL",), + "image": ("IMAGE",), + "threshold": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("SEGS", ) + FUNCTION = "doit" + + CATEGORY = "ImpactPack/Detector" + + OUTPUT_NODE = True + + def doit(self, onnx_model, image, threshold): + # prepare image + pil = tensor2pil(image) + image = np.ascontiguousarray(pil) + image = image[:, :, ::-1] # to BGR image + image = image.astype(np.float32) + image -= [103.939, 116.779, 123.68] # 'caffe' mode image preprocessing + + # do detection + onnx_model = onnxruntime.InferenceSession(onnx_model) + outputs = onnx_model.run( + [s_i.name for s_i in onnx_model.get_outputs()], + {onnx_model.get_inputs()[0].name: np.expand_dims(image, axis=0)}, + ) + + labels = [op for op in outputs if op.dtype == "int32"][0] + scores = [op for op in outputs if isinstance(op[0][0], np.float32)][0] + boxes = [op for op in outputs if isinstance(op[0][0], np.ndarray)][0] + + # filter-out useless item + idx = np.where(labels[0] == -1)[0][0] + + labels = labels[0][:idx] + scores = scores[0][:idx] + boxes = boxes[0][:idx].astype(np.uint32) + + # collect feasible item + result = [] + + for i in range(len(labels)): + if scores[i] > threshold: + x1, y1, x2, y2 = boxes[i] + + mask = np.ones((y2-y1,x2-x1)) + item = (None, mask, scores[i], boxes[i], boxes[i]) + result.append(item) + + return (result,) + class DetailerForEach: @classmethod def INPUT_TYPES(s): @@ -1133,9 +1219,12 @@ class SubtractMask: NODE_CLASS_MAPPINGS = { "MMDetLoader": MMDetLoader, "SAMLoader": SAMLoader, + "ONNXLoader": ONNXLoader, "BboxDetectorForEach": BboxDetectorForEach, "SegmDetectorForEach": SegmDetectorForEach, + "ONNXDetectorForEach": ONNXDetectorForEach, + "BitwiseAndMaskForEach": BitwiseAndMaskForEach, "DetailerForEach": DetailerForEach, diff --git a/misc/onnx.json b/misc/onnx.json new file mode 100644 index 0000000..a5799fc --- /dev/null +++ b/misc/onnx.json @@ -0,0 +1,416 @@ +{ + "last_node_id": 13, + "last_link_id": 15, + "nodes": [ + { + "id": 5, + "type": "SegsMaskCombine", + "pos": [ + 327, + 530 + ], + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "segs", + "type": "SEGS", + "link": 15 + }, + { + "name": "image", + "type": "IMAGE", + "link": 4 + } + ], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [ + 6 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "SegsMaskCombine" + } + }, + { + "id": 6, + "type": "VAEEncodeForInpaint", + "pos": [ + 609, + 493 + ], + "size": { + "0": 226.8000030517578, + "1": 66 + }, + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "pixels", + "type": "IMAGE", + "link": 7 + }, + { + "name": "vae", + "type": "VAE", + "link": 8 + }, + { + "name": "mask", + "type": "MASK", + "link": 6 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 9 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAE Encode (for Inpainting)" + } + }, + { + "id": 8, + "type": "VAEDecode", + "pos": [ + 858, + 495 + ], + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 9 + }, + { + "name": "vae", + "type": "VAE", + "link": 10 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 11 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAE Decode" + } + }, + { + "id": 9, + "type": "SaveImage", + "pos": [ + 1107, + 494 + ], + "size": { + "0": 315, + "1": 250 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 11 + } + ], + "properties": {}, + "widgets_values": [ + "ComfyUI" + ] + }, + { + "id": 7, + "type": "VAELoader", + "pos": [ + 327, + 684 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 8, + 10 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "Load VAE" + }, + "widgets_values": [ + "vae-ft-mse-840000-ema-pruned.safetensors" + ] + }, + { + "id": 11, + "type": "ONNXLoader", + "pos": [ + 672.1527099609375, + 363.65423583984375 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 1, + "mode": 0, + "outputs": [ + { + "name": "ONNX_MODEL", + "type": "ONNX_MODEL", + "links": null + } + ], + "properties": { + "Node name for S&R": "ONNXLoader" + }, + "widgets_values": [ + "detector_v2_base_checkpoint.onnx" + ] + }, + { + "id": 1, + "type": "ONNXLoader", + "pos": [ + -25, + 389 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 2, + "mode": 0, + "outputs": [ + { + "name": "ONNX_MODEL", + "type": "ONNX_MODEL", + "links": [ + 13 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ONNXLoader" + }, + "widgets_values": [ + "detector_v2_base_checkpoint.onnx" + ] + }, + { + "id": 2, + "type": "LoadImage", + "pos": [ + -23, + 504 + ], + "size": { + "0": 312.3335266113281, + "1": 241.66659545898438 + }, + "flags": {}, + "order": 3, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 4, + 7, + 14 + ], + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "Node name for S&R": "Load Image" + }, + "widgets_values": [ + "rga6845.jpeg", + "image" + ] + }, + { + "id": 13, + "type": "ONNXDetectorForEach", + "pos": [ + 325, + 391 + ], + "size": [ + 210, + 78 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "onnx_model", + "type": "ONNX_MODEL", + "link": 13 + }, + { + "name": "image", + "type": "IMAGE", + "link": 14 + } + ], + "outputs": [ + { + "name": "SEGS", + "type": "SEGS", + "links": [ + 15 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ONNXDetectorForEach" + }, + "widgets_values": [ + 0.5 + ] + } + ], + "links": [ + [ + 4, + 2, + 0, + 5, + 1, + "IMAGE" + ], + [ + 6, + 5, + 0, + 6, + 2, + "MASK" + ], + [ + 7, + 2, + 0, + 6, + 0, + "IMAGE" + ], + [ + 8, + 7, + 0, + 6, + 1, + "VAE" + ], + [ + 9, + 6, + 0, + 8, + 0, + "LATENT" + ], + [ + 10, + 7, + 0, + 8, + 1, + "VAE" + ], + [ + 11, + 8, + 0, + 9, + 0, + "IMAGE" + ], + [ + 13, + 1, + 0, + 13, + 0, + "ONNX_MODEL" + ], + [ + 14, + 2, + 0, + 13, + 1, + "IMAGE" + ], + [ + 15, + 13, + 0, + 5, + 0, + "SEGS" + ] + ], + "groups": [], + "config": {}, + "extra": {}, + "version": 0.4 +} \ No newline at end of file