feature_extractor.h 2.8 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
#include "paddle_api.h" // NOLINT
#include "json/json.h"
#include <arm_neon.h>
#include <chrono>
#include <fstream>
D
dongshuilong 已提交
21
#include <include/preprocess_op.h>
D
dongshuilong 已提交
22 23 24
#include <iostream>
#include <math.h>
#include <opencv2/opencv.hpp>
D
dongshuilong 已提交
25
#include <stdlib.h>
D
dongshuilong 已提交
26 27 28 29 30 31 32 33 34 35 36 37 38 39
#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;
};

L
lubin 已提交
40
class FeatureExtract {
D
dongshuilong 已提交
41
public:
L
lubin 已提交
42
  explicit FeatureExtract(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

    SetPreProcessParam(config_file["RecPreProcess"]["transform_ops"]);
L
lubin 已提交
53
    printf("feature extract model create!\n");
D
dongshuilong 已提交
54 55 56 57 58 59
  }

  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 已提交
60
        this->size = item["size"].as<int>();
D
dongshuilong 已提交
61
      } else if (op_name == "NormalizeImage") {
D
dongshuilong 已提交
62 63
        this->mean.clear();
        this->std.clear();
D
dongshuilong 已提交
64
        for (auto tmp : item["mean"]) {
D
dongshuilong 已提交
65
          this->mean.emplace_back(tmp.as<float>());
D
dongshuilong 已提交
66 67
        }
        for (auto tmp : item["std"]) {
D
dongshuilong 已提交
68
          this->std.emplace_back(tmp.as<float>());
D
dongshuilong 已提交
69
        }
D
dongshuilong 已提交
70
        this->scale = item["scale"].as<double>();
D
dongshuilong 已提交
71 72
      }
    }
D
dongshuilong 已提交
73
  }
D
dongshuilong 已提交
74

D
dongshuilong 已提交
75 76 77
  void RunRecModel(const cv::Mat &img, double &cost_time,
                   std::vector<float> &feature);
  // void PostProcess(std::vector<float> &feature);
D
dongshuilong 已提交
78
  void FeatureNorm(std::vector<float> &featuer);
D
dongshuilong 已提交
79

D
dongshuilong 已提交
80 81
private:
  std::shared_ptr<PaddlePredictor> predictor;
D
dongshuilong 已提交
82
  // std::vector<std::string> label_list;
D
dongshuilong 已提交
83
  std::vector<float> mean = {0.485f, 0.456f, 0.406f};
D
dongshuilong 已提交
84
  std::vector<float> std = {0.229f, 0.224f, 0.225f};
D
dongshuilong 已提交
85
  double scale = 0.00392157;
D
dongshuilong 已提交
86 87 88 89 90 91
  int size = 224;

  // pre-process
  Resize resize_op_;
  NormalizeImage normalize_op_;
  Permute permute_op_;
D
dongshuilong 已提交
92 93
};
} // namespace PPShiTu