Update inference.py

inference.py

@@ -155,11 +155,6 @@ def do_detection(dataloader, model, device, gp=None, batch_size=BATCH_SIZE, verb
                 inf_out, _ = model(img, augment=False) 
 
 
-            32, 30240, 6
-            print(inf_out.shape)
-            print(inf_out[1, 1, :])
-            w = inf_out[:, :, 2] - inf_out[:, :, 0]
-            print(w.shape)
 
             # Save shapes for resizing to original shape
             batch_shape = []
@@ -399,9 +394,11 @@ def json_dump_round_float(some_object, out_path, num_digits=4):
     
 def non_max_suppression(
         prediction,
+        pix2w,
         conf_thres=0.25,
         iou_thres=0.45,
         max_det=300,
+        max_length=1.5
 ):
     """Non-Maximum Suppression (NMS) on inference results to reject overlapping detections
 
@@ -423,6 +420,12 @@ def non_max_suppression(
         prediction = prediction.cpu()
     bs = prediction.shape[0]  # batch size
     xc = prediction[..., 4] > conf_thres  # candidates
+    print(xc.shape)
+    wc = (prediction[..., 2] - prediction[..., 0])*pix2w < max_length
+    print(wc.shape)
+    xc = xc and wc
+    print(xc.shape)
+
 
     # Settings
     # min_wh = 2  # (pixels) minimum box width and height