support directory predict

This commit is contained in:
LDOUBLEV 2021-05-26 06:49:17 +00:00
parent 5d24736a62
commit 4ae833c519
7 changed files with 90 additions and 29 deletions

View File

@ -44,6 +44,9 @@ public:
inline static size_t argmax(ForwardIterator first, ForwardIterator last) { inline static size_t argmax(ForwardIterator first, ForwardIterator last) {
return std::distance(first, std::max_element(first, 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 } // namespace PaddleOCR

View File

@ -27,9 +27,12 @@
#include <fstream> #include <fstream>
#include <numeric> #include <numeric>
#include <glog/logging.h>
#include <include/config.h> #include <include/config.h>
#include <include/ocr_det.h> #include <include/ocr_det.h>
#include <include/ocr_rec.h> #include <include/ocr_rec.h>
#include <include/utility.h>
#include <sys/stat.h>
using namespace std; using namespace std;
using namespace cv; using namespace cv;
@ -47,13 +50,9 @@ int main(int argc, char **argv) {
config.PrintConfigInfo(); config.PrintConfigInfo();
std::string img_path(argv[2]); 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, DBDetector det(config.det_model_dir, config.use_gpu, config.gpu_id,
config.gpu_mem, config.cpu_math_library_num_threads, config.gpu_mem, config.cpu_math_library_num_threads,
@ -76,9 +75,25 @@ int main(int argc, char **argv) {
config.use_tensorrt, config.use_fp16); config.use_tensorrt, config.use_fp16);
auto start = std::chrono::system_clock::now(); auto start = std::chrono::system_clock::now();
std::vector<std::vector<std::vector<int>>> boxes;
det.Run(srcimg, boxes);
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<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); rec.Run(boxes, srcimg, cls);
auto end = std::chrono::system_clock::now(); auto end = std::chrono::system_clock::now();
auto duration = auto duration =
@ -88,6 +103,7 @@ int main(int argc, char **argv) {
std::chrono::microseconds::period::num / std::chrono::microseconds::period::num /
std::chrono::microseconds::period::den std::chrono::microseconds::period::den
<< "s" << std::endl; << "s" << std::endl;
}
return 0; return 0;
} }

View File

@ -88,6 +88,10 @@ void CRNNRecognizer::Run(std::vector<std::vector<std::vector<int>>> boxes,
std::cout << str_res[i]; std::cout << str_res[i];
} }
std::cout << "\tscore: " << score << std::endl; 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;
} }
} }

View File

@ -12,12 +12,14 @@
// See the License for the specific language governing permissions and // See the License for the specific language governing permissions and
// limitations under the License. // limitations under the License.
#include <dirent.h>
#include <include/utility.h>
#include <iostream> #include <iostream>
#include <ostream> #include <ostream>
#include <sys/stat.h>
#include <sys/types.h>
#include <vector> #include <vector>
#include <include/utility.h>
namespace PaddleOCR { namespace PaddleOCR {
std::vector<std::string> Utility::ReadDict(const std::string &path) { std::vector<std::string> Utility::ReadDict(const std::string &path) {
@ -57,4 +59,37 @@ void Utility::VisualizeBboxes(
<< std::endl; << 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 } // namespace PaddleOCR

View File

@ -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 BUILD_DIR=build
rm -rf ${BUILD_DIR} rm -rf ${BUILD_DIR}
@ -12,9 +14,10 @@ cmake .. \
-DWITH_MKL=ON \ -DWITH_MKL=ON \
-DWITH_GPU=OFF \ -DWITH_GPU=OFF \
-DWITH_STATIC_LIB=OFF \ -DWITH_STATIC_LIB=OFF \
-DUSE_TENSORRT=OFF \ -DWITH_TENSORRT=ON \
-DOPENCV_DIR=${OPENCV_DIR} \ -DOPENCV_DIR=${OPENCV_DIR} \
-DCUDNN_LIB=${CUDNN_LIB_DIR} \ -DCUDNN_LIB=${CUDNN_LIB_DIR} \
-DCUDA_LIB=${CUDA_LIB_DIR} \ -DCUDA_LIB=${CUDA_LIB_DIR} \
-DTENSORRT_DIR=${TENSORRT_DIR} \
make -j make -j

View File

@ -1,5 +1,5 @@
# model load config # model load config
use_gpu 0 use_gpu 1
gpu_id 0 gpu_id 0
gpu_mem 4000 gpu_mem 4000
cpu_math_library_num_threads 10 cpu_math_library_num_threads 10
@ -23,7 +23,7 @@ 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 # show the detection results
visualize 1 visualize 0
# use_tensorrt # use_tensorrt
use_tensorrt 0 use_tensorrt 0

View File

@ -1,2 +1,2 @@
./build/ocr_system ./tools/config.txt ../../doc/imgs/12.jpg ./build/ocr_system ./tools/config.txt ../../doc/imgs/