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.utilities import free_gpu_resources from my_model.KBVQA import KBVQA, prepare_kbvqa_model def answer_question(image, question, caption, detected_objects_str, model): answer = model.generate_answer(question, caption, detected_objects_str) st.image(image) st.write(caption) st.write("----------------") st.write(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, show_processed_image=False): img = copy.deepcopy(image) caption = model.get_caption(img) image_with_boxes, detected_objects_str = model.detect_objects(img) if show_processed_image: st.image(image_with_boxes) return caption, detected_objects_str def image_qa_app(kbvqa): # Initialize session state for storing the current image and its Q&A history if 'current_image' not in st.session_state: st.session_state['current_image'] = None if 'qa_history' not in st.session_state: st.session_state['qa_history'] = [] if 'analysis_done' not in st.session_state: st.session_state['analysis_done'] = False if 'answer_in_progress' not in st.session_state: st.session_state['answer_in_progress'] = False # 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}'): st.session_state['current_image'] = image st.session_state['qa_history'] = [] st.session_state['analysis_done'] = False st.session_state['answer_in_progress'] = False # Image uploader uploaded_image = st.file_uploader("Or upload an Image", type=["png", "jpg", "jpeg"]) if uploaded_image is not None: image = Image.open(uploaded_image) st.session_state['current_image'] = image st.session_state['qa_history'] = [] st.session_state['analysis_done'] = False st.session_state['answer_in_progress'] = False if st.session_state.get('current_image') and not st.session_state.get('analysis_done', False): if st.button('Analyze Image'): caption, detected_objects_str = analyze_image(st.session_state['current_image'], kbvqa) st.session_state['caption'] = caption st.session_state['detected_objects_str'] = detected_objects_str st.session_state['analysis_done'] = True # Get Answer button if st.session_state.get('analysis_done', False): question = st.text_input("Ask a question about this image:") if st.button('Get Answer'): answer = answer_question(st.session_state['current_image'], question, st.session_state.get('caption', ''), st.session_state.get('detected_objects_str', ''), kbvqa) st.session_state['qa_history'].append((question, answer)) # Display all Q&A for q, a in st.session_state.get('qa_history', []): st.text(f"Q: {q}\nA: {a}\n") # Reset the answer_in_progress flag after displaying the answer if st.session_state['answer_in_progress']: st.session_state['answer_in_progress'] = 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 ) # 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 KBVQA 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.session_state['kbvqa'] = prepare_kbvqa_model(detection_model) if st.session_state['kbvqa']: st.write("Model is ready for inference.") # 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 Detection Confidence Level", min_value=0.1, max_value=0.9, value=default_confidence, step=0.1 ) st.session_state['kbvqa'].detection_confidence = confidence_level if st.session_state['kbvqa']: image_qa_app(st.session_state['kbvqa']) else: st.write(f'{method} 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"]) 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 == "Object Detection": run_object_detection() if __name__ == "__main__": main()