def do_search(table_name, molecular_name, metric, top_k):
try:
feats = []
index_client = milvus_client()
feat = smiles_to_vec(molecular_name)
feats.append(feat)
# print(feats)
time1 = time.time()
vectors = search_vectors(index_client, table_name, feats, metric,
top_k)
time2 = time.time()
print("milvus search:", time2 - time1)
vids = [x.id for x in vectors[0]]
# print("-----------------", vids)
conn = connect_postgres_server()
cur = conn.cursor()
res_smi = []
for i in vids:
# index = search_loc_in_pg(cur, i, table_name)
if i < 0:
break
index = search_loc_in_pg(cur, i)
res_smi.append(index)
# print(res_smi)
return res_smi
except Exception as e:
logging.error(e)
return "Fail with error {}".format(e)
finally:
if conn:
cur.close()
conn.close()
def op_search(table_name,
img_path,
top_k,
model,
graph,
sess,
mycol,
partition=None):
try:
feats = []
client = milvus_client()
if partition:
status, ok = has_partition(client, table_name, partition)
if not ok:
return False, ''
feat = vgg_extract_feat(img_path, model, graph, sess)
feats.append(feat)
_, vectors = search_vectors(client, table_name, feats, top_k,
partition)
if len(vectors) > 0:
vids = [x.id for x in vectors[0]]
res = filter_data(vids, mycol, vectors[0])
return True, res
else:
return True, vectors
except Exception as e:
logging.error(e)
return '发生错误'.format(e)
def do_search(image_encoder, table_name, img_path, top_k):
try:
#print(top_k)
detector = Detector()
run(detector, img_path)
vect, obj_images = get_object_vector(image_encoder,
img_path + '/object')
#print("search...after detect:", len(vect), obj_images)
index_client = milvus_client()
#vect = normaliz_vec(vect)
status, results = search_vectors(index_client, table_name, vect, top_k)
# print(status, results)
vids = []
dis = []
for result in results:
for j in result:
vids.append(j.id)
dis.append(j.distance)
res_id = [x for x in query_name_from_ids(vids)]
#print("------------------res", vids, dis, res_id)
shutil.rmtree(img_path + '/object')
return res_id, dis
except Exception as e:
logging.error(e)
return "Fail with error {}".format(e)
def do_search(table_name, molecular_name, top_k):
try:
feats = []
index_client = milvus_client()
feat = smiles_to_vec(molecular_name)
feats.append(feat)
status, vectors = search_vectors(index_client, table_name, feats,
top_k)
print(status)
vids = [x.id for x in vectors[0]]
conn = connect_postgres_server()
cur = conn.cursor()
res_smi = []
for i in vids:
# index = search_loc_in_pg(cur, i, table_name)
if i < 0:
break
index = search_loc_in_pg(cur, i)
res_smi.append(index)
return res_smi
except Exception as e:
logging.error(e)
return "Fail with error {}".format(e)
finally:
if index_client:
index_client.disconnect()
if conn:
cur.close()
conn.close()
def do_search_face(face_encoder, index_client, conn, cursor, table_name, filename):
if not table_name:
table_name = FACE_TABLE
faces = face_encoder.execute(filename)
results = search_vectors(index_client, table_name, faces, "L2")
info = get_face_info(conn, cursor, table_name, results)
return info
def do_only_her(face_encoder, index_client, conn, cursor, table_name, filename):
if not table_name:
table_name = FACE_TABLE
prefix = filename.split("/")[2].split(".")[0] + "-" + uuid.uuid4().hex
images = extract_frame(filename, 1, prefix)
vectors, face_images = get_face_vector(face_encoder, DATA_PATH + '/' + prefix)
global_info = []
for faces in vectors:
results = search_vectors(index_client, table_name, faces, "L2")
info = get_face_info(conn, cursor, table_name, results)
global_info.append(info)
return global_info
def do_search_logo(image_encoder, index_client, conn, cursor, table_name, filename):
detector = Detector()
if not table_name:
table_name = LOGO_TABLE
prefix = filename.split("/")[2].split(".")[0] + "-" + uuid.uuid4().hex
images = extract_frame(filename, 1, prefix)
run(detector, DATA_PATH + '/' + prefix)
vectors, obj_images = get_object_vector(image_encoder, DATA_PATH + '/' + prefix + '/object')
results = search_vectors(index_client, table_name, vectors, "L2")
info, times = get_object_info(conn, cursor, table_name, results, obj_images)
return info, times
def do_search(img, voice):
try:
conn = None
feats_img = img_to_vec(img)
feats_voc = voc_to_vec(voice)
# print(feats_voc,feats_img)
res = ['false', -1, '-1']
index_client = milvus_client()
_, re_img = search_vectors(index_client, IMG_TABLE, [feats_img], 1)
_, re_voc = search_vectors(index_client, VOC_TABLE, [feats_voc], 1)
ids_img = re_img[0][0].id
ids_voc = re_voc[0][0].id
dis_img = float(re_img[0][0].distance)
dis_voc = float(re_voc[0][0].distance)
print(ids_img, ids_voc, dis_img)
if dis_img > 0.75 and dis_voc > 0.65 and ids_img == ids_voc:
conn = connect_postgres_server()
cur = conn.cursor()
print(ids_voc)
index = search_loc_in_pg(cur, ids_voc)
res = ['true', ids_img, index]
print("-----search:", res)
return res
except Exception as e:
logging.error(e)
print("Fail with error {}".format(e))
return "Fail with error {}".format(
"please confirm only one face in camera")
finally:
if conn:
cur.close()
conn.close()
def do_search(index_client, conn, cursor, img_to_vec, img_list, table_name):
if not table_name:
table_name = DEFAULT_TABLE
vectors_img = img_to_vec(img_list)
print("the num of search images:", len(vectors_img))
print("doing search, table_name:", table_name)
status, ids_milvus = search_vectors(index_client, table_name, vectors_img)
# print("distance_array:", ids_milvus.distance_array)
re_ids_img = []
for ids in ids_milvus:
vids = [x.id for x in ids]
re_ids = search_by_milvus_ids(conn, cursor, vids, table_name)
re_ids_img.append(re_ids)
return re_ids_img, ids_milvus.distance_array
def do_search(index_client, conn, cursor, Text, Image):
#_, vector_item = get_vector_by_ids(index_client, Text, Image)
status, results = search_vectors(index_client, Text, vector_item)
print("-----milvus search status------", status)
results_ids = []
for results_id in results.id_array:
for i in results_id:
img = str(i) + '.jpg'
if img in img_list and i not in search_id:
results_ids.append(img)
# print(results_ids)
try:
list_ids = random.sample(results_ids, 10)
except:
list_ids = results_ids
return list_ids
def do_search(index_client, conn, cursor, img_list, search_id, table_name):
if not table_name:
table_name = MILVUS_TABLE
_, vector_item = get_vector_by_ids(index_client, table_name, search_id)
status, results = search_vectors(index_client, table_name, vector_item)
print("-----milvus search status------", status)
results_ids = []
for results_id in results.id_array:
for i in results_id:
img = str(i) + '.jpg'
if img in img_list and i not in search_id:
results_ids.append(img)
# print(results_ids)
try:
list_ids = random.sample(results_ids, 100)
except:
list_ids = results_ids
return list_ids
def do_search(table_name, molecular_name, top_k):
try:
feats = []
index_client = milvus_client()
feat = smiles_to_vec(molecular_name)
feats.append(feat)
_, vectors = search_vectors(index_client, table_name, feats, top_k)
vids = [x.id for x in vectors[0]]
# print(vids)
res_smi = [x.decode('utf-8') for x in query_smi_from_ids(vids)]
# print("vids:",vids)
res_distance = [x.distance for x in vectors[0]]
res_ids = [x.id for x in vectors[0]]
# print(res_distance,res_smi)
return res_smi, res_distance, res_ids
except Exception as e:
logging.error(e)
return "Fail with error {}".format(e)
def do_search(table_name, caption, topk, model, args):
try:
index_client = milvus_client()
feat = extract_caption_feat(caption, model, args)
feat = (feat / feat.norm()).tolist()
_, vectors = search_vectors(index_client, table_name, [feat], 2 * topk)
print('cap_vec:', vectors)
vids = [x.id for x in vectors[0]]
# print(vids)
# res = [x.decode('utf-8') for x in query_name_from_ids(vids)]
res_id = [x.decode('utf-8') for x in query_name_from_ids(vids)]
print(res_id)
res_distance = [x.distance for x in vectors[0]]
# print(res_distance)
# res = dict(zip(res_id,distance))
return res_id, res_distance
except Exception as e:
logging.error(e)
return "Fail with error {}".format(e)
def do_search(table_name, img_path, top_k, model, graph, sess):
try:
feats = []
index_client = milvus_client()
feat = extract_feat(img_path, model, graph, sess)
feats.append(feat)
_, vectors = search_vectors(index_client, table_name, feats, top_k)
vids = [x.id for x in vectors[0]]
# print(vids)
# res = [x.decode('utf-8') for x in query_name_from_ids(vids)]
res_id = [x.decode('utf-8') for x in query_name_from_ids(vids)]
# print(res_id)
res_distance = [x.distance for x in vectors[0]]
# print(res_distance)
# res = dict(zip(res_id,distance))
return res_id, res_distance
except Exception as e:
logging.error(e)
return "Fail with error {}".format(e)
def do_search(index_client, conn, cursor, img_list, search_id, table_name):
if not table_name:
table_name = MILVUS_TABLE
_, vector_item = get_vector_by_ids(index_client, table_name, search_id)
status, results = search_vectors(index_client, table_name, vector_item)
print("-----milvus search status------", status)
results_ids = []
search_num = len(search_id)
num = 100 / search_num
print("-----num:", num)
for results_id in results.id_array:
k = 0
for i in results_id:
if k >= num:
break
img = str(i) + '.jpg'
if img in img_list and i not in search_id:
results_ids.append(img)
k += 1
# print(results_ids)
return results_ids
def do_search(index_client, conn, cursor, Text, Image, table_name, img_list,
host):
if not table_name:
table_name = MILVUS_TABLE
def parse_opt():
"""Parses the input arguments."""
parser = argparse.ArgumentParser()
parser.add_argument('-f', type=str, default='')
parser.add_argument('--comment', type=str, default='test_notebook')
parser.add_argument('--dataset', type=str, default='css3d')
parser.add_argument('--dataset_path', type=str, default='./css')
parser.add_argument('--model', type=str, default='tirg')
parser.add_argument('--embed_dim', type=int, default=512)
parser.add_argument('--learning_rate', type=float, default=1e-2)
parser.add_argument('--learning_rate_decay_frequency',
type=int,
default=9999999)
parser.add_argument('--batch_size', type=int, default=32)
parser.add_argument('--weight_decay', type=float, default=1e-6)
parser.add_argument('--num_iters', type=int, default=10000)
parser.add_argument('--loss', type=str, default='soft_triplet')
parser.add_argument('--loader_num_workers', type=int, default=4)
args = parser.parse_args()
return args
one_queries = []
opt = parse_opt()
opt.model == 'tirg'
opt = parse_opt()
transform = torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
])
def normaliz_vec(vec_list):
for i in range(len(vec_list)):
vec = vec_list[i]
square_sum = reduce(lambda x, y: x + y, map(lambda x: x * x, vec))
sqrt_square_sum = np.sqrt(square_sum)
coef = 1 / sqrt_square_sum
vec = list(map(lambda x: x * coef, vec))
vec_list[i] = vec
return vec_list
texts = [Text]
img_path = Image
img = PIL.Image.open(img_path)
img = img.convert('RGB')
imgs = transform(img)
imgs = Variable(torch.unsqueeze(imgs, dim=0).float(), requires_grad=False)
model = img_text_composition_models.TIRG(texts, embed_dim=opt.embed_dim)
pthfile = '/data1/workspace/jingjing/ctirg/tirg/webserver/src/tirg/runs/Aug/latest_checkpoint.pth'
model.load_state_dict(torch.load(pthfile), strict=False)
model.eval()
model.cuda()
f = model.compose_img_text(imgs.cuda(), texts).data.cpu().numpy()
one_queries += [f]
one_queries = np.concatenate(one_queries)
# for i in range(one_queries.shape[0]):
# one_queries[i, :] /= np.linalg.norm(one_queries[i, :])
one_queries = normaliz_vec(one_queries.tolist())
status, results = search_vectors(index_client, table_name, one_queries)
print("-----milvus search status------", status, results)
results_ids = []
for results_id in results.id_array:
for i in results_id:
img = 'css_test_' + str(i).rjust(6, '0') + '.png'
if img in img_list:
res = "http://" + str(host) + "/getImage?img="
results_ids.append(res + img)
list_ids = results_ids
return list_ids
