/ main.py
main.py
 1  import cv2
 2  import cvzone
 3  from collections import deque
 4  
 5  thres = 0.55  # Umbral de confianza
 6  nmsThres = 0.2  # Umbral de NMS (Non-Maxima Suppression)
 7  
 8  cap = cv2.VideoCapture('v1.mp4')
 9  cap.set(3, 640)
10  cap.set(4, 480)
11  
12  # Carga los nombres de las clases
13  classNames = []
14  classFile = 'coco.names'
15  with open(classFile, 'rt') as f:
16      classNames = f.read().strip().split('\n')
17  
18  # Cargar la configuración y los pesos de la red
19  configPath = 'ssd_mobilenet_v3_large_coco_2020_01_14.pbtxt'
20  weightsPath = "frozen_inference_graph.pb"
21  
22  net = cv2.dnn_DetectionModel(weightsPath, configPath)
23  net.setInputSize(320, 320)
24  net.setInputScale(1.0 / 127.5)
25  net.setInputMean((127.5, 127.5, 127.5))
26  net.setInputSwapRB(True)
27  
28  # Inicializa el contador de personas y una lista de seguimiento
29  person_count = 0
30  trackers = []  # Lista de trackers activos
31  max_distance = 50  # Distancia máxima entre dos detecciones para considerar que es la misma persona
32  detections = deque(maxlen=20)  # Cola para almacenar las últimas detecciones de personas
33  
34  while True:
35      success, img = cap.read()
36      if not success:
37          break  # Sale del bucle si no se pueden leer más frames
38  
39      # Detección de objetos
40      classIds, confs, bbox = net.detect(img, confThreshold=thres, nmsThreshold=nmsThres)
41  
42      # Verifica si la detección no está vacía
43      if len(classIds) != 0:
44          for classId, conf, box in zip(classIds.flatten(), confs.flatten(), bbox):
45              if 0 < classId <= len(classNames):
46                  if classNames[classId - 1].lower() == "person":
47                      # Verifica si la persona ya fue contada en los últimos frames
48                      x, y, w, h = box
49                      detected = False
50  
51                      for prev_box in detections:
52                          prev_x, prev_y, prev_w, prev_h = prev_box
53                          distance = ((x - prev_x) ** 2 + (y - prev_y) ** 2) ** 0.5
54                          if distance < max_distance:
55                              detected = True
56                              break
57                      
58                      if not detected:
59                          person_count += 1  # Incrementa el contador solo si no ha sido detectada recientemente
60                          detections.append(box)  # Agrega la nueva detección a la cola
61  
62                      # Dibuja el cuadro alrededor de la persona
63                      cvzone.cornerRect(img, box)
64                      cv2.putText(img, f'PERSON {round(conf * 100, 2)}%',
65                                  (box[0] + 10, box[1] + 30), cv2.FONT_HERSHEY_COMPLEX_SMALL,
66                                  1, (0, 255, 0), 2)
67  
68      cv2.imshow("Image", img)
69      cv2.waitKey(1)
70  
71  # Al finalizar el video, muestra el conteo de personas
72  print(f'Cantidad total de personas detectadas: {person_count}')