GLASSES_AI_SERVER/make_vector_files.py

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)