src/model/classifier.py

"""Builds the EfficientNetV2B0-based pest classifier and loads trained weights."""
import tensorflow as tf
from tensorflow.keras import layers, regularizers

from src.configs import settings


def build_model() -> tf.keras.Model:
    base_model = tf.keras.applications.efficientnet_v2.EfficientNetV2B0(include_top=False)
    base_model.trainable = False

    inputs = layers.Input(shape=settings.INPUT_SHAPE, name="input_layer")
    x = base_model(inputs, training=False)
    x = layers.GlobalAveragePooling2D(name="pooling_layer")(x)
    x = layers.Dense(len(settings.CLASS_LABELS), kernel_regularizer=regularizers.l2(0.001))(x)
    outputs = layers.Activation("softmax", dtype=tf.float32, name="softmax_float32")(x)

    model = tf.keras.Model(inputs, outputs)
    model.compile(
        loss="categorical_crossentropy",
        optimizer=tf.keras.optimizers.Adam(),
        metrics=["accuracy"],
    )
    return model


def load_trained_model() -> tf.keras.Model:
    model = build_model()
    model.load_weights(str(settings.WEIGHTS_PATH))
    return model