gx

chenyc

2025-06-12 7b72ac13a83764a662159d4a49b7fffb90476ecb

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
import * as tf from '../../dist/tfjs.esm';
 
import { fullyConnectedLayer } from '../common/fullyConnectedLayer';
import { NetInput } from '../dom/index';
import { FaceFeatureExtractorParams, IFaceFeatureExtractor, TinyFaceFeatureExtractorParams } from '../faceFeatureExtractor/types';
import { NeuralNetwork } from '../NeuralNetwork';
import { extractParams } from './extractParams';
import { extractParamsFromWeightMap } from './extractParamsFromWeightMap';
import { NetParams } from './types';
import { seperateWeightMaps } from './util';
 
export abstract class FaceProcessor<
  TExtractorParams extends FaceFeatureExtractorParams | TinyFaceFeatureExtractorParams
>
  extends NeuralNetwork<NetParams> {
  protected _faceFeatureExtractor: IFaceFeatureExtractor<TExtractorParams>;
 
  constructor(_name: string, faceFeatureExtractor: IFaceFeatureExtractor<TExtractorParams>) {
    super(_name);
    this._faceFeatureExtractor = faceFeatureExtractor;
  }
 
  public get faceFeatureExtractor(): IFaceFeatureExtractor<TExtractorParams> {
    return this._faceFeatureExtractor;
  }
 
  protected abstract override getDefaultModelName(): string
 
  protected abstract getClassifierChannelsIn(): number
 
  protected abstract getClassifierChannelsOut(): number
 
  public runNet(input: NetInput | tf.Tensor4D): tf.Tensor2D {
    const { params } = this;
 
    if (!params) {
      throw new Error(`${this._name} - load model before inference`);
    }
 
    return tf.tidy(() => {
      const bottleneckFeatures = input instanceof NetInput
        ? this.faceFeatureExtractor.forwardInput(input)
        : input;
      return fullyConnectedLayer(bottleneckFeatures.as2D(bottleneckFeatures.shape[0], -1), params.fc);
    });
  }
 
  public override dispose(throwOnRedispose = true) {
    this.faceFeatureExtractor.dispose(throwOnRedispose);
    super.dispose(throwOnRedispose);
  }
 
  public loadClassifierParams(weights: Float32Array) {
    const { params, paramMappings } = this.extractClassifierParams(weights);
    this._params = params;
    this._paramMappings = paramMappings;
  }
 
  public extractClassifierParams(weights: Float32Array) {
    return extractParams(weights, this.getClassifierChannelsIn(), this.getClassifierChannelsOut());
  }
 
  protected extractParamsFromWeightMap(weightMap: tf.NamedTensorMap) {
    const { featureExtractorMap, classifierMap } = seperateWeightMaps(weightMap);
 
    this.faceFeatureExtractor.loadFromWeightMap(featureExtractorMap);
 
    return extractParamsFromWeightMap(classifierMap);
  }
 
  protected extractParams(weights: Float32Array) {
    const cIn = this.getClassifierChannelsIn();
    const cOut = this.getClassifierChannelsOut();
    const classifierWeightSize = (cOut * cIn) + cOut;
 
    const featureExtractorWeights = weights.slice(0, weights.length - classifierWeightSize);
    const classifierWeights = weights.slice(weights.length - classifierWeightSize);
 
    this.faceFeatureExtractor.extractWeights(featureExtractorWeights);
    return this.extractClassifierParams(classifierWeights);
  }
}