edit : 벡터파일 생성 관련 분리

2025-12-02 14:15:49 +09:00
parent 18b24dcafa
commit 8b7ad2e973
15 changed files with 188 additions and 71 deletions
--- a/make_vector_files.py
+++ b/make_vector_files.py
@@ -0,0 +1,117 @@
+# from custom_apps.FEATURE_VECTOR_SIMILARITY_FAISS.faiss_functions import get_clip_info
+import os
+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
+
+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(IMG_LIST_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,
+                                save_txt_dir=IMG_LIST_PATH,
+                                item_info=item_info,
+                                index_type=model.value[1].index_type)
+        
+    else:
+        temp_path = os.path.join(IMG_LIST_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,
+                                    save_txt_dir=IMG_LIST_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__':
+    
+    import time
+    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)
+        
+        time.sleep(5) # huggingface api 요청 제한 회피 위해 대기 TODO(jwkim) huggingface 로그인은 한번만 진행하게 변경