From c44543aeaa9c33025dce442bd51239d07af18554 Mon Sep 17 00:00:00 2001
From: FlyingQianMM <245467267@qq.com>
Date: Thu, 11 Nov 2021 10:21:48 +0800
Subject: [PATCH] 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
---
 deploy/pptracking/python/mot/utils.py     |  2 +-
 deploy/pptracking/python/mot_sde_infer.py | 32 ++++++++++++++++++++++-
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/deploy/pptracking/python/mot/utils.py b/deploy/pptracking/python/mot/utils.py
index 3cee6b2f6..8b6be3995 100644
--- a/deploy/pptracking/python/mot/utils.py
+++ b/deploy/pptracking/python/mot/utils.py
@@ -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.
diff --git a/deploy/pptracking/python/mot_sde_infer.py b/deploy/pptracking/python/mot_sde_infer.py
index b25dab4f1..0f1837778 100644
--- a/deploy/pptracking/python/mot_sde_infer.py
+++ b/deploy/pptracking/python/mot_sde_infer.py
@@ -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,
-- 
GitLab