ml/models/cnn_tensorflow.py

import tensorflow as tf
from tensorflow.keras import layers, models

def build_model(num_classes=6, input_shape=(150, 150, 3)):
    inputs = layers.Input(shape=input_shape)

    def block(x, filters):
        x = layers.Conv2D(filters, 3, padding="same", activation=None)(x)
        x = layers.BatchNormalization()(x)
        x = layers.ReLU()(x)
        x = layers.MaxPooling2D()(x)
        return x

    x = block(inputs, 32)
    x = block(x, 64)
    x = block(x, 128)
    x = block(x, 256)

    x = layers.GlobalAveragePooling2D()(x)
    x = layers.Dropout(0.5)(x)

    outputs = layers.Dense(num_classes)(x)

    model = models.Model(inputs, outputs)

    model.compile(
        optimizer="adam",
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=["accuracy"]
    )

    return model