def do_load(table_name, database_path):
if not table_name:
table_name = DEFAULT_TABLE
cache = Cache(default_cache_dir)
try:
vectors, names = feature_extract(table_name, database_path)
print("start connetc to milvus")
index_client = milvus_client()
status, ok = has_table(index_client, table_name)
if not ok:
print("create table.")
create_table(index_client, table_name=table_name)
print("insert into:", table_name)
# status, ids = insert_vectors(index_client, table_name, vectors)
total_ids = []
ids_lens = 0
while ids_lens<len(vectors) :
try:
status, ids = insert_vectors(index_client, table_name, vectors[ids_lens:ids_lens+100000])
except:
status, ids = insert_vectors(index_client, table_name, vectors[ids_lens:len(vectors)])
ids_lens += 100000
total_ids += ids
print("ids:",len(ids))
create_index(index_client, table_name)
for i in range(len(names)):
cache[total_ids[i]] = names[i]
print("FP finished")
return "FP finished"
except Exception as e:
logging.error(e)
return "Error with {}".format(e)
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 do_train(table_name, database_path):
if not table_name:
table_name = DEFAULT_TABLE
cache = Cache(default_cache_dir)
try:
vectors, names = feature_extract(database_path, VGGNet())
print("start connetc to milvus")
index_client = milvus_client()
# delete_table(index_client, table_name=table_name)
# time.sleep(1)
status, ok = has_table(index_client, table_name)
if not ok:
print("create table.")
create_table(index_client, table_name=table_name)
print("insert into:", table_name)
status, ids = insert_vectors(index_client, table_name, vectors)
create_index(index_client, table_name)
for i in range(len(names)):
# cache[names[i]] = ids[i]
cache[ids[i]] = names[i]
print("Train finished")
return "Train finished"
except Exception as e:
logging.error(e)
return "Error with {}".format(e)
def delete_imgs(table_name, ids, mycol):
print(table_name, ids, '内部')
if not table_name:
return False
try:
client = milvus_client()
status, ok = has_table(client, table_name)
if not ok:
return False, '表不存在'
if ok:
list = []
for id in ids:
item = mycol.find_one({'img': id})
if item:
list.append(int(item['id']))
if list:
dstatus = delete_entity(client, table_name, list)
for id in ids:
mycol.delete_one({"img": id})
if dstatus.code == 0:
return True, '删除成功'
else:
return True, '该图片未进行索引或已被删除'
except Exception as e:
logging.error(e)
return False
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_train(table_name, data_loader, model, args):
if not table_name:
table_name = DEFAULT_TABLE
cache = Cache(default_cache_dir)
try:
vectors, names = feature_extract(data_loader, model, args)
vectors = vectors / vectors.norm()
print(vectors.shape)
vectors = vectors.tolist()
index_client = milvus_client()
# delete_table(index_client, table_name=table_name)
# time.sleep(1)
status, ok = has_table(index_client, table_name)
if not ok:
print("create table.")
create_table(index_client, table_name=table_name)
print("insert into:", table_name)
status, ids = insert_vectors(index_client, table_name, vectors)
create_index(index_client, table_name)
for i in range(len(names)):
# cache[names[i]] = ids[i]
cache[ids[i]] = names[i]
print("Train finished")
return "Train finished"
except Exception as e:
logging.error(e)
return "Error with {}".format(e)
def collection_info(table_name):
client = milvus_client()
tablestatus, ok = has_table(client, table_name)
if not ok:
return False, '', ''
status, info = get_collection_info(client, table_name)
return True, status, info
def get_partitions_list(table_name):
client = milvus_client()
tablestatus, ok = has_table(client, table_name)
if not ok:
return ok, '', ''
status, info = get_list_partitions(client, table_name)
return ok, status, info
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 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 create_collection(table_name):
print(table_name, '1111')
client = milvus_client()
tablestatus, ok = has_table(client, table_name)
print(tablestatus, ok)
if ok:
return False, ''
status = create_table(client, table_name)
return True, status
def main():
parser = argparse.ArgumentParser()
parser.add_argument('dataset_path', type=str)
args = parser.parse_args()
index_client = milvus_client()
conn = connect_mysql()
cursor = conn.cursor()
init_table(index_client, conn, cursor)
insert_data(index_client, conn, cursor, args.dataset_path)
def do_delete(table_name):
if not table_name:
table_name = DEFAULT_TABLE
try:
index_client = milvus_client()
status = delete_table(index_client, table_name=table_name)
return status
except Exception as e:
logging.error(e)
return "Error with {}".format(e)
def do_count(table_name):
if not table_name:
table_name = DEFAULT_TABLE
try:
index_client = milvus_client()
print("get table rows:", table_name)
num = count_table(index_client, table_name=table_name)
return num
except Exception as e:
logging.error(e)
return "Error with {}".format(e)
def create_p(table_name, partition_name):
client = milvus_client()
try:
status, ok = has_table(client, table_name)
if not ok:
create_table(client, table_name)
statusp = create_partition(client, table_name, partition_name)
return statusp
except Exception as e:
logging.error(e)
return '发生错误'.format(e)
def count_tab(table_name):
if not table_name:
return False
try:
client = milvus_client()
status, ok = has_table(client, table_name)
if not ok:
return False, ''
val = count_table(client, table_name)
return True, val
except Exception as e:
logging.error(e)
return '发生错误'.format(e)
def delete_collection(table_name):
if not table_name:
return False
try:
client = milvus_client()
status, ok = has_table(client, table_name)
if not ok:
return False
if ok:
status = delete_table(client, table_name)
if status.code == 0:
return True
except Exception as e:
logging.error(e)
return '错误'.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)
_, 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 insert_data_to_milvus(ids, img, voc):
index_client = milvus_client()
status, ok = has_table(index_client, IMG_TABLE)
if not ok:
print("create table.")
create_table(index_client, table_name=IMG_TABLE)
create_table(index_client, table_name=VOC_TABLE)
vectors_img = img_to_vec(img)
vectors_voc = voc_to_vec(voc)
status = {}
if not vectors_img:
status = {'status': 'faile', 'message': 'there is no file data'}
return status
try:
insert_vectors(index_client, IMG_TABLE, [vectors_img], [ids])
insert_vectors(index_client, VOC_TABLE, [vectors_voc], [ids])
status = {'status': 'success'}
return status
except Exception as e:
logging.error(e)
return "Fail with error {}".format(e)
def do_train(table_name, database_path):
detector = Detector()
if not table_name:
table_name = DEFAULT_TABLE
cache = Cache(default_cache_dir)
try:
result_images, object_num = run(detector, database_path)
#print("after detect:", object_num)
vectors, obj_images = get_object_vector(cache, image_encoder,
database_path + "/object")
#print("after detect:", len(vectors), obj_images)
index_client = milvus_client()
status, ok = has_table(index_client, table_name)
if not ok:
print("create table.")
create_table(index_client, table_name=table_name)
print("insert into:", table_name)
# vectors = normaliz_vec(vectors)
status, ids = insert_vectors(index_client, table_name, vectors)
#print(status,ids)
create_index(index_client, table_name)
shutil.rmtree(database_path + "/object")
imgs = os.listdir(database_path)
imgs.sort()
#print("-----imgs", imgs)
k = 0
ids = list(reversed(ids))
#print("ids", ids)
for num in object_num:
for i in range(num):
a = ids.pop()
#print("real;;;;;;;;;",a, imgs[k])
cache[a] = imgs[k]
k += 1
return print("train finished")
except Exception as e:
logging.error(e)
return "Error with {}".format(e)
def curd(vectors, img_name, mycol, partition=None, table_name=DEFAULT_TABLE):
try:
client = milvus_client()
if not table_name:
table_name = DEFAULT_TABLE
status, ok = has_table(client, table_name)
if not ok:
print('开始创建table')
create_table(client, table_name)
if partition:
status, ok = has_partition(client, table_name, partition)
if not ok:
create_partition(client, table_name, partition)
status, id = insert_vectors(client, table_name, vectors, partition)
# 存入缓存 以便后续进行反查
# redis.hset(REDIS_NAME, id[0], img_name)
mycol.insert_one({'id': id[0], 'img': img_name, 'table': table_name})
create_index(client, table_name)
print('OK 了')
return True, ''
except Exception as e:
logging.error(e)
return False, e
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 get_collections():
client = milvus_client()
status, info = get_table_list(client)
return status, info
def init_conn():
conn = connect_mysql()
cursor = conn.cursor()
index_client = milvus_client()
return index_client, conn, cursor
