add faiss lib to do retrieval; upate reademe.md

deploy/lite_shitu/Makefile

@@ -30,6 +30,8 @@ OPENCV_LIBS = ${THIRD_PARTY_DIR}/${OPENCV_VERSION}/${ARM_PLAT}/libs/libopencv_im
               ${THIRD_PARTY_DIR}/${OPENCV_VERSION}/${ARM_PLAT}/3rdparty/libs/libtbb.a \
               ${THIRD_PARTY_DIR}/${OPENCV_VERSION}/${ARM_PLAT}/3rdparty/libs/libcpufeatures.a
 
+FAISS_VERSION=faiss1.5.3
+FAISS_LIBS = ${THIRD_PARTY_DIR}/${FAISS_VERSION}/libs/${ARM_PLAT}/libfaiss.a
 
 LITE_LIBS = -L${LITE_ROOT}/cxx/lib/ -lpaddle_light_api_shared
 ###############################################################
@@ -43,7 +45,7 @@ LITE_LIBS = -L${LITE_ROOT}/cxx/lib/ -lpaddle_light_api_shared
 # 2. Undo comment below line using `libpaddle_api_light_bundled.a`
 # LITE_LIBS = ${LITE_ROOT}/cxx/lib/libpaddle_api_light_bundled.a
 
-CXX_LIBS = $(LITE_LIBS) ${OPENCV_LIBS} $(SYSTEM_LIBS)
+CXX_LIBS = $(LITE_LIBS) ${OPENCV_LIBS} ${FAISS_LIBS} $(SYSTEM_LIBS)
 
 LOCAL_DIRSRCS=$(wildcard src/*.cc)
 LOCAL_SRCS=$(notdir $(LOCAL_DIRSRCS))
@@ -51,9 +53,17 @@ LOCAL_OBJS=$(patsubst %.cpp, %.o, $(patsubst %.cc, %.o, $(LOCAL_SRCS)))
 
 JSON_OBJS = json_reader.o json_value.o json_writer.o
 
-pp_shitu: $(LOCAL_OBJS) $(JSON_OBJS) fetch_opencv
+pp_shitu: $(LOCAL_OBJS) $(JSON_OBJS) fetch_opencv fetch_faiss
 	$(CC) $(SYSROOT_LINK) $(CXXFLAGS_LINK) $(LOCAL_OBJS) $(JSON_OBJS) -o pp_shitu $(CXX_LIBS) $(LDFLAGS)
 
+fetch_faiss:
+	@ test -d ${THIRD_PARTY_DIR} ||  mkdir ${THIRD_PARTY_DIR}
+	@ test -e ${THIRD_PARTY_DIR}/${FAISS_VERSION}.tar.gz || \
+      (echo "fetch faiss libs" && \
+      wget -P ${THIRD_PARTY_DIR} https://paddle-inference-dist.bj.bcebos.com/${FAISS_VERSION}.tar.gz)
+	@ test -d ${THIRD_PARTY_DIR}/${FAISS_VERSION} || \
+      tar -xf ${THIRD_PARTY_DIR}/${FAISS_VERSION}.tar.gz -C ${THIRD_PARTY_DIR}
+
 fetch_opencv:
 	@ test -d ${THIRD_PARTY_DIR} ||  mkdir ${THIRD_PARTY_DIR}
 	@ test -e ${THIRD_PARTY_DIR}/${OPENCV_VERSION}.tar.gz || \
@@ -72,11 +82,12 @@ fetch_json_code:
 
 LOCAL_INCLUDES = -I./ -Iinclude
 OPENCV_INCLUDE = -I${THIRD_PARTY_DIR}/${OPENCV_VERSION}/${ARM_PLAT}/include
+FAISS_INCLUDE = -I${THIRD_PARTY_DIR}/${FAISS_VERSION}/include
 JSON_INCLUDE = -I${THIRD_PARTY_DIR}/jsoncpp_code/include
-CXX_INCLUDES = ${LOCAL_INCLUDES} ${INCLUDES} ${OPENCV_INCLUDE} ${JSON_INCLUDE} -I$(LITE_ROOT)/cxx/include
+CXX_INCLUDES = ${LOCAL_INCLUDES} ${INCLUDES} ${OPENCV_INCLUDE} ${FAISS_INCLUDE} ${JSON_INCLUDE} -I$(LITE_ROOT)/cxx/include
 
 
-$(LOCAL_OBJS): %.o: src/%.cc fetch_opencv fetch_json_code
+$(LOCAL_OBJS): %.o: src/%.cc fetch_opencv fetch_json_code fetch_faiss
 	$(CC) $(SYSROOT_COMPLILE) $(CXX_DEFINES) $(CXX_INCLUDES) $(CXX_FLAGS) -c $< -o $@
 
 $(JSON_OBJS): %.o: ${THIRD_PARTY_DIR}/jsoncpp_code/%.cpp fetch_json_code

deploy/lite_shitu/README.md

@@ -92,9 +92,9 @@ PaddleClas 提供了转换并优化后的推理模型，可以直接参考下方
 ```shell
 # 进入lite_ppshitu目录
 cd $PaddleClas/deploy/lite_shitu
-wget https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/lite/ppshitu_lite_models_v1.0.tar
-tar -xf ppshitu_lite_models_v1.0.tar
-rm -f ppshitu_lite_models_v1.0.tar
+wget https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/lite/ppshitu_lite_models_v1.0.tar.gz
+tar -xf ppshitu_lite_models_v1.0.tar.gz
+rm -f ppshitu_lite_models_v1.0.tar.gz
 ```
 
 #### 2.1.2 使用其他模型
@@ -191,16 +191,28 @@ cd deploy/lite_shitu
 
 ```shell
 # 如果测试单张图像
-python generate_json_config.py --det_model_path ppshitu_lite_models_v1.0/mainbody_PPLCNet_x2_5_640_quant_v1.0_lite.nb  --rec_model_path ppshitu_lite_models_v1.0/general_PPLCNet_x2_5_quant_v1.0_lite.nb --rec_label_path ppshitu_lite_models_v1.0/label.txt --img_path images/demo.jpg
+python generate_json_config.py --det_model_path ppshitu_lite_models_v1.0/mainbody_PPLCNet_x2_5_640_quant_v1.0_lite.nb  --rec_model_path ppshitu_lite_models_v1.0/general_PPLCNet_x2_5_lite_v1.0_infer.nb --img_path images/demo.jpg
 # or
 # 如果测试多张图像
-python generate_json_config.py --det_model_path ppshitu_lite_models_v1.0/mainbody_PPLCNet_x2_5_640_quant_v1.0_lite.nb  --rec_model_path ppshitu_lite_models_v1.0/general_PPLCNet_x2_5_quant_v1.0_lite.nb --rec_label_path ppshitu_lite_models_v1.0/label.txt --img_dir images
-
+python generate_json_config.py --det_model_path ppshitu_lite_models_v1.0/mainbody_PPLCNet_x2_5_640_quant_v1.0_lite.nb  --rec_model_path ppshitu_lite_models_v1.0/general_PPLCNet_x2_5_lite_v1.0_infer.nb --img_dir images
 # 执行完成后，会在lit_shitu下生成shitu_config.json配置文件
+```
+
+### 2.3 index字典转换
+由于python的检索库字典，使用`pickle`进行的序列化存储，导致C++不方便读取，因此需要进行转换
+
+```shell
+# 下载瓶装饮料数据集
+wget https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/rec/data/drink_dataset_v1.0.tar && tar -xf drink_dataset_v1.0.tar
+rm -rf drink_dataset_v1.0.tar
 
+# 转化id_map.pkl为id_map.txt
+python transform_id_map.py -c ../configs/inference_drink.yaml
 ```
+转换成功后，会在`IndexProcess.index_dir`目录下生成`id_map.txt`。
+
 
-### 2.3 与手机联调
+### 2.4 与手机联调
 
 首先需要进行一些准备工作。
 1. 准备一台arm8的安卓手机，如果编译的预测库是armv7，则需要arm7的手机，并修改Makefile中`ARM_ABI=arm7`。
@@ -252,6 +264,7 @@ make ARM_ABI=arm8
 ```shell
 mkdir deploy
 mv ppshitu_lite_models_v1.0 deploy/
+mv drink_dataset_v1.0 deploy/
 mv images deploy/
 mv shitu_config.json deploy/
 cp pp_shitu deploy/
@@ -265,12 +278,12 @@ cp ../../../cxx/lib/libpaddle_light_api_shared.so deploy/
 ```shell
 deploy/
 |-- ppshitu_lite_models_v1.0/
-|   |--mainbody_PPLCNet_x2_5_640_v1.0_lite.nb        优化后的主体检测模型文件
-|   |--general_PPLCNet_x2_5_quant_v1.0_lite.nb       优化后的识别模型文件
-|   |--label.txt                                     识别模型的label文件
+|   |--mainbody_PPLCNet_x2_5_lite_v1.0_infer.nb        优化后的主体检测模型文件
+|   |--general_PPLCNet_x2_5_quant_v1.0_lite.nb         优化后的识别模型文件
 |-- images/
 |   |--demo.jpg                                      图片文件
-|   ...                                              图片文件
+|-- drink_dataset_v1/                                瓶装饮料demo数据
+|   |--index                                         检索index目录
 |-- pp_shitu                                         生成的移动端执行文件
 |-- shitu_config.json                                执行时参数配置文件
 |-- libpaddle_light_api_shared.so                    Paddle-Lite库文件
@@ -298,8 +311,8 @@ chmod 777 pp_shitu
 如果对代码做了修改，则需要重新编译并push到手机上。
 
 运行效果如下：
-
-![](../../docs/images/ppshitu_lite_demo.png)
+```
+````
 
 ## FAQ
 Q1：如果想更换模型怎么办，需要重新按照流程走一遍吗？

deploy/lite_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 "json/json.h"
+#include <cstring>
+#include <faiss/Index.h>
+#include <faiss/index_io.h>
+#include <map>
+
+namespace PPShiTu {
+struct SearchResult {
+  std::vector<faiss::Index::idx_t> I;
+  std::vector<float> D;
+  int return_k;
+};
+
+class VectorSearch {
+public:
+  explicit VectorSearch(const Json::Value &config) {
+    // IndexProcess
+    this->index_dir = config["IndexProcess"]["index_dir"].as<std::string>();
+    this->return_k = config["IndexProcess"]["return_k"].as<int>();
+    this->score_thres = config["IndexProcess"]["score_thres"].as<float>();
+    this->max_query_number = config["Global"]["max_det_results"].as<int>() + 1;
+
+    LoadIdMap();
+    LoadIndexFile();
+    this->I.resize(this->return_k * this->max_query_number);
+    this->D.resize(this->return_k * this->max_query_number);
+  };
+
+  void LoadIdMap();
+
+  void LoadIndexFile();
+
+  const SearchResult &Search(float *feature, int query_number);
+
+  const std::string &GetLabel(faiss::Index::idx_t ind);
+
+  const float &GetThreshold() { return this->score_thres; }
+
+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;
+  std::vector<float> D;
+  std::vector<faiss::Index::idx_t> I;
+  SearchResult sr;
+};
+}

deploy/lite_shitu/src/main.cc

@@ -24,9 +24,10 @@
 #include <vector>
 
 #include "include/config_parser.h"
+#include "include/feature_extractor.h"
 #include "include/object_detector.h"
 #include "include/preprocess_op.h"
-#include "include/feature_extractor.h"
+#include "include/vector_search.h"
 #include "json/json.h"
 
 Json::Value RT_Config;
@@ -111,14 +112,18 @@ void DetPredictImage(const std::vector<cv::Mat> &batch_imgs,
   }
 }
 
-void PrintResult(const std::string &image_path,
-                 std::vector<PPShiTu::ObjectResult> &det_result) {
-  printf("%s:\n", image_path.c_str());
+void PrintResult(std::string &img_path,
+                 std::vector<PPShiTu::ObjectResult> &det_result,
+                 PPShiTu::VectorSearch &vector_search,
+                 PPShiTu::SearchResult &search_result) {
+  printf("%s:\n", img_path.c_str());
   for (int i = 0; i < det_result.size(); ++i) {
+    int t = i;
     printf("\tresult%d: bbox[%d, %d, %d, %d], score: %f, label: %s\n", i,
-           det_result[i].rect[0], det_result[i].rect[1], det_result[i].rect[2],
-           det_result[i].rect[3], det_result[i].rec_result[0].score,
-           det_result[i].rec_result[0].class_name.c_str());
+           det_result[t].rect[0], det_result[t].rect[1], det_result[t].rect[2],
+           det_result[t].rect[3], det_result[t].confidence,
+           vector_search.GetLabel(search_result.I[search_result.return_k * t])
+               .c_str());
   }
 }
 
@@ -160,12 +165,13 @@ int main(int argc, char **argv) {
       RT_Config["Global"]["batch_size"].as<int>());
   // create rec model
   PPShiTu::FeatureExtract rec(RT_Config);
+  PPShiTu::VectorSearch searcher(RT_Config);
   // Do inference on input image
 
   std::vector<PPShiTu::ObjectResult> det_result;
   std::vector<cv::Mat> batch_imgs;
 
-  //for vector search
+  // for vector search
   std::vector<float> feature;
   std::vector<float> features;
   double rec_time;
@@ -182,8 +188,7 @@ int main(int argc, char **argv) {
         return -1;
       }
     } else {
-      cv::glob(img_dir,
-               cv_all_img_paths);
+      cv::glob(img_dir, cv_all_img_paths);
       for (const auto &img_path : cv_all_img_paths) {
         all_img_paths.push_back(img_path);
       }
@@ -208,6 +213,7 @@ int main(int argc, char **argv) {
       det_result.push_back(result_whole_img);
 
       // get rec result
+      PPShiTu::SearchResult search_result;
       for (int j = 0; j < det_result.size(); ++j) {
         int w = det_result[j].rect[2] - det_result[j].rect[0];
         int h = det_result[j].rect[3] - det_result[j].rect[1];
@@ -217,11 +223,10 @@ int main(int argc, char **argv) {
         features.insert(features.end(), feature.begin(), feature.end());
       }
 
-      std::cout << "feature len is:  " << features.size() << std::endl;
-      // rec nms
-      // PPShiTu::nms(det_result,
-      //              RT_Config["Global"]["rec_nms_thresold"].as<float>(), true);
-      // PrintResult(img_path, det_result);
+      // do vectore search
+      search_result = searcher.Search(features.data(), det_result.size());
+      PrintResult(img_path, det_result, searcher, search_result);
+
       batch_imgs.clear();
       det_result.clear();
     }

deploy/lite_shitu/src/vector_search.cc

+// 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 "include/vector_search.h"
+#include <cstdio>
+#include <faiss/index_io.h>
+#include <fstream>
+#include <iostream>
+#include <regex>
+
+namespace PPShiTu {
+// load the vector.index
+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);
+}
+
+// load id_map.txt
+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]));
+    }
+  }
+}
+
+// doing search
+const SearchResult &VectorSearch::Search(float *feature, int query_number) {
+  this->D.resize(this->return_k * query_number);
+  this->I.resize(this->return_k * query_number);
+  this->index->search(query_number, feature, return_k, D.data(), I.data());
+  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);
+}
+}
\ No newline at end of file

deploy/lite_shitu/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.read(), yaml.FullLoader)
+    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()