ddetailer work?

scripts/ddetailer.py

@@ -651,6 +651,7 @@ def create_segmask_preview(results, image):
     labels = results[0]
     bboxes = results[1]
     segms = results[2]
+    scores = results[3]
 
     cv2_image = np.array(image)
     cv2_image = cv2_image[:, :, ::-1].copy()
@@ -670,7 +671,7 @@ def create_segmask_preview(results, image):
         cv2_image = np.where(cv2_mask_rgb == 255, color_image, cv2_image)
         text_color = tuple([int(x) for x in (color[0][0] - 100)])
         name = labels[i]
-        score = bboxes[i][4]
+        score = scores[i]
         score = str(score)[:4]
         text = name + ":" + score
         cv2.putText(
@@ -792,7 +793,6 @@ def create_segmasks(results):
     return segmasks
 
 
-import mmcv
 from mmdet.apis import inference_detector, init_detector
 from mmdet.evaluation import get_classes
 
@@ -819,25 +819,24 @@ def inference_mmdet_segm(image, modelname, conf_thres, label):
     model = init_detector(
         model_config, model_checkpoint, palette="random", device=model_device
     )
-    mmdet_results = inference_detector(model, np.array(image))
-    bbox_results, segm_results = mmdet_results
+    mmdet_results = inference_detector(model, np.array(image)).pred_instances
+    bboxes = mmdet_results.bboxes.numpy()
+    segms = mmdet_results.masks.numpy()
+    scores = mmdet_results.scores.numpy()
     dataset = modeldataset(modelname)
     classes = get_classes(dataset)
-    labels = [
-        np.full(bbox.shape[0], i, dtype=np.int32) for i, bbox in enumerate(bbox_results)
-    ]
-    n, m = bbox_results[0].shape
+
+    n, m = bboxes.shape
     if n == 0:
-        return [[], [], []]
-    labels = np.concatenate(labels)
-    bboxes = np.vstack(bbox_results)
-    segms = mmcv.concat_list(segm_results)
-    filter_inds = np.where(bboxes[:, -1] > conf_thres)[0]
-    results = [[], [], []]
+        return [[], [], [], []]
+    labels = mmdet_results.labels
+    filter_inds = np.where(mmdet_results.scores > conf_thres)[0]
+    results = [[], [], [], []]
     for i in filter_inds:
         results[0].append(label + "-" + classes[labels[i]])
         results[1].append(bboxes[i])
         results[2].append(segms[i])
+        results[3].append(scores[i])
 
     return results
 
@@ -849,28 +848,30 @@ def inference_mmdet_bbox(image, modelname, conf_thres, label):
     model = init_detector(
         model_config, model_checkpoint, palette="random", device=model_device
     )
-    results = inference_detector(model, np.array(image))
+    output = inference_detector(model, np.array(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, _conf in results[0]:
+    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 = results[0].shape
+    n, m = output.bboxes.shape
     if n == 0:
-        return [[], [], []]
-    bboxes = np.vstack(results[0])
-    filter_inds = np.where(bboxes[:, -1] > conf_thres)[0]
-    results = [[], [], []]
+        return [[], [], [], []]
+    bboxes = output.bboxes.numpy()
+    scores = output.scores.numpy()
+    filter_inds = np.where(scores > conf_thres)[0]
+    results = [[], [], [], []]
     for i in filter_inds:
         results[0].append(label)
         results[1].append(bboxes[i])
         results[2].append(segms[i])
+        results[3].append(scores[i])
 
     return results