write directly to video now

src/demo.py

@@ -5,36 +5,81 @@
 import _init_paths
 
 import logging
+import cv2
 import os
 import os.path as osp
 from opts import opts
 from tracking_utils.utils import mkdir_if_missing
 from tracking_utils.log import logger
 import datasets.dataset.jde as datasets
 from track import eval_seq
+import csv
 
 
 logger.setLevel(logging.INFO)
 
 
+def get_onlyseconds(match_csv):
+    lines = []
+    fp = open(match_csv, newline='')
+    reader = csv.reader(fp, delimiter=',')
+
+    first_line = reader.__next__()
+    field2idx = {}
+    for idx, field in enumerate(first_line):
+        field2idx[field] = idx
+
+    onlyseconds = {}
+    in_play = False
+    start_time = 0
+    last_time = 0
+
+    for idx, fields in enumerate(reader):
+        video_time = int(fields[field2idx['video_time']])
+        skill      = fields[field2idx['skill']]
+        evaluation = fields[field2idx['evaluation']]
+
+        if skill == 'Serve':
+            if in_play:
+                for i in range(start_time, last_time+2):
+                    onlyseconds[i] = 1
+            start_time = video_time
+            in_play = True
+
+        last_time = video_time
+
+    return onlyseconds
+
+
+
 def demo(opt):
+    if opt.match_csv is not None:
+        onlyseconds = get_onlyseconds(opt.match_csv)
+    else:
+        onlyseconds = None
+
     result_root = opt.output_root if opt.output_root != '' else '.'
     mkdir_if_missing(result_root)
 
+    print(f'img_size {opt.img_size}')
     logger.info('Starting tracking...')
     dataloader = datasets.LoadVideo(opt.input_video, opt.img_size)
     result_filename = os.path.join(result_root, 'results.txt')
     frame_rate = dataloader.frame_rate
 
-    frame_dir = None if opt.output_format == 'text' else osp.join(result_root, 'frame')
-    eval_seq(opt, dataloader, 'mot', result_filename,
-             save_dir=frame_dir, show_image=False, frame_rate=frame_rate,
-             use_cuda=opt.gpus!=[-1])
-
     if opt.output_format == 'video':
-        output_video_path = osp.join(result_root, 'MOT16-03-results.mp4')
-        cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg -b 5000k -c:v mpeg4 {}'.format(osp.join(result_root, 'frame'), output_video_path)
-        os.system(cmd_str)
+        basename = osp.basename(opt.input_video)
+        output_video_path = osp.join(result_root, basename)
+
+        vid_writer = cv2.VideoWriter(
+            output_video_path, cv2.VideoWriter_fourcc(*"mp4v"),
+            dataloader.frame_rate, (dataloader.w, dataloader.h)
+        )
+    else:
+        vid_writer = None
+
+    eval_seq(opt, dataloader, 'mot', result_filename, vid_writer=vid_writer,
+             frame_rate=frame_rate, use_cuda=opt.gpus!=[-1], onlyseconds=onlyseconds)
 
 
 if __name__ == '__main__':

src/demo.sh

@@ -0,0 +1,15 @@
+
+DATA_ROOT=/home/atao/data/vball
+VID=womens_clip_2plays
+MODEL=fairmot_dla34
+CONF=0.3
+
+LOGROOT=/home/atao/output/FairMOT/${VID}_${MODEL}_${CONF}
+
+OF=video  # set to text for no video
+
+MATCH_CSV="--match_csv ${DATA_ROOT}/${VID}_data.csv"
+MATCH_CSV=""
+MAX_SECONDS=""
+
+python demo.py mot --load_model ../models/${MODEL}.pth --conf_thres $CONF --input-video ${DATA_ROOT}/${VID}.mp4 --output-root ${LOGROOT} --output-format $OF $MATCH_CSV $MAX_SECONDS

src/dump_tracklets.py

@@ -29,21 +29,24 @@
     help="output root for image crops"
 )
 parser.add_argument(
-    "--crop_dims", default=(64,128),
+    "--crop_dims", default=(128,256),
     help="target resolution for player crops"
 )
 parser.add_argument(
     "--tracker_scale", default=1.5, type=float,
     help="how to scale tracker values down to original video resolution"
 )
 parser.add_argument(
-    "--sampling_rate", default=5, type=int,
+    "--sampling_rate", default=15, type=int,
     help="Only sample every N frames"
 )
 args = parser.parse_args()
 
 Court = namedtuple('court', ['top', 'bot', 'left', 'right'])
+# half court:
 court_bounds = Court(320, 680, 32, 1214)
+# full court:
+court_bounds = Court(193, 680, 32, 1214)
 entry_fields = ['fnum', 'tid', 'x', 'y', 'w', 'h']
 num_entries = len(entry_fields)
 TrackEntry = namedtuple('TrackEntry', entry_fields)

src/lib/datasets/dataset/jde.py

@@ -86,7 +86,8 @@ def __len__(self):
 class LoadVideo:  # for inference
     def __init__(self, path, img_size=(1088, 608)):
         self.cap = cv2.VideoCapture(path)
-        self.frame_rate = int(round(self.cap.get(cv2.CAP_PROP_FPS)))
+        # self.frame_rate = int(round(self.cap.get(cv2.CAP_PROP_FPS)))
+        self.frame_rate = self.cap.get(cv2.CAP_PROP_FPS)
         self.vw = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         self.vh = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         self.vn = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))

src/lib/opts.py

@@ -12,6 +12,9 @@ def __init__(self):
     # basic experiment setting
     self.parser.add_argument('task', default='mot', help='mot')
     self.parser.add_argument('--dataset', default='jde', help='jde')
+    self.parser.add_argument('--max_seconds', default=None, type=int,
+                             help='max seconds to process from video')
+    self.parser.add_argument('--match_csv', default=None, help='annotated match data')
     self.parser.add_argument('--exp_id', default='default')
     self.parser.add_argument('--test', action='store_true')
     #self.parser.add_argument('--load_model', default='../models/ctdet_coco_dla_2x.pth',

src/lib/tracker/multitracker.py

@@ -230,6 +230,9 @@ def merge_outputs(self, detections):
         return results
 
     def update(self, im_blob, img0):
+        """
+        im_blob
+        """
         self.frame_id += 1
         activated_starcks = []
         refind_stracks = []

src/track.py

@@ -5,7 +5,6 @@
 import _init_paths
 import os
 import os.path as osp
-import cv2
 import logging
 import argparse
 import motmetrics as mm
@@ -67,20 +66,32 @@ def write_results_score(filename, results, data_type):
     logger.info('save results to {}'.format(filename))
 
 
-def eval_seq(opt, dataloader, data_type, result_filename, save_dir=None, show_image=True, frame_rate=30, use_cuda=True):
-    if save_dir:
-        mkdir_if_missing(save_dir)
+def eval_seq(opt, dataloader, data_type, result_filename, vid_writer, frame_rate=30.038,
+             use_cuda=True, onlyseconds=None):
+
     tracker = JDETracker(opt, frame_rate=frame_rate)
     timer = Timer()
     results = []
+    img_id = 0
     frame_id = 0
-    #for path, img, img0 in dataloader:
-    for i, (path, img, img0) in enumerate(dataloader):
-        #if i % 8 != 0:
-            #continue
+    for frame_id, (path, img, img0) in enumerate(dataloader):
+        # only process times listed within onlyseconds
+        second = int(frame_id / frame_rate)
+        if onlyseconds is not None and second not in onlyseconds:
+            #if vid_writer is not None:
+            #    vid_writer.write(img0)
+            second = frame_id // int(frame_rate)
+            if frame_id == (second // 10) * 10 * int(frame_rate):
+                print(f'skipping @{second}s')
+            continue
+
         if frame_id % 20 == 0:
             logger.info('Processing frame {} ({:.2f} fps)'.format(frame_id, 1. / max(1e-5, timer.average_time)))
 
+        if opt.max_seconds is not None and second > opt.max_seconds:
+            print('Exiting!')
+            break
+
         # run tracking
         timer.tic()
         if use_cuda:
@@ -90,27 +101,25 @@ def eval_seq(opt, dataloader, data_type, result_filename, save_dir=None, show_im
         online_targets = tracker.update(blob, img0)
         online_tlwhs = []
         online_ids = []
-        #online_scores = []
         for t in online_targets:
             tlwh = t.tlwh
             tid = t.track_id
             vertical = tlwh[2] / tlwh[3] > 1.6
             if tlwh[2] * tlwh[3] > opt.min_box_area and not vertical:
                 online_tlwhs.append(tlwh)
                 online_ids.append(tid)
-                #online_scores.append(t.score)
         timer.toc()
+
         # save results
         results.append((frame_id + 1, online_tlwhs, online_ids))
-        #results.append((frame_id + 1, online_tlwhs, online_ids, online_scores))
-        if show_image or save_dir is not None:
+
+        if vid_writer is not None:
             online_im = vis.plot_tracking(img0, online_tlwhs, online_ids, frame_id=frame_id,
                                           fps=1. / timer.average_time)
-        if show_image:
-            cv2.imshow('online_im', online_im)
-        if save_dir is not None:
-            cv2.imwrite(os.path.join(save_dir, '{:05d}.jpg'.format(frame_id)), online_im)
-        frame_id += 1
+            vid_writer.write(online_im)
+
+        img_id += 1
+
     # save results
     write_results(result_filename, results, data_type)
     #write_results_score(result_filename, results, data_type)