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

[MOT] visualize entrance line (#4554)

* add implementment for FlowStatics

* rename count to entrance_count; rename count_set to id_set

* add a todo for entrance line setting

* rename count to do_entrance_counting

* convert print_save_file to the origin line

* visualize entrance line

* fix conflicts

* move entrance_line show in plot_track
上级 c902676e
...@@ -272,7 +272,10 @@ def predict_video(detector, camera_id): ...@@ -272,7 +272,10 @@ def predict_video(detector, camera_id):
online_scores, online_scores,
frame_id=frame_id, frame_id=frame_id,
fps=fps, fps=fps,
ids2names=ids2names) ids2names=ids2names,
do_entrance_counting=FLAGS.do_entrance_counting,
entrance=entrance)
if FLAGS.save_images: if FLAGS.save_images:
save_dir = os.path.join(FLAGS.output_dir, video_name.split('.')[-2]) save_dir = os.path.join(FLAGS.output_dir, video_name.split('.')[-2])
if not os.path.exists(save_dir): if not os.path.exists(save_dir):
......
...@@ -596,7 +596,9 @@ def predict_video(detector, reid_model, camera_id): ...@@ -596,7 +596,9 @@ def predict_video(detector, reid_model, camera_id):
online_ids, online_ids,
online_scores, online_scores,
frame_id=frame_id, frame_id=frame_id,
fps=fps) fps=fps,
do_entrance_counting=FLAGS.do_entrance_counting,
entrance=entrance)
if FLAGS.save_images: if FLAGS.save_images:
save_dir = os.path.join(FLAGS.output_dir, video_name.split('.')[-2]) save_dir = os.path.join(FLAGS.output_dir, video_name.split('.')[-2])
......
...@@ -128,7 +128,9 @@ def plot_tracking(image, ...@@ -128,7 +128,9 @@ def plot_tracking(image,
scores=None, scores=None,
frame_id=0, frame_id=0,
fps=0., fps=0.,
ids2names=[]): ids2names=[],
do_entrance_counting=False,
entrance=None):
im = np.ascontiguousarray(np.copy(image)) im = np.ascontiguousarray(np.copy(image))
im_h, im_w = im.shape[:2] im_h, im_w = im.shape[:2]
...@@ -153,7 +155,8 @@ def plot_tracking(image, ...@@ -153,7 +155,8 @@ def plot_tracking(image,
obj_id = int(obj_ids[i]) obj_id = int(obj_ids[i])
id_text = '{}'.format(int(obj_id)) id_text = '{}'.format(int(obj_id))
if ids2names != []: if ids2names != []:
assert len(ids2names) == 1, "plot_tracking only supports single classes." assert len(
ids2names) == 1, "plot_tracking only supports single classes."
id_text = '{}_'.format(ids2names[0]) + id_text id_text = '{}_'.format(ids2names[0]) + id_text
_line_thickness = 1 if obj_id <= 0 else line_thickness _line_thickness = 1 if obj_id <= 0 else line_thickness
color = get_color(abs(obj_id)) color = get_color(abs(obj_id))
...@@ -174,6 +177,15 @@ def plot_tracking(image, ...@@ -174,6 +177,15 @@ def plot_tracking(image,
cv2.FONT_HERSHEY_PLAIN, cv2.FONT_HERSHEY_PLAIN,
text_scale, (0, 255, 255), text_scale, (0, 255, 255),
thickness=text_thickness) thickness=text_thickness)
if do_entrance_counting:
entrance_line = tuple(map(int, entrance))
cv2.rectangle(
im,
entrance_line[0:2],
entrance_line[2:4],
color=(0, 255, 255),
thickness=line_thickness)
return im return im
...@@ -184,7 +196,9 @@ def plot_tracking_dict(image, ...@@ -184,7 +196,9 @@ def plot_tracking_dict(image,
scores_dict, scores_dict,
frame_id=0, frame_id=0,
fps=0., fps=0.,
ids2names=[]): ids2names=[],
do_entrance_counting=False,
entrance=None):
im = np.ascontiguousarray(np.copy(image)) im = np.ascontiguousarray(np.copy(image))
im_h, im_w = im.shape[:2] im_h, im_w = im.shape[:2]
...@@ -242,4 +256,12 @@ def plot_tracking_dict(image, ...@@ -242,4 +256,12 @@ def plot_tracking_dict(image,
cv2.FONT_HERSHEY_PLAIN, cv2.FONT_HERSHEY_PLAIN,
text_scale, (0, 255, 255), text_scale, (0, 255, 255),
thickness=text_thickness) thickness=text_thickness)
if num_classes == 1 and do_entrance_counting:
entrance_line = tuple(map(int, entrance))
cv2.rectangle(
im,
entrance_line[0:2],
entrance_line[2:4],
color=(0, 255, 255),
thickness=line_thickness)
return im return im
...@@ -228,6 +228,14 @@ void Pipeline::PredictMOT(const std::string& video_path) { ...@@ -228,6 +228,14 @@ void Pipeline::PredictMOT(const std::string& video_path) {
if (save_result_) { if (save_result_) {
PaddleDetection::SaveMOTResult(result, frame_id, &records); PaddleDetection::SaveMOTResult(result, frame_id, &records);
} }
// Draw the entrance line
if (do_entrance_counting_) {
float line_thickness = std::max(1, static_cast<int>(video_width / 500.));
cv::Point pt1 = cv::Point(entrance.left, entrance.top);
cv::Point pt2 = cv::Point(entrance.right, entrance.bottom);
cv::line(out_img, pt1, pt2, cv::Scalar(0, 255, 255), line_thickness);
}
video_out.write(out_img); video_out.write(out_img);
} }
capture.release(); capture.release();
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册