Spaces:
Sleeping
Sleeping
| import cv2 | |
| import gradio as gr | |
| from ultralytics import YOLO | |
| # ── Config ───────────────────────────────────────────── | |
| MODEL_PATH = "yolov8n.pt" # modelo pre-entrenado (clase “person”) | |
| CONF_THRES = 0.3 # confianza mínima detección | |
| LINE_RATIO = 0.5 # posición de la línea virtual (50 % altura) | |
| # ─────────────────────────────────────────────────────── | |
| model = YOLO(MODEL_PATH) | |
| # estado global | |
| memory = {} # {track_id: previous_cy} | |
| in_count = 0 | |
| out_count = 0 | |
| def count_people(frame): | |
| """ | |
| Recibe un frame RGB (numpy) -> procesa -> devuelve frame RGB anotado y string. | |
| Se llama de forma continua porque el input tiene `streaming=True`. | |
| """ | |
| global memory, in_count, out_count | |
| if frame is None: | |
| return None, "" | |
| # ── paso 1: RGB ➜ BGR para OpenCV/YOLO ───────────── | |
| frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) | |
| h, w = frame_bgr.shape[:2] | |
| line_y = int(h * LINE_RATIO) | |
| # ── detección + tracking ─────────────────────────── | |
| results = model.track( | |
| frame_bgr, | |
| classes=[0], # solo “person” | |
| conf=CONF_THRES, | |
| persist=True, | |
| verbose=False | |
| ) | |
| annotated = frame_bgr.copy() | |
| cv2.line(annotated, (0, line_y), (w, line_y), (0, 255, 255), 2) | |
| if results: | |
| for box in results[0].boxes: | |
| x1, y1, x2, y2 = map(int, box.xyxy[0]) | |
| cx, cy = int((x1 + x2) / 2), int((y1 + y2) / 2) | |
| tid = int(box.id[0]) if box.id is not None else -1 | |
| # cruces entrada / salida | |
| prev_cy = memory.get(tid, cy) | |
| if prev_cy < line_y <= cy: # cruzó de arriba → abajo (ENTRA) | |
| in_count += 1 | |
| elif prev_cy > line_y >= cy: # abajo → arriba (SALE) | |
| out_count += 1 | |
| memory[tid] = cy | |
| # dibujitos | |
| cv2.rectangle(annotated, (x1, y1), (x2, y2), (0, 255, 0), 1) | |
| cv2.circle(annotated, (cx, cy), 3, (0, 0, 255), -1) | |
| cv2.putText(annotated, str(tid), (x1, y1 - 5), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1) | |
| total = in_count - out_count | |
| label = f"In: {in_count} | Out: {out_count} | Ocupación: {total}" | |
| # ── paso 2: BGR ➜ RGB para mostrar en Gradio ─────── | |
| annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB) | |
| return annotated_rgb, label | |
| def reset_counts(): | |
| """Callback para botón ‘Limpiar’.""" | |
| global memory, in_count, out_count | |
| memory = {} | |
| in_count = 0 | |
| out_count = 0 | |
| return None, "" | |
| with gr.Blocks(title="Contador de personas (entrada única)") as demo: | |
| gr.Markdown("# Contador de personas (entrada única)") | |
| with gr.Row(): | |
| cam = gr.Image(sources=["webcam"], streaming=True, label="frame") | |
| out_img = gr.Image(label="Video") | |
| out_lbl = gr.Text(label="Contador") | |
| btn_clear = gr.Button("Limpiar") | |
| # wire-up | |
| cam.stream(fn=count_people, outputs=[out_img, out_lbl]) | |
| btn_clear.click(fn=reset_counts, outputs=[out_img, out_lbl]) | |
| if __name__ == "__main__": | |
| demo.launch() | |