From 9378909c392a970fbbe0eca110d7812cb275677c Mon Sep 17 00:00:00 2001 From: wangguanzhong Date: Sat, 6 Nov 2021 17:48:53 +0800 Subject: [PATCH] [MOT] Add stream input (#4489) * add stream input * add stream input & add save_result --- deploy/pptracking/CMakeLists.txt | 2 +- deploy/pptracking/include/pipeline.h | 35 ++++-- deploy/pptracking/include/postprocess.h | 2 +- deploy/pptracking/src/main.cc | 14 ++- deploy/pptracking/src/pipeline.cc | 156 +++++++++++++++++++----- deploy/pptracking/src/postprocess.cc | 27 +++- 6 files changed, 186 insertions(+), 50 deletions(-) diff --git a/deploy/pptracking/CMakeLists.txt b/deploy/pptracking/CMakeLists.txt index e93799325..303b3c47d 100644 --- a/deploy/pptracking/CMakeLists.txt +++ b/deploy/pptracking/CMakeLists.txt @@ -19,7 +19,7 @@ include_directories("${CMAKE_SOURCE_DIR}/") include_directories("${CMAKE_CURRENT_BINARY_DIR}/ext/yaml-cpp/src/ext-yaml-cpp/include") link_directories("${CMAKE_CURRENT_BINARY_DIR}/ext/yaml-cpp/lib") -set(SRCS src/main.cc src/preprocess_op.cc src/pipeline.cc src/predictor.cc src/jde_predictor.cc src/sde_predictor.cc src/tracker.cc src/trajectory.cc src/lapjv.cpp src/base_predictor.cc src/postprocess.cc) +set(SRCS src/main.cc src/preprocess_op.cc src/pipeline.cc src/jde_predictor.cc src/sde_predictor.cc src/tracker.cc src/trajectory.cc src/lapjv.cpp src/base_predictor.cc src/postprocess.cc) macro(safe_set_static_flag) foreach(flag_var diff --git a/deploy/pptracking/include/pipeline.h b/deploy/pptracking/include/pipeline.h index 8b736b5bf..dd1d00491 100644 --- a/deploy/pptracking/include/pipeline.h +++ b/deploy/pptracking/include/pipeline.h @@ -31,7 +31,8 @@ #include #endif -#include "include/predictor.h" +#include "include/jde_predictor.h" +#include "include/sde_predictor.h" namespace PaddleDetection { @@ -46,7 +47,10 @@ class Pipeline { const int cpu_threads=1, const bool trt_calib_mode=false, const bool count=false, - const bool save_result=false) { + const bool save_result=false, + const std::string& scene="pedestrian", + const bool tiny_obj=false, + const bool is_mtmct=false) { std::vector input; this->input_ = input; this->device_ = device; @@ -59,34 +63,45 @@ class Pipeline { this->trt_calib_mode_ = trt_calib_mode; this->count_ = count; this->save_result_ = save_result; + SelectModel(scene, tiny_obj, is_mtmct); + InitPredictor(); } - // 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_mct=false); // Set input, it must execute before Run() void SetInput(std::string& input_video); + void ClearInput(); - // Run pipeline + // Run pipeline in video void Run(); + void PredictMOT(const std::string& video_path); + void PredictMTMCT(const std::vector video_inputs); - void PredictSCT(const std::string& video_path); - void PredictMCT(const std::vector video_inputs); + // Run pipeline in stream + void RunMOTStream(const cv::Mat img, const int frame_id, cv::Mat out_img, std::vector& records, std::vector& count_list, std::vector& in_count_list, std::vector& out_count_list); + void RunMTMCTStream(const std::vector imgs, std::vector& records); void PrintBenchmarkLog(std::vector det_time, int img_num); private: + // 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); + void InitPredictor(); + + std::shared_ptr jde_sct_; + std::shared_ptr sde_sct_; + std::vector input_; + std::vector stream_; std::string device_; double threshold_; std::string output_dir_; std::string track_model_dir_; std::string det_model_dir_; std::string reid_model_dir_; - std::string mct_model_dir_; std::string run_mode_ = "fluid"; int gpu_id_ = 0; bool use_mkldnn_ = false; diff --git a/deploy/pptracking/include/postprocess.h b/deploy/pptracking/include/postprocess.h index 933aeb4b5..64a49bb75 100644 --- a/deploy/pptracking/include/postprocess.h +++ b/deploy/pptracking/include/postprocess.h @@ -43,6 +43,6 @@ void FlowStatistic(const MOTResult& results, const int frame_id, std::vector* out_count_list); // Save Tracking Results -void SaveResult(const MOTResult& results, const std::string& output_dir); +void SaveMOTResult(const MOTResult& results, const int frame_id, std::vector& records); } // namespace PaddleDetection diff --git a/deploy/pptracking/src/main.cc b/deploy/pptracking/src/main.cc index a3645b167..21538a3de 100644 --- a/deploy/pptracking/src/main.cc +++ b/deploy/pptracking/src/main.cc @@ -36,7 +36,7 @@ DEFINE_string(video_file, "", "Path of input video."); -DEFINE_string(video_other_file, "", "Path of other input video used for MCT."); +DEFINE_string(video_other_file, "", "Path of other input video used for MTMCT."); DEFINE_string(device, "CPU", "Choose the device you want to run, it can be: CPU/GPU/XPU, default is CPU."); DEFINE_double(threshold, 0.5, "Threshold of score."); DEFINE_string(output_dir, "output", "Directory of output visualization files."); @@ -49,7 +49,7 @@ DEFINE_bool(tiny_obj, false, "Whether tracking tiny object"); DEFINE_bool(count, false, "Whether counting after tracking"); DEFINE_bool(save_result, false, "Whether saving result after tracking"); DEFINE_string(scene, "", "scene of tracking system, it can be : pedestrian/vehicle/multiclass"); -DEFINE_bool(is_mct, false, "Whether use multi-camera tracking"); +DEFINE_bool(is_mtmct, false, "Whether use multi-target multi-camera tracking"); static std::string DirName(const std::string &filepath) { auto pos = filepath.rfind(OS_PATH_SEP); @@ -112,13 +112,17 @@ int main(int argc, char** argv) { return -1; } + if (!PathExists(FLAGS_output_dir)) { + MkDirs(FLAGS_output_dir); + } + PaddleDetection::Pipeline pipeline( FLAGS_device, FLAGS_threshold, FLAGS_output_dir, FLAGS_run_mode, FLAGS_gpu_id, - FLAGS_use_mkldnn, FLAGS_cpu_threads, - FLAGS_count, FLAGS_save_result); + FLAGS_use_mkldnn, FLAGS_cpu_threads, FLAGS_trt_calib_mode, + FLAGS_count, FLAGS_save_result, FLAGS_scene, FLAGS_tiny_obj, + FLAGS_is_mtmct); - pipeline.SelectModel(FLAGS_scene, FLAGS_tiny_obj, FLAGS_is_mct); pipeline.SetInput(FLAGS_video_file); if (!FLAGS_video_other_file.empty()) { pipeline.SetInput(FLAGS_video_other_file); diff --git a/deploy/pptracking/src/pipeline.cc b/deploy/pptracking/src/pipeline.cc index 9457358a2..e36add5c9 100644 --- a/deploy/pptracking/src/pipeline.cc +++ b/deploy/pptracking/src/pipeline.cc @@ -29,34 +29,64 @@ void Pipeline::SetInput(std::string& input_video) { input_.push_back(input_video); } +void Pipeline::ClearInput() { + input_.clear(); + stream_.clear(); +} + void Pipeline::SelectModel(const std::string& scene, const bool tiny_obj, - const bool is_mct) { - // Single camera model - // use deepsort for multiclass tracking - // use fairmot for single class tracking + const bool is_mtmct) { + // Single camera model, based on FairMot if (scene == "pedestrian") { + if (tiny_obj) { + track_model_dir_ = "../pedestrian_track_tiny"; + } else { track_model_dir_ = "../pedestrian_track"; + } } else if (scene != "vehicle") { + if (tiny_obj) { + track_model_dir_ = "../vehicle_track_tiny"; + } else { track_model_dir_ = "../vehicle_track"; + } } else if (scene == "multiclass") { - det_model_dir_ = "../multiclass_det"; - reid_model_dir_ = "../multiclass_reid"; + if (tiny_obj) { + track_model_dir_ = "../multiclass_track_tiny"; + } else { + track_model_dir_ = "../multiclass_track"; + } } - // Multi-camera model - if (is_mct && scene == "pedestrian") { - mct_model_dir_ = "../pedestrian_mct"; - } else if (is_mct && scene == "vehicle") { - mct_model_dir_ = "../vehicle_mct"; - } else if (is_mct && scene == "multiclass") { + // Multi-camera model, based on PicoDet & LCNet + if (is_mtmct && scene == "pedestrian") { + det_model_dir_ = "../pedestrian_det"; + reid_model_dir_ = "../pedestrian_reid"; + } else if (is_mtmct && scene == "vehicle") { + det_model_dir_ = "../vehicle_det"; + reid_model_dir_ = "../vehicle_reid"; + } else if (is_mtmct && scene == "multiclass") { throw "Multi-camera tracking is not supported in multiclass scene now."; - } + } } +void Pipeline::InitPredictor() { + if (track_model_dir_.empty() && det_model_dir_.empty()) { + throw "Predictor must receive track_model or det_model!"; + } + + if (!track_model_dir_.empty()) { + jde_sct_ = std::make_shared(device_, track_model_dir_, threshold_, run_mode_, gpu_id_, use_mkldnn_, cpu_threads_, trt_calib_mode_); + } + if (!det_model_dir_.empty()) { + sde_sct_ = std::make_shared(device_, det_model_dir_, reid_model_dir_, threshold_, run_mode_, gpu_id_, use_mkldnn_, cpu_threads_, trt_calib_mode_); + } +} + + void Pipeline::Run() { - if (track_model_dir_.empty()) { + if (track_model_dir_.empty() && det_model_dir_.empty()) { std::cout << "Pipeline must use SelectModel before Run"; return; } @@ -65,24 +95,24 @@ void Pipeline::Run() { return; } - if (mct_model_dir_.empty()) { + if (!track_model_dir_.empty()) { // single camera if (input_.size() > 1) { throw "Single camera tracking except single video, but received %d", input_.size(); } - PredictSCT(input_[0]); + PredictMOT(input_[0]); } else { // multi cameras if (input_.size() != 2) { throw "Multi camera tracking except two videos, but received %d", input_.size(); } - PredictMCT(input_); + PredictMTMCT(input_); } } -void Pipeline::PredictSCT(const std::string& video_path) { - PaddleDetection::Predictor sct(device_, track_model_dir_, det_model_dir_, reid_model_dir_, threshold_, run_mode_, gpu_id_, use_mkldnn_, cpu_threads_, trt_calib_mode_); +void Pipeline::PredictMOT(const std::string& video_path) { + // Open video cv::VideoCapture capture; capture.open(video_path.c_str()); @@ -96,9 +126,14 @@ void Pipeline::PredictSCT(const std::string& video_path) { int video_height = static_cast(capture.get(CV_CAP_PROP_FRAME_HEIGHT)); int video_fps = static_cast(capture.get(CV_CAP_PROP_FPS)); + LOG(INFO) << "----------------------- Input info -----------------------"; + LOG(INFO) << "video_width: " << video_width; + LOG(INFO) << "video_height: " << video_height; + LOG(INFO) << "input fps: " << video_fps; + // Create VideoWriter for output cv::VideoWriter video_out; - std::string video_out_path = "mot_output.mp4"; + std::string video_out_path = output_dir_ + OS_PATH_SEP + "mot_output.mp4"; int fcc = cv::VideoWriter::fourcc('m','p','4','v'); video_out.open(video_out_path.c_str(), fcc, //0x00000021, @@ -116,21 +151,30 @@ void Pipeline::PredictSCT(const std::string& video_path) { std::vector in_count_list; std::vector out_count_list; double times; + double total_time; // Capture all frames and do inference cv::Mat frame; int frame_id = 0; + + std::vector records; + records.push_back("result format: frame_id, track_id, x1, y1, w, h\n"); + + LOG(INFO) << "------------------- Predict info ------------------------"; while (capture.read(frame)) { if (frame.empty()) { break; } std::vector imgs; imgs.push_back(frame); - printf("frame_id: %d\n", frame_id); - sct.Predict(imgs, threshold_, &result, &det_times); + jde_sct_->Predict(imgs, threshold_, &result, &det_times); frame_id += 1; - times = std::accumulate(det_times.begin(), det_times.end(), 0) / frame_id; + total_time = std::accumulate(det_times.begin(), det_times.end(), 0.); + times = total_time / frame_id; + + LOG(INFO) << "frame_id: " << frame_id + << " predict time(s): "<< total_time / 1000; - cv::Mat out_im = PaddleDetection::VisualizeTrackResult( + cv::Mat out_img = PaddleDetection::VisualizeTrackResult( frame, result, 1000./times, frame_id); if (count_) { @@ -139,17 +183,69 @@ void Pipeline::PredictSCT(const std::string& video_path) { PaddleDetection::FlowStatistic(result, frame_id, &count_list, &in_count_list, &out_count_list); } if (save_result_) { - PaddleDetection::SaveResult(result, output_dir_); + PaddleDetection::SaveMOTResult(result, frame_id, records); } - video_out.write(out_im); + video_out.write(out_img); } capture.release(); video_out.release(); PrintBenchmarkLog(det_times, frame_id); - printf("Visualized output saved as %s\n", video_out_path.c_str()); + LOG(INFO) << "-------------------- Final Output info -------------------"; + LOG(INFO) << "Total frame: " << frame_id; + LOG(INFO) << "Visualized output saved as " << video_out_path.c_str(); + if (save_result_) { + FILE * fp; + + std::string result_output_path = output_dir_ + OS_PATH_SEP + "mot_output.txt"; + if((fp = fopen(result_output_path.c_str(), "w+")) == NULL) { + printf("Open %s error.\n", result_output_path.c_str()); + return; + } + for (int l; l < records.size(); ++l) { + fprintf(fp, records[l].c_str()); + } + + fclose(fp); + LOG(INFO) << "txt result output saved as " << result_output_path.c_str(); + } +} + +void Pipeline::PredictMTMCT(const std::vector video_path) { + throw "Not Implement!"; +} + +void Pipeline::RunMOTStream(const cv::Mat img, const int frame_id, cv::Mat out_img, std::vector& records, std::vector& count_list, std::vector& in_count_list, std::vector& out_count_list) { + PaddleDetection::MOTResult result; + std::vector det_times(3); + double times; + double total_time; + + LOG(INFO) << "------------------- Predict info ------------------------"; + std::vector imgs; + imgs.push_back(img); + jde_sct_->Predict(imgs, threshold_, &result, &det_times); + total_time = std::accumulate(det_times.begin(), det_times.end(), 0.); + times = total_time / frame_id; + + LOG(INFO) << "frame_id: " << frame_id + << " predict time(s): "<< total_time / 1000; + + out_img = PaddleDetection::VisualizeTrackResult( + img, result, 1000./times, frame_id); + + if (count_) { + // Count total number + // Count in & out number + PaddleDetection::FlowStatistic(result, frame_id, &count_list, &in_count_list, &out_count_list); + } + + PrintBenchmarkLog(det_times, frame_id); + if (save_result_) { + PaddleDetection::SaveMOTResult(result, frame_id, records); + } } -void Pipeline::PredictMCT(const std::vector video_path) { +void Pipeline::RunMTMCTStream(const std::vector imgs, std::vector& records) { throw "Not Implement!"; } @@ -171,8 +267,8 @@ void Pipeline::PrintBenchmarkLog(std::vector det_time, int img_num){ LOG(INFO) << "cpu_math_library_num_threads: " << cpu_threads_; LOG(INFO) << "----------------------- Perf info ------------------------"; LOG(INFO) << "Total number of predicted data: " << img_num - << " and total time spent(ms): " - << std::accumulate(det_time.begin(), det_time.end(), 0.); + << " and total time spent(s): " + << std::accumulate(det_time.begin(), det_time.end(), 0.) / 1000; img_num = std::max(1, img_num); LOG(INFO) << "preproce_time(ms): " << det_time[0] / img_num << ", inference_time(ms): " << det_time[1] / img_num diff --git a/deploy/pptracking/src/postprocess.cc b/deploy/pptracking/src/postprocess.cc index efa6c4cba..9614e6065 100644 --- a/deploy/pptracking/src/postprocess.cc +++ b/deploy/pptracking/src/postprocess.cc @@ -104,9 +104,30 @@ void FlowStatistic(const MOTResult& results, const int frame_id, throw "Not Implement"; } -void SaveResult(const MOTResult& results, const std::string& output_dir) { - throw "Not Implement"; +void SaveMOTResult(const MOTResult& results, const int frame_id, std::vector& records) { + // result format: frame_id, track_id, x1, y1, w, h + std::string record; + for (int i = 0; i < results.size(); ++i) { + MOTTrack mot_track = results[i]; + int ids = mot_track.ids; + float score = mot_track.score; + Rect rects = mot_track.rects; + float x1 = rects.left; + float y1 = rects.top; + float x2 = rects.right; + float y2 = rects.bottom; + float w = x2 - x1; + float h = y2 - y1; + if (w == 0 || h == 0) { + continue; + } + std::ostringstream os; + os << frame_id << " " << ids << "" + << x1 << " " << y1 << " " + << w << " " << h <<"\n"; + record = os.str(); + records.push_back(record); + } } - } // namespace PaddleDetection -- GitLab