gx

chenyc

2025-02-12 ea42ff3ebee1eeb3fb29423aa848a249441db81c

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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
import * as tf from '../../dist/tfjs.esm.js';
import { fullyConnectedLayer } from '../common/fullyConnectedLayer';
import { seperateWeightMaps } from '../faceProcessor/util';
import { TinyXception } from '../xception/TinyXception';
import { extractParams } from './extractParams';
import { extractParamsFromWeightMap } from './extractParamsFromWeightMap';
import { AgeAndGenderPrediction, Gender, NetOutput, NetParams } from './types';
import { NeuralNetwork } from '../NeuralNetwork';
import { NetInput, TNetInput, toNetInput } from '../dom/index';
 
export class AgeGenderNet extends NeuralNetwork<NetParams> {
  private _faceFeatureExtractor: TinyXception;
 
  constructor(faceFeatureExtractor: TinyXception = new TinyXception(2)) {
    super('AgeGenderNet');
    this._faceFeatureExtractor = faceFeatureExtractor;
  }
 
  public get faceFeatureExtractor(): TinyXception {
    return this._faceFeatureExtractor;
  }
 
  public runNet(input: NetInput | tf.Tensor4D): NetOutput {
    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;
 
      const pooled = tf.avgPool(bottleneckFeatures, [7, 7], [2, 2], 'valid').as2D(bottleneckFeatures.shape[0], -1);
      const age = fullyConnectedLayer(pooled, params.fc.age).as1D();
      const gender = fullyConnectedLayer(pooled, params.fc.gender);
      return { age, gender };
    });
  }
 
  public forwardInput(input: NetInput | tf.Tensor4D): NetOutput {
    return tf.tidy(() => {
      const { age, gender } = this.runNet(input);
      return { age, gender: tf.softmax(gender) };
    });
  }
 
  public async forward(input: TNetInput): Promise<NetOutput> {
    return this.forwardInput(await toNetInput(input));
  }
 
  public async predictAgeAndGender(input: TNetInput): Promise<AgeAndGenderPrediction | AgeAndGenderPrediction[]> {
    const netInput = await toNetInput(input);
    const out = await this.forwardInput(netInput);
 
    const ages = tf.unstack(out.age);
    const genders = tf.unstack(out.gender);
    const ageAndGenderTensors = ages.map((ageTensor, i) => ({
      ageTensor,
      genderTensor: genders[i],
    }));
 
    const predictionsByBatch = await Promise.all(
      ageAndGenderTensors.map(async ({ ageTensor, genderTensor }) => {
        const age = (ageTensor.dataSync())[0];
        const probMale = (genderTensor.dataSync())[0];
        const isMale = probMale > 0.5;
        const gender = isMale ? Gender.MALE : Gender.FEMALE;
        const genderProbability = isMale ? probMale : (1 - probMale);
 
        ageTensor.dispose();
        genderTensor.dispose();
        return { age, gender, genderProbability };
      }),
    );
    out.age.dispose();
    out.gender.dispose();
 
    return netInput.isBatchInput ? predictionsByBatch as AgeAndGenderPrediction[] : predictionsByBatch[0] as AgeAndGenderPrediction;
  }
 
  protected getDefaultModelName(): string {
    return 'age_gender_model';
  }
 
  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);
  }
 
  protected extractParamsFromWeightMap(weightMap: tf.NamedTensorMap) {
    const { featureExtractorMap, classifierMap } = seperateWeightMaps(weightMap);
 
    this.faceFeatureExtractor.loadFromWeightMap(featureExtractorMap);
 
    return extractParamsFromWeightMap(classifierMap);
  }
 
  protected extractParams(weights: Float32Array) {
    const classifierWeightSize = (512 * 1 + 1) + (512 * 2 + 2);
 
    const featureExtractorWeights = weights.slice(0, weights.length - classifierWeightSize);
    const classifierWeights = weights.slice(weights.length - classifierWeightSize);
 
    this.faceFeatureExtractor.extractWeights(featureExtractorWeights);
    return this.extractClassifierParams(classifierWeights);
  }
}