Update app.py

app.py

@@ -29,29 +29,16 @@ sample_images = ["Files/sample1.jpg", "Files/sample2.jpg", "Files/sample3.jpg",
                  "Files/sample4.jpg", "Files/sample5.jpg", "Files/sample6.jpg", 
                  "Files/sample7.jpg"]
 
-def run_inference():
-    st.title("Run Inference")
-
-    # Initialize session state for the 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('yolov5')  # Replace with your model
 
-        if st.session_state['kbvqa']:
-            st.write("Model is ready for inference.")
-        else:
-            st.write("Please load the model first")
-
-    if st.session_state['kbvqa']:
-        image_qa_app(st.session_state['kbvqa'])
 
+def analyze_image(image, model):
+    # Placeholder for your analysis function
+    # This function should prepare captions, detect objects, etc.
+    # For example:
+    # caption = model.get_caption(image)
+    # detected_objects = model.detect_objects(image)
+    # return caption, detected_objects
+    pass
 
 def image_qa_app(kbvqa):
     # Initialize session state for storing the current image and its Q&A history
@@ -59,43 +46,71 @@ def image_qa_app(kbvqa):
         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
 
     # 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]:
-            # Display each sample image with a button
             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['processed_image'] = copy.deepcopy(image)
                 st.session_state['qa_history'] = []
+                st.session_state['analysis_done'] = 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['processed_image'] = copy.deepcopy(image)
         st.session_state['qa_history'] = []
+        st.session_state['analysis_done'] = False
+
+    # Analyze Image button
+    if st.session_state.get('current_image') and not st.session_state['analysis_done']:
+        if st.button('Analyze Image'):
+            # Perform analysis on the image
+            analyze_image(st.session_state['current_image'], kbvqa)
+            st.session_state['analysis_done'] = True
 
     # Display the current image (unaltered)
-    if st.session_state.get('current_image') is not None:
+    if st.session_state.get('current_image'):
         st.image(st.session_state['current_image'], caption='Uploaded Image.', use_column_width=True)
 
-        # Question input and processing
+    # Get Answer button
+    if st.session_state['analysis_done']:
         question = st.text_input("Ask a question about this image:")
         if st.button('Get Answer'):
-            processed_image = st.session_state.get('processed_image')
-            if processed_image:
-                answer = answer_question(processed_image, question, model=kbvqa)
-                st.session_state['qa_history'].append((question, answer))
-
-            # Display all Q&A
-            for q, a in st.session_state['qa_history']:
-                st.text(f"Q: {q}\nA: {a}\n")
+            answer = answer_question(st.session_state['current_image'], question, model=kbvqa)
+            st.session_state['qa_history'].append((question, answer))
+
+        # Display all Q&A
+        for q, a in st.session_state['qa_history']:
+            st.text(f"Q: {q}\nA: {a}\n")
+
+def run_inference():
+    st.title("Run Inference")
+
+    # Initialize session state for the 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('yolov5')  # Replace with your model
+
+        if st.session_state['kbvqa']:
+            st.write("Model is ready for inference.")
+
+    if st.session_state['kbvqa']:
+        image_qa_app(st.session_state['kbvqa'])
 
             
 # Main function