File size: 5,828 Bytes
125214f 1812270 eff41fa 125214f 1948116 eaa7a81 1948116 eaa7a81 1948116 8bd4ecd 1948116 d4b85b8 1948116 29f316e d080299 29f316e 1948116 d4b85b8 1948116 29f316e 1948116 d4b85b8 7c8c861 1948116 7c8c861 1948116 d6dc998 1948116 7c8c861 1948116 11781a6 0d909f9 11781a6 ab428b9 11781a6 2ec2d03 c327e49 ab428b9 11781a6 1948116 7c8c861 1948116 5ac6aec 7c8c861 5ac6aec 1948116 5ac6aec 1948116 5ac6aec 1948116 7c8c861 |
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 |
import streamlit as st
import torch
import bitsandbytes
import accelerate
import scipy
import copy
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
class InferenceRunner(StateManager):
def __init__(self):
super().__init__()
self.initialize_state()
self.sample_images = [
"Files/sample1.jpg", "Files/sample2.jpg", "Files/sample3.jpg"]
def answer_question(self, caption, detected_objects_str, question, model):
free_gpu_resources()
answer = model.generate_answer(question, caption, detected_objects_str)
free_gpu_resources()
return answer
def image_qa_app(self, kbvqa):
# Display sample images as clickable thumbnails
self.col1.write("Choose from sample images:")
cols = self.col1.columns(len(self.sample_images))
for idx, sample_image_path in enumerate(self.sample_images):
with cols[idx]:
image = Image.open(sample_image_path)
image_for_display = copy.deepcopy(image) # resize just for the display control without changing original image
image_for_display = self.resize_image(image_for_display, 500, 500)
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
for image_key, image_data in self.get_images_data().items():
image_for_display = copy.deepcopy(image_data['image']) # resize just for the display control without changing original image
image_for_display = self.resize_image(image_for_display, 800, 800)
self.col2.image(image_for_display, caption=f'Uploaded Image: {image_key[-11:]}')
if not image_data['analysis_done']:
self.col2.text("Cool image, please click 'Analyze Image'..")
if self.col2.button('Analyze Image', key=f'analyze_{image_key}'):
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)
# Initialize qa_history for each image
qa_history = image_data.get('qa_history', [])
if image_data['analysis_done']:
question = self.col2.text_input(f"Ask a question about this image ({image_key[-11:]}):", key=f'question_{image_key}')
if self.col2.button('Get Answer', key=f'answer_{image_key}'):
if question not in [q for q, _ in qa_history]:
answer = self.answer_question(image_data['caption'], image_data['detected_objects_str'], question, kbvqa)
self.add_to_qa_history(image_key, question, answer)
# Display Q&A history for each image
for q, a in qa_history:
self.col2.text(f"Q: {q}\nA: {a}\n")
def run_inference(self):
self.set_up_widgets()
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 .. ")
###################################
self.col1.write("Choose from sample images:")
cols = self.col1.columns(len(self.sample_images))
for idx, sample_image_path in enumerate(self.sample_images):
with cols[idx]:
image = Image.open(sample_image_path)
image_for_display = copy.deepcopy(image) # resize just for the display control without changing original image
image_for_display = self.resize_image(image_for_display, 500, 500)
st.image(image_for_display)
# Image uploader
uploaded_image = self.col1.file_uploader("Or upload an Image", type=["png", "jpg", "jpeg"])
image_for_display = copy.deepcopy(uploaded_image) # resize just for the display control without changing original image
image_for_display = self.resize_image(image_for_display, 500, 500)
############################################
st.session_state.button_label = "Reload Model" if self.is_model_loaded() and self.settings_changed else "Load Model"
if st.session_state.method == "Fine-Tuned Model":
if self.col1.button(st.session_state.button_label):
if st.session_state.button_label == "Load Model":
if self.is_model_loaded():
free_gpu_resources()
self.col1.text("Model already loaded and no settings were changed:)")
else:
free_gpu_resources()
self.load_model()
else:
free_gpu_resources()
self.reload_detection_model()
if self.is_model_loaded() and st.session_state.kbvqa.all_models_loaded:
free_gpu_resources()
self.image_qa_app(self.get_model())
else:
self.col1.warning(f'Model using {st.session_state.method} is not deployed yet, will be ready later.') |