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 model 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) def compile_classifier(model, config: Config): optimizer = keras.optimizers.Adam(learning_rate=config.learning_rate) model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])