From a1e373b9208172346b3ed5cc4ef68f4ba83dfa8d Mon Sep 17 00:00:00 2001 From: nemonameless Date: Wed, 28 Jul 2021 08:12:24 +0000 Subject: [PATCH] fix mot infer video --- deploy/python/mot_jde_infer.py | 21 +++++++++++++++------ deploy/python/mot_keypoint_unite_infer.py | 17 +++++++++++++---- deploy/python/mot_sde_infer.py | 21 +++++++++++++++------ 3 files changed, 43 insertions(+), 16 deletions(-) diff --git a/deploy/python/mot_jde_infer.py b/deploy/python/mot_jde_infer.py index 021e5aea707..2c2b901b644 100644 --- a/deploy/python/mot_jde_infer.py +++ b/deploy/python/mot_jde_infer.py @@ -220,7 +220,8 @@ def predict_video(detector, camera_id): if not os.path.exists(FLAGS.output_dir): os.makedirs(FLAGS.output_dir) out_path = os.path.join(FLAGS.output_dir, video_name) - writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height)) + if not FLAGS.save_images: + writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height)) frame_id = 0 timer = MOTTimer() results = [] @@ -235,7 +236,7 @@ def predict_video(detector, camera_id): results.append((frame_id + 1, online_tlwhs, online_scores, online_ids)) fps = 1. / timer.average_time - online_im = mot_vis.plot_tracking( + im = mot_vis.plot_tracking( frame, online_tlwhs, online_ids, @@ -248,11 +249,11 @@ def predict_video(detector, camera_id): os.makedirs(save_dir) cv2.imwrite( os.path.join(save_dir, '{:05d}.jpg'.format(frame_id)), - online_im) + im) + else: + writer.write(im) frame_id += 1 print('detect frame:%d' % (frame_id)) - im = np.array(online_im) - writer.write(im) if camera_id != -1: cv2.imshow('Tracking Detection', im) if cv2.waitKey(1) & 0xFF == ord('q'): @@ -261,7 +262,15 @@ def predict_video(detector, camera_id): result_filename = os.path.join(FLAGS.output_dir, video_name.split('.')[-2] + '.txt') write_mot_results(result_filename, results) - writer.release() + + if FLAGS.save_images: + save_dir = os.path.join(FLAGS.output_dir, video_name.split('.')[-2]) + cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg -vf "scale=trunc(iw/2)*2:trunc(ih/2)*2" {}'.format( + save_dir, out_path) + os.system(cmd_str) + print('Save video in {}.'.format(out_path)) + else: + writer.release() def main(): diff --git a/deploy/python/mot_keypoint_unite_infer.py b/deploy/python/mot_keypoint_unite_infer.py index 46ecefe78cb..743e4566954 100644 --- a/deploy/python/mot_keypoint_unite_infer.py +++ b/deploy/python/mot_keypoint_unite_infer.py @@ -136,7 +136,8 @@ def mot_keypoint_unite_predict_video(mot_model, if not os.path.exists(FLAGS.output_dir): os.makedirs(FLAGS.output_dir) out_path = os.path.join(FLAGS.output_dir, video_name) - writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height)) + if not FLAGS.save_images: + writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height)) frame_id = 0 timer_mot = FPSTimer() timer_kp = FPSTimer() @@ -201,8 +202,8 @@ def mot_keypoint_unite_predict_video(mot_model, os.makedirs(save_dir) cv2.imwrite( os.path.join(save_dir, '{:05d}.jpg'.format(frame_id)), im) - - writer.write(im) + else: + writer.write(im) if camera_id != -1: cv2.imshow('Tracking and keypoint results', im) if cv2.waitKey(1) & 0xFF == ord('q'): @@ -211,7 +212,15 @@ def mot_keypoint_unite_predict_video(mot_model, result_filename = os.path.join(FLAGS.output_dir, video_name.split('.')[-2] + '.txt') write_mot_results(result_filename, mot_results) - writer.release() + + if FLAGS.save_images: + save_dir = os.path.join(FLAGS.output_dir, video_name.split('.')[-2]) + cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg -vf "scale=trunc(iw/2)*2:trunc(ih/2)*2" {}'.format( + save_dir, out_path) + os.system(cmd_str) + print('Save video in {}.'.format(out_path)) + else: + writer.release() def main(): diff --git a/deploy/python/mot_sde_infer.py b/deploy/python/mot_sde_infer.py index 55dbef8c9eb..76c00bd1068 100644 --- a/deploy/python/mot_sde_infer.py +++ b/deploy/python/mot_sde_infer.py @@ -355,7 +355,8 @@ def predict_video(detector, reid_model, camera_id): if not os.path.exists(FLAGS.output_dir): os.makedirs(FLAGS.output_dir) out_path = os.path.join(FLAGS.output_dir, video_name) - writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height)) + if not FLAGS.save_images: + writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height)) frame_id = 0 timer = MOTTimer() results = [] @@ -378,7 +379,7 @@ def predict_video(detector, reid_model, camera_id): results.append((frame_id + 1, online_tlwhs, online_scores, online_ids)) fps = 1. / timer.average_time - online_im = mot_vis.plot_tracking( + im = mot_vis.plot_tracking( frame, online_tlwhs, online_ids, @@ -391,11 +392,11 @@ def predict_video(detector, reid_model, camera_id): os.makedirs(save_dir) cv2.imwrite( os.path.join(save_dir, '{:05d}.jpg'.format(frame_id)), - online_im) + im) + else: + writer.write(im) frame_id += 1 print('detect frame:%d' % (frame_id)) - im = np.array(online_im) - writer.write(im) if camera_id != -1: cv2.imshow('Tracking Detection', im) if cv2.waitKey(1) & 0xFF == ord('q'): @@ -404,7 +405,15 @@ def predict_video(detector, reid_model, camera_id): result_filename = os.path.join(FLAGS.output_dir, video_name.split('.')[-2] + '.txt') write_mot_results(result_filename, results) - writer.release() + + if FLAGS.save_images: + save_dir = os.path.join(FLAGS.output_dir, video_name.split('.')[-2]) + cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg -vf "scale=trunc(iw/2)*2:trunc(ih/2)*2" {}'.format( + save_dir, out_path) + os.system(cmd_str) + print('Save video in {}.'.format(out_path)) + else: + writer.release() def main():