import os import asyncio import time from typing import List, Any, cast from enum import Enum from fastapi import APIRouter, FastAPI, Request, WebSocket from fastapi.encoders import jsonable_encoder from pydantic import BaseModel import cv2 import logging from pathlib import Path import numpy as np from app.config import is_running_in_windows from app.wsmesssages.wsmessages import FaceData, FaceRecognitionData, Resolution from app.config import ServerConfiguration logger = logging.getLogger(__name__) # --- Initialisation des drivers et importations --- try: from ultralytics import YOLO as yolo_cls # type: ignore[reportMissingImports] except Exception: logger.warning("ultralytics n'est pas installé. Le détecteur YOLO ne sera pas disponible.") yolo_cls = None # Variables globales pour stocker le délégué Coral unique g_coral_delegate: Any = None tflite_module: Any = None def _init_coral_driver() -> Any: """Initialise et retourne le délégué Coral de manière unique (Singleton) avec support LiteRT.""" global g_coral_delegate, tflite_module if g_coral_delegate is not None: return g_coral_delegate if is_running_in_windows(): try: # 1. On tente d'abord d'utiliser le nouveau package officiel LiteRT import ai_edge_litert.interpreter as litert # type: ignore tflite_module = litert dll_path = "edgetpu.dll" if hasattr(litert, 'load_delegate'): g_coral_delegate = litert.load_delegate(dll_path) else: # Fallback au cas où l'implémentation binaire sous-jacente diffère from ai_edge_litert.interpreter.wrapper import load_delegate as _load_delegate # type: ignore g_coral_delegate = _load_delegate(dll_path) logger.info("Délégué Coral Windows (edgetpu.dll) initialisé via LiteRT.") except ImportError: # 2. Repli sur l'ancien TensorFlow Lite si le nouveau package n'est pas installé try: import tensorflow.lite as tf_lite # type: ignore import tensorflow as tf tflite_module = tf_lite dll_path = "edgetpu.dll" if hasattr(tf_lite, 'load_delegate'): g_coral_delegate = tf_lite.load_delegate(dll_path) elif hasattr(tf.lite.experimental, 'load_delegate'): g_coral_delegate = tf.lite.experimental.load_delegate(dll_path) else: from tensorflow.python.interpreter.wrapper import load_delegate as _load_delegate g_coral_delegate = _load_delegate(dll_path) logger.info("Délégué Coral Windows initialisé via l'ancien tf.lite (déprécié).") except ImportError: logger.warning("Ni ai-edge-litert ni tensorflow ne sont installés. Coral indisponible sous Windows.") else: # Branche Linux (tflite_runtime reste pour le moment le standard pycoral) try: import tflite_runtime.interpreter as tflite_rt # type: ignore from pycoral.utils.edgetpu import load_edgedpu_delegate # type: ignore tflite_module = tflite_rt g_coral_delegate = load_edgedpu_delegate() logger.info("Délégué Coral Linux (libedgetpu) initialisé avec succès.") except ImportError: logger.warning("tflite_runtime ou pycoral n'est pas installé. Coral indisponible sous Linux.") return g_coral_delegate url = '/facerecog' swag_tags: List[str | Enum] = ["Face detection"] router = APIRouter(tags=swag_tags, prefix=url) # --- Modèles Pydantic --- class FaceRecogStartCmd(BaseModel): url: str camera: str detector: str | None = 'cv2' class FaceRecogStartResponse(BaseModel): camera: str url: str status: str message: str detector: str class FaceRecogStopCmd(BaseModel): camera: str url: str class FaceRecogStopResponse(BaseModel): camera: str url: str status: str message: str class FaceRecogStatusCmd(BaseModel): camera: str url: str class FaceRecogStatusResponse(BaseModel): camera: str url: str status: str active: bool task_exists: bool task_done: bool task_cancelled: bool detector: str | None = None # --- État global --- g_active_streams: dict[str, bool] = {} g_cameras_by_stream: dict[str, str] = {} g_active_processes: dict[str, asyncio.Task[Any]] = {} g_stream_detectors: dict[str, str] = {} g_yolo_models_by_stream: dict[str, Any] = {} g_yolo_face_class_ids_by_stream: dict[str, list[int]] = {} # Chargement OpenCV Haar Cascade haar_root = getattr(getattr(cv2, "data", None), "haarcascades", "") face_cascade = cv2.CascadeClassifier(os.path.join(haar_root, 'haarcascade_frontalface_default.xml')) def _on_stream_task_done(url: str, task: asyncio.Task[Any]) -> None: try: if not task.cancelled(): exc = task.exception() if exc is not None: logger.error(f"Tâche d'analyse terminée en erreur pour {url}: {exc}") except Exception as callback_exc: logger.error(f"Erreur dans le callback de fin de tâche pour {url}: {callback_exc}") finally: g_active_processes.pop(url, None) def _resolve_detector(detector: str | None = None) -> str: selected = (detector or os.getenv("SVL_FACE_DETECTOR", "cv2")).strip().lower() if selected not in {"cv2", "yolo"}: raise ValueError("detector must be 'cv2' or 'yolo'") return selected def _get_yolo_model(stream_key: str, device: str = "cpu") -> Any: """Récupère ou initialise le modèle selon la plateforme et le device cible.""" model = g_yolo_models_by_stream.get(stream_key) if model is not None: return model is_windows = is_running_in_windows() default_filename = ( "ssd_mobilenet_v2_face_quant_postprocess_edgetpu.tflite" if device in {"coral", "edgetpu"} else "yolov8n-face.pt" ) # Adaptation automatique du nom sous Linux pour votre version spécifique if not is_windows and device in {"coral", "edgetpu"}: default_filename = "yolov8n-face_full_integer_quant_edgetpu.tflite" default_model_path = Path(__file__).resolve().parent.parent / f"yolo/models/{default_filename}" model_path = os.getenv("SVL_YOLO_MODEL", str(default_model_path)) # Cas Edge TPU (Coral) if device in {"coral", "edgetpu"}: delegate = _init_coral_driver() if delegate is None or tflite_module is None: raise RuntimeError("Le détecteur Coral a été demandé mais les drivers requis sont absents.") try: interpreter = tflite_module.Interpreter( model_path=model_path, experimental_delegates=[delegate] ) interpreter.allocate_tensors() g_yolo_models_by_stream[stream_key] = interpreter logger.info(f"Modèle Edge TPU Coral chargé avec succès depuis {model_path}") return interpreter except Exception as e: logger.error(f"Erreur lors de l'initialisation de l'interpréteur Coral : {e}") raise RuntimeError(f"Impossible d'initialiser le Coral USB : {e}") # Cas Standard (Ultralytics CPU/GPU) if yolo_cls is None: raise RuntimeError("Détecteur YOLO demandé mais 'ultralytics' n'est pas installé.") model = cast(Any, yolo_cls(model_path)) g_yolo_models_by_stream[stream_key] = model logger.info(f"Modèle Ultralytics classique chargé depuis {model_path}") return model def _detect_faces_cv2(frame: Any) -> list[FaceData]: gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30)) return [FaceData(x=int(x), y=int(y), w=int(w), h=int(h)) for (x, y, w, h) in faces] def _detect_faces_yolo(app: FastAPI, frame: Any, stream_key: str) -> list[FaceData]: config: ServerConfiguration | None = getattr(app.state, "config", None) if config is None: logger.error("config indisponible dans app.state") return [] device = config.yolo.device.strip().lower() conf_threshold = config.yolo.confidence max_faces = config.yolo.max_faces iou_threshold = config.yolo.iou model_or_interpreter = _get_yolo_model(stream_key, device=device) faces_coords: list[FaceData] = [] frame_h, frame_w = frame.shape[:2] # ------------------------------------------------------------- # BRANCHE A : INFERENCE SUR GOOGLE CORAL (TFLITE) # ------------------------------------------------------------- if device in {"coral", "edgetpu"}: interpreter = model_or_interpreter input_details = interpreter.get_input_details() output_details = interpreter.get_output_details() input_shape = input_details[0]['shape'] model_h, model_w = input_shape[1], input_shape[2] # Prétraitement optimisé img_resized = cv2.resize(frame, (model_w, model_h), interpolation=cv2.INTER_LINEAR) img_rgb = cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB) input_data = np.expand_dims(img_rgb, axis=0).astype(np.uint8) interpreter.set_tensor(input_details[0]['index'], input_data) interpreter.invoke() # Récupération et parsing selon l'architecture du modèle de sortie # Gestion dynamique pour supporter à la fois les sorties YOLOv8-tflite et SSD MobileNet if len(output_details) > 1: # Format Typique SSD MobileNet (Boxes, Classes, Scores, Count) raw_boxes = np.squeeze(interpreter.get_tensor(output_details[0]['index'])) raw_scores = np.squeeze(interpreter.get_tensor(output_details[2]['index'])) # Filtrage vectorisé NumPy rapide valid_indices = np.where(raw_scores >= conf_threshold)[0] boxes: list[list[int]] = [] confidences: list[float] = [] for idx in valid_indices: # Format SSD souvent normalisé [ymin, xmin, ymax, xmax] ymin, xmin, ymax, xmax = raw_boxes[idx] x1 = int(xmin * frame_w) y1 = int(ymin * frame_h) w = int((xmax - xmin) * frame_w) h = int((ymax - ymin) * frame_h) if w > 0 and h > 0: boxes.append([x1, y1, w, h]) confidences.append(float(raw_scores[idx])) else: # Format Typique YOLOv8 (Tenseur unique transposé ou non) predictions = np.squeeze(interpreter.get_tensor(output_details[0]['index'])) if predictions.shape[0] > predictions.shape[1]: predictions = predictions.T # YOLOv8 face : [x_center, y_center, w, h, score_face, ...] # Utilisation de NumPy pour filtrer les milliers de lignes d'un coup scores = predictions[:, 4] valid_indices = np.where(scores >= conf_threshold)[0] boxes: list[list[int]] = [] confidences: list[float] = [] for idx in valid_indices: pred = predictions[idx] x_center, y_center, w_model, h_model = pred[0:4] x1 = int((x_center - w_model / 2) * (frame_w / model_w)) y1 = int((y_center - h_model / 2) * (frame_h / model_h)) w = int(w_model * (frame_w / model_w)) h = int(h_model * (frame_h / model_h)) if w > 0 and h > 0: boxes.append([x1, y1, w, h]) confidences.append(float(pred[4])) if not boxes: return [] # Application obligatoire du Non-Maximum Suppression (NMS) pour éviter les boîtes multiples indices_nms = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, iou_threshold) if len(indices_nms) > 0: for idx in indices_nms.flatten()[:max_faces]: x1, y1, w, h = boxes[idx] faces_coords.append(FaceData(x=max(0, x1), y=max(0, y1), w=w, h=h)) return faces_coords # ------------------------------------------------------------- # BRANCHE B : INFERENCE STANDARD ULTRALYTICS (CPU/CUDA) # ------------------------------------------------------------- model = model_or_interpreter frame_area = max(1, frame_w * frame_h) min_area_ratio = config.yolo.min_area_ratio max_area_ratio = config.yolo.max_area_ratio use_half = config.yolo.use_half face_class_ids = g_yolo_face_class_ids_by_stream.get(stream_key) if face_class_ids is None: model_names = cast(dict[Any, Any], getattr(model, "names", {})) face_class_ids = [int(k) for k, v in model_names.items() if str(v).strip().lower() == "face"] g_yolo_face_class_ids_by_stream[stream_key] = face_class_ids if not face_class_ids: raise RuntimeError("Le modèle YOLO n'a pas de classe 'face'.") results = model.predict( source=frame, verbose=False, conf=conf_threshold, iou=iou_threshold, classes=face_class_ids, imgsz=config.yolo.imgsz, max_det=max_faces, device=device, half=use_half, ) for result in results: boxes = getattr(result, "boxes", None) if boxes is None: continue for box in boxes: xyxy = box.xyxy[0].tolist() x1, y1, x2, y2 = [int(v) for v in xyxy] w = max(0, x2 - x1) h = max(0, y2 - y1) area_ratio = (w * h) / frame_area if w > 0 and h > 0 and min_area_ratio <= area_ratio <= max_area_ratio: faces_coords.append(FaceData(x=x1, y=y1, w=w, h=h)) if len(faces_coords) >= max_faces: return faces_coords return faces_coords def _detect_faces(app: FastAPI, frame: Any, detector: str, stream_key: str) -> list[FaceData]: if detector == "cv2": return _detect_faces_cv2(frame) if detector == "yolo": return _detect_faces_yolo(app, frame, stream_key) raise ValueError("Unsupported detector") def _rescale_faces_to_original(faces_data: list[FaceData], scale_x: float, scale_y: float) -> list[FaceData]: if scale_x == 1.0 and scale_y == 1.0: return faces_data return [ FaceData( x=int(face.x * scale_x), y=int(face.y * scale_y), w=max(1, int(face.w * scale_x)), h=max(1, int(face.h * scale_y)), ) for face in faces_data ] def _get_active_websockets(app: FastAPI) -> set[WebSocket] | None: raw_websockets = getattr(app.state, "active_websockets", None) if raw_websockets is None: return None return cast(set[WebSocket], raw_websockets) async def broadcast_faces(app: FastAPI, stream_url: str, camera: str, faces_data: list[FaceData], frame_width: int, frame_height: int): frd = FaceRecognitionData( stream_url=stream_url, camera=camera, resolution=Resolution(width=frame_width, height=frame_height), data=faces_data ) payload = jsonable_encoder(frd.toDict) active_websockets = _get_active_websockets(app) if not active_websockets: return tasks = [asyncio.create_task(ws.send_json(payload)) for ws in active_websockets] results = await asyncio.gather(*tasks, return_exceptions=True) for result in results: if isinstance(result, Exception): logger.error(f"Erreur lors de l'envoi à un client WebSocket : {result}") async def process_video_stream(app: FastAPI, url: str, camera: str): config: ServerConfiguration | None = getattr(app.state, "config", None) if config is None: logger.error("config indisponible dans app.state") return while g_active_streams.get(url, False): detector = g_stream_detectors.get(url, _resolve_detector()) os.environ["OPENCV_FFMPEG_CAPTURE_OPTIONS"] = "fflags;nobuffer|flags;low_delay" cap = cv2.VideoCapture(url, cv2.CAP_FFMPEG) cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) if not cap.isOpened(): logger.error(f"Erreur : Impossible d'ouvrir le flux {url}") g_active_streams[url] = False return detect_width = config.facerecog.detect_width process_every_n_frames = config.facerecog.process_every_n_frames drop_grab_count = config.facerecog.drop_grab_count max_broadcast_fps = config.facerecog.max_broadcast_fps loop_sleep_seconds = config.facerecog.loop_sleep_seconds min_broadcast_interval = 1.0 / max_broadcast_fps last_broadcast_time = 0.0 frame_index = 0 while cap.isOpened(): if not g_active_streams.get(url, False): break for _ in range(drop_grab_count): if not g_active_streams.get(url, False): break grabbed = await asyncio.to_thread(cap.grab) if not grabbed: break if not g_active_streams.get(url, False): break ret, frame = await asyncio.to_thread(cap.read) if not g_active_streams.get(url, False) or not ret: break frame_index += 1 if process_every_n_frames > 1 and frame_index % process_every_n_frames != 0: if loop_sleep_seconds > 0.0: await asyncio.sleep(loop_sleep_seconds) continue frame_height, frame_width, _ = frame.shape detect_frame = frame scale_x, scale_y = 1.0, 1.0 if detect_width > 0 and frame_width > detect_width: detect_height = max(1, int(frame_height * (detect_width / frame_width))) detect_frame = cv2.resize(frame, (detect_width, detect_height), interpolation=cv2.INTER_LINEAR) scale_x = frame_width / detect_width scale_y = frame_height / detect_height try: faces_coords = await asyncio.to_thread(_detect_faces, app, detect_frame, detector, url) if not g_active_streams.get(url, False): break faces_coords = _rescale_faces_to_original(faces_coords, scale_x, scale_y) except Exception as exc: logger.error(f"Erreur détecteur '{detector}' ({exc}), fallback vers cv2.") g_stream_detectors[url] = "cv2" detector = "cv2" faces_coords = await asyncio.to_thread(_detect_faces, app, detect_frame, "cv2", url) if not g_active_streams.get(url, False): break faces_coords = _rescale_faces_to_original(faces_coords, scale_x, scale_y) now = time.perf_counter() if faces_coords and (now - last_broadcast_time) >= min_broadcast_interval: await broadcast_faces(app, url, camera, faces_coords, frame_width, frame_height) last_broadcast_time = now if loop_sleep_seconds > 0.0: await asyncio.sleep(loop_sleep_seconds) cap.release() if not g_active_streams.get(url, False): break g_active_streams.pop(url, None) g_cameras_by_stream.pop(url, None) g_stream_detectors.pop(url, None) g_yolo_models_by_stream.pop(url, None) g_yolo_face_class_ids_by_stream.pop(url, None) @router.post("/start", response_model=FaceRecogStartResponse) async def start_stream(params: FaceRecogStartCmd, request: Request): url = params.url camera = params.camera detector = _resolve_detector(params.detector) running_task = g_active_processes.get(url) if g_active_streams.get(url) and running_task is not None and not running_task.done(): return FaceRecogStartResponse( camera=camera, url=url, status="info", message=f"Le flux {url} est déjà en cours d'analyse.", detector=detector ) g_active_streams[url] = True g_stream_detectors[url] = detector g_cameras_by_stream[url] = camera task = asyncio.create_task(process_video_stream(request.app, url, camera), name=f"facerecog:{camera}:{url}") task.add_done_callback(lambda done_task, stream_url=url: _on_stream_task_done(stream_url, done_task)) g_active_processes[url] = task return FaceRecogStartResponse( camera=camera, url=url, status="success", message=f"Analyse démarrée pour le flux {url}", detector=detector ) @router.post("/stop", response_model=FaceRecogStopResponse) async def stop_stream(params: FaceRecogStopCmd): url = params.url camera = params.camera if g_active_streams.get(url): g_active_streams[url] = False task = g_active_processes.get(url) if task is not None: try: await asyncio.wait_for(task, timeout=10.0) except asyncio.TimeoutError: logger.warning(f"Le processus d'analyse pour {url} n'a pas terminé à temps.") return FaceRecogStopResponse(camera=camera, url=url, status="success", message=f"Arrêt de l'analyse pour {url}") return FaceRecogStopResponse(camera=camera, url=url, status="error", message="Flux introuvable ou déjà arrêté.") @router.post("/status", response_model=FaceRecogStatusResponse) async def stream_status(params: FaceRecogStatusCmd): url = params.url camera = params.camera task = g_active_processes.get(url) is_active = bool(g_active_streams.get(url, False)) task_exists = task is not None task_done = bool(task.done()) if task is not None else True task_cancelled = bool(task.cancelled()) if task is not None else False status = "running" if (is_active and task_exists and not task_done) else "stopped" return FaceRecogStatusResponse( camera=camera, url=url, status=status, active=is_active, task_exists=task_exists, task_done=task_done, task_cancelled=task_cancelled, detector=g_stream_detectors.get(url) ) @router.post("/stopall", response_model=List[FaceRecogStopResponse]) async def stop_all_streams() -> List[FaceRecogStopResponse]: responses: List[FaceRecogStopResponse] = [] for url in list(g_active_streams.keys()): camera = g_cameras_by_stream.get(url, "") params = FaceRecogStopCmd(camera=camera, url=url) response = await stop_stream(params) responses.append(response) return responses