79 lines
3.0 KiB
Python
79 lines
3.0 KiB
Python
from tools.Interface import AiInterface, AiClass
|
|
# from Interface import AiInterface, AiClass
|
|
from config import cfg
|
|
import numpy as np
|
|
from PIL import Image
|
|
import pdb
|
|
from track_ import *
|
|
import time
|
|
|
|
'''
|
|
跟踪与对比输入
|
|
'''
|
|
ai_obj = AiClass()
|
|
|
|
|
|
def deal_similarity_data(message):
|
|
# print('message --- > {}'.format(message))
|
|
car_mac = message['videoIds'].split('_')[-3]
|
|
similar = {'add_flag': True if message['action'] in {'1', '3'} else False,
|
|
'barcode_flag': True if not message['barcode'] == 'null' else False}
|
|
if similar['add_flag'] and similar['barcode_flag']: # 加购有barcode
|
|
for Id, image_list in zip(message['trackIdList'],
|
|
message['images']):
|
|
similar[car_mac + '_' + message['barcode'] + '_' + Id] = image_list
|
|
similar['barcode_list'] = message['barcodeList']
|
|
elif similar['add_flag'] and (not similar['barcode_flag']): # 加购无barcode
|
|
for Id, image_list in zip(message['trackIdList'],
|
|
message['images']):
|
|
similar[car_mac + '_' + Id] = image_list
|
|
similar['barcode_list'] = message['barcodeList']
|
|
elif (not similar['add_flag']) and similar['barcode_flag']: # 退购有barcode
|
|
for Id, image_list in zip(message['trackIdList'],
|
|
message['images']):
|
|
similar[car_mac + '_' + message['barcode'] + '_' + Id] = image_list
|
|
similar['barcode_list'] = message['barcodeList']
|
|
else: # 退购无barcode
|
|
for Id, image_list in zip(message['trackIdList'],
|
|
message['images']):
|
|
similar[car_mac + '_' + Id] = image_list
|
|
similar['barcode_list'] = message['barcodeList']
|
|
similar['sequenceId'] = message['sequenceId']
|
|
return similar
|
|
|
|
|
|
def get_similarity_result(obj: AiInterface, videopth, model, camera_id, message):
|
|
"""
|
|
获取相似度结果。
|
|
|
|
:param videopth:
|
|
:param obj: AiInterface 对象，用于获取跟踪框和相似度数据。
|
|
:param message: dict, 可选参数，包含跟踪数据和相似度处理结果。
|
|
:return: 相似度结果。
|
|
"""
|
|
try:
|
|
track_boxes, features_dict, frame_id_img = run(model, source=videopth)
|
|
allimages, trackIdList = obj.getTrackingBox(track_boxes, features_dict, camera_id, frame_id_img)
|
|
message['trackIdList'] = trackIdList
|
|
message['images'] = allimages
|
|
message = deal_similarity_data(message)
|
|
similarityRes = obj.getSimilarity(model, message)
|
|
except ValueError as ve:
|
|
print('ve >>>> ', ve)
|
|
similarityRes = {'top10': {},
|
|
'top1': {},
|
|
'topn': {},
|
|
'tempLibList': [],
|
|
'sequenceId': message['sequenceId']
|
|
}
|
|
return similarityRes
|
|
|
|
|
|
if __name__ == '__main__':
|
|
message = {
|
|
'action': '1',
|
|
'barcode': '084501446314',
|
|
'sequenceId': 'test'
|
|
}
|
|
get_similarity_result(ai_obj, message)
|