diff --git a/vision_agent/agent/vision_agent_prompts.py b/vision_agent/agent/vision_agent_prompts.py
index 7e06bfb2..66d6614e 100644
--- a/vision_agent/agent/vision_agent_prompts.py
+++ b/vision_agent/agent/vision_agent_prompts.py
@@ -54,7 +54,7 @@
false
OBSERVATION:
-[Artifact dog_detector.py]
+[Artifact dog_detector.py (5 lines total)]
0|from vision_agent.tools import load_image, owl_v2
1|def detect_dogs(image_path: str):
2| image = load_image(image_path)
@@ -95,7 +95,7 @@
false
OBSERVATION:
-[Artifact dog_detector.py]
+[Artifact dog_detector.py (5 lines total)]
0|from vision_agent.tools import load_image, owl_v2
1|def detect_dogs(image_path: str):
2| image = load_image(image_path)
@@ -167,7 +167,7 @@
false
OBSERVATION:
-[Artifact code.py]
+[Artifact code.py (17 lines total)]
0|from vision_agent.tools import load_image, florence2_phrase_grounding, closest_box_distance, overlay_bounding_boxes, save_image
1|def count_workers_with_helmets(image_path: str, output_path: str):
2| image = load_image(image_path)
@@ -202,7 +202,9 @@
AGENT: Two workers with helmets are detected, I will show this to the user and ask them if the result looks good.
The code to detect workers with helmets is saved in code.py and the visualization under 'workers_viz.png'.
true
+"""
+EXAMPLES_CODE3_EXTRA1 = """
USER: The detections are slightly off. Can you fine tune florence2 using these labels? "[{'image_path': 'image1.png': 'labels': ['worker', 'helmet'], 'bboxes': [[235, 118, 294, 241], [232, 118, 294, 128]]}, {'image_path': 'image2.png': 'labels': ['worker', 'helmet'], 'bboxes': [[235, 118, 294, 241], [232, 118, 294, 128]]}]"
AGENT: Because the user has supplied me with labels I can call object_detection_fine_tuning on their behalf to fine tune the model
@@ -250,3 +252,60 @@ def count_workers_with_helmets(image_path: str, output_path: str):
I have updated the code using the fine tuned florence2_phrase_grounding model to count the workers wearing helmets in code.py and saved the visualization under 'workers_viz.png'.
true
"""
+
+EXAMPLES_CODE3_EXTRA2 = """
+USER: Can you write a streamlit app that allows me to upload an image and count the workers with helmets?
+
+AGENT: I will use edit_code_artifact to generate a streamlit app that allows the user to upload an image and count the workers with helmets.
+
+CODE = '''
+import streamlit as st
+from code imoprt count_workers_with_helmets
+import os
+
+def main():
+ st.title("Worker with Helmet Counter")
+ uploaded_file = st.file_uploader("Choose image...", type=["jpeg", "jpg", "png"])
+ if uploaded_file is not None:
+ with open("temp_image.jpg", "wb") as f:
+ f.write(uploaded_file.read())
+
+ count = count_workers_with_helmets("temp_image.jpg", "temp_image_viz.jpg")
+ st.image("temp_image_viz.jpg")
+ st.write(f"Number of workers with helmets: {count}")
+ os.remove("temp_image.jpg")
+ os.remove("temp_image_viz.jpg")
+
+if __name__ == "__main__":
+ main()
+'''
+edit_code_artifact(artifacts, 'streamlit_app.py', CODE, 0, 0)
+
+
+OBSERVATION:
+[Artifact streamlit_app.py (19 lines total)]
+0|import streamlit as st
+1|from code imoprt count_workers_with_helmets
+2|import os
+3|
+4|def main():
+5| st.title("Worker with Helmet Counter")
+6| uploaded_file = st.file_uploader("Choose image...", type=["jpeg", "jpg", "png"])
+7| if uploaded_file is not None:
+8| with open("temp_image.jpg", "wb") as f:
+9| f.write(uploaded_file.read())
+10|
+11| count = count_workers_with_helmets("temp_image.jpg", "temp_image_viz.jpg")
+12| st.image("temp_image_viz.jpg")
+13| st.write(f"Number of workers with helmets: {count}")
+14| os.remove("temp_image.jpg")
+15| os.remove("temp_image_viz.jpg")
+16|
+17|if __name__ == "__main__":
+18| main()
+[End of artifact]
+
+AGENT: I have generated the streamlit app that allows the user to upload an image and count the workers
+The streamlit app is saved in streamlit_app.py, you can run the app by running `streamlit run streamlit_app.py`.
+true
+"""