From 4ae833c519d445c8cad417b1abdf7b3c5bb8ab8b Mon Sep 17 00:00:00 2001 From: LDOUBLEV Date: Wed, 26 May 2021 06:49:17 +0000 Subject: [PATCH] support directory predict --- deploy/cpp_infer/include/utility.h | 3 ++ deploy/cpp_infer/src/main.cpp | 54 +++++++++++++++++++----------- deploy/cpp_infer/src/ocr_rec.cpp | 4 +++ deploy/cpp_infer/src/utility.cpp | 39 +++++++++++++++++++-- deploy/cpp_infer/tools/build.sh | 13 ++++--- deploy/cpp_infer/tools/config.txt | 6 ++-- deploy/cpp_infer/tools/run.sh | 2 +- 7 files changed, 91 insertions(+), 30 deletions(-) diff --git a/deploy/cpp_infer/include/utility.h b/deploy/cpp_infer/include/utility.h index 367e37e4..6e8173e0 100644 --- a/deploy/cpp_infer/include/utility.h +++ b/deploy/cpp_infer/include/utility.h @@ -44,6 +44,9 @@ public: inline static size_t argmax(ForwardIterator first, ForwardIterator last) { return std::distance(first, std::max_element(first, last)); } + + static void GetAllFiles(const char *dir_name, + std::vector &all_inputs); }; } // namespace PaddleOCR \ No newline at end of file diff --git a/deploy/cpp_infer/src/main.cpp b/deploy/cpp_infer/src/main.cpp index 588c8374..96b5e895 100644 --- a/deploy/cpp_infer/src/main.cpp +++ b/deploy/cpp_infer/src/main.cpp @@ -27,9 +27,12 @@ #include #include +#include #include #include #include +#include +#include using namespace std; using namespace cv; @@ -47,13 +50,9 @@ int main(int argc, char **argv) { config.PrintConfigInfo(); std::string img_path(argv[2]); - - cv::Mat srcimg = cv::imread(img_path, cv::IMREAD_COLOR); - - if (!srcimg.data) { - std::cerr << "[ERROR] image read failed! image path: " << img_path << "\n"; - exit(1); - } + std::vector all_img_names; + // cv::Mat srcimg = cv::imread(img_path, cv::IMREAD_COLOR); + Utility::GetAllFiles((char *)img_path.c_str(), all_img_names); DBDetector det(config.det_model_dir, config.use_gpu, config.gpu_id, config.gpu_mem, config.cpu_math_library_num_threads, @@ -76,18 +75,35 @@ int main(int argc, char **argv) { config.use_tensorrt, config.use_fp16); auto start = std::chrono::system_clock::now(); - std::vector>> boxes; - det.Run(srcimg, boxes); - - rec.Run(boxes, srcimg, cls); - auto end = std::chrono::system_clock::now(); - auto duration = - std::chrono::duration_cast(end - start); - std::cout << "Cost " - << double(duration.count()) * - std::chrono::microseconds::period::num / - std::chrono::microseconds::period::den - << "s" << std::endl; + + for (auto img_dir : all_img_names) { + LOG(INFO) << "The predict img: " << img_dir; + + cv::Mat srcimg = cv::imread(img_dir, cv::IMREAD_COLOR); + if (!srcimg.data) { + std::cerr << "[ERROR] image read failed! image path: " << img_path + << "\n"; + exit(1); + } + std::vector>> boxes; + + det.Run(srcimg, boxes); + // for (auto box : boxes){ + // std::cout << "box: " << box[0][0] << " " << box[0][1] << " " + // << box[1][0] << " " << box[1][1] << " " + // << box[2][0] << " " << box[2][1] << " " + // << box[3][0] << " " << box[3][1] << " " << std::endl; + // } + rec.Run(boxes, srcimg, cls); + auto end = std::chrono::system_clock::now(); + auto duration = + std::chrono::duration_cast(end - start); + std::cout << "Cost " + << double(duration.count()) * + std::chrono::microseconds::period::num / + std::chrono::microseconds::period::den + << "s" << std::endl; + } return 0; } diff --git a/deploy/cpp_infer/src/ocr_rec.cpp b/deploy/cpp_infer/src/ocr_rec.cpp index 76873dad..29b5618c 100644 --- a/deploy/cpp_infer/src/ocr_rec.cpp +++ b/deploy/cpp_infer/src/ocr_rec.cpp @@ -88,6 +88,10 @@ void CRNNRecognizer::Run(std::vector>> boxes, std::cout << str_res[i]; } std::cout << "\tscore: " << score << std::endl; + auto box = boxes[i]; + std::cout << "box: " << box[0][0] << " " << box[0][1] << " " << box[1][0] + << " " << box[1][1] << " " << box[2][0] << " " << box[2][1] + << " " << box[3][0] << " " << box[3][1] << " " << std::endl; } } diff --git a/deploy/cpp_infer/src/utility.cpp b/deploy/cpp_infer/src/utility.cpp index c1c9d938..2cd84f7e 100644 --- a/deploy/cpp_infer/src/utility.cpp +++ b/deploy/cpp_infer/src/utility.cpp @@ -12,12 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include +#include #include #include +#include +#include #include -#include - namespace PaddleOCR { std::vector Utility::ReadDict(const std::string &path) { @@ -57,4 +59,37 @@ void Utility::VisualizeBboxes( << std::endl; } +// list all files under a directory +void Utility::GetAllFiles(const char *dir_name, + std::vector &all_inputs) { + if (NULL == dir_name) { + std::cout << " dir_name is null ! " << std::endl; + return; + } + struct stat s; + lstat(dir_name, &s); + if (!S_ISDIR(s.st_mode)) { + std::cout << "dir_name is not a valid directory !" << std::endl; + all_inputs.push_back(dir_name); + return; + } else { + struct dirent *filename; // return value for readdir() + DIR *dir; // return value for opendir() + dir = opendir(dir_name); + if (NULL == dir) { + std::cout << "Can not open dir " << dir_name << std::endl; + return; + } + std::cout << "Successfully opened the dir !" << std::endl; + while ((filename = readdir(dir)) != NULL) { + if (strcmp(filename->d_name, ".") == 0 || + strcmp(filename->d_name, "..") == 0) + continue; + // img_dir + std::string("/") + all_inputs[0]; + all_inputs.push_back(dir_name + std::string("/") + + std::string(filename->d_name)); + } + } +} + } // namespace PaddleOCR \ No newline at end of file diff --git a/deploy/cpp_infer/tools/build.sh b/deploy/cpp_infer/tools/build.sh index 60653948..0c99f0e9 100755 --- a/deploy/cpp_infer/tools/build.sh +++ b/deploy/cpp_infer/tools/build.sh @@ -1,7 +1,9 @@ -OPENCV_DIR=your_opencv_dir -LIB_DIR=your_paddle_inference_dir -CUDA_LIB_DIR=your_cuda_lib_dir -CUDNN_LIB_DIR=your_cudnn_lib_dir +OPENCV_DIR=/paddle/Paddle/opencv-3.4.7/opencv3 +LIB_DIR=/paddle/OCR/debug/paddle_inference +#LIB_DIR=/paddle/Paddle/inference/2.0.2/paddle_inference +CUDA_LIB_DIR=/usr/local/cuda/lib64 +CUDNN_LIB_DIR=/usr/lib/x86_64-linux-gnu +TENSORRT_DIR=/paddle/Paddle/package/TensorRT/TensorRT-6.0.1.5/ BUILD_DIR=build rm -rf ${BUILD_DIR} @@ -12,9 +14,10 @@ cmake .. \ -DWITH_MKL=ON \ -DWITH_GPU=OFF \ -DWITH_STATIC_LIB=OFF \ - -DUSE_TENSORRT=OFF \ + -DWITH_TENSORRT=ON \ -DOPENCV_DIR=${OPENCV_DIR} \ -DCUDNN_LIB=${CUDNN_LIB_DIR} \ -DCUDA_LIB=${CUDA_LIB_DIR} \ + -DTENSORRT_DIR=${TENSORRT_DIR} \ make -j diff --git a/deploy/cpp_infer/tools/config.txt b/deploy/cpp_infer/tools/config.txt index 0e5f8472..9715fab3 100644 --- a/deploy/cpp_infer/tools/config.txt +++ b/deploy/cpp_infer/tools/config.txt @@ -1,5 +1,5 @@ # model load config -use_gpu 0 +use_gpu 1 gpu_id 0 gpu_mem 4000 cpu_math_library_num_threads 10 @@ -20,10 +20,10 @@ cls_thresh 0.9 # rec config rec_model_dir ./inference/ch_ppocr_mobile_v2.0_rec_infer/ -char_list_file ../../ppocr/utils/ppocr_keys_v1.txt +char_list_file ../../ppocr/utils/ppocr_keys_v1.txt # show the detection results -visualize 1 +visualize 0 # use_tensorrt use_tensorrt 0 diff --git a/deploy/cpp_infer/tools/run.sh b/deploy/cpp_infer/tools/run.sh index fa61da75..73118b71 100755 --- a/deploy/cpp_infer/tools/run.sh +++ b/deploy/cpp_infer/tools/run.sh @@ -1,2 +1,2 @@ -./build/ocr_system ./tools/config.txt ../../doc/imgs/12.jpg +./build/ocr_system ./tools/config.txt ../../doc/imgs/ -- GitLab