未验证 提交 c44543ae 编写于 作者: F FlyingQianMM 提交者: GitHub

add python flow_statistic for sde (#4542)

* add flow_statistic for mot_jde

* raise NotImplementedError when multi-class entrance counting is chosen

* revert * 'mp4v' to *'mp4v'

* add python flow_statistic for sde
上级 8e13eedb
......@@ -220,7 +220,7 @@ def flow_statistic(result,
out_id_list,
prev_center,
records,
data_type,
data_type='mot',
num_classes=1):
# Count in and out number:
# Use horizontal center line as the entrance just for simplification.
......
......@@ -31,7 +31,7 @@ from benchmark_utils import PaddleInferBenchmark
from visualize import plot_tracking
from mot.tracker import DeepSORTTracker
from mot.utils import MOTTimer, write_mot_results
from mot.utils import MOTTimer, write_mot_results, flow_statistic
# Global dictionary
MOT_SUPPORT_MODELS = {'DeepSORT'}
......@@ -545,6 +545,15 @@ def predict_video(detector, reid_model, camera_id):
frame_id = 0
timer = MOTTimer()
results = defaultdict(list)
id_set = set()
interval_id_set = set()
in_id_list = list()
out_id_list = list()
prev_center = dict()
records = list()
entrance = [0, height / 2., width, height / 2.]
video_fps = fps
while (1):
ret, frame = capture.read()
if not ret:
......@@ -565,6 +574,19 @@ def predict_video(detector, reid_model, camera_id):
crops, pred_dets)
results[0].append(
(frame_id + 1, online_tlwhs, online_scores, online_ids))
# NOTE: just implement flow statistic for one class
result = (frame_id + 1, online_tlwhs, online_scores, online_ids)
statistic = flow_statistic(
result, FLAGS.secs_interval, FLAGS.do_entrance_counting,
video_fps, entrance, id_set, interval_id_set, in_id_list,
out_id_list, prev_center, records)
id_set = statistic['id_set']
interval_id_set = statistic['interval_id_set']
in_id_list = statistic['in_id_list']
out_id_list = statistic['out_id_list']
prev_center = statistic['prev_center']
records = statistic['records']
timer.toc()
fps = 1. / timer.duration
......@@ -598,6 +620,14 @@ def predict_video(detector, reid_model, camera_id):
video_name.split('.')[-2] + '.txt')
write_mot_results(result_filename, results)
result_filename = os.path.join(
FLAGS.output_dir, video_name.split('.')[-2] + '_flow_statistic.txt')
f = open(result_filename, 'w')
for line in records:
f.write(line)
print('Flow statistic save in {}'.format(result_filename))
f.close()
if FLAGS.save_images:
save_dir = os.path.join(FLAGS.output_dir, video_name.split('.')[-2])
cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg {}'.format(save_dir,
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册