update

image_quality.py

@@ -0,0 +1,104 @@
+import os
+
+import numpy as np
+
+from ytracking.track_ import *
+from tools.Interface import AiInterface, AiClass
+from tools.operate_usearch import create_base_index, search_in_index
+from tools.initModel import models
+from imgcompare import get_feature_list, compute_similarity_matrix
+import pickle
+models.initModel()
+ai_obj = AiClass()
+
+
+def get_img_lists(pth):
+    imglist, imglists = [], []
+    for root, dirs, files in os.walk(pth):
+        if not any(dirs):
+            for file in files:
+                if file.endswith('.jpg'):
+                    imglist.append(os.sep.join([root, file]))
+            imglists.append(imglist)
+            imglist = []
+    return imglists
+
+def get_standard_image(cosine_similarities, similarity_threshold=0.6):
+    """
+    :param cosine_similarities:
+    :return:
+    """
+    target_indexs = []
+    max_similarity = {}
+    mask = (cosine_similarities > similarity_threshold)
+    counts = mask.sum(axis=1)
+    for key in range(counts.shape[0]):
+        max_similarity[key] = counts[key]
+    sorted_dict_desc = dict(sorted(max_similarity.items(), key=lambda item: item[1], reverse=True))
+    keys = list(sorted_dict_desc.keys())
+    while len(keys) > 10:
+        target_indexs.append(keys[0])
+        single_line = cosine_similarities[keys[0], :]
+        rows = np.where((single_line > similarity_threshold))
+        if len(rows[0]) < 2:
+            break
+        for row in rows[0]:
+            try:
+                keys.remove(row)
+            except Exception as e:
+                continue
+    # print(target_indexs)
+    return target_indexs
+
+def create_feature_library(pth, save_index_name, index_file_pth=None):
+    target_feature_lists, target_barcode_lists = [], []
+    imglists = get_img_lists(pth)
+    for imglist in imglists:
+        feature_list = get_feature_list(imglist, False)
+        cosine_similarities = compute_similarity_matrix(feature_list)
+        target_indexs = get_standard_image(cosine_similarities)
+        target_feature_lists.append([feature_list[i] for i in target_indexs])
+        target_barcode_lists.append([os.path.basename(imglist[i]).split('_')[0] for i in target_indexs])
+    create_base_index(save_index_name=save_index_name,
+                      barcodes=target_barcode_lists,
+                      features=target_feature_lists,
+                      index_file_pth=index_file_pth)
+    with open('search_library/target_barcode_lists.pkl', 'wb') as f:
+        pickle.dump(target_barcode_lists, f)
+
+def search_top_in_index(test_image_pth, index_name):  #1:N
+    s_barcode, s_similarity = [], []
+    img_lists = [os.sep.join([test_image_pth, name]) for name in os.listdir(test_image_pth)]
+    feature_lists = get_feature_list(img_lists, False)
+    for feature in feature_lists:
+        result = search_in_index(query=np.array(feature), index_name=index_name)
+        s_barcode.append(result.keys)
+        s_similarity.append(1-result.distances)
+    s_barcode = np.array(s_barcode)
+    s_similarity = np.array(s_similarity)
+    return s_barcode, s_similarity
+
+def search_one_in_index(test_image_pth, index_name): # 1:1
+    barcodes = [int(os.path.basename(name).split('_')[0]) for name in os.listdir(test_image_pth)]
+    barcodes = list(set(barcodes))
+    # barcodes = ['6934364805640']
+    img_lists = [os.sep.join([test_image_pth, name]) for name in os.listdir(test_image_pth)]
+    feature_lists = get_feature_list(img_lists, False)
+    result = search_in_index(barcode=barcodes,
+                             query=feature_lists,
+                             index_name=index_name,
+                             temp_index=False)
+    print(feature_lists)
+
+
+if __name__ == '__main__':
+    pth = 'imageQualityData/test_data'
+    save_index_name = 'search_library/test_index_10_simple_0717.usearch'
+    create_feature_library(pth,
+                           save_index_name=save_index_name)
+
+    # test_images_pth = 'D:/Project/ieemoo/image_quality_assessment/imageQualityData/test_images'
+    # # index_name = 'D:/Project/ieemoo/image_quality_assessment/search_library/test_index_10_normal_0717.usearch'
+    # index_name = 'D:/Project/ieemoo/image_quality_assessment/search_library/test_index_10_simple_0717.usearch'
+    # # search_top_in_index(test_images_pth, index_name)
+    # search_one_in_index(test_images_pth, index_name)