Spaces:
Running
Running
# app.py | |
import cv2 | |
import gradio as gr | |
from ultralytics import YOLO | |
# ── Config ───────────────────────────────────────────── | |
MODEL_PATH = "yolov8n.pt" # modelo pre-entrenado (COCO, clase “person”) | |
CONF_THRES = 0.3 # confianza mínima | |
LINE_RATIO = 0.5 # línea virtual en mitad de la altura | |
# ─────────────────────────────────────────────────────── | |
model = YOLO(MODEL_PATH) | |
# Estado global para entradas/salidas | |
memory = {} # {track_id: previous_cy} | |
in_count = 0 | |
out_count = 0 | |
def count_people(frame): | |
global memory, in_count, out_count | |
h, w = frame.shape[:2] | |
line_y = int(h * LINE_RATIO) | |
# detección + tracking interno (ByteTrack) | |
results = model.track( | |
frame, | |
classes=[0], # solo “person” | |
conf=CONF_THRES, | |
persist=True, | |
verbose=False | |
) | |
annotated = frame.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 | |
# lógica de cruce | |
prev_cy = memory.get(tid, cy) | |
if prev_cy < line_y <= cy: # entra | |
in_count += 1 | |
elif prev_cy > line_y >= cy: # sale | |
out_count += 1 | |
memory[tid] = cy | |
# dibujo | |
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 | |
cv2.putText( | |
annotated, | |
f"In: {in_count} Out: {out_count} Ocupación: {total}", | |
(10, 30), | |
cv2.FONT_HERSHEY_SIMPLEX, | |
1, | |
(255, 255, 255), | |
2 | |
) | |
return annotated, f"In: {in_count} | Out: {out_count} | Ocupación: {total}" | |
demo = gr.Interface( | |
fn=count_people, | |
inputs=gr.Image(sources=["webcam"], streaming=True), # 👈 parámetro correcto | |
outputs=[gr.Image(label="Video"), gr.Text(label="Contador")], | |
live=True, | |
title="Contador de personas (entrada única)" | |
) | |
if __name__ == "__main__": | |
demo.launch() | |