cesar commited on
Commit
0fdbd20
·
verified ·
1 Parent(s): f0b735d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +19 -34
app.py CHANGED
@@ -3,39 +3,29 @@ import gradio as gr
3
  from ultralytics import YOLO
4
 
5
  # ── Config ─────────────────────────────────────────────
6
- MODEL_PATH = "yolov8n.pt" # peso pre-entrenado (clase “person”)
7
- CONF_THRES = 0.30 # confianza mínima
8
- LINE_RATIO = 0.50 # posición de la línea virtual (50 % de la altura)
9
  # ───────────────────────────────────────────────────────
10
 
11
- # modelo — se carga una sola vez
12
  model = YOLO(MODEL_PATH)
13
 
14
- # estado global del contador
15
  memory, in_count, out_count = {}, 0, 0
16
 
17
 
18
  def count_people(frame_rgb):
19
- """
20
- ▸ Recibe un frame RGB (numpy) desde la webcam.
21
- ▸ Devuelve el frame anotado (RGB) y el texto del contador.
22
- """
23
  global memory, in_count, out_count
24
 
25
  if frame_rgb is None:
26
  return None, ""
27
 
28
- # ── convertir RGB ➜ BGR (OpenCV/YOLO) ──────────────
29
  frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
30
- h, w = frame_bgr.shape[:2]
31
- line_y = int(h * LINE_RATIO)
32
 
33
- # ── detección + tracking (ByteTrack interno) ───────
34
  results = model.track(frame_bgr,
35
- classes=[0], # solo “person”
36
- conf=CONF_THRES,
37
- persist=True,
38
- verbose=False)
39
 
40
  annotated = frame_bgr.copy()
41
  cv2.line(annotated, (0, line_y), (w, line_y), (0, 255, 255), 2)
@@ -43,14 +33,14 @@ def count_people(frame_rgb):
43
  if results:
44
  for box in results[0].boxes:
45
  x1, y1, x2, y2 = map(int, box.xyxy[0])
46
- cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
47
- tid = int(box.id[0]) if box.id is not None else -1
48
 
49
  prev_cy = memory.get(tid, cy)
50
- if prev_cy < line_y <= cy: # cruza de arriba → abajo (ENTRA)
51
- in_count += 1
52
- elif prev_cy > line_y >= cy: # cruza de abajo → arriba (SALE)
53
- out_count += 1
54
  memory[tid] = cy
55
 
56
  cv2.rectangle(annotated, (x1, y1), (x2, y2), (0, 255, 0), 1)
@@ -61,34 +51,29 @@ def count_people(frame_rgb):
61
  total = in_count - out_count
62
  label = f"In: {in_count} | Out: {out_count} | Ocupación: {total}"
63
 
64
- # ── devolver en RGB para que Gradio lo muestre bien ─
65
  annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
66
  return annotated_rgb, label
67
 
68
 
69
  def reset_counts():
70
- """Resetea contador y memoria."""
71
  global memory, in_count, out_count
72
  memory, in_count, out_count = {}, 0, 0
73
  return None, ""
74
 
75
 
76
- # ── Interfaz Gradio ────────────────────────────────────
77
  with gr.Blocks(title="Contador de personas (entrada única)") as demo:
78
  gr.Markdown("# Contador de personas (entrada única)")
79
 
80
  with gr.Row():
81
- cam = gr.Image(sources=["webcam"], streaming=True, label="frame")
82
- out_img = gr.Image(label="Video")
83
  out_lbl = gr.Text(label="Contador")
84
 
85
  btn_clear = gr.Button("Limpiar")
86
 
87
- # conectar streaming y botón
88
- cam.stream(fn=count_people, inputs=[cam], outputs=[out_img, out_lbl])
89
- btn_clear.click(fn=reset_counts, outputs=[out_img, out_lbl])
90
 
91
- # ── Cola y lanzamiento ─────────────────────────────────
92
- demo.queue(concurrency_count=2, max_size=60) # 2 hilos, cola máx 60 frames
93
  demo.launch()
94
-
 
3
  from ultralytics import YOLO
4
 
5
  # ── Config ─────────────────────────────────────────────
6
+ MODEL_PATH = "yolov8n.pt"
7
+ CONF_THRES = 0.30
8
+ LINE_RATIO = 0.50
9
  # ───────────────────────────────────────────────────────
10
 
 
11
  model = YOLO(MODEL_PATH)
12
 
 
13
  memory, in_count, out_count = {}, 0, 0
14
 
15
 
16
  def count_people(frame_rgb):
 
 
 
 
17
  global memory, in_count, out_count
18
 
19
  if frame_rgb is None:
20
  return None, ""
21
 
 
22
  frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
23
+ h, w = frame_bgr.shape[:2]
24
+ line_y = int(h * LINE_RATIO)
25
 
 
26
  results = model.track(frame_bgr,
27
+ classes=[0], conf=CONF_THRES,
28
+ persist=True, verbose=False)
 
 
29
 
30
  annotated = frame_bgr.copy()
31
  cv2.line(annotated, (0, line_y), (w, line_y), (0, 255, 255), 2)
 
33
  if results:
34
  for box in results[0].boxes:
35
  x1, y1, x2, y2 = map(int, box.xyxy[0])
36
+ cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
37
+ tid = int(box.id[0]) if box.id is not None else -1
38
 
39
  prev_cy = memory.get(tid, cy)
40
+ if prev_cy < line_y <= cy:
41
+ in_count += 1 # entró
42
+ elif prev_cy > line_y >= cy:
43
+ out_count += 1 # salió
44
  memory[tid] = cy
45
 
46
  cv2.rectangle(annotated, (x1, y1), (x2, y2), (0, 255, 0), 1)
 
51
  total = in_count - out_count
52
  label = f"In: {in_count} | Out: {out_count} | Ocupación: {total}"
53
 
 
54
  annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
55
  return annotated_rgb, label
56
 
57
 
58
  def reset_counts():
 
59
  global memory, in_count, out_count
60
  memory, in_count, out_count = {}, 0, 0
61
  return None, ""
62
 
63
 
 
64
  with gr.Blocks(title="Contador de personas (entrada única)") as demo:
65
  gr.Markdown("# Contador de personas (entrada única)")
66
 
67
  with gr.Row():
68
+ cam = gr.Image(sources=["webcam"], streaming=True, label="frame")
69
+ out_img = gr.Image(label="Video")
70
  out_lbl = gr.Text(label="Contador")
71
 
72
  btn_clear = gr.Button("Limpiar")
73
 
74
+ cam.stream(count_people, inputs=[cam], outputs=[out_img, out_lbl])
75
+ btn_clear.click(reset_counts, outputs=[out_img, out_lbl])
 
76
 
77
+ # límite de cola (para que no se llene si baja el FPS)
78
+ demo.queue(max_size=60)
79
  demo.launch()