02-ago: [proyecto] vigilancia - reconocimiento facial (YuNet+ArcFace) y corporal (ResNet18) + registro por webcam (enroll.py) + integracion en main.py

vigilancia/reconocedor.py

@@ -0,0 +1,184 @@
+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):
+        x, y, w, h = (int(v) for v in caja[:4])
+        x, y = max(0, x), max(0, y)
+        recorte = frame[y:y + h, x:x + w]
+        if recorte.size == 0:
+            return None
+        blob = cv2.dnn.blobFromImage(
+            recorte, 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
+
+    # ------------------------------------------------------------------
+    # 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()