Update my_model/tabs/run_inference.py

my_model/tabs/run_inference.py

@@ -0,0 +1,188 @@
+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.gen_utilities import free_gpu_resources
+from my_model.KBVQA import KBVQA, prepare_kbvqa_model
+from my_model.utilities.st_utils import UIManager, StateManager
+
+
+
+def answer_question(caption, detected_objects_str, question, model):
+
+    answer = model.generate_answer(question, caption, detected_objects_str)
+    return answer
+
+
+# Sample images (assuming these are paths to your sample images)
+sample_images = ["Files/sample1.jpg", "Files/sample2.jpg", "Files/sample3.jpg", 
+                 "Files/sample4.jpg", "Files/sample5.jpg", "Files/sample6.jpg", 
+                 "Files/sample7.jpg"]
+
+
+
+def analyze_image(image, model):
+    
+    img = copy.deepcopy(image)  # we dont wanna apply changes to the original image
+    caption = model.get_caption(img)
+    image_with_boxes, detected_objects_str = model.detect_objects(img)
+    st.text("I am ready, let's talk!")
+    free_gpu_resources()
+    
+    return caption, detected_objects_str, image_with_boxes
+    
+
+def image_qa_app(kbvqa):
+    if 'images_data' not in st.session_state:
+        st.session_state['images_data'] = {}
+
+    # Display sample images as clickable thumbnails
+    st.write("Choose from sample images:")
+    cols = st.columns(len(sample_images))
+    for idx, sample_image_path in enumerate(sample_images):
+        with cols[idx]:
+            image = Image.open(sample_image_path)
+            st.image(image, use_column_width=True)
+            if st.button(f'Select Sample Image {idx + 1}', key=f'sample_{idx}'):
+                process_new_image(sample_image_path, image, kbvqa)
+
+    # Image uploader
+    uploaded_image = st.file_uploader("Or upload an Image", type=["png", "jpg", "jpeg"])
+    if uploaded_image is not None:
+        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 st.session_state['images_data'].items():
+        st.image(image_data['image'], caption=f'Uploaded Image: {image_key[-11:]}', use_column_width=True)
+        if not image_data['analysis_done']:
+            st.text("Cool image, please click 'Analyze Image'..")
+            if st.button('Analyze Image', key=f'analyze_{image_key}'):
+                caption, detected_objects_str, image_with_boxes = analyze_image(image_data['image'], kbvqa)  # we can use the image_with_boxes later if we want to show it. 
+                image_data['caption'] = caption
+                image_data['detected_objects_str'] = detected_objects_str
+                image_data['analysis_done'] = True
+
+        # Initialize qa_history for each image
+        qa_history = image_data.get('qa_history', [])
+
+        if image_data['analysis_done']:
+            question = st.text_input(f"Ask a question about this image ({image_key[-11:]}):", key=f'question_{image_key}')
+            if st.button('Get Answer', key=f'answer_{image_key}'):
+                if question not in [q for q, _ in qa_history]:
+                    answer = answer_question(image_data['caption'], image_data['detected_objects_str'], question, kbvqa)
+                    qa_history.append((question, answer))
+                    image_data['qa_history'] = qa_history
+                else:
+                    st.info("This question has already been asked.")
+
+        # Display Q&A history for each image
+        for q, a in qa_history:
+            st.text(f"Q: {q}\nA: {a}\n")
+
+
+def process_new_image(image_key, image, kbvqa):
+    """Process a new image and update the session state."""
+    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 run_inference():
+    st.title("Run Inference")
+    st.write("Please note that this is not a general purpose model, it is specifically trained on OK-VQA dataset and is designed to give direct and short answers to the given questions.")
+
+    method = st.selectbox(
+        "Choose a method:",
+        ["Fine-Tuned Model", "In-Context Learning (n-shots)"],
+        index=0
+    )
+
+    detection_model = st.selectbox(
+        "Choose a model for objects detection:",
+        ["yolov5", "detic"],
+        index=1  # "detic" is selected by default
+    )
+
+    default_confidence = 0.2 if detection_model == "yolov5" else 0.4
+    confidence_level = st.slider(
+        "Select minimum detection confidence level",
+        min_value=0.1,
+        max_value=0.9,
+        value=default_confidence,
+        step=0.1
+    )
+
+    if 'model_settings' not in st.session_state:
+        st.session_state['model_settings'] = {'detection_model': detection_model, 'confidence_level': confidence_level}
+
+    settings_changed = (st.session_state['model_settings']['detection_model'] != detection_model or
+                        st.session_state['model_settings']['confidence_level'] != confidence_level)
+
+    need_model_reload = settings_changed and 'kbvqa' in st.session_state and st.session_state['kbvqa'] is not None
+    
+    if need_model_reload:
+        st.text("Model Settings have changed, please reload the model, this will take no time :)")
+
+    button_label = "Reload Model" if need_model_reload else "Load Model"
+
+    if method == "Fine-Tuned Model":
+        if 'kbvqa' not in st.session_state:
+            st.session_state['kbvqa'] = None
+
+        if st.button(button_label):
+            
+            free_gpu_resources()
+            if st.session_state['kbvqa'] is not None:
+                if not settings_changed:
+                    st.write("Model already loaded.")
+                else: 
+                    free_gpu_resources()
+                    detection_model = st.session_state['model_settings']['detection_model']
+                    confidence_level = st.session_state['model_settings']['confidence_level']
+                    prepare_kbvqa_model(detection_model, only_reload_detection_model=True)  # only reload detection model with new settings
+                    st.session_state['kbvqa'].detection_confidence = confidence_level
+                    free_gpu_resources()
+            else:
+                st.text("Loading the model will take no more than a few minutes . .")
+                st.session_state['kbvqa'] = prepare_kbvqa_model(detection_model)
+                st.session_state['kbvqa'].detection_confidence = confidence_level
+                st.session_state['model_settings'] = {'detection_model': detection_model, 'confidence_level': confidence_level}
+                st.write("Model is ready for inference.")
+                free_gpu_resources()
+
+                
+
+        if st.session_state['kbvqa']:
+            display_model_settings()
+            display_session_state()
+            image_qa_app(st.session_state['kbvqa'])
+
+    else:
+        st.write('Model is not ready yet, will be updated later.')
+
+
+def display_model_settings():
+    st.write("### Current Model Settings:")
+    st.table(pd.DataFrame(st.session_state['model_settings'], index=[0]))
+
+def display_session_state():
+    st.write("### Current Session State:")
+    # Convert session state to a list of dictionaries, each representing a row
+    data = [{'Key': key, 'Value': str(value)} for key, value in st.session_state.items()]
+    # Create a DataFrame from the list
+    df = pd.DataFrame(data)
+    st.table(df)
+
+
+