Update inference.py

inference.py

@@ -64,11 +64,11 @@ def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None
     if config.associative_tracker == TrackerType.BYTETRACK:
 
         # Find low confidence detections
-        low_outputs = do_suppression(inference, conf_thres=config.byte_low_conf, iou_thres=config.nms_iou, gp=gp)
+        low_outputs = do_suppression(inference, image_meter_width, width, conf_thres=config.byte_low_conf, iou_thres=config.nms_iou, gp=gp)
         low_preds, real_width, real_height = format_predictions(image_shapes, low_outputs, width, height, gp=gp)
 
         # Find high confidence detections
-        high_outputs = do_suppression(inference, conf_thres=config.byte_high_conf, iou_thres=config.nms_iou, gp=gp)
+        high_outputs = do_suppression(inference, image_meter_width, width, conf_thres=config.byte_high_conf, iou_thres=config.nms_iou, gp=gp)
         high_preds, real_width, real_height = format_predictions(image_shapes, high_outputs, width, height, gp=gp)
 
         # Perform associative tracking (ByteTrack)
@@ -80,7 +80,7 @@ def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None
     else: 
 
         # Find confident detections
-        outputs = do_suppression(inference, conf_thres=config.conf_thresh, iou_thres=config.nms_iou, gp=gp)
+        outputs = do_suppression(inference, image_meter_width, width, conf_thres=config.conf_thresh, iou_thres=config.nms_iou, gp=gp)
 
         if config.associative_tracker == TrackerType.CONF_BOOST:
 
@@ -88,7 +88,7 @@ def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None
             do_confidence_boost(inference, outputs, boost_power=config.boost_power, boost_decay=config.boost_decay, gp=gp)
 
             # Find confident detections from boosted list
-            outputs = do_suppression(inference, conf_thres=config.conf_thresh, iou_thres=config.nms_iou, gp=gp)
+            outputs = do_suppression(inference, image_meter_width, width, conf_thres=config.conf_thresh, iou_thres=config.nms_iou, gp=gp)
 
         # Format confident detections
         all_preds, real_width, real_height = format_predictions(image_shapes, outputs, width, height, gp=gp)
@@ -169,7 +169,7 @@ def do_detection(dataloader, model, device, gp=None, batch_size=BATCH_SIZE, verb
 
     return inference, image_shapes, width, height
 
-def do_suppression(inference, gp=None, batch_size=BATCH_SIZE, conf_thres=CONF_THRES, iou_thres=NMS_IOU, verbose=True):
+def do_suppression(inference, image_meter_width, image_pixel_width, gp=None, batch_size=BATCH_SIZE, conf_thres=CONF_THRES, iou_thres=NMS_IOU, max_length=1.5, verbose=True):
     """
     Args:
         frames_dir: a directory containing frames to be evaluated
@@ -188,7 +188,7 @@ def do_suppression(inference, gp=None, batch_size=BATCH_SIZE, conf_thres=CONF_TH
             if gp: gp(batch_i / len(inference), pbar.__str__())
 
             with torch.no_grad():
-                output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres)
+                output = non_max_suppression(inf_out, image_meter_width, image_pixel_width, conf_thres=conf_thres, iou_thres=iou_thres, max_length=max_length)
 
 
             outputs.append(output)
@@ -422,13 +422,13 @@ def filter_detection_size(inference, image_meter_width, width, max_length):
         print(wc.shape)
         bs = batch.shape[0]  # batches
 
-        output = [torch.zeros((0, 6), device=batch.device)] * bs
+        output = torch.zeros((bs, 0, 6), device=batch.device)
         print("wc")
         print(batch.shape)
         for xi, x in enumerate(batch):
             x = x[wc[xi]]  # confidence
             print(x.shape)
-            output[xi] = x
+            output[xi, :, :] = x
 
         output = torch.tensor(output)
         print("output len", output.shape)
@@ -439,6 +439,9 @@ def filter_detection_size(inference, image_meter_width, width, max_length):
 
 def non_max_suppression(
         prediction,
+        image_meter_width,
+        image_pixel_width,
+        max_length=1.5,
         conf_thres=0.25,
         iou_thres=0.45,
         max_det=300
@@ -463,6 +466,9 @@ def non_max_suppression(
         prediction = prediction.cpu()
     bs = prediction.shape[0]  # batch size
     xc = prediction[..., 4] > conf_thres  # candidates
+    pix2width = image_meter_width/width
+    width = prediction[..., 2]*pix2width
+    wc = width < max_length
 
 
     # Settings
@@ -476,7 +482,7 @@ def non_max_suppression(
 
             
         # Keep boxes that pass confidence threshold
-        x = x[xc[xi]]  # confidence
+        x = x[xc[xi] * wc[xi]]  # confidence
 
         # If none remain process next image
         if not x.shape[0]: