Spaces:
Sleeping
Sleeping
File size: 1,900 Bytes
5b889a9 acfc7fb 960edbd 5b889a9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import tensorflow as tf
from tensorflow import keras
from ultralytics import YOLO
from official.projects.movinet.modeling.movinet import Movinet
from official.projects.movinet.modeling.movinet_model import MovinetClassifier
from configuration import Config
class AttentionDenseClassifierHead(keras.layers.Layer):
def __init__(self, attention_heads, dense_units, dropout_rate=0.2, **kwargs):
super().__init__(**kwargs)
self.attention = keras.layers.MultiHeadAttention(num_heads=attention_heads, key_dim=1)
self.normalization = keras.layers.LayerNormalization(epsilon=1e-6)
self.dropout = keras.layers.Dropout(dropout_rate)
self.dense = keras.layers.Dense(dense_units, activation='softmax')
def call(self, x, training):
y = tf.expand_dims(x, -1)
y = self.attention(query=y, key=y, value=y)
y = tf.squeeze(y, axis=-1)
y = self.dropout(y, training=training)
y = self.normalization(x + y*0.01)
y = self.dense(y)
return y
def build_movinet(output_size, config: Config):
model = MovinetClassifier(
backbone=Movinet(model_id=config.model_id),
num_classes=output_size)
model.build(config.input_shape)
return model
def build_classifier_head(input_size, config: Config):
# inputs = keras.Input(shape=(input_size,))
# classifier = AttentionDenseClassifierHead(2, config.num_classes)(inputs)
# model = keras.Model(inputs=inputs, outputs=classifier)
return keras.layers.Dense(config.num_classes, activation='softmax')
def build_model(movinet, classifier_head):
return keras.models.Sequential([movinet, classifier_head])
def load_classifier(config: Config):
movinet = build_movinet(600, config)
classifier_head = build_classifier_head(600, config)
model = build_model(movinet, classifier_head)
model.load_weights(config.classifier_path)
return model
def load_detector(config: Config):
return YOLO(config.detector_path)
|