# -*- coding: utf-8 -*- import pdb import random import json import time import torch from PIL import Image from contrast.model import resnet18, MobileNetV3_Large # import pymilvus # from pymilvus import ( # connections, # utility, # FieldSchema, CollectionSchema, DataType, # Collection, # Milvus # ) # from config import config as conf from contrast.search import ImgSearch from contrast.img_data import queueImgs_add import sys from threading import Thread sys.path.append('../tools') from tools.config import cfg as conf from tools.config import gvalue def test_preprocess(images: list, actionModel) -> torch.Tensor: res = [] for img in images: # print(img) try: im = conf.test_transform(img) if actionModel else conf.test_transform(Image.open(img)) res.append(im) except: continue data = torch.stack(res) return data def inference(images, model, actionModel): data = test_preprocess(images, actionModel) if torch.cuda.is_available(): data = data.to(conf.device) features = model(data) return features def group_image(images, batch=64) -> list: """Group image paths by batch size""" size = len(images) res = [] for i in range(0, size, batch): end = min(batch + i, size) res.append(images[i:end]) return res def barcode_state(barcodeIDList): with open('contrast/main_barcodes.json', 'r') as file: data = json.load(file) main_barcode = list(data.values())[0] barIdList_true = [] barIdList_false = [] for barId in barcodeIDList: bar = barId.split('_')[1] if bar in main_barcode: barIdList_true.append(barId) else: barIdList_false.append(barId) return barIdList_true, barIdList_false def getFeatureList(barList, imgList, model, actionModel): featList = [[] for i in range(len(barList))] for index, feat in enumerate(imgList): groups = group_image(feat) for group in groups: feat_tensor = inference(group, model, actionModel) for fe in feat_tensor: if fe.device == 'cpu': fe_np = fe.squeeze().detach().numpy() else: fe_np = fe.squeeze().detach().cpu().numpy() featList[index].append(fe_np) return featList def img2feature(imgs_dict, model, actionModel, barcode_flag): if not len(imgs_dict) > 0: raise ValueError("Tracking fail no images files provided") queBarIdList = list(imgs_dict.keys()) if barcode_flag: # # ========判断barcode是否在特征库============ queBarIdList_t, barIdList_f = barcode_state(queBarIdList) queFeatList_t = [] if len(queBarIdList_t) == 0: print(f"All barcodes are not in the main_library: {barIdList_f}") return queBarIdList_t, queFeatList_t else: if len(barIdList_f) > 0: ## 将不在barcode库中的barcode及图片删除 print(f"These barcodes are not in the main_library: {barIdList_f}") for bar_f in barIdList_f: del imgs_dict[bar_f] queImgList_t = list(imgs_dict.values()) queFeatList_t = getFeatureList(queBarIdList_t, queImgList_t, model, actionModel) return queBarIdList_t, queFeatList_t else: queImgsList = list(imgs_dict.values()) queFeatList = getFeatureList(queBarIdList, queImgsList, model, actionModel) return queBarIdList, queFeatList # def create_milvus(collection_name, host, port, barcode_list, features): # # 1. connect to Milvus # fmt = "\n=== {:30} ===\n" # connections.connect('default', host=host, port=port) # 连接到 Milvus 服务器 # has = utility.has_collection(collection_name) ##检查collection_name是否存在milvus中 # print(f"Does collection {collection_name} exist in Milvus: {has}") # # if has: ## 删除collection_name的库 # # utility.drop_collection(collection_name) # # # 2. create colllection # fields = [ # FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=False, max_length=100), ###图片路径 # FieldSchema(name="embeddings", dtype=DataType.FLOAT_VECTOR, dim=256) # ] # schema = CollectionSchema(fields) # print(fmt.format(f"Create collection {collection_name}")) # hello_milvus = Collection(collection_name, schema, consistency_level="Strong") # # 3. insert data # for i in range(len(features)): # entities = [ # # provide the pk field because `auto_id` is set to False # [barcode_list[i]] * len(features[i]), ## 图片维度和向量维度需匹配 每个向量都生成一个barcode # features[i], # ] # print(fmt.format("Start inserting entities")) # insert_result = hello_milvus.insert(entities) # hello_milvus.flush() # print(f"Number of entities in {collection_name}: {hello_milvus.num_entities}") # check the num_entities # return hello_milvus # def load_collection(collection_name): # collection = Collection(collection_name) # # collection.release() ### 将collection从加载状态变成未加载 # # collection.drop_index() ### 删除索引 # # index_params = { # "index_type": "IVF_FLAT", # # "index_type": "IVF_SQ8", # # "index_type": "GPU_IVF_FLAT", # "metric_type": "COSINE", # "params": { # "nlist": 10000 # } # } # #### 准确率低 # # index_params = { # # "index_type": "IVF_PQ", # # "metric_type": "COSINE", # # "params": { # # "nlist": 99, # # "m": 2, # # "nbits": 8 # # } # # } # collection.create_index( # field_name="embeddings", # index_params=index_params, # index_name="SQ8" # ) # collection.load() # return collection # def similarity(queImgsDict, add_flag, barcode_flag, main_milvus, model, barcode_list, actionModel): # searchImg = ImgSearch() ## 相似度比较 # # 将输入图片加入临时库 # if add_flag: # if actionModel: # queBarIdList, queBarIdFeatures = img2feature(dict(list(queImgsDict.items())[2:-2]), model, actionModel, barcode_flag) # else: # queBarIdList, queBarIdFeatures = img2feature(dict(list(queImgsDict.items())[:-2]), model, actionModel, barcode_flag) # # if barcode_flag: ### 加购 有barcode -> 输出top10和top1 # if len(queBarIdList) == 0: # top10, top1 = {}, {} # else: # for bar in queBarIdList: # # gvalue.tempLibList.append(bar) ## 临时特征库key值为macID_barcode_trackID # if gvalue.tempLibLists.get(gvalue.mac_id) is not None: # gvalue.tempLibLists[gvalue.mac_id] += [bar] ## 临时特征库key值为macID_barcode_trackID # else: # gvalue.tempLibLists[gvalue.mac_id] = [bar] # # 存入临时特征库 # # create_milvus('temp_features', conf.host, conf.port, queBarIdList, queBarIdFeatures) # # thread = Thread(target=create_milvus, kwargs={'collection_name': 'temp_features', # 'host': conf.host, # 'port': conf.port, # 'barcode_list': queBarIdList, # 'features': queBarIdFeatures}) # thread.start() # start1 = time.time() # top10 = searchImg.mainSearch10(main_milvus, queBarIdList, queBarIdFeatures) # start2 = time.time() # print('search top10 time>>>> {}'.format(start2-start1)) # top1 = searchImg.mainSearch1(main_milvus, queBarIdList, queBarIdFeatures) # start3 = time.time() # print('search top1 time>>>>> {}'.format(start3-start2)) # return top10, top1, gvalue.tempLibLists # else: # 加购 无barcode -> 输出top10 # # 无barcode时，生成随机数作为字典key值 # queBarIdList_rand = [] # for i in range(len(queBarIdList)): # random_number = ''.join(random.choices('0123456789', k=10)) # queBarIdList_rand.append(str(random_number)) # # gvalue.tempLibList.append(str(random_number)) # if gvalue.tempLibLists.get(gvalue.mac_id) is not None: # gvalue.tempLibLists[gvalue.mac_id] += [str(random_number)] ## 临时特征库key值为macID_barcode_trackID # else: # gvalue.tempLibLists[gvalue.mac_id] = [str(random_number)] # # create_milvus('temp_features', conf.host, conf.port, queBarIdList_rand, queBarIdFeatures) # thread = Thread(target=create_milvus, kwargs={'collection_name': 'temp_features', # 'host': conf.host, # 'port': conf.port, # 'barcode_list': queBarIdList_rand, # 'features': queBarIdFeatures}) # thread.start() # top10 = searchImg.mainSearch10(main_milvus, queBarIdList, queBarIdFeatures) # # print(f'top10: {top10}') # return top10, gvalue.tempLibLists # else: # 退购 -> 输出top10和topn # if gvalue.tempLibLists.get(gvalue.mac_id) is None: # gvalue.tempLibList = [] # else: # gvalue.tempLibList = gvalue.tempLibLists[gvalue.mac_id] # ## 加载临时特征库 # tempMilvusName = "temp_features" # has = utility.has_collection(tempMilvusName) # print(f"Does collection {tempMilvusName} exist in Milvus: {has}") # tempMilvus = load_collection(tempMilvusName) # print(f"Number of entities in {tempMilvusName}: {tempMilvus.num_entities}") # if actionModel: # barcode_list = barcode_list # else: # barcode_list = queueImgs_add['barcode_list'] # if actionModel: # queBarIdList, queBarIdFeatures = img2feature(dict(list(queImgsDict.items())[2:-1]), model, actionModel, barcode_flag) # else: # queBarIdList, queBarIdFeatures = img2feature(dict(list(queImgsDict.items())[:-3]), model, actionModel, barcode_flag) # if barcode_flag: # if len(queBarIdList) == 0: # top10, top1, top_n = {}, {}, {} # else: # start1 = time.time() # top1 = searchImg.mainSearch1(main_milvus, queBarIdList, queBarIdFeatures) # start2 = time.time() # print('search top1 time>>>> {}'.format(start2 - start1)) # top10 = searchImg.mainSearch10(main_milvus, queBarIdList, queBarIdFeatures) # start3 = time.time() # print('search top10 time>>>> {}'.format(start3 - start2)) # top_n = searchImg.tempSearch(tempMilvus, queBarIdList, queBarIdFeatures, barcode_list, gvalue.tempLibList) # # print(f'top10: {top10}, top1: {top1}, topn: {top_n}') # return top10, top1, top_n # else: # top10 = searchImg.mainSearch10(main_milvus, queBarIdList, queBarIdFeatures) # top_n = searchImg.tempSearch(tempMilvus, queBarIdList, queBarIdFeatures, barcode_list, gvalue.tempLibList) # # print(f'top10: {top10}, topn: {top_n}') # return top10, top_n def similarity_interface(dataCollection): queImgsDict = dataCollection.queImgsDict add_flag = dataCollection.add_flag barcode_flag = dataCollection.barcode_flag main_milvus = dataCollection.mainMilvus #tempLibList = dataCollection.tempLibList model = dataCollection.model actionModel = dataCollection.actionModel barcode_list = dataCollection.barcode_list #return similarity(queImgsDict, add_flag, barcode_flag, main_milvus, tempLibList, model, barcode_list, actionModel) return 0 if __name__ == '__main__': pass # connections.connect('default', host=conf.host, port=conf.port) # # 加载主特征库 # mainMilvusName = "main_features" # has = utility.has_collection(mainMilvusName) # print(f"Does collection {mainMilvusName} exist in Milvus: {has}") # mainMilvus = Collection(mainMilvusName) # mainMilvus.load() # model = initModel() # # queueImgs_add queueImgs_back 分别为加购和退购时的入参 # add_flag = queueImgs_add['add_flag'] # barcode_flag = queueImgs_add['barcode_flag'] # tempLibList = [] # 临时特征库的barcodeId_list # # tempLibList = ['3500610085338_01', '4260290263776_01'] ##test # if add_flag: # if barcode_flag: # 加购 有barcode -> 输出top10和top1 # top10, top1, tempLibList = similarity(queueImgs_add, add_flag, barcode_flag, mainMilvus, tempLibList, model) # print(f"top10: {top10}\ntop1: {top1}") # else: # 加购 无barcode -> 输出top10 # top10, tempLibList = similarity(queueImgs_add, add_flag, barcode_flag, mainMilvus, tempLibList, model) # else: # 退购 -> 输出top10和topn # top10, topn = similarity(queueImgs_back, add_flag, barcode_flag, mainMilvus, tempLibList, model)