cesar commited on
Commit
982fb7b
·
verified ·
1 Parent(s): c66d0d6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +20 -36
app.py CHANGED
@@ -3,42 +3,32 @@ import gradio as gr
3
  from ultralytics import YOLO
4
 
5
  # ── Config ─────────────────────────────────────────────
6
- MODEL_PATH = "yolov8n.pt" # modelo pre-entrenado (clase “person”)
7
- CONF_THRES = 0.3 # confianza mínima detección
8
- LINE_RATIO = 0.5 # posición de la línea virtual (50 % altura)
9
  # ───────────────────────────────────────────────────────
10
 
 
11
  model = YOLO(MODEL_PATH)
12
 
13
  # estado global
14
- memory = {} # {track_id: previous_cy}
15
- in_count = 0
16
- out_count = 0
17
 
18
 
19
- def count_people(frame):
20
- """
21
- Recibe un frame RGB (numpy) -> procesa -> devuelve frame RGB anotado y string.
22
- Se llama de forma continua porque el input tiene `streaming=True`.
23
- """
24
  global memory, in_count, out_count
25
 
26
- if frame is None:
27
  return None, ""
28
 
29
- # ── paso 1: RGB ➜ BGR para OpenCV/YOLO ─────────────
30
- frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
31
  h, w = frame_bgr.shape[:2]
32
  line_y = int(h * LINE_RATIO)
33
 
34
- # ── detección + tracking ───────────────────────────
35
- results = model.track(
36
- frame_bgr,
37
- classes=[0], # solo “person”
38
- conf=CONF_THRES,
39
- persist=True,
40
- verbose=False
41
- )
42
 
43
  annotated = frame_bgr.copy()
44
  cv2.line(annotated, (0, line_y), (w, line_y), (0, 255, 255), 2)
@@ -46,18 +36,16 @@ def count_people(frame):
46
  if results:
47
  for box in results[0].boxes:
48
  x1, y1, x2, y2 = map(int, box.xyxy[0])
49
- cx, cy = int((x1 + x2) / 2), int((y1 + y2) / 2)
50
  tid = int(box.id[0]) if box.id is not None else -1
51
 
52
- # cruces entrada / salida
53
  prev_cy = memory.get(tid, cy)
54
- if prev_cy < line_y <= cy: # cruzó de arriba → abajo (ENTRA)
55
  in_count += 1
56
- elif prev_cy > line_y >= cy: # abajo → arriba (SALE)
57
  out_count += 1
58
  memory[tid] = cy
59
 
60
- # dibujitos
61
  cv2.rectangle(annotated, (x1, y1), (x2, y2), (0, 255, 0), 1)
62
  cv2.circle(annotated, (cx, cy), 3, (0, 0, 255), -1)
63
  cv2.putText(annotated, str(tid), (x1, y1 - 5),
@@ -66,17 +54,13 @@ def count_people(frame):
66
  total = in_count - out_count
67
  label = f"In: {in_count} | Out: {out_count} | Ocupación: {total}"
68
 
69
- # ── paso 2: BGR ➜ RGB para mostrar en Gradio ───────
70
  annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
71
  return annotated_rgb, label
72
 
73
 
74
  def reset_counts():
75
- """Callback para botón ‘Limpiar’."""
76
  global memory, in_count, out_count
77
- memory = {}
78
- in_count = 0
79
- out_count = 0
80
  return None, ""
81
 
82
 
@@ -90,9 +74,9 @@ with gr.Blocks(title="Contador de personas (entrada única)") as demo:
90
 
91
  btn_clear = gr.Button("Limpiar")
92
 
93
- # wire-up
94
- cam.stream(fn=count_people, outputs=[out_img, out_lbl])
95
- btn_clear.click(fn=reset_counts, outputs=[out_img, out_lbl])
96
 
97
  if __name__ == "__main__":
98
- demo.launch()
 
3
  from ultralytics import YOLO
4
 
5
  # ── Config ─────────────────────────────────────────────
6
+ MODEL_PATH = "yolov8n.pt" # peso pre-entrenado (clase “person”)
7
+ CONF_THRES = 0.3
8
+ LINE_RATIO = 0.5 # línea virtual al 50 % de la altura
9
  # ───────────────────────────────────────────────────────
10
 
11
+ # carga una sola vez
12
  model = YOLO(MODEL_PATH)
13
 
14
  # estado global
15
+ memory, in_count, out_count = {}, 0, 0
 
 
16
 
17
 
18
+ def count_people(frame_rgb):
19
+ """Procesa un frame RGB -> devuelve frame RGB anotado + texto contador."""
 
 
 
20
  global memory, in_count, out_count
21
 
22
+ if frame_rgb is None:
23
  return None, ""
24
 
25
+ frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
 
26
  h, w = frame_bgr.shape[:2]
27
  line_y = int(h * LINE_RATIO)
28
 
29
+ # detección + tracking
30
+ results = model.track(frame_bgr, classes=[0], conf=CONF_THRES,
31
+ persist=True, verbose=False)
 
 
 
 
 
32
 
33
  annotated = frame_bgr.copy()
34
  cv2.line(annotated, (0, line_y), (w, line_y), (0, 255, 255), 2)
 
36
  if results:
37
  for box in results[0].boxes:
38
  x1, y1, x2, y2 = map(int, box.xyxy[0])
39
+ cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
40
  tid = int(box.id[0]) if box.id is not None else -1
41
 
 
42
  prev_cy = memory.get(tid, cy)
43
+ if prev_cy < line_y <= cy: # entra
44
  in_count += 1
45
+ elif prev_cy > line_y >= cy: # sale
46
  out_count += 1
47
  memory[tid] = cy
48
 
 
49
  cv2.rectangle(annotated, (x1, y1), (x2, y2), (0, 255, 0), 1)
50
  cv2.circle(annotated, (cx, cy), 3, (0, 0, 255), -1)
51
  cv2.putText(annotated, str(tid), (x1, y1 - 5),
 
54
  total = in_count - out_count
55
  label = f"In: {in_count} | Out: {out_count} | Ocupación: {total}"
56
 
 
57
  annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
58
  return annotated_rgb, label
59
 
60
 
61
  def reset_counts():
 
62
  global memory, in_count, out_count
63
+ memory, in_count, out_count = {}, 0, 0
 
 
64
  return None, ""
65
 
66
 
 
74
 
75
  btn_clear = gr.Button("Limpiar")
76
 
77
+ # 💡 la clave: inputs=[cam]
78
+ cam.stream(count_people, inputs=[cam], outputs=[out_img, out_lbl])
79
+ btn_clear.click(reset_counts, outputs=[out_img, out_lbl])
80
 
81
  if __name__ == "__main__":
82
+ demo.launch(max_threads=2, max_queue=60) # cola chica para que no se tape