Update utils.py

utils.py

@@ -25,26 +25,18 @@ def preprocess(image):
     return image
 
 def model_arc():
-    model = Sequential()
-
-    # Convolution blocks
-    model.add(Conv2D(32, kernel_size=(3, 3), padding='same', input_shape=(256, 256, 3), activation='relu'))
-    model.add(MaxPooling2D(pool_size=2))
-
-    model.add(Conv2D(64, kernel_size=(3, 3), padding='same', activation='relu'))
-    model.add(MaxPooling2D(pool_size=2))
-
-    model.add(Conv2D(32, kernel_size=(3, 3), padding='same', activation='relu'))
-    model.add(MaxPooling2D(pool_size=2))
-
-    # Classification layers
-    model.add(Flatten())
-
-    model.add(Dense(64, activation='relu'))
-    model.add(Dropout(0.2))
-    model.add(Dense(32, activation='relu'))
-
-    model.add(Dropout(0.2))
-    model.add(Dense(6, activation='softmax'))
+    model = tf.keras.Sequential([
+        data_augmentation,
+        base_model,
+        tf.keras.layers.GlobalAveragePooling2D(),
+        tf.keras.layers.Dense(6, activation='softmax')
+    ])
+
+    learning_rate = 0.00001
+    model.compile(
+        loss='sparse_categorical_crossentropy',
+        optimizer=tf.keras.optimizers.Adam(learning_rate),
+        metrics=['accuracy']
+    )
 
     return model