Optimize ocr system code (#209)

* Support PPYOLOE plus model

* Optimize ocr system code

* modify example code

* fix patchelf of openvino

* optimize demo code of ocr

* remove debug code

* update demo code of ocr

Co-authored-by: Jack Zhou <zhoushunjie@baidu.com>
This commit is contained in:
Jason
2022-09-14 09:46:03 +08:00
committed by GitHub
parent 1452275efe
commit 0dd9ecee65
39 changed files with 1041 additions and 909 deletions

View File

@@ -19,11 +19,7 @@ const char sep = '\\';
const char sep = '/';
#endif
void CpuInfer(const std::string& det_model_dir,
const std::string& cls_model_dir,
const std::string& rec_model_dir,
const std::string& rec_label_file,
const std::string& image_file) {
void InitAndInfer(const std::string& det_model_dir, const std::string& cls_model_dir, const std::string& rec_model_dir, const std::string& rec_label_file, const std::string& image_file, const fastdeploy::RuntimeOption& option) {
auto det_model_file = det_model_dir + sep + "inference.pdmodel";
auto det_params_file = det_model_dir + sep + "inference.pdiparams";
@@ -32,235 +28,32 @@ void CpuInfer(const std::string& det_model_dir,
auto rec_model_file = rec_model_dir + sep + "inference.pdmodel";
auto rec_params_file = rec_model_dir + sep + "inference.pdiparams";
auto rec_label = rec_label_file;
fastdeploy::vision::ocr::DBDetector det_model;
fastdeploy::vision::ocr::Classifier cls_model;
fastdeploy::vision::ocr::Recognizer rec_model;
auto det_model = fastdeploy::vision::ocr::DBDetector(det_model_file, det_params_file, option);
auto cls_model = fastdeploy::vision::ocr::Classifier(cls_model_file, cls_params_file, option);
auto rec_model = fastdeploy::vision::ocr::Recognizer(rec_model_file, rec_params_file, rec_label_file, option);
if (!det_model_dir.empty()) {
auto det_option = fastdeploy::RuntimeOption();
det_option.UseCpu();
det_model = fastdeploy::vision::ocr::DBDetector(
det_model_file, det_params_file, det_option);
assert(det_model.Initialized());
assert(cls_model.Initialized());
assert(rec_model.Initialized());
if (!det_model.Initialized()) {
std::cerr << "Failed to initialize det_model." << std::endl;
return;
}
}
if (!cls_model_dir.empty()) {
auto cls_option = fastdeploy::RuntimeOption();
cls_option.UseCpu();
cls_model = fastdeploy::vision::ocr::Classifier(
cls_model_file, cls_params_file, cls_option);
if (!cls_model.Initialized()) {
std::cerr << "Failed to initialize cls_model." << std::endl;
return;
}
}
if (!rec_model_dir.empty()) {
auto rec_option = fastdeploy::RuntimeOption();
rec_option.UseCpu();
rec_model = fastdeploy::vision::ocr::Recognizer(
rec_model_file, rec_params_file, rec_label, rec_option);
if (!rec_model.Initialized()) {
std::cerr << "Failed to initialize rec_model." << std::endl;
return;
}
}
auto ocrv3_app = fastdeploy::application::ocrsystem::PPOCRSystemv3(
&det_model, &cls_model, &rec_model);
// 其中分类模型可选因此也可使用如下方式串联OCR系统
// auto ocr_system_v3 = fastdeploy::application::ocrsystem::PPOCRSystemv3(&det_model, &rec_model);
auto ocr_system_v3 = fastdeploy::application::ocrsystem::PPOCRSystemv3(&det_model, &cls_model, &rec_model);
auto im = cv::imread(image_file);
auto im_bak = im.clone();
fastdeploy::vision::OCRResult res;
//开始预测
if (!ocrv3_app.Predict(&im, &res)) {
fastdeploy::vision::OCRResult result;
if (!ocr_system_v3.Predict(&im, &result)) {
std::cerr << "Failed to predict." << std::endl;
return;
}
//输出预测信息
std::cout << res.Str() << std::endl;
std::cout << result.Str() << std::endl;
//可视化
auto vis_img = fastdeploy::vision::Visualize::VisOcr(im_bak, res);
cv::imwrite("vis_result.jpg", vis_img);
std::cout << "Visualized result saved in ./vis_result.jpg" << std::endl;
}
void GpuInfer(const std::string& det_model_dir,
const std::string& cls_model_dir,
const std::string& rec_model_dir,
const std::string& rec_label_file,
const std::string& image_file) {
auto det_model_file = det_model_dir + sep + "inference.pdmodel";
auto det_params_file = det_model_dir + sep + "inference.pdiparams";
auto cls_model_file = cls_model_dir + sep + "inference.pdmodel";
auto cls_params_file = cls_model_dir + sep + "inference.pdiparams";
auto rec_model_file = rec_model_dir + sep + "inference.pdmodel";
auto rec_params_file = rec_model_dir + sep + "inference.pdiparams";
auto rec_label = rec_label_file;
fastdeploy::vision::ocr::DBDetector det_model;
fastdeploy::vision::ocr::Classifier cls_model;
fastdeploy::vision::ocr::Recognizer rec_model;
//准备模型
if (!det_model_dir.empty()) {
auto det_option = fastdeploy::RuntimeOption();
det_option.UseGpu();
det_model = fastdeploy::vision::ocr::DBDetector(
det_model_file, det_params_file, det_option);
if (!det_model.Initialized()) {
std::cerr << "Failed to initialize det_model." << std::endl;
return;
}
}
if (!cls_model_dir.empty()) {
auto cls_option = fastdeploy::RuntimeOption();
cls_option.UseGpu();
cls_model = fastdeploy::vision::ocr::Classifier(
cls_model_file, cls_params_file, cls_option);
if (!cls_model.Initialized()) {
std::cerr << "Failed to initialize cls_model." << std::endl;
return;
}
}
if (!rec_model_dir.empty()) {
auto rec_option = fastdeploy::RuntimeOption();
rec_option.UseGpu();
rec_model = fastdeploy::vision::ocr::Recognizer(
rec_model_file, rec_params_file, rec_label, rec_option);
if (!rec_model.Initialized()) {
std::cerr << "Failed to initialize rec_model." << std::endl;
return;
}
}
auto ocrv3_app = fastdeploy::application::ocrsystem::PPOCRSystemv3(
&det_model, &cls_model, &rec_model);
auto im = cv::imread(image_file);
auto im_bak = im.clone();
fastdeploy::vision::OCRResult res;
//开始预测
if (!ocrv3_app.Predict(&im, &res)) {
std::cerr << "Failed to predict." << std::endl;
return;
}
//输出预测信息
std::cout << res.Str() << std::endl;
//可视化
auto vis_img = fastdeploy::vision::Visualize::VisOcr(im_bak, res);
cv::imwrite("vis_result.jpg", vis_img);
std::cout << "Visualized result saved in ./vis_result.jpg" << std::endl;
}
void TrtInfer(const std::string& det_model_dir,
const std::string& cls_model_dir,
const std::string& rec_model_dir,
const std::string& rec_label_file,
const std::string& image_file) {
auto det_model_file = det_model_dir + sep + "inference.pdmodel";
auto det_params_file = det_model_dir + sep + "inference.pdiparams";
auto cls_model_file = cls_model_dir + sep + "inference.pdmodel";
auto cls_params_file = cls_model_dir + sep + "inference.pdiparams";
auto rec_model_file = rec_model_dir + sep + "inference.pdmodel";
auto rec_params_file = rec_model_dir + sep + "inference.pdiparams";
auto rec_label = rec_label_file;
fastdeploy::vision::ocr::DBDetector det_model;
fastdeploy::vision::ocr::Classifier cls_model;
fastdeploy::vision::ocr::Recognizer rec_model;
//准备模型
if (!det_model_dir.empty()) {
auto det_option = fastdeploy::RuntimeOption();
det_option.UseGpu();
det_option.UseTrtBackend();
det_option.SetTrtInputShape("x", {1, 3, 50, 50}, {1, 3, 640, 640},
{1, 3, 960, 960});
det_model = fastdeploy::vision::ocr::DBDetector(
det_model_file, det_params_file, det_option);
if (!det_model.Initialized()) {
std::cerr << "Failed to initialize det_model." << std::endl;
return;
}
}
if (!cls_model_dir.empty()) {
auto cls_option = fastdeploy::RuntimeOption();
cls_option.UseGpu();
cls_option.UseTrtBackend();
cls_option.SetTrtInputShape("x", {1, 3, 48, 192});
cls_model = fastdeploy::vision::ocr::Classifier(
cls_model_file, cls_params_file, cls_option);
if (!cls_model.Initialized()) {
std::cerr << "Failed to initialize cls_model." << std::endl;
return;
}
}
if (!rec_model_dir.empty()) {
auto rec_option = fastdeploy::RuntimeOption();
rec_option.UseGpu();
rec_option.UseTrtBackend();
rec_option.SetTrtInputShape("x", {1, 3, 48, 10}, {1, 3, 48, 320},
{1, 3, 48, 2000});
rec_model = fastdeploy::vision::ocr::Recognizer(
rec_model_file, rec_params_file, rec_label, rec_option);
if (!rec_model.Initialized()) {
std::cerr << "Failed to initialize rec_model." << std::endl;
return;
}
}
auto ocrv3_app = fastdeploy::application::ocrsystem::PPOCRSystemv3(
&det_model, &cls_model, &rec_model);
auto im = cv::imread(image_file);
auto im_bak = im.clone();
fastdeploy::vision::OCRResult res;
//开始预测
if (!ocrv3_app.Predict(&im, &res)) {
std::cerr << "Failed to predict." << std::endl;
return;
}
//输出预测信息
std::cout << res.Str() << std::endl;
//可视化
auto vis_img = fastdeploy::vision::Visualize::VisOcr(im_bak, res);
cv::imwrite("vis_result.jpg", vis_img);
auto vis_im = fastdeploy::vision::Visualize::VisOcr(im_bak, result);
cv::imwrite("vis_result.jpg", vis_im);
std::cout << "Visualized result saved in ./vis_result.jpg" << std::endl;
}
@@ -279,12 +72,23 @@ int main(int argc, char* argv[]) {
return -1;
}
if (std::atoi(argv[6]) == 0) {
CpuInfer(argv[1], argv[2], argv[3], argv[4], argv[5]);
} else if (std::atoi(argv[6]) == 1) {
GpuInfer(argv[1], argv[2], argv[3], argv[4], argv[5]);
} else if (std::atoi(argv[6]) == 2) {
TrtInfer(argv[1], argv[2], argv[3], argv[4], argv[5]);
fastdeploy::RuntimeOption option;
int flag = std::atoi(argv[6]);
if (flag == 0) {
option.UseCpu();
} else if (flag == 1) {
option.UseGpu();
} else if (flag == 2) {
option.UseGpu();
option.UseTrtBackend();
}
std::string det_model_dir = argv[1];
std::string cls_model_dir = argv[2];
std::string rec_model_dir = argv[3];
std::string rec_label_file = argv[4];
std::string test_image = argv[5];
InitAndInfer(det_model_dir, cls_model_dir, rec_model_dir, rec_label_file, test_image, option);
return 0;
}
}