Skip to content

Commit 0bfaa03

Browse files
justadudewhohacksjustadudewhohacks
committed
check in retrained tiny 68 point landmark model
1 parent 7126f5c commit 0bfaa03

32 files changed

+1406
-1398
lines changed

‎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 { DenseBlockParams, TinyNetParams } from './types';
12+
13+
function denseBlock(
14+
x: tf.Tensor4D,
15+
denseBlockParams: DenseBlockParams,
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/extractParamsTiny.ts

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
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';
5+
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+
}
49+
50+
export function extractParamsTiny(weights: Float32Array): { params: TinyNetParams, paramMappings: ParamMapping[] } {
51+
52+
const paramMappings: ParamMapping[] = []
53+
54+
const {
55+
extractWeights,
56+
getRemainingWeights
57+
} = extractWeightsFactory(weights)
58+
59+
const {
60+
extractSeparableConvParams,
61+
extractFCParams
62+
} = extractorsFactory(extractWeights, paramMappings)
63+
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')
80+
const fc = extractFCParams(128, 136, 'fc')
81+
82+
if (getRemainingWeights().length !== 0) {
83+
throw new Error(`weights remaing after extract: ${getRemainingWeights().length}`)
84+
}
85+
86+
return {
87+
paramMappings,
88+
params: { dense0, dense1, dense2, fc }
89+
}
90+
}

‎src/faceLandmarkNet/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { FaceLandmark68Net } from './FaceLandmark68Net';
22

33
export * from './FaceLandmark68Net';
4+
export * from './FaceLandmark68TinyNet';
45

56
export class FaceLandmarkNet extends FaceLandmark68Net {}
67

‎src/faceLandmarkNet/loadQuantizedParamsTiny.ts

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
import * as tf from '@tensorflow/tfjs-core';
2+
import {
3+
disposeUnusedWeightTensors,
4+
extractWeightEntryFactory,
5+
loadWeightMap,
6+
ParamMapping,
7+
} from 'tfjs-image-recognition-base';
8+
import { ConvParams, FCParams } from 'tfjs-tiny-yolov2';
9+
import { SeparableConvParams } from 'tfjs-tiny-yolov2/build/tinyYolov2/types';
10+
11+
import { DenseBlockParams, TinyNetParams } from './types';
12+
13+
const DEFAULT_MODEL_NAME = 'face_landmark_68_tiny_model'
14+
15+
function extractorsFactory(weightMap: any, paramMappings: ParamMapping[]) {
16+
17+
const extractWeightEntry = extractWeightEntryFactory(weightMap, paramMappings)
18+
19+
function extractConvParams(prefix: string): ConvParams {
20+
const filters = extractWeightEntry<tf.Tensor4D>(`${prefix}/filters`, 4)
21+
const bias = extractWeightEntry<tf.Tensor1D>(`${prefix}/bias`, 1)
22+
23+
return { filters, bias }
24+
}
25+
26+
function extractSeparableConvParams(prefix: string): SeparableConvParams {
27+
const depthwise_filter = extractWeightEntry<tf.Tensor4D>(`${prefix}/depthwise_filter`, 4)
28+
const pointwise_filter = extractWeightEntry<tf.Tensor4D>(`${prefix}/pointwise_filter`, 4)
29+
const bias = extractWeightEntry<tf.Tensor1D>(`${prefix}/bias`, 1)
30+
31+
return new SeparableConvParams(
32+
depthwise_filter,
33+
pointwise_filter,
34+
bias
35+
)
36+
}
37+
38+
function extractDenseBlockParams(prefix: string, isFirstLayer: boolean = false): DenseBlockParams {
39+
const conv0 = isFirstLayer
40+
? extractConvParams(`${prefix}/conv0`)
41+
: extractSeparableConvParams(`${prefix}/conv0`)
42+
const conv1 = extractSeparableConvParams(`${prefix}/conv1`)
43+
const conv2 = extractSeparableConvParams(`${prefix}/conv2`)
44+
45+
return { conv0, conv1, conv2 }
46+
}
47+
48+
function extractFcParams(prefix: string): FCParams {
49+
const weights = extractWeightEntry<tf.Tensor2D>(`${prefix}/weights`, 2)
50+
const bias = extractWeightEntry<tf.Tensor1D>(`${prefix}/bias`, 1)
51+
52+
return { weights, bias }
53+
}
54+
55+
return {
56+
extractDenseBlockParams,
57+
extractFcParams
58+
}
59+
}
60+
61+
export async function loadQuantizedParamsTiny(
62+
uri: string | undefined
63+
): Promise<{ params: TinyNetParams, paramMappings: ParamMapping[] }> {
64+
65+
const weightMap = await loadWeightMap(uri, DEFAULT_MODEL_NAME)
66+
const paramMappings: ParamMapping[] = []
67+
68+
const {
69+
extractDenseBlockParams,
70+
extractFcParams
71+
} = extractorsFactory(weightMap, paramMappings)
72+
73+
const params = {
74+
dense0: extractDenseBlockParams('dense0', true),
75+
dense1: extractDenseBlockParams('dense1'),
76+
dense2: extractDenseBlockParams('dense2'),
77+
fc: extractFcParams('fc')
78+
}
79+
80+
disposeUnusedWeightTensors(weightMap, paramMappings)
81+
82+
return { params, paramMappings }
83+
}

‎src/faceLandmarkNet/types.ts

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1+
import * as tf from '@tensorflow/tfjs-core';
12
import { ConvParams, FCParams } from 'tfjs-tiny-yolov2';
3+
import { SeparableConvParams } from 'tfjs-tiny-yolov2/build/tinyYolov2/types';
24

35
export type NetParams = {
46
conv0: ConvParams
@@ -11,4 +13,49 @@ export type NetParams = {
1113
conv7: ConvParams
1214
fc0: FCParams
1315
fc1: FCParams
14-
}
16+
}
17+
18+
export type ConvWithBatchNormParams = BatchNormParams & {
19+
filter: tf.Tensor4D
20+
}
21+
22+
export type BatchNormParams = {
23+
mean: tf.Tensor1D
24+
variance: tf.Tensor1D
25+
scale: tf.Tensor1D
26+
offset: tf.Tensor1D
27+
}
28+
29+
export type SeparableConvWithBatchNormParams = {
30+
depthwise: ConvWithBatchNormParams
31+
pointwise: ConvWithBatchNormParams
32+
}
33+
34+
export declare type FCWithBatchNormParams = BatchNormParams & {
35+
weights: tf.Tensor2D
36+
}
37+
38+
export type DenseBlockParams = {
39+
conv0: SeparableConvParams | ConvParams
40+
conv1: SeparableConvParams
41+
conv2: SeparableConvParams
42+
//conv3: SeparableConvParams
43+
}
44+
45+
export type TinyNetParams = {
46+
dense0: DenseBlockParams
47+
dense1: DenseBlockParams
48+
dense2: DenseBlockParams
49+
fc: FCParams
50+
}
51+
52+
export type MobileResnetParams = {
53+
conv0: SeparableConvParams
54+
conv1: SeparableConvParams
55+
conv2: SeparableConvParams
56+
conv3: SeparableConvParams
57+
conv4: SeparableConvParams
58+
conv5: SeparableConvParams
59+
fc: FCParams
60+
}
61+

‎tools/quantize/views/quantizeModel.html

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,9 @@
1010
<script>
1111
tf = faceapi.tf
1212

13-
const modelName = 'tiny_yolov2_separable_conv'
14-
const uncompressedWeightsUri = `tiny_yolov2_separable_conv_model_v1.weights`
15-
const net = new faceapi.TinyYolov2(true)
13+
const modelName = 'face_landmark_68_tiny_model'
14+
const uncompressedWeightsUri = `face_landmark_68_tiny_model.weights`
15+
const net = new faceapi.FaceLandmark68TinyNet()
1616

1717
async function loadNetWeights(uri) {
1818
return new Float32Array(await (await fetch(uri)).arrayBuffer())

0 commit comments

Comments
 (0)