modules/image_recommender/create_index.py

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
    )