Skip to content

Commit ef3ecfb

Browse files
justadudewhohacksjustadudewhohacks
authored
Merge pull request justadudewhohacks#97 from justadudewhohacks/train-landmark-models
trained two new 68 point face landmark models from scratch, which are more accurate, much faster and much smaller in size
2 parents 8601af9 + b3d2589 commit ef3ecfb

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

54 files changed

+1961
-1609
lines changed

‎README.md

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,8 @@ Table of Contents:
4343

4444
### Face Landmarks
4545

46-
![preview_face_landmarks_boxes](https://user-images.githubusercontent.com/31125521/41507933-65f9b642-723c-11e8-8f4e-aab13303e7ff.jpg)
46+
![face_landmarks_boxes_1](https://user-images.githubusercontent.com/31125521/46063403-fff9f480-c16c-11e8-900f-e4b7a3828d1d.jpg)
47+
![face_landmarks_boxes_2](https://user-images.githubusercontent.com/31125521/46063404-00928b00-c16d-11e8-8f29-e9c50afd2bc8.jpg)
4748

4849
![preview_face_landmarks](https://user-images.githubusercontent.com/31125521/41507950-e121b05e-723c-11e8-89f2-d8f9348a8e86.png)
4950

@@ -57,10 +58,6 @@ Table of Contents:
5758

5859
![mtcnn-preview](https://user-images.githubusercontent.com/31125521/42756818-0a41edaa-88fe-11e8-9033-8cd141b0fa09.gif)
5960

60-
### Face Alignment
61-
62-
![preview_face_alignment](https://user-images.githubusercontent.com/31125521/41526994-1a690818-72e6-11e8-8f3c-d2cf31fe517b.jpg)
63-
6461
<a name="running-the-examples"></a>
6562

6663
## Running the Examples
@@ -89,7 +86,7 @@ The face detection model has been trained on the [WIDERFACE dataset](http://mmla
8986

9087
### Face Detection - Tiny Yolo v2
9188

92-
The Tiny Yolo v2 implementation is a very performant face detector, which can easily adapt to different input image sizes, thus can be used as an alternative to SSD Mobilenet v1 to trade off accuracy for performance (inference time). In general the models ability to locate smaller face bounding boxes is not as accurate as SSD Mobilenet v1.
89+
The Tiny Yolo v2 implementation is a very performant face detector, which can easily adapt to different input image sizes, thus can be used as an alternative to SSD Mobilenet v1 to trade off accuracy for performance (inference time). In general the models ability to locate smaller face bounding boxes is not as accurate as SSD Mobilenet v1.
9390

9491
The face detector has been trained on a custom dataset of ~10K images labeled with bounding boxes and uses depthwise separable convolutions instead of regular convolutions, which ensures very fast inference and allows to have a quantized model size of only 1.7MB making the model extremely mobile and web friendly. Thus, the Tiny Yolo v2 face detector should be your GO-TO face detector on mobile devices.
9592

@@ -113,9 +110,7 @@ The neural net is equivalent to the **FaceRecognizerNet** used in [face-recognit
113110

114111
### 68 Point Face Landmark Detection
115112

116-
This package implements a CNN to detect the 68 point face landmarks for a given face image.
117-
118-
The model has been trained on a variety of public datasets and the model weights are provided by [yinguobing](https://github.com/yinguobing) in [this](https://github.com/yinguobing/head-pose-estimation) repo.
113+
This package implements a very lightweight and fast, yet accurate 68 point face landmark detector. The default model has a size of only 350kb and the tiny model is only 80kb. Both models employ the ideas of depthwise separable convolutions as well as densely connected blocks. The models have been trained on a dataset of ~35k face images labeled with 68 face landmark points.
119114

120115
<a name="usage"></a>
121116

@@ -145,6 +140,7 @@ Assuming the models reside in **public/models**:
145140
await faceapi.loadFaceDetectionModel('/models')
146141
// accordingly for the other models:
147142
// await faceapi.loadFaceLandmarkModel('/models')
143+
// await faceapi.loadFaceLandmarkTinyModel('/models')
148144
// await faceapi.loadFaceRecognitionModel('/models')
149145
// await faceapi.loadMtcnnModel('/models')
150146
// await faceapi.loadTinyYolov2Model('/models')
@@ -155,19 +151,18 @@ As an alternative, you can also create instance of the neural nets:
155151
``` javascript
156152
const net = new faceapi.FaceDetectionNet()
157153
// accordingly for the other models:
158-
// const net = new faceapi.FaceLandmarkNet()
154+
// const net = new faceapi.FaceLandmark68Net()
155+
// const net = new faceapi.FaceLandmark68TinyNet()
159156
// const net = new faceapi.FaceRecognitionNet()
160157
// const net = new faceapi.Mtcnn()
161158
// const net = new faceapi.TinyYolov2()
162159

163160
await net.load('/models/face_detection_model-weights_manifest.json')
164161
// await net.load('/models/face_landmark_68_model-weights_manifest.json')
162+
// await net.load('/models/face_landmark_68_tiny_model-weights_manifest.json')
165163
// await net.load('/models/face_recognition_model-weights_manifest.json')
166164
// await net.load('/models/mtcnn_model-weights_manifest.json')
167165
// await net.load('/models/tiny_yolov2_separable_conv_model-weights_manifest.json')
168-
169-
// or simply load all models
170-
await net.load('/models')
171166
```
172167

173168
Using instances, you can also load the weights as a Float32Array (in case you want to use the uncompressed models):

‎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: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
import * as tf from '@tensorflow/tfjs-core';
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';
5+
6+
import { depthwiseSeparableConv } from './depthwiseSeparableConv';
7+
import { extractParamsTiny } from './extractParamsTiny';
8+
import { FaceLandmark68NetBase } from './FaceLandmark68NetBase';
9+
import { fullyConnectedLayer } from './fullyConnectedLayer';
10+
import { loadQuantizedParamsTiny } from './loadQuantizedParamsTiny';
11+
import { DenseBlock3Params, TinyNetParams } from './types';
12+
13+
function denseBlock(
14+
x: tf.Tensor4D,
15+
denseBlockParams: DenseBlock3Params,
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+
return tf.relu(tf.add(out1, tf.add(out2, out3))) as tf.Tensor4D
33+
})
34+
}
35+
36+
export class FaceLandmark68TinyNet extends FaceLandmark68NetBase<TinyNetParams> {
37+
38+
constructor() {
39+
super('FaceLandmark68TinyNet')
40+
}
41+
42+
public runNet(input: NetInput): tf.Tensor2D {
43+
44+
const { params } = this
45+
46+
if (!params) {
47+
throw new Error('FaceLandmark68TinyNet - 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 = tf.avgPool(out, [14, 14], [2, 2], 'valid')
59+
60+
return fullyConnectedLayer(out.as2D(out.shape[0], -1), params.fc)
61+
})
62+
}
63+
64+
protected loadQuantizedParams(uri: string | undefined) {
65+
return loadQuantizedParamsTiny(uri)
66+
}
67+
68+
protected extractParams(weights: Float32Array) {
69+
return extractParamsTiny(weights)
70+
}
71+
}

‎src/faceLandmarkNet/depthwiseSeparableConv.ts

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
import * as tf from '@tensorflow/tfjs-core';
2+
import { SeparableConvParams } from 'tfjs-tiny-yolov2/build/tinyYolov2/types';
3+
4+
export function depthwiseSeparableConv(
5+
x: tf.Tensor4D,
6+
params: SeparableConvParams,
7+
stride: [number, number]
8+
): tf.Tensor4D {
9+
return tf.tidy(() => {
10+
let out = tf.separableConv2d(x, params.depthwise_filter, params.pointwise_filter, stride, 'same')
11+
out = tf.add(out, params.bias)
12+
return out
13+
})
14+
}

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

0 commit comments

Comments
 (0)