recognition.h 3.6 KB
Newer Older
D
dongshuilong 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

15
#pragma once
D
dongshuilong 已提交
16 17 18 19 20 21 22 23
#include "paddle_api.h" // NOLINT
#include "json/json.h"
#include <arm_neon.h>
#include <chrono>
#include <fstream>
#include <iostream>
#include <math.h>
#include <opencv2/opencv.hpp>
D
dongshuilong 已提交
24
#include <stdlib.h>
D
dongshuilong 已提交
25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41
#include <sys/time.h>
#include <vector>

using namespace paddle::lite_api; // NOLINT
using namespace std;

namespace PPShiTu {

struct RESULT {
  std::string class_name;
  int class_id;
  float score;
};

class Recognition {

public:
D
dongshuilong 已提交
42
  explicit Recognition(const Json::Value &config_file) {
D
dongshuilong 已提交
43 44 45
    MobileConfig config;
    if (config_file["Global"]["rec_model_path"].as<std::string>().empty()) {
      std::cout << "Please set [rec_model_path] in config file" << std::endl;
D
dongshuilong 已提交
46
      exit(-1);
D
dongshuilong 已提交
47 48 49
    }
    config.set_model_from_file(
        config_file["Global"]["rec_model_path"].as<std::string>());
D
dongshuilong 已提交
50
    this->predictor = CreatePaddlePredictor<MobileConfig>(config);
D
dongshuilong 已提交
51 52 53

    if (config_file["Global"]["rec_label_path"].as<std::string>().empty()) {
      std::cout << "Please set [rec_label_path] in config file" << std::endl;
D
dongshuilong 已提交
54
      exit(-1);
D
dongshuilong 已提交
55 56 57
    }
    LoadLabel(config_file["Global"]["rec_label_path"].as<std::string>());
    SetPreProcessParam(config_file["RecPreProcess"]["transform_ops"]);
58
    if (!config_file["Global"].isMember("return_k")){
D
dongshuilong 已提交
59
      this->topk = config_file["Global"]["return_k"].as<int>();
60 61
    }
    printf("rec model create!\n");
D
dongshuilong 已提交
62 63 64 65 66 67 68 69 70 71 72 73 74
  }

  void LoadLabel(std::string path) {
    std::ifstream file;
    std::vector<std::string> label_list;
    file.open(path);
    while (file) {
      std::string line;
      std::getline(file, line);
      std::string::size_type pos = line.find(" ");
      if (pos != std::string::npos) {
        line = line.substr(pos);
      }
D
dongshuilong 已提交
75
      this->label_list.push_back(line);
D
dongshuilong 已提交
76 77 78 79 80 81 82 83 84
    }
    file.clear();
    file.close();
  }

  void SetPreProcessParam(const Json::Value &config_file) {
    for (const auto &item : config_file) {
      auto op_name = item["type"].as<std::string>();
      if (op_name == "ResizeImage") {
D
dongshuilong 已提交
85
        this->size = item["size"].as<int>();
D
dongshuilong 已提交
86
      } else if (op_name == "NormalizeImage") {
D
dongshuilong 已提交
87 88
        this->mean.clear();
        this->std.clear();
D
dongshuilong 已提交
89
        for (auto tmp : item["mean"]) {
D
dongshuilong 已提交
90
          this->mean.emplace_back(tmp.as<float>());
D
dongshuilong 已提交
91 92
        }
        for (auto tmp : item["std"]) {
D
dongshuilong 已提交
93
          this->std.emplace_back(1 / tmp.as<float>());
D
dongshuilong 已提交
94
        }
D
dongshuilong 已提交
95
        this->scale = item["scale"].as<double>();
D
dongshuilong 已提交
96 97
      }
    }
D
dongshuilong 已提交
98
  }
D
dongshuilong 已提交
99

D
dongshuilong 已提交
100 101 102 103 104
  std::vector<RESULT> RunRecModel(const cv::Mat &img, double &cost_time);
  std::vector<RESULT> PostProcess(const float *output_data, int output_size,
                                  cv::Mat &output_image);
  cv::Mat ResizeImage(const cv::Mat &img);
  void NeonMeanScale(const float *din, float *dout, int size);
D
dongshuilong 已提交
105

D
dongshuilong 已提交
106 107 108 109 110 111 112 113 114 115
private:
  std::shared_ptr<PaddlePredictor> predictor;
  std::vector<std::string> label_list;
  std::vector<float> mean = {0.485f, 0.456f, 0.406f};
  std::vector<float> std = {1 / 0.229f, 1 / 0.224f, 1 / 0.225f};
  double scale = 0.00392157;
  float size = 224;
  int topk = 5;
};
} // namespace PPShiTu