import argparse import asyncio import logging from typing import List, Optional import gradio as gr from gradio_image_annotation import image_annotator from PIL import Image from client.retrieval_client import RemoteRetrievalClient from utils.image_utils import base64_to_image, resize_images from utils.utils import get_timestamp, load_yaml, save_json # Set up logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) def parse_args(): parser = argparse.ArgumentParser(description="Retrieval Demo Client") 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" ) parser.add_argument( "--server_url", type=str, default="http://localhost:8000", help="URL of the retrieval server" ) parser.add_argument( "--timeout", type=int, default=60, help="Request timeout in seconds" ) parser.add_argument( "--port", type=int, default=7860, help="Port for the Gradio interface" ) parser.add_argument( "--share", action="store_true", help="Create a public link for the interface" ) return parser.parse_args() args = parse_args() # Load configuration config = load_yaml(args.config_path) captioning_model_config = load_yaml(args.captioning_model_config_path) logger.info(f"Loaded config from {args.config_path}") # Initialize retrieval client retrieval_client = RemoteRetrievalClient(server_url=args.server_url) logger.info(f"Initialized remote client for {args.server_url}") # Initialize logging logs = { "start_timestamp": get_timestamp(), "config_path": args.config_path, "captioning_model_config_path": args.captioning_model_config_path, "server_url": args.server_url, "experiments": {}, } retrieval_round = 1 experiment_id = 0 # Store processed feedback embeddings processed_feedback_embeddings = { "positive_embeddings": None, "negative_embeddings": None } # Functions calling the server: image search async def image_search(search_query: str, top_k: int = 5): """Retrieve images based on text query""" global retrieval_round, experiment_id experiment_id += 1 logs["experiments"][experiment_id] = [] try: logger.info(f"Searching for: {search_query}") images, scores, retrieved_image_paths = await retrieval_client.search_images(search_query, top_k) update_logs_retrieval( experiment_id, retrieval_round, search_query, top_k, retrieved_image_paths, scores ) logger.info(f"Search completed successfully, found {len(images)} images") return images, scores, retrieved_image_paths except Exception as e: logger.error(f"Search failed: {str(e)}") return [], [], [] # Functions calling the server: process feedback async def process_feedback( feedback_query: str, top_k: int, image_paths: List[str], annotator_boxes: List, user_prompt: Optional[str] = None ): """Process feedback from the annotator and store embeddings for later application""" global processed_feedback_embeddings try: logger.info(f"Processing feedback for query: {feedback_query}") relevance_feedback_results = await retrieval_client.process_feedback( query=feedback_query, relevant_image_paths=image_paths, annotator_json_boxes_list=annotator_boxes, visualization=True, top_k_feedback=top_k, user_prompt=user_prompt, prompt=captioning_model_config.get("PROMPT", None) ) processed_feedback_embeddings["positive_embeddings"] = relevance_feedback_results.get("positive") processed_feedback_embeddings["negative_embeddings"] = relevance_feedback_results.get("negative") logger.info("Feedback processed successfully and embeddings stored") return relevance_feedback_results except Exception as e: logger.error(f"Process feedback failed: {str(e)}") return [] # Functions calling the server: apply feedback async def apply_feedback( feedback_query: str, top_k: int, relevant_captions: Optional[List[str]] = None, irrelevant_captions: Optional[List[str]] = None, fuse_query: bool = False, use_stored_embeddings: bool = True ): """Apply feedback to the image search using stored processed embeddings""" global retrieval_round, processed_feedback_embeddings try: logger.info(f"Applying feedback for query: {feedback_query}") images, scores, retrieved_image_paths = await retrieval_client.apply_feedback( query=feedback_query, top_k=top_k, relevant_captions=relevant_captions, irrelevant_captions=irrelevant_captions, fuse_initial_query=fuse_query ) retrieval_round += 1 update_logs_retrieval( experiment_id, retrieval_round, feedback_query, top_k, retrieved_image_paths, scores ) logger.info(f"Feedback applied successfully, found {len(images)} images") return images, scores, retrieved_image_paths except Exception as e: logger.error(f"Apply feedback failed: {str(e)}") return [], [], [] # Update logs after retrieval def update_logs_retrieval( experiment_id: int, retrieval_round: int, user_query: str, top_k: int, retrieved_image_paths: List[str], scores: List[float], ): """Update logs with retrieval information""" 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, }) try: save_json(logs, config["RETRIEVAL_LOGS_PATH"]) except Exception as e: logger.warning(f"Failed to save logs: {str(e)}") # Update logs after feedback def update_logs_feedback( exp_id: int, round_num: int, user_query: str, annotations: List, relevant_features: Optional[str] = None, irrelevant_features: Optional[str] = None ): """Update logs with feedback information""" logs["experiments"][exp_id].append({ "timestamp": get_timestamp(), "type": "feedback", "round": round_num, "user_query": user_query, "annotations": annotations, "relevant_textual_features": relevant_features.split(", ") if relevant_features else [], "irrelevant_textual_features": irrelevant_features.split(", ") if irrelevant_features else [], }) try: save_json(logs, config["RETRIEVAL_LOGS_PATH"]) except Exception as e: logger.warning(f"Failed to save logs: {str(e)}") def get_boxes_json(annotations): """Get bounding boxes from annotator""" return annotations["boxes"] if annotations["boxes"] else None def format_outputs_image_search(images: List, scores: List[float], retrieved_image_paths: List[str]): """Format outputs for image search""" outputs_annotators = [] outputs_gallery = [] outputs_retrieved_image_paths = [] outputs_images_with_saliency = None images = resize_images(images, config) for idx in range(len(images)): outputs_annotators.append({"image": images[idx]}) outputs_gallery.append((images[idx], f"Relevance score: {scores[idx]:.4f}")) outputs_retrieved_image_paths.append(retrieved_image_paths[idx]) final_outputs = [outputs_gallery] + [outputs_retrieved_image_paths] + [outputs_images_with_saliency] + outputs_annotators return final_outputs def format_outputs_process_feedback( positive: List[float], negative: List[float], relevant_captions: str, irrelevant_captions: str, explanation: List[Image.Image] ): """Format outputs for process feedback""" outputs_explanation = [] for idx in range(len(explanation)): outputs_explanation.append(explanation[idx]) # Clean up captions if relevant_captions: if isinstance(relevant_captions, str): relevant_captions_list = relevant_captions.split(", ") else: relevant_captions_list = relevant_captions for idx, caption in enumerate(relevant_captions_list): if caption.endswith("."): relevant_captions_list[idx] = caption[:-1] outputs_relevant_captions = ", ".join(relevant_captions_list) else: outputs_relevant_captions = "" if irrelevant_captions: if isinstance(irrelevant_captions, str): irrelevant_captions_list = irrelevant_captions.split(", ") else: irrelevant_captions_list = irrelevant_captions for idx, caption in enumerate(irrelevant_captions_list): if caption.endswith("."): irrelevant_captions_list[idx] = caption[:-1] outputs_irrelevant_captions = ", ".join(irrelevant_captions_list) else: outputs_irrelevant_captions = "" final_outputs = [outputs_relevant_captions] + [outputs_irrelevant_captions] + [outputs_explanation] return final_outputs def format_outputs_feedback( images: List, scores: List[float], retrieved_image_paths: List[str], images_with_saliency: List[Image.Image], explanation: List[Image.Image] ): """Format outputs for feedback""" outputs_annotators = [] outputs_gallery = [] outputs_retrieved_image_paths = [] images = resize_images(images, config) for idx in range(len(images)): outputs_annotators.append({"image": images[idx], "boxes": []}) outputs_gallery.append((images[idx], f"Relevance score: {scores[idx]:.4f}")) outputs_retrieved_image_paths.append(retrieved_image_paths[idx]) final_outputs = [outputs_gallery] + [outputs_retrieved_image_paths] + outputs_annotators return final_outputs # Start of the Gradio interface css = """ #warning {background-color: #FFCCCB} .feedback {font-size: 20px !important;} .feedback textarea {font-size: 20px !important;} .server-status {background-color: #E8F5E8; padding: 10px; border-radius: 5px; margin: 10px 0;} .error-message {background-color: #FFE6E6; padding: 10px; border-radius: 5px; margin: 10px 0;} """ with gr.Blocks(title="VisualReF: GenAI Captioning", css=css) as demo: gr.Markdown("# Text-to-Image Search (GenAI Captioning)") image_top_k = gr.State(value=config.get("TOP_K", 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:", placeholder="Enter your search query here..." ) image_search_btn = gr.Button("Search Images", variant="primary") error_display = gr.HTML(visible=False) with gr.Row(): image_gallery = gr.Gallery( label="Retrieved Images", columns=5, rows=1, visible=config["SHOW_IMAGE_GALLERY"], show_label=True ) # Annotators for feedback annotators = [] annotator_json_boxes_list = [] 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) relevant_image_paths = gr.State(value=None) with gr.Row(): user_prompt_text = gr.Textbox( label="Instructions for the captioning model", visible=True, interactive=True, placeholder="Enter instructions for the captioning model..." ) with gr.Row(): process_feedback_btn = gr.Button("Process Feedback", variant="secondary") with gr.Row(): feedback_explanation_gallery = gr.Gallery( label="Feedback Explanations (Previous Round)", columns=5, rows=1, visible=config["SHOW_IMAGE_GALLERY"] ) async def handle_image_search(search_query, top_k): try: images, scores, retrieved_image_paths = await image_search(search_query, top_k) formatted_outputs = format_outputs_image_search(images, scores, retrieved_image_paths) return [gr.HTML(visible=False)] + formatted_outputs except Exception as e: logger.error(f"Image search error: {str(e)}") image_search_btn.click( fn=handle_image_search, inputs=[query, image_top_k], outputs=[error_display, image_gallery, relevant_image_paths, feedback_explanation_gallery, *annotators], ) # Textual features input with gr.Row(): with gr.Column(): relevant_features = gr.Textbox( label="Relevant features", visible=True, interactive=True, elem_classes=["feedback"], placeholder="Enter relevant features separated by commas..." ) with gr.Column(): irrelevant_features = gr.Textbox( label="Irrelevant features", visible=True, interactive=True, elem_classes=["feedback"], placeholder="Enter irrelevant features separated by commas..." ) # Process feedback button handler async def handle_process_feedback( feedback_query, top_k, image_paths, user_prompt, *annotator_boxes ): try: if not feedback_query.strip(): return ["", "", []] logger.info(f"{feedback_query}, {top_k}, {image_paths}, {list(annotator_boxes)}") relevance_feedback_results = await process_feedback( feedback_query, top_k, image_paths, list(annotator_boxes), user_prompt) explanation = relevance_feedback_results.get("explanation", []) if explanation is not None: explanation = [base64_to_image(img) for img in explanation] relevance_feedback_results["explanation"] = explanation logger.info(f"Relevance feedback results: {relevance_feedback_results}") return format_outputs_process_feedback( relevance_feedback_results.get("positive", []), relevance_feedback_results.get("negative", []), relevance_feedback_results.get("relevant_captions", ""), relevance_feedback_results.get("irrelevant_captions", ""), relevance_feedback_results.get("explanation", []) ) except Exception as e: logger.error(f"Process feedback error: {str(e)}") return ["", "", []] process_feedback_btn.click( fn=handle_process_feedback, inputs=[query, image_top_k, relevant_image_paths, user_prompt_text, *annotator_json_boxes_list], outputs=[relevant_features, irrelevant_features, feedback_explanation_gallery], ) # Apply feedback button with gr.Row(): apply_feedback_btn = gr.Button("Apply Feedback", variant="primary") # Apply feedback button handler async def handle_apply_feedback( feedback_query, top_k, relevant_captions, irrelevant_captions, fuse_query ): try: images, scores, retrieved_image_paths = await apply_feedback( feedback_query, top_k, relevant_captions, irrelevant_captions, fuse_query ) formatted_outputs = format_outputs_feedback( images, scores, retrieved_image_paths, [], # images_with_saliency - not used in current implementation [] # explanation - not used in current implementation ) # formatted_outputs = [gallery, retrieved_image_paths, *annotators] # Expected outputs: [error_display, image_gallery, relevant_image_paths, *annotators] gallery, retrieved_paths, *annotator_outputs = formatted_outputs return [gr.HTML(visible=False), gallery, retrieved_paths] + annotator_outputs except Exception as e: logger.error(f"Apply feedback error: {str(e)}") apply_feedback_btn.click( fn=handle_apply_feedback, inputs=[query, image_top_k, relevant_features, irrelevant_features, fuse_initial_query], outputs=[error_display, image_gallery, relevant_image_paths, *annotators], ).then( fn=lambda: [None for _ in annotator_json_boxes_list], inputs=None, outputs=[*annotator_json_boxes_list] ) # Server health check with gr.Tab("Server Status"): gr.Markdown("## Server Information") async def check_server_health(): try: response = await retrieval_client.health() return response except Exception as e: logger.error(f"Server health check error: {str(e)}") return None health_check_btn = gr.Button("Check Server Health") health_display = gr.JSON() health_check_btn.click( fn=check_server_health, outputs=[health_display] ) # Cleanup function async def cleanup(): """Cleanup resources""" try: await retrieval_client.close() logger.info("Client cleanup completed") except Exception as e: logger.error(f"Cleanup error: {str(e)}") if __name__ == "__main__": try: logger.info(f"Starting client on port {args.port}") demo.launch( server_port=args.port, share=args.share, show_error=True ) except KeyboardInterrupt: logger.info("Shutting down client...") except Exception as e: logger.error(f"Client startup error: {str(e)}") finally: # Cleanup asyncio.run(cleanup())