fix mot infer video (#3823)

deploy/python/mot_jde_infer.py

@@ -221,6 +221,7 @@ 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)
+    if not FLAGS.save_images:
         writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
     frame_id = 0
     timer = MOTTimer()
@@ -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,6 +263,14 @@ 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)
+
+    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()
 
 

deploy/python/mot_keypoint_unite_infer.py

@@ -137,6 +137,7 @@ 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)
+    if not FLAGS.save_images:
         writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
     frame_id = 0
     timer_mot = FPSTimer()
@@ -202,7 +203,7 @@ 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)
-
+        else:
             writer.write(im)
         if camera_id != -1:
             cv2.imshow('Tracking and keypoint results', im)
@@ -212,6 +213,14 @@ 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)
+
+    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()
 
 

deploy/python/mot_sde_infer.py

@@ -356,6 +356,7 @@ 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)
+    if not FLAGS.save_images:
         writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
     frame_id = 0
     timer = MOTTimer()
@@ -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,6 +406,14 @@ 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)
+
+    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()