Spaces:
Sleeping
Sleeping
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) | |
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']) | |