update

tools/getResult.py

@@ -0,0 +1,78 @@
+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)