File size: 9,840 Bytes
125214f
 
 
 
 
 
7e54217
125214f
 
 
 
 
1812270
eff41fa
40d77a8
125214f
 
1948116
7c3b785
 
 
 
 
 
 
 
1948116
7c3b785
 
 
eaa7a81
1948116
eaa7a81
40d77a8
1948116
 
7c3b785
 
 
 
 
 
 
 
 
 
 
 
1948116
 
 
 
 
 
 
7c3b785
 
 
 
 
 
 
 
1948116
d4b85b8
7c3b785
 
1948116
5667733
 
29f316e
1948116
 
 
 
d4b85b8
1948116
 
 
 
677e938
7316948
0bac0de
 
180f51e
8cedf13
 
 
 
824fe45
522b5f3
677e938
824fe45
 
2b70006
8cedf13
 
 
 
 
c791f22
40d77a8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8cedf13
bb17a59
 
1948116
3394a6e
1948116
7c3b785
 
 
 
 
40d77a8
2fdc9bb
3b61686
 
 
 
 
1948116
 
7c8c861
6a338ab
1948116
5600c91
 
 
 
 
3a20d92
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
 
 
7e54217
6b844f6
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
202
203
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, model):
        """
        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.
            model (KBVQA): The loaded KBVQA model used for generating the answer.

        Returns:
            str: The generated answer to the question.
        """
        free_gpu_resources()
        answer = model.generate_answer(question, caption, detected_objects_str)
        free_gpu_resources()
        return answer


    def image_qa_app(self, kbvqa):
        """
        Main application interface for image-based question answering. It handles displaying
        of sample images, uploading of new images, and facilitates the QA process.

        Args:
            kbvqa (KBVQA): The loaded KBVQA model used for image analysis and question answering.
        """
        
        # 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, kbvqa)

        # 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), kbvqa)

        # 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:]}')
                    if not image_data['analysis_done']:
                        nested_col22.text("Please click 'Analyze Image'..")
                        with nested_col22:
                            if st.button('Analyze Image', key=f'analyze_{image_key}', on_click=self.disable_widgets, disabled=self.is_widget_disabled):
                                
                                caption, detected_objects_str, image_with_boxes = self.analyze_image(image_data['image'], kbvqa)
                                self.update_image_data(image_key, caption, detected_objects_str, True)
                            st.session_state['loading_in_progress'] = False
        
                    # Initialize qa_history for each image
                    qa_history = image_data.get('qa_history', [])
        
                    if image_data['analysis_done']:
                        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 not question:
                            nested_col22.warning("Please select or enter a question.")
                        else:
                            if question in [q for q, _ in qa_history]:
                                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):

                                    answer = self.answer_question(image_data['caption'], image_data['detected_objects_str'], question, kbvqa)
                                    st.session_state['loading_in_progress'] = False
                                    self.add_to_qa_history(image_key, question, answer)
        
                    # Display Q&A history for each image
                    for num, (q, a) in enumerate(qa_history):
                        nested_col22.text(f"Q{num+1}: {q}\nA{num+1}: {a}\n")

        
    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
        
        st.session_state['settings_changed'] = self.has_state_changed()
        if st.session_state['settings_changed']:
            self.col1.warning("Model settings have changed, please reload the model, this will take a second .. ")
           
        st.session_state.button_label = "Reload Model" if self.is_model_loaded() and self.settings_changed 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(self.get_model())