support directory predict

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<std::string> &all_inputs);
 };
 
 } // namespace PaddleOCR

deploy/cpp_infer/src/main.cpp

@@ -27,9 +27,12 @@
 #include <fstream>
 #include <numeric>
 
+#include <glog/logging.h>
 #include <include/config.h>
 #include <include/ocr_det.h>
 #include <include/ocr_rec.h>
+#include <include/utility.h>
+#include <sys/stat.h>
 
 using namespace std;
 using namespace cv;
@@ -47,13 +50,9 @@ int main(int argc, char **argv) {
   config.PrintConfigInfo();
 
   std::string img_path(argv[2]);
-
+  std::vector<std::string> all_img_names;
-  cv::Mat srcimg = cv::imread(img_path, cv::IMREAD_COLOR);
+  // cv::Mat srcimg = cv::imread(img_path, cv::IMREAD_COLOR);
-
+  Utility::GetAllFiles((char *)img_path.c_str(), all_img_names);
-  if (!srcimg.data) {
-    std::cerr << "[ERROR] image read failed! image path: " << img_path << "\n";
-    exit(1);
-  }
 
   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<std::vector<std::vector<int>>> boxes;
-  det.Run(srcimg, boxes);
 
-  rec.Run(boxes, srcimg, cls);
+  for (auto img_dir : all_img_names) {
-  auto end = std::chrono::system_clock::now();
+    LOG(INFO) << "The predict img: " << img_dir;
-  auto duration =
+
-      std::chrono::duration_cast<std::chrono::microseconds>(end - start);
+    cv::Mat srcimg = cv::imread(img_dir, cv::IMREAD_COLOR);
-  std::cout << "Cost  "
+    if (!srcimg.data) {
-            << double(duration.count()) *
+      std::cerr << "[ERROR] image read failed! image path: " << img_path
-                   std::chrono::microseconds::period::num /
+                << "\n";
-                   std::chrono::microseconds::period::den
+      exit(1);
-            << "s" << std::endl;
+    }
+    std::vector<std::vector<std::vector<int>>> 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<std::chrono::microseconds>(end - start);
+    std::cout << "Cost  "
+              << double(duration.count()) *
+                     std::chrono::microseconds::period::num /
+                     std::chrono::microseconds::period::den
+              << "s" << std::endl;
+  }
 
   return 0;
 }

deploy/cpp_infer/src/ocr_rec.cpp

@@ -88,6 +88,10 @@ void CRNNRecognizer::Run(std::vector<std::vector<std::vector<int>>> 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;
   }
 }
 

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 <dirent.h>
+#include <include/utility.h>
 #include <iostream>
 #include <ostream>
+#include <sys/stat.h>
+#include <sys/types.h>
 #include <vector>
 
-#include <include/utility.h>
-
 namespace PaddleOCR {
 
 std::vector<std::string> 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<std::string> &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

deploy/cpp_infer/tools/build.sh

@@ -1,7 +1,9 @@
-OPENCV_DIR=your_opencv_dir
+OPENCV_DIR=/paddle/Paddle/opencv-3.4.7/opencv3
-LIB_DIR=your_paddle_inference_dir
+LIB_DIR=/paddle/OCR/debug/paddle_inference
-CUDA_LIB_DIR=your_cuda_lib_dir
+#LIB_DIR=/paddle/Paddle/inference/2.0.2/paddle_inference
-CUDNN_LIB_DIR=your_cudnn_lib_dir
+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

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

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/