soft nms support

Author: GlassyWing

yolo3/detect/img_detect.py

@@ -15,7 +15,8 @@
 
 from yolo3.dataset.dataset import pad_to_square, resize
 from yolo3.utils.helper import load_classes
-from yolo3.utils.model_build import non_max_suppression, rescale_boxes, xywh2p1p2, resize_boxes
+from yolo3.utils.model_build import non_max_suppression, rescale_boxes, xywh2p1p2, resize_boxes, \
+    soft_non_max_suppression
 
 
 def scale(image, shape, max_size):
@@ -83,7 +84,7 @@ def detect(self, img):
             prev_time = time.time()
             with torch.no_grad():
                 detections = self.model(image)
-                detections = non_max_suppression(detections, self.thres, self.nms_thres)
+                detections = soft_non_max_suppression(detections, self.thres, self.nms_thres)
                 detections = detections[0]
                 if detections is not None:
                     # detections = rescale_boxes(detections, self.model.img_size, (h, w))
@@ -140,7 +141,9 @@ def detect(self, img):
                 # (1, n, 5 + num_class)
                 rescaled_detections = torch.cat(rescaled_detections, 0).unsqueeze(0)
 
-                detections = non_max_suppression(rescaled_detections, self.thres, self.nms_thres, is_p1p2=True)
+                detections = soft_non_max_suppression(rescaled_detections, self.thres, self.nms_thres,
+                                                      merge=True,
+                                                      is_p1p2=True)
                 detections = detections[0]
 
             current_time = time.time()

yolo3/utils/label_draw.py

@@ -4,12 +4,13 @@
 import numpy as np
 import time
 
+
 def _get_statistic_info(detections, unique_labels, classes):
     """获得统计信息"""
     statistic_info = {}
     for label in unique_labels:
         statistic_info[classes[int(label)]] = (
-            detections[:, -1] == label).sum().item()
+                detections[:, -1] == label).sum().item()
     return statistic_info
 
 
@@ -52,12 +53,12 @@ def draw_rects_and_labels(img, dets, colors, labels, thickness, font_size, font=
             text_size, _ = cv2.getTextSize(
                 labels[i], cv2.FONT_HERSHEY_COMPLEX_SMALL, font_size, 1)
             font_w, font_h = text_size
-            cv2.rectangle(img, (c1[0], max(0, int( c1[1] - 3 - 18 * font_size))),
+            cv2.rectangle(img, (c1[0], max(0, int(c1[1] - 3 - 18 * font_size))),
                           (c1[0] + font_w, max(c1[1], int(3 + 18 * font_size))), colors[cls], -1)
             cv2.putText(img,
                         labels[i],
                         (c1[0], max(c1[1] - 3, font_h)), cv2.FONT_HERSHEY_COMPLEX_SMALL, font_size,
-                        (0, 0, 0), 2)
+                        (0, 0, 0), 1)
     return img
 
 
@@ -79,14 +80,18 @@ def draw_single_img(img, detections, img_size,
             unique_labels = detections[:, -1].unique()
             statistic_info = _get_statistic_info(
                 detections, unique_labels, classes)
-        
+
         if only_rect:
             draw_rects(img, detections, colors, thickness)
         else:
             labels = []
             for detection in detections:
-                labels.append(classes[int(detection[-1])] +
-                              ' (' + str(round(detection[-3] * detection[-2] * 100, 2)) + '%)')
+                if len(detection) == 7:
+                    labels.append(classes[int(detection[-1])] +
+                                  ' (' + str(round(detection[-3] * detection[-2] * 100, 2)) + '%)')
+                else:
+                    labels.append(classes[int(detection[-1])] +
+                                  ' (' + str(round(detection[-2] * 100, 2)) + '%)')
             draw_rects_and_labels(img, detections, colors,
                                   labels, thickness, font_size, font)
 
@@ -168,10 +173,10 @@ def draw_labels_by_trackers(self, img, detections, only_rect):
             for detection in detections:
                 if self.id2label is not None and str(int(detection[4])) in self.id2label:
                     label = str(int(detection[4])) + ":" + \
-                        self.id2label[str(int(detection[4]))]
+                            self.id2label[str(int(detection[4]))]
                 else:
                     label = str(int(detection[4])) + ":" + \
-                        self.classes[int(detection[-1])]
+                            self.classes[int(detection[-1])]
                 labels.append(label)
 
             # 绘制所有标签

yolo3/utils/model_build.py

@@ -1,5 +1,8 @@
+import time
+
 import torch
 import numpy as np
+import torchvision
 import tqdm
 from torchvision.ops.boxes import batched_nms
 
@@ -46,6 +49,94 @@ def rescale_boxes(boxes, current_dim, original_shape):
     return boxes
 
 
+def soft_non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, merge=False, classes=None, agnostic=False,
+                             is_p1p2=False):
+    """Performs Non-Maximum Suppression (NMS) on inference results
+
+    Returns:
+         detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
+    """
+    if prediction.dtype is torch.float16:
+        prediction = prediction.float()  # to FP32
+
+    nc = prediction[0].shape[1] - 5  # number of classes
+    xc = prediction[..., 4] > conf_thres  # candidates
+
+    # Settings
+    min_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and height
+    max_det = 300  # maximum number of detections per image
+    time_limit = 10.0  # seconds to quit after
+    redundant = True  # require redundant detections
+    multi_label = nc > 1  # multiple labels per box (adds 0.5ms/img)
+
+    t = time.time()
+    output = [None] * prediction.shape[0]
+    for xi, x in enumerate(prediction):  # image index, image inference
+        # Apply constraints
+        # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0  # width-height
+        x = x[xc[xi]]  # confidence
+
+        # If none remain process next image
+        if not x.shape[0]:
+            continue
+
+        # Compute conf
+        x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf
+
+        # Box (center x, center y, width, height) to (x1, y1, x2, y2)
+        if not is_p1p2:
+            box = xywh2p1p2(x[:, :4])
+        else:
+            box = x[:, :4]
+
+        # Detections matrix nx6 (xyxy, conf, cls)
+        if multi_label:
+            i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
+            x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
+        else:  # best class only
+            conf, j = x[:, 5:].max(1, keepdim=True)
+            x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
+
+        # Filter by class
+        if classes:
+            x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
+
+        # Apply finite constraint
+        # if not torch.isfinite(x).all():
+        #     x = x[torch.isfinite(x).all(1)]
+
+        # If none remain process next image
+        n = x.shape[0]  # number of boxes
+        if not n:
+            continue
+
+        # Sort by confidence
+        # x = x[x[:, 4].argsort(descending=True)]
+
+        # Batched NMS
+        c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
+        boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
+        i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)
+        if i.shape[0] > max_det:  # limit detections
+            i = i[:max_det]
+        if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)
+            try:  # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
+                iou = bbox_iou(boxes[i], boxes) > iou_thres  # iou matrix
+                weights = iou * scores[None]  # box weights
+                x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
+                if redundant:
+                    i = i[iou.sum(1) > 1]  # require redundancy
+            except:  # possible CUDA error https://github.com/ultralytics/yolov3/issues/1139
+                print(x, i, x.shape, i.shape)
+                pass
+
+        output[xi] = x[i]
+        if (time.time() - t) > time_limit:
+            break  # time limit exceeded
+
+    return output
+
+
 def non_max_suppression(prediction, thres=0.5, nms_thres=0.4, is_p1p2=False):
     """
     Removes detections with lower object confidence score than 'conf_thres' and performs
@@ -73,7 +164,6 @@ def non_max_suppression(prediction, thres=0.5, nms_thres=0.4, is_p1p2=False):
         # If none anchor are remaining => process next image
         if not image_pred.size(0):
             continue
-        
 
         detections = torch.cat((image_pred[:, :5],
                                 class_confs.type(prediction.dtype),