security_vision/Towards-Realtime-MOT
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

save to temp folder, in higher quality

Browse source
This commit is contained in:
Ruben van de Ven 2023-04-24 18:20:22 +02:00
parent 48a31a9bd9
commit 7c3433e456
2 changed files with 26 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
demo.py
View file

@ -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__':

5
track.py
View file

@ -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