Fixed the video doohickey and made --score_threshold faster.

eval.py

@@ -102,7 +102,7 @@ def parse_args(argv=None):
                         help='A path to a video to evaluate on.')
     parser.add_argument('--video_multiframe', default=1, type=int,
                         help='The number of frames to evaluate in parallel to make videos play at higher fps.')
-    parser.add_argument('--score_threshold', default=0, type=float ,
+    parser.add_argument('--score_threshold', default=0, type=float,
                         help='Detections with a score under this threshold will not be considered. This currently only works in display mode.')
     parser.add_argument('--dataset', default=None, type=str,
                         help='If specified, override the dataset specified in the config with this one (example: coco2017_dataset).')
@@ -136,7 +136,7 @@ def prep_display(dets_out, img, gt, gt_masks, h, w, undo_transform=True, class_c
         h, w, _ = img.shape
 
     with timer.env('Postprocess'):
-        t = postprocess(dets_out, w, h, visualize_lincomb=args.display_lincomb, crop_masks=args.crop)
+        t = postprocess(dets_out, w, h, visualize_lincomb=args.display_lincomb, crop_masks=args.crop, score_threshold=args.score_threshold)
         torch.cuda.synchronize()
 
     with timer.env('Copy'):
@@ -202,7 +202,7 @@ def get_color(j):
 
 def prep_benchmark(dets_out, h, w):
     with timer.env('Postprocess'):
-        t = postprocess(dets_out, w, h, crop_masks=args.crop)
+        t = postprocess(dets_out, w, h, crop_masks=args.crop, score_threshold=args.score_threshold)
 
     with timer.env('Copy'):
         classes, scores, boxes, masks = [x[:args.top_k].cpu().numpy() for x in t]
@@ -351,7 +351,7 @@ def prep_metrics(ap_data, dets, img, gt, gt_masks, h, w, num_crowd, image_id, de
                 crowd_classes, gt_classes = split(gt_classes)
 
     with timer.env('Postprocess'):
-        classes, scores, boxes, masks = postprocess(dets, w, h, crop_masks=args.crop)
+        classes, scores, boxes, masks = postprocess(dets, w, h, crop_masks=args.crop, score_threshold=args.score_threshold)
 
         if classes.size(0) == 0:
             return
@@ -596,7 +596,7 @@ def prep_frame(inp):
             frame, preds = inp
             return prep_display(preds, frame, None, None, None, None, undo_transform=False, class_color=True)
 
-    extract_frame = lambda x, i: (x[0][i], {k: v[i].unsqueeze(0) for k, v in x[1].items()})
+    extract_frame = lambda x, i: (x[0][i], [x[1][i]])
 
     # Prime the network on the first frame because I do some thread unsafe things otherwise
     print('Initializing model... ', end='')

layers/output_utils.py

@@ -13,7 +13,7 @@
 from .box_utils import crop, sanitize_coordinates
 
 def postprocess(det_output, w, h, batch_idx=0, interpolation_mode='bilinear',
-                visualize_lincomb=False, crop_masks=True):
+                visualize_lincomb=False, crop_masks=True, score_threshold=0):
     """
     Postprocesses the output of Yolact on testing mode into a format that makes sense,
     accounting for all the possible configuration settings.
@@ -37,6 +37,16 @@ def postprocess(det_output, w, h, batch_idx=0, interpolation_mode='bilinear',
     if dets is None:
         return [torch.Tensor()] * 4 # Warning, this is 4 copies of the same thing
 
+    if score_threshold > 0:
+        keep = dets['score'] > score_threshold
+
+        for k in dets:
+            if k != 'proto':
+                dets[k] = dets[k][keep]
+
+        if dets['score'].size(0) == 0:
+            return [torch.Tensor()] * 4
+
     # im_w and im_h when it concerns bboxes. This is a workaround hack for preserve_aspect_ratio
     b_w, b_h = (w, h)