Update my_model/tabs/run_inference.py

my_model/tabs/run_inference.py

@@ -47,30 +47,31 @@ class InferenceRunner(StateManager):
 
         # Display and interact with each uploaded/selected image
         with self.col2:
-            nested_col21, nested_col22 = st.columns([0.7, 0.3])
-            for image_key, image_data in self.get_images_data().items():
-                image_for_display = self.resize_image(image_data['image'], 600)
-                nested_col21.image(image_for_display, caption=f'Uploaded Image: {image_key[-11:]}')
-                if not image_data['analysis_done']:
-                    nested_col22.text("Please click 'Analyze Image'..")
-                    if nested_col22.button('Analyze Image', key=f'analyze_{image_key}'):
-                        caption, detected_objects_str, image_with_boxes = self.analyze_image(image_data['image'], kbvqa)
-                        self.update_image_data(image_key, caption, detected_objects_str, True)
-    
-                # Initialize qa_history for each image
-                qa_history = image_data.get('qa_history', [])
-    
-                if image_data['analysis_done']:
-                    question = nested_col22.text_input(f"Ask a question about this image ({image_key[-11:]}):", key=f'question_{image_key}')
-                    if nested_col22.button('Get Answer', key=f'answer_{image_key}'):
-                        if question not in [q for q, _ in qa_history]:
-                            answer = self.answer_question(image_data['caption'], image_data['detected_objects_str'], question, kbvqa)
-                            self.add_to_qa_history(image_key, question, answer)
-                        else: nested_col22.warning("This questions has already been answered.")
-    
-                # Display Q&A history for each image
-                for q, a in qa_history:
-                    nested_col22.text(f"Q: {q}\nA: {a}\n")
+            with st.container():
+                nested_col21, nested_col22 = st.columns([0.7, 0.3])
+                for image_key, image_data in self.get_images_data().items():
+                    image_for_display = self.resize_image(image_data['image'], 600)
+                    nested_col21.image(image_for_display, caption=f'Uploaded Image: {image_key[-11:]}')
+                    if not image_data['analysis_done']:
+                        nested_col22.text("Please click 'Analyze Image'..")
+                        if nested_col22.button('Analyze Image', key=f'analyze_{image_key}'):
+                            caption, detected_objects_str, image_with_boxes = self.analyze_image(image_data['image'], kbvqa)
+                            self.update_image_data(image_key, caption, detected_objects_str, True)
+        
+                    # Initialize qa_history for each image
+                    qa_history = image_data.get('qa_history', [])
+        
+                    if image_data['analysis_done']:
+                        question = nested_col22.text_input(f"Ask a question about this image ({image_key[-11:]}):", key=f'question_{image_key}')
+                        if nested_col22.button('Get Answer', key=f'answer_{image_key}'):
+                            if question not in [q for q, _ in qa_history]:
+                                answer = self.answer_question(image_data['caption'], image_data['detected_objects_str'], question, kbvqa)
+                                self.add_to_qa_history(image_key, question, answer)
+                            else: nested_col22.warning("This questions has already been answered.")
+        
+                    # Display Q&A history for each image
+                    for q, a in qa_history:
+                        nested_col22.text(f"Q: {q}\nA: {a}\n")
 
 
     def run_inference(self):