Update app.py

app.py

@@ -92,29 +92,18 @@ def process_new_image(image_key, image, kbvqa):
 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
-    )
-
-    # 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 minimum detection confidence level",
-        min_value=0.1,
-        max_value=0.9,
-        value=default_confidence,
-        step=0.1
-    )
+    method = st.selectbox("Choose a method:", ["Fine-Tuned Model", "In-Context Learning (n-shots)"], index=0)
+    detection_model = st.selectbox("Choose a model for object detection:", ["yolov5", "detic"], index=0)
+    confidence_level = st.slider("Select minimum detection confidence level", min_value=0.1, max_value=0.9, value=0.2 if detection_model == "yolov5" else 0.4, step=0.1)
+
+    # Check for changes in model or confidence level
+    model_changed = (st.session_state.get('detection_model') != detection_model)
+    confidence_changed = (st.session_state.get('confidence_level') != confidence_level)
+
+    if model_changed or confidence_changed:
+        st.session_state['detection_model'] = detection_model
+        st.session_state['confidence_level'] = confidence_level
+        st.warning("Detection model or confidence level changed. Please reload the model.")
 
 
     
@@ -140,8 +129,23 @@ def run_inference():
         if st.session_state['kbvqa']:
             image_qa_app(st.session_state['kbvqa'])
 
+
+        if st.button('Load Model'):
+            if st.session_state.get('kbvqa') and not model_changed and not confidence_changed:
+                st.write("Model already loaded.")
+            else:
+                st.text("Loading the model will take no more than a few minutes . .")
+                st.session_state['kbvqa'] = prepare_kbvqa_model(detection_model)
+                st.session_state['kbvqa'].detection_confidence = confidence_level
+                st.success("Model loaded with updated settings.")
+
+        if st.session_state.get('kbvqa'):
+            st.write("Model is ready for inference.")
+            image_qa_app(st.session_state['kbvqa'])
+
     else: 
         st.write('Model is not ready for inference yet')
+        # here goes the code for n-shot learning
 
             
 # Main function