Upload 19 files

src/CLIPExtractor.py

@@ -0,0 +1,93 @@
+import os
+import torch
+from transformers import CLIPProcessor, CLIPModel
+import cv2
+from PIL import Image
+import numpy as np
+
+
+
+class CLIPExtractor:
+    def __init__(self, model_name="openai/clip-vit-large-patch14", cache_dir=None):
+
+        # 设置代理环境变量
+        os.environ['HTTP_PROXY'] = 'http://localhost:8234'
+        os.environ['HTTPS_PROXY'] = 'http://localhost:8234'
+
+        # 设置环境变量
+        os.environ["HF_ENDPOINT"] = "https://hf-api.gitee.com"
+        os.environ["HF_HOME"] = os.path.expanduser("models/")
+
+        if not cache_dir:
+            # 指定缓存目录
+            cache_dir = "models"
+            if not os.path.exists(cache_dir) and os.path.exists("../models"):
+                cache_dir = "../models"
+
+        # Initialize the model and processor with specified values
+        self.model = CLIPModel.from_pretrained(model_name, cache_dir=cache_dir)
+        self.processor = CLIPProcessor.from_pretrained(model_name, cache_dir=cache_dir)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
+
+    def extract_image(self, frame):
+        # Convert frame (from OpenCV) to PIL Image
+        image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+        images = [image]
+        
+        # Process the image and extract features
+        inputs = self.processor(images=images, return_tensors="pt").to(self.device)
+        with torch.no_grad():
+            outputs = self.model.get_image_features(**inputs)
+        
+        ans = outputs.cpu().numpy()
+        return ans[0]
+
+    def extract_image_from_file(self, file_name):
+        if not os.path.exists(file_name):
+            raise FileNotFoundError(f"File {file_name} not found.")
+        
+        images = [Image.open(file_name).convert("RGB")]
+        
+        # Process the image and extract features
+        inputs = self.processor(images=images, return_tensors="pt").to(self.device)
+        with torch.no_grad():
+            outputs = self.model.get_image_features(**inputs)
+        
+        ans = outputs.cpu().numpy()
+        return ans[0]
+    
+    def extract_text(self, text):
+        if not isinstance(text, str) or not text:
+            raise ValueError("Input text should be a non-empty string.")
+        
+         # Tokenize the text
+        inputs = self.processor.tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=77).to(self.device)
+        
+
+        # Process the text and extract features
+        # inputs = self.processor(text=[text], return_tensors="pt", padding=True).to(self.device)
+        
+        with torch.no_grad():
+            outputs = self.model.get_text_features(**inputs)
+        
+        ans = outputs.cpu().numpy()
+        return ans[0]
+
+
+if __name__ == "__main__":
+
+    clip_extractor = CLIPExtractor()
+
+    sample_image = "images/狐狸.jpg"
+    # 提取图像特征
+    image_feature = clip_extractor.extract_image_from_file(sample_image)
+
+
+    # 提取文本特征
+    sample_text = "A photo of fox"
+    text_feature = clip_extractor.extract_text(sample_text)
+
+    # consine similarity
+    cosine_similarity = np.dot(image_feature, text_feature) / (np.linalg.norm(image_feature) * np.linalg.norm(text_feature))
+    print(cosine_similarity)

src/Captioner.py

@@ -0,0 +1,148 @@
+from PIL import Image
+import base64
+from io import BytesIO
+import os
+from openai import OpenAI
+import json
+
+class Captioner:
+    def __init__(self, api_key_path = None, proxy=None, api_base="https://api.lingyiwanwu.com/v1"):
+        
+        if api_key_path is None:
+            # try find datas/01_key.txt and ../datas/01_key.txt
+            cand_paths = ['datas/01_key.txt', '../datas/01_key.txt']
+            flag = False
+            for path in cand_paths:
+                if os.path.exists(path):
+                    api_key_path = path
+                    flag = True
+                    break
+
+            if not flag:
+                raise ValueError("Please provide the path to the API key file.")
+
+
+        self.api_key = self.load_access_token(api_key_path)
+        self.api_base = api_base
+        if proxy:
+            os.environ['HTTP_PROXY'] = proxy
+            os.environ['HTTPS_PROXY'] = proxy
+        self.client = OpenAI(
+            api_key=self.api_key,
+            base_url=self.api_base
+        )
+
+        self.history = {}
+        self.history_file = None
+
+        self.load_history()
+
+    def load_access_token(self, file_path):
+        with open(file_path, 'r') as file:
+            return file.read().strip()
+
+    def image2base64(self, image_path):
+        # 打开图像
+        with Image.open(image_path) as img:
+            # 检查图像高度是否超过480
+            if img.height > 480:
+                # 计算调整后的宽度，以保持宽高比不变
+                aspect_ratio = img.width / img.height
+                new_height = 480
+                new_width = int(new_height * aspect_ratio)
+                img = img.resize((new_width, new_height), Image.ANTIALIAS)
+            
+            # 使用BytesIO在内存中保存调整大小后的图像
+            buffered = BytesIO()
+            img.save(buffered, format="JPEG")
+            buffered.seek(0)
+
+            # 将图像转换为Base64编码字符串
+            img_base64 = "data:image/jpeg;base64," + base64.b64encode(buffered.read()).decode('utf-8')
+        
+        return img_base64
+    
+    def load_history(self, jsonl_file_name=None):
+        if jsonl_file_name is None:
+            jsonl_file_name = "datas/caption_history.jsonl"
+        
+        self.history_file = jsonl_file_name
+        
+        if os.path.exists(jsonl_file_name):
+            with open(jsonl_file_name, 'r', encoding='utf-8') as f:
+                for line in f:
+                    data = json.loads(line)
+                    self.history[data['file_name']] = data['response']
+
+    def search_from_history(self, file_name):
+        return self.history.get(file_name, None)
+    
+    def save_history(self, jsonl_file_name=None):
+        if jsonl_file_name is None:
+            jsonl_file_name = self.history_file
+        
+        if jsonl_file_name:
+            with open(jsonl_file_name, 'w', encoding='utf-8') as f:
+                for file_name, response in self.history.items():
+                    json.dump({'file_name': file_name, 'response': response}, f, ensure_ascii=False)
+                    f.write('\n')
+
+        # print(f"History saved to {jsonl_file_name}")
+
+    def add_to_history(self, file_name, response):
+        self.history[file_name] = response
+
+    def caption(self, image_name):
+
+        # Check if the caption is already in the history
+        cached_response = self.search_from_history(image_name)
+        if cached_response:
+            # print("return the cache")
+            return cached_response
+        
+        prompt = """Analyze the image and output in JSON format, including the following fields:
+            - "detailed_description": A detailed description of the image content.
+            - "major_object": Determine the main object/scene in the image based on the description, output with a simple word
+            - "Chinese_name": 判断图片中主要物体的中文名
+            - "real_or_composite":  Determine whether this image was taken with a camera or created/modifed by a computer, output with real or composite."""
+        
+        img_base64 = self.image2base64(image_name)
+        
+        completion = self.client.chat.completions.create(
+            model="yi-vision",
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": prompt
+                        },
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": img_base64
+                            }
+                        }
+                    ]
+                }
+            ],
+            stream=False
+        )
+
+        response = completion.choices[0].message.content
+
+        # Add the new response to history
+        self.add_to_history(image_name, response)
+        # Save history after adding the new entry
+        self.save_history()
+
+        return response
+
+if __name__ == "__main__":
+    import os
+    os.environ['HTTP_PROXY'] = 'http://localhost:8234'
+    os.environ['HTTPS_PROXY'] = 'http://localhost:8234'
+    captioner = Captioner()
+    test_image = "temp_images/3zjz9b3l.jpg"
+    print(captioner.caption(test_image))

src/Database.py

@@ -0,0 +1,278 @@
+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)

src/Founder.py

@@ -0,0 +1,69 @@
+import json
+from collections import defaultdict
+
+class Founder:
+    def __init__(self, filepath='datas/founder.jsonl'):
+        self.filepath = filepath
+        self.datas = {}
+        self.founder2items = defaultdict(list)
+
+        try:
+            self.load_founder()
+        except FileNotFoundError:
+            self.datas = {}
+
+        # Initialize the reverse mapping
+        for word, founder in self.datas.items():
+            self.founder2items[founder].append(word)
+
+    def load_founder(self):
+        """Load founder data from a jsonl file."""
+        with open(self.filepath, 'r', encoding='utf-8') as file:
+            for line in file:
+                data = json.loads(line.strip())
+                self.datas.update(data)
+
+    def save_founder(self):
+        """Save founder data to a jsonl file."""
+        with open(self.filepath, 'w', encoding='utf-8') as file:
+            for word, founder in self.datas.items():
+                file.write(json.dumps({word: founder}, ensure_ascii=False) + '\n')
+
+    def get_founder(self, word):
+        """Get the founder of a given word."""
+        return self.datas.get(word, None)
+
+    def set_founder(self, word, founder, enforce=False):
+        """Set the founder of a word if it's not already set or if enforce is True."""
+        if word in self.datas and not enforce:
+            print(f"Warning: {word} already has a founder: {self.datas[word]}. Use enforce=True to override.")
+        else:
+            self.datas[word] = founder
+            self.founder2items[founder].append(word)
+            self.save_founder()
+
+    def get_all_items_from_founder(self, founder):
+        """Get all words discovered by a specific founder."""
+        return self.founder2items.get(founder, [])
+
+    def get_top_rank(self, top_k=20):
+        """Get the top_k founders with the most discovered words."""
+        sorted_founders = sorted(self.founder2items.items(), key=lambda x: len(x[1]), reverse=True)
+        return sorted_founders[:top_k]
+
+# Example usage:
+# founder = Founder()
+# founder.set_founder('apple', 'Alice')
+# founder.set_founder('banana', 'Bob')
+# print(founder.get_founder('apple'))
+# print(founder.get_all_items_from_founder('Alice'))
+# print(founder.get_top_rank())
+
+if __name__ == '__main__':
+    founder = Founder()
+    founder.set_founder('test_apple', '鲁鲁道祖')
+    founder.set_founder('test_banana', '鲁鲁道祖')
+    founder.set_founder('test_orange', "文钊道祖")
+    print(founder.get_founder('test_apple'))
+    print(founder.get_all_items_from_founder('Alice'))
+    print(founder.get_top_rank())

src/GameMaster.py

@@ -0,0 +1,189 @@
+import os
+from glob import glob
+
+try:
+    from src.Database import Database
+    from src.Captioner import Captioner
+    from src.ImageBase import Imagebase
+    from src.get_major_object import get_major_object, verify_keyword_in_base
+    from src.generate_cultivation import generate_cultivation_with_rag
+except:
+    from Database import Database
+    from Captioner import Captioner
+    from ImageBase import Imagebase
+    from get_major_object import get_major_object, verify_keyword_in_base
+    from generate_cultivation import generate_cultivation_with_rag
+
+
+class GameMaster:
+    def __init__( self ):
+        self.textdb = self.init_textdb()
+
+        self.clip_extractor = self.textdb.clip_extractor
+
+        self.imgdb = self.init_imgdb()
+
+        self.captioner = Captioner()
+
+        self.minimal_image_threshold = 0.9
+
+    def init_textdb( self ):
+        text_db = Database()
+        text_db.init_bge_extractor()
+        text_db.init_clip_extractor()
+        return text_db
+    
+    def init_imgdb( self ):
+        img_db = Imagebase()
+        return img_db
+    
+    def random_image_text_data( self, n = 12 ):
+        random_img_datas = self.imgdb.random_sample(n)
+        # keep image_name and keywords only
+        image_names = [img_data['image_name'] for img_data in random_img_datas]
+        blank_image_path = "datas/blank_item.jpg"
+        for i in range(len(image_names)):
+            if not os.path.exists(image_names[i]):
+                image_names[i] = blank_image_path
+                
+        keywords_zh = [img_data['keyword'] for img_data in random_img_datas]
+        keywords = [img_data['translated_word'] for img_data in random_img_datas]
+        descriptions = []
+
+        for keyword, keyword_zh in zip(keywords, keywords_zh):
+            result = self.textdb.search_by_en_keyword(keyword)
+            if result and "description_in_cultivation" in result:
+                description = result['description_in_cultivation']
+                if "name_in_cultivation" in result:
+                    description = result['name_in_cultivation'] + "--" + description
+                descriptions.append(description)
+            else:
+                descriptions.append("")
+
+        #return tuple of imapge path and description
+        return zip(image_names, descriptions)
+
+    
+    def search_with_path( self, image_path , threshold = None ):
+        # this is a relatively light weight search
+        image_feature = self.clip_extractor.extract_image_from_file(image_path)
+
+        # image_search_result = img_db.search_with_image_name(image_path)
+        image_search_result = self.imgdb.top_k_search(image_feature, top_k=1)
+
+        search_result = None
+
+        if threshold is None:
+            threshold = self.minimal_image_threshold
+
+        if image_search_result and len(image_search_result)>0 and image_search_result[0]['similarity'] > threshold:
+        
+            # try find data with translated_word
+            result = self.textdb.search_by_en_keyword(image_search_result[0]['translated_word'])
+            if result and "name_in_cultivation" in result:
+                search_result = result
+                search_result['similarity'] = image_search_result[0]['similarity']
+            else:
+                print("Warning! Unfound keyword: ", image_search_result[0]['translated_word'])
+
+        # backup_results = None
+        # if search_result is None:
+            # try search with textdb
+        backup_results = self.textdb.top_k_search(image_feature, 'clip_feature', top_k = 5)
+
+        return search_result, backup_results, image_feature
+    
+    def generate_cultivation_data( self, image_path , image_feature, text_search_result ):
+        # this is very expensive
+
+        cultivation_data = None
+        
+        try:
+            caption_response = self.captioner.caption(image_path)
+        except:
+            print("Error occurred while captioning the image ", image_path)
+            return cultivation_data
+        
+        if text_search_result is None:
+            # complete text search
+            text_search_result = self.textdb.top_k_search(image_feature, 'clip_feature', top_k = 5)
+
+        seen = set()
+        keywords = [res['translated_word'] for res in text_search_result if not (res['translated_word'] in seen or seen.add(res['translated_word']))]
+
+        try:
+            json_response = get_major_object(caption_response , keywords)
+        except:
+            print("Error occurred while getting major object from caption ", caption_response)
+            return cultivation_data
+        
+        in_base_data , alt_data = verify_keyword_in_base(json_response , self.textdb )
+
+        if in_base_data is not None:
+            cultivation_data = in_base_data
+
+            # 这意味着找到了一张新的图片，不需要生成额外的词条
+            # required_fields = ['image_name', 'keyword', 'translated_word']
+            image_data = {
+                'image_name': image_path,
+                'keyword': in_base_data['keyword'],
+                'translated_word': in_base_data['translated_word']
+            }
+            self.imgdb.add_image( image_data, True, image_feature )
+        elif alt_data is not None:
+            try:
+                cultivation_data = generate_cultivation_with_rag(alt_data, text_search_result)
+            except:
+                print("Error occurred while generating cultivation data")
+                return cultivation_data
+            
+            new_data = {
+                "keyword": alt_data['keyword'],
+                "name_in_cultivation": cultivation_data['new_name'],
+                "description_in_cultivation": cultivation_data['final_enhanced_description'],
+                "translated_word": alt_data['translated_word'],
+                "description": alt_data['description']
+            }
+            self.textdb.add_data(new_data)
+            print("Added new data to textdb: ", new_data["name_in_cultivation"])
+
+            image_data = {
+                'image_name': image_path,
+                'keyword': new_data['keyword'],
+                'translated_word': new_data['translated_word']
+            }
+            self.imgdb.add_image( image_data, True, image_feature )
+            print("Added new image to imgdb: ", image_data["keyword"])
+
+            cultivation_data = new_data
+
+        return cultivation_data
+
+
+    
+if __name__ == "__main__":
+    os.environ['HTTP_PROXY'] = 'http://localhost:8234'
+    os.environ['HTTPS_PROXY'] = 'http://localhost:8234'
+
+    game_master = GameMaster()
+
+    target_folder="temp_images"
+
+    image_files = glob(os.path.join(target_folder, "*.jpg"))
+
+    for index, image_path in enumerate(image_files):
+        print("index:" , index )
+
+        search_result, backup_results, image_feature = game_master.search_with_path(image_path)
+
+        if search_result:
+            print(search_result)
+
+        break
+
+    test_image_path = "temp_images/向日葵.jpg"
+
+    search_result, backup_results, image_feature = game_master.search_with_path(test_image_path)
+    cultivation_data = game_master.generate_cultivation_data( \
+        test_image_path, image_feature, backup_results )
+    print(cultivation_data)

src/ImageBase.py

@@ -0,0 +1,162 @@
+import pandas as pd
+import os
+from tqdm import tqdm
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
+
+class Imagebase:
+    def __init__(self, parquet_path=None):
+        self.default_parquet_path = 'datas/imagebase.parquet'
+        self.parquet_path = parquet_path or self.default_parquet_path
+        self.datas = None
+        
+        if os.path.exists(self.parquet_path):
+            self.load_from_parquet(self.parquet_path)
+
+        self.clip_extractor = None
+
+    def random_sample(self, num_samples=12):
+        if self.datas is not None:
+            return self.datas.sample(num_samples).to_dict(orient='records')
+        else:
+            return []
+
+    def load_from_parquet(self, parquet_path):
+        self.datas = pd.read_parquet(parquet_path)
+    
+    def save_to_parquet(self, parquet_path=None):
+        parquet_path = parquet_path or self.default_parquet_path
+        if self.datas is not None:
+            self.datas.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 = "D:\\aistudio\\LubaoGithub\\models"
+            self.clip_extractor = CLIPExtractor(model_name="openai/clip-vit-large-patch14", cache_dir=cache_dir)
+    
+    def top_k_search(self, query_feature, top_k=15):
+        if self.datas is None:
+            return []
+        if 'clip_feature' not in self.datas.columns:
+            raise ValueError("clip_feature column not found in the data.")
+        
+        query_feature = np.array(query_feature).reshape(1, -1)
+        attribute_features = np.stack(self.datas['clip_feature'].dropna().values)
+        
+        similarities = cosine_similarity(query_feature, attribute_features)[0]
+        
+        top_k_indices = np.argsort(similarities)[-top_k:][::-1]
+        
+        top_k_results = self.datas.iloc[top_k_indices].copy()
+        
+        top_k_results['similarity'] = similarities[top_k_indices]
+        
+        # Drop the 'clip_feature' column
+        top_k_results = top_k_results.drop(columns=['clip_feature'])
+        
+        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)
+
+    def search_with_image(self, image, if_opencv=False):
+        self.init_clip_extractor()
+
+        img_feature = self.clip_extractor.extract_image(image, if_opencv=if_opencv)
+
+        return self.top_k_search(img_feature)
+    
+    def add_image(self, data, if_save = True, image_feature = None):
+        required_fields = ['image_name', 'keyword', 'translated_word']
+        if not all(field in data for field in required_fields):
+            raise ValueError(f"Data must contain the following fields: {required_fields}")
+        
+        
+        
+        image_name = data['image_name']
+        if image_feature is None:
+            self.init_clip_extractor()
+            data['clip_feature'] = self.clip_extractor.extract_image_from_file(image_name)
+        else:
+            data['clip_feature'] = image_feature
+            
+        if self.datas is None:
+            self.datas = pd.DataFrame([data])
+        else:
+            self.datas = pd.concat([self.datas, pd.DataFrame([data])], ignore_index=True)
+        if if_save:
+            self.save_to_parquet()
+    
+    def add_images(self, datas):
+        for data in datas:
+            self.add_image(data, if_save=False)
+        self.save_to_parquet()
+
+import os
+from glob import glob
+
+def scan_and_update_imagebase(db, target_folder="temp_images"):
+    # 获取target_folder目录下所有.jpg文件
+    image_files = glob(os.path.join(target_folder, "*.jpg"))
+
+    duplicate_count = 0
+    added_count = 0
+
+    for image_path in image_files:
+        # 使用文件名作为keyword
+        keyword = os.path.basename(image_path).rsplit('.', 1)[0]
+        translated_word = keyword  # 可以根据需要调整translated_word
+
+        # 搜索数据库中是否有相似的图片
+        results = db.search_with_image_name(image_path)
+
+        if results and results[0]['similarity'] > 0.9:
+            print(f"Image '{image_path}' is considered a duplicate.")
+            duplicate_count += 1
+        else:
+            new_image_data = {
+                'image_name': image_path,
+                'keyword': keyword,
+                'translated_word': translated_word
+            }
+            db.add_image(new_image_data)
+            print(f"Image '{image_path}' added to the database.")
+            added_count += 1
+
+    print(f"Total duplicate images found: {duplicate_count}")
+    print(f"Total new images added to the database: {added_count}")
+
+if __name__ == '__main__':
+    img_db = Imagebase()
+
+    # 目���目录
+    target_folder = "temp_images"
+    
+    # 扫描并更新数据库
+    scan_and_update_imagebase(img_db, target_folder)
+    
+    # Usage example
+    # img_db = Imagebase()
+
+    # new_image_data = {
+    #     'image_name': "datas/老虎.jpg",
+    #     'keyword': 'tiger',
+    #     'translated_word': '老虎'
+    # }
+    
+    # img_db.add_image(new_image_data)
+
+    # image_path = "datas/老虎.jpg"
+    # results = img_db.search_with_image_name(image_path)
+    # for result in results[:3]:
+    #     print(result)

src/ZhipuClient.py

@@ -0,0 +1,35 @@
+from zhipuai import ZhipuAI
+import os
+
+class ZhipuClient:
+    def __init__(self, api_key_file_path = None):
+        if api_key_file_path is None:
+            cands = ['./datas/zhipu_key.txt', '../datas/zhipu_key.txt']
+            flag = False
+            for cand in cands:
+                if os.path.exists(cand):
+                    api_key_file_path = cand
+                    flag = True
+                    break
+            if not flag:
+                raise ValueError("No valid api key file found.")
+
+        self.api_key = self._load_access_token(api_key_file_path)
+        self.client = ZhipuAI(api_key=self.api_key)
+
+    def _load_access_token(self, file_path):
+        with open(file_path, 'r') as file:
+            return file.read().strip()
+
+    def prompt2response(self, prompt):
+        response = self.client.chat.completions.create(
+            model="glm-4",  # 填写需要调用的模型名称
+            messages=[
+                {"role": "user", "content": prompt}
+            ],
+        )
+        return response.choices[0].message.content
+
+# Usage:
+# zhipu_client = ZhipuClient('../datas/zhipu_key.txt')
+# response = zhipu_client.prompt2response('Your prompt here')

src/generate_cultivation.py

@@ -0,0 +1,183 @@
+import json
+
+def data2reference( top_k_items, output_n = 3 ):
+    outputted_items = set()
+
+    output_str = "#Reference:\n"
+
+    for item in top_k_items:
+        item_in_life = item["keyword"]
+        if item_in_life in outputted_items:
+            continue
+        name_in_cultivation = item["name_in_cultivation"]
+        description_in_cultivation = item["description_in_cultivation"]
+        # output_str += f"name_in_life: {item_in_life}\n"
+        # output_str += f"name_in_cultivation: {name_in_cultivation}\n"
+        # output_str += f"description_in_cultivation: {description_in_cultivation}\n\n"
+        # output with into json format
+        output_data = {
+            "name_in_life": item_in_life,
+            "name_in_cultivation": name_in_cultivation,
+            "description_in_cultivation": description_in_cultivation
+        }
+        output_str += json.dumps(output_data, ensure_ascii=False) + "\n\n"
+
+        outputted_items.add(item_in_life)
+        if len(outputted_items) >= output_n:
+            break
+    return output_str.strip()
+
+
+
+def data2prompt(query_item , top_k_items):
+
+    reference_prompt = data2reference(top_k_items, 3)
+
+    task_prompt1 = "\n请参考Reference中的物品描述，将Input中的输入物品，联系改写成修仙世界中的对应物品\n"
+
+    input_prompt = "# Input:\n"
+    if "keyword" in query_item:
+        input_prompt += f"input_name:{query_item['keyword']}\n"
+        if "description" in query_item:
+            input_prompt += f"description_in_life:{query_item['description']}\n"
+    else:
+        # directly dump query_item
+        input_prompt += json.dumps(query_item, ensure_ascii=False) + "\n"
+
+    CoT_prompt = \
+"""Let's think it step by step，以json形式输出逐个字段。包含以下字段
+- name_in_life: 进一步明确要生成描述的物品名称
+- name_in_cultivation_1: 尝试编写物品在修仙界对应的名称
+- description_in_cultivation_1: 尝试编写物品在修仙界对应的描述
+- echo_1: "我将分析description_in_cultivation_1与Reference中的差异，分析description_in_cultivation_1是否已经足够生动"
+- critique: 相比于Reference中的描述，分析description_in_cultivation_1在哪些方面有所欠缺
+- echo_2: "根据input_name和description_in_cultivation_1，我将分析从物体的哪些属性，可以进一步加强、夸张和修改描述"
+- analysis: 分析从物体的哪些属性，可以进一步加强、夸张和修改描述
+- echo_3: "我将尝试3次，从不同角度加强description_in_cultivation_1的描述"
+- candidate_descriptions: 从不同角度，输出3次不同的加强后的描述
+- analysis_candidates: 分析各个candidates有什么优点
+- echo_4: "根据analysis_candidates，我将merge出一个最终的描述"
+- final_enhanced_description: 通过各个candidates的优点， merge出一个最终的描述
+- echo_5: "我将分析根据final_description,是否简易将物品名称替换为新的名词"
+- name_fit_analysis: 分析item_name是否还匹配final_description的描述，是否需要给input_name起一个更响亮的名字
+- new_name: 如果需要，给input_name起一个更响亮的名字， 如果不需要，则仍然输出name_in_cultivation_1
+"""
+
+    return reference_prompt + task_prompt1 + input_prompt + CoT_prompt
+
+try:
+    from src.ZhipuClient import ZhipuClient
+except:
+    from ZhipuClient import ZhipuClient
+
+zhipu_client = None
+
+
+import json
+
+def markdown_to_json(markdown_str):
+    # 移除Markdown语法中可能存在的标记，如代码块标记等
+    if markdown_str.startswith("```json"):
+        markdown_str = markdown_str[7:-3].strip()
+    elif markdown_str.startswith("```"):
+        markdown_str = markdown_str[3:-3].strip()
+
+    # 将字符串转换为JSON字典
+    json_dict = json.loads(markdown_str)
+
+    return json_dict
+
+import re
+
+def forced_extract(input_str, keywords):
+    result = {key: "" for key in keywords}
+
+    for key in keywords:
+        # 使用正则表达式来查找关键词-值对
+        pattern = f'"{key}":\s*"(.*?)"'
+        match = re.search(pattern, input_str)
+        if match:
+            result[key] = match.group(1)
+    
+    return result
+
+def generate_cultivation_with_rag( query_item, search_result ):
+    global zhipu_client
+    if zhipu_client is None:
+        zhipu_client = ZhipuClient()
+    prompt = data2prompt(query_item, search_result)
+    response = zhipu_client.prompt2response(prompt)
+
+    try:
+        json_response = markdown_to_json(response)
+    except:
+        keyword_list = ["name_in_life", "name_in_cultivation_1","description_in_cultivation_1", "final_enhanced_description", "new_name"]
+        json_response = forced_extract(response, keyword_list)
+
+    if "new_name" not in json_response or json_response["new_name"] == "":
+        if "name_in_cultivation_1" in json_response:
+            json_response["new_name"] = json_response["name_in_cultivation_1"]
+        else:
+            json_response["new_name"] = ""
+
+    if "final_enhanced_description" not in json_response or json_response["final_enhanced_description"] == "":
+        if "description_in_cultivation_1" in json_response:
+            json_response["final_enhanced_description"] = json_response["description_in_cultivation_1"]
+        else:
+            json_response["final_enhanced_description"] = json_response["new_name"]
+    
+
+    return json_response
+
+if __name__ == '__main__':
+    try:
+        from src.Database import Database
+    except:
+        from Database import Database
+    
+    db = Database()
+
+    try:
+        from src.Captioner import Captioner
+    except:
+        from Captioner import Captioner
+
+    import os
+    os.environ['HTTP_PROXY'] = 'http://localhost:8234'
+    os.environ['HTTPS_PROXY'] = 'http://localhost:8234'
+    
+    
+    captioner = Captioner()
+
+    test_image = "temp_images/3or47vg0.jpg"
+    caption_response = captioner.caption(test_image)
+
+    # print(caption_response)
+
+    search_result = db.search_with_image_name( test_image )
+
+    # print(search_result[0].keys())
+    # reference_str = data2reference(search_result, output_n = 3)
+    # print(reference_str)
+
+    seen = set()
+    keywords = [res['translated_word'] for res in search_result if not (res['translated_word'] in seen or seen.add(res['translated_word']))]
+    # print(keywords)
+
+    # prompt = data2prompt(caption_response , keywords)
+    # print(prompt)
+    from get_major_object import get_major_object, verify_keyword_in_base
+
+    json_response = get_major_object(caption_response , keywords)
+    
+    print(json_response)
+
+    print()
+
+    in_base_data , alt_data = verify_keyword_in_base(json_response , db)
+
+    if alt_data is not None:
+        result = generate_cultivation_with_rag(alt_data , search_result)
+        print(result)
+
+

src/get_comments_from_level.py

@@ -0,0 +1,29 @@
+def get_comments_from_level(inbase_similarity_level, inlibrary_similarity_level):
+    name = "{name}"  # Placeholder for the item name
+
+    # Define comments based on the similarity levels
+    comments_mapping = {
+        (0, 0): f"道友，在下才疏学浅不太认识这个东西，我感觉这个东西有点像古籍上所说的{name}，还要找经验更丰富的长老来确定下",
+        (0, 1): f"这个东西似乎有些眼熟，但天机阁尚未收录，古籍上的描述也有些模糊，可能需要长老们进一步鉴定，我猜是{name}",
+        (0, 2): f"这东西颇为罕见，天机阁未曾有过记录，但古籍中的描述与{name}颇为相似，我将呈给长老们以作鉴定",
+        (0, 3): f"想必这一定是{name}吧，虽然天机阁还没有收录过这个东西，倒是和修仙古籍上的记载非常相像。我赶紧拿给长老再鉴定下",
+        (1, 0): f"天机阁的记录中似乎没有这个东西，但我依稀记得古籍中提到过{name}，还需长老进一步确认",
+        (1, 1): f"这个物品有些特别，天机阁的记录不多，古籍中的描述也只是一笔带过，可能是{name}，还需长老鉴定",
+        (1, 2): f"此物颇为罕见，天机阁记录较少，但古籍中的描述与{name}有一定相似之处，长老们或能给出答案",
+        (1, 3): f"虽然古籍中对{name}的描述详细，但天机阁中却鲜有记录，或许这是一件稀世之宝",
+        (2, 0): f"天机阁中对此物知之甚少，但古籍中曾提到{name}，这件物品或许不简单，需长老们鉴定",
+        (2, 1): f"天机阁中对此物的记录不多，古籍中对{name}的描述也有限，但似乎是一件非凡之物",
+        (2, 2): f"这件物品在古籍中有所记载，天机阁也有少量收录，看来是{name}无疑，但还需长老确认",
+        (2, 3): f"虽然在古籍中有记载，天机阁过往有一点点收录，但也算稀世珍宝，{name}确实非凡",
+        (3, 0): f"天机阁中没有记录，但古籍中对{name}的描述颇为详细，这件物品可能是个谜",
+        (3, 1): f"天机阁中记录较少，但古籍中对{name}的描述详尽，这件物品或许有着不同寻常的来历",
+        (3, 2): f"古籍中记载{name}颇多，天机阁中也有所收录，看来这东西并不罕见",
+        (3, 3): f"{name}这种东西很常见啊，天机阁的库房里面都有不少呢"
+    }
+
+    # Return the appropriate comment based on the similarity levels
+    return comments_mapping.get((inbase_similarity_level, inlibrary_similarity_level), "道友，我会给出初步的鉴定") + "。"
+
+# Example usage:
+# comments = get_comments_from_level(2, 3)
+# print(comments)

src/get_major_object.py

@@ -0,0 +1,195 @@
+def data2prompt(caption_response, ref_words):
+
+    ref_word_str = ",".join(ref_words[:5])
+
+    task_prompt = "Based on the following Caption Response, you will output a description of the Major Object's name."
+
+    input_str = "# Caption Response:\n" + caption_response + "\n"
+
+    CoT_prompt = \
+f"""
+Let's think it step by step. Output each field in JSON format. Include the following fields:
+- major_object: From the caption response, identify the major_object. If not present, extract it again from the detailed_description or caption_response.
+- better_major_object: Reread the description in the caption response to see if there's a more suitable word for the major object. If not, still output major_object.
+- echo_1: "I will generate a simple description in about 200 words in English for the better_major_object, introducing what the input object is."
+- description: Generate a WIKI description for the better_major_object (explain what is the better_major_object).
+- major_object_chinese: Translate the better_major_object into Chinese.
+- echo_2: "I will check whether there is synonym of the major_object_chinese in the '{ref_word_str}'."
+- synonym: If present, output the synonym directly; otherwise, output "NOT_INCLUDED."
+- recheck: Based on the content of the Caption Response, determine whether the synonym is accurate. If accurate, output "ACCURATE"; otherwise, output "NOT_ACCURATE."
+"""
+    return task_prompt + input_str + CoT_prompt
+
+
+# def data2prompt(caption_response , ref_words ):
+
+
+#     ref_word_str = ",".join(ref_words[:5])
+
+#     ref_str = "# Reference Word:\n"+ref_word_str+"\n\n"
+
+#     task_prompt = "你将根据下面的Caption Response，输出Major Object的名称描述"
+
+#     input_str = "# Caption Response:\n"+caption_response+"\n"
+
+#     CoT_prompt = \
+# """
+# Let's think it step by step，以json形式输出逐个字段。包含以下字段
+# - major_object: 从caption response中，确认major_object，如果没有，则从detailed_description或者caption_response中重新抽取
+# - better_major_object: 重新阅读caption response中的描述，看看是否有更合适的major object的词语，如果没有则仍然输出major_object
+# - echo_1: "I will generate a simple description in about 200 words in English for the input word, introducing what the input object is"
+# - description: generate the description for the input object
+# - major_object_chinese: 将major_object翻译为中文
+# - echo_2: "我将判断reference word中，是否存在major_object的同义词"
+# - 同义词: 如果存在，则直接输出同义词，否则输出"NOT_INCLUDED"
+# - recheck: 结合Caption Response的内容，判断同义词是否准确，如果准确，则输出"ACCURATE"，否则输出"NOT_ACCURATE"
+# """
+#     return ref_str+task_prompt+input_str+CoT_prompt
+
+try:
+    from src.ZhipuClient import ZhipuClient
+except:
+    from ZhipuClient import ZhipuClient
+
+zhipu_client = None
+
+import json
+
+def markdown_to_json(markdown_str):
+    # 移除Markdown语法中可能存在的标记，如代码块标记等
+    if markdown_str.startswith("```json"):
+        markdown_str = markdown_str[7:-3].strip()
+    elif markdown_str.startswith("```"):
+        markdown_str = markdown_str[3:-3].strip()
+
+    # 将字符串转换为JSON字典
+    json_dict = json.loads(markdown_str)
+
+    return json_dict
+
+import re
+
+def forced_extract(input_str, keywords):
+    result = {key: "" for key in keywords}
+
+    for key in keywords:
+        # 使用正则表达式来查找关键词-值对
+        pattern = f'"{key}":\s*"(.*?)"'
+        match = re.search(pattern, input_str)
+        if match:
+            result[key] = match.group(1)
+    
+    return result
+
+def get_major_object(caption_response, ref_words):
+    global zhipu_client
+    if zhipu_client is None:
+        zhipu_client = ZhipuClient()
+    prompt = data2prompt(caption_response , ref_words)
+    response = zhipu_client.prompt2response(prompt)
+
+    try:
+        json_response = markdown_to_json(response)
+    except:
+        keyword_list = ["major_object", "better_major_object", "description", "major_object_chinese", "synonym", "recheck"]
+        json_response = forced_extract(response, keyword_list)
+
+    return json_response
+
+def verify_keyword_in_base( json_response , database ):
+
+    keyword2verify = []
+    if "better_major_object" in json_response:
+        keyword2verify.append(json_response["better_major_object"].lower())
+
+    if "major_object" in json_response:
+        keyword2verify.append(json_response["major_object"].lower())
+
+    if "recheck" in json_response and json_response["recheck"] == "ACCURATE":
+        if "synonym" in json_response and json_response["synonym"] != "NOT_INCLUDED":
+            keyword2verify.append(json_response["synonym"].lower())
+
+    ans = None
+
+    for keyword in keyword2verify:
+        res = database.search_by_en_keyword(keyword)
+        if res is None:
+            continue
+        ans = res
+        return ans, None
+    
+    if len(keyword2verify) == 0:
+        return None, None
+    
+    # 这里我们需要一个新的data， keyword是中文名， translated_word是英文名，description是英文描述
+    description = keyword2verify[0]
+    if "description" in json_response:
+        description = json_response["description"]
+
+    translated_word = keyword2verify[0]
+
+    keyword = translated_word
+    if "major_object_chinese" in json_response:
+        keyword = json_response["major_object_chinese"]
+
+    data = {
+        "keyword": keyword,
+        "translated_word": translated_word,
+        "description": description
+    }
+
+    return None, data
+
+
+
+
+if __name__ == '__main__':
+    try:
+        from src.Database import Database
+    except:
+        from Database import Database
+    
+    db = Database()
+
+    try:
+        from src.Captioner import Captioner
+    except:
+        from Captioner import Captioner
+
+    import os
+    os.environ['HTTP_PROXY'] = 'http://localhost:8234'
+    os.environ['HTTPS_PROXY'] = 'http://localhost:8234'
+    
+    
+    captioner = Captioner()
+
+    test_image = "temp_images/3or47vg0.jpg"
+    caption_response = captioner.caption(test_image)
+
+    # print(caption_response)
+
+    search_result = db.search_with_image_name( test_image )
+    
+    seen = set()
+    keywords = [res['translated_word'] for res in search_result if not (res['translated_word'] in seen or seen.add(res['translated_word']))]
+    # print(keywords)
+
+    # prompt = data2prompt(caption_response , keywords)
+    # print(prompt)
+
+    json_response = get_major_object(caption_response , keywords)
+    
+    print(json_response)
+
+    print()
+
+    in_base_data , alt_data = verify_keyword_in_base(json_response , db)
+
+    if in_base_data is not None:
+        print(in_base_data)
+
+    if alt_data is not None:
+        print(alt_data)
+
+# {'keyword': '埃菲尔铁塔', 'translated_word': 'eiffel tower', 'description': "The Eiffel Tower is an iconic symbol of Paris and one of the most recognizable stru
+# ower', 'description': "The Eiffel Tower is an iconic symbol of Paris and one of the most recognizable structures in the world. Designed and constructed by the engineer Gustave Eiffel and his company for the 1889 Exposition Universelle (World's Fair) to celebrate the 100th anniversary of the French Revolution, the tower was initially criticized by some of France's leading artists and intellectuals. However, it quickly became a beloved landmark and a symbol of French pride. Standing 324 meters tall, the tower is made of wrought iron and consists of thousands of metal parts, including over 18,000 individual iron rivets. It is renowned for its architectural and engineering design, and it is visited by millions of people each year, making it one of the most visited paid monuments in the world."}

src/text_embedding.py

@@ -0,0 +1,289 @@
+import torch
+from transformers import AutoTokenizer, AutoModel
+import os
+
+class TextExtractor:
+    def __init__(self, model_name, proxy=None):
+        """
+        Initialize the TextExtractor with a specified model and optional proxy settings.
+
+        Parameters:
+        - model_name (str): The name of the pre-trained model to load from HuggingFace Hub.
+        - proxy (str, optional): The proxy address to use for HTTP and HTTPS requests.
+        """
+        if proxy is None:
+            proxy = 'http://localhost:8234'
+
+        if proxy:
+            os.environ['HTTP_PROXY'] = proxy
+            os.environ['HTTPS_PROXY'] = proxy
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+            self.model = AutoModel.from_pretrained(model_name)
+        except:
+            print('try switch on local_files_only')
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=True)
+            self.model = AutoModel.from_pretrained(model_name, local_files_only=True)
+
+        self.model.eval()
+
+    def extract(self, sentences):
+        """
+        Extract sentence embeddings for the provided sentences.
+
+        Parameters:
+        - sentences (list of str): A list of sentences to extract embeddings for.
+
+        Returns:
+        - torch.Tensor: The normalized sentence embeddings.
+        """
+        encoded_input = self.tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
+        
+        with torch.no_grad():
+            model_output = self.model(**encoded_input)
+            sentence_embeddings = model_output[0][:, 0]
+        
+        sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
+        return sentence_embeddings
+    
+import pandas as pd
+def get_qas(excel_file = None):
+
+    defaule_excel_file = 'data/output_fixid.xlsx'
+    if excel_file is None:
+        excel_file = defaule_excel_file
+
+    # 读取Excel文件
+    df = pd.read_excel(excel_file)
+
+    df = df[df["question"].notna()]
+    df = df[df["summary"].notna()]
+
+    datas = []
+
+    # 遍历DataFrame的每一行
+    for index, row in df.iterrows():
+        id = row['id']
+        question = row['question']
+        short_answer = row['summary']
+        category = row['category']
+
+        texts = [question, short_answer]
+
+        data_value = {
+            "texts":texts,
+        }
+        
+        data = {
+            "id":id,
+            "value":data_value
+        }
+
+        datas.append(data)
+
+    return datas
+
+from tqdm import tqdm
+
+def extract_embedding(datas, text_extractor):
+    """
+    Extract embeddings for each item in the provided data.
+
+    Parameters:
+    - datas (list of dict): A list of dictionaries containing text data.
+
+    Returns:
+    - list of dict: The input data with added embeddings.
+    """
+    for data in tqdm(datas):
+        texts = data["value"]["texts"]
+        text = "。".join(texts)
+        embeddings = text_extractor.extract(text)
+        embeddings_list = embeddings.tolist()  # Convert tensor to list of lists
+        data["value"]["embedding"] = embeddings_list
+    return datas
+
+def save_parquet(datas, file_path):
+    """
+    Save the provided data to a Parquet file.
+
+    Parameters:
+    - datas (list of dict): A list of dictionaries containing text data and embeddings.
+    - file_path (str): The path to the output Parquet file.
+    """
+    # Flatten the data for easier conversion to DataFrame
+    flattened_data = []
+    for data in datas:
+        id = data["id"]
+        texts = data["value"]["texts"]
+        text = "。".join(texts)
+        embedding = data["value"]["embedding"]
+        flattened_data.append({
+            "id": id,
+            "text": text,
+            "embedding": embedding
+        })
+    
+    # Create DataFrame
+    df = pd.DataFrame(flattened_data)
+    
+    # Save DataFrame to Parquet
+    df.to_parquet(file_path, index=False)
+
+import pandas as pd
+import os
+
+def get_id2embedding(regen=False, parquet_file='datas/qa_with_embedding.parquet'):
+    """
+    Get a dictionary mapping IDs to embeddings. Regenerate embeddings if specified.
+
+    Parameters:
+    - parquet_file (str): The path to the Parquet file.
+    - regen (bool): Whether to regenerate embeddings.
+
+    Returns:
+    - dict: A dictionary mapping IDs to list of float embeddings.
+    """
+    if regen or not os.path.exists(parquet_file):
+        print("Regenerating embeddings...")
+        # Example usage:
+        model_name = 'BAAI/bge-small-zh-v1.5'
+        text_extractor = TextExtractor(model_name)
+        
+        datas = get_qas()
+        print("Extracting embeddings for", len(datas), "data items")
+        
+        datas = extract_embedding(datas, text_extractor)
+        save_parquet(datas, parquet_file)
+    
+    df = pd.read_parquet(parquet_file)
+    
+    id2embedding = {}
+    for index, row in df.iterrows():
+        id = row['id']
+        embedding = row['embedding']
+        id2embedding[id] = embedding[0]
+    
+    return id2embedding
+
+import torch
+from sklearn.metrics.pairwise import cosine_similarity
+import heapq
+
+def __get_id2top30map(id2embedding):
+    """
+    Get a dictionary mapping IDs to their top 30 nearest neighbors based on cosine similarity.
+
+    Parameters:
+    - id2embedding (dict): A dictionary mapping IDs to list of float embeddings.
+
+    Returns:
+    - dict: A dictionary mapping each ID to a list of the top 30 nearest neighbor IDs.
+    """
+    ids = list(id2embedding.keys())
+    embeddings = torch.tensor([id2embedding[id] for id in ids])
+
+    # Compute cosine similarity matrix
+    cos_sim_matrix = cosine_similarity(embeddings)
+
+    id2top30map = {}
+    for i, id in enumerate(ids):
+        # Get the similarity scores for the current ID
+        sim_scores = cos_sim_matrix[i]
+        
+        # Get the top 30 indices (excluding the current ID itself)
+        top_indices = heapq.nlargest(31, range(len(sim_scores)), key=lambda x: sim_scores[x])
+        top_indices.remove(i)  # Remove the index of the current ID
+        
+        # Map the indices back to IDs
+        top_30_ids = [ids[idx] for idx in top_indices[:30]]
+        
+        id2top30map[id] = top_30_ids
+    
+    return id2top30map
+
+import pickle
+
+def get_id2top30map( id2embedding = None ):
+    default_save_pkl = "data/id2top30map.pkl"
+    if id2embedding is None:
+        if os.path.exists(default_save_pkl):
+            with open(default_save_pkl, 'rb') as f:
+                id2top30map = pickle.load(f)
+        else:
+            print("No embedding found, generating new one...")
+            id2embedding = get_id2embedding(regen=False)
+            id2top30map = __get_id2top30map(id2embedding)
+            with open(default_save_pkl, 'wb') as f:
+                pickle.dump(id2top30map, f)
+    else:
+        id2top30map = __get_id2top30map(id2embedding)
+
+    return id2top30map
+        
+
+
+if __name__ == '__main__':
+    if False:
+        # Example usage:
+        model_name = 'BAAI/bge-small-zh-v1.5'
+        sentences = ["样例数据-1", "样例数据-2"]
+
+        text_extractor = TextExtractor(model_name)
+        embeddings = text_extractor.extract(sentences)
+        print("Sentence embeddings:", embeddings)
+
+        datas = get_qas()
+
+        print("extract embedding for ", len(datas), " datas")
+
+        datas = extract_embedding(datas, text_extractor )
+
+        default_parquet_save_name = "data/qa_with_embedding.parquet"
+
+        save_parquet(datas, default_parquet_save_name)
+    if True:
+        id2embedding = get_id2embedding(regen=False)
+        print(len(id2embedding[4]))
+        id2top30map = get_id2top30map( None )
+        print("ID to Top 30 Neighbors dictionary:", id2top30map[4])
+
+    if True:
+
+        start_id = 332
+        visited_ids = [start_id]
+        current_queue = [start_id]
+
+        expend_num = 5
+
+        for iteration in range(10):
+            current_node = current_queue.pop(0)
+            top30 = id2top30map[current_node]
+            current_expend = []
+            for id in top30:
+                if id not in visited_ids:
+                    visited_ids.append(id)
+                    current_queue.append(id)
+                    current_expend.append(id)
+                    if len(current_expend) >= expend_num:
+                        break
+            display_text = f"{current_node} | ->" + ",".join([str(i) for i in current_expend])
+            print(display_text)
+
+        from get_qa_and_image import get_qa_and_image
+        image_datas = get_qa_and_image()
+
+        id2index = {}
+
+        for i, data in enumerate(image_datas):
+            id2index[data['id']] = i
+
+        indexes = [id2index[i] for i in visited_ids if i in id2index]
+        image_names = [image_datas[index]['value']['image'] for index in indexes]
+
+        target_copy_folder = "data/asso_collection"
+        
+        import shutil
+        # copy image into target_copy_folder
+        for image_name in image_names:
+            shutil.copy(image_name, target_copy_folder)