tracking / app.py
fcernafukuzaki's picture
Update app.py
fa0f406 verified
import streamlit as st
from streamlit.runtime.scriptrunner.script_run_context import add_script_run_ctx
import threading
import cv2
from ultralytics import YOLO
model = YOLO('best.pt')
model.fuse()
# Configurar cámaras
camaras = {
'camara1': 'IMG_0559.mov', #0,
'camara2': 'IMG_0559.mov', #0,
}
def tracking(source, iou, conf):
result = model.track(source=source, iou=iou, conf=conf,
persist=True,
device='mps',
tracker='botsort.yaml',
verbose=False)
return result
def inferir_camara(index, source, iou, conf, st_frame, st_cantidad, stop_button_pressed):
camera1 = cv2.VideoCapture(source) # Open camera 0
while camera1.isOpened() and not stop_button_pressed:
ret, frame = camera1.read()
if not ret:
st.write(f"Video {str(index)} Capture Ended")
break
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
resultado = tracking(frame, iou, conf)
# Extract bounding boxes (replace this with the specific method from your library)
objects_detected = resultado[0].boxes.cls # Boxes object for bbox outputs
object_count = len(objects_detected)
st_cantidad.write(f'Cantidad detectados en cámara {str(index)}: {object_count}')
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:
if stop_button_pressed:
break
# Release cameras
#camera1.release()
#cv2.destroyAllWindows()
# Interfaz
st.set_page_config(page_title="Tracking YOLOv8")
st.title("Tracking YOLOv8")
iou = float(st.sidebar.slider("NMS IoU threshold", 30, 100, 80)) / 100 # Umbral de intersección sobre unión (IoU) para NMS
conf = float(st.sidebar.slider("Umbral o threshold", 30, 100, 80)) / 100 # Select Model Confidence
col1, col2 = st.columns(2)
play_button_pressed = st.empty()
with col1:
stop_button_pressed_1 = st.button("Detener cámara 1")
play_button_1 = st.empty()
st_frame_camara_1 = st.empty()
st_cantidad_1 = st.empty()
tracker_thread1 = threading.Thread(target=inferir_camara,
args=("1", camaras.get('camara1', 0),
iou, conf,
st_frame_camara_1, st_cantidad_1,
stop_button_pressed_1),
daemon=True)
add_script_run_ctx(tracker_thread1)
play_button_1 = st.button("Iniciar cámara 1", on_click=tracker_thread1.start())
with col2:
stop_button_pressed_2 = st.button("Detener cámara 2")
st_frame_camara_2 = st.empty()
st_cantidad_2 = st.empty()
tracker_thread2 = threading.Thread(target=inferir_camara,
args=("2", camaras.get('camara2', 0),
iou, conf,
st_frame_camara_2, st_cantidad_2,
stop_button_pressed_2),
daemon=True)
add_script_run_ctx(tracker_thread2)
play_button_2 = st.button("Iniciar cámara 2", on_click=tracker_thread2.start())
# Wait for the tracker threads to finish
tracker_thread1.join()
tracker_thread2.join()