Update inference.py

inference.py

@@ -330,45 +330,36 @@ def do_tracking(all_preds, image_meter_width, image_meter_height, gp=None, max_a
 
     return json_data
 
-def do_associative_tracking(raw_detections, image_meter_width, image_meter_height, gp=None, conf_thresh=0.2, max_age=MAX_AGE, iou_thres=IOU_THRES, min_hits=MIN_HITS, min_length=MIN_LENGTH, verbose=True):
+def do_associative_tracking(inference, image_shapes, width, height, image_meter_width, image_meter_height, gp=None, low_thresh=0.001, high_threshold=0.2, max_age=MAX_AGE, iou_thres=IOU_THRES, min_hits=MIN_HITS, min_length=MIN_LENGTH, batch_size=BATCH_SIZE, verbose=True):
 
     if (gp): gp(0, "Tracking...")
 
     print("Preprocessing")
 
-    low_dets = []
-    high_dets = []
-    with tqdm(total=len(raw_detections), desc="Running tracking", ncols=0, disable=not verbose) as pbar:
-        for batch in raw_detections:
-            for frame in batch:
-                low_frame = []
-                high_frame = []
-                for bbox in frame:
-                    if bbox[4] > conf_thresh:
-                        high_frame.append(bbox)
-                    else:
-                        low_frame.append(bbox)
-                low_dets.append(low_frame)
-                high_dets.append(high_frame)
-            pbar.update(1)
+
+    low_outputs = do_suppression(inference, conf_thres=low_thresh, iou_thres=iou_thres, gp=gp)
+    low_preds, real_width, real_height = format_predictions(image_shapes, low_outputs, width, height, verbose=verbose)
+
+    high_outputs = do_suppression(inference, conf_thres=high_threshold, iou_thres=iou_thres, gp=gp)
+    high_preds, real_width, real_height = format_predictions(image_shapes, high_outputs, width, height, verbose=verbose)
 
     print("Preprocess done")
 
     # Initialize tracker
     clip_info = {
         'start_frame': 0,
-        'end_frame': len(raw_detections),
+        'end_frame': len(low_preds),
         'image_meter_width': image_meter_width,
         'image_meter_height': image_meter_height
     }
     tracker = Tracker(clip_info, algorithm=Associate, args={ 'max_age': max_age, 'min_hits': 0, 'iou_threshold': iou_thres}, min_hits=min_hits)
     
     # Run tracking
-    with tqdm(total=len(low_dets), desc="Running tracking", ncols=0, disable=not verbose) as pbar:
-        for i in range(len(low_dets)):
-            if gp: gp(i / len(low_dets), pbar.__str__())
-            low_boxes = low_dets[i]
-            high_boxes = high_dets[i]
+    with tqdm(total=len(low_preds), desc="Running tracking", ncols=0, disable=not verbose) as pbar:
+        for i in range(len(low_preds)):
+            if gp: gp(i / len(low_preds), pbar.__str__())
+            low_boxes = low_preds[i]
+            high_boxes = high_preds[i]
             boxes = (low_boxes, high_boxes)
             if len(low_boxes) + len(high_boxes) > 0:
                 tracker.update(boxes)
@@ -455,6 +446,90 @@ def non_max_suppression(
         x = torch.cat((box, conf, j.float(), mask), 1)[conf.view(-1) > conf_thres]
 
 
+        # Check shape
+        n = x.shape[0]  # number of boxes
+        if not n:  # no boxes
+            continue
+        x = x[x[:, 4].argsort(descending=True)[:max_nms]]  # sort by confidence and remove excess boxes
+
+        # Batched NMS
+        boxes  = x[:, :4]  # boxes (offset by class), scores
+        scores = x[:, 4]
+        i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
+
+        i = i[:max_det]  # limit detections
+        if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)
+            # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
+            iou = box_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
+
+        output[xi] = x[i]
+        if mps:
+            output[xi] = output[xi].to(device)
+        
+        logging = False
+
+    return output
+
+
+def no_suppression(
+        prediction,
+        conf_thres=0.25,
+        iou_thres=0.45,
+        max_det=300,
+):
+    """Non-Maximum Suppression (NMS) on inference results to reject overlapping detections
+
+    Returns:
+         list of detections, on (n,6) tensor per image [xyxy, conf, cls]
+    """
+
+    # Checks
+    assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
+    assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'
+    if isinstance(prediction, (list, tuple)):  # YOLOv5 model in validation model, output = (inference_out, loss_out)
+        prediction = prediction[0]  # select only inference output
+
+    device = prediction.device
+    mps = 'mps' in device.type  # Apple MPS
+    if mps:  # MPS not fully supported yet, convert tensors to CPU before NMS
+        prediction = prediction.cpu()
+    bs = prediction.shape[0]  # batch size
+    xc = prediction[..., 4] > conf_thres  # candidates
+
+    # Settings
+    # min_wh = 2  # (pixels) minimum box width and height
+    max_nms = 30000  # maximum number of boxes into torchvision.ops.nms()
+    redundant = True  # require redundant detections
+    merge = False  # use merge-NMS
+
+    output = [torch.zeros((0, 6), device=prediction.device)] * bs
+    for xi, x in enumerate(prediction):  # image index, image inference
+
+            
+        # Keep boxes that pass confidence threshold
+        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/Mask
+        box = xywh2xyxy(x[:, :4])  # center_x, center_y, width, height) to (x1, y1, x2, y2)
+        mask = x[:, 6:]  # zero columns if no masks
+
+        # Detections matrix nx6 (xyxy, conf, cls)
+        conf, j = x[:, 5:6].max(1, keepdim=True)
+        x = torch.cat((box, conf, j.float(), mask), 1)[conf.view(-1) > conf_thres]
+
+
         # Check shape
         n = x.shape[0]  # number of boxes
         if not n:  # no boxes