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