15
15
* =============================================================================
16
16
*/
17
17
import * as posedetection from '@tensorflow-models/pose-detection' ;
18
- import { type } from 'os' ;
19
18
20
19
import * as params from './params' ;
21
20
@@ -39,77 +38,111 @@ export function setupDatGui(urlParams) {
39
38
const modelFolder = gui . addFolder ( 'Model' ) ;
40
39
41
40
const model = urlParams . get ( 'model' ) ;
42
- const type = urlParams . get ( 'type' ) ;
41
+ let type = urlParams . get ( 'type' ) ;
42
+
43
+ let modelConfigFolder ;
44
+
43
45
switch ( model ) {
44
46
case 'posenet' :
45
- addPoseNetControllers ( modelFolder ) ;
47
+ params . STATE . model = posedetection . SupportedModels . PoseNet ;
46
48
break ;
47
49
case 'movenet' :
48
- addMoveNetControllers ( modelFolder , type ) ;
50
+ params . STATE . model = posedetection . SupportedModels . MoveNet ;
51
+ if ( type !== 'lightning' && type !== 'thunder' ) {
52
+ // Nulify invalid value.
53
+ type = null ;
54
+ }
49
55
break ;
50
56
case 'blazepose' :
51
- addBlazePoseControllers ( modelFolder , type ) ;
57
+ params . STATE . model = type === 'upperbody' ?
58
+ posedetection . SupportedModels . MediapipeBlazeposeUpperBody :
59
+ posedetection . SupportedModels . MediapipeBlazeposeFullBody ;
60
+ if ( type !== 'fullbody' && type !== 'upperbody' ) {
61
+ // Nulify invalid value.
62
+ type = null ;
63
+ }
52
64
break ;
53
65
default :
54
66
alert ( `${ urlParams . get ( 'model' ) } ` ) ;
55
67
break ;
56
68
}
57
69
58
- modelFolder . open ( ) ;
70
+ const modelController = modelFolder . add (
71
+ params . STATE , 'model' , Object . values ( posedetection . SupportedModels ) ) ;
59
72
60
- return gui ;
61
- }
73
+ modelController . onChange ( model => {
74
+ params . STATE . changeToModel = model ;
62
75
63
- // The MoveNet model config folder contains options for MoveNet config
64
- // settings.
65
- function addMoveNetControllers ( modelFolder , type ) {
66
- params . STATE . model = {
67
- model : posedetection . SupportedModels . MoveNet ,
68
- ...params . MOVENET_CONFIG
69
- } ;
70
-
71
- params . STATE . model . type =
72
- type !== 'thunder' && type !== 'lightning' ? 'thunder' : type ;
73
-
74
- const typeController =
75
- modelFolder . add ( params . STATE . model , 'type' , [ 'thunder' , 'lightning' ] ) ;
76
- typeController . onChange ( type => {
77
- params . STATE . changeToModel = type ;
76
+ // We don't pass in type, so that it will use default type when switching
77
+ // models.
78
+ modelConfigFolder = updateModelConfigFolder ( gui , model , modelConfigFolder ) ;
79
+
80
+ modelConfigFolder . open ( ) ;
78
81
} ) ;
79
82
80
- modelFolder . add ( params . STATE . model , 'scoreThreshold' , 0 , 1 ) ;
81
- }
83
+ modelFolder . open ( ) ;
82
84
83
- // The Blazepose model config folder contains options for Blazepose config
84
- // settings.
85
- function addBlazePoseControllers ( modelFolder , type ) {
86
- params . STATE . model = { ...params . BLAZEPOSE_CONFIG } ;
85
+ // For initialization, pass in type from url.
86
+ modelConfigFolder =
87
+ updateModelConfigFolder ( gui , params . STATE . model , modelConfigFolder , type ) ;
88
+
89
+ modelConfigFolder . open ( ) ;
90
+
91
+ return gui ;
92
+ }
87
93
88
- params . STATE . model . model = type === 'upperbody' ?
89
- posedetection . SupportedModels . MediapipeBlazeposeUpperBody :
90
- posedetection . SupportedModels . MediapipeBlazeposeFullBody ;
94
+ function updateModelConfigFolder ( gui , model , modelConfigFolder , type ) {
95
+ if ( modelConfigFolder != null ) {
96
+ gui . removeFolder ( modelConfigFolder ) ;
97
+ }
91
98
92
- params . STATE . model . type = type === 'upperbody' ? 'upperbody' : 'fullbody' ;
99
+ const newModelConfigFolder = gui . addFolder ( 'Model Config' ) ;
93
100
94
- const typeController =
95
- modelFolder . add ( params . STATE . model , 'type' , [ 'fullbody' , 'upperbody' ] ) ;
96
- typeController . onChange ( type => {
97
- params . STATE . changeToModel = type ;
98
- params . STATE . model . model = type === 'upperbody' ?
99
- posedetection . SupportedModels . MediapipeBlazeposeUpperBody :
100
- posedetection . SupportedModels . MediapipeBlazeposeFullBody ;
101
- } )
101
+ switch ( model ) {
102
+ case posedetection . SupportedModels . PoseNet :
103
+ addPoseNetControllers ( newModelConfigFolder ) ;
104
+ break ;
105
+ case posedetection . SupportedModels . MoveNet :
106
+ addMoveNetControllers ( newModelConfigFolder , type ) ;
107
+ break ;
108
+ case posedetection . SupportedModels . MediapipeBlazeposeUpperBody :
109
+ case posedetection . SupportedModels . MediapipeBlazeposeFullBody :
110
+ addBlazePoseControllers ( newModelConfigFolder ) ;
111
+ break ;
112
+ default :
113
+ alert ( `Model ${ model } is not supported.` ) ;
114
+ }
102
115
103
- modelFolder . add ( params . STATE . model , 'scoreThreshold' , 0 , 1 ) ;
116
+ return newModelConfigFolder ;
104
117
}
105
118
106
119
// The PoseNet model config folder contains options for PoseNet config
107
120
// settings.
108
- function addPoseNetControllers ( modelFolder ) {
109
- params . STATE . model = {
110
- model : posedetection . SupportedModels . PoseNet ,
111
- ...params . POSENET_CONFIG
112
- } ;
121
+ function addPoseNetControllers ( modelConfigFolder ) {
122
+ params . STATE . modelConfig = { ...params . POSENET_CONFIG } ;
123
+ modelConfigFolder . add ( params . STATE . modelConfig , 'scoreThreshold' , 0 , 1 ) ;
124
+ }
125
+
126
+ // The MoveNet model config folder contains options for MoveNet config
127
+ // settings.
128
+ function addMoveNetControllers ( modelConfigFolder , type ) {
129
+ params . STATE . modelConfig = { ...params . MOVENET_CONFIG } ;
130
+ params . STATE . modelConfig . type = type != null ? type : 'thunder' ;
131
+
132
+ const typeController = modelConfigFolder . add (
133
+ params . STATE . modelConfig , 'type' , [ 'thunder' , 'lightning' ] ) ;
134
+ typeController . onChange ( _ => {
135
+ // Set changeToModel to non-null, so that we don't render any result when
136
+ // changeToModel is non-null.
137
+ params . STATE . changeToModel = params . STATE . model ;
138
+ } ) ;
113
139
114
- modelFolder . add ( params . STATE . model , 'scoreThreshold' , 0 , 1 ) ;
140
+ modelConfigFolder . add ( params . STATE . modelConfig , 'scoreThreshold' , 0 , 1 ) ;
141
+ }
142
+
143
+ // The Blazepose model config folder contains options for Blazepose config
144
+ // settings.
145
+ function addBlazePoseControllers ( modelConfigFolder ) {
146
+ params . STATE . modelConfig = { ...params . BLAZEPOSE_CONFIG } ;
147
+ modelConfigFolder . add ( params . STATE . modelConfig , 'scoreThreshold' , 0 , 1 ) ;
115
148
}
0 commit comments