#!/usr/bin/env python

import argparse
import os
import sys
import cv2
from vdetlib.vdet.dataset import imagenet_vdet_class_idx
from vdetlib.utils.common import imread, imwrite
from vdetlib.utils.protocol import proto_load, proto_dump, frame_path_at, frame_top_detections
from vdetlib.utils.visual import add_bbox
from vdetlib.utils.cython_nms import nms
import scipy.io as sio
import numpy as np

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('vid_file')
    parser.add_argument('det_root')
    parser.add_argument('--cls', choices=imagenet_vdet_class_idx.keys())
    parser.add_argument('--save_dir', default=None)
    parser.add_argument('--top_k', default=10)
    args = parser.parse_args()

    vid_proto = proto_load(args.vid_file)
    if args.save_dir and not os.path.isdir(args.save_dir):
        os.mkdir(args.save_dir)
    cls_index = imagenet_vdet_class_idx[args.cls]

    for frame in vid_proto['frames']:
        det_file = os.path.join(args.det_root,
            "{}.mat".format(os.path.splitext(frame['path'])[0]))
        det = sio.loadmat(det_file)
        frame_idx = frame['frame']
        img = imread(frame_path_at(vid_proto, frame_idx))
        boxes = det['boxes'][:,cls_index,:]
        scores = det['zs'][:,cls_index]
        keep = nms(np.hstack((boxes,scores[:,np.newaxis])), 0.3)
        det_img = add_bbox(img, [boxes[i,:] for i in keep], [scores[i] for i in keep])
        cv2.imshow('detection', det_img)
        if cv2.waitKey(0) == ord('q'):
            cv2.destroyAllWindows()
            sys.exit(0)
        cv2.destroyAllWindows()
        if args.save_dir:
            imwrite(os.path.join(args.save_dir, "{:04d}.jpg".format(frame_idx)), det_img)


@@ -699,18 +699,132 @@
 fault=10
+, type=int)%0A    parser.add_argument('--no_nms', action='store_false', dest='nms')%0A    parser.set_defaults(nms=True
 )%0A
-
     args
@@ -1353,16 +1353,33 @@
 index,:%5D
+.astype('single')
 %0A       
@@ -1410,17 +1410,59 @@
 s_index%5D
-%0A
+.astype('single')%0A        if args.nms:%0A    
         
@@ -1530,70 +1530,361 @@
     
-det_img = add_bbox(img, %5Bboxes%5Bi,:%5D for i in keep%5D, %5Bscores
+else:%0A            keep = range(len(boxes))%0A        kept_boxes = %5Bboxes%5Bi,:%5D for i in keep%5D%0A        kept_scores = %5Bscores%5Bi%5D for i in keep%5D%0A        top_idx = np.argsort(np.asarray(kept_scores))%5B::-1%5D%0A        top_boxes = %5Bkept_boxes%5Btop_idx%5Bi%5D%5D for i in %5C%0A            xrange(min(args.top_k, len(kept_boxes)))%5D%0A        top_scores = %5Bkept_scores%5Btop_idx
 %5Bi%5D
+%5D
  for
@@ -1889,21 +1889,125 @@
 or i in 
-keep%5D
+%5C%0A            xrange(min(args.top_k, len(kept_boxes)))%5D%0A        det_img = add_bbox(img, top_boxes, top_scores
 )%0A      

Add option to disable NMS when showing video detection results. Fix bugs in show_video_matfile_detection.py