File size: 9,816 Bytes
125214f 7e54217 125214f 1812270 eff41fa 40d77a8 125214f 1948116 7c3b785 1948116 7c3b785 eaa7a81 1948116 40d77a8 1948116 df65239 7c3b785 1948116 6657600 76f3951 1948116 0b0cffd 1948116 8f568ec d4b85b8 7c3b785 1948116 5667733 29f316e 1948116 df65239 1948116 0f0c882 d4b85b8 1948116 df65239 1948116 256561a 0f0c882 f283cf3 0f0c882 f283cf3 0f0c882 1d707c1 0f0c882 256561a 0f0c882 f283cf3 094a401 0f0c882 1d707c1 0f0c882 7316948 0bac0de 82e0944 256561a 0f0c882 c20559f 1948116 3394a6e 1948116 7c3b785 40d77a8 2fdc9bb 4a6ec58 3b61686 4a6ec58 da1ffd0 7c8c861 c809ab0 4a6ec58 6a338ab 53586ab 4a6ec58 5600c91 3a20d92 5600c91 e36aa58 5600c91 3394a6e 5600c91 3394a6e 3db717c 3394a6e 3a20d92 3394a6e fc89ea0 3394a6e c45c1b1 fc89ea0 904c909 8b8fe48 c791f22 904c909 802de9d c791f22 904c909 c791f22 6b844f6 904c909 802de9d fc89ea0 904c909 8b8fe48 c791f22 7e54217 7c3b785 3394a6e 6b844f6 c791f22 6b844f6 2111435 6b844f6 7e54217 df65239 677e938 6b844f6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 |
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__()
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 display_sample_images(self):
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)
def handle_image_upload(self):
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))
def display_image_and_analysis(self, image_key, image_data, nested_col21, nested_col22):
image_for_display = self.resize_image(image_data['image'], 600)
nested_col21.image(image_for_display, caption=f'Uploaded Image: {image_key[-11:]}')
self.handle_analysis_button(image_key, image_data, nested_col22)
def handle_analysis_button(self, image_key, image_data, nested_col22):
if not image_data['analysis_done'] or self.settings_changed or self.confidance_change:
nested_col22.text("Please click 'Analyze Image'..")
analyze_button_key = f'analyze_{image_key}_{st.session_state.detection_model}_{st.session_state.confidence_level}'
if nested_col22.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
def handle_question_answering(self, image_key, image_data, nested_col22):
# Initialize qa_history for each image
#qa_history = image_data.get('qa_history', [])
if image_data['analysis_done']:
self.display_question_answering_interface(image_key, image_data, nested_col22)
if self.settings_changed or self.confidance_change:
nested_col22.warning("Confidence level changed, please click 'Analyze Image'.")
def display_question_answering_interface(self, image_key, image_data, nested_col22):
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}')
custom_question = nested_col22.text_input("Or ask your own question:", key=f'custom_question_{image_key}')
question = custom_question if selected_question == "Custom question..." else selected_question
self.process_question(image_key, question, image_data, nested_col22)
qa_history = image_data.get('qa_history', [])
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")
def process_question(self, image_key, question, image_data, nested_col22):
qa_history = image_data.get('qa_history', [])
if question and (question not in [q for q, _, _ in qa_history] or self.settings_changed or self.confidance_change):
if nested_col22.button('Get Answer', key=f'answer_{image_key}', disabled=self.is_widget_disabled):
answer, prompt_length = self.answer_question(image_data['caption'], image_data['detected_objects_str'], question)
self.add_to_qa_history(image_key, question, answer, prompt_length)
# nested_col22.text(f"Q: {question}\nA: {answer}\nPrompt Length: {prompt_length}")
def image_qa_app(self):
self.display_sample_images()
self.handle_image_upload()
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])
self.display_image_and_analysis(image_key, image_data, nested_col21, nested_col22)
self.handle_question_answering(image_key, image_data, nested_col22)
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()
|