Spaces:
Sleeping
Sleeping
| import cv2 | |
| import gradio as gr | |
| from ultralytics import YOLO | |
| # ── Config ───────────────────────────────────────────── | |
| MODEL_PATH = "yolov8n.pt" # peso pre-entrenado (clase “person”) | |
| CONF_THRES = 0.30 # confianza mínima | |
| LINE_RATIO = 0.50 # posición de la línea virtual (50 % de la altura) | |
| # ─────────────────────────────────────────────────────── | |
| # modelo — se carga una sola vez | |
| model = YOLO(MODEL_PATH) | |
| # estado global del contador | |
| memory, in_count, out_count = {}, 0, 0 | |
| def count_people(frame_rgb): | |
| """ | |
| ▸ Recibe un frame RGB (numpy) desde la webcam. | |
| ▸ Devuelve el frame anotado (RGB) y el texto del contador. | |
| """ | |
| global memory, in_count, out_count | |
| if frame_rgb is None: | |
| return None, "" | |
| # ── convertir RGB ➜ BGR (OpenCV/YOLO) ────────────── | |
| frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR) | |
| h, w = frame_bgr.shape[:2] | |
| line_y = int(h * LINE_RATIO) | |
| # ── detección + tracking (ByteTrack interno) ─────── | |
| 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 = (x1 + x2) // 2, (y1 + y2) // 2 | |
| tid = int(box.id[0]) if box.id is not None else -1 | |
| prev_cy = memory.get(tid, cy) | |
| if prev_cy < line_y <= cy: # cruza de arriba → abajo (ENTRA) | |
| in_count += 1 | |
| elif prev_cy > line_y >= cy: # cruza de abajo → arriba (SALE) | |
| out_count += 1 | |
| memory[tid] = cy | |
| 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}" | |
| # ── devolver en RGB para que Gradio lo muestre bien ─ | |
| annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB) | |
| return annotated_rgb, label | |
| def reset_counts(): | |
| """Resetea contador y memoria.""" | |
| global memory, in_count, out_count | |
| memory, in_count, out_count = {}, 0, 0 | |
| return None, "" | |
| # ── Interfaz Gradio ──────────────────────────────────── | |
| 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") | |
| # conectar streaming y botón | |
| cam.stream(fn=count_people, inputs=[cam], outputs=[out_img, out_lbl]) | |
| btn_clear.click(fn=reset_counts, outputs=[out_img, out_lbl]) | |
| # ── Cola y lanzamiento ───────────────────────────────── | |
| demo.queue(concurrency_count=2, max_size=60) # 2 hilos, cola máx 60 frames | |
| demo.launch() | |