Skip to content

Commit fdb7065

Browse files
justadudewhohacksjustadudewhohacks
committed
fixed postprocessing of 68 point face landmarks + init a mobilenet for 68 point face landmarks
1 parent fdfedbe commit fdb7065

15 files changed

+434
-140
lines changed

‎src/allFacesFactory.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import { FaceLandmarks68 } from './classes/FaceLandmarks68';
77
import { FullFaceDescription } from './classes/FullFaceDescription';
88
import { extractFaceTensors } from './dom';
99
import { FaceDetectionNet } from './faceDetectionNet/FaceDetectionNet';
10-
import { FaceLandmarkNet } from './faceLandmarkNet/FaceLandmarkNet';
10+
import { FaceLandmarkNet } from './faceLandmarkNet';
1111
import { FaceRecognitionNet } from './faceRecognitionNet/FaceRecognitionNet';
1212
import { Mtcnn } from './mtcnn/Mtcnn';
1313
import { MtcnnForwardParams } from './mtcnn/types';

‎src/faceLandmarkNet/FaceLandmark68MobileNet.ts

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
import * as tf from '@tensorflow/tfjs-core';
2+
import { NetInput } from 'tfjs-image-recognition-base';
3+
4+
import { dephtwiseSeparableConv } from './dephtwiseSeparableConv';
5+
import { extractMobilenetParams } from './extractMobilenetParams';
6+
import { FaceLandmark68NetBase } from './FaceLandmark68NetBase';
7+
import { fullyConnectedLayer } from './fullyConnectedLayer';
8+
import { MobilenetParams } from './types';
9+
10+
export class FaceLandmark68MobileNet extends FaceLandmark68NetBase<MobilenetParams> {
11+
12+
constructor() {
13+
super('FaceLandmark68MobileNet')
14+
}
15+
16+
public runNet(input: NetInput): tf.Tensor2D {
17+
18+
const { params } = this
19+
20+
if (!params) {
21+
throw new Error('FaceLandmark68MobileNet - load model before inference')
22+
}
23+
24+
return tf.tidy(() => {
25+
const batchTensor = input.toBatchTensor(112, true)
26+
27+
let out = dephtwiseSeparableConv(batchTensor, params.conv0, [2, 2])
28+
out = dephtwiseSeparableConv(batchTensor, params.conv1, [1, 1])
29+
out = dephtwiseSeparableConv(batchTensor, params.conv2, [2, 2])
30+
out = dephtwiseSeparableConv(batchTensor, params.conv3, [1, 1])
31+
out = dephtwiseSeparableConv(batchTensor, params.conv4, [2, 2])
32+
out = dephtwiseSeparableConv(batchTensor, params.conv5, [1, 1])
33+
out = dephtwiseSeparableConv(batchTensor, params.conv6, [2, 2])
34+
out = dephtwiseSeparableConv(batchTensor, params.conv7, [1, 1])
35+
out = tf.avgPool(out, [7, 7], [1, 1], 'valid')
36+
37+
return fullyConnectedLayer(out.as2D(out.shape[0], -1), params.fc)
38+
})
39+
}
40+
/*
41+
protected loadQuantizedParams(uri: string | undefined) {
42+
return loadQuantizedParams(uri)
43+
}
44+
*/
45+
protected extractParams(weights: Float32Array) {
46+
return extractMobilenetParams(weights)
47+
}
48+
}

‎src/faceLandmarkNet/FaceLandmark68Net.ts

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
import * as tf from '@tensorflow/tfjs-core';
2+
import { NetInput } from 'tfjs-image-recognition-base';
3+
import { convLayer, ConvParams } from 'tfjs-tiny-yolov2';
4+
5+
import { extractParams } from './extractParams';
6+
import { FaceLandmark68NetBase } from './FaceLandmark68NetBase';
7+
import { fullyConnectedLayer } from './fullyConnectedLayer';
8+
import { loadQuantizedParams } from './loadQuantizedParams';
9+
import { NetParams } from './types';
10+
11+
function conv(x: tf.Tensor4D, params: ConvParams): tf.Tensor4D {
12+
return convLayer(x, params, 'valid', true)
13+
}
14+
15+
function maxPool(x: tf.Tensor4D, strides: [number, number] = [2, 2]): tf.Tensor4D {
16+
return tf.maxPool(x, [2, 2], strides, 'valid')
17+
}
18+
19+
export class FaceLandmark68Net extends FaceLandmark68NetBase<NetParams> {
20+
21+
constructor() {
22+
super('FaceLandmark68Net')
23+
}
24+
25+
public runNet(input: NetInput): tf.Tensor2D {
26+
27+
const { params } = this
28+
29+
if (!params) {
30+
throw new Error('FaceLandmark68Net - load model before inference')
31+
}
32+
33+
return tf.tidy(() => {
34+
const batchTensor = input.toBatchTensor(128, true)
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)
51+
})
52+
}
53+
54+
protected loadQuantizedParams(uri: string | undefined) {
55+
return loadQuantizedParams(uri)
56+
}
57+
58+
protected extractParams(weights: Float32Array) {
59+
return extractParams(weights)
60+
}
61+
}

‎src/faceLandmarkNet/FaceLandmark68NetBase.ts

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
import * as tf from '@tensorflow/tfjs-core';
2+
import { isEven, NetInput, NeuralNetwork, Point, TNetInput, toNetInput, Dimensions } from 'tfjs-image-recognition-base';
3+
4+
import { FaceLandmarks68 } from '../classes/FaceLandmarks68';
5+
6+
export class FaceLandmark68NetBase<NetParams> extends NeuralNetwork<NetParams> {
7+
8+
// TODO: make super.name protected
9+
private __name: string
10+
11+
constructor(_name: string) {
12+
super(_name)
13+
this.__name = _name
14+
}
15+
16+
public runNet(_: NetInput): tf.Tensor2D {
17+
throw new Error(`${this.__name} - runNet not implemented`)
18+
}
19+
20+
public postProcess(output: tf.Tensor2D, inputSize: number, originalDimensions: Dimensions[]): tf.Tensor2D {
21+
22+
const inputDimensions = originalDimensions.map(({ width, height }) => {
23+
const scale = inputSize / Math.max(height, width)
24+
return {
25+
width: width * scale,
26+
height: height * scale
27+
}
28+
})
29+
30+
const batchSize = inputDimensions.length
31+
32+
return tf.tidy(() => {
33+
const createInterleavedTensor = (fillX: number, fillY: number) =>
34+
tf.stack([
35+
tf.fill([68], fillX),
36+
tf.fill([68], fillY)
37+
], 1).as2D(1, 136).as1D()
38+
39+
const getPadding = (batchIdx: number, cond: (w: number, h: number) => boolean): number => {
40+
const { width, height } = inputDimensions[batchIdx]
41+
return cond(width, height) ? Math.abs(width - height) / 2 : 0
42+
}
43+
const getPaddingX = (batchIdx: number) => getPadding(batchIdx, (w, h) => w < h)
44+
const getPaddingY = (batchIdx: number) => getPadding(batchIdx, (w, h) => h < w)
45+
46+
const landmarkTensors = output
47+
.mul(tf.fill([batchSize, 136], inputSize))
48+
.sub(tf.stack(Array.from(Array(batchSize), (_, batchIdx) =>
49+
createInterleavedTensor(
50+
getPaddingX(batchIdx),
51+
getPaddingY(batchIdx)
52+
)
53+
)))
54+
.div(tf.stack(Array.from(Array(batchSize), (_, batchIdx) =>
55+
createInterleavedTensor(
56+
inputDimensions[batchIdx].width,
57+
inputDimensions[batchIdx].height
58+
)
59+
)))
60+
61+
return landmarkTensors as tf.Tensor2D
62+
})
63+
}
64+
65+
public forwardInput(input: NetInput): tf.Tensor2D {
66+
return tf.tidy(() => {
67+
const out = this.runNet(input)
68+
return this.postProcess(
69+
out,
70+
input.inputSize,
71+
input.inputDimensions.map(([height, width]) => ({ height, width }))
72+
)
73+
})
74+
}
75+
76+
public async forward(input: TNetInput): Promise<tf.Tensor2D> {
77+
return this.forwardInput(await toNetInput(input, true))
78+
}
79+
80+
public async detectLandmarks(input: TNetInput): Promise<FaceLandmarks68 | FaceLandmarks68[]> {
81+
const netInput = await toNetInput(input, true)
82+
83+
const landmarkTensors = tf.tidy(
84+
() => tf.unstack(this.forwardInput(netInput))
85+
)
86+
87+
const landmarksForBatch = await Promise.all(landmarkTensors.map(
88+
async (landmarkTensor, batchIdx) => {
89+
const landmarksArray = Array.from(await landmarkTensor.data())
90+
const xCoords = landmarksArray.filter((_, i) => isEven(i))
91+
const yCoords = landmarksArray.filter((_, i) => !isEven(i))
92+
93+
return new FaceLandmarks68(
94+
Array(68).fill(0).map((_, i) => new Point(xCoords[i], yCoords[i])),
95+
{
96+
height: netInput.getInputHeight(batchIdx),
97+
width : netInput.getInputWidth(batchIdx),
98+
}
99+
)
100+
}
101+
))
102+
103+
landmarkTensors.forEach(t => t.dispose())
104+
105+
return netInput.isBatchInput
106+
? landmarksForBatch
107+
: landmarksForBatch[0]
108+
}
109+
}

‎src/faceLandmarkNet/FaceLandmarkNet.ts

Lines changed: 0 additions & 127 deletions
This file was deleted.

‎src/faceLandmarkNet/dephtwiseSeparableConv.ts

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

0 commit comments

Comments
 (0)