Skip to content

Commit 178904c

Browse files
justadudewhohacksjustadudewhohacks
committed
seperate feature extraction nets from landmark nets to make them reusable
1 parent 7e766eb commit 178904c

19 files changed

+399
-200
lines changed

‎src/faceFeatureExtractor/FaceFeatureExtractor.ts

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
import * as tf from '@tensorflow/tfjs-core';
2+
import { NetInput, NeuralNetwork, normalize } from 'tfjs-image-recognition-base';
3+
import { ConvParams, SeparableConvParams } from 'tfjs-tiny-yolov2';
4+
5+
import { depthwiseSeparableConv } from './depthwiseSeparableConv';
6+
import { extractParams } from './extractParams';
7+
import { extractParamsFromWeigthMap } from './extractParamsFromWeigthMap';
8+
import { DenseBlock4Params, NetParams } from './types';
9+
10+
function denseBlock(
11+
x: tf.Tensor4D,
12+
denseBlockParams: DenseBlock4Params,
13+
isFirstLayer: boolean = false
14+
): tf.Tensor4D {
15+
return tf.tidy(() => {
16+
const out1 = tf.relu(
17+
isFirstLayer
18+
? tf.add(
19+
tf.conv2d(x, (denseBlockParams.conv0 as ConvParams).filters, [2, 2], 'same'),
20+
denseBlockParams.conv0.bias
21+
)
22+
: depthwiseSeparableConv(x, denseBlockParams.conv0 as SeparableConvParams, [2, 2])
23+
) as tf.Tensor4D
24+
const out2 = depthwiseSeparableConv(out1, denseBlockParams.conv1, [1, 1])
25+
26+
const in3 = tf.relu(tf.add(out1, out2)) as tf.Tensor4D
27+
const out3 = depthwiseSeparableConv(in3, denseBlockParams.conv2, [1, 1])
28+
29+
const in4 = tf.relu(tf.add(out1, tf.add(out2, out3))) as tf.Tensor4D
30+
const out4 = depthwiseSeparableConv(in4, denseBlockParams.conv3, [1, 1])
31+
32+
return tf.relu(tf.add(out1, tf.add(out2, tf.add(out3, out4)))) as tf.Tensor4D
33+
})
34+
}
35+
36+
export class FaceFeatureExtractor extends NeuralNetwork<NetParams> {
37+
38+
constructor() {
39+
super('FaceFeatureExtractor')
40+
}
41+
42+
public forward(input: NetInput): tf.Tensor4D {
43+
44+
const { params } = this
45+
46+
if (!params) {
47+
throw new Error('FaceFeatureExtractor - load model before inference')
48+
}
49+
50+
return tf.tidy(() => {
51+
const batchTensor = input.toBatchTensor(112, true)
52+
const meanRgb = [122.782, 117.001, 104.298]
53+
const normalized = normalize(batchTensor, meanRgb).div(tf.scalar(255)) as tf.Tensor4D
54+
55+
let out = denseBlock(normalized, params.dense0, true)
56+
out = denseBlock(out, params.dense1)
57+
out = denseBlock(out, params.dense2)
58+
out = denseBlock(out, params.dense3)
59+
out = tf.avgPool(out, [7, 7], [2, 2], 'valid')
60+
61+
return out
62+
})
63+
}
64+
65+
protected getDefaultModelName(): string {
66+
return 'face_feature_extractor_model'
67+
}
68+
69+
protected extractParamsFromWeigthMap(weightMap: tf.NamedTensorMap) {
70+
return extractParamsFromWeigthMap(weightMap)
71+
}
72+
73+
protected extractParams(weights: Float32Array) {
74+
return extractParams(weights)
75+
}
76+
}

‎src/faceFeatureExtractor/TinyFaceFeatureExtractor.ts

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
import * as tf from '@tensorflow/tfjs-core';
2+
import { NetInput, NeuralNetwork, normalize } from 'tfjs-image-recognition-base';
3+
import { ConvParams, SeparableConvParams } from 'tfjs-tiny-yolov2';
4+
5+
import { depthwiseSeparableConv } from './depthwiseSeparableConv';
6+
import { extractParamsFromWeigthMapTiny } from './extractParamsFromWeigthMapTiny';
7+
import { extractParamsTiny } from './extractParamsTiny';
8+
import { DenseBlock3Params, TinyNetParams } from './types';
9+
10+
function denseBlock(
11+
x: tf.Tensor4D,
12+
denseBlockParams: DenseBlock3Params,
13+
isFirstLayer: boolean = false
14+
): tf.Tensor4D {
15+
return tf.tidy(() => {
16+
const out1 = tf.relu(
17+
isFirstLayer
18+
? tf.add(
19+
tf.conv2d(x, (denseBlockParams.conv0 as ConvParams).filters, [2, 2], 'same'),
20+
denseBlockParams.conv0.bias
21+
)
22+
: depthwiseSeparableConv(x, denseBlockParams.conv0 as SeparableConvParams, [2, 2])
23+
) as tf.Tensor4D
24+
const out2 = depthwiseSeparableConv(out1, denseBlockParams.conv1, [1, 1])
25+
26+
const in3 = tf.relu(tf.add(out1, out2)) as tf.Tensor4D
27+
const out3 = depthwiseSeparableConv(in3, denseBlockParams.conv2, [1, 1])
28+
29+
return tf.relu(tf.add(out1, tf.add(out2, out3))) as tf.Tensor4D
30+
})
31+
}
32+
33+
export class TinyFaceFeatureExtractor extends NeuralNetwork<TinyNetParams> {
34+
35+
constructor() {
36+
super('TinyFaceFeatureExtractor')
37+
}
38+
39+
public forward(input: NetInput): tf.Tensor4D {
40+
41+
const { params } = this
42+
43+
if (!params) {
44+
throw new Error('TinyFaceFeatureExtractor - load model before inference')
45+
}
46+
47+
return tf.tidy(() => {
48+
const batchTensor = input.toBatchTensor(112, true)
49+
const meanRgb = [122.782, 117.001, 104.298]
50+
const normalized = normalize(batchTensor, meanRgb).div(tf.scalar(255)) as tf.Tensor4D
51+
52+
let out = denseBlock(normalized, params.dense0, true)
53+
out = denseBlock(out, params.dense1)
54+
out = denseBlock(out, params.dense2)
55+
out = tf.avgPool(out, [14, 14], [2, 2], 'valid')
56+
57+
return out
58+
})
59+
}
60+
61+
protected getDefaultModelName(): string {
62+
return 'face_landmark_68_tiny_model'
63+
}
64+
65+
protected extractParamsFromWeigthMap(weightMap: tf.NamedTensorMap) {
66+
return extractParamsFromWeigthMapTiny(weightMap)
67+
}
68+
69+
protected extractParams(weights: Float32Array) {
70+
return extractParamsTiny(weights)
71+
}
72+
}

‎src/faceLandmarkNet/depthwiseSeparableConv.ts renamed to ‎src/faceFeatureExtractor/depthwiseSeparableConv.ts

File renamed without changes.

‎src/faceFeatureExtractor/extractParams.ts

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
import { extractWeightsFactory, ParamMapping } from 'tfjs-image-recognition-base';
2+
3+
import { extractorsFactory } from './extractorsFactory';
4+
import { NetParams } from './types';
5+
6+
export function extractParams(weights: Float32Array): { params: NetParams, paramMappings: ParamMapping[] } {
7+
8+
const paramMappings: ParamMapping[] = []
9+
10+
const {
11+
extractWeights,
12+
getRemainingWeights
13+
} = extractWeightsFactory(weights)
14+
15+
const {
16+
extractDenseBlock4Params
17+
} = extractorsFactory(extractWeights, paramMappings)
18+
19+
const dense0 = extractDenseBlock4Params(3, 32, 'dense0', true)
20+
const dense1 = extractDenseBlock4Params(32, 64, 'dense1')
21+
const dense2 = extractDenseBlock4Params(64, 128, 'dense2')
22+
const dense3 = extractDenseBlock4Params(128, 256, 'dense3')
23+
24+
if (getRemainingWeights().length !== 0) {
25+
throw new Error(`weights remaing after extract: ${getRemainingWeights().length}`)
26+
}
27+
28+
return {
29+
paramMappings,
30+
params: { dense0, dense1, dense2, dense3 }
31+
}
32+
}

‎src/faceFeatureExtractor/extractParamsFromWeigthMap.ts

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
import * as tf from '@tensorflow/tfjs-core';
2+
import { disposeUnusedWeightTensors, ParamMapping } from 'tfjs-image-recognition-base';
3+
4+
import { loadParamsFactory } from './loadParamsFactory';
5+
import { NetParams } from './types';
6+
7+
export function extractParamsFromWeigthMap(
8+
weightMap: tf.NamedTensorMap
9+
): { params: NetParams, paramMappings: ParamMapping[] } {
10+
11+
const paramMappings: ParamMapping[] = []
12+
13+
const {
14+
extractDenseBlock4Params
15+
} = loadParamsFactory(weightMap, paramMappings)
16+
17+
const params = {
18+
dense0: extractDenseBlock4Params('dense0', true),
19+
dense1: extractDenseBlock4Params('dense1'),
20+
dense2: extractDenseBlock4Params('dense2'),
21+
dense3: extractDenseBlock4Params('dense3')
22+
}
23+
24+
disposeUnusedWeightTensors(weightMap, paramMappings)
25+
26+
return { params, paramMappings }
27+
}

‎src/faceLandmarkNet/extractParamsFromWeigthMapTiny.ts renamed to ‎src/faceFeatureExtractor/extractParamsFromWeigthMapTiny.ts

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,15 +11,13 @@ export function extractParamsFromWeigthMapTiny(
1111
const paramMappings: ParamMapping[] = []
1212

1313
const {
14-
extractDenseBlock3Params,
15-
extractFcParams
14+
extractDenseBlock3Params
1615
} = loadParamsFactory(weightMap, paramMappings)
1716

1817
const params = {
1918
dense0: extractDenseBlock3Params('dense0', true),
2019
dense1: extractDenseBlock3Params('dense1'),
21-
dense2: extractDenseBlock3Params('dense2'),
22-
fc: extractFcParams('fc')
20+
dense2: extractDenseBlock3Params('dense2')
2321
}
2422

2523
disposeUnusedWeightTensors(weightMap, paramMappings)

‎src/faceLandmarkNet/extractParamsTiny.ts renamed to ‎src/faceFeatureExtractor/extractParamsTiny.ts

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,21 +13,19 @@ export function extractParamsTiny(weights: Float32Array): { params: TinyNetParam
1313
} = extractWeightsFactory(weights)
1414

1515
const {
16-
extractDenseBlock3Params,
17-
extractFCParams
16+
extractDenseBlock3Params
1817
} = extractorsFactory(extractWeights, paramMappings)
1918

2019
const dense0 = extractDenseBlock3Params(3, 32, 'dense0', true)
2120
const dense1 = extractDenseBlock3Params(32, 64, 'dense1')
2221
const dense2 = extractDenseBlock3Params(64, 128, 'dense2')
23-
const fc = extractFCParams(128, 136, 'fc')
2422

2523
if (getRemainingWeights().length !== 0) {
2624
throw new Error(`weights remaing after extract: ${getRemainingWeights().length}`)
2725
}
2826

2927
return {
3028
paramMappings,
31-
params: { dense0, dense1, dense2, fc }
29+
params: { dense0, dense1, dense2 }
3230
}
3331
}

‎src/faceLandmarkNet/extractorsFactory.ts renamed to ‎src/faceFeatureExtractor/extractorsFactory.ts

Lines changed: 1 addition & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -24,21 +24,6 @@ export function extractorsFactory(extractWeights: ExtractWeightsFunction, paramM
2424
)
2525
}
2626

27-
function extractFCParams(channelsIn: number, channelsOut: number, mappedPrefix: string): FCParams {
28-
const weights = tf.tensor2d(extractWeights(channelsIn * channelsOut), [channelsIn, channelsOut])
29-
const bias = tf.tensor1d(extractWeights(channelsOut))
30-
31-
paramMappings.push(
32-
{ paramPath: `${mappedPrefix}/weights` },
33-
{ paramPath: `${mappedPrefix}/bias` }
34-
)
35-
36-
return {
37-
weights,
38-
bias
39-
}
40-
}
41-
4227
const extractConvParams = extractConvParamsFactory(extractWeights, paramMappings)
4328

4429
function extractDenseBlock3Params(channelsIn: number, channelsOut: number, mappedPrefix: string, isFirstLayer: boolean = false): DenseBlock3Params {
@@ -62,8 +47,7 @@ export function extractorsFactory(extractWeights: ExtractWeightsFunction, paramM
6247

6348
return {
6449
extractDenseBlock3Params,
65-
extractDenseBlock4Params,
66-
extractFCParams
50+
extractDenseBlock4Params
6751
}
6852

6953
}

‎src/faceFeatureExtractor/index.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
export * from './FaceFeatureExtractor';
2+
export * from './TinyFaceFeatureExtractor';

‎src/faceLandmarkNet/loadParamsFactory.ts renamed to ‎src/faceFeatureExtractor/loadParamsFactory.ts

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import * as tf from '@tensorflow/tfjs-core';
22
import { extractWeightEntryFactory, ParamMapping } from 'tfjs-image-recognition-base';
3-
import { ConvParams, FCParams, SeparableConvParams } from 'tfjs-tiny-yolov2';
3+
import { ConvParams, SeparableConvParams } from 'tfjs-tiny-yolov2';
44

55
import { DenseBlock3Params, DenseBlock4Params } from './types';
66

@@ -48,16 +48,8 @@ export function loadParamsFactory(weightMap: any, paramMappings: ParamMapping[])
4848
return { conv0, conv1, conv2, conv3 }
4949
}
5050

51-
function extractFcParams(prefix: string): FCParams {
52-
const weights = extractWeightEntry<tf.Tensor2D>(`${prefix}/weights`, 2)
53-
const bias = extractWeightEntry<tf.Tensor1D>(`${prefix}/bias`, 1)
54-
55-
return { weights, bias }
56-
}
57-
5851
return {
5952
extractDenseBlock3Params,
60-
extractDenseBlock4Params,
61-
extractFcParams
53+
extractDenseBlock4Params
6254
}
6355
}

‎src/faceFeatureExtractor/types.ts

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
import * as tf from '@tensorflow/tfjs-core';
2+
import { ConvParams, SeparableConvParams } from 'tfjs-tiny-yolov2';
3+
4+
export type ConvWithBatchNormParams = BatchNormParams & {
5+
filter: tf.Tensor4D
6+
}
7+
8+
export type BatchNormParams = {
9+
mean: tf.Tensor1D
10+
variance: tf.Tensor1D
11+
scale: tf.Tensor1D
12+
offset: tf.Tensor1D
13+
}
14+
15+
export type SeparableConvWithBatchNormParams = {
16+
depthwise: ConvWithBatchNormParams
17+
pointwise: ConvWithBatchNormParams
18+
}
19+
20+
export type DenseBlock3Params = {
21+
conv0: SeparableConvParams | ConvParams
22+
conv1: SeparableConvParams
23+
conv2: SeparableConvParams
24+
}
25+
26+
export type DenseBlock4Params = DenseBlock3Params & {
27+
conv3: SeparableConvParams
28+
}
29+
30+
export type TinyNetParams = {
31+
dense0: DenseBlock3Params
32+
dense1: DenseBlock3Params
33+
dense2: DenseBlock3Params
34+
}
35+
36+
export type NetParams = {
37+
dense0: DenseBlock4Params
38+
dense1: DenseBlock4Params
39+
dense2: DenseBlock4Params
40+
dense3: DenseBlock4Params
41+
}
42+

0 commit comments

Comments
 (0)