submodel.py 2.8 KB
from tensorflow import keras
from .. import initializers
from .. import layers
from ..utils.anchors import AnchorParameters
from . import assert_training_model
from . import retinanet

def custom_classification_model(
    num_classes,
    num_anchors,
    pyramid_feature_size=256,
    prior_probability=0.01,
    classification_feature_size=256,
    name='classification_submodel'
):
     # set input
    if keras.backend.image_data_format() == "channels_first":
        inputs = keras.layers.Input(shape=(pyramid_feature_size, None, None))
    else:
        inputs = keras.layers.Input(shape=(None, None, pyramid_feature_size))

    outputs = inputs

    # 3 layer
    for i in range(3):

        # 각 층의 output
        outputs = keras.layers.Conv2D(
            filters=classification_feature_size,
            activation="relu",
            name="pyramid_classification_{}".format(i),
            kernel_initializer=keras.initializers.RandomNormal(
                mean=0.0, stddev=0.01, seed=None
            ),  # 정규분포에 따라 텐서를 생성하는 초기값 설정
            bias_initializer="zeros",
            **options
        )(outputs)

    # 마지막 layer는 다른 필터로 다른 conv layer를 생성
    outputs = keras.layers.Conv2D(
        filters=num_classes * num_anchors,
        kernel_initializer=keras.initializers.RandomNormal(
            mean=0.0, stddev=0.01, seed=None
        ),
        bias_initializer=initializers.PriorProbability(probability=prior_probability),
        name="pyramid_classification",
        **options
    )(outputs)

    # reshape output and apply sigmoid
    if keras.backend.image_data_format() == "channels_first":
        outputs = keras.layers.Permute(
            (2, 3, 1), name="pyramid_classification_permute"
        )(outputs)

    # reshape : 2차원 > 1차원
    outputs = keras.layers.Reshape(
        (-1, num_classes), name="pyramid_classification_reshape"
    )(outputs)

    # output layer activation : sigmoid
    outputs = keras.layers.Activation("sigmoid", name="pyramid_classification_sigmoid")(
        outputs
    )

    return keras.models.Model(inputs=inputs, outputs=outputs, name=name)

def custom_regression_model(num_values, num_anchors, pyramid_feature_size=256, regression_feature_size=256, name='regression_submodel'):
    if num_anchors is None:
        num_anchors = AnchorParameters.default.num_anchors()
    model = retinanet.default_regression_model(num_values, num_anchors, pyramid_feature_size, regression_feature_size, name)
    return model


def custom_submodels(num_classes, num_anchors):
    if num_anchors is None:
        num_anchors = AnchorParameters.default.num_anchors()
    return [
        ("regression", custom_regression_model(4, num_anchors)),
        ("classification", custom_classification_model(num_classes, num_anchors)),
    ]