Update app.py

app.py

@@ -3,6 +3,7 @@ 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
@@ -64,10 +65,12 @@ def image_qa_app(kbvqa):
     cols = st.columns(len(sample_images))
     for idx, sample_image_path in enumerate(sample_images):
         with cols[idx]:
+            # Display each sample image with a button
             image = Image.open(sample_image_path)
-            # Use button for each image
-            if st.button(f'Sample Image {idx + 1}', key=f'sample_{idx}'):
+            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['processed_image'] = copy.deepcopy(image)
                 st.session_state['qa_history'] = []
 
     # Image uploader
@@ -77,7 +80,6 @@ def image_qa_app(kbvqa):
         st.session_state['current_image'] = image
         st.session_state['processed_image'] = copy.deepcopy(image)
         st.session_state['qa_history'] = []
-        print("Image uploaded and processed image set.")
 
     # Display the current image (unaltered)
     if st.session_state.get('current_image') is not None:
@@ -88,17 +90,13 @@ def image_qa_app(kbvqa):
         if st.button('Get Answer'):
             processed_image = st.session_state.get('processed_image')
             if processed_image:
-                print("Processing image for question:", question)
                 answer = answer_question(processed_image, question, model=kbvqa)
                 st.session_state['qa_history'].append((question, answer))
-                print("Answer generated:", answer)
 
             # Display all Q&A
             for q, a in st.session_state['qa_history']:
                 st.text(f"Q: {q}\nA: {a}\n")
 
-            st.session_state['processed_image'] = processed_image
-
             
 # Main function
 def main():