Skip to content

Commit 34ac598

Browse files
lina128lina128
authored
[pose-detection]Use url querystring to control which model to load. (#656)
PROCESS
1 parent 14e31bc commit 34ac598

File tree

5 files changed

+98
-84
lines changed

5 files changed

+98
-84
lines changed

‎pose-detection/demo/src/camera.js

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -127,8 +127,7 @@ export class Camera {
127127
drawKeypoint(keypoint) {
128128
// If score is null, just show the keypoint.
129129
const score = keypoint.score != null ? keypoint.score : 1;
130-
const scoreThreshold =
131-
params.STATE.model[params.STATE.model.model].scoreThreshold || 0;
130+
const scoreThreshold = params.STATE.model.scoreThreshold || 0;
132131

133132
if (score >= scoreThreshold) {
134133
const circle = new Path2D();
@@ -155,8 +154,7 @@ export class Camera {
155154
// If score is null, just show the keypoint.
156155
const score1 = kp1.score != null ? kp1.score : 1;
157156
const score2 = kp2.score != null ? kp2.score : 1;
158-
const scoreThreshold =
159-
params.STATE.model[params.STATE.model.model].scoreThreshold || 0;
157+
const scoreThreshold = params.STATE.model.scoreThreshold || 0;
160158

161159
if (score1 >= scoreThreshold && score2 >= scoreThreshold) {
162160
this.ctx.beginPath();

‎pose-detection/demo/src/index.js

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,8 @@ import {setupStats} from './stats_panel';
2727

2828
let detector, camera, stats;
2929

30-
async function createDetector(model) {
31-
switch (model) {
30+
async function createDetector() {
31+
switch (STATE.model.model) {
3232
case posedetection.SupportedModels.PoseNet:
3333
return posedetection.createDetector(STATE.model.model, {
3434
quantBytes: 4,
@@ -41,12 +41,10 @@ async function createDetector(model) {
4141
return posedetection.createDetector(
4242
STATE.model.model, {quantBytes: 4, upperBodyOnly: false});
4343
case posedetection.SupportedModels.MoveNet:
44-
const modelType =
45-
STATE.model[STATE.model.model].modelType == 'Lightning' ?
44+
const modelType = STATE.model.type == 'lightning' ?
4645
posedetection.movenet.modelType.SINGLEPOSE_LIGHTNING :
4746
posedetection.movenet.modelType.SINGLEPOSE_THUNDER;
48-
return posedetection.createDetector(
49-
STATE.model.model, {modelType: modelType});
47+
return posedetection.createDetector(STATE.model.model, {modelType});
5048
}
5149
}
5250

@@ -93,11 +91,21 @@ async function renderPrediction() {
9391

9492
async function app() {
9593
await tf.setBackend('webgl');
96-
setupDatGui();
94+
95+
// Gui content will change depending on which model is in the query string.
96+
const urlParams = new URLSearchParams(window.location.search);
97+
if (!urlParams.has('model')) {
98+
alert('Cannot find model in the query string.');
99+
return;
100+
}
101+
102+
setupDatGui(urlParams);
103+
97104
stats = setupStats();
105+
98106
camera = await Camera.setupCamera(STATE.camera);
99107

100-
detector = await createDetector(STATE.model.model);
108+
detector = await createDetector();
101109

102110
renderPrediction();
103111
};

‎pose-detection/demo/src/option_panel.js

Lines changed: 70 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -16,80 +16,91 @@
1616
*/
1717
import * as posedetection from '@tensorflow-models/pose-detection';
1818

19-
import {STATE, VIDEO_SIZE} from './params';
19+
import * as params from './params';
2020

21-
export function setupDatGui() {
21+
export function setupDatGui(urlParams) {
2222
const gui = new dat.GUI({width: 300});
2323

2424
// The camera folder contains options for video settings.
2525
const cameraFolder = gui.addFolder('Camera');
26-
const fpsController = cameraFolder.add(STATE.camera, 'targetFPS');
26+
const fpsController = cameraFolder.add(params.STATE.camera, 'targetFPS');
2727
fpsController.onFinishChange((targetFPS) => {
28-
STATE.changeToTargetFPS = +targetFPS;
28+
params.STATE.changeToTargetFPS = +targetFPS;
2929
});
30-
const sizeController =
31-
cameraFolder.add(STATE.camera, 'sizeOption', Object.keys(VIDEO_SIZE));
30+
const sizeController = cameraFolder.add(
31+
params.STATE.camera, 'sizeOption', Object.keys(params.VIDEO_SIZE));
3232
sizeController.onChange(option => {
33-
STATE.changeToSizeOption = option;
33+
params.STATE.changeToSizeOption = option;
3434
});
3535
cameraFolder.open();
3636

3737
// The model folder contains options for model selection.
3838
const modelFolder = gui.addFolder('Model');
39-
const modelController = modelFolder.add(
40-
STATE.model, 'model', Object.values(posedetection.SupportedModels));
41-
modelController.onChange(model => {
42-
STATE.changeToModel = model;
43-
switch (model) {
44-
case posedetection.SupportedModels.PoseNet:
45-
poseNetFolder.open();
46-
blazePoseFolder.close();
47-
moveNetFolder.close();
48-
break;
49-
case posedetection.SupportedModels.MediapipeBlazepose:
50-
blazePoseFolder.open();
51-
poseNetFolder.close();
52-
moveNetFolder.close();
53-
break;
54-
case posedetection.SupportedModels.MoveNet:
55-
blazePoseFolder.close();
56-
poseNetFolder.close();
57-
moveNetFolder.open();
58-
break;
59-
default:
60-
throw new Error(`${model} is not supported.`);
61-
}
62-
});
39+
40+
const model = urlParams.get('model');
41+
const type = urlParams.get('type');
42+
switch (model) {
43+
case 'posenet':
44+
addPoseNetControllers(modelFolder);
45+
break;
46+
case 'movenet':
47+
addMoveNetControllers(modelFolder, type);
48+
break;
49+
case 'blazepose':
50+
addBlazePoseControllers(modelFolder, type);
51+
break;
52+
default:
53+
alert(`${urlParams.get('model')}`);
54+
break;
55+
}
56+
6357
modelFolder.open();
6458

65-
// The MoveNet model config folder contains options for MoveNet config
66-
// settings.
67-
const moveNetFolder = gui.addFolder('MoveNet Config');
68-
const moveNetTypeController = moveNetFolder.add(
69-
STATE.model[posedetection.SupportedModels.MoveNet], 'modelType',
70-
['Thunder', 'Lightning']);
71-
moveNetTypeController.onChange(type => {
72-
STATE.changeToModel = type;
59+
return gui;
60+
}
61+
62+
// The MoveNet model config folder contains options for MoveNet config
63+
// settings.
64+
function addMoveNetControllers(modelFolder, type) {
65+
params.STATE.model = {
66+
model: posedetection.SupportedModels.MoveNet,
67+
...params.MOVENET_CONFIG
68+
};
69+
70+
let $type = type != null ? type : 'thunder';
71+
if ($type !== 'thunder' && $type !== 'lightning') {
72+
$type = 'thunder';
73+
}
74+
75+
params.STATE.model.type = $type;
76+
77+
const typeController =
78+
modelFolder.add(params.STATE.model, 'type', ['thunder', 'lightning']);
79+
typeController.onChange(type => {
80+
params.STATE.changeToModel = type;
7381
});
74-
moveNetFolder.add(
75-
STATE.model[posedetection.SupportedModels.MoveNet], 'scoreThreshold', 0,
76-
1);
77-
78-
// The Blazepose model config folder contains options for Blazepose config
79-
// settings.
80-
const blazePoseFolder = gui.addFolder('MediapipeBlazepose Config');
81-
blazePoseFolder.add(
82-
STATE.model[posedetection.SupportedModels.MediapipeBlazepose],
83-
'scoreThreshold', 0, 1);
84-
85-
// The PoseNet model config folder contains options for PoseNet config
86-
// settings.
87-
const poseNetFolder = gui.addFolder('PoseNet Config');
88-
poseNetFolder.add(
89-
STATE.model[posedetection.SupportedModels.PoseNet], 'scoreThreshold', 0,
90-
1);
91-
92-
moveNetFolder.open();
9382

94-
return gui;
83+
modelFolder.add(params.STATE.model, 'scoreThreshold', 0, 1);
84+
}
85+
86+
// The Blazepose model config folder contains options for Blazepose config
87+
// settings.
88+
function addBlazePoseControllers(modelFolder) {
89+
params.STATE.model = {
90+
model: posedetection.SupportedModels.MediapipeBlazepose,
91+
...params.BLAZEPOSE_CONFIG
92+
};
93+
94+
modelFolder.add(params.STATE.model, 'scoreThreshold', 0, 1);
95+
}
96+
97+
// The PoseNet model config folder contains options for PoseNet config
98+
// settings.
99+
function addPoseNetControllers(modelFolder) {
100+
params.STATE.model = {
101+
model: posedetection.SupportedModels.PoseNet,
102+
...params.POSENET_CONFIG
103+
};
104+
105+
modelFolder.add(params.STATE.model, 'scoreThreshold', 0, 1);
95106
}

‎pose-detection/demo/src/params.js

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -24,18 +24,14 @@ export const VIDEO_SIZE = {
2424
'640 X 360': {width: 640, height: 360}
2525
};
2626
export const STATE = {
27-
camera: {targetFPS: 60, sizeOption: '640 X 480'},
28-
model: {
29-
model: posedetection.SupportedModels.MoveNet,
30-
}
27+
camera: {targetFPS: 60, sizeOption: '640 X 480'}
3128
};
32-
STATE.model[posedetection.SupportedModels.MediapipeBlazepose] = {
29+
export const BLAZEPOSE_CONFIG = {
3330
scoreThreshold: 0.65
3431
};
35-
STATE.model[posedetection.SupportedModels.PoseNet] = {
32+
export const POSENET_CONFIG = {
3633
scoreThreshold: 0.5
3734
};
38-
STATE.model[posedetection.SupportedModels.MoveNet] = {
39-
modelType: 'Thunder',
35+
export const MOVENET_CONFIG = {
4036
scoreThreshold: 0.3
4137
};

‎pose-detection/src/constants.ts

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -47,11 +47,12 @@ export const COCO_CONNECTED_KEYPOINTS_PAIRS = [
4747
[7, 9], [8, 10], [11, 12], [11, 13], [12, 14], [13, 15], [14, 16]
4848
];
4949
export const BLAZEPOSE_CONNECTED_KEYPOINTS_PAIRS = [
50-
[0, 1], [0, 4], [1, 2], [2, 3], [3, 7], [4, 5], [5, 6],
51-
[6, 8], [9, 10], [11, 12], [11, 13], [11, 23], [12, 14], [14, 16],
52-
[12, 24], [13, 15], [15, 17], [16, 18], [15, 21], [16, 22], [17, 19],
53-
[18, 20], [23, 25], [23, 24], [24, 26], [25, 27], [26, 28], [27, 29],
54-
[28, 30], [27, 31], [28, 32], [29, 31], [30, 32]
50+
[0, 1], [0, 4], [1, 2], [2, 3], [3, 7], [4, 5],
51+
[5, 6], [6, 8], [9, 10], [11, 12], [11, 13], [11, 23],
52+
[12, 14], [14, 16], [12, 24], [13, 15], [15, 17], [16, 18],
53+
[16, 20], [15, 17], [15, 19], [15, 21], [16, 22], [17, 19],
54+
[18, 20], [23, 25], [23, 24], [24, 26], [25, 27], [26, 28],
55+
[27, 29], [28, 30], [27, 31], [28, 32], [29, 31], [30, 32]
5556
];
5657
export const COCO_KEYPOINTS_NAMED_MAP: {[index: string]: number} = {
5758
nose: 0,

0 commit comments

Comments
 (0)