From ce04964e8e54526c1ab1dd8ea7f6d2c737198222 Mon Sep 17 00:00:00 2001
From: FlyingQianMM <245467267@qq.com>
Date: Fri, 12 Nov 2021 00:19:28 +0800
Subject: [PATCH] [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
---
 deploy/pptracking/python/mot_jde_infer.py |  5 +++-
 deploy/pptracking/python/mot_sde_infer.py |  4 +++-
 deploy/pptracking/python/visualize.py     | 28 ++++++++++++++++++++---
 deploy/pptracking/src/pipeline.cc         |  8 +++++++
 4 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/deploy/pptracking/python/mot_jde_infer.py b/deploy/pptracking/python/mot_jde_infer.py
index d7edaeeab..41011f5f2 100644
--- a/deploy/pptracking/python/mot_jde_infer.py
+++ b/deploy/pptracking/python/mot_jde_infer.py
@@ -272,7 +272,10 @@ def predict_video(detector, camera_id):
             online_scores,
             frame_id=frame_id,
             fps=fps,
-            ids2names=ids2names)
+            ids2names=ids2names,
+            do_entrance_counting=FLAGS.do_entrance_counting,
+            entrance=entrance)
+
         if FLAGS.save_images:
             save_dir = os.path.join(FLAGS.output_dir, video_name.split('.')[-2])
             if not os.path.exists(save_dir):
diff --git a/deploy/pptracking/python/mot_sde_infer.py b/deploy/pptracking/python/mot_sde_infer.py
index 0f1837778..7f5c3ed78 100644
--- a/deploy/pptracking/python/mot_sde_infer.py
+++ b/deploy/pptracking/python/mot_sde_infer.py
@@ -596,7 +596,9 @@ def predict_video(detector, reid_model, camera_id):
                 online_ids,
                 online_scores,
                 frame_id=frame_id,
-                fps=fps)
+                fps=fps,
+                do_entrance_counting=FLAGS.do_entrance_counting,
+                entrance=entrance)
 
         if FLAGS.save_images:
             save_dir = os.path.join(FLAGS.output_dir, video_name.split('.')[-2])
diff --git a/deploy/pptracking/python/visualize.py b/deploy/pptracking/python/visualize.py
index 79e72ce2e..86d9fc599 100644
--- a/deploy/pptracking/python/visualize.py
+++ b/deploy/pptracking/python/visualize.py
@@ -128,7 +128,9 @@ def plot_tracking(image,
                   scores=None,
                   frame_id=0,
                   fps=0.,
-                  ids2names=[]):
+                  ids2names=[],
+                  do_entrance_counting=False,
+                  entrance=None):
     im = np.ascontiguousarray(np.copy(image))
     im_h, im_w = im.shape[:2]
 
@@ -153,7 +155,8 @@ def plot_tracking(image,
         obj_id = int(obj_ids[i])
         id_text = '{}'.format(int(obj_id))
         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
         _line_thickness = 1 if obj_id <= 0 else line_thickness
         color = get_color(abs(obj_id))
@@ -174,6 +177,15 @@ def plot_tracking(image,
                 cv2.FONT_HERSHEY_PLAIN,
                 text_scale, (0, 255, 255),
                 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
 
 
@@ -184,7 +196,9 @@ def plot_tracking_dict(image,
                        scores_dict,
                        frame_id=0,
                        fps=0.,
-                       ids2names=[]):
+                       ids2names=[],
+                       do_entrance_counting=False,
+                       entrance=None):
     im = np.ascontiguousarray(np.copy(image))
     im_h, im_w = im.shape[:2]
 
@@ -242,4 +256,12 @@ def plot_tracking_dict(image,
                     cv2.FONT_HERSHEY_PLAIN,
                     text_scale, (0, 255, 255),
                     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
diff --git a/deploy/pptracking/src/pipeline.cc b/deploy/pptracking/src/pipeline.cc
index 0ed852fff..9606f65a6 100644
--- a/deploy/pptracking/src/pipeline.cc
+++ b/deploy/pptracking/src/pipeline.cc
@@ -228,6 +228,14 @@ void Pipeline::PredictMOT(const std::string& video_path) {
     if (save_result_) {
       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);
   }
   capture.release();
-- 
GitLab