未验证 提交 e4259d42 编写于 作者: G George Ni 提交者: GitHub

fix mot infer video (#3823)

上级 a3e2b2ea
......@@ -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()
......
......@@ -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()
......
......@@ -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()
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册