From e4259d42c16c6ecff869da07a85cfca4edbbec3d Mon Sep 17 00:00:00 2001 From: George Ni Date: Thu, 29 Jul 2021 19:14:00 +0800 Subject: [PATCH] fix mot infer video (#3823) --- 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 103cf4585..d35a39629 100644 --- a/deploy/python/mot_jde_infer.py +++ b/deploy/python/mot_jde_infer.py @@ -221,7 +221,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 = [] @@ -236,7 +237,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, @@ -249,11 +250,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'): @@ -262,7 +263,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 2db100c2d..706b61127 100644 --- a/deploy/python/mot_keypoint_unite_infer.py +++ b/deploy/python/mot_keypoint_unite_infer.py @@ -137,7 +137,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() @@ -202,8 +203,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'): @@ -212,7 +213,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 5dfb944c6..60a3fd7db 100644 --- a/deploy/python/mot_sde_infer.py +++ b/deploy/python/mot_sde_infer.py @@ -356,7 +356,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 = [] @@ -379,7 +380,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, @@ -392,11 +393,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'): @@ -405,7 +406,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(): -- GitLab