Update inference.py

inference.py

@@ -59,14 +59,16 @@ def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None
     # Detect boxes in frames
     inference, image_shapes, width, height = do_detection(dataloader, model, device, gp=gp)
 
+    inference = filter_detection_size(inference, image_meter_width, width, 1.5)
+
     if config.associative_tracker == TrackerType.BYTETRACK:
 
         # Find low confidence detections
-        low_outputs = do_suppression(inference, image_meter_width/width, conf_thres=config.byte_low_conf, iou_thres=config.nms_iou, gp=gp)
+        low_outputs = do_suppression(inference, 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, image_meter_width/width, conf_thres=config.byte_high_conf, iou_thres=config.nms_iou, gp=gp)
+        high_outputs = do_suppression(inference, 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)
@@ -167,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, pix2w, gp=None, batch_size=BATCH_SIZE, conf_thres=CONF_THRES, iou_thres=NMS_IOU, verbose=True):
+def do_suppression(inference, gp=None, batch_size=BATCH_SIZE, conf_thres=CONF_THRES, iou_thres=NMS_IOU, verbose=True):
     """
     Args:
         frames_dir: a directory containing frames to be evaluated
@@ -186,7 +188,7 @@ def do_suppression(inference, pix2w, gp=None, batch_size=BATCH_SIZE, conf_thres=
             if gp: gp(batch_i / len(inference), pbar.__str__())
 
             with torch.no_grad():
-                output = non_max_suppression(inf_out, pix2w, conf_thres=conf_thres, iou_thres=iou_thres)
+                output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres)
 
 
             outputs.append(output)
@@ -392,13 +394,32 @@ def json_dump_round_float(some_object, out_path, num_digits=4):
     with patch('json.encoder._make_iterencode', wraps=inner):
         return json.dump(some_object, open(out_path, 'w'), indent=2)
     
+
+def filter_detection_size(inference, image_meter_width, width, max_length):
+
+    outputs = []
+    for prediction in inference:
+
+        pix2width = image_meter_width/width
+        print("pix2w", pix2width)
+        width = prediction[..., 2]*pix2width
+        wc = width < max_length
+        bs = prediction.shape[0]  # batches
+
+        output = [torch.zeros((0, 6), device=prediction.device)] * bs
+        for xi, x in enumerate(prediction):
+            x = x[wc[xi]]  # confidence
+            output[xi] = x
+
+        outputs.append(output)
+         
+    return outputs
+
 def non_max_suppression(
         prediction,
-        pix2w,
         conf_thres=0.25,
         iou_thres=0.45,
-        max_det=300,
-        max_length=0.8
+        max_det=300
 ):
     """Non-Maximum Suppression (NMS) on inference results to reject overlapping detections
 
@@ -420,19 +441,6 @@ def non_max_suppression(
         prediction = prediction.cpu()
     bs = prediction.shape[0]  # batch size
     xc = prediction[..., 4] > conf_thres  # candidates
-    print("pix2w", pix2w)
-    width = prediction[..., 2]*pix2w
-    wc = width < max_length
-
-    print(0.1, sum(sum(width < 0.1)))
-    print(0.8, sum(sum(width < 0.8)))
-    print(1.5, sum(sum(width < 1.5)))
-    print(10,  sum(sum(width < 10 )))
-    print(10,  sum(sum(width < 10 )))
-    print(prediction[1,512,:])
-    print("mean", torch.mean(width))
-
-
 
 
     # Settings
@@ -446,7 +454,7 @@ def non_max_suppression(
 
             
         # Keep boxes that pass confidence threshold
-        x = x[xc[xi] * wc[xi]]  # confidence
+        x = x[xc[xi]]  # confidence
 
         # If none remain process next image
         if not x.shape[0]: