120 lines
5.0 KiB
Python
120 lines
5.0 KiB
Python
import os
|
|
import logging
|
|
import matplotlib.pyplot as plt
|
|
|
|
from custom_apps.FEATURE_VECTOR_SIMILARITY_FAISS.faiss_similarity_search import VectorSimilarity
|
|
from custom_apps.FEATURE_VECTOR_SIMILARITY_FAISS.const import *
|
|
from custom_apps.FEATURE_VECTOR_SIMILARITY_FAISS.faiss_functions import get_models, find_glass_folder_images, find_parts_folder_images
|
|
from vector_rest.app import models as VM
|
|
|
|
|
|
#log level
|
|
os.environ["HF_HUB_VERBOSITY"] = "info"
|
|
matplotlib_logger = logging.getLogger("matplotlib")
|
|
matplotlib_logger.setLevel(logging.INFO)
|
|
|
|
def make_image_files(item_info, index_type:VM.VitIndexType, model_type:VM.VitModelType):
|
|
model = get_models(index_type=index_type, model_type=model_type)
|
|
|
|
txt_file_path = os.path.join(FAISS_VECTOR_PATH,f'{model.name}_{os.path.basename(model.value[1].trained_model)}_{model.value[1].index_type}_{item_info}.txt')
|
|
|
|
image_lists = []
|
|
if not os.path.exists(txt_file_path):
|
|
if item_info == VectorSearchItem.glass:
|
|
image_lists = find_glass_folder_images(IMG_LIST_PATH)
|
|
else:
|
|
image_lists = find_parts_folder_images(IMG_LIST_PATH, item_info)
|
|
|
|
if image_lists:
|
|
with open(txt_file_path, 'w') as f:
|
|
for img_path in image_lists:
|
|
f.write(f"{img_path}\n")
|
|
else:
|
|
temp_path = os.path.join(FAISS_VECTOR_PATH,f'{model.name}_{os.path.basename(model.value[1].trained_model)}_{model.value[1].index_type}_{item_info}.bak')
|
|
|
|
try:
|
|
os.rename(txt_file_path, temp_path)
|
|
|
|
if item_info == VectorSearchItem.glass:
|
|
image_lists = find_glass_folder_images(IMG_LIST_PATH)
|
|
else:
|
|
image_lists = find_parts_folder_images(IMG_LIST_PATH, item_info)
|
|
|
|
if image_lists:
|
|
with open(txt_file_path, 'w') as f:
|
|
for img_path in image_lists:
|
|
f.write(f"{img_path}\n")
|
|
|
|
except Exception as e:
|
|
print(f"Error updating image file list: {e}")
|
|
|
|
finally:
|
|
if os.path.exists(txt_file_path):
|
|
os.remove(temp_path)
|
|
else:
|
|
os.rename(temp_path, txt_file_path)
|
|
|
|
return txt_file_path
|
|
|
|
def make_vector_files(item_info, index_type:VM.VitIndexType, model_type:VM.VitModelType):
|
|
model = get_models(index_type=index_type, model_type=model_type)
|
|
|
|
txt_file_path = make_image_files(item_info=item_info, index_type=index_type, model_type=model_type)
|
|
|
|
index_file_path = os.path.join(FAISS_VECTOR_PATH,f'{model.name}_{os.path.basename(model.value[1].trained_model)}_{model.value[1].index_type}_{item_info}_{DEFAULT_INDEX_NAME_SUFFIX}')
|
|
|
|
if not os.path.exists(index_file_path):
|
|
|
|
vector_model = VectorSimilarity(model_info=model,
|
|
index_file_path=index_file_path,
|
|
txt_file_path=txt_file_path)
|
|
|
|
with open(txt_file_path, 'r') as f:
|
|
image_lists = [line.strip() for line in f.readlines()]
|
|
|
|
vector_model.save_index_from_files(images_path_lists=image_lists,
|
|
save_index_dir=FAISS_VECTOR_PATH,
|
|
item_info=item_info,
|
|
index_type=model.value[1].index_type)
|
|
|
|
else:
|
|
temp_path = os.path.join(FAISS_VECTOR_PATH,f'{model.name}_{os.path.basename(model.value[1].trained_model)}_{model.value[1].index_type}_{item_info}_{DEFAULT_INDEX_NAME_SUFFIX}.bak')
|
|
|
|
try:
|
|
os.rename(index_file_path, temp_path)
|
|
|
|
vector_model = VectorSimilarity(model_info=model,
|
|
index_file_path=index_file_path,
|
|
txt_file_path=txt_file_path)
|
|
|
|
with open(txt_file_path, 'r') as f:
|
|
image_lists = [line.strip() for line in f.readlines()]
|
|
|
|
vector_model.save_index_from_files(images_path_lists=image_lists,
|
|
save_index_dir=FAISS_VECTOR_PATH,
|
|
item_info=item_info,
|
|
index_type=model.value[1].index_type)
|
|
except Exception as e:
|
|
print(f"Error updating vector file: {e}")
|
|
finally:
|
|
if os.path.exists(index_file_path):
|
|
os.remove(temp_path)
|
|
else:
|
|
os.rename(temp_path, index_file_path)
|
|
|
|
print(f"Vector file created: {index_file_path} :::::::::: item_info: {item_info}, index_type: {index_type}, model_type: {model_type}")
|
|
|
|
return index_file_path
|
|
|
|
if __name__ == '__main__':
|
|
|
|
class_attributes = dict(VectorSearchItem.__dict__)
|
|
|
|
pure_data_dict = {
|
|
key: value
|
|
for key, value in class_attributes.items()
|
|
if not key.startswith('__')
|
|
}
|
|
|
|
for item_key, item_value in pure_data_dict.items():
|
|
make_vector_files(item_info=item_value, index_type=VM.VitIndexType.l2, model_type=VM.VitModelType.b32) |