214 lines
7.5 KiB
Python
214 lines
7.5 KiB
Python
import os
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
import torchvision.models as models
|
|
from torchvision import transforms
|
|
|
|
BASE = os.path.dirname(os.path.abspath(__file__))
|
|
MODELOS = os.path.join(BASE, "modelos")
|
|
CONOCIDOS = os.path.join(BASE, "conocidos")
|
|
YUNET_PATH = os.path.join(MODELOS, "face_detection_yunet.onnx")
|
|
ARCFACE_PATH = os.path.join(MODELOS, "buffalo_s", "w600k_mbf.onnx")
|
|
YOLO_PATH = os.path.join(BASE, "yolov8n.pt")
|
|
|
|
UMBRAL_CARA = 0.40
|
|
UMBRAL_CUERPO = 0.55
|
|
|
|
|
|
class Reconocedor:
|
|
def __init__(self):
|
|
self.yunet = cv2.FaceDetectorYN_create(YUNET_PATH, "", (320, 320))
|
|
self.yunet.setScoreThreshold(0.6)
|
|
self.arcface = cv2.dnn.readNetFromONNX(ARCFACE_PATH)
|
|
self.body = self._cargar_body()
|
|
self.tf_body = transforms.Compose([
|
|
transforms.ToPILImage(),
|
|
transforms.Resize((224, 224)),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
|
|
])
|
|
self.conocidos = self._cargar_conocidos()
|
|
|
|
def _cargar_body(self):
|
|
m = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
|
|
m.fc = torch.nn.Identity()
|
|
m.eval()
|
|
return m
|
|
|
|
def _cargar_conocidos(self):
|
|
conocidos = {}
|
|
if not os.path.isdir(CONOCIDOS):
|
|
return conocidos
|
|
for nombre in os.listdir(CONOCIDOS):
|
|
carpeta = os.path.join(CONOCIDOS, nombre)
|
|
if not os.path.isdir(carpeta):
|
|
continue
|
|
emb_cara = self._cargar_npy(carpeta, "face")
|
|
emb_cuerpo = self._cargar_npy(carpeta, "body")
|
|
if emb_cara is not None or emb_cuerpo is not None:
|
|
conocidos[nombre] = (emb_cara, emb_cuerpo)
|
|
return conocidos
|
|
|
|
def _cargar_npy(self, carpeta, tipo):
|
|
ruta = os.path.join(carpeta, f"{tipo}.npy")
|
|
if os.path.exists(ruta):
|
|
return np.load(ruta)
|
|
return None
|
|
|
|
# ------------------------------------------------------------------
|
|
# Cara
|
|
# ------------------------------------------------------------------
|
|
def detectar_caras(self, frame):
|
|
h, w = frame.shape[:2]
|
|
self.yunet.setInputSize((w, h))
|
|
_, caras = self.yunet.detect(frame)
|
|
if caras is None:
|
|
return []
|
|
return [c for c in caras if c[-1] >= 0.6]
|
|
|
|
def embedding_cara(self, frame, caja):
|
|
alineada = self._alinear_cara(frame, caja)
|
|
if alineada is None:
|
|
return None
|
|
blob = cv2.dnn.blobFromImage(
|
|
alineada, 1.0 / 128, (112, 112), (127.5, 127.5, 127.5), swapRB=True
|
|
)
|
|
self.arcface.setInput(blob)
|
|
emb = self.arcface.forward().flatten()
|
|
norm = np.linalg.norm(emb)
|
|
return emb / norm if norm > 0 else None
|
|
|
|
@staticmethod
|
|
def _alinear_cara(frame, caja):
|
|
"""Alinea el rostro con los landmarks (ArcFace espera 112x112 alineado)."""
|
|
if len(caja) < 15:
|
|
return None
|
|
pts = caja[4:14].reshape(5, 2)
|
|
if pts.min() < 0:
|
|
return None
|
|
# Template estandar de ArcFace (112x112)
|
|
dst = np.array([
|
|
[38.2946, 51.6963],
|
|
[73.5318, 51.5014],
|
|
[56.0252, 71.7366],
|
|
[41.5493, 92.3655],
|
|
[70.7299, 92.2041],
|
|
], dtype=np.float32)
|
|
# Orden YuNet: ojo der, ojo izq, nariz, comisura der, comisura izq
|
|
src = np.array([
|
|
pts[1], pts[0], pts[2], pts[4], pts[3],
|
|
], dtype=np.float32)
|
|
try:
|
|
M, _ = cv2.estimateAffinePartial2D(src, dst, method=cv2.LMEDS)
|
|
except Exception:
|
|
return None
|
|
if M is None:
|
|
return None
|
|
alineada = cv2.warpAffine(
|
|
frame, M, (112, 112), borderValue=0.0
|
|
)
|
|
return alineada
|
|
|
|
# ------------------------------------------------------------------
|
|
# Cuerpo
|
|
# ------------------------------------------------------------------
|
|
def embedding_cuerpo(self, recorte_bgr):
|
|
if recorte_bgr.size == 0:
|
|
return None
|
|
rgb = cv2.cvtColor(recorte_bgr, cv2.COLOR_BGR2RGB)
|
|
tensor = self.tf_body(rgb).unsqueeze(0)
|
|
with torch.no_grad():
|
|
emb = self.body(tensor).numpy().flatten()
|
|
norm = np.linalg.norm(emb)
|
|
return emb / norm if norm > 0 else None
|
|
|
|
# ------------------------------------------------------------------
|
|
# Identificacion
|
|
# ------------------------------------------------------------------
|
|
def identificar_cara(self, emb):
|
|
if emb is None or not self.conocidos:
|
|
return None, 0.0
|
|
mejor, score = None, -1.0
|
|
for nombre, (emb_cara, _) in self.conocidos.items():
|
|
if emb_cara is not None:
|
|
sim = float(np.dot(emb, emb_cara))
|
|
if sim > score:
|
|
score = sim
|
|
mejor = nombre
|
|
return (mejor, score) if score >= UMBRAL_CARA else (None, score)
|
|
|
|
def identificar_cuerpo(self, emb):
|
|
if emb is None or not self.conocidos:
|
|
return None, 0.0
|
|
mejor, score = None, -1.0
|
|
for nombre, (_, emb_cuerpo) in self.conocidos.items():
|
|
if emb_cuerpo is not None:
|
|
sim = float(np.dot(emb, emb_cuerpo))
|
|
if sim > score:
|
|
score = sim
|
|
mejor = nombre
|
|
return (mejor, score) if score >= UMBRAL_CUERPO else (None, score)
|
|
|
|
def analizar_frame(self, frame):
|
|
"""Devuelve [{nombre:'X'} | {'desconocido':True}] por cada persona del frame."""
|
|
personas = self._detectar_personas(frame)
|
|
if not personas:
|
|
return []
|
|
resultados = []
|
|
for p in personas:
|
|
x, y, w, h = p
|
|
persona = self._identificar_recorte(frame, x, y, w, h)
|
|
resultados.append(persona)
|
|
return resultados
|
|
|
|
def _detectar_personas(self, frame):
|
|
try:
|
|
from ultralytics import YOLO
|
|
model = self._modelo_yolo()
|
|
r = model.predict(frame, verbose=False, conf=0.4)
|
|
cajas = []
|
|
for box in r[0].boxes:
|
|
if int(box.cls) == 0:
|
|
x1, y1, x2, y2 = box.xyxy[0].tolist()
|
|
cajas.append([int(x1), int(y1), int(x2 - x1), int(y2 - y1)])
|
|
return cajas
|
|
except Exception:
|
|
return []
|
|
|
|
def _modelo_yolo(self):
|
|
if not hasattr(self, "_yolo"):
|
|
from ultralytics import YOLO
|
|
self._yolo = YOLO(YOLO_PATH)
|
|
return self._yolo
|
|
|
|
def _identificar_recorte(self, frame, x, y, w, h):
|
|
x, y = max(0, x), max(0, y)
|
|
recorte = frame[y:y + h, x:x + w]
|
|
if recorte.size == 0:
|
|
return {"desconocido": True}
|
|
|
|
mejor_nombre, mejor_score = None, 0.0
|
|
for cara in self.detectar_caras(frame):
|
|
cx, cy, cw, ch = (int(v) for v in cara[:4])
|
|
if cx < x + w and cx + cw > x and cy < y + h and cy + ch > y:
|
|
emb = self.embedding_cara(frame, cara)
|
|
nombre, score = self.identificar_cara(emb)
|
|
if nombre and score > mejor_score:
|
|
mejor_nombre, mejor_score = nombre, score
|
|
|
|
if mejor_nombre:
|
|
return {"nombre": mejor_nombre, "score": mejor_score}
|
|
|
|
emb_cuerpo = self.embedding_cuerpo(recorte)
|
|
nombre, score = self.identificar_cuerpo(emb_cuerpo)
|
|
if nombre:
|
|
return {"nombre": nombre, "score": score, "via": "cuerpo"}
|
|
|
|
return {"desconocido": True}
|
|
|
|
|
|
def cargar_reconocedor():
|
|
return Reconocedor()
|