Skip to content

Commit 84e287c

Browse files
justadudewhohacksjustadudewhohacks
committed
postprocessing for tiny-yolov2
1 parent e82cc5a commit 84e287c

File tree

10 files changed

+134
-24
lines changed

10 files changed

+134
-24
lines changed

‎src/mtcnn/BoundingBox.ts renamed to ‎src/BoundingBox.ts

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import { Rect } from './Rect';
2+
13
export class BoundingBox {
24
constructor(
35
private _left: number,
@@ -31,7 +33,6 @@ export class BoundingBox {
3133
return this.bottom - this.top
3234
}
3335

34-
3536
public toSquare(): BoundingBox {
3637
let { left, top, right, bottom } = this
3738

@@ -98,4 +99,8 @@ export class BoundingBox {
9899
this.bottom + (region.bottom * this.height)
99100
).toSquare().round()
100101
}
102+
103+
public toRect(): Rect {
104+
return new Rect(this.left, this.top, this.width, this.height)
105+
}
101106
}

‎src/mtcnn/nms.ts renamed to ‎src/commons/nonMaxSuppression.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
1-
import { BoundingBox } from './BoundingBox';
1+
import { BoundingBox } from '../BoundingBox';
22

3-
export function nms(
3+
export function nonMaxSuppression(
44
boxes: BoundingBox[],
55
scores: number[],
66
iouThreshold: number,

‎src/mtcnn/extractImagePatches.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
import * as tf from '@tensorflow/tfjs-core';
22

3+
import { BoundingBox } from '../BoundingBox';
34
import { Dimensions } from '../types';
45
import { createCanvas, getContext2dOrThrow } from '../utils';
5-
import { BoundingBox } from './BoundingBox';
66
import { normalize } from './normalize';
77

88
export async function extractImagePatches(

‎src/mtcnn/stage1.ts

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
11
import * as tf from '@tensorflow/tfjs-core';
22

3+
import { BoundingBox } from '../BoundingBox';
4+
import { nonMaxSuppression } from '../commons/nonMaxSuppression';
35
import { Point } from '../Point';
4-
import { BoundingBox } from './BoundingBox';
56
import { CELL_SIZE, CELL_STRIDE } from './config';
6-
import { nms } from './nms';
7+
import { getSizesForScale } from './getSizesForScale';
78
import { normalize } from './normalize';
89
import { PNet } from './PNet';
910
import { PNetParams } from './types';
10-
import { getSizesForScale } from './getSizesForScale';
1111

1212
function rescaleAndNormalize(x: tf.Tensor4D, scale: number): tf.Tensor4D {
1313
return tf.tidy(() => {
@@ -109,7 +109,7 @@ export function stage1(
109109
}
110110

111111
let ts = Date.now()
112-
const indices = nms(
112+
const indices = nonMaxSuppression(
113113
boundingBoxes.map(bbox => bbox.cell),
114114
boundingBoxes.map(bbox => bbox.score),
115115
0.5
@@ -130,7 +130,7 @@ export function stage1(
130130

131131
if (allBoxes.length > 0) {
132132
let ts = Date.now()
133-
const indices = nms(
133+
const indices = nonMaxSuppression(
134134
allBoxes.map(bbox => bbox.cell),
135135
allBoxes.map(bbox => bbox.score),
136136
0.7

‎src/mtcnn/stage2.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
import * as tf from '@tensorflow/tfjs-core';
22

3-
import { BoundingBox } from './BoundingBox';
3+
import { BoundingBox } from '../BoundingBox';
4+
import { nonMaxSuppression } from '../commons/nonMaxSuppression';
45
import { extractImagePatches } from './extractImagePatches';
5-
import { nms } from './nms';
66
import { RNet } from './RNet';
77
import { RNetParams } from './types';
88

@@ -47,7 +47,7 @@ export async function stage2(
4747

4848
if (filteredBoxes.length > 0) {
4949
ts = Date.now()
50-
const indicesNms = nms(
50+
const indicesNms = nonMaxSuppression(
5151
filteredBoxes,
5252
filteredScores,
5353
0.7

‎src/mtcnn/stage3.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
import * as tf from '@tensorflow/tfjs-core';
22

3+
import { BoundingBox } from '../BoundingBox';
4+
import { nonMaxSuppression } from '../commons/nonMaxSuppression';
35
import { Point } from '../Point';
4-
import { BoundingBox } from './BoundingBox';
56
import { extractImagePatches } from './extractImagePatches';
6-
import { nms } from './nms';
77
import { ONet } from './ONet';
88
import { ONetParams } from './types';
99

@@ -57,7 +57,7 @@ export async function stage3(
5757
if (filteredBoxes.length > 0) {
5858

5959
ts = Date.now()
60-
const indicesNms = nms(
60+
const indicesNms = nonMaxSuppression(
6161
filteredBoxes,
6262
filteredScores,
6363
0.7,

‎src/tinyYolov2/TinyYolov2.ts

Lines changed: 77 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,26 @@
11
import * as tf from '@tensorflow/tfjs-core';
22

3+
import { BoundingBox } from '../BoundingBox';
34
import { convLayer } from '../commons/convLayer';
45
import { NeuralNetwork } from '../commons/NeuralNetwork';
6+
import { nonMaxSuppression } from '../commons/nonMaxSuppression';
7+
import { FaceDetection } from '../FaceDetection';
58
import { NetInput } from '../NetInput';
69
import { toNetInput } from '../toNetInput';
710
import { TNetInput } from '../types';
11+
import { BOX_ANCHORS, INPUT_SIZES, NUM_BOXES, NUM_CELLS } from './config';
812
import { convWithBatchNorm } from './convWithBatchNorm';
913
import { extractParams } from './extractParams';
10-
import { NetParams } from './types';
11-
14+
import { getDefaultParams } from './getDefaultParams';
15+
import { NetParams, TinyYolov2ForwardParams } from './types';
1216

1317
export class TinyYolov2 extends NeuralNetwork<NetParams> {
1418

1519
constructor() {
1620
super('TinyYolov2')
1721
}
1822

19-
public async forwardInput(input: NetInput): Promise<any> {
23+
public forwardInput(input: NetInput, inputSize: number): tf.Tensor4D {
2024

2125
const { params } = this
2226

@@ -25,10 +29,7 @@ export class TinyYolov2 extends NeuralNetwork<NetParams> {
2529
}
2630

2731
const out = tf.tidy(() => {
28-
//const batchTensor = input.toBatchTensor(416).div(tf.scalar(255)).toFloat() as tf.Tensor4D
29-
30-
// TODO: fix boxes after padding
31-
const batchTensor = tf.image.resizeBilinear(input.inputs[0], [416, 416]).toFloat().div(tf.scalar(255)).expandDims() as tf.Tensor4D
32+
const batchTensor = input.toBatchTensor(inputSize, false).div(tf.scalar(255)).toFloat() as tf.Tensor4D
3233

3334
let out = convWithBatchNorm(batchTensor, params.conv0)
3435
out = tf.maxPool(out, [2, 2], [2, 2], 'same')
@@ -52,8 +53,75 @@ export class TinyYolov2 extends NeuralNetwork<NetParams> {
5253
return out
5354
}
5455

55-
public async forward(input: TNetInput): Promise<any> {
56-
return await this.forwardInput(await toNetInput(input, true, true))
56+
public async forward(input: TNetInput, inputSize: number): Promise<tf.Tensor4D> {
57+
return await this.forwardInput(await toNetInput(input, true, true), inputSize)
58+
}
59+
60+
public async locateFaces(input: TNetInput, forwardParams: TinyYolov2ForwardParams = {}): Promise<FaceDetection[]> {
61+
62+
const { sizeType, scoreThreshold } = getDefaultParams(forwardParams)
63+
64+
65+
const inputSize = INPUT_SIZES[sizeType]
66+
const numCells = NUM_CELLS[sizeType]
67+
68+
if (!inputSize) {
69+
throw new Error(`TinyYolov2 - unkown sizeType: ${sizeType}, expected one of: xs | sm | md | lg`)
70+
}
71+
72+
const netInput = await toNetInput(input, true)
73+
const out = await this.forwardInput(netInput, inputSize)
74+
75+
const [boxesTensor, scoresTensor] = tf.tidy(() => {
76+
const reshaped = out.reshape([numCells, numCells, NUM_BOXES, 6])
77+
out.dispose()
78+
79+
const boxes = reshaped.slice([0, 0, 0, 0], [numCells, numCells, NUM_BOXES, 4])
80+
const scores = reshaped.slice([0, 0, 0, 4], [numCells, numCells, NUM_BOXES, 1])
81+
return [boxes, scores]
82+
})
83+
84+
const expit = (x: number): number => 1 / (1 + Math.exp(-x))
85+
86+
const paddedHeightRelative = (netInput.getPaddings(0).y + netInput.getInputHeight(0)) / netInput.getInputHeight(0)
87+
const paddedWidthRelative = (netInput.getPaddings(0).x + netInput.getInputWidth(0)) / netInput.getInputWidth(0)
88+
89+
const boxes: BoundingBox[] = []
90+
const scores: number[] = []
91+
92+
for (let row = 0; row < numCells; row ++) {
93+
for (let col = 0; col < numCells; col ++) {
94+
for (let box = 0; box < NUM_BOXES; box ++) {
95+
const score = expit(scoresTensor.get(row, col, box, 0))
96+
if (score > scoreThreshold) {
97+
const ctX = ((col + expit(boxesTensor.get(row, col, box, 0))) / numCells) * paddedWidthRelative
98+
const ctY = ((row + expit(boxesTensor.get(row, col, box, 1))) / numCells) * paddedHeightRelative
99+
const width = ((Math.exp(boxesTensor.get(row, col, box, 2)) * BOX_ANCHORS[box].x) / numCells) * paddedWidthRelative
100+
const height = ((Math.exp(boxesTensor.get(row, col, box, 3)) * BOX_ANCHORS[box].y) / numCells) * paddedHeightRelative
101+
102+
const x = (ctX - (width / 2))
103+
const y = (ctY - (height / 2))
104+
boxes.push(new BoundingBox(x, y, x + width, y + height))
105+
scores.push(score)
106+
}
107+
}
108+
}
109+
}
110+
111+
boxesTensor.dispose()
112+
scoresTensor.dispose()
113+
114+
const indices = nonMaxSuppression(boxes, scores, 0.4, true)
115+
116+
const detections = indices.map(idx =>
117+
new FaceDetection(
118+
scores[idx],
119+
boxes[idx].toRect(),
120+
{ width: netInput.getInputWidth(0), height: netInput.getInputHeight(0) }
121+
)
122+
)
123+
124+
return detections
57125
}
58126

59127
/* TODO

‎src/tinyYolov2/config.ts

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
import { Point } from '../Point';
2+
3+
export const INPUT_SIZES = { xs: 224, sm: 320, md: 416, lg: 608 }
4+
export const NUM_CELLS = { xs: 7, sm: 10, md: 13, lg: 19 }
5+
export const NUM_BOXES = 5
6+
7+
export const BOX_ANCHORS = [
8+
new Point(0.738768, 0.874946),
9+
new Point(2.42204, 2.65704),
10+
new Point(4.30971, 7.04493),
11+
new Point(10.246, 4.59428),
12+
new Point(12.6868, 11.8741)
13+
]

‎src/tinyYolov2/getDefaultParams.ts

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
import { SizeType, TinyYolov2ForwardParams } from './types';
2+
3+
export function getDefaultParams(params: TinyYolov2ForwardParams) {
4+
return Object.assign(
5+
{},
6+
{
7+
sizeType: SizeType.MD,
8+
scoreThreshold: 0.5
9+
},
10+
params
11+
)
12+
}

‎src/tinyYolov2/types.ts

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,4 +22,16 @@ export type NetParams = {
2222
conv6: ConvWithBatchNorm
2323
conv7: ConvWithBatchNorm
2424
conv8: ConvParams
25+
}
26+
27+
export enum SizeType {
28+
XS = 'xs',
29+
SM = 'sm',
30+
MD = 'md',
31+
LG = 'lg'
32+
}
33+
34+
export type TinyYolov2ForwardParams = {
35+
sizeType?: SizeType
36+
scoreThreshold?: number
2537
}

0 commit comments

Comments
 (0)