import tensorflow as tf
from tensorflow.keras import layers

input_shape = (448, 448, 3)
def yolov1_model(input_shape):
    inputs = layers.Input(shape=input_shape)
    x = (layers.Conv2D(64, (7, 7), strides=(2, 2), padding='same', input_shape=input_shape))(inputs)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)

    x = (layers.MaxPooling2D())(x)

    x = (layers.Conv2D(192, (3, 3), padding='same'))(x)
    x = x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)

    x = (layers.MaxPooling2D())(x)

    x = (layers.Conv2D(128, (1, 1), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(256, (3, 3), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(256, (1, 1), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(512, (3, 3), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)

    x = (layers.MaxPooling2D())(x)

    x = (layers.Conv2D(256, (1, 1), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(512, (3, 3), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(256, (1, 1), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(512, (3, 3), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(256, (1, 1), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(512, (3, 3), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(256, (1, 1), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(512, (3, 3), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(512, (1, 1), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(1024, (3, 3), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)

    x = (layers.MaxPooling2D())(x)

    x = (layers.Conv2D(512, (1, 1), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(1024, (3, 3), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(512, (1, 1), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(1024, (3, 3), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(1024, (3, 3), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(1024, (3, 3), padding='same', strides=(2, 2)))(x)
    x = x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)

    x = (layers.Conv2D(1024, (3, 3), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = (layers.Conv2D(1024, (3, 3), padding='same'))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)

    x = (layers.GlobalAveragePooling2D())(x)
    x = (layers.Dense(1024))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    x = x = (layers.Dense(4096))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)

    x = (layers.Dense(7 * 7 * 30))(x)
    x = (layers.BatchNormalization())(x)
    x = (layers.Activation('relu'))(x)
    outputs = layers.Reshape((7, 7, 30), name='outputs')(x)
    yolov1_model = tf.keras.models.Model(inputs=inputs, outputs=outputs)
    return yolov1_model

yolov1_model(input_shape).summary()



