diff --git a/deploy/pptracking/include/pipeline.h b/deploy/pptracking/include/pipeline.h index 185cd894100ed228d390b833afa1330087c938bc..442af09aeda8bb6c686603f07ef2b67ad5679b7c 100644 --- a/deploy/pptracking/include/pipeline.h +++ b/deploy/pptracking/include/pipeline.h @@ -53,7 +53,10 @@ class Pipeline { const std::string& scene = "pedestrian", const bool tiny_obj = false, const bool is_mtmct = false, - const int secs_interval = 10) { + const int secs_interval = 10, + const std::string track_model_dir = "", + const std::string det_model_dir = "", + const std::string reid_model_dir = "") { std::vector input; this->input_ = input; this->device_ = device; @@ -67,7 +70,12 @@ class Pipeline { this->do_entrance_counting_ = do_entrance_counting; this->secs_interval_ = secs_interval_; this->save_result_ = save_result; - SelectModel(scene, tiny_obj, is_mtmct); + SelectModel(scene, + tiny_obj, + is_mtmct, + track_model_dir, + det_model_dir, + reid_model_dir); InitPredictor(); } @@ -102,7 +110,10 @@ class Pipeline { // Select model according to scenes, it must execute before Run() void SelectModel(const std::string& scene = "pedestrian", const bool tiny_obj = false, - const bool is_mtmct = false); + const bool is_mtmct = false, + const std::string track_model_dir = "", + const std::string det_model_dir = "", + const std::string reid_model_dir = ""); void InitPredictor(); std::shared_ptr jde_sct_; diff --git a/deploy/pptracking/src/main.cc b/deploy/pptracking/src/main.cc index 1f54c87e1bc327bfcc626383b57a40665f343ce5..9861eecb91d334feecfd0bc88f4e9dc856ca152b 100644 --- a/deploy/pptracking/src/main.cc +++ b/deploy/pptracking/src/main.cc @@ -54,7 +54,8 @@ DEFINE_bool(trt_calib_mode, "If the model is produced by TRT offline quantitative calibration, " "trt_calib_mode need to set True"); DEFINE_bool(tiny_obj, false, "Whether tracking tiny object"); -DEFINE_bool(do_entrance_counting, false, +DEFINE_bool(do_entrance_counting, + false, "Whether counting the numbers of identifiers entering " "or getting out from the entrance."); DEFINE_int32(secs_interval, 10, "The seconds interval to count after tracking"); @@ -64,6 +65,9 @@ DEFINE_string( "", "scene of tracking system, it can be : pedestrian/vehicle/multiclass"); DEFINE_bool(is_mtmct, false, "Whether use multi-target multi-camera tracking"); +DEFINE_string(track_model_dir, "", "Path of tracking model"); +DEFINE_string(det_model_dir, "", "Path of detection model"); +DEFINE_string(reid_model_dir, "", "Path of reid model"); static std::string DirName(const std::string& filepath) { auto pos = filepath.rfind(OS_PATH_SEP); @@ -109,15 +113,21 @@ static void MkDirs(const std::string& path) { int main(int argc, char** argv) { // Parsing command-line google::ParseCommandLineFlags(&argc, &argv, true); - if (FLAGS_video_file.empty() || FLAGS_scene.empty()) { - std::cout << "Usage: ./main " - << "-video_file=/PATH/TO/INPUT/IMAGE/ " - << "-scene=pedestrian/vehicle/multiclass" << std::endl; + bool has_model_dir = + !(FLAGS_track_model_dir.empty() && FLAGS_det_model_dir.empty() && + FLAGS_reid_model_dir.empty()); + if (FLAGS_video_file.empty() || (FLAGS_scene.empty() && !has_model_dir)) { + LOG(ERROR) << "Usage: \n" + << "1. ./main -video_file=/PATH/TO/INPUT/IMAGE/ " + << "-scene=pedestrian/vehicle/multiclass\n" + << "2. ./main -video_file=/PATH/TO/INPUT/IMAGE/ " + << "-track_model_dir=/PATH/TO/MODEL_DIR" << std::endl; + return -1; } if (!(FLAGS_run_mode == "fluid" || FLAGS_run_mode == "trt_fp32" || FLAGS_run_mode == "trt_fp16" || FLAGS_run_mode == "trt_int8")) { - std::cout + LOG(ERROR) << "run_mode should be 'fluid', 'trt_fp32', 'trt_fp16' or 'trt_int8'."; return -1; } @@ -127,7 +137,7 @@ int main(int argc, char** argv) { ::toupper); if (!(FLAGS_device == "CPU" || FLAGS_device == "GPU" || FLAGS_device == "XPU")) { - std::cout << "device should be 'CPU', 'GPU' or 'XPU'."; + LOG(ERROR) << "device should be 'CPU', 'GPU' or 'XPU'."; return -1; } @@ -148,7 +158,10 @@ int main(int argc, char** argv) { FLAGS_scene, FLAGS_tiny_obj, FLAGS_is_mtmct, - FLAGS_secs_interval); + FLAGS_secs_interval, + FLAGS_track_model_dir, + FLAGS_det_model_dir, + FLAGS_reid_model_dir); pipeline.SetInput(FLAGS_video_file); if (!FLAGS_video_other_file.empty()) { diff --git a/deploy/pptracking/src/pipeline.cc b/deploy/pptracking/src/pipeline.cc index d508650877b0cb54dab6faa67360adbc2e5730cb..0ed852ffffda88801a84bf9d4a1a7ee7d03ee55c 100644 --- a/deploy/pptracking/src/pipeline.cc +++ b/deploy/pptracking/src/pipeline.cc @@ -36,7 +36,21 @@ void Pipeline::ClearInput() { void Pipeline::SelectModel(const std::string& scene, const bool tiny_obj, - const bool is_mtmct) { + const bool is_mtmct, + const std::string track_model_dir, + const std::string det_model_dir, + const std::string reid_model_dir) { + // model_dir has higher priority + if (!track_model_dir.empty()) { + track_model_dir_ = track_model_dir; + return; + } + if (!det_model_dir.empty() && !reid_model_dir.empty()) { + det_model_dir_ = det_model_dir; + reid_model_dir_ = reid_model_dir; + return; + } + // Single camera model, based on FairMot if (scene == "pedestrian") { if (tiny_obj) { @@ -100,11 +114,11 @@ void Pipeline::InitPredictor() { void Pipeline::Run() { if (track_model_dir_.empty() && det_model_dir_.empty()) { - std::cout << "Pipeline must use SelectModel before Run"; + LOG(ERROR) << "Pipeline must use SelectModel before Run"; return; } if (input_.size() == 0) { - std::cout << "Pipeline must use SetInput before Run"; + LOG(ERROR) << "Pipeline must use SetInput before Run"; return; } @@ -165,7 +179,8 @@ void Pipeline::PredictMOT(const std::string& video_path) { std::vector in_id_list; std::vector out_id_list; std::map> prev_center; - Rect entrance = {0, static_cast(video_height) / 2, + Rect entrance = {0, + static_cast(video_height) / 2, static_cast(video_width), static_cast(video_height) / 2}; double times; @@ -195,12 +210,20 @@ void Pipeline::PredictMOT(const std::string& video_path) { cv::Mat out_img = PaddleDetection::VisualizeTrackResult( frame, result, 1000. / times, frame_id); - - // TODO: the entrance line can be set by users - PaddleDetection::FlowStatistic( - result, frame_id, secs_interval_, do_entrance_counting_, video_fps, entrance, - &id_set, &interval_id_set, &in_id_list, &out_id_list, - &prev_center, &flow_records); + + // TODO(qianhui): the entrance line can be set by users + PaddleDetection::FlowStatistic(result, + frame_id, + secs_interval_, + do_entrance_counting_, + video_fps, + entrance, + &id_set, + &interval_id_set, + &in_id_list, + &out_id_list, + &prev_center, + &flow_records); if (save_result_) { PaddleDetection::SaveMOTResult(result, frame_id, &records); @@ -228,7 +251,7 @@ void Pipeline::PredictMOT(const std::string& video_path) { fclose(fp); LOG(INFO) << "txt result output saved as " << result_output_path.c_str(); - + result_output_path = output_dir_ + OS_PATH_SEP + "flow_statistic.txt"; if ((fp = fopen(result_output_path.c_str(), "w+")) == NULL) { printf("Open %s error.\n", result_output_path); @@ -273,15 +296,23 @@ void Pipeline::RunMOTStream(const cv::Mat img, LOG(INFO) << "frame_id: " << frame_id << " predict time(s): " << total_time / 1000; - out_img = PaddleDetection::VisualizeTrackResult(img, result, 1000. / times, - frame_id); + out_img = PaddleDetection::VisualizeTrackResult( + img, result, 1000. / times, frame_id); // Count total number // Count in & out number - PaddleDetection::FlowStatistic(result, frame_id, secs_interval_, do_entrance_counting_, - video_fps, entrance, id_set, - interval_id_set, in_id_list, - out_id_list, prev_center, flow_records); + PaddleDetection::FlowStatistic(result, + frame_id, + secs_interval_, + do_entrance_counting_, + video_fps, + entrance, + id_set, + interval_id_set, + in_id_list, + out_id_list, + prev_center, + flow_records); PrintBenchmarkLog(det_times, frame_id); if (save_result_) { @@ -326,4 +357,4 @@ void Pipeline::PrintBenchmarkLog(const std::vector det_time, << ", postprocess_time(ms): " << det_time[2] / num; } -} // namespace PaddleDetection +} // namespace PaddleDetection