|
import folder_paths
|
|
from impact.core import *
|
|
import os
|
|
|
|
import mmcv
|
|
from mmdet.apis import (inference_detector, init_detector)
|
|
from mmdet.evaluation import get_classes
|
|
|
|
|
|
def load_mmdet(model_path):
|
|
model_config = os.path.splitext(model_path)[0] + ".py"
|
|
model = init_detector(model_config, model_path, device="cpu")
|
|
return model
|
|
|
|
|
|
def inference_segm_old(model, image, conf_threshold):
|
|
image = image.numpy()[0] * 255
|
|
mmdet_results = inference_detector(model, image)
|
|
|
|
bbox_results, segm_results = mmdet_results
|
|
label = "A"
|
|
|
|
classes = get_classes("coco")
|
|
labels = [
|
|
np.full(bbox.shape[0], i, dtype=np.int32)
|
|
for i, bbox in enumerate(bbox_results)
|
|
]
|
|
n, m = bbox_results[0].shape
|
|
if n == 0:
|
|
return [[], [], []]
|
|
labels = np.concatenate(labels)
|
|
bboxes = np.vstack(bbox_results)
|
|
segms = mmcv.concat_list(segm_results)
|
|
filter_idxs = np.where(bboxes[:, -1] > conf_threshold)[0]
|
|
results = [[], [], []]
|
|
for i in filter_idxs:
|
|
results[0].append(label + "-" + classes[labels[i]])
|
|
results[1].append(bboxes[i])
|
|
results[2].append(segms[i])
|
|
|
|
return results
|
|
|
|
|
|
def inference_segm(image, modelname, conf_thres, lab="A"):
|
|
image = image.numpy()[0] * 255
|
|
mmdet_results = inference_detector(modelname, image).pred_instances
|
|
bboxes = mmdet_results.bboxes.numpy()
|
|
segms = mmdet_results.masks.numpy()
|
|
scores = mmdet_results.scores.numpy()
|
|
|
|
classes = get_classes("coco")
|
|
|
|
n, m = bboxes.shape
|
|
if n == 0:
|
|
return [[], [], [], []]
|
|
labels = mmdet_results.labels
|
|
filter_inds = np.where(mmdet_results.scores > conf_thres)[0]
|
|
results = [[], [], [], []]
|
|
for i in filter_inds:
|
|
results[0].append(lab + "-" + classes[labels[i]])
|
|
results[1].append(bboxes[i])
|
|
results[2].append(segms[i])
|
|
results[3].append(scores[i])
|
|
|
|
return results
|
|
|
|
|
|
def inference_bbox(modelname, image, conf_threshold):
|
|
image = image.numpy()[0] * 255
|
|
label = "A"
|
|
output = inference_detector(modelname, image).pred_instances
|
|
cv2_image = np.array(image)
|
|
cv2_image = cv2_image[:, :, ::-1].copy()
|
|
cv2_gray = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2GRAY)
|
|
|
|
segms = []
|
|
for x0, y0, x1, y1 in output.bboxes:
|
|
cv2_mask = np.zeros(cv2_gray.shape, np.uint8)
|
|
cv2.rectangle(cv2_mask, (int(x0), int(y0)), (int(x1), int(y1)), 255, -1)
|
|
cv2_mask_bool = cv2_mask.astype(bool)
|
|
segms.append(cv2_mask_bool)
|
|
|
|
n, m = output.bboxes.shape
|
|
if n == 0:
|
|
return [[], [], [], []]
|
|
|
|
bboxes = output.bboxes.numpy()
|
|
scores = output.scores.numpy()
|
|
filter_idxs = np.where(scores > conf_threshold)[0]
|
|
results = [[], [], [], []]
|
|
for i in filter_idxs:
|
|
results[0].append(label)
|
|
results[1].append(bboxes[i])
|
|
results[2].append(segms[i])
|
|
results[3].append(scores[i])
|
|
|
|
return results
|
|
|
|
|
|
class BBoxDetector:
|
|
bbox_model = None
|
|
|
|
def __init__(self, bbox_model):
|
|
self.bbox_model = bbox_model
|
|
|
|
def detect(self, image, threshold, dilation, crop_factor, drop_size=1, detailer_hook=None):
|
|
drop_size = max(drop_size, 1)
|
|
mmdet_results = inference_bbox(self.bbox_model, image, threshold)
|
|
segmasks = create_segmasks(mmdet_results)
|
|
|
|
if dilation > 0:
|
|
segmasks = dilate_masks(segmasks, dilation)
|
|
|
|
items = []
|
|
h = image.shape[1]
|
|
w = image.shape[2]
|
|
|
|
for x in segmasks:
|
|
item_bbox = x[0]
|
|
item_mask = x[1]
|
|
|
|
y1, x1, y2, x2 = item_bbox
|
|
|
|
if x2 - x1 > drop_size and y2 - y1 > drop_size:
|
|
crop_region = make_crop_region(w, h, item_bbox, crop_factor)
|
|
cropped_image = crop_image(image, crop_region)
|
|
cropped_mask = crop_ndarray2(item_mask, crop_region)
|
|
confidence = x[2]
|
|
|
|
|
|
item = SEG(cropped_image, cropped_mask, confidence, crop_region, item_bbox, None, None)
|
|
|
|
items.append(item)
|
|
|
|
shape = image.shape[1], image.shape[2]
|
|
return shape, items
|
|
|
|
def detect_combined(self, image, threshold, dilation):
|
|
mmdet_results = inference_bbox(self.bbox_model, image, threshold)
|
|
segmasks = create_segmasks(mmdet_results)
|
|
if dilation > 0:
|
|
segmasks = dilate_masks(segmasks, dilation)
|
|
|
|
return combine_masks(segmasks)
|
|
|
|
def setAux(self, x):
|
|
pass
|
|
|
|
|
|
class SegmDetector(BBoxDetector):
|
|
segm_model = None
|
|
|
|
def __init__(self, segm_model):
|
|
self.segm_model = segm_model
|
|
|
|
def detect(self, image, threshold, dilation, crop_factor, drop_size=1, detailer_hook=None):
|
|
drop_size = max(drop_size, 1)
|
|
mmdet_results = inference_segm(image, self.segm_model, threshold)
|
|
segmasks = create_segmasks(mmdet_results)
|
|
|
|
if dilation > 0:
|
|
segmasks = dilate_masks(segmasks, dilation)
|
|
|
|
items = []
|
|
h = image.shape[1]
|
|
w = image.shape[2]
|
|
for x in segmasks:
|
|
item_bbox = x[0]
|
|
item_mask = x[1]
|
|
|
|
y1, x1, y2, x2 = item_bbox
|
|
|
|
if x2 - x1 > drop_size and y2 - y1 > drop_size:
|
|
crop_region = make_crop_region(w, h, item_bbox, crop_factor)
|
|
cropped_image = crop_image(image, crop_region)
|
|
cropped_mask = crop_ndarray2(item_mask, crop_region)
|
|
confidence = x[2]
|
|
|
|
item = SEG(cropped_image, cropped_mask, confidence, crop_region, item_bbox, None, None)
|
|
items.append(item)
|
|
|
|
segs = image.shape, items
|
|
|
|
if detailer_hook is not None and hasattr(detailer_hook, "post_detection"):
|
|
segs = detailer_hook.post_detection(segs)
|
|
|
|
return segs
|
|
|
|
def detect_combined(self, image, threshold, dilation):
|
|
mmdet_results = inference_bbox(self.bbox_model, image, threshold)
|
|
segmasks = create_segmasks(mmdet_results)
|
|
if dilation > 0:
|
|
segmasks = dilate_masks(segmasks, dilation)
|
|
|
|
return combine_masks(segmasks)
|
|
|
|
def setAux(self, x):
|
|
pass
|
|
|
|
|
|
class MMDetDetectorProvider:
|
|
@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_DETECTOR", "SEGM_DETECTOR")
|
|
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 BBoxDetector(model), NO_SEGM_DETECTOR()
|
|
else:
|
|
return NO_BBOX_DETECTOR(), model |
