demo/app.py

import argparse
import os
from typing import Any, Dict, List, Optional

import gradio as gr
import numpy as np
import torch
from gradio_image_annotation import image_annotator
from PIL import Image

import faiss
from models.configs import get_model_config
from models.llava import init_llava
from models.relevance_feedback import CaptionVLMRelevanceFeedback, RocchioUpdate
from utils.image_utils import resize_images
from utils.utils import get_timestamp, load_yaml, save_json

import os

print("CWD:", os.getcwd())
print("ROOT /app:", os.listdir("/app"))

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--config_path", 
        type=str, 
        required=True,
        help="Path to the main configuration file"
    )
    parser.add_argument(
        "--captioning_model_config_path", 
        type=str, 
        required=True,
        help="Path to captioning model config file"
    )
    return parser.parse_args()

args = parse_args()

CONFIG_PATH = args.config_path
CAPTIONING_MODEL_CONFIG_PATH = args.captioning_model_config_path

logs = {
    "start_timestamp": get_timestamp(),
    "config_path": CONFIG_PATH,
    "captioning_model_config_path": CAPTIONING_MODEL_CONFIG_PATH,
    "experiments": {},
}

retrieval_round = 1
experiment_id = 0

accumulated_query_embeddings = {"query_embedding": None}

config = load_yaml(CONFIG_PATH)
default_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# VLM: Retrieval Backbone
model_config = get_model_config(config["VLM_MODEL_FAMILY"], config["VLM_MODEL_NAME"])
processor = model_config["processor_class"].from_pretrained(model_config["model_id"])
model = model_config["model_class"].from_pretrained(model_config["model_id"])
model.eval()
wrapper = model_config["wrapper_class"](model=model, processor=processor)

# Read image index: candidate images
index = faiss.read_index(config["INDEX_PATH"])
with open(os.path.join(os.path.dirname(config["INDEX_PATH"]), "image_paths.txt"), "r") as f:
    candidate_image_paths = [line.strip().replace('\\', '/') for line in f.readlines()]

# Initialize relevance feedback model
captioning_model_config = load_yaml(CAPTIONING_MODEL_CONFIG_PATH)
captioning_device = torch.device(captioning_model_config.get("DEVICE", default_device.type))
use_8bit = bool(captioning_model_config.get("USE_8BIT", False))
if captioning_device.type == "cpu":
    use_8bit = False
model_config = get_model_config(
    captioning_model_config["MODEL_FAMILY"], 
    captioning_model_config["MODEL_ID"]
)
# --- FORCE OVERRIDE ---
model_config["model_id"] = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf"
# ----------------------
if captioning_model_config["MODEL_FAMILY"] == "llava":
    captioning_model = init_llava(
        model_config=model_config,
        device=captioning_device,
        use_8bit=use_8bit
    )
else:
    raise ValueError(f"Captioning model family {captioning_model_config['model_family']} not supported")
captioning_relevance_feedback = CaptionVLMRelevanceFeedback(
    vlm_wrapper_retrieval=wrapper,
    vlm_wrapper_captioning=captioning_model,
)
rocchio_update = RocchioUpdate(alpha=0.6, beta=0.2, gamma=0.2)


def update_logs_retrieval(experiment_id, retrieval_round, user_query, top_k, retrieved_image_paths, scores):
    logs["experiments"][experiment_id].append(
        {
            "timestamp": get_timestamp(),
            "type": "retrieval",
            "round": retrieval_round,
            "user_query": user_query,
            "top_k": top_k,
            "retrieved_image_paths": retrieved_image_paths,
            "scores": scores,
        }
    )
    save_json(logs, config["RETRIEVAL_LOGS_PATH"])


def update_logs_feedback(
        experiment_id: str,
        retrieval_round: int,
        user_query: str,
        annotations: List[Dict[str, Any]],
        relevant_textual_features: Optional[str] = None,
        irrelevant_textual_features: Optional[str] = None
    ):
    if relevant_textual_features is None:
        relevant_textual_features = ""
    if irrelevant_textual_features is None:
        irrelevant_textual_features = ""
    logs["experiments"][experiment_id].append(
        {
            "timestamp": get_timestamp(),
            "type": "feedback",
            "round": retrieval_round,
            "user_query": user_query,
            "annotations": annotations,
            "relevant_textual_features": relevant_textual_features.split(", "),
            "irrelevant_textual_features": irrelevant_textual_features.split(", "),
        }
    )
    save_json(logs, config["RETRIEVAL_LOGS_PATH"])


def image_search(query, top_k=5):
    """Retrieve images based on text query"""
    global retrieval_round
    global experiment_id
    experiment_id += 1
    logs["experiments"][experiment_id] = []

    processed_query = wrapper.process_inputs(text=query)
    with torch.no_grad():
        query_embedding = wrapper.get_text_embeddings(processed_query)

    accumulated_query_embeddings["query_embedding"] = query_embedding

    scores, img_ids = index.search(query_embedding, top_k)
    scores = scores.squeeze().tolist()
    img_ids = img_ids.squeeze().tolist()
    retrieved_image_paths = [candidate_image_paths[i] for i in img_ids]
    retrieved_images = [Image.open(path) for path in retrieved_image_paths]
    retrieved_images = resize_images(retrieved_images, config)

    update_logs_retrieval(experiment_id, retrieval_round, query, top_k, retrieved_image_paths, scores)

    return retrieved_images, scores, retrieved_image_paths


def process_feedback(query, top_k, image_paths, annotator_json_boxes_list):
    """Process feedback from the annotator"""

    relevance_feedback_results = captioning_relevance_feedback(
        query=query,
        relevant_image_paths=image_paths,
        visualization=True,
        top_k_feedback=top_k,
        annotator_json_boxes_list=annotator_json_boxes_list,
        prompt_based_on_query=False,
        prompt=captioning_model_config.get("PROMPT", None)
    )
    
    return (
        relevance_feedback_results["positive"],
        relevance_feedback_results["negative"],
        relevance_feedback_results.get("relevant_captions", ""),
        relevance_feedback_results.get("irrelevant_captions", ""),
        relevance_feedback_results["explanation"],
    )


def feedback_loop(
    query, 
    top_k, 
    image_paths, 
    annotator_json_boxes_list,
    relevant_textual_features: Optional[str] = None,
    irrelevant_textual_features: Optional[str] = None,
    fuse_initial_query: bool = False
):
    """Apply feedback to the image search"""
    print(annotator_json_boxes_list)
    global retrieval_round
    
    processed_query = wrapper.process_inputs(text=query)
    with torch.no_grad():
        query_embedding = wrapper.get_text_embeddings(processed_query)

    relevance_feedback_results = captioning_relevance_feedback(
        query=query,
        relevant_image_paths=image_paths,
        visualization=False,
        top_k_feedback=top_k,
        annotator_json_boxes_list=annotator_json_boxes_list,
        prompt_based_on_query=False,
        relevant_captions=relevant_textual_features,
        irrelevant_captions=irrelevant_textual_features
    )

    rocchio_query_embedding = (accumulated_query_embeddings["query_embedding"] + query_embedding) / 2 if (
        fuse_initial_query
    ) else accumulated_query_embeddings["query_embedding"]
    accumulated_query_embeddings["query_embedding"] = rocchio_update(
        query_embeddings=rocchio_query_embedding,
        positive_embeddings=relevance_feedback_results["positive"],
        negative_embeddings=relevance_feedback_results["negative"]
    )

    scores, img_ids = index.search(accumulated_query_embeddings["query_embedding"], top_k)
    scores = scores.squeeze().tolist()
    img_ids = img_ids.squeeze().tolist()
    retrieved_image_paths = [candidate_image_paths[i] for i in img_ids]
    retrieved_images = [Image.open(path) for path in retrieved_image_paths]
    retrieved_images = resize_images(retrieved_images, config)

    update_logs_feedback(
        experiment_id,
        retrieval_round,
        query,
        annotator_json_boxes_list,
        relevant_textual_features,
        irrelevant_textual_features
    )

    retrieval_round += 1

    update_logs_retrieval(experiment_id, retrieval_round, query, top_k, retrieved_image_paths, scores)

    return (
        retrieved_images, 
        scores, 
        retrieved_image_paths, 
        relevance_feedback_results["explanation"]
    )


def get_boxes_json(annotations):
    """Get bounding boxes from annotator"""
    return annotations["boxes"] if annotations["boxes"] else None


css = """
#warning {background-color: #FFCCCB} 
.feedback {font-size: 20px !important;}
.feedback textarea {font-size: 20px !important;}
"""

with gr.Blocks(title="Multimodal Retrieval Demo", css=css) as demo:
    gr.Markdown("# Text-to-Image Search")

    image_top_k = gr.State(value=5)
    fuse_initial_query = gr.State(value=config.get("FUSE_INITIAL_QUERY", True))

    with gr.Tab("Image Search"):
        with gr.Row():
            with gr.Column():
                query = gr.Textbox(label="Describe the image you would like to find:")
                image_search_btn = gr.Button("Search Images")

        annotators = []
        annotator_json_boxes_list = []

        with gr.Row():
            image_gallery = gr.Gallery(
                label="Retrieved Images", columns=5, rows=1, visible=config["SHOW_IMAGE_GALLERY"]
            )

        with gr.Row():
            for i in range(image_top_k.value):
                with gr.Column():
                    annotator = image_annotator(
                        value=None,
                        label_list=["Relevant", "Irrelevant"],
                        label_colors=[(0, 255, 0), (255, 0, 0)],
                        label=f"Result {i + 1}",
                        visible=config["SHOW_ANNOTATORS"],
                        sources=[],
                    )
                    annotators.append(annotator)
                    button_get = gr.Button(f"Get bounding boxes for Result {i + 1}")
                    annotator_json_boxes = gr.JSON(visible=True)
                    annotator_json_boxes_list.append(annotator_json_boxes)
                    button_get.click(get_boxes_json, inputs=annotator, outputs=annotator_json_boxes)

        def format_outputs_image_search(images, scores, retrieved_image_paths):
            outputs_annotators = []
            outputs_gallery = []
            outputs_retrieved_image_paths = []
            outputs_images_with_saliency = None
            for i in range(len(images)):
                outputs_annotators.append({"image": images[i]})
                outputs_gallery.append((images[i], f"Relevance score: {scores[i]}"))
                outputs_retrieved_image_paths.append(retrieved_image_paths[i])
            final_outputs = [outputs_gallery] \
                + [outputs_retrieved_image_paths] \
                + [outputs_images_with_saliency] \
                + outputs_annotators
            return final_outputs
        
        relevant_image_paths = gr.State(value=None)

        with gr.Row():
            process_feedback_btn = gr.Button("Process feedback")
            
        with gr.Row():
            feedback_explanation_gallery = gr.Gallery(
                label="Feedback Explanations (Previous Round)", columns=5, rows=1, visible=config["SHOW_IMAGE_GALLERY"]
            )

        image_search_btn.click(
            fn=lambda query, top_k: format_outputs_image_search(*image_search(query, top_k)),
            inputs=[query, image_top_k],
            outputs=[image_gallery, relevant_image_paths, feedback_explanation_gallery, *annotators],
        )

        with gr.Row():
            with gr.Column():
                    relevant_features = gr.Textbox(
                        label="Relevant features",
                        visible=True if CAPTIONING_MODEL_CONFIG_PATH is not None else False,
                        interactive=True,
                        elem_classes=["feedback"]
                    )
            with gr.Column():
                    irrelevant_features = gr.Textbox(
                        label="Irrelevant features",
                        visible=True if CAPTIONING_MODEL_CONFIG_PATH is not None else False,
                        interactive=True,
                        elem_classes=["feedback"]
                    )


        def format_outputs_process_feedback(positive, negative, relevant_captions, irrelevant_captions, explanation):
            outputs_explanation = []
            for i in range(len(explanation)):
                outputs_explanation.append(explanation[i])
            for i, caption in enumerate(relevant_captions):
                if caption.endswith("."):
                    relevant_captions[i] = caption[:-1]
            outputs_relevant_captions = ", ".join(relevant_captions)
            for i, caption in enumerate(irrelevant_captions):
                if caption.endswith("."):
                    irrelevant_captions[i] = caption[:-1]
            outputs_irrelevant_captions = ", ".join(irrelevant_captions)
            final_outputs =  [outputs_relevant_captions] \
                + [outputs_irrelevant_captions] \
                + [outputs_explanation]
            return final_outputs
            

        process_feedback_btn.click(
            fn=lambda query, top_k, image_paths, *annotator_json_boxes_list: format_outputs_process_feedback(
                *process_feedback(query, top_k, image_paths, annotator_json_boxes_list)
            ),
            inputs=[query, image_top_k, relevant_image_paths, *annotator_json_boxes_list],
            outputs=[relevant_features, irrelevant_features, feedback_explanation_gallery],
        )

        with gr.Row():
            apply_feedback_btn = gr.Button("Apply Feedback")

        def format_outputs_feedback(images, scores, retrieved_image_paths, images_with_saliency):
            outputs_annotators = []
            outputs_gallery = []
            outputs_retrieved_image_paths = []  
            for i in range(len(images)):
                outputs_annotators.append({"image": images[i], "boxes": []})
                outputs_gallery.append((images[i], f"Relevance score: {scores[i]}"))
                outputs_retrieved_image_paths.append(retrieved_image_paths[i])
            final_outputs = [outputs_gallery] \
                + [outputs_retrieved_image_paths] \
                + outputs_annotators
            return final_outputs


        def feedback_interface(
            query,
            top_k,
            image_paths,
            relevant_features,
            irrelevant_features,
            fuse_initial_query,
            *annotator_json_boxes_list,
        ):
            results = feedback_loop(
                query,
                top_k,
                image_paths,
                annotator_json_boxes_list,
                relevant_features,
                irrelevant_features,
                fuse_initial_query
            )
            return format_outputs_feedback(*results)
        
        apply_feedback_btn.click(
            fn=feedback_interface,
            inputs=[query, image_top_k, relevant_image_paths, relevant_features, irrelevant_features, fuse_initial_query, *annotator_json_boxes_list],
            outputs=[image_gallery, relevant_image_paths, *annotators],
        ).then(
            fn=lambda: [None for _ in annotator_json_boxes_list],
            inputs=None,
            outputs=[*annotator_json_boxes_list]
        )


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)