|
import pandas as pd |
|
import copy |
|
import streamlit as st |
|
from my_model.gen_utilities import free_gpu_resources |
|
from my_model.KBVQA import KBVQA, prepare_kbvqa_model |
|
|
|
|
|
class StateManager: |
|
def __init__(self): |
|
self.initialize_state() |
|
|
|
def initialize_state(self): |
|
if 'images_data' not in st.session_state: |
|
st.session_state['images_data'] = {} |
|
if 'method' not in st.session_state: |
|
st.session_state['method'] = None |
|
if 'detection_model' not in st.session_state: |
|
st.session_state['detection_model'] = None |
|
if 'kbvqa' not in st.session_state: |
|
st.session_state['kbvqa'] = None |
|
if 'confidence_level' not in st.session_state: |
|
st.session_state['confidence_level'] = None |
|
|
|
|
|
|
|
|
|
def set_slider_value(self, text, min_value, max_value, value, step, slider_key_name): |
|
|
|
return st.slider(text, min_value, max_value, value, step, key=slider_key_name) |
|
|
|
|
|
def check_settings_changed(self, current_selected_method, current_detection_model, current_confidence_level): |
|
return (st.session_state['model_settings']['detection_model'] != current_detection_model or |
|
st.session_state['model_settings']['confidence_level'] != current_confidence_level or |
|
st.session_state['model_settings']['selected_method'] != current_selected_method) |
|
|
|
def display_model_settings(self): |
|
st.write("### Current Model Settings:") |
|
data = [{'Key': key, 'Value': str(value)} for key, value in st.session_state.items() if key in ["confidence_level", 'detection_model', 'method', 'kbvqa']] |
|
st.table(pd.DataFrame(data)) |
|
|
|
def display_session_state(self): |
|
st.write("### Current Model:") |
|
data = [{'Key': key, 'Value': str(value)} for key, value in st.session_state.items()] |
|
df = pd.DataFrame(data) |
|
st.table(df) |
|
|
|
def load_model(self): |
|
"""Load the KBVQA model with specified settings.""" |
|
try: |
|
free_gpu_resources() |
|
st.text("Loading the model, this should take no more than a few minutes, please wait...") |
|
st.session_state['kbvqa'] = prepare_kbvqa_model(st.session_state.detection_model) |
|
st.session_state['kbvqa'].detection_confidence = st.session_state.confidence_level |
|
|
|
st.text("Model is ready for inference.") |
|
free_gpu_resources() |
|
except Exception as e: |
|
st.error(f"Error loading model: {e}") |
|
|
|
def get_model(self): |
|
"""Retrieve the KBVQA model from the session state.""" |
|
return st.session_state.get('kbvqa', None) |
|
|
|
def is_model_loaded(self): |
|
return 'kbvqa' in st.session_state and st.session_state['kbvqa'] is not None |
|
|
|
def reload_detection_model(self, detection_model, confidence_level): |
|
try: |
|
free_gpu_resources() |
|
if self.is_model_loaded(): |
|
prepare_kbvqa_model(detection_model, only_reload_detection_model=True) |
|
st.session_state['kbvqa'].detection_confidence = confidence_level |
|
|
|
free_gpu_resources() |
|
except Exception as e: |
|
st.error(f"Error reloading detection model: {e}") |
|
|
|
|
|
def process_new_image(self, image_key, image, kbvqa): |
|
if image_key not in st.session_state['images_data']: |
|
st.session_state['images_data'][image_key] = { |
|
'image': image, |
|
'caption': '', |
|
'detected_objects_str': '', |
|
'qa_history': [], |
|
'analysis_done': False |
|
} |
|
|
|
def analyze_image(self, image, kbvqa): |
|
img = copy.deepcopy(image) |
|
caption = kbvqa.get_caption(img) |
|
image_with_boxes, detected_objects_str = kbvqa.detect_objects(img) |
|
return caption, detected_objects_str, image_with_boxes |
|
|
|
def add_to_qa_history(self, image_key, question, answer): |
|
if image_key in st.session_state['images_data']: |
|
st.session_state['images_data'][image_key]['qa_history'].append((question, answer)) |
|
|
|
def get_images_data(self): |
|
return st.session_state['images_data'] |
|
|
|
def update_image_data(self, image_key, caption, detected_objects_str, analysis_done): |
|
if image_key in st.session_state['images_data']: |
|
st.session_state['images_data'][image_key].update({ |
|
'caption': caption, |
|
'detected_objects_str': detected_objects_str, |
|
'analysis_done': analysis_done |
|
}) |
|
|
