app updated

app.py

@@ -27,12 +27,6 @@ api_key = st.secrets["PALM_API_KEY"]
 palm.configure(api_key=api_key)
 
 
-# Load YOLO pipeline
-yolo_pipe = pipeline(
-    "object-detection", model="eagle0504/detr-finetuned-balloon-v2"
-)  # "hustvl/yolos-small"
-
-
 # Main function of the Streamlit app
 def main():
     st.title("Generative AI Demo on Camera Input/Image/PDF 💻")
@@ -151,6 +145,23 @@ def main():
     if image is not None:
         st.sidebar.success("Check the following box to run YOLO algorithm if desired!")
         use_yolo = st.sidebar.checkbox("Use YOLO!", value=False)
+        if use_yolo:
+            yolo_option = st.selectbox(
+                "Which YOLO algorithm would you like?",
+                ("hustvl/yolos-small", "eagle0504/detr-finetuned-balloon-v2"),
+            )
+        else:
+            yolo_option = None
+
+        # Load YOLO pipeline
+        if yolo_option == "hustvl/yolos-small":
+            yolo_pipe = pipeline("object-detection", model="hustvl/yolos-small")
+        elif yolo_option == "eagle0504/detr-finetuned-balloon-v2":
+            yolo_pipe = pipeline(
+                "object-detection", model="eagle0504/detr-finetuned-balloon-v2"
+            )
+        else:
+            yolo_pipe = pipeline("object-detection", model="hustvl/yolos-small")
 
         if use_yolo:
             # Process image with YOLO