id-scanner-lib/src/modules/face/face-model-loader.ts

Browser-based ID card, QR code, and face recognition scanner with liveness detection

/**
 * @file 人脸模型加载器
 * @description 统一管理 face-api 模型的懒加载和缓存
 * @module modules/face/face-model-loader
 */

import * as faceapi from '@vladmandic/face-api';
import { Logger } from '../../core/logger';
import { ResourceLoadError } from '../../core/errors';
import { FaceModelType } from './face-detector';

/**
 * 已加载模型记录
 */
interface LoadedModelRecord {
  /** 模型名称 */
  name: string;
  /** 加载时间戳 */
  loadedAt: number;
}

/**
 * 模型加载器配置
 */
export interface ModelLoaderConfig {
  /** 模型路径 */
  modelPath: string;
  /** 检测模型类型 */
  detectionModel: FaceModelType;
  /** 关键点模型类型 */
  landmarksModel: 'tiny' | '68_points';
  /** 是否检测关键点 */
  detectLandmarks: boolean;
  /** 是否检测表情 */
  detectExpressions: boolean;
  /** 是否检测年龄和性别 */
  detectAgeGender: boolean;
  /** 是否提取人脸特征向量 */
  extractEmbeddings: boolean;
}

/**
 * 人脸模型加载器
 *
 * 负责按需加载 face-api 模型，避免一次性加载所有模型造成内存浪费
 */
export class FaceModelLoader {
  /** 日志记录器 */
  private logger: Logger;

  /** 已加载模型集合 */
  private loadedModels: Set<string> = new Set();

  /** 模型是否已全部加载 */
  private modelsLoaded: boolean = false;

  /** 配置 */
  private config: ModelLoaderConfig;

  /**
   * 构造函数
   * @param config 模型加载配置
   */
  constructor(config: ModelLoaderConfig) {
    this.logger = Logger.getInstance();
    this.config = config;
  }

  /**
   * 获取检测模型选项
   */
  private getDetectorOptions(minConfidence: number): faceapi.SsdMobilenetv1Options | faceapi.TinyFaceDetectorOptions | faceapi.MtcnnOptions {
    switch (this.config.detectionModel) {
      case FaceModelType.SSD_MOBILENET:
        return new faceapi.SsdMobilenetv1Options({ minConfidence });
      case FaceModelType.TINY_FACE:
        return new faceapi.TinyFaceDetectorOptions({ scoreThreshold: minConfidence });
      case FaceModelType.MTCNN:
        return new faceapi.MtcnnOptions({ minConfidence });
      default:
        return new faceapi.SsdMobilenetv1Options({ minConfidence });
    }
  }

  /**
   * 懒加载单个模型
   * @param modelType 模型类型
   * @param modelPath 模型路径
   */
  async lazyLoadModel(modelType: string, modelPath: string): Promise<void> {
    if (this.loadedModels.has(modelType)) {
      return;
    }

    this.logger.info('FaceModelLoader', `懒加载模型: ${modelType}`);

    try {
      switch (modelType) {
        case 'ssdMobilenetv1':
          await faceapi.nets.ssdMobilenetv1.loadFromUri(modelPath);
          break;
        case 'tinyFaceDetector':
          await faceapi.nets.tinyFaceDetector.loadFromUri(modelPath);
          break;
        case 'faceLandmark68Net':
          await faceapi.nets.faceLandmark68Net.loadFromUri(modelPath);
          break;
        case 'faceLandmark68TinyNet':
          await faceapi.nets.faceLandmark68TinyNet.loadFromUri(modelPath);
          break;
        case 'faceExpressionNet':
          await faceapi.nets.faceExpressionNet.loadFromUri(modelPath);
          break;
        case 'ageGenderNet':
          await faceapi.nets.ageGenderNet.loadFromUri(modelPath);
          break;
        case 'faceRecognitionNet':
          await faceapi.nets.faceRecognitionNet.loadFromUri(modelPath);
          break;
        default:
          this.logger.warn('FaceModelLoader', `未知模型类型: ${modelType}`);
          return;
      }

      this.loadedModels.add(modelType);
      this.logger.info('FaceModelLoader', `模型加载完成: ${modelType}`);
    } catch (error) {
      this.logger.error('FaceModelLoader', `模型加载失败: ${modelType}`, error as Error);
      throw new ResourceLoadError(modelType, `模型加载失败: ${error}`);
    }
  }

  /**
   * 根据需求加载模型
   * @param options 检测选项
   */
  async loadModelsOnDemand(options: {
    withLandmarks?: boolean;
    withAttributes?: boolean;
    withEmbedding?: boolean;
  }): Promise<void> {
    const { modelPath, landmarksModel, detectExpressions, detectAgeGender, extractEmbeddings } = this.config;

    // 基础检测模型
    const detectionModelKey = this.config.detectionModel === FaceModelType.TINY_FACE
      ? 'tinyFaceDetector'
      : 'ssdMobilenetv1';
    await this.lazyLoadModel(detectionModelKey, modelPath);

    // 关键点模型
    if (options.withLandmarks || this.config.detectLandmarks) {
      await this.lazyLoadModel(
        landmarksModel === '68_points' ? 'faceLandmark68Net' : 'faceLandmark68TinyNet',
        modelPath
      );
    }

    // 表情模型
    if (options.withAttributes || detectExpressions) {
      await this.lazyLoadModel('faceExpressionNet', modelPath);
    }

    // 年龄性别模型
    if (options.withAttributes || detectAgeGender) {
      await this.lazyLoadModel('ageGenderNet', modelPath);
    }

    // 人脸识别模型
    if (options.withEmbedding || extractEmbeddings) {
      await this.lazyLoadModel('faceRecognitionNet', modelPath);
    }

    this.modelsLoaded = true;
  }

  /**
   * 确保模型已加载（用于需要所有模型的场景）
   */
  async ensureModelsLoaded(): Promise<void> {
    if (this.modelsLoaded) return;
    await this.loadModelsOnDemand({});
  }

  /**
   * 获取模型加载状态
   */
  isModelsLoaded(): boolean {
    return this.modelsLoaded;
  }

  /**
   * 获取已加载模型列表
   */
  getLoadedModels(): string[] {
    return Array.from(this.loadedModels);
  }

  /**
   * 释放所有模型（释放 TensorFlow.js 内存）
   */
  async dispose(): Promise<void> {
    if (this.loadedModels.size > 0) {
      try {
        // face-api 使用 TensorFlow.js，需要显式释放
        const tf = await import('@tensorflow/tfjs');
        tf.dispose();
        this.loadedModels.clear();
        this.modelsLoaded = false;
        this.logger.debug('FaceModelLoader', '模型资源已释放');
      } catch (error) {
        this.logger.error('FaceModelLoader', `释放模型失败: ${error}`);
      }
    }
  }

  /**
   * 获取检测器选项（用于 detectAllFaces）
   */
  getFaceDetectionOptions(minConfidence: number): faceapi.SsdMobilenetv1Options | faceapi.TinyFaceDetectorOptions | faceapi.MtcnnOptions {
    return this.getDetectorOptions(minConfidence);
  }
}