app.py

import gradio as gr
import torch
import numpy as np
from PIL import Image
import pandas as pd
from pathlib import Path
from sklearn.metrics.pairwise import cosine_similarity

# Import your model and necessary functions
from src.config import ConfigManager
from src.token_classifier import load_token_classifier, predict
from your_model_file import YourModel  # Replace with your actual model import

# Load model and configurations
def load_model():
    model = YourModel()  # Initialize your model
    model.eval()
    return model

def load_dataset():
    # Load your default dataset
    database_df = pd.read_csv('database.csv')  # Adjust path as needed
    return database_df

def process_single_query(model, query_image_path, query_text, database_embeddings, database_df):
    device = ConfigManager().get("training")["device"]
    
    # Process query image
    query_img = model.processor(Image.open(query_image_path)).unsqueeze(0).to(device)
    
    # Get token classifier
    token_classifier, token_classifier_tokenizer = load_token_classifier(
        ConfigManager().get("paths")["pretrained_token_classifier_path"],
        device
    )
    
    with torch.no_grad():
        query_img_embd = model.feature_extractor.encode_image(query_img)
        
        # Process text query
        predictions = predict(
            tokens=query_text,
            model=token_classifier,
            tokenizer=token_classifier_tokenizer,
            device=device,
            max_length=128
        )
        
        # Process positive and negative objects
        pos = []
        neg = []
        last_tag = ''
        for token, label in predictions:
            if label == '<positive_object>':
                if last_tag != '<positive_object>':
                    pos.append(f"a photo of a {token}.")
                else:
                    pos[-1] = pos[-1][:-1] + f" {token}."
            elif label == '<negative_object>':
                if last_tag != '<negative_object>':
                    neg.append(f"a photo of a {token}.")
                else:
                    neg[-1] = neg[-1][:-1] + f" {token}."
            last_tag = label
            
        # Combine embeddings
        for obj in pos:
            query_img_embd += model.feature_extractor.encode_text(
                model.tokenizer(obj).to(device)
            )[0]
        for obj in neg:
            query_img_embd -= model.feature_extractor.encode_text(
                model.tokenizer(obj).to(device)
            )[0]
            
        query_img_embd = torch.nn.functional.normalize(query_img_embd, dim=1, p=2)
        
    # Calculate similarities
    query_embedding = query_img_embd.cpu().numpy()
    similarities = cosine_similarity(query_embedding, database_embeddings)[0]
    
    # Get most similar image
    most_similar_idx = np.argmax(similarities)
    most_similar_image_path = database_df.iloc[most_similar_idx]['target_image']
    
    return most_similar_image_path

# Initialize model and database
model = load_model()
database_df = load_dataset()
database_embeddings = encode_database(model, database_df)  # Using your existing function

def interface_fn(selected_image, query_text):
    result_image_path = process_single_query(
        model, 
        selected_image, 
        query_text, 
        database_embeddings, 
        database_df
    )
    return Image.open(result_image_path)

# Create Gradio interface
demo = gr.Interface(
    fn=interface_fn,
    inputs=[
        gr.Image(type="filepath", label="Select Query Image"),
        gr.Textbox(label="Enter Query Text")
    ],
    outputs=gr.Image(label="Retrieved Image"),
    title="Compositional Image Retrieval",
    description="Select an image and enter a text query to find the most similar image.",
    examples=[
        ["example_images/image1.jpg", "a red car"],
        ["example_images/image2.jpg", "a blue house"]
    ]
)

if __name__ == "__main__":
    demo.launch()