Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import torch | |
| import bitsandbytes | |
| import accelerate | |
| import scipy | |
| import copy | |
| import time | |
| from PIL import Image | |
| import torch.nn as nn | |
| import pandas as pd | |
| from my_model.object_detection import detect_and_draw_objects | |
| from my_model.captioner.image_captioning import get_caption | |
| from my_model.utilities.gen_utilities import free_gpu_resources | |
| from my_model.state_manager import StateManager | |
| from my_model.config import inference_config as config | |
| class InferenceRunner(StateManager): | |
| """ | |
| InferenceRunner manages the user interface and interactions for a Streamlit-based | |
| Knowledge-Based Visual Question Answering (KBVQA) application. It handles image uploads, | |
| displays sample images, and facilitates the question-answering process using the KBVQA model. | |
| it inherits the StateManager class. | |
| """ | |
| def __init__(self): | |
| """ | |
| Initializes the InferenceRunner instance, setting up the necessary state. | |
| """ | |
| super().__init__() | |
| # self.initialize_state() | |
| def answer_question(self, caption, detected_objects_str, question): | |
| """ | |
| Generates an answer to a given question based on the image's caption and detected objects. | |
| Args: | |
| caption (str): The caption generated for the image. | |
| detected_objects_str (str): String representation of objects detected in the image. | |
| question (str): The user's question about the image. | |
| Returns: | |
| str: The generated answer to the question. | |
| """ | |
| free_gpu_resources() | |
| answer = st.session_state.kbvqa.generate_answer(question, caption, detected_objects_str) | |
| prompt_length = st.session_state.kbvqa.current_prompt_length | |
| free_gpu_resources() | |
| return answer, prompt_length | |
| def image_qa_app(self): | |
| """ | |
| Main application interface for image-based question answering. It handles displaying | |
| of sample images, uploading of new images, and facilitates the QA process. | |
| """ | |
| # Display sample images as clickable thumbnails | |
| self.col1.write("Choose from sample images:") | |
| cols = self.col1.columns(len(config.SAMPLE_IMAGES)) | |
| for idx, sample_image_path in enumerate(config.SAMPLE_IMAGES): | |
| with cols[idx]: | |
| image = Image.open(sample_image_path) | |
| image_for_display = self.resize_image(sample_image_path, 80, 80) | |
| st.image(image_for_display) | |
| if st.button(f'Select Sample Image {idx + 1}', key=f'sample_{idx}'): | |
| self.process_new_image(sample_image_path, image) | |
| # Image uploader | |
| uploaded_image = self.col1.file_uploader("Or upload an Image", type=["png", "jpg", "jpeg"]) | |
| if uploaded_image is not None: | |
| self.process_new_image(uploaded_image.name, Image.open(uploaded_image)) | |
| # Display and interact with each uploaded/selected image | |
| self.display_session_state() | |
| with self.col2: | |
| for image_key, image_data in self.get_images_data().items(): | |
| with st.container(): | |
| nested_col21, nested_col22 = st.columns([0.65, 0.35]) | |
| image_for_display = self.resize_image(image_data['image'], 600) | |
| nested_col21.image(image_for_display, caption=f'Uploaded Image: {image_key[-11:]}') | |
| nested_col21.write(image_data['analysis_done'] , self.settings_changed , self.confidance_change) | |
| if not image_data['analysis_done'] or self.settings_changed or self.confidance_change: # if not done analysis before or even done but settings changed, then we need to analyze again | |
| nested_col22.text("Please click 'Analyze Image'..") | |
| free_gpu_resources() | |
| with nested_col22: | |
| analyze_button_key = f'analyze_{image_key}_{st.session_state.detection_model}_{st.session_state.confidence_level}' # unique key for each click | |
| if st.button('Analyze Image', key=analyze_button_key, on_click=self.disable_widgets, disabled=self.is_widget_disabled): | |
| caption, detected_objects_str, image_with_boxes = self.analyze_image(image_data['image']) | |
| self.update_image_data(image_key, caption, detected_objects_str, True) | |
| st.session_state['loading_in_progress'] = False | |
| free_gpu_resources() | |
| # Initialize qa_history for each image | |
| qa_history = image_data.get('qa_history', []) | |
| if image_data['analysis_done']: | |
| free_gpu_resources() | |
| if self.confidance_change: | |
| nested_col22.warning("If you change the Confidence level, please click analyze again.") | |
| st.session_state['loading_in_progress'] = False | |
| sample_questions = config.SAMPLE_QUESTIONS.get(image_key, []) | |
| selected_question = nested_col22.selectbox( | |
| "Select a sample question or type your own:", | |
| ["Custom question..."] + sample_questions, | |
| key=f'sample_question_{image_key}') | |
| # Text input for custom question | |
| custom_question = nested_col22.text_input( | |
| "Or ask your own question:", | |
| key=f'custom_question_{image_key}') | |
| # Use the selected sample question or the custom question | |
| question = custom_question if selected_question == "Custom question..." else selected_question | |
| if question in [q for q, _, _ in qa_history] and not self.settings_changed and not self.confidance_change: | |
| nested_col22.warning("This question has already been answered.") | |
| else: | |
| if nested_col22.button('Get Answer', key=f'answer_{image_key}', disabled=self.is_widget_disabled): | |
| free_gpu_resources() | |
| answer, prompt_length = self.answer_question(image_data['caption'], image_data['detected_objects_str'], question) | |
| st.session_state['loading_in_progress'] = False | |
| self.add_to_qa_history(image_key, question, answer, prompt_length) | |
| # Display Q&A history and prompts lengths for each image | |
| for num, (q, a, p) in enumerate(qa_history): | |
| nested_col22.text(f"Q{num+1}: {q}\nA{num+1}: {a}\nPrompt Length: {p}\n") | |
| free_gpu_resources() | |
| def run_inference(self): | |
| """ | |
| Sets up the widgets and manages the inference process. This method handles model loading, | |
| reloading, and the overall flow of the inference process based on user interactions. | |
| """ | |
| self.set_up_widgets() | |
| load_fine_tuned_model = False | |
| fine_tuned_model_already_loaded = False | |
| reload_detection_model = False | |
| force_reload_full_model = False | |
| if self.is_model_loaded and self.settings_changed: | |
| self.col1.warning("Model settings have changed, please reload the model, this will take a second .. ") | |
| self.update_prev_state() | |
| st.session_state.button_label = "Reload Model" if self.is_model_loaded and st.session_state.kbvqa.detection_model != st.session_state['detection_model'] else "Load Model" | |
| with self.col1: | |
| if st.session_state.method == "Fine-Tuned Model": | |
| with st.container(): | |
| nested_col11, nested_col12 = st.columns([0.5, 0.5]) | |
| if nested_col11.button(st.session_state.button_label, on_click=self.disable_widgets, disabled=self.is_widget_disabled): | |
| if st.session_state.button_label == "Load Model": | |
| if self.is_model_loaded: | |
| free_gpu_resources() | |
| fine_tuned_model_already_loaded = True | |
| else: | |
| load_fine_tuned_model = True | |
| else: | |
| reload_detection_model = True | |
| if nested_col12.button("Force Reload", on_click=self.disable_widgets, disabled=self.is_widget_disabled): | |
| force_reload_full_model = True | |
| if load_fine_tuned_model: | |
| t1=time.time() | |
| free_gpu_resources() | |
| self.load_model() | |
| st.session_state['time_taken_to_load_model'] = int(time.time()-t1) | |
| st.session_state['loading_in_progress'] = False | |
| elif fine_tuned_model_already_loaded: | |
| free_gpu_resources() | |
| self.col1.text("Model already loaded and no settings were changed:)") | |
| st.session_state['loading_in_progress'] = False | |
| elif reload_detection_model: | |
| free_gpu_resources() | |
| self.reload_detection_model() | |
| st.session_state['loading_in_progress'] = False | |
| elif force_reload_full_model: | |
| free_gpu_resources() | |
| t1=time.time() | |
| self.force_reload_model() | |
| st.session_state['time_taken_to_load_model'] = int(time.time()-t1) | |
| st.session_state['loading_in_progress'] = False | |
| st.session_state['model_loaded'] = True | |
| elif st.session_state.method == "In-Context Learning (n-shots)": | |
| self.col1.warning(f'Model using {st.session_state.method} is not deployed yet, will be ready later.') | |
| st.session_state['loading_in_progress'] = False | |
| if self.is_model_loaded: | |
| free_gpu_resources() | |
| st.session_state['loading_in_progress'] = False | |
| self.image_qa_app() | |