import faiss
|
import cv2
|
import time
|
import sys
|
import numpy as np
|
import pandas as pd
|
from flask import Flask
|
from flask import request
|
from flask import render_template
|
|
### search for vehicle or face
|
if len(sys.argv) != 2:
|
print("search for vehicle or face?")
|
print("'python3 search_server.py vehicle' OR 'python3 search_server.py face'")
|
quit()
|
search_for_vehicle = True if sys.argv[1] == "vehicle" else False
|
###
|
|
|
'''
|
create Index Object of faiss
|
'''
|
# load csv from disk
|
embedding_df = pd.read_csv("embeddings.csv", sep='|', index_col=0)
|
print(f"load csv snapshot, contains {len(embedding_df)} embeddings!")
|
print(embedding_df)
|
|
# create Index Object
|
embeddings_list = [[e for e in str(embedding).split(',')] for embedding in embedding_df["embedding"].to_list()]
|
embeddings = np.array(embeddings_list)
|
print(f"load embeddings successfully, shape is {embeddings.shape}")
|
dimension = embeddings.shape[1]
|
# you can try different Index type
|
search_index = faiss.IndexFlatIP(dimension)
|
# add all embeddings into Index, use it in search action later
|
search_index.add(embeddings)
|
|
|
'''
|
generate embedding data for query image
|
'''
|
# vehicle embedding model keep as same as vehicle_encoding_pipeline.cpp
|
vehicle_embedding_model_path = "../../../vp_data/models/trt/vehicle/dim256_batch=1_0916_norm.onnx"
|
# face embedding model keep as same as face_encoding_pipeline.cpp
|
face_embedding_model_path = "../../../vp_data/models/face/face_recognition_sface_2021dec.onnx"
|
|
embedding_model_path = vehicle_embedding_model_path if search_for_vehicle else face_embedding_model_path
|
input_size = (256, 256) if search_for_vehicle else (112, 112)
|
net = cv2.dnn.readNetFromONNX(embedding_model_path)
|
output_layer_names = net.getUnconnectedOutLayersNames()
|
def generate_embedding(query_image):
|
input_image = cv2.resize(query_image, input_size)
|
blob = cv2.dnn.blobFromImage(input_image)
|
net.setInput(blob)
|
out = net.forward(output_layer_names)
|
embedding = np.array(out[0])
|
return embedding
|
|
|
'''
|
create flask app
|
'''
|
app = Flask(__name__)
|
|
# home page
|
@app.route("/")
|
def home():
|
search_for = 'vehicles' if search_for_vehicle else 'faces'
|
return render_template("home.html", total_images = search_index.ntotal, search_for = search_for)
|
|
# search action
|
@app.route("/search", methods=["POST"])
|
def search():
|
# no exception catch
|
query_image = request.files['query_image']
|
topK = int(request.form['topK'])
|
|
# convert to cv2.mat and save to disk
|
img_array = np.asarray(bytearray(query_image.read()), dtype=np.uint8)
|
query_mat = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
|
query_image_path = "static/query_images/" + str(time.time()) + ".jpg"
|
cv2.imwrite(query_image_path, query_mat)
|
|
# generate embedding for query image and do search action
|
query_embedding = generate_embedding(query_mat)
|
D, I = search_index.search(query_embedding, topK)
|
|
# print search info
|
print("query embedding:")
|
print(query_embedding)
|
print("similiarity distance:")
|
print(D)
|
print("similiarity uid:")
|
print(I)
|
|
# shoot data
|
shoot_df = embedding_df.iloc[I[0]]
|
shoot_path = shoot_df['path']
|
shoot_distance = D[0]
|
shoot_uid = I[0]
|
|
# render search results
|
search_results = []
|
for i,p,d in zip(shoot_uid, shoot_path, shoot_distance):
|
search_results.append({"uid": i, "image_url": p, "distance":d})
|
return render_template("search_result.html", topK=topK, search_results=search_results, query_image_url=query_image_path)
|
|
|
if __name__ == "__main__":
|
app.run(host="192.168.77.68", port=9999)
|
