diff --git a/deploy/pptracking/python/mot/utils.py b/deploy/pptracking/python/mot/utils.py index 7b0da82746b1774d667145d33fe73349d44840d8..1fd6548d81099292e8113fd5b18af16869fd0114 100644 --- a/deploy/pptracking/python/mot/utils.py +++ b/deploy/pptracking/python/mot/utils.py @@ -19,14 +19,8 @@ import paddle import numpy as np __all__ = [ - 'MOTTimer', - 'Detection', - 'write_mot_results', - 'load_det_results', - 'preprocess_reid', - 'get_crops', - 'clip_box', - 'scale_coords', + 'MOTTimer', 'Detection', 'write_mot_results', 'load_det_results', + 'preprocess_reid', 'get_crops', 'clip_box', 'scale_coords', 'flow_statistic' ] @@ -219,3 +213,81 @@ def preprocess_reid(imgs, im_batch.append(img) im_batch = np.concatenate(im_batch, 0) return im_batch + + +def flow_statistic(result, + secs_interval, + do_entrance_counting, + video_fps, + entrance, + id_set, + interval_id_set, + in_id_list, + out_id_list, + prev_center, + records, + data_type, + num_classes=1): + # Count in and out number: + # Use horizontal center line as the entrance just for simplification. + # If a person located in the above the horizontal center line + # at the previous frame and is in the below the line at the current frame, + # the in number is increased by one. + # If a person was in the below the horizontal center line + # at the previous frame and locates in the below the line at the current frame, + # the out number is increased by one. + # TODO: if the entrance is not the horizontal center line, + # the counting method should be optimized. + if do_entrance_counting: + entrance_y = entrance[1] # xmin, ymin, xmax, ymax + frame_id, tlwhs, tscores, track_ids = result + for tlwh, score, track_id in zip(tlwhs, tscores, track_ids): + if track_id < 0: continue + if data_type == 'kitti': + frame_id -= 1 + + x1, y1, w, h = tlwh + center_x = x1 + w / 2. + center_y = y1 + h / 2. + if track_id in prev_center: + if prev_center[track_id][1] <= entrance_y and \ + center_y > entrance_y: + in_id_list.append(track_id) + if prev_center[track_id][1] >= entrance_y and \ + center_y < entrance_y: + out_id_list.append(track_id) + prev_center[track_id][0] = center_x + prev_center[track_id][1] = center_y + else: + prev_center[track_id] = [center_x, center_y] + # Count totol number, number at a manual-setting interval + frame_id, tlwhs, tscores, track_ids = result + for tlwh, score, track_id in zip(tlwhs, tscores, track_ids): + if track_id < 0: continue + id_set.add(track_id) + interval_id_set.add(track_id) + + # Reset counting at the interval beginning + if frame_id % video_fps == 0 and frame_id / video_fps % secs_interval == 0: + curr_interval_count = len(interval_id_set) + interval_id_set.clear() + info = "Frame id: {}, Total count: {}".format(frame_id, len(id_set)) + if do_entrance_counting: + info += ", In count: {}, Out count: {}".format( + len(in_id_list), len(out_id_list)) + if frame_id % video_fps == 0 and frame_id / video_fps % secs_interval == 0: + info += ", Count during {} secs: {}".format(secs_interval, + curr_interval_count) + interval_id_set.clear() + print(info) + info += "\n" + records.append(info) + + return { + "id_set": id_set, + "interval_id_set": interval_id_set, + "in_id_list": in_id_list, + "out_id_list": out_id_list, + "prev_center": prev_center, + "records": records + } diff --git a/deploy/pptracking/python/mot_jde_infer.py b/deploy/pptracking/python/mot_jde_infer.py index 0f9220d3fe38a5ee56bcceac4bf129420fc26b99..a79cc4512d6a1c2be0b338f8e5d9b4f461ec4b52 100644 --- a/deploy/pptracking/python/mot_jde_infer.py +++ b/deploy/pptracking/python/mot_jde_infer.py @@ -29,7 +29,7 @@ from benchmark_utils import PaddleInferBenchmark from visualize import plot_tracking_dict from mot.tracker import JDETracker -from mot.utils import MOTTimer, write_mot_results +from mot.utils import MOTTimer, write_mot_results, flow_statistic # Global dictionary MOT_SUPPORT_MODELS = { @@ -220,6 +220,16 @@ def predict_video(detector, camera_id): data_type = 'mcmot' if num_classes > 1 else 'mot' ids2names = detector.pred_config.labels + if num_classes == 1: + 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: @@ -233,6 +243,25 @@ def predict_video(detector, camera_id): results[cls_id].append((frame_id + 1, online_tlwhs[cls_id], online_scores[cls_id], online_ids[cls_id])) + # NOTE: just implement flow statistic for one class + if num_classes == 1: + result = (frame_id + 1, online_tlwhs[0], online_scores[0], + online_ids[0]) + 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, data_type, num_classes) + 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'] + + elif num_classes > 1 and do_entrance_counting: + raise NotImplementedError( + 'Multi-class flow counting is not implemented now!') + fps = 1. / timer.average_time im = plot_tracking_dict( frame, @@ -264,6 +293,16 @@ def predict_video(detector, camera_id): write_mot_results(result_filename, results, data_type, num_classes) + if num_classes == 1: + 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, diff --git a/deploy/pptracking/python/utils.py b/deploy/pptracking/python/utils.py index 9af16ba7643c42143bd4c81d448b2fbf6e949693..b3f2852e90f23d4ec8241c140f63c66e44cdabde 100644 --- a/deploy/pptracking/python/utils.py +++ b/deploy/pptracking/python/utils.py @@ -120,6 +120,17 @@ def argsparser(): type=int, default=50, help="max batch_size for reid model inference.") + parser.add_argument( + "--do_entrance_counting", + action='store_true', + help="Whether counting the numbers of identifiers entering " + "or getting out from the entrance. Note that only support one-class" + "counting, multi-class counting is coming soon.") + parser.add_argument( + "--secs_interval", + type=int, + default=10, + help="The seconds interval to count after tracking") return parser