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) | |