diff --git a/demo.py b/demo.py index 3cb7e75..ac6fe35 100644 --- a/demo.py +++ b/demo.py @@ -31,6 +31,8 @@ from utils.log import logger from utils.timer import Timer from utils.parse_config import parse_model_cfg import utils.datasets as datasets +from pathlib import Path +import tempfile from track import eval_seq @@ -49,26 +51,29 @@ def track(opt): n_frame = 0 logger.info('Starting tracking...') - if os.path.isdir(opt.input_video): - print('Use image sequence') - dataloader = datasets.LoadImages(opt.input_video, opt.img_size) - frame_rate = 30 # hack for now; see https://motchallenge.net/data/MOT16/ - else: - dataloader = datasets.LoadVideo(opt.input_video, opt.img_size) - frame_rate = dataloader.frame_rate - result_filename = os.path.join(result_root, 'results.txt') + + with tempfile.TemporaryDirectory() as tmpdirname: + if os.path.isdir(opt.input_video): + print('Use image sequence') + dataloader = datasets.LoadImages(opt.input_video, opt.img_size) + frame_rate = 30 # hack for now; see https://motchallenge.net/data/MOT16/ + else: + dataloader = datasets.LoadVideo(opt.input_video, opt.img_size) + frame_rate = dataloader.frame_rate + result_filename = os.path.join(result_root, 'results.txt') - frame_dir = None if opt.output_format=='text' else osp.join(result_root, 'frame') - try: - eval_seq(opt, dataloader, 'mot', result_filename, - save_dir=frame_dir, show_image=False, frame_rate=frame_rate) - except Exception as e: - logger.info(e) + frame_dir = None if opt.output_format=='text' else tmpdirname + try: + eval_seq(opt, dataloader, 'mot', result_filename, + save_dir=frame_dir, show_image=False, save_img=True, frame_rate=frame_rate) + except Exception as e: + logger.info(e) - if opt.output_format == 'video': - output_video_path = osp.join(result_root, 'result.mp4') - cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg -c:v copy {}'.format(osp.join(result_root, 'frame'), output_video_path) - os.system(cmd_str) + if opt.output_format == 'video': + name = 'result-' + Path(opt.input_video).stem + '.mp4' + output_video_path = osp.join(result_root, name) + cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg -c:v mjpeg -q:v 1 {}'.format(frame_dir, output_video_path) + os.system(cmd_str) if __name__ == '__main__': diff --git a/track.py b/track.py index a7804c7..3f635aa 100644 --- a/track.py +++ b/track.py @@ -100,7 +100,8 @@ def eval_seq(opt, dataloader, data_type, result_filename, save_dir=None, save_im online_targets, frame_embeddings = tracker.update(blob, img0) online_tlwhs = [] online_ids = [] - for t in online_targets: + # for t in online_targets: + for t in tracker.tracked_stracks: tlwh = t.tlwh tid = t.track_id vertical = tlwh[2] / tlwh[3] > 1.6 @@ -118,7 +119,7 @@ def eval_seq(opt, dataloader, data_type, result_filename, save_dir=None, save_im if save_dir is not None: base_fn = os.path.join(save_dir, '{:05d}'.format(frame_id)) if save_img: - cv2.imwrite(base_fn+'.jpg', online_im) + cv2.imwrite(base_fn+'.jpg', online_im, [cv2.IMWRITE_JPEG_QUALITY, 100]) if save_figures: for i, fe in enumerate(frame_embeddings): tlwh, curr_feat = fe