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)