import streamlit as st import torch import bitsandbytes import accelerate import scipy import copy from PIL import Image import torch.nn as nn 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 def answer_question(caption, detected_objects_str, question, model): answer = model.generate_answer(question, caption, detected_objects_str) return answer def get_caption(image): return "Generated caption for the image" def free_gpu_resources(): pass # 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): st.write("Analyzing . . .") caption = model.get_caption(image) image_with_boxes, detected_objects_str = model.detect_objects(image) return caption, detected_objects_str 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[-8:]}', use_column_width=True) if not image_data['analysis_done']: if st.button('Analyze Image', key=f'analyze_{image_key}'): caption, detected_objects_str = analyze_image(image_data['image'], kbvqa) image_data['caption'] = caption image_data['detected_objects_str'] = detected_objects_str image_data['analysis_done'] = True if image_data['analysis_done']: question = st.text_input(f"Ask a question about this image ({image_key}):", key=f'question_{image_key}') if st.button('Get Answer', key=f'answer_{image_key}'): answer = answer_question(image_data['caption'], image_data['detected_objects_str'], question, kbvqa) image_data['qa_history'].append((question, answer)) for q, a in image_data['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") method = st.selectbox( "Choose a method:", ["Fine-Tuned Model", "In-Context Learning (n-shots)"], index=0 # Default to the first option ) detection_model = st.selectbox( "Choose a model for object detection:", ["yolov5", "detic"], index=0 # Default to the first option ) # Set default confidence based on the selected model default_confidence = 0.2 if detection_model == "yolov5" else 0.4 # Slider for confidence level confidence_level = st.slider( "Select minimum detection confidence level", min_value=0.1, max_value=0.9, value=default_confidence, step=0.1 ) # Initialize session state for the model if method == "Fine-Tuned Model": if 'kbvqa' not in st.session_state: st.session_state['kbvqa'] = None # Button to load KBVQA models if st.button('Load Model'): if st.session_state['kbvqa'] is not None: st.write("Model already loaded.") else: # Call the function to load models and show progress 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 if st.session_state['kbvqa']: st.write("Model is ready for inference.") if st.session_state['kbvqa']: image_qa_app(st.session_state['kbvqa']) else: st.write('Model is not ready for inference yet') # Main function def main(): st.sidebar.title("Navigation") selection = st.sidebar.radio("Go to", ["Home", "Dataset Analysis", "Evaluation Results", "Run Inference", "Dissertation Report"]) st.sidebar.write("More Pages will follow .. ") if selection == "Home": st.title("MultiModal Learning for Knowledg-Based Visual Question Answering") st.write("Home page content goes here...") elif selection == "Dissertation Report": st.title("Dissertation Report") st.write("Click the link below to view the PDF.") # Example to display a link to a PDF st.download_button( label="Download PDF", data=open("Files/Dissertation Report.pdf", "rb"), file_name="example.pdf", mime="application/octet-stream" ) elif selection == "Evaluation Results": st.title("Evaluation Results") st.write("This is a Place Holder until the contents are uploaded.") elif selection == "Dataset Analysis": st.title("OK-VQA Dataset Analysis") st.write("This is a Place Holder until the contents are uploaded.") elif selection == "Run Inference": run_inference() elif selection == "More Pages will follow .. ": st.title("Staye Tuned") st.write("This is a Place Holder until the contents are uploaded.") if __name__ == "__main__": main()