split output model to bbox_model, segm_model
show error message for invalid model
This commit is contained in:
Dr.Lt.Data
+32
-7
@@ -63,7 +63,10 @@ from mmdet.apis import (inference_detector,
|
||||
init_detector)
|
||||
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
|
||||
@@ -154,13 +157,19 @@ 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)
|
||||
|
||||
class NO_BBOX_MODEL:
|
||||
ERROR = ""
|
||||
|
||||
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 = ("MODEL", )
|
||||
RETURN_TYPES = ("BBOX_MODEL", "SEGM_MODEL")
|
||||
FUNCTION = "load_mmdet"
|
||||
|
||||
CATEGORY = "ImpactPack"
|
||||
@@ -168,7 +177,11 @@ class MMDetLoader:
|
||||
def load_mmdet(self, model_name):
|
||||
mmdet_path = folder_paths.get_full_path("mmdets", model_name)
|
||||
model = load_mmdet(mmdet_path)
|
||||
return (model, )
|
||||
|
||||
if model_name.startswith("bbox"):
|
||||
return (model, NO_SEGM_MODEL())
|
||||
else:
|
||||
return (NO_BBOX_MODEL(), model)
|
||||
|
||||
class SegmDetector:
|
||||
input_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input")
|
||||
@@ -176,7 +189,7 @@ class SegmDetector:
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{
|
||||
"model": ("MODEL", ),
|
||||
"segm_model": ("SEGM_MODEL", ),
|
||||
"image": ("IMAGE", ),
|
||||
"threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
}
|
||||
@@ -187,8 +200,8 @@ class SegmDetector:
|
||||
|
||||
CATEGORY = "ImpactPack"
|
||||
|
||||
def doit(self, model, image, threshold):
|
||||
mmdet_results = inference_segm(model, image, threshold)
|
||||
def doit(self, segm_model, image, threshold):
|
||||
mmdet_results = inference_segm(segm_model, image, threshold)
|
||||
segmasks = create_segmasks(mmdet_results)
|
||||
mask = combine_masks(segmasks)
|
||||
|
||||
@@ -196,10 +209,22 @@ class SegmDetector:
|
||||
return (image,)
|
||||
|
||||
class BboxDetector(SegmDetector):
|
||||
def doit(self, model, image, threshold):
|
||||
mmdet_results = inference_bbox(model, image, threshold)
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{
|
||||
"bbox_model": ("BBOX_MODEL", ),
|
||||
"image": ("IMAGE", ),
|
||||
"threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def doit(self, bbox_model, image, threshold):
|
||||
mmdet_results = inference_bbox(bbox_model, image, threshold)
|
||||
segmasks = create_segmasks(mmdet_results)
|
||||
mask = combine_masks(segmasks)
|
||||
|
||||
image = pil2tensor(mask)
|
||||
return (image,)
|
||||
|
||||
|
||||
+221
-213
@@ -1,153 +1,7 @@
|
||||
{
|
||||
"last_node_id": 9,
|
||||
"last_link_id": 10,
|
||||
"last_node_id": 21,
|
||||
"last_link_id": 23,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 7,
|
||||
"type": "MMDetLoader",
|
||||
"pos": [
|
||||
515,
|
||||
771
|
||||
],
|
||||
"size": {
|
||||
"0": 443,
|
||||
"1": 58
|
||||
},
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "MODEL",
|
||||
"type": "MODEL",
|
||||
"links": [
|
||||
7
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "MMDetLoader"
|
||||
},
|
||||
"widgets_values": [
|
||||
"segm/mmdet_dd-person_mask2former.pth"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"type": "MMDetLoader",
|
||||
"pos": [
|
||||
515,
|
||||
180
|
||||
],
|
||||
"size": {
|
||||
"0": 439,
|
||||
"1": 62
|
||||
},
|
||||
"flags": {
|
||||
"collapsed": false
|
||||
},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "MODEL",
|
||||
"type": "MODEL",
|
||||
"links": [
|
||||
4
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "MMDetLoader"
|
||||
},
|
||||
"widgets_values": [
|
||||
"bbox/mmdet_anime-face_yolov3.pth"
|
||||
],
|
||||
"shape": 2
|
||||
},
|
||||
{
|
||||
"id": 6,
|
||||
"type": "SegmDetector",
|
||||
"pos": [
|
||||
178,
|
||||
765
|
||||
],
|
||||
"size": {
|
||||
"0": 210,
|
||||
"1": 78
|
||||
},
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "model",
|
||||
"type": "MODEL",
|
||||
"link": 7
|
||||
},
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"link": 8
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
10
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "SegmDetector"
|
||||
},
|
||||
"widgets_values": [
|
||||
0.49999999999999956
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 1,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
510,
|
||||
314
|
||||
],
|
||||
"size": {
|
||||
"0": 448,
|
||||
"1": 387
|
||||
},
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
6,
|
||||
8
|
||||
],
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"snh.jpg",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"type": "PreviewImage",
|
||||
@@ -166,7 +20,7 @@
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 5
|
||||
"link": 20
|
||||
}
|
||||
],
|
||||
"title": "PreviewImage - Faces",
|
||||
@@ -174,49 +28,6 @@
|
||||
"Node name for S&R": "PreviewImage"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 5,
|
||||
"type": "BboxDetector",
|
||||
"pos": [
|
||||
1048,
|
||||
180
|
||||
],
|
||||
"size": {
|
||||
"0": 210,
|
||||
"1": 78
|
||||
},
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "model",
|
||||
"type": "MODEL",
|
||||
"link": 4
|
||||
},
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"link": 6
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
5
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "BboxDetector"
|
||||
},
|
||||
"widgets_values": [
|
||||
0.49999999999999956
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 9,
|
||||
"type": "PreviewImage",
|
||||
@@ -235,59 +46,256 @@
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 10
|
||||
"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
|
||||
],
|
||||
"size": {
|
||||
"0": 437,
|
||||
"1": 87
|
||||
},
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "BBOX_MODEL",
|
||||
"type": "BBOX_MODEL",
|
||||
"links": [
|
||||
18
|
||||
],
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "SEGM_MODEL",
|
||||
"type": "SEGM_MODEL",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "MMDetLoader"
|
||||
},
|
||||
"widgets_values": [
|
||||
"bbox/mmdet_anime-face_yolov3.pth"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 20,
|
||||
"type": "MMDetLoader",
|
||||
"pos": [
|
||||
523,
|
||||
801
|
||||
],
|
||||
"size": {
|
||||
"0": 437,
|
||||
"1": 87
|
||||
},
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "BBOX_MODEL",
|
||||
"type": "BBOX_MODEL",
|
||||
"links": [],
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "SEGM_MODEL",
|
||||
"type": "SEGM_MODEL",
|
||||
"links": [
|
||||
19
|
||||
],
|
||||
"slot_index": 1
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "MMDetLoader"
|
||||
},
|
||||
"widgets_values": [
|
||||
"segm/mmdet_dd-person_mask2former.pth"
|
||||
]
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
4,
|
||||
2,
|
||||
18,
|
||||
18,
|
||||
0,
|
||||
5,
|
||||
17,
|
||||
0,
|
||||
"MODEL"
|
||||
"BBOX_MODEL"
|
||||
],
|
||||
[
|
||||
5,
|
||||
5,
|
||||
19,
|
||||
20,
|
||||
1,
|
||||
21,
|
||||
0,
|
||||
"SEGM_MODEL"
|
||||
],
|
||||
[
|
||||
20,
|
||||
17,
|
||||
0,
|
||||
4,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
6,
|
||||
21,
|
||||
1,
|
||||
0,
|
||||
5,
|
||||
17,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
7,
|
||||
7,
|
||||
0,
|
||||
6,
|
||||
0,
|
||||
"MODEL"
|
||||
],
|
||||
[
|
||||
8,
|
||||
22,
|
||||
1,
|
||||
0,
|
||||
6,
|
||||
21,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
10,
|
||||
6,
|
||||
23,
|
||||
21,
|
||||
0,
|
||||
9,
|
||||
0,
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 194 KiB After Width: | Height: | Size: 204 KiB |
Reference in New Issue
Block a user