/ 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}')