| 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 |
