tosanoob/mediapipe_fastapi_20classes
0
1# --------- Define auxiliary classes ---------
2
3import os
4import keras
5import tensorflow as tf
6
7@keras.saving.register_keras_serializable(package="1DCNN_Transformer")
8class ECA(tf.keras.layers.Layer):
9 def __init__(self, kernel_size=5, **kwargs):
10 super().__init__(**kwargs)
11 self.supports_masking = True
12 self.kernel_size = kernel_size
13 self.conv = tf.keras.layers.Conv1D(1, kernel_size=kernel_size, strides=1, padding="same", use_bias=False)
14
15 def call(self, inputs, mask=None):
16 nn = tf.keras.layers.GlobalAveragePooling1D()(inputs, mask=mask)
17 nn = tf.expand_dims(nn, -1)
18 nn = self.conv(nn)
19 nn = tf.squeeze(nn, -1)
20 nn = tf.nn.sigmoid(nn)
21 nn = nn[:,None,:]
22 return inputs * nn
23
24 def get_config(self):
25 base_config = super().get_config()
26 config = {
27 # "supports_masking" : keras.saving.serialize_keras_object(self.supports_masking),
28 "kernel_size" : keras.saving.serialize_keras_object(self.kernel_size)
29 }
30 return {**base_config, **config}
31
32 @classmethod
33 def from_config(cls,config):
34 kernel_size_config = config.pop("kernel_size")
35 kernel_size = keras.saving.deserialize_keras_object(kernel_size_config)
36 return cls(kernel_size, **config)
37
38@keras.saving.register_keras_serializable(package="1DCNN_Transformer")
39class LateDropout(tf.keras.layers.Layer):
40 def __init__(self, rate, noise_shape=None, start_step=0, **kwargs):
41 super().__init__(**kwargs)
42 self.supports_masking = True
43 self.rate = rate
44 self.noise_shape = noise_shape
45 self.start_step = start_step
46 self.dropout = tf.keras.layers.Dropout(rate, noise_shape=noise_shape)
47
48 def build(self, input_shape):
49 super().build(input_shape)
50 agg = tf.VariableAggregation.ONLY_FIRST_REPLICA
51 self._train_counter = tf.Variable(0, dtype="int64", aggregation=agg, trainable=False)
52
53 def call(self, inputs, training=False):
54 x = tf.cond(self._train_counter < self.start_step, lambda:inputs, lambda:self.dropout(inputs, training=training))
55 if training:
56 self._train_counter.assign_add(1)
57 return x
58
59 def get_config(self):
60 base_config = super().get_config()
61 config = {
62 # "supports_masking" : keras.saving.serialize_keras_object(self.supports_masking),
63 "rate" : keras.saving.serialize_keras_object(self.rate),
64 "start_step" : keras.saving.serialize_keras_object(self.start_step),
65 "noise_shape" : keras.saving.serialize_keras_object(self.noise_shape),
66 }
67 return {**base_config, **config}
68
69 @classmethod
70 def from_config(cls,config):
71 rate_config = config.pop("rate")
72 rate = keras.saving.deserialize_keras_object(rate_config)
73 start_step_config = config.pop("start_step")
74 start_step = keras.saving.deserialize_keras_object(start_step_config)
75 noise_shape_config = config.pop("noise_shape")
76 noise_shape = keras.saving.deserialize_keras_object(noise_shape_config)
77 return cls(rate, noise_shape, start_step, **config)
78
79@keras.saving.register_keras_serializable(package="1DCNN_Transformer")
80class CausalDWConv1D(tf.keras.layers.Layer):
81 def __init__(self,
82 kernel_size=17,
83 dilation_rate=1,
84 use_bias=False,
85 depthwise_initializer='glorot_uniform',
86 name='', **kwargs):
87 super().__init__(name=name,**kwargs)
88 self.kernel_size = kernel_size
89 self.dilation_rate = dilation_rate
90 self.use_bias = use_bias
91 self.depthwise_initializer=depthwise_initializer
92 self.lname=name
93
94 self.causal_pad = tf.keras.layers.ZeroPadding1D((dilation_rate*(kernel_size-1),0),name=name + '_pad')
95 self.dw_conv = tf.keras.layers.DepthwiseConv1D(
96 kernel_size,
97 strides=1,
98 dilation_rate=dilation_rate,
99 padding='valid',
100 use_bias=use_bias,
101 depthwise_initializer=depthwise_initializer,
102 name=name + '_dwconv')
103 self.supports_masking = True
104
105 def call(self, inputs):
106 x = self.causal_pad(inputs)
107 x = self.dw_conv(x)
108 return x
109
110 def get_config(self):
111 base_config = super().get_config()
112 config = {
113 "kernel_size" : keras.saving.serialize_keras_object(self.kernel_size),
114 "dilation_rate" : keras.saving.serialize_keras_object(self.dilation_rate),
115 "use_bias" : keras.saving.serialize_keras_object(self.use_bias),
116 "depthwise_initializer" : keras.saving.serialize_keras_object(self.depthwise_initializer),
117 "name" : keras.saving.serialize_keras_object(self.lname),
118 }
119 return {**base_config, **config}
120
121 @classmethod
122 def from_config(cls,config):
123 kernel_size_config = config.pop("kernel_size")
124 kernel_size = keras.saving.deserialize_keras_object(kernel_size_config)
125 dilation_rate_config = config.pop("dilation_rate")
126 dilation_rate = keras.saving.deserialize_keras_object(dilation_rate_config)
127 bias_config = config.pop("use_bias")
128 bias = keras.saving.deserialize_keras_object(bias_config)
129 depthwise_config = config.pop("depthwise_initializer")
130 depthwise = keras.saving.deserialize_keras_object(depthwise_config)
131 name_config = config.pop("name")
132 name = keras.saving.deserialize_keras_object(name_config)
133
134 return cls(kernel_size,dilation_rate,bias,depthwise,name, **config)
135
136def Conv1DBlock(channel_size,
137 kernel_size,
138 dilation_rate=1,
139 drop_rate=0.0,
140 expand_ratio=2,
141 se_ratio=0.25,
142 activation='swish',
143 name=None):
144 '''
145 efficient conv1d block, @hoyso48
146 '''
147 if name is None:
148 name = str(tf.keras.backend.get_uid("mbblock"))
149 # Expansion phase
150 def apply(inputs):
151 channels_in = tf.keras.backend.int_shape(inputs)[-1]
152 channels_expand = channels_in * expand_ratio
153
154 skip = inputs
155
156 x = tf.keras.layers.Dense(
157 channels_expand,
158 use_bias=True,
159 activation=activation,
160 name=name + '_expand_conv')(inputs)
161
162 # Depthwise Convolution
163 x = CausalDWConv1D(kernel_size,
164 dilation_rate=dilation_rate,
165 use_bias=False,
166 name=name + '_dwconv')(x)
167
168 x = tf.keras.layers.BatchNormalization(momentum=0.95, name=name + '_bn')(x)
169
170 x = ECA()(x)
171
172 x = tf.keras.layers.Dense(
173 channel_size,
174 use_bias=True,
175 name=name + '_project_conv')(x)
176
177 if drop_rate > 0:
178 x = tf.keras.layers.Dropout(drop_rate, noise_shape=(None,1,1), name=name + '_drop')(x)
179
180 if (channels_in == channel_size):
181 x = tf.keras.layers.add([x, skip], name=name + '_add')
182 return x
183
184 return apply
185
186
187@keras.saving.register_keras_serializable(package="1DCNN_Transformer")
188class MultiHeadSelfAttention(tf.keras.layers.Layer):
189 def __init__(self, dim=256, num_heads=4, dropout=0, **kwargs):
190 super().__init__(**kwargs)
191 self.dim = dim
192 self.scale = self.dim ** -0.5
193 self.num_heads = num_heads
194 self.dropout = dropout
195 self.qkv = tf.keras.layers.Dense(3 * dim, use_bias=False)
196 self.drop1 = tf.keras.layers.Dropout(dropout)
197 self.proj = tf.keras.layers.Dense(dim, use_bias=False)
198 self.supports_masking = True
199
200 def call(self, inputs, mask=None):
201 qkv = self.qkv(inputs)
202 qkv = tf.keras.layers.Permute((2, 1, 3))(tf.keras.layers.Reshape((-1, self.num_heads, self.dim * 3 // self.num_heads))(qkv))
203 q, k, v = tf.split(qkv, [self.dim // self.num_heads] * 3, axis=-1)
204
205 attn = tf.matmul(q, k, transpose_b=True) * self.scale
206
207 if mask is not None:
208 mask = mask[:, None, None, :]
209
210 attn = tf.keras.layers.Softmax(axis=-1)(attn, mask=mask)
211 attn = self.drop1(attn)
212
213 x = attn @ v
214 x = tf.keras.layers.Reshape((-1, self.dim))(tf.keras.layers.Permute((2, 1, 3))(x))
215 x = self.proj(x)
216 return x
217
218 def get_config(self):
219 base_config = super().get_config()
220 config = {
221 "dim" : self.dim,
222 "num_heads" : self.num_heads,
223 "dropout" : self.dropout,
224 }
225 return {**base_config, **config}
226
227 @classmethod
228 def from_config(cls,config):
229 dim_config = config.pop("dim")
230 dim = keras.saving.deserialize_keras_object(dim_config)
231 num_heads_config = config.pop("num_heads")
232 num_heads = keras.saving.deserialize_keras_object(num_heads_config)
233 dropout_config = config.pop("dropout")
234 dropout = keras.saving.deserialize_keras_object(dropout_config)
235 return cls(dim,num_heads,dropout)
236
237def TransformerBlock(dim=256, num_heads=4, expand=4, attn_dropout=0.2, drop_rate=0.2, activation='swish'):
238 def apply(inputs):
239 x = inputs
240 x = tf.keras.layers.BatchNormalization(momentum=0.95)(x)
241 x = MultiHeadSelfAttention(dim=dim,num_heads=num_heads,dropout=attn_dropout)(x)
242 x = tf.keras.layers.Dropout(drop_rate, noise_shape=(None,1,1))(x)
243 x = tf.keras.layers.Add()([inputs, x])
244 attn_out = x
245
246 x = tf.keras.layers.BatchNormalization(momentum=0.95)(x)
247 x = tf.keras.layers.Dense(dim*expand, use_bias=False, activation=activation)(x)
248 x = tf.keras.layers.Dense(dim, use_bias=False)(x)
249 x = tf.keras.layers.Dropout(drop_rate, noise_shape=(None,1,1))(x)
250 x = tf.keras.layers.Add()([attn_out, x])
251 return x
252 return apply
253
254MAX_LEN = 30 # number of frame
255CHANNELS = 258 # number of keypoint value
256NUM_CLASSES = 20
257PAD = -100
258
259# ----------------------------------------- DEFINE MODEL -----------------------------
260def get_model(max_len=MAX_LEN, dropout_step=0, dim=256):
261 inp = tf.keras.Input((max_len,CHANNELS))
262 # x = tf.keras.layers.Masking(mask_value=PAD,input_shape=(max_len,CHANNELS))(inp) #we don't need masking layer with inference
263 x = inp
264 ksize = 3
265 x = tf.keras.layers.Permute((2,1))(x)
266 x = tf.keras.layers.Dense(dim, use_bias=False,name='stem_conv')(x)
267 x = tf.keras.layers.BatchNormalization(momentum=0.95,name='stem_bn')(x)
268
269 x = Conv1DBlock(dim,ksize,drop_rate=0.2)(x)
270 x = Conv1DBlock(dim,ksize,drop_rate=0.2)(x)
271 x = Conv1DBlock(dim,ksize,drop_rate=0.2)(x)
272 x = TransformerBlock(dim,expand=2)(x)
273
274 x = Conv1DBlock(dim,ksize,drop_rate=0.2)(x)
275 x = Conv1DBlock(dim,ksize,drop_rate=0.2)(x)
276 x = Conv1DBlock(dim,ksize,drop_rate=0.2)(x)
277 x = TransformerBlock(dim,expand=2)(x)
278
279 if dim == 384: #for the 4x sized model
280 x = Conv1DBlock(dim,ksize,drop_rate=0.2)(x)
281 x = Conv1DBlock(dim,ksize,drop_rate=0.2)(x)
282 x = Conv1DBlock(dim,ksize,drop_rate=0.2)(x)
283 x = TransformerBlock(dim,expand=2)(x)
284
285 x = Conv1DBlock(dim,ksize,drop_rate=0.2)(x)
286 x = Conv1DBlock(dim,ksize,drop_rate=0.2)(x)
287 x = Conv1DBlock(dim,ksize,drop_rate=0.2)(x)
288 x = TransformerBlock(dim,expand=2)(x)
289
290 x = tf.keras.layers.Dense(dim*2,activation=None,name='top_conv')(x)
291 x = tf.keras.layers.GlobalAveragePooling1D()(x)
292 # x = LateDropout(0.5, start_step=dropout_step)(x)
293 x = tf.keras.layers.Dense(NUM_CLASSES,name='classifier',activation="softmax")(x)
294 return tf.keras.Model(inp, x)
295
296def load_model(path='1DCNN_Transformer_L-dim256_train8_1405_checkpoint.weights.h5'):
297 model = get_model()
298 module_dir = os.path.dirname(os.path.abspath(__file__))
299 model_path = os.path.join(module_dir,path)
300 model.load_weights(model_path)
301 return model