# 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. import os import sys from PyQt5 import QtCore, QtGui, QtWidgets import mod.mainwindow from paddleclas.deploy.utils import config, logger from paddleclas.deploy.python.predict_rec import RecPredictor from fastapi import FastAPI import uvicorn import numpy as np import faiss from typing import List import pickle import cv2 import socket import json import operator from multiprocessing import Process """ 完整的index库如下: root_path/ # 库存储目录 |-- image_list.txt # 图像列表，每行：image_path label。由前端生成及修改。后端只读 |-- features.pkl # 建库之后，保存的embedding向量，后端生成，前端无需操作 |-- images # 图像存储目录，由前端生成及增删查等操作。后端只读 | |-- md5.jpg | |-- md5.jpg | |-- …… |-- index # 真正的生成的index库存储目录，后端生成及操作，前端无需操作。 | |-- vector.index # faiss生成的索引库 | |-- id_map.pkl # 索引文件 """ class ShiTuIndexManager(object): def __init__(self, config): self.root_path = None self.image_list_path = "image_list.txt" self.image_dir = "images" self.index_path = "index/vector.index" self.id_map_path = "index/id_map.pkl" self.features_path = "features.pkl" self.index = None self.id_map = None self.features = None self.config = config self.predictor = RecPredictor(config) def _load_pickle(self, path): if os.path.exists(path): return pickle.load(open(path, 'rb')) else: return None def _save_pickle(self, path, data): if not os.path.exists(os.path.dirname(path)): os.makedirs(os.path.dirname(path), exist_ok=True) with open(path, 'wb') as fd: pickle.dump(data, fd) def _load_index(self): self.index = faiss.read_index( os.path.join(self.root_path, self.index_path)) self.id_map = self._load_pickle( os.path.join(self.root_path, self.id_map_path)) self.features = self._load_pickle( os.path.join(self.root_path, self.features_path)) def _save_index(self, index, id_map, features): faiss.write_index(index, os.path.join(self.root_path, self.index_path)) self._save_pickle(os.path.join(self.root_path, self.id_map_path), id_map) self._save_pickle(os.path.join(self.root_path, self.features_path), features) def _update_path(self, root_path, image_list_path=None): if root_path == self.root_path: pass else: self.root_path = root_path if not os.path.exists(os.path.join(root_path, "index")): os.mkdir(os.path.join(root_path, "index")) if image_list_path is not None: self.image_list_path = image_list_path def _cal_featrue(self, image_list): batch_images = [] featrures = None cnt = 0 for idx, image_path in enumerate(image_list): image = cv2.imread(image_path) if image is None: return "{} is broken or not exist. Stop" else: image = image[:, :, ::-1] batch_images.append(image) cnt += 1 if cnt % self.config["Global"]["batch_size"] == 0 or ( idx + 1) == len(image_list): if len(batch_images) == 0: continue batch_results = self.predictor.predict(batch_images) featrures = batch_results if featrures is None else np.concatenate( (featrures, batch_results), axis=0) batch_images = [] return featrures def _split_datafile(self, data_file, image_root): ''' data_file: image path and info, which can be splitted by spacer image_root: image path root delimiter: delimiter ''' gallery_images = [] gallery_docs = [] gallery_ids = [] with open(data_file, 'r', encoding='utf-8') as f: lines = f.readlines() for _, ori_line in enumerate(lines): line = ori_line.strip().split() text_num = len(line) assert text_num >= 2, f"line({ori_line}) must be splitted into at least 2 parts, but got {text_num}" image_file = os.path.join(image_root, line[0]) gallery_images.append(image_file) gallery_docs.append(ori_line.strip()) gallery_ids.append(os.path.basename(line[0]).split(".")[0]) return gallery_images, gallery_docs, gallery_ids def create_index(self, image_list: str, index_method: str = "HNSW32", image_root: str = None): if not os.path.exists(image_list): return "{} is not exist".format(image_list) if index_method.lower() not in ['hnsw32', 'ivf', 'flat']: return "The index method Only support: HNSW32, IVF, Flat" self._update_path(os.path.dirname(image_list), image_list) # get image_paths image_root = image_root if image_root is not None else self.root_path gallery_images, gallery_docs, image_ids = self._split_datafile( image_list, image_root) # gernerate index if index_method == "IVF": index_method = index_method + str( min(max(int(len(gallery_images) // 32), 2), 65536)) + ",Flat" index = faiss.index_factory( self.config["IndexProcess"]["embedding_size"], index_method, faiss.METRIC_INNER_PRODUCT) self.index = faiss.IndexIDMap2(index) features = self._cal_featrue(gallery_images) self.index.train(features) index_ids = np.arange(0, len(gallery_images)).astype(np.int64) self.index.add_with_ids(features, index_ids) self.id_map = dict() for i, d in zip(list(index_ids), gallery_docs): self.id_map[i] = d self.features = { "features": features, "index_method": index_method, "image_ids": image_ids, "index_ids": index_ids.tolist() } self._save_index(self.index, self.id_map, self.features) def open_index(self, root_path: str, image_list_path: str) -> str: self._update_path(root_path) _, _, image_ids = self._split_datafile(image_list_path, root_path) if os.path.exists(os.path.join(self.root_path, self.index_path)) and \ os.path.exists(os.path.join(self.root_path, self.id_map_path)) and \ os.path.exists(os.path.join(self.root_path, self.features_path)): self._update_path(root_path) self._load_index() if operator.eq(set(image_ids), set(self.features['image_ids'])): return "" else: return "The image list is different from index, Please update index" else: return "File not exist: features.pkl, vector.index, id_map.pkl" def update_index(self, image_list: str, image_root: str = None) -> str: if self.index and self.id_map and self.features: image_paths, image_docs, image_ids = self._split_datafile( image_list, image_root if image_root is not None else self.root_path) # for add image add_ids = list( set(image_ids).difference(set(self.features["image_ids"]))) add_indexes = [i for i, x in enumerate(image_ids) if x in add_ids] add_image_paths = [image_paths[i] for i in add_indexes] add_image_docs = [image_docs[i] for i in add_indexes] add_image_ids = [image_ids[i] for i in add_indexes] self._add_index(add_image_paths, add_image_docs, add_image_ids) # delete images delete_ids = list( set(self.features["image_ids"]).difference(set(image_ids))) self._delete_index(delete_ids) self._save_index(self.index, self.id_map, self.features) return "" else: return "Failed. Please create or open index first" def _add_index(self, image_list: List, image_docs: List, image_ids: List): if len(image_ids) == 0: return featrures = self._cal_featrue(image_list) index_ids = (np.arange(0, len(image_list)) + max(self.id_map.keys()) + 1).astype(np.int64) self.index.add_with_ids(featrures, index_ids) for i, d in zip(index_ids, image_docs): self.id_map[i] = d self.features['features'] = np.concatenate( [self.features['features'], featrures], axis=0) self.features['image_ids'].extend(image_ids) self.features['index_ids'].extend(index_ids.tolist()) def _delete_index(self, image_ids: List): if len(image_ids) == 0: return indexes = [ i for i, x in enumerate(self.features['image_ids']) if x in image_ids ] self.features["features"] = np.delete(self.features["features"], indexes, axis=0) self.features["image_ids"] = np.delete(np.asarray( self.features["image_ids"]), indexes, axis=0).tolist() index_ids = np.delete(np.asarray(self.features["index_ids"]), indexes, axis=0).tolist() id_map_values = [self.id_map[i] for i in index_ids] self.index.reset() ids = np.arange(0, len(id_map_values)).astype(np.int64) self.index.add_with_ids(self.features['features'], ids) self.id_map.clear() for i, d in zip(ids, id_map_values): self.id_map[i] = d self.features["index_ids"] = ids app = FastAPI() @app.get("/new_index") def new_index(image_list_path: str, index_method: str = "HNSW32", index_root_path: str = None, force: bool = False): result = "" try: if index_root_path is not None: image_list_path = os.path.join(index_root_path, image_list_path) index_path = os.path.join(index_root_path, "index", "vector.index") id_map_path = os.path.join(index_root_path, "index", "id_map.pkl") if not (os.path.exists(index_path) and os.path.exists(id_map_path)) or force: manager.create_index(image_list_path, index_method, index_root_path) else: result = "There alrealy has index in {}".format(index_root_path) except Exception as e: result = e.__str__() data = {"error_message": result} return json.dumps(data).encode() @app.get("/open_index") def open_index(index_root_path: str, image_list_path: str): result = "" try: image_list_path = os.path.join(index_root_path, image_list_path) result = manager.open_index(index_root_path, image_list_path) except Exception as e: result = e.__str__() data = {"error_message": result} return json.dumps(data).encode() @app.get("/update_index") def update_index(image_list_path: str, index_root_path: str = None): result = "" try: if index_root_path is not None: image_list_path = os.path.join(index_root_path, image_list_path) result = manager.update_index(image_list=image_list_path, image_root=index_root_path) except Exception as e: result = e.__str__() data = {"error_message": result} return json.dumps(data).encode() def FrontInterface(server_process=None): front = QtWidgets.QApplication([]) main_window = mod.mainwindow.MainWindow(process=server_process) main_window.showMaximized() sys.exit(front.exec_()) def Server(args): [app, host, port] = args uvicorn.run(app, host=host, port=port) if __name__ == '__main__': args = config.parse_args() model_config = config.get_config(args.config, overrides=args.override, show=True) manager = ShiTuIndexManager(model_config) try: ip = socket.gethostbyname(socket.gethostname()) except: ip = '127.0.0.1' port = 8000 p_server = Process(target=Server, args=([app, ip, port],)) p_server.start() # p_client = Process(target=FrontInterface, args=()) # p_client.start() # p_client.join() FrontInterface(p_server) p_server.terminate() sys.exit(0)