Updated object detection

app.py

@@ -53,11 +53,11 @@ def initialise_object_detection_model():
     return detector
     
 # Function to get result from object detection
-def get_object_detection_results(detector, image_path):
+def get_object_detection_results(detector, image_path, object_labels):
     image = Image.open(image_path)
     predictions = detector(
         image,
-        candidate_labels=["fly", "human face", "insect", "flies"],
+        candidate_labels=object_labels,
     ) 
     draw = ImageDraw.Draw(image)
     for prediction in predictions:
@@ -208,7 +208,8 @@ if st.session_state['authentication_status']:
                 st.write('\n')
                 st.write('\n')
                 st.write("Try OWL-ViT:")
-                object_labels = st.multiselect("Enter your labels for the model to detect", options=["insect"], default="insect")
+                object_labels = st.text_input("Enter your labels for the model to detect (comma-separated)", value="insect")
+                labels = object_labels.split(", ")
                 image_file = st.file_uploader("Upload an image", type=["jpg", "png"])
                 demo_image = st.checkbox("Load in demo image")
                 if image_file:
@@ -217,14 +218,14 @@ if st.session_state['authentication_status']:
                 if image_file and object_labels:
                     st.write('After:')
                     with st.spinner("Detecting"):
-                        st.image(image = get_object_detection_results(detector, image_file))
+                        st.image(image = get_object_detection_results(detector, image_file, labels))
                 if demo_image:
                     st.write('Before:')
                     st.image("images/fly.jpg")
                 if demo_image and object_labels:
                     st.write('After:')
                     with st.spinner("Detecting"):
-                        st.image(image = get_object_detection_results(detector, "images/fly.jpg"))
+                        st.image(image = get_object_detection_results(detector, "images/fly.jpg", labels))
                 st.write('\n')
                 st.write('\n')
                 st.write("2) Behaviour Analysis - By comparing consecutive frames, the system can extract data such as the trajectory, speed, and direction of each fly's movement. Training the system on these data can improve the system's detection of flies.")