<filename>tfcv/modeling/backbones/vgg.py
import tensorflow as tf

from tfcv.modeling.layers import conv2d


CONFIGS = {
    11: [[64], [128], [256, 256], [512, 512], [512, 512]],
    13: [[64, 64], [128, 128], [256, 256], [512, 512], [512, 512]],
    16: [[64, 64], [128, 128], [256, 256, 256], [512, 512, 512], [512, 512, 512]],
    19: [
        [64, 64],
        [128, 128],
        [256, 256, 256, 256],
        [512, 512, 512, 512],
        [512, 512, 512, 512],
    ],
}


class VGG(tf.keras.Model):
    """ VGG.

    Architecture follows [2] which replaces the dense layers with global
    pooling resulting in a much smaller model size while improving the
    localization ability of the network.

    [1] Very Deep Convolutional Networks for Large-Scale Image Recognition
          https://arxiv.org/abs/1409.1556

    [2] Learning Deep Features for Discriminative Localization
          https://arxiv.org/abs/1512.04150
    """

    def __init__(
        self,
        input_spec: tf.keras.layers.InputSpec,
        depth: int,
        weight_decay: float,
        **kwargs,
    ):
        self._input_spec = input_spec
        self._depth = depth
        self._weight_decay = weight_decay

        x = inputs = tf.keras.Input(shape=self._input_spec.shape[1:])
        for i, channels in enumerate(CONFIGS[depth]):
            for j, c in enumerate(channels):
                x = conv2d(
                    x, c, 3, self._weight_decay, name=f"conv_block{i + 1}_{j + 1}"
                )
            if i < 4:
                x = tf.keras.layers.MaxPool2D(2, 2, "same")(x)
        x = conv2d(x, 1024, 3, self._weight_decay, name="conv6")
        super().__init__(inputs=inputs, outputs=x, name=f"VGG{depth}", **kwargs)

    def get_config(self):
        return {
            "depth": self._depth,
            "input_spec": self._input_spec,
            "weight_decay": self._weight_decay,
        }


def build_vgg(
    input_spec: tf.keras.layers.InputSpec, depth: int, weight_decay: float
) -> tf.keras.Model:
    model = VGG(input_spec, depth, weight_decay)
    return model


if __name__ == "__main__":
    input_spec = tf.keras.layers.InputSpec(dtype=tf.float32, shape=(None, 224, 224, 3))
    weight_decay = 1e-4

    model = build_vgg(input_spec, 11, weight_decay)
    tf.keras.utils.plot_model(model, "vgg11.png", True, show_layer_names=True)

    model = build_vgg(input_spec, 13, weight_decay)
    tf.keras.utils.plot_model(model, "vgg13.png", True, show_layer_names=True)

    model = build_vgg(input_spec, 16, weight_decay)
    tf.keras.utils.plot_model(model, "vgg16.png", True, show_layer_names=True)

    model = build_vgg(input_spec, 19, weight_decay)
    tf.keras.utils.plot_model(model, "vgg19.png", True, show_layer_names=True)
