Upload folder using huggingface_hub

modules/image_recommender/__init__.py

@@ -0,0 +1,3 @@
+"""
+Image Recommender package initialization.
+"""

modules/image_recommender/analysis.py

@@ -0,0 +1,73 @@
+# check results_m folder and get distribution
+import os, json
+json_path = './results_m/'
+
+available_list = ['image_content', 'topic', 'data_facts', 'color_style', 'size']
+data_fact_list = ['increasing', 'decreasing', 'highlight', 'deny', 'maximum', 'minimum', 'comparison', 'none']
+color_style_list = ['monochrome', 'colorful', 'grayscale', 'dual-color']
+size_list = ['icon', 'clipart', 'background']
+
+distribution = {
+    'data_facts': {},
+    'color_style': {},
+    'size': {}
+}
+
+cases = {
+    'data_facts': {},
+    'color_style': {},
+    'size': {}
+}
+ct = 0
+case_limit = 10
+image_pathes = []
+with open('image_pathes.json', 'r') as f:
+    image_pathes = json.load(f)
+
+
+for i, file in enumerate(os.listdir(json_path)):
+    if not file.endswith('.json'):
+        continue
+    
+    # load json
+    with open(f'{json_path}{file}', 'r') as f:
+        data = json.load(f)
+    
+    ct += 1
+    distribution['color_style'][data['color_style']] = distribution['color_style'].get(data['color_style'], 0) + 1
+    distribution['size'][data['size']] = distribution['size'].get(data['size'], 0) + 1
+    distribution['data_facts'][data['data_facts']] = distribution['data_facts'].get(data['data_facts'], 0) + 1
+
+    file_id = int(file.split('.')[0])
+    cases['color_style'][data['color_style']] = cases['color_style'].get(data['color_style'], [])
+    if len(cases['color_style'][data['color_style']]) < case_limit:
+        cases['color_style'][data['color_style']].append((file_id, image_pathes[file_id]))
+    cases['size'][data['size']] = cases['size'].get(data['size'], [])
+    if len(cases['size'][data['size']]) < case_limit:
+        cases['size'][data['size']].append((file_id, image_pathes[file_id]))
+    cases['data_facts'][data['data_facts']] = cases['data_facts'].get(data['data_facts'], [])
+    if len(cases['data_facts'][data['data_facts']]) < case_limit:
+        cases['data_facts'][data['data_facts']].append((file_id, image_pathes[file_id]))
+
+# save json
+with open('distribution.json', 'w') as f:
+    json.dump(distribution, f, indent=4)
+with open('cases.json', 'w') as f:
+    json.dump(cases, f, indent=4)
+
+# move cases to './cases/'
+cases_path = './cases/'
+if not os.path.exists(cases_path):
+    os.makedirs(cases_path)
+for key in cases:
+    for case in cases[key]:
+        for i, info in enumerate(cases[key][case]):
+            file_id, image_path = info
+            target_path = f'{cases_path}{key}/{case}/'
+            if not os.path.exists(target_path):
+                os.makedirs(target_path)
+            os.system(f'cp {image_path} {target_path}{file_id}.jpg')
+
+print(ct)
+
+

modules/image_recommender/check.py

@@ -0,0 +1,194 @@
+# 1. 扫描当前目录下所有image文件 包含png jpg jpeg webp，进行编号 存储下list
+import os, json
+import uuid  # 添加导入uuid模块
+image_pathes = []
+root = '/data/lizhen/resources/image'
+image_root = os.path.join(root, 'images')
+image_path_file = os.path.join(root, 'image_pathes.txt')
+result_map_file = os.path.join(root, 'result_map.txt')
+
+# Load existing result mappings
+result_map = {}
+if os.path.exists(result_map_file):
+    with open(result_map_file, 'r') as f:
+        for line in f:
+            filename, result_file = line.strip().split(',')
+            result_map[filename] = result_file
+
+existing_paths = set()
+if os.path.exists(image_path_file):
+    with open(image_path_file, 'r') as f:
+        existing_paths = set(line.strip() for line in f.readlines())
+        image_pathes = [os.path.join(image_root, path) for path in existing_paths]
+        print(f"Loaded {len(image_pathes)} existing image paths")
+
+print('Scanning images...')
+new_paths = []
+for root_dir, dirs, files in os.walk(image_root):
+    for file in files:
+        if file.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
+            abs_path = os.path.join(root_dir, file)
+            rel_path = os.path.relpath(abs_path, image_root)
+            if rel_path not in existing_paths:
+                new_paths.append(rel_path)
+                image_pathes.append(abs_path)
+
+if new_paths:
+    print(f"Found {len(new_paths)} new images")
+    with open(image_path_file, 'a') as f:
+        for path in new_paths:
+            f.write(path + '\n')
+
+print(f"Total images: {len(image_pathes)}")
+
+from openai import OpenAI
+from PIL import Image
+import base64
+from io import BytesIO
+import requests
+from concurrent.futures import ThreadPoolExecutor
+import threading
+
+client = OpenAI(
+    api_key=os.getenv("OPENAI_API_KEY") or os.getenv("AIHUBMIX_API_KEY", ""),
+    base_url=os.getenv("OPENAI_BASE_URL", "https://aihubmix.com/v1")
+)
+
+def resize_image(img, max_size=512):
+    width, height = img.size
+    ratio = min(max_size / width, max_size / height)
+    if ratio >= 1:
+        return img
+    new_width = int(width * ratio)
+    new_height = int(height * ratio)
+    resized_img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+    return resized_img
+
+def repaint_image(img):
+    # rapaint transparent area with white color
+    img = img.convert('RGBA')
+    data = img.getdata()
+    new_data = []
+    for item in data:
+        if item[3] == 0:
+            new_data.append((255, 255, 255, 255))
+        else:
+            new_data.append(item)
+    img.putdata(new_data)
+    img = img.convert('RGB')
+    # img.save('temp.png')
+    return img
+    
+def image_to_base64(image_path, show=False, target_size=512):
+    with Image.open(image_path) as img:
+        # img = img.resize(size)
+        img = resize_image(img, 512)
+        img = repaint_image(img)
+        buffered = BytesIO()
+        img.save(buffered, format="PNG")
+        img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
+        return img_base64
+
+wwxxhh = 0
+def ask_image(prompt, image_data):
+    number_of_trials = 0
+    while number_of_trials < 5:
+        try:
+            response = requests.post(
+                "https://aihubmix.com/v1/chat/completions",
+                headers={
+                    "Authorization": f"Bearer {os.getenv('OPENAI_API_KEY') or os.getenv('AIHUBMIX_API_KEY', '')}",
+                    "Content-Type": "application/json"
+                },
+                json={
+                    "model": "gemini-2.0-flash",
+                    "messages": [{
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": prompt},
+                            {
+                                "type": "image_url",
+                                "image_url": {
+                                    "url": f"data:image/jpeg;base64,{image_data}"
+                                }
+                            }
+                        ]
+                    }]
+                }
+            )
+            if response.status_code == 200:
+                return response.json()['choices'][0]['message']['content']
+            else:
+                print(f"Error status code: {response.status_code}")
+                number_of_trials += 1
+        except Exception as e:
+            print(f"Request error: {e}")
+            number_of_trials += 1
+    
+    return 'Error!'
+
+# 2. 读取prompt.json文件，读取整个作为字符串，逐个读取image文件，调用ask_image函数，将返回的结果存储下来
+import json
+with open('modules/image_recommender/prompt.json', 'r') as f:
+    prompt = f.read()
+# print(prompt)
+
+results_path = os.path.join(root, 'results')
+if not os.path.exists(results_path):
+    os.makedirs(results_path)
+
+def process_image(args):
+    i, image_path, prompt, results_path = args
+    rel_path = os.path.relpath(image_path, image_root)
+    
+    # 使用UUID生成��机文件名，而不是使用索引
+    random_filename = str(uuid.uuid4())
+    target_path = os.path.join(results_path, f'{random_filename}.json')
+    
+    # Skip if already processed
+    if rel_path in result_map:
+        print(f'Skipping {i+1}/{len(image_pathes)} (already exists in result map)')
+        return
+    
+    print(f'Processing {i+1}/{len(image_pathes)}')
+    try:
+        image_data = image_to_base64(image_path)
+        result = ask_image(prompt, image_data)
+        try:
+            result = json.loads(result)
+        except:
+            result = result.replace('```json', '').replace('```', '')
+            result = json.loads(result)
+        
+        # Add filename and remove explanation
+        result['filename'] = rel_path
+        if 'explanation' in result:
+            del result['explanation']
+        
+        with open(target_path, 'w') as f:
+            json.dump(result, f)
+            
+        # Update result mapping
+        with open(result_map_file, 'a') as f:
+            f.write(f"{rel_path},{target_path}\n")
+        result_map[rel_path] = target_path
+    except Exception as e:
+        print(f'Failed to process {i+1}/{len(image_pathes)}: {str(e)}')
+
+# Pre-scan for existing results
+print("Pre-scanning for existing results...")
+results_path = os.path.join(root, 'results')
+if not os.path.exists(results_path):
+    os.makedirs(results_path)
+
+# Filter out already processed images
+image_pathes = [path for path in image_pathes if os.path.relpath(path, image_root) not in result_map]
+print(f"Remaining images to process: {len(image_pathes)}")
+
+# Main processing loop with thread pool
+num_threads = 20
+with ThreadPoolExecutor(max_workers=num_threads) as executor:
+    tasks = [
+        (i, image_pathes[i], prompt, results_path) for i in range(len(image_pathes))
+    ]
+    executor.map(process_image, tasks)

modules/image_recommender/create_index.py

@@ -0,0 +1,243 @@
+import os
+import sys
+import json
+import numpy as np
+import faiss
+from tqdm import tqdm
+from utils.model_loader import ModelLoader
+
+# Add the project root directory to Python path
+project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.append(project_root)
+
+from typing import Optional, List, Dict, Union
+
+class ImageRecommender:
+    def __init__(self, embed_model_path: str):
+        self.model = ModelLoader.get_model(embed_model_path)
+        self.index = None
+        self.image_paths = []
+        self.image_data = []
+        self.icon_indices = []  # Store indices of icon images
+        self.clipart_indices = []  # Store indices of clipart images
+        
+    def create_index(self, image_list_path: str, image_resource_path: str):
+        """Create FAISS index from image embeddings"""
+        # Read image paths
+        with open(image_list_path, 'r') as f:
+            self.image_paths = [line.strip().split(',') for line in f.readlines()]
+            
+        # Load and process each image's data
+        embeddings = []
+        for idx, (image_path, json_path) in enumerate(tqdm(self.image_paths)):
+            # Construct full path to the image's JSON data
+            #json_path = os.path.join(image_resource_path, f'results/{idx}.json')
+            
+            if not os.path.exists(json_path):
+                continue
+            
+            try:    
+                with open(json_path, 'r') as f:
+                    try:
+                        image_info = json.load(f)
+                    except json.JSONDecodeError as e:
+                        print(f"Error decoding JSON at index {idx}: {str(e)}")
+                        continue
+                    
+                # Combine multiple fields for better semantic representation
+                semantic_text = ""
+                try:
+                    image_content = image_info.get('image_content', '')
+                    topic = image_info.get('topic', '')
+                    explanation = image_info.get('explanation', '')
+                    
+                    if image_content:
+                        semantic_text += f"{image_content}"
+                    if topic:
+                        semantic_text += f". {topic}"
+                    if explanation:
+                        semantic_text += f". {explanation}"
+                    
+                    if not semantic_text:
+                        semantic_text = "Image without description"
+                        print(f"Warning: Missing semantic information for image at index {idx}")
+                except Exception as e:
+                    semantic_text = "Image without description"
+                    print(f"Error processing semantic text at index {idx}: {str(e)}")
+                
+                # Generate embedding
+                try:
+                    embedding = self.model.encode(semantic_text)
+                    embeddings.append(embedding)
+                    self.image_data.append(image_info)
+                    
+                    # Store indices based on image type
+                    if image_info.get('icon_or_clipart') == 'icon':
+                        self.icon_indices.append(len(embeddings) - 1)
+                    elif image_info.get('icon_or_clipart') == 'clipart':
+                        self.clipart_indices.append(len(embeddings) - 1)
+                except Exception as e:
+                    print(f"Error generating embedding at index {idx}: {str(e)}")
+                    continue
+                    
+            except Exception as e:
+                print(f"Error processing image at index {idx}: {str(e)}")
+                continue
+            
+        self.image_paths = [image_path for image_path, _ in self.image_paths]
+        embeddings = np.array(embeddings).astype('float32')
+        
+        # Create and train FAISS index
+        dimension = embeddings.shape[1]
+        self.index = faiss.IndexFlatL2(dimension)
+        self.index.add(embeddings)
+        
+    def save_index(self, index_path, data_path):
+        """Save the FAISS index and associated data"""
+        faiss.write_index(self.index, index_path)
+        
+        # Save image data and indices
+        with open(data_path, 'w') as f:
+            json.dump({
+                'paths': self.image_paths,
+                'data': self.image_data,
+                'icon_indices': self.icon_indices,
+                'clipart_indices': self.clipart_indices
+            }, f)
+            
+    def load_index(self, index_path, data_path):
+        """Load the FAISS index and associated data"""
+        self.index = faiss.read_index(index_path)
+        
+        with open(data_path, 'r') as f:
+            data = json.load(f)
+            self.image_paths = data['paths']
+            self.image_data = data['data']
+            self.icon_indices = data['icon_indices']
+            self.clipart_indices = data['clipart_indices']
+            
+    def search(self, query_text: str, new_index = None, new_data = None, top_k: int = 5, image_type: Optional[str] = None) -> List[Dict]:
+        """
+        Search for similar images based on query text
+        
+        Args:
+            query_text: The text query to search for
+            new_index: Optional new FAISS index to search in addition
+            new_data: Optional new data associated with new_index
+            top_k: Number of results to return
+            image_type: Optional filter for image type ('icon' or 'clipart')
+            
+        Returns:
+            List of dictionaries containing image information and similarity scores
+        """
+        if self.index is None:
+            raise ValueError("Index not loaded. Please load the index first.")
+            
+        # Generate query embedding
+        query_embedding = self.model.encode(query_text)
+        query_embedding = np.array([query_embedding]).astype('float32')
+        
+        results = []
+        
+        # 搜索旧索引
+        # Determine which indices to search in
+        if image_type == 'icon':
+            search_indices = self.icon_indices
+        elif image_type == 'clipart':
+            search_indices = self.clipart_indices
+        else:
+            search_indices = None
+            
+        if search_indices:
+            # Create a subset index for the specific image type
+            subset_index = faiss.IndexFlatL2(self.index.d)
+            subset_index.add(self.index.reconstruct_n(0, self.index.ntotal)[search_indices])
+            
+            # Search in the subset index
+            distances, indices = subset_index.search(query_embedding, top_k)
+            # Map back to original indices
+            indices = [search_indices[i] for i in indices[0]]
+            distances = distances[0]
+        else:
+            # Search in the full index
+            distances, indices = self.index.search(query_embedding, top_k)
+            indices = indices[0]
+            distances = distances[0]
+        
+        # 添加旧索引结果
+        for idx, distance in zip(indices, distances):
+            if idx < len(self.image_paths):  # Ensure index is valid
+                results.append({
+                    'image_path': self.image_paths[idx],
+                    'image_data': self.image_data[idx],
+                    'distance': float(distance)
+                })
+                
+        # 如果是icon类型且有新索引,搜索新索引
+        if image_type == 'icon' and new_index is not None and new_data is not None:
+            new_distances, new_indices = new_index.search(query_embedding, top_k)
+            new_indices = new_indices[0]
+            new_distances = new_distances[0]
+            
+            # 添加新索引结果
+            for idx, distance in zip(new_indices, new_distances):
+                if idx < len(new_data['index']):
+                    data = new_data['index'][str(idx)]
+                    # print("data: ", data)
+                    results.append({
+                        'image_path': data["path"],
+                        'image_data': data["data"],
+                        'distance': float(distance) + 0.1
+                    })
+                    
+        # 按距离排序并返回前top_k个结果
+        results.sort(key=lambda x: x['distance'])
+        return results[:top_k]
+
+def main(image_list_path: str = None, 
+         image_resource_path: str = None,
+         index_path: str = None,
+         data_path: str = None,
+         embed_model_path: str = None,
+         force: bool = False):
+    """
+    Create and save the image index
+    
+    Args:
+        image_list_path: Path to the file containing list of image paths
+        image_resource_path: Path to the directory containing image resources
+        index_path: Path to save the FAISS index
+        data_path: Path to save the image data
+        embed_model_path: Path to the sentence embedding model
+        force: Whether to force rebuild the index if it exists
+    """
+    # Check if index exists and handle force flag
+    if os.path.exists(index_path) and not force:
+        print(f"Index file {index_path} already exists. Use --force to rebuild.")
+        return 0
+        
+    recommender = ImageRecommender(embed_model_path)
+    recommender.create_index(image_list_path, image_resource_path)
+    recommender.save_index(index_path, data_path)
+    print("Index built successfully!")
+    return 0
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser(description='Build image index using FAISS and SentenceTransformer')
+    parser.add_argument('--image_list_path', type=str, required=True, help='Path to the file containing list of image paths')
+    parser.add_argument('--image_resource_path', type=str, required=True, help='Path to the directory containing image resources')
+    parser.add_argument('--index_path', type=str, required=True, help='Path to save the FAISS index')
+    parser.add_argument('--data_path', type=str, required=True, help='Path to save the image data')
+    parser.add_argument('--embed_model_path', type=str, required=True, help='Path to the sentence embedding model')
+    parser.add_argument('--force', action='store_true', help='Force rebuild even if index exists')
+    
+    args = parser.parse_args()
+    main(
+        image_list_path=args.image_list_path,
+        image_resource_path=args.image_resource_path,
+        index_path=args.index_path,
+        data_path=args.data_path,
+        embed_model_path=args.embed_model_path,
+        force=args.force
+    )

modules/image_recommender/image_recommender.py

@@ -0,0 +1,732 @@
+import json
+import os
+import requests
+from typing import Dict, List, Optional
+import pandas as pd
+import faiss
+from logging import getLogger
+from utils.model_loader import ModelLoader
+logger = getLogger(__name__)
+
+class ImageRecommender:
+    def __init__(self, embed_model_path: str = None, resource_path: str = None, data_path: str = None, index_path: str = None, base_url: str = None, api_key: str = None):
+        self.index_builder = None
+        if embed_model_path and data_path and index_path:
+            from .create_index import ImageRecommender as IndexBuilder
+            self.index_builder = IndexBuilder(embed_model_path)
+            self.index_builder.load_index(index_path, data_path)
+        self.model = ModelLoader.get_model(embed_model_path)
+        self.base_url = base_url
+        self.request_url = self.base_url + '/chat/completions'
+        self.api_key = api_key
+        self.resource_path = resource_path
+        self.normal_icons = os.path.join(resource_path, 'images')
+        self.special_icons = os.path.join(resource_path, 'special_icons')
+        self.special_icon_index = os.path.join(self.special_icons, 'data.json')
+        self.newicon_path = os.path.join(resource_path, 'attribute_icons')
+        self.newicon_data_path = os.path.join(self.newicon_path, 'index.json')
+        self.newicon_data = json.load(open(self.newicon_data_path))
+        self.newicon_faiss = os.path.join(self.newicon_path, 'faiss.index')
+        self.newicon_index = faiss.read_index(self.newicon_faiss)
+        self.special_categories = ["country", "emotion"]
+        
+    def create_query_text_for_value(self, input_data: Dict, group_value: str, group_col: str) -> str:
+        """Create a text query from input data for finding similar images."""
+        # titles = input_data.get("titles", {})
+        text_parts = []
+
+        text_parts.append(group_value)
+        columns = [col["name"] + " (" + col["description"] + ")" for col in input_data.get("data", {}).get("columns", []) if col["name"] == group_col][0]
+        text_parts.append(" in " + columns)
+
+        # if "main_title" in titles:
+        #    text_parts.append(titles["main_title"])
+
+        return "; ".join(text_parts)
+    
+
+    def create_query_text(self, input_data: Dict) -> str:
+        """Create a text query from input data for finding similar images."""
+        titles = input_data.get("titles", {})
+        data_facts = input_data.get("datafacts", [])
+        
+        text_parts = []
+        
+        # Add metadata
+        if "main_title" in titles:
+            text_parts.append(titles["main_title"])
+        if "sub_title" in titles:
+            text_parts.append(titles["sub_title"])
+        # columns = [col["name"] + " (" + col["description"] + ")" for col in input_data.get("data", {}).get("columns", [])]
+        # text_parts.append("Columns: " + ", ".join(columns))
+
+        for fact in data_facts[:1]:
+            text_parts.append(f'{fact["subtype"]} {fact["type"]}, {fact["annotation"]}')
+            
+        return "; ".join(text_parts)
+
+    def select_optimal_icons(self, group_icons: Dict, embeddings: Dict) -> Dict:
+        """
+        Select optimal icons for each group value maximizing pairwise similarity.
+        
+        Args:
+            group_icons: Dictionary mapping group values to lists of icon candidates
+            embeddings: Dictionary mapping image paths to their embeddings
+        
+        Returns:
+            Dictionary mapping group values to selected optimal icons
+        """
+        from scipy.spatial.distance import cosine
+        import numpy as np
+        
+        selected_icons = {}
+        used_images = set()
+        
+        # Sort group values by number of candidates (ascending)
+        sorted_groups = sorted(group_icons.keys(), key=lambda x: len(group_icons[x]))
+        
+        for group_value in sorted_groups:
+            candidates = group_icons[group_value]
+            best_score = float('-inf')
+            best_candidate = None
+            
+            for candidate in candidates:
+                if candidate['image_path'] in used_images:
+                    continue
+                    
+                # If this is the first selection, just use similarity score
+                if not selected_icons:
+                    if candidate['similarity_score'] > best_score:
+                        best_score = candidate['similarity_score']
+                        best_candidate = candidate
+                    continue
+                
+                # Calculate average dissimilarity with previously selected icons
+                candidate_embedding = embeddings[candidate['image_path']]
+                dissimilarities = []
+                for selected in selected_icons.values():
+                    selected_embedding = embeddings[selected['image_path']]
+                    # Use cosine distance directly as dissimilarity measure
+                    dissimilarity = cosine(candidate_embedding, selected_embedding)
+                    dissimilarities.append(dissimilarity)
+                
+                # Combine original similarity score with average dissimilarity
+                # Higher dissimilarity is better for diversity
+                avg_dissimilarity = np.mean(dissimilarities)
+                combined_score = 0.3 * candidate['similarity_score'] + 0.7 * avg_dissimilarity
+                
+                if combined_score > best_score:
+                    best_score = combined_score
+                    best_candidate = candidate
+            
+            if best_candidate:
+                selected_icons[group_value] = best_candidate
+                used_images.add(best_candidate['image_path'])
+        return selected_icons
+
+    def get_semantic_text(self, image_data: Dict) -> str:
+        """
+        Get semantic text from image data
+        """
+        semantic_text = ""
+        image_content = image_data.get('image_content', '')
+        topic = image_data.get('topic', '')
+        color_style = image_data.get('color_style', '')
+        name = image_data.get('name', '')
+
+        if image_content:
+            semantic_text += f"{image_content}"
+        if topic:
+            semantic_text += f". {topic}"
+        if color_style:
+            semantic_text += f". {color_style}"
+        if name:
+            semantic_text += f". {name}"
+            
+        return semantic_text
+
+    def process_special_icons(self, category: str, unique_values: List) -> Optional[Dict]:
+        """
+        Process icons for special categories (country, sports)
+        
+        Args:
+            category: Category name (e.g. "country", "sports")
+            unique_values: List of unique values to match icons for
+            
+        Returns:
+            Dictionary mapping values to icons, or None if processing fails
+        """
+        try:
+            # Load special icons data
+            with open(self.special_icon_index, 'r') as f:
+                special_icons = json.load(f)
+                
+            if category not in special_icons:
+                return None
+                
+            # Randomly select a type
+            import random
+            available_types = list(special_icons[category].keys())
+            if category == "country":
+                selected_type = "circle"
+            else:
+                selected_type = random.choice(available_types)
+            icons = special_icons[category][selected_type]
+            
+            model = self.model
+            
+            icon_names = [icon["name"] for icon in icons]
+            icon_embeddings = model.encode(icon_names)
+            
+            result = {}
+            for value in unique_values:
+                # Convert numpy types to Python native types
+                value_str = str(value.item() if hasattr(value, 'item') else value)
+                value_embedding = model.encode([value_str])[0]
+                
+                # Calculate similarities
+                from scipy.spatial.distance import cosine
+                similarities = [1 - cosine(value_embedding, icon_emb) for icon_emb in icon_embeddings]
+                best_idx = max(range(len(similarities)), key=lambda i: similarities[i])
+                
+                result[value_str] = {
+                    "image_path": os.path.join(self.special_icons, icons[best_idx]["path"]),
+                    "similarity_score": similarities[best_idx]
+                }
+            
+            return result
+            
+        except Exception as e:
+            logger.error(f"Error processing special icons: {str(e)}")
+            return None
+
+    def process_normal_icons(self, input_data: Dict, unique_values: List, group_col: str) -> Dict:
+        """
+        Process icons for normal categories using image search
+        
+        Args:
+            input_data: Input data dictionary
+            unique_values: List of unique values to find icons for
+            
+        Returns:
+            Dictionary mapping values to icons
+        """
+        all_group_icons = {}
+        for value in unique_values:
+            group_query = self.create_query_text_for_value(input_data, str(value), group_col)
+            group_icons = self.index_builder.search(group_query, top_k=20, new_index=self.newicon_index, new_data=self.newicon_data, image_type='icon')
+
+            all_group_icons[str(value)] = [
+                {
+                    "image_path": img["image_path"],
+                    "image_data": img["image_data"],
+                    "similarity_score": 1.0 / (1.0 + float(img["distance"]))
+                }
+                for img in group_icons
+            ]
+        
+        # Get embeddings for all candidate images using the model directly
+        all_images = {img["image_path"]: img["image_data"]
+                     for group_candidates in all_group_icons.values()
+                     for img in group_candidates}
+        
+        for key in all_images:
+            all_images[key] = self.model.encode(self.get_semantic_text(all_images[key]))
+        
+        # Select optimal icons
+        return self.select_optimal_icons(all_group_icons, all_images)
+    def post_process_image(self, image_path: str, max_size: int = 768) -> str:
+        """
+        Post-process the recommended image:
+        1. Convert relative path to absolute path
+        2. Remove white background that connects to edges while preserving inner fills
+        3. Resize to specified resolution
+        4. Convert to base64 string
+        
+        Args:
+            image_path: Relative path to the image
+            max_size: Maximum dimension for resizing (default: 768)
+                
+        Returns:
+            Base64 encoded string of the processed image
+        """
+        try:
+            from PIL import Image
+            import numpy as np
+            import base64
+            import io
+            from scipy.ndimage import label
+            
+            # Get absolute path
+            abs_path = os.path.join(self.normal_icons, image_path)
+            if not os.path.exists(abs_path):
+                abs_path = os.path.join(self.newicon_path, image_path)
+                if not os.path.exists(abs_path):
+                    logger.error(f"Image not found: {abs_path}")
+                    return ""
+        
+            # Open and convert image
+            img = Image.open(abs_path)
+            
+            # Resize image while maintaining aspect ratio
+            if img.width > max_size or img.height > max_size:
+                ratio = min(max_size / img.width, max_size / img.height)
+                new_width = int(img.width * ratio)
+                new_height = int(img.height * ratio)
+                img = img.resize((new_width, new_height), Image.LANCZOS)
+            
+            if img.mode != 'RGBA':
+                img = img.convert('RGBA')
+            
+            # Convert to numpy array for processing
+            data = np.array(img)
+            
+            # Create mask for white background (with tolerance)
+            white_color = np.array([255, 255, 255])
+            tolerance = 30
+            color_dists = np.sqrt(np.sum((data[..., :3] - white_color)**2, axis=2))
+            bg_mask = color_dists < tolerance
+            
+            # Find connected components
+            labeled, num_components = label(bg_mask)
+            
+            # Create border mask
+            border_mask = np.zeros_like(bg_mask, dtype=bool)
+            border_mask[0, :] = border_mask[-1, :] = True
+            border_mask[:, 0] = border_mask[:, -1] = True
+            
+            # Find components that touch the border
+            border_components = set(labeled[border_mask & (labeled > 0)])
+            
+            # Create mask for border-connected white areas
+            outer_mask = np.zeros_like(bg_mask, dtype=bool)
+            for component in border_components:
+                outer_mask = outer_mask | (labeled == component)
+            
+            # Set alpha to 0 for border-connected white areas
+            data[outer_mask, 3] = 0
+            
+            # Convert back to PIL Image
+            processed_img = Image.fromarray(data)
+            
+            # Convert to base64 string
+            buffered = io.BytesIO()
+            processed_img.save(buffered, format="PNG")
+            img_str = base64.b64encode(buffered.getvalue()).decode()
+            
+            return f"data:image/png;base64,{img_str}"
+            
+        except Exception as e:
+            logger.error(f"Error processing image {image_path}: {str(e)}")
+            return ""
+
+    def identify_categorical_columns(self, columns: List[Dict], input_data: Dict) -> List[Dict]:
+        """
+        识别所有categorical列，为每个列确定是否应使用图标以及图标类别。
+        
+        Args:
+            columns: 列信息列表
+            input_data: 输入数据
+            
+        Returns:
+            列表，包含所有应该使用图标的categorical列信息
+        """
+        categorical_columns = []
+        
+        for column in columns:
+            if column["data_type"] == "categorical":
+                # 检查是否应该使用图标
+                result = self._should_use_icons(column["name"], input_data)
+                if result:
+                    categorical_columns.append(result)
+        
+        return categorical_columns
+    
+    def _should_use_icons(self, column_name: str, input_data: Dict) -> Optional[Dict]:
+        """
+        Determine whether icons should be used for a specific column
+        
+        Args:
+            column_name: Column name
+            input_data: Input data
+            
+        Returns:
+            Dictionary containing column name and category, or None if icons should not be used
+        """
+        if not self.base_url or not self.api_key:
+            return None
+            
+        # Get unique values for the column
+        data_rows = input_data.get("data", {}).get("data", [])
+        unique_values = list(set([row[column_name] for row in data_rows]))
+        titles = input_data.get("metadata", {}).get("titles", {})
+
+        prompt = f"""Please analyze whether icons should be used to distinguish between different groups in this chart.
+
+Chart Context:
+Title: {titles.get('main_title', '')}
+Subtitle: {titles.get('sub_title', '')}
+Column Name: {column_name}
+Unique Values: {', '.join(map(str, unique_values[:10]))}{"..." if len(unique_values) > 10 else ""}
+
+Categorize these values into one of the following types:
+1. country (use country flags, including historical countries and regions)
+2. industry (e.g., tech, finance, healthcare, manufacturing)
+3. weather (e.g., sunny, rainy, cloudy, stormy)
+4. emotion (e.g., happy, sad, neutral, excited, like, dislike, support, oppose, agree, disagree)
+5. transport (e.g., car, plane, train, ship)
+6. nature (e.g., animals, plants, landscapes)
+7. sports (e.g., football, basketball, tennis)
+8. politics (e.g., president, prime minister, king, queen, party, government)
+9. other (e.g., food, drink, art, music, science, technology, history, geography, etc.) suitable for icons
+10. abstract (not suitable for icons)
+
+Consider whether icons would enhance understanding of the data. If the values don't clearly fit into any specific category above, or if using icons would not add meaningful information, select 'abstract' and set 'use_icons' to false.
+
+Please respond in JSON format:
+{{
+    "use_icons": true/false,
+    "category": "country/industry/weather/emotion/transport/nature/sports/abstract"
+}}"""
+
+        try:
+            response = requests.post(
+                self.request_url,
+                headers={
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json"
+                },
+                json={
+                    "model": "gpt-4o-mini",
+                    "messages": [{
+                        "role": "user",
+                        "content": prompt
+                    }]
+                }
+            )
+            if response.status_code == 200:
+                content = response.json()['choices'][0]['message']['content']
+                content = content.replace('```json\n', '').replace('\n```', '').strip()
+                try:
+                    decision = json.loads(content)
+                    logger.info(f"Column {column_name} icon usage decision: {decision}")
+                    
+                    if not decision['use_icons']:
+                        return None
+                        
+                    return {
+                        "column": column_name,
+                        "category": decision['category']
+                    }
+                    
+                except json.JSONDecodeError:
+                    logger.error("Failed to parse LLM response as JSON")
+                    return None
+            else:
+                logger.error(f"LLM API request failed with status code: {response.status_code}")
+                return None
+            
+        except Exception as e:
+            logger.error(f"Error calling LLM API: {str(e)}")
+            return None
+
+    def recommend_images(self, input_data: Dict) -> Dict:
+        """
+        Recommend images based on the input data.
+        """
+        if self.index_builder is None:
+            raise ValueError("Index builder not initialized")
+        
+        # Extract necessary information
+        data_dict = input_data.get("data", {})
+        columns = data_dict.get("columns", [])
+        
+        # Step 1: 识别所有应该使用图标的categorical列
+        categorical_columns = self.identify_categorical_columns(columns, input_data)
+        # print(categorical_columns)
+        
+        # Step 2: Get topic-level clipart recommendations
+        query_text = self.create_query_text(input_data)
+        topic_clipart = self.index_builder.search(query_text, top_k=5, image_type='clipart')
+        
+        # Prepare the result
+        result = {
+            "topic_clipart": [
+                {
+                    "image_path": img["image_path"],
+                    "image_data": img["image_data"],
+                    "similarity_score": 1.0 / (1.0 + img["distance"])
+                }
+                for img in topic_clipart
+            ],
+            "group_icons": {}
+        }
+
+        # Step 3: 为每个categorical列处理图标
+        for column_info in categorical_columns:
+            group_col = column_info["column"]
+            category = column_info["category"]
+            data_rows = input_data.get("data", {}).get("data", [])
+            unique_values = list(set([row[group_col] for row in data_rows]))
+            
+            # 创建一个包含列名的键，用于存储该列的图标
+            column_key = f"{group_col}"
+            result["group_icons"][column_key] = {}
+            
+            # Try special categories first
+            if category in self.special_categories:
+                icons = self.process_special_icons(category, unique_values)
+                if icons:
+                    result["group_icons"][column_key] = icons
+            else:
+                # print(f"Processing normal icons for column {group_col}")
+                # 为该列创建专门的输入数据
+                column_input = {
+                    "data": data_dict,
+                    "titles": input_data.get("metadata", {}).get("titles", {})
+                }
+                icons = self.process_normal_icons(column_input, unique_values, group_col)
+                result["group_icons"][column_key] = icons
+        
+        return self.process_recommendation_result(result)
+
+    def process_recommendation_result(self, result: Dict) -> Dict:
+        """
+        Process the recommendation result:
+        1. Remove unnecessary fields
+        2. Convert paths to absolute paths
+        3. Post-process images
+        4. Random select from lists
+        
+        Args:
+            result: Original recommendation result
+            
+        Returns:
+            Processed recommendation result
+        """
+        import random
+        
+        processed = {"other": {}, "field": {}}
+        
+        # Process topic clipart
+        if result["topic_clipart"]:
+            # Randomly select one clipart
+            clipart = random.choice(result["topic_clipart"])
+            processed["other"]["primary"] = self.post_process_image(clipart["image_path"], max_size=768)
+        
+        # Process group icons - 新结构支持多列的图标
+        for column_name, column_icons in result["group_icons"].items():
+            key_icon_pairs = {}
+            for group, icons in column_icons.items():
+                field_key = group
+                if isinstance(icons, list):
+                    # For normal icons (list of candidates)
+                    icon = random.choice(icons)
+                    processed["field"][field_key] = self.post_process_image(icon["image_path"], max_size=256)
+                    key_icon_pairs[group] = icon
+                else:
+                    # For special icons (single icon)
+                    processed["field"][field_key] = self.post_process_image(icons["image_path"], max_size=256)
+                    key_icon_pairs[group] = icons
+            # processed["field"] = self.process_stretch_icons(self.input_data, processed["field"], column_name, key_icon_pairs)
+        
+        return processed
+
+    def process_stretch_icons(self, input_data: Dict, icons: Dict, column_name: str, key_icon_pair: Dict) -> Dict:
+        """
+        Process stretch icons based on the input data.
+        """
+        data_rows = input_data.get("data", {}).get("data", [])
+        numerical_column = {}
+        for column in input_data.get("data", {}).get("columns", []):
+            if column["role"] == "y":
+                numerical_column = column
+                break
+        # print(numerical_column)
+        if not numerical_column:
+            return icons
+        data_values = {}
+        for key, icon in key_icon_pair.items():
+            for row in data_rows:
+                if row[column_name] == key:
+                    data_values[key] = {
+                        "value": row[numerical_column["name"]],
+                        "icon": icon
+                    }
+        # print(data_values)
+        # 检查data_values是否是数值
+        if not all(isinstance(value["value"], (int, float)) for value in data_values.values()):
+            return icons
+        
+        width = 30
+        min_height = 5
+        max_height = 300
+        min_value = 1e5
+        max_value = -1e5
+        for key, value in data_values.items():
+            if value["value"] < min_value:
+                min_value = value["value"]
+            if value["value"] > max_value:
+                max_value = value["value"]
+        height_range = max_value - min_value
+        for key, value in data_values.items():
+            value["height"] = min(max(min_height, width * value["value"] / height_range), max_height)
+            from PIL import Image
+            import numpy as np
+            import base64
+            import io
+            # Get absolute path
+            abs_path = os.path.join(self.normal_icons, value["icon"]["image_path"])
+            if not os.path.exists(abs_path):
+                logger.error(f"Image not found: {abs_path}")
+                return ""
+            
+            # Convert white background to transparent
+            img = Image.open(abs_path)
+            if img.mode != 'RGBA':
+                img = img.convert('RGBA')
+            
+            data = np.array(img)
+            # Convert white-ish pixels to transparent
+            # RGB all > 240 is considered white-ish
+            white_mask = (data[..., :3] > 240).all(axis=2)
+            data[white_mask, 3] = 0
+            
+            processed_img = Image.fromarray(data)
+            
+            # Convert to base64 string
+            buffered = io.BytesIO()
+            processed_img.save(buffered, format="PNG")
+            img_str = base64.b64encode(buffered.getvalue()).decode()
+            request_data = {
+                "image": f"data:image/png;base64,{img_str}",
+                "scale": value["height"] / 30,
+                'hrz': 'false'
+            }
+            response = requests.post(
+                "http://166.111.86.168:5000/scale",
+                json=request_data
+            )
+            response.raise_for_status()
+            result = response.content
+            # 解码base64数据并保存图片
+            img_data = base64.b64decode(result.split(b"data:image/png;base64,")[1])
+            with open(f"./tmp/stretch_icon_{key}.png", "wb") as f:
+                f.write(img_data)
+            # value["icon"]["image_path"] = result["image_path"]
+        return data_values
+
+def scale_icon(icon: Dict, height: float) -> Dict:
+    """
+    Scale the icon to the given height.
+    """
+    from PIL import Image
+    import numpy as np
+    import base64
+    import io
+    
+    # Get absolute path
+    abs_path = os.path.join(self.normal_icons, icon["image_path"])
+    if not os.path.exists(abs_path):
+        logger.error(f"Image not found: {abs_path}")
+        return ""
+
+def process(input: str, output: str, embed_model_path: str = None, resource_path: str = None, data_path: str = None, index_path: str = None, base_url: str = None, api_key: str = None) -> bool:
+    """
+    Pipeline入口函数，处理单个文件的图像推荐
+    
+    Args:
+        input_path: 输入JSON文件路径
+        output_path: 输出JSON文件路径
+        embed_model_path: 嵌入模型路径
+        data_path: 图像数据路径
+        index_path: 索引文件路径
+    """
+    # print("process")
+    try:
+        # 读取输入文件
+        with open(input, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        # print("input")
+        # 预处理数据
+        processed_data = preprocess_data(data)
+        # print("processed_data")
+        # 生成图像推荐
+        recommender = ImageRecommender(
+            embed_model_path=embed_model_path,
+            data_path=data_path,
+            index_path=index_path,
+            resource_path=resource_path,
+            base_url=base_url,
+            api_key=api_key
+        )
+        recommender.input_data = processed_data
+        image_result = recommender.recommend_images(processed_data)
+        # print("image_result")
+        # 添加图像推荐到数据中
+        processed_data["images"] = image_result
+        
+        # 保存结果
+        with open(output, "w", encoding="utf-8") as f:
+            json.dump(processed_data, f, indent=2, ensure_ascii=False)
+            
+        return True
+            
+    except Exception as e:
+        logger.error(f"图像推荐失败: {str(e)}")
+        return False
+
+def preprocess_data(data: Dict) -> Dict:
+    """
+    预处理数据，确保格式正确
+    """
+    try:
+        # 深拷贝避免修改原始数据
+        processed = data.copy()
+        
+        # 确保metadata字段存在
+        if "metadata" not in processed:
+            processed["metadata"] = {}
+            
+        # 确保titles字段存在
+        if "titles" not in processed["metadata"]:
+            processed["metadata"]["titles"] = {}
+            
+        # 确保data_facts字段存在
+        if "data_facts" not in processed["metadata"]:
+            processed["metadata"]["data_facts"] = []
+        return processed
+        
+    except Exception as e:
+        logger.error(f"数据预处理失败: {str(e)}")
+        raise
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser(description="Image Recommender")
+    parser.add_argument("--input", type=str, required=True, help="Input JSON file path")
+    parser.add_argument("--output", type=str, required=True, help="Output JSON file path")
+    parser.add_argument("--embed_model_path", type=str, help="Path to the embedding model")
+    parser.add_argument("--data_path", type=str, help="Path to the image data file")
+    parser.add_argument("--index_path", type=str, help="Path to the index file")
+    parser.add_argument("--base_url", type=str, help="Base URL for LLM API")
+    parser.add_argument("--api_key", type=str, help="API key for LLM")
+    args = parser.parse_args()
+
+    success = process(
+        input=args.input,
+        output=args.output,
+        embed_model_path=args.embed_model_path,
+        data_path=args.data_path,
+        index_path=args.index_path,
+        base_url=args.base_url,
+        api_key=args.api_key
+    )
+
+    if success:
+        print("Processing json successed.")
+    else:
+        print("Processing json failed.")
+
+if __name__ == "__main__":
+    main() 

modules/image_recommender/modify.py

@@ -0,0 +1,70 @@
+# check all json file in results_path
+import json, os
+results_path = './results/'
+target_path = './results_m/'
+# remove target_path
+if os.path.exists(target_path):
+    os.system(f'rm -r {target_path}')
+if not os.path.exists(target_path):
+    os.makedirs(target_path)
+
+available_list = ['image_content', 'topic', 'data_facts', 'color_style', 'icon_or_clipart']
+data_fact_list = ['increasing', 'decreasing', 'highlight', 'deny', 'maximum', 'minimum', 'comparison', 'none']
+color_style_list = ['monochrome', 'colorful', 'grayscale', 'dual-color']
+scale_and_complexity_level_list = ['icon', 'clipart', 'background']
+for i, file in enumerate(os.listdir(results_path)):
+    if not file.endswith('.json'):
+        continue
+    
+    # load json
+    with open(f'{results_path}{file}', 'r') as f:
+        data = json.load(f)
+
+    if 'output_case' in data:
+        data = data['output_case']
+
+    if 'Topic' in data:
+        data['topic'] = data['Topic']
+        del data['Topic']
+    
+    if 'key_words' in data:
+        data['topic'] = data['key_words']
+        del data['key_words']
+    
+    # key error
+    error = False
+    for key in available_list:
+        if key not in data:
+            error = True
+            break
+    if error:
+        print(f'{file} is not complete')
+        continue
+
+    # data_facts error
+    if data['data_facts'] == 'increase':
+        data['data_facts'] = 'increasing'
+    if data['data_facts'] not in data_fact_list:
+        print(f'{file} data_facts is error')
+        continue
+
+    # color_style error
+    if data['color_style'] == 'gray-scale' or data['color_style'] == 'black and white':
+        data['color_style'] = 'grayscale'
+    if data['color_style'] == 'duo-color' or data['color_style'] == 'duel-color':
+        data['color_style'] = 'dual-color'
+    if data['color_style'] not in color_style_list:
+        print(f'{file} color_style is error')
+        continue
+
+    # scale_and_complexity_level error
+    if data['icon_or_clipart'] not in scale_and_complexity_level_list:
+        print(f'{file} icon_or_clipart is error')
+        continue
+    data['size'] = data['icon_or_clipart']
+    del data['icon_or_clipart']
+    
+    # save json
+    target_file = f'{target_path}{file}'
+    with open(target_file, 'w') as f:
+        json.dump(data, f)

modules/image_recommender/prompt.json

@@ -0,0 +1,135 @@
+{
+    "task": "Analyze the semantic and style modes in the given image and provide explanations.",
+    "output_format": {
+      "image_content": {
+        "type": "string",
+        "description": "Primary content in image (less than 10 words)",
+        "examples": [
+          {
+            "example_value": "A single red apple",
+            "explanation": "A concise description with 5 words focusing on 'what' the image is."
+          },
+          {
+            "example_value": "Two people shaking hands",
+            "explanation": "An example describing the main objects and action within 10 words."
+          }
+        ]
+      },
+      "topic": {
+        "type": "string",
+        "description": "Relevant topic (3-5 words)",
+        "examples": [
+          {
+            "example_value": "fruit, environment, healthy eating",
+            "explanation": "Three to five words separated by commas, capturing the core theme."
+          },
+          {
+            "example_value": "teamwork, collaboration, business partnership",
+            "explanation": "Three to five key words describing the main idea of the image."
+          }
+        ]
+      },
+      "data_facts": {
+        "type": "enum",
+        "description": "Possible usage to reflect any data fact (choose one from the list or 'none')",
+        "choices": [
+          "increasing",
+          "decreasing",
+          "highlight",
+          "deny",
+          "maximum",
+          "minimum",
+          "none"
+        ],
+        "examples": [
+          {
+            "example_value": "increasing",
+            "explanation": "Use if the image shows an upward arrow or indicates growth."
+          },
+          {
+            "example_value": "decreasing",
+            "explanation": "Use if the image shows a downward arrow or indicates decline."
+          },
+          {
+            "example_value": "highlight",
+            "explanation": "Use if the image emphasizes something align with it (e.g., a light bulb)."
+          },
+          {
+            "example_value": "deny",
+            "explanation": "Use if the image represents a 'no' or 'wrong' symbol."
+          },
+          {
+            "example_value": "maximum",
+            "explanation": "Use if the image indicates a top/peak point."
+          },
+          {
+            "example_value": "minimum",
+            "explanation": "Use if the image indicates a bottom/lowest point."
+          },
+          {
+            "example_value": "none",
+            "explanation": "Use if none of the above data facts apply."
+          }
+        ]
+      },
+      "color_style": {
+        "type": "enum",
+        "description": "Dominant color style, if image is black and white, return grayscale instead of monochrome",
+        "choices": [
+          "grayscale",
+          "monochrome",
+          "dual-color",
+          "colorful"
+        ],
+        "examples": [
+          {
+            "example_value": "grayscale",
+            "explanation": "Use if the image is purely in black, white, and shades of gray."
+          },
+          {
+            "example_value": "monochrome",
+            "explanation": "Use if the image uses colors with same hue but not grayscale."
+          },
+          {
+            "example_value": "dual-color",
+            "explanation": "Use if the image uses two sets of colors on the color wheel."
+          },
+          {
+            "example_value": "colorful",
+            "explanation": "Use if the image has multiple bright or varied colors."
+          }
+        ]
+      },
+      "icon_or_clipart": {
+        "type": "enum",
+        "description": "judge the image whether it is an small icon, medium clipart or large background",
+        "choices": [
+          "icon",
+          "clipart",
+          "background"
+        ],
+        "examples": [
+          {
+            "example_value": "icon",
+            "explanation": "Icon showing basic shapes and always small."
+          },
+          {
+            "example_value": "clipart",
+            "explanation": "Clipart showing detailed objects."
+          },
+          {
+            "example_value": "background",
+            "explanation": "Large scale image or clipart that suitable to be used as a background."
+          }
+        ]
+      }
+    },
+    "output_example": {
+      "image_content": "A single red apple",
+      "topic": "fruit, environment",
+      "data_facts": "none",
+      "color_style": "colorful",
+      "icon_or_clipart": "icon",
+      "explanation": "The image shows a red apple, focusing on a fruit in nature. It is colorful and icon-sized. No specific data fact is represented."
+    }
+  }

modules/image_recommender/test.py

@@ -0,0 +1,112 @@
+# 1. 扫描当前目录下所有image文件 包含png jpg jpeg webp，进行编号 存储下list
+import os, json
+image_pathes = []
+root = './images/'
+image_path_file = 'image_pathes.json'
+
+if os.path.exists(image_path_file):
+    with open(image_path_file, 'r') as f:
+        image_pathes = json.load(f)
+        print(len(image_pathes))
+else:
+    print('error')
+
+from openai import OpenAI
+from PIL import Image
+import base64
+from io import BytesIO
+
+client = OpenAI(
+    api_key=os.getenv("OPENAI_API_KEY") or os.getenv("AIHUBMIX_API_KEY", ""),
+    base_url=os.getenv("OPENAI_BASE_URL", "https://aihubmix.com/v1")
+)
+
+def resize_image(img, max_size=512):
+    width, height = img.size
+    ratio = min(max_size / width, max_size / height)
+    if ratio >= 1:
+        return img
+    new_width = int(width * ratio)
+    new_height = int(height * ratio)
+    resized_img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+    return resized_img
+
+def repaint_image(img):
+    # rapaint transparent area with white color
+    img = img.convert('RGBA')
+    data = img.getdata()
+    new_data = []
+    for item in data:
+        if item[3] == 0:
+            new_data.append((255, 255, 255, 255))
+        else:
+            new_data.append(item)
+    img.putdata(new_data)
+    img = img.convert('RGB')
+    # img.save('temp.png')
+    return img
+    
+def image_to_base64(image_path, show=False, target_size=512):
+    with Image.open(image_path) as img:
+        # img = img.resize(size)
+        img = resize_image(img, 512)
+        img = repaint_image(img)
+        buffered = BytesIO()
+        img.save(buffered, format="PNG")
+        img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
+        return img_base64
+
+wwxxhh = 0
+def ask_image(prompt, image_data):
+    global wwxxhh
+    number_of_trials = 0
+    while number_of_trials < 5:
+        try:
+            response = client.chat.completions.create(
+            #   model="gpt-4o-mini",
+            model="gemini-2.0-flash",
+            messages=[
+                {
+                  "role": "user",
+                  "content": [
+                    {   
+                        "type": "text", 
+                        "text": prompt},
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:image/jpeg;base64,{image_data}"
+                        },
+                    },
+                  ],
+                }
+              ]
+            )
+            wwxxhh += response.usage.total_tokens
+            return response.choices[0].message.content
+
+        except Exception as e:
+            number_of_trials += 1
+            print(e)
+
+    return 'Error!'
+
+# 2. 读取prompt.json文件，读取整个作为字符串，逐个读取image文件，调用ask_image函数，将返回的结果存储下来
+import json
+with open('prompt.json', 'r') as f:
+    prompt = f.read()
+# print(prompt)
+
+
+
+ques_id = 15394
+results_path = './results/'
+image_path = image_pathes[ques_id]
+print(image_path)
+
+image_data = image_to_base64(image_path)
+result = ask_image(prompt, image_data)
+print(result)
+
+# from IPython import embed
+# embed()