add faiss for ppshitu

deploy/cpp_shitu/CMakeLists.txt

@@ -5,15 +5,23 @@ option(WITH_MKL        "Compile demo with MKL/OpenBlas support, default use MKL.
 option(WITH_GPU        "Compile demo with GPU/CPU, default use CPU."                    OFF)
 option(WITH_STATIC_LIB "Compile demo with static/shared library, default use static."   ON)
 option(WITH_TENSORRT "Compile demo with TensorRT."   OFF)
+option(FAISS_WITH_MKL "Faiss Compile demo with MKL."   OFF)
 
 SET(PADDLE_LIB "" CACHE PATH "Location of libraries")
 SET(OPENCV_DIR "" CACHE PATH "Location of libraries")
 SET(CUDA_LIB "" CACHE PATH "Location of libraries")
 SET(CUDNN_LIB "" CACHE PATH "Location of libraries")
 SET(TENSORRT_DIR "" CACHE PATH "Compile demo with TensorRT")
+SET(FAISS_DIR "" CACHE PATH "Location of libraries")
 
 set(DEMO_NAME "pp_shitu")
 
+if (FAISS_WITH_MKL)
+   SET(BLAS_NAME "mklml_intel")
+else()
+   SET(BLAS_NAME "openblas")
+endif()
+
 include(external-cmake/yaml-cpp.cmake)
 include_directories("${CMAKE_SOURCE_DIR}/")
 include_directories("${CMAKE_CURRENT_BINARY_DIR}/ext/yaml-cpp/src/ext-yaml-cpp/include")
@@ -34,6 +42,10 @@ if (WITH_MKL)
     ADD_DEFINITIONS(-DUSE_MKL)
 endif()
 
+if(NOT DEFINED FAISS_DIR)
+  message(FATAL_ERROR "please set FAISS_DIR with -DFAISS_DIR=/path/faiss")
+endif()
+
 if(NOT DEFINED PADDLE_LIB)
   message(FATAL_ERROR "please set PADDLE_LIB with -DPADDLE_LIB=/path/paddle/lib")
 endif()
@@ -42,6 +54,15 @@ if(NOT DEFINED OPENCV_DIR)
     message(FATAL_ERROR "please set OPENCV_DIR with -DOPENCV_DIR=/path/opencv")
 endif()
 
+if (WIN32)
+  include_directories("${FAISS_DIR}/include")
+  link_directories("${FAISS_DIR}/lib")
+  find_package(faiss REQUIRED PATHS ${FAISS_DIR}/share/faiss/ NO_DEFAULT_PATH)
+else()
+  find_package(faiss REQUIRED PATHS ${FAISS_DIR}/share/faiss NO_DEFAULT_PATH)
+  include_directories("${FAISS_DIR}/include")
+  link_directories("${FAISS_DIR}/lib")
+endif()
 
 if (WIN32)
   include_directories("${PADDLE_LIB}/paddle/fluid/inference")
@@ -213,6 +234,8 @@ AUX_SOURCE_DIRECTORY(./src SRCS)
 add_executable(${DEMO_NAME} ${SRCS})
 ADD_DEPENDENCIES(${DEMO_NAME} ext-yaml-cpp)
 target_link_libraries(${DEMO_NAME} ${DEPS})
+target_link_libraries(${DEMO_NAME} ${FAISS_DIR}/lib/libfaiss${CMAKE_STATIC_LIBRARY_SUFFIX})
+target_link_libraries(${DEMO_NAME} ${BLAS_NAME}${CMAKE_SHARED_LIBRARY_SUFFIX})
 
 if (WIN32 AND WITH_MKL)
     add_custom_command(TARGET ${DEMO_NAME} POST_BUILD

deploy/cpp_shitu/include/vector_search.h

+// Copyright (c) 2020 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.
+
+#pragma once
+#ifdef WIN32
+#define OS_PATH_SEP "\\"
+#else
+#define OS_PATH_SEP "/"
+#endif
+
+#include <faiss/Index.h>
+#include <faiss/index_io.h>
+#include <cstring>
+#include <map>
+#include "yaml-cpp/yaml.h"
+
+struct SearchResult{
+  faiss::Index::idx_t* I;
+  float* D;
+  int query_number;
+  int return_k;
+};
+
+class VectorSearch{
+public:
+  explicit VectorSearch(const YAML::Node &config_file){
+    // IndexProcess
+    this->index_dir = config_file["IndexProcess"]["index_dir"].as<std::string>();
+    this->return_k = config_file["IndexProcess"]["return_k"].as<int>();
+    this->score_thres = config_file["IndexProcess"]["score_thres"].as<float>();
+    this->max_query_number = config_file["Global"]["max_det_results"].as<int>() + 1;
+    LoadIdMap();
+    LoadIndexFile();
+    this->I = new faiss::Index::idx_t[this->return_k * this->max_query_number];
+    this->D = new float[this->return_k * this->max_query_number];
+  }
+  ~VectorSearch(){
+    delete[] I;
+    delete[] D;
+  };
+  void LoadIdMap();
+  void LoadIndexFile();
+  void Search(float* feature, int query_number);
+  const SearchResult& GetSearchResult();
+  const std::string& GetLabel(faiss::Index::idx_t ind);
+
+private:
+  std::string index_dir;
+  int return_k = 5;
+  float score_thres = 0.5;
+  std::map <long int, std::string> id_map;
+  faiss::Index * index;
+  int max_query_number = 6;
+  int real_query_number = 0;
+  float *D = NULL;
+  faiss::Index::idx_t* I = NULL;
+  SearchResult sr;
+};

deploy/cpp_shitu/include/yaml_config.h

@@ -39,7 +39,6 @@
 #include "opencv2/imgproc.hpp"
 #include "yaml-cpp/yaml.h"
 
-namespace PaddleClas {
 
 class YamlConfig {
 public:
@@ -52,5 +51,3 @@ public:
   void PrintConfigInfo();
   YAML::Node config_file;
 };
-
-} // namespace PaddleClas

deploy/cpp_shitu/readme.md

@@ -134,6 +134,15 @@ tar -xvf paddle_inference.tgz
 
 最终会在当前的文件夹中生成`paddle_inference/`的子文件夹。
 
+### 1.3 安装faiss库
+
+```shell
+ git clone https://github.com/facebookresearch/faiss.git
+ cd faiss 
+ cmake -B build . -DFAISS_ENABLE_PYTHON=OFF -DCMAKE_INSTALL_PREFIX=${faiss_install_path}
+ make -C build -j faiss
+ make -C build install
+```
 
 ## 2 开始运行
 

deploy/cpp_shitu/src/main.cpp

@@ -30,6 +30,7 @@
 #include <include/cls.h>
 #include <include/object_detector.h>
 #include <include/yaml_config.h>
+#include <include/vector_search.h>
 
 using namespace std;
 using namespace cv;
@@ -133,7 +134,7 @@ int main(int argc, char **argv) {
     exit(1);
   }
 
-  PaddleClas::YamlConfig config(argv[1]);
+  YamlConfig config(argv[1]);
   config.PrintConfigInfo();
 
   // config

deploy/cpp_shitu/src/vector_search.cpp

+// Copyright (c) 2020 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.
+#include <faiss/index_factory.h>
+#include <faiss/index_io.h>
+#include <fstream>
+#include <regex>
+#include <iostream>
+#include <cstdio>
+#include "include/vector_search.h"
+
+void VectorSearch::LoadIndexFile(){
+  std::string file_path = this->index_dir + OS_PATH_SEP + "vector.index"; 
+  const char* fname = file_path.c_str();
+  this->index = faiss::read_index(fname, 0);
+}
+
+void VectorSearch::LoadIdMap(){
+  std::string file_path = this->index_dir + OS_PATH_SEP + "id_map.txt";
+  std::ifstream in(file_path);
+  std::string line;
+  std::vector<std::string> m_vec;
+  if (in){
+    while (getline(in, line)){
+      std::regex ws_re("\\s+");
+      std::vector<std::string> v(
+	std::sregex_token_iterator(line.begin(), line.end(), ws_re, -1),
+	std::sregex_token_iterator());
+      if (v.size() !=2){
+	std::cout << "The number of element for each line in : " << file_path
+		  << "must be 2, exit the program..." << std::endl;
+	exit(1);
+      }else
+        this->id_map.insert(std::pair<long int, std::string>(std::stol(v[0], nullptr, 10), v[1]));
+    }
+  }
+}
+
+void VectorSearch::Search(float *feature, int query_number){
+  this->index->search(query_number, feature, return_k, D, I);
+  this->real_query_number = query_number;
+}
+
+const SearchResult& VectorSearch::GetSearchResult(){
+  this->sr.query_number = this->real_query_number;
+  this->sr.return_k = this->return_k;
+  this->sr.D = this->D;
+  this->sr.I = this->I;
+  return this->sr;
+}
+
+const std::string& VectorSearch::GetLabel(faiss::Index::idx_t ind){
+  return this->id_map.at(ind);
+}

deploy/cpp_shitu/src/yaml_config.cpp

@@ -18,7 +18,6 @@
 
 #include <include/yaml_config.h>
 
-namespace PaddleClas {
 
 std::vector<std::string> YamlConfig::ReadDict(const std::string &path) {
   std::ifstream in(path);
@@ -77,5 +76,3 @@ void YamlConfig::PrintConfigInfo() {
   //   std::cout << it->as<std::string>() << "\n";
   //   }
 }
-
-} // namespace PaddleClas

deploy/cpp_shitu/tools/build.sh

@@ -2,6 +2,7 @@ OPENCV_DIR=/work/project/project/cpp_infer/opencv-3.4.7/opencv3
 LIB_DIR=/work/project/project/cpp_infer/paddle_inference/
 CUDA_LIB_DIR=/usr/local/cuda/lib64
 CUDNN_LIB_DIR=/usr/lib/x86_64-linux-gnu/
+FAISS_DIR=/work/project/project/cpp_infer/faiss/faiss_install
 
 BUILD_DIR=build
 rm -rf ${BUILD_DIR}
@@ -16,5 +17,7 @@ cmake .. \
     -DOPENCV_DIR=${OPENCV_DIR} \
     -DCUDNN_LIB=${CUDNN_LIB_DIR} \
     -DCUDA_LIB=${CUDA_LIB_DIR} \
+    -DFAISS_DIR=${FAISS_DIR} \
+    -DFAISS_WITH_MKL=OFF
 
 make -j

deploy/cpp_shitu/tools/transform_id_map.py

+import argparse
+import os
+import pickle
+
+import yaml
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-c', '--config', type=str, required=True)
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    with open(args.config)as fd:
+        config = yaml.load(fd)
+    index_dir = ""
+    try:
+        index_dir = config["IndexProcess"]["index_dir"]
+    except Exception as e:
+        print("The IndexProcess.index_dir in config_file dose not exist") 
+        exit(1)
+    id_map_path = os.path.join(index_dir, "id_map.pkl")
+    assert os.path.exists(id_map_path), "The id_map file dose not exist: {}".format(id_map_path)
+
+    with open(id_map_path, "rb")as fd:
+        ids = pickle.load(fd)
+    with open(os.path.join(index_dir, "id_map.txt"), "w")as fd:
+        for k, v in ids.items():
+            v = v.split("\t")[1]
+            fd.write(str(k) + " " + v + "\n")
+    print('Transform id_map sucess')
+
+
+if __name__ == "__main__":
+    main()