Files
biblioteca_conocimiento_lab…/vigilancia/enroll.py

122 lines
3.7 KiB
Python

import os
import sys
import cv2
import numpy as np
import reconocedor as rec
CONOCIDOS = rec.CONOCIDOS
WEB_CAM = 0
def escribir(texto, frame):
cv2.putText(frame, texto, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
def capturar_cara(r, frame):
caras = r.detectar_caras(frame)
if not caras:
return None, None
mejor = max(caras, key=lambda c: c[-1])
emb = r.embedding_cara(frame, mejor)
return emb, mejor
def capturar_cuerpo(r, frame):
personas = r._detectar_personas(frame)
if not personas:
return None, None
mayor = max(personas, key=lambda p: p[2] * p[3])
x, y, w, h = mayor
recorte = frame[y:y + h, x:x + w]
emb = r.embedding_cuerpo(recorte)
return emb, mayor
def main():
nombres = sys.argv[1:3]
while len(nombres) < 2:
nombres.append(f"Hija{len(nombres) + 1}")
print(f"Registrando: {nombres}")
r = rec.cargar_reconocedor()
os.makedirs(CONOCIDOS, exist_ok=True)
cap = cv2.VideoCapture(WEB_CAM)
if not cap.isOpened():
print("ERROR: no se pudo abrir la webcam")
sys.exit(1)
cv2.namedWindow("Registro", cv2.WINDOW_NORMAL)
for nombre in nombres:
carpeta = os.path.join(CONOCIDOS, nombre)
os.makedirs(carpeta, exist_ok=True)
print(f"\n>>> Hija: {nombre}")
fase = "cara"
while True:
ok, frame = cap.read()
if not ok:
continue
caras = r.detectar_caras(frame)
personas = r._detectar_personas(frame)
for c in caras:
x, y, w, h = (int(v) for v in c[:4])
cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
for p in personas:
x, y, w, h = p
cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 2)
if fase == "cara":
escribir("CARA: pon la cara mirando a la camara y pulsa ESPACIO", frame)
else:
escribir("CUERPO: alejate (cuerpo entero) y pulsa B", frame)
escribir(f"Persona: {nombre}", frame)
cv2.imshow("Registro", frame)
tecla = cv2.waitKey(1) & 0xFF
if tecla == ord("q"):
cap.release()
cv2.destroyAllWindows()
print("Abortado por el usuario")
sys.exit(0)
if fase == "cara" and tecla == ord(" "):
emb, caja = capturar_cara(r, frame)
if emb is None:
print(" No se vio ninguna cara clara. Reintenta.")
continue
np.save(os.path.join(carpeta, "face.npy"), emb)
x, y, w, h = (int(v) for v in caja[:4])
cv2.imwrite(os.path.join(carpeta, "cara_ref.jpg"),
frame[max(0, y - 20):y + h + 20, max(0, x - 20):x + w + 20])
print(f" Cara de {nombre} registrada (score {caja[-1]:.2f})")
fase = "cuerpo"
elif fase == "cuerpo" and tecla == ord("b"):
emb, caja = capturar_cuerpo(r, frame)
if emb is None:
print(" No se detecto ninguna persona. Reintenta.")
continue
np.save(os.path.join(carpeta, "body.npy"), emb)
x, y, w, h = caja
cv2.imwrite(os.path.join(carpeta, "cuerpo_ref.jpg"),
frame[max(0, y):y + h, max(0, x):x + w])
print(f" Cuerpo de {nombre} registrado")
break
cap.release()
cv2.destroyAllWindows()
print("\nRegistro completado. Conocidos guardados:")
for n in nombres:
print(f" - {n} (face.npy + body.npy)")
if __name__ == "__main__":
main()