Update app.py

app.py

@@ -1,24 +1,37 @@
 import streamlit as st
 import cv2
+from ultralytics import YOLO
 
-def main():
-    st.set_page_config(page_title="Streamlit WebCam App")
-    st.title("Webcam Display Steamlit App")
-    st.caption("Powered by OpenCV, Streamlit")
-    cap = cv2.VideoCapture(0)
-    frame_placeholder = st.empty()
-    stop_button_pressed = st.button("Stop")
-    while cap.isOpened() and not stop_button_pressed:
-        ret, frame = cap.read()
-        if not ret:
-            st.write("Video Capture Ended")
-            break
-        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        frame_placeholder.image(frame,channels="RGB")
-        if cv2.waitKey(1) & 0xFF == ord("q") or stop_button_pressed:
-            break
-    cap.release()
-    cv2.destroyAllWindows()
+model = YOLO('best.pt')
 
-if __name__ == "__main__":
-    main()
+def track(source):
+    result = model.track(source=source)
+    return result
+
+st.set_page_config(page_title="Streamlit WebCam App")
+st.title("Webcam Display Steamlit App")
+cap = cv2.VideoCapture(0)
+frame_placeholder = st.empty()
+st_frame = st.empty()
+stop_button_pressed = st.button("Stop")
+while cap.isOpened() and not stop_button_pressed:
+    ret, frame = cap.read()
+    if not ret:
+        st.write("Video Capture Ended")
+        break
+    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    frame_placeholder.image(frame, channels="RGB")
+
+    resultado = track(frame)
+    res_plotted = resultado[0].plot()
+    
+    st_frame.image(res_plotted,
+                       caption='Detected Video',
+                       channels="RGB",
+                       #use_column_width=True
+                       )
+
+    if cv2.waitKey(1) & 0xFF == ord("q") or stop_button_pressed:
+        break
+cap.release()
+cv2.destroyAllWindows()