ONNXLoader, ONNXDetector added
This commit is contained in:
Dr.Lt.Data
@@ -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)
|
||||
|
||||
+130
-41
@@ -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,
|
||||
|
||||
+416
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user