import pandas as pd import os from tqdm import tqdm import numpy as np from sklearn.metrics.pairwise import cosine_similarity class Database: def __init__(self, parquet_path=None, customized_parquets = None): self.default_parquet_path = 'datas/database_4000.parquet' self.parquet_path = parquet_path or self.default_parquet_path self.default_customized_parquets = ["datas/customized_database_0.parquet"] self.customized_parquets = customized_parquets or self.default_customized_parquets self.datas = None self.last_save_table = None if os.path.exists(self.parquet_path): self.load_from_parquet(self.parquet_path) self.load_from_customized(self.customized_parquets) self.clip_extractor = None self.bge_extractor = None self.en_keyword2data = {} def build_en_keyword2index(self): # build in lower case self.en_keyword2data = {row['translated_word'].lower(): row for i, row in self.datas.iterrows()} def search_by_en_keyword(self, keyword): if len(self.en_keyword2data) == 0: self.build_en_keyword2index() keyword = keyword.lower() if keyword in self.en_keyword2data: ans = self.en_keyword2data[keyword].to_dict() del ans["clip_feature"] del ans["bge_feature"] return ans else: return None def load_from_parquet(self, parquet_path): self.datas = pd.read_parquet(parquet_path) def load_from_customized(self, customized_parquets=None): customized_parquets = customized_parquets or self.customized_parquets # Load each parquet file and concatenate them into the self.datas DataFrame for index, parquet_file in enumerate(customized_parquets): if os.path.exists(parquet_file): temp_df = pd.read_parquet(parquet_file) if self.datas is None: self.datas = temp_df else: self.datas = pd.concat([self.datas, temp_df], ignore_index=True) # if last parquet file if index == len(customized_parquets) - 1: self.last_save_table = temp_df # if customized_parquets: # Record the last parquet file's contents as self.last_save_table def add_data(self, data, if_save=True): required_columns = ['keyword', 'name_in_cultivation', 'description_in_cultivation', 'translated_word', 'description'] for column in required_columns: if column not in data: raise ValueError(f"Missing required field: {column}") # Optional field if 'founder' not in data: data['founder'] = "" # Extract features if self.clip_extractor is None: self.init_clip_extractor() if self.bge_extractor is None: self.init_bge_extractor() data['clip_feature'] = self.clip_extractor.extract_text(data['translated_word'] + '.' + data['description']) data['bge_feature'] = self.bge_extractor.extract([data['keyword']])[0].tolist() # Convert to DataFrame and add to self.datas data_df = pd.DataFrame([data]) if self.datas is None: self.datas = data_df else: self.datas = pd.concat([self.datas, data_df], ignore_index=True) # set self.en_keyword2data to last row of self.datas self.en_keyword2data[data['translated_word'].lower()] = self.datas.iloc[-1] # Add to last_save_table if self.last_save_table is None: # self.last_save_table = data_df # create a new DataFrame with the same columns as self.datas self.last_save_table = pd.DataFrame(columns=self.datas.columns) self.last_save_table = pd.concat([self.last_save_table, data_df], ignore_index=True) if if_save: self.save_to_parquet(self.customized_parquets[-1], self.last_save_table ) def add_datas(self, datas, if_save=True): for data in datas: self.add_data(data, if_save=False) if if_save: self.save_to_parquet(self.customized_parquets[-1], self.last_save_table) def init_from_excel(self, excel_path): df = pd.read_excel(excel_path) # Drop rows with any empty cell in the required columns df.dropna(subset=['keyword', 'name_in_cultivation', 'description_in_cultivation', 'translated_word', 'description'], inplace=True) # Add the new columns df['clip_feature'] = None df['bge_feature'] = None self.datas = df self.extract_clip() self.extract_bge() def save_to_parquet(self, parquet_path=None, df = None): parquet_path = parquet_path or self.default_parquet_path if df is None: if self.datas is not None: self.datas.to_parquet(parquet_path) else: df.to_parquet(parquet_path) def init_clip_extractor(self): if self.clip_extractor is None: try: from CLIPExtractor import CLIPExtractor except: from src.CLIPExtractor import CLIPExtractor cache_dir = "models" self.clip_extractor = CLIPExtractor(model_name = "openai/clip-vit-large-patch14",cache_dir = cache_dir) def extract_clip(self): if self.clip_extractor is None: self.init_clip_extractor() clip_features = [] # for text in tqdm(self.datas['keyword'], desc='Extracting CLIP features'): for index, row in tqdm(self.datas.iterrows(), desc='Extracting CLIP features', total=len(self.datas)): text = row['translated_word'] + '.' + row['description'] if text: feature = self.clip_extractor.extract_text(text) else: feature = None clip_features.append(feature) self.datas['clip_feature'] = clip_features def init_bge_extractor(self): if self.bge_extractor is None: try: from text_embedding import TextExtractor except: from src.text_embedding import TextExtractor self.bge_extractor = TextExtractor('BAAI/bge-small-zh-v1.5') def top_k_search(self, query_feature, attribute, top_k=15): # Ensure the attribute exists in the dataframe if attribute not in self.datas.columns: raise ValueError(f"Attribute {attribute} not found in the data.") # Convert query feature and attribute features to numpy arrays query_feature = np.array(query_feature).reshape(1, -1) attribute_features = np.stack(self.datas[attribute].dropna().values) # Compute cosine similarity between query and all attributes similarities = cosine_similarity(query_feature, attribute_features)[0] # Get the top_k indices based on similarity top_k_indices = np.argsort(similarities)[-top_k:][::-1] # Retrieve the top_k most similar items top_k_results = self.datas.iloc[top_k_indices].copy() top_k_results = top_k_results.drop(columns=['clip_feature', 'bge_feature']) top_k_results['similarity'] = similarities[top_k_indices] return top_k_results.to_dict(orient='records') def search_with_image_name(self, image_name): self.init_clip_extractor() img_feature = self.clip_extractor.extract_image_from_file(image_name) return self.top_k_search(img_feature, 'clip_feature') def search_with_image(self, image, if_opencv = False ): if self.clip_extractor is None: self.init_clip_extractor() img_feature = self.clip_extractor.extract_image(image, if_opencv = if_opencv) return self.top_k_search(img_feature, 'clip_feature') def search_with_chinese(self, text): if self.bge_extractor is None: self.init_bge_extractor() text_feature = self.bge_extractor.extract([text])[0].tolist() return self.top_k_search(text_feature, 'bge_feature') def extract_bge(self): if self.bge_extractor is None: self.init_bge_extractor() # Extract features for each row and store them in the bge_feature column bge_features = [] for text in tqdm(self.datas['keyword'], desc='Extracting BGE features'): if text: feature = self.bge_extractor.extract([text])[0].tolist() else: feature = None bge_features.append(feature) self.datas['bge_feature'] = bge_features if __name__ == '__main__': # Usage example db = Database() re_extract = False if db.datas is None or re_extract: print("Rebuilding database from excel file") db.init_from_excel('datas/database_4000.xlsx') db.save_to_parquet() # print(db.datas[0].keys()) query_text = "钢琴" results = db.search_with_chinese(query_text) print(results[0].keys()) for result in results[:3]: print(result) image_path = "datas/老虎.jpg" results = db.search_with_image_name(image_path) for result in results[:3]: print(result) # 'keyword': '老虎狗', 'name_in_cultivation': '灵虎犬神', 'description_in_cultivation': '在九天灵脉汇聚的仙山之巅，灵虎犬神身披星图 # 斑纹，汲取日月精华，以雷霆之力守护仙脉，其双眼中映照着轮回之道，是修仙者追寻天地真理的指引，也是象征极致灵性的神秘灵兽。', 'translated_word': 'Tiger Dog', 'description': 'A Tiger Dog is a term that might refer to a mythical creature or a breed of dog with a distinctive and unusual appearance, resembling the features of a tiger. It could be characterized by its striking coat with patterns similar to those of a tiger, or by having a demeanor that is fierce and majestic like a tiger. This term is not commonly used in # conventional contexts and might be found in stories, folktales, or in the names of unique dog breeds that have been bred to exhibit such features.', 'founder': '' # test_new_data = { # "keyword": "老虎狗2", # "name_in_cultivation": "灵虎犬神", # "description_in_cultivation": "在九天灵脉汇聚的仙山之巅，灵虎犬神身披星图斑纹，汲取日月精华，以雷霆之力守护仙脉，其双眼中映照着轮回之道，是修仙者追寻天地真理的指引，也是象征极致灵性的神秘灵兽。", # "translated_word": "Tiger Dog", # "description":"A Tiger Dog is a term that might refer to a mythical creature or a breed of dog with a distinctive and unusual appearance, resembling the features of a tiger. It could be characterized by its striking coat with patterns similar to those of a tiger, or by having a demeanor that is fierce and majestic like a tiger. This term is not commonly used in conventional contexts and might be found in stories, folktales, or in the names of unique dog breeds that have been bred to exhibit such features." # } # db.add_data(test_new_data)