Skip to content

Commit 4d3706f

Browse files
justadudewhohacksjustadudewhohacks
committed
retrained landmark 68 and tiny landmark 68 point models
1 parent 0bfaa03 commit 4d3706f

22 files changed

+299
-285
lines changed

‎examples/public/styles.css

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,12 @@
6767
left: 0;
6868
}
6969

70+
.overlay {
71+
position: absolute;
72+
top: 0;
73+
left: 0;
74+
}
75+
7076
#facesContainer canvas {
7177
margin: 10px;
7278
}

‎src/faceLandmarkNet/FaceLandmark68Net.ts

Lines changed: 42 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,60 +1,75 @@
11
import * as tf from '@tensorflow/tfjs-core';
2-
import { NetInput } from 'tfjs-image-recognition-base';
3-
import { convLayer, ConvParams } from 'tfjs-tiny-yolov2';
2+
import { NetInput, normalize } from 'tfjs-image-recognition-base';
3+
import { ConvParams } from 'tfjs-tiny-yolov2';
4+
import { SeparableConvParams } from 'tfjs-tiny-yolov2/build/tinyYolov2/types';
45

6+
import { depthwiseSeparableConv } from './depthwiseSeparableConv';
57
import { extractParams } from './extractParams';
68
import { FaceLandmark68NetBase } from './FaceLandmark68NetBase';
79
import { fullyConnectedLayer } from './fullyConnectedLayer';
810
import { loadQuantizedParams } from './loadQuantizedParams';
9-
import { NetParams } from './types';
11+
import { DenseBlock4Params, NetParams } from './types';
1012

11-
function conv(x: tf.Tensor4D, params: ConvParams): tf.Tensor4D {
12-
return convLayer(x, params, 'valid', true)
13-
}
13+
function denseBlock(
14+
x: tf.Tensor4D,
15+
denseBlockParams: DenseBlock4Params,
16+
isFirstLayer: boolean = false
17+
): tf.Tensor4D {
18+
return tf.tidy(() => {
19+
const out1 = tf.relu(
20+
isFirstLayer
21+
? tf.add(
22+
tf.conv2d(x, (denseBlockParams.conv0 as ConvParams).filters, [2, 2], 'same'),
23+
denseBlockParams.conv0.bias
24+
)
25+
: depthwiseSeparableConv(x, denseBlockParams.conv0 as SeparableConvParams, [2, 2])
26+
) as tf.Tensor4D
27+
const out2 = depthwiseSeparableConv(out1, denseBlockParams.conv1, [1, 1])
28+
29+
const in3 = tf.relu(tf.add(out1, out2)) as tf.Tensor4D
30+
const out3 = depthwiseSeparableConv(in3, denseBlockParams.conv2, [1, 1])
31+
32+
const in4 = tf.relu(tf.add(out1, tf.add(out2, out3))) as tf.Tensor4D
33+
const out4 = depthwiseSeparableConv(in4, denseBlockParams.conv3, [1, 1])
1434

15-
function maxPool(x: tf.Tensor4D, strides: [number, number] = [2, 2]): tf.Tensor4D {
16-
return tf.maxPool(x, [2, 2], strides, 'valid')
35+
return tf.relu(tf.add(out1, tf.add(out2, tf.add(out3, out4)))) as tf.Tensor4D
36+
})
1737
}
1838

1939
export class FaceLandmark68Net extends FaceLandmark68NetBase<NetParams> {
2040

2141
constructor() {
22-
super('FaceLandmark68Net')
42+
super('FaceLandmark68LargeNet')
2343
}
2444

2545
public runNet(input: NetInput): tf.Tensor2D {
2646

2747
const { params } = this
2848

2949
if (!params) {
30-
throw new Error('FaceLandmark68Net - load model before inference')
50+
throw new Error('FaceLandmark68LargeNet - load model before inference')
3151
}
3252

3353
return tf.tidy(() => {
34-
const batchTensor = input.toBatchTensor(128, true).toFloat() as tf.Tensor4D
35-
36-
let out = conv(batchTensor, params.conv0)
37-
out = maxPool(out)
38-
out = conv(out, params.conv1)
39-
out = conv(out, params.conv2)
40-
out = maxPool(out)
41-
out = conv(out, params.conv3)
42-
out = conv(out, params.conv4)
43-
out = maxPool(out)
44-
out = conv(out, params.conv5)
45-
out = conv(out, params.conv6)
46-
out = maxPool(out, [1, 1])
47-
out = conv(out, params.conv7)
48-
const fc0 = tf.relu(fullyConnectedLayer(out.as2D(out.shape[0], -1), params.fc0))
49-
50-
return fullyConnectedLayer(fc0, params.fc1)
54+
const batchTensor = input.toBatchTensor(112, true)
55+
const meanRgb = [122.782, 117.001, 104.298]
56+
const normalized = normalize(batchTensor, meanRgb).div(tf.scalar(255)) as tf.Tensor4D
57+
58+
let out = denseBlock(normalized, params.dense0, true)
59+
out = denseBlock(out, params.dense1)
60+
out = denseBlock(out, params.dense2)
61+
out = denseBlock(out, params.dense3)
62+
out = tf.avgPool(out, [7, 7], [2, 2], 'valid')
63+
64+
return fullyConnectedLayer(out.as2D(out.shape[0], -1), params.fc)
5165
})
5266
}
5367

5468
protected loadQuantizedParams(uri: string | undefined) {
5569
return loadQuantizedParams(uri)
5670
}
5771

72+
5873
protected extractParams(weights: Float32Array) {
5974
return extractParams(weights)
6075
}

‎src/faceLandmarkNet/FaceLandmark68TinyNet.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,11 @@ import { extractParamsTiny } from './extractParamsTiny';
88
import { FaceLandmark68NetBase } from './FaceLandmark68NetBase';
99
import { fullyConnectedLayer } from './fullyConnectedLayer';
1010
import { loadQuantizedParamsTiny } from './loadQuantizedParamsTiny';
11-
import { DenseBlockParams, TinyNetParams } from './types';
11+
import { DenseBlock3Params, TinyNetParams } from './types';
1212

1313
function denseBlock(
1414
x: tf.Tensor4D,
15-
denseBlockParams: DenseBlockParams,
15+
denseBlockParams: DenseBlock3Params,
1616
isFirstLayer: boolean = false
1717
): tf.Tensor4D {
1818
return tf.tidy(() => {

‎src/faceLandmarkNet/extractParams.ts

Lines changed: 11 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { extractWeightsFactory, ParamMapping } from 'tfjs-image-recognition-base';
2-
import { extractConvParamsFactory, extractFCParamsFactory } from 'tfjs-tiny-yolov2';
32

3+
import { extractorsFactory } from './extractorsFactory';
44
import { NetParams } from './types';
55

66
export function extractParams(weights: Float32Array): { params: NetParams, paramMappings: ParamMapping[] } {
@@ -12,37 +12,23 @@ export function extractParams(weights: Float32Array): { params: NetParams, param
1212
getRemainingWeights
1313
} = extractWeightsFactory(weights)
1414

15-
const extractConvParams = extractConvParamsFactory(extractWeights, paramMappings)
16-
const extractFCParams = extractFCParamsFactory(extractWeights, paramMappings)
15+
const {
16+
extractDenseBlock4Params,
17+
extractFCParams
18+
} = extractorsFactory(extractWeights, paramMappings)
1719

18-
const conv0 = extractConvParams(3, 32, 3, 'conv0')
19-
const conv1 = extractConvParams(32, 64, 3, 'conv1')
20-
const conv2 = extractConvParams(64, 64, 3, 'conv2')
21-
const conv3 = extractConvParams(64, 64, 3, 'conv3')
22-
const conv4 = extractConvParams(64, 64, 3, 'conv4')
23-
const conv5 = extractConvParams(64, 128, 3, 'conv5')
24-
const conv6 = extractConvParams(128, 128, 3, 'conv6')
25-
const conv7 = extractConvParams(128, 256, 3, 'conv7')
26-
const fc0 = extractFCParams(6400, 1024, 'fc0')
27-
const fc1 = extractFCParams(1024, 136, 'fc1')
20+
const dense0 = extractDenseBlock4Params(3, 32, 'dense0', true)
21+
const dense1 = extractDenseBlock4Params(32, 64, 'dense1')
22+
const dense2 = extractDenseBlock4Params(64, 128, 'dense2')
23+
const dense3 = extractDenseBlock4Params(128, 256, 'dense3')
24+
const fc = extractFCParams(256, 136, 'fc')
2825

2926
if (getRemainingWeights().length !== 0) {
3027
throw new Error(`weights remaing after extract: ${getRemainingWeights().length}`)
3128
}
3229

3330
return {
3431
paramMappings,
35-
params: {
36-
conv0,
37-
conv1,
38-
conv2,
39-
conv3,
40-
conv4,
41-
conv5,
42-
conv6,
43-
conv7,
44-
fc0,
45-
fc1
46-
}
32+
params: { dense0, dense1, dense2, dense3, fc }
4733
}
4834
}

‎src/faceLandmarkNet/extractParamsTiny.ts

Lines changed: 7 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -1,51 +1,7 @@
1-
import * as tf from '@tensorflow/tfjs-core';
2-
import { extractWeightsFactory, ExtractWeightsFunction, ParamMapping } from 'tfjs-image-recognition-base';
3-
import { extractConvParamsFactory, FCParams } from 'tfjs-tiny-yolov2';
4-
import { SeparableConvParams } from 'tfjs-tiny-yolov2/build/tinyYolov2/types';
1+
import { extractWeightsFactory, ParamMapping } from 'tfjs-image-recognition-base';
52

6-
import { DenseBlockParams, TinyNetParams } from './types';
7-
8-
function extractorsFactory(extractWeights: ExtractWeightsFunction, paramMappings: ParamMapping[]) {
9-
10-
function extractSeparableConvParams(channelsIn: number, channelsOut: number, mappedPrefix: string): SeparableConvParams {
11-
const depthwise_filter = tf.tensor4d(extractWeights(3 * 3 * channelsIn), [3, 3, channelsIn, 1])
12-
const pointwise_filter = tf.tensor4d(extractWeights(channelsIn * channelsOut), [1, 1, channelsIn, channelsOut])
13-
const bias = tf.tensor1d(extractWeights(channelsOut))
14-
15-
paramMappings.push(
16-
{ paramPath: `${mappedPrefix}/depthwise_filter` },
17-
{ paramPath: `${mappedPrefix}/pointwise_filter` },
18-
{ paramPath: `${mappedPrefix}/bias` }
19-
)
20-
21-
return new SeparableConvParams(
22-
depthwise_filter,
23-
pointwise_filter,
24-
bias
25-
)
26-
}
27-
28-
function extractFCParams(channelsIn: number, channelsOut: number, mappedPrefix: string): FCParams {
29-
const weights = tf.tensor2d(extractWeights(channelsIn * channelsOut), [channelsIn, channelsOut])
30-
const bias = tf.tensor1d(extractWeights(channelsOut))
31-
32-
paramMappings.push(
33-
{ paramPath: `${mappedPrefix}/weights` },
34-
{ paramPath: `${mappedPrefix}/bias` }
35-
)
36-
37-
return {
38-
weights,
39-
bias
40-
}
41-
}
42-
43-
return {
44-
extractSeparableConvParams,
45-
extractFCParams
46-
}
47-
48-
}
3+
import { extractorsFactory } from './extractorsFactory';
4+
import { TinyNetParams } from './types';
495

506
export function extractParamsTiny(weights: Float32Array): { params: TinyNetParams, paramMappings: ParamMapping[] } {
517

@@ -57,26 +13,13 @@ export function extractParamsTiny(weights: Float32Array): { params: TinyNetParam
5713
} = extractWeightsFactory(weights)
5814

5915
const {
60-
extractSeparableConvParams,
16+
extractDenseBlock3Params,
6117
extractFCParams
6218
} = extractorsFactory(extractWeights, paramMappings)
6319

64-
const extractConvParams = extractConvParamsFactory(extractWeights, paramMappings)
65-
66-
function extractDenseBlockParams(channelsIn: number, channelsOut: number, mappedPrefix: string, isFirstLayer: boolean = false): DenseBlockParams {
67-
68-
const conv0 = isFirstLayer
69-
? extractConvParams(channelsIn, channelsOut, 3, `${mappedPrefix}/conv0`)
70-
: extractSeparableConvParams(channelsIn, channelsOut, `${mappedPrefix}/conv0`)
71-
const conv1 = extractSeparableConvParams(channelsOut, channelsOut, `${mappedPrefix}/conv1`)
72-
const conv2 = extractSeparableConvParams(channelsOut, channelsOut, `${mappedPrefix}/conv2`)
73-
74-
return { conv0, conv1, conv2 }
75-
}
76-
77-
const dense0 = extractDenseBlockParams(3, 32, 'dense0', true)
78-
const dense1 = extractDenseBlockParams(32, 64, 'dense1')
79-
const dense2 = extractDenseBlockParams(64, 128, 'dense2')
20+
const dense0 = extractDenseBlock3Params(3, 32, 'dense0', true)
21+
const dense1 = extractDenseBlock3Params(32, 64, 'dense1')
22+
const dense2 = extractDenseBlock3Params(64, 128, 'dense2')
8023
const fc = extractFCParams(128, 136, 'fc')
8124

8225
if (getRemainingWeights().length !== 0) {

‎src/faceLandmarkNet/extractorsFactory.ts

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
import * as tf from '@tensorflow/tfjs-core';
2+
import { ExtractWeightsFunction, ParamMapping } from 'tfjs-image-recognition-base';
3+
import { extractConvParamsFactory, FCParams } from 'tfjs-tiny-yolov2';
4+
import { SeparableConvParams } from 'tfjs-tiny-yolov2/build/tinyYolov2/types';
5+
6+
import { DenseBlock3Params, DenseBlock4Params } from './types';
7+
8+
export function extractorsFactory(extractWeights: ExtractWeightsFunction, paramMappings: ParamMapping[]) {
9+
10+
function extractSeparableConvParams(channelsIn: number, channelsOut: number, mappedPrefix: string): SeparableConvParams {
11+
const depthwise_filter = tf.tensor4d(extractWeights(3 * 3 * channelsIn), [3, 3, channelsIn, 1])
12+
const pointwise_filter = tf.tensor4d(extractWeights(channelsIn * channelsOut), [1, 1, channelsIn, channelsOut])
13+
const bias = tf.tensor1d(extractWeights(channelsOut))
14+
15+
paramMappings.push(
16+
{ paramPath: `${mappedPrefix}/depthwise_filter` },
17+
{ paramPath: `${mappedPrefix}/pointwise_filter` },
18+
{ paramPath: `${mappedPrefix}/bias` }
19+
)
20+
21+
return new SeparableConvParams(
22+
depthwise_filter,
23+
pointwise_filter,
24+
bias
25+
)
26+
}
27+
28+
function extractFCParams(channelsIn: number, channelsOut: number, mappedPrefix: string): FCParams {
29+
const weights = tf.tensor2d(extractWeights(channelsIn * channelsOut), [channelsIn, channelsOut])
30+
const bias = tf.tensor1d(extractWeights(channelsOut))
31+
32+
paramMappings.push(
33+
{ paramPath: `${mappedPrefix}/weights` },
34+
{ paramPath: `${mappedPrefix}/bias` }
35+
)
36+
37+
return {
38+
weights,
39+
bias
40+
}
41+
}
42+
43+
const extractConvParams = extractConvParamsFactory(extractWeights, paramMappings)
44+
45+
function extractDenseBlock3Params(channelsIn: number, channelsOut: number, mappedPrefix: string, isFirstLayer: boolean = false): DenseBlock3Params {
46+
47+
const conv0 = isFirstLayer
48+
? extractConvParams(channelsIn, channelsOut, 3, `${mappedPrefix}/conv0`)
49+
: extractSeparableConvParams(channelsIn, channelsOut, `${mappedPrefix}/conv0`)
50+
const conv1 = extractSeparableConvParams(channelsOut, channelsOut, `${mappedPrefix}/conv1`)
51+
const conv2 = extractSeparableConvParams(channelsOut, channelsOut, `${mappedPrefix}/conv2`)
52+
53+
return { conv0, conv1, conv2 }
54+
}
55+
56+
function extractDenseBlock4Params(channelsIn: number, channelsOut: number, mappedPrefix: string, isFirstLayer: boolean = false): DenseBlock4Params {
57+
58+
const { conv0, conv1, conv2 } = extractDenseBlock3Params(channelsIn, channelsOut, mappedPrefix, isFirstLayer)
59+
const conv3 = extractSeparableConvParams(channelsOut, channelsOut, `${mappedPrefix}/conv3`)
60+
61+
return { conv0, conv1, conv2, conv3 }
62+
}
63+
64+
return {
65+
extractDenseBlock3Params,
66+
extractDenseBlock4Params,
67+
extractFCParams
68+
}
69+
70+
}

0 commit comments

Comments
 (0)