Update my_model/tabs/run_inference.py

my_model/tabs/run_inference.py

@@ -32,8 +32,8 @@ class InferenceRunner(StateManager):
 
     def image_qa_app(self, kbvqa):
         # Display sample images as clickable thumbnails
-        st.write("Choose from sample images:")
-        cols = st.columns(len(self.sample_images))
+        self.col1.write("Choose from sample images:")
+        cols = self.col1.columns(len(self.sample_images))
         for idx, sample_image_path in enumerate(self.sample_images):
             with cols[idx]:
                 image = Image.open(sample_image_path)
@@ -42,16 +42,16 @@ class InferenceRunner(StateManager):
                     self.process_new_image(sample_image_path, image, kbvqa)
 
         # Image uploader
-        uploaded_image = st.file_uploader("Or upload an Image", type=["png", "jpg", "jpeg"])
+        uploaded_image = self.col1.file_uploader("Or upload an Image", type=["png", "jpg", "jpeg"])
         if uploaded_image is not None:
             self.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 self.get_images_data().items():
-            st.image(image_data['image'], caption=f'Uploaded Image: {image_key[-11:]}', use_column_width=True)
+            col2.image(image_data['image'], caption=f'Uploaded Image: {image_key[-11:]}', use_column_width=True)
             if not image_data['analysis_done']:
-                st.text("Cool image, please click 'Analyze Image'..")
-                if st.button('Analyze Image', key=f'analyze_{image_key}'):
+                self.col2.text("Cool image, please click 'Analyze Image'..")
+                if col2.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)
 
@@ -59,8 +59,8 @@ class InferenceRunner(StateManager):
             qa_history = image_data.get('qa_history', [])
 
             if image_data['analysis_done']:
-                question = st.text_input(f"Ask a question about this image ({image_key[-11:]}):", key=f'question_{image_key}')
-                if st.button('Get Answer', key=f'answer_{image_key}'):
+                question = col2.text_input(f"Ask a question about this image ({image_key[-11:]}):", key=f'question_{image_key}')
+                if col2.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)
@@ -77,15 +77,15 @@ class InferenceRunner(StateManager):
         self.set_up_widgets()
         st.session_state['settings_changed'] = self.has_state_changed()
         if st.session_state['settings_changed']:
-            st.warning("Model settings have changed, please reload the model, this will take a second .. ")
+            col1.warning("Model settings have changed, please reload the model, this will take a second .. ")
 
         st.session_state.button_label = "Reload Model" if self.is_model_loaded() and self.settings_changed else "Load Model"
         
         if st.session_state.method == "Fine-Tuned Model":
-            if st.button(st.session_state.button_label):
+            if col1.button(st.session_state.button_label):
                 if st.session_state.button_label == "Load Model":
                     if self.is_model_loaded():
-                        st.text("Model already loaded and no settings were changed:)")
+                        col1.text("Model already loaded and no settings were changed:)")
                     else: 
                         self.load_model()
                 else: