src/GameMaster.py

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)