Update trainer.py

trainer.py

@@ -13,7 +13,7 @@ from tqdm import tqdm
 sac = os.getenv('accesstoken')
 sn.set(font_scale=1.4)
 
-class_names = ['buildings', 'forest', 'glacier', 'mountain', 'sea', 'street']
+class_names = ['buildings', 'forest', 'glacier']
 class_names_label = {class_name: i for i, class_name in enumerate(class_names)}
 nb_classes = len(class_names)
 print(class_names_label)
@@ -53,32 +53,32 @@ def load_data():
 train_images, train_labels = shuffle(train_images, train_labels, random_state=25)
 print("Train: ", train_images.shape, train_labels.shape)
 print("Test: ", test_images.shape, test_labels.shape)
-#
-# model = tf.keras.models.Sequential([
-#     tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(150, 150, 3)),
-#     tf.keras.layers.MaxPooling2D(2, 2),
-#     tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
-#     tf.keras.layers.MaxPooling2D(2, 2),
-#     tf.keras.layers.Conv2D(128, (3, 3), activation='relu'),
-#     tf.keras.layers.MaxPooling2D(2, 2),
-#     tf.keras.layers.Conv2D(128, (3, 3), activation='relu'),
-#     tf.keras.layers.MaxPooling2D(2, 2),
-#     tf.keras.layers.Flatten(),
-#     tf.keras.layers.Dense(512, activation='relu'),
-#     tf.keras.layers.Dense(6, activation='softmax')
-# ])
-model = tf.keras.Sequential([
+
+model = tf.keras.models.Sequential([
     tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(150, 150, 3)),
     tf.keras.layers.MaxPooling2D(2, 2),
-    tf.keras.layers.Conv2D(32, (3, 3), activation='relu'),
+    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
+    tf.keras.layers.MaxPooling2D(2, 2),
+    tf.keras.layers.Conv2D(128, (3, 3), activation='relu'),
+    tf.keras.layers.MaxPooling2D(2, 2),
+    tf.keras.layers.Conv2D(128, (3, 3), activation='relu'),
     tf.keras.layers.MaxPooling2D(2, 2),
     tf.keras.layers.Flatten(),
-    tf.keras.layers.Dense(128, activation=tf.nn.relu),
-    tf.keras.layers.Dense(6, activation=tf.nn.softmax)
+    tf.keras.layers.Dense(512, activation='relu'),
+    tf.keras.layers.Dense(6, activation='softmax')
 ])
+# model = tf.keras.Sequential([
+#     tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(150, 150, 3)),
+#     tf.keras.layers.MaxPooling2D(2, 2),
+#     tf.keras.layers.Conv2D(32, (3, 3), activation='relu'),
+#     tf.keras.layers.MaxPooling2D(2, 2),
+#     tf.keras.layers.Flatten(),
+#     tf.keras.layers.Dense(128, activation=tf.nn.relu),
+#     tf.keras.layers.Dense(6, activation=tf.nn.softmax)
+# ])
 
 model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
-model.fit(train_images, train_labels, epochs=6, validation_split=0.2)
+model.fit(train_images, train_labels, epochs=10, validation_split=0.2)
 
 # Evaluate the model
 model.evaluate(test_images, test_labels)
@@ -93,7 +93,7 @@ model.save("model.keras")
 # Upload the model to your Hugging Face space repository
 push_to_hub_keras(
     model,
-    repo_id="okeowo1014/imgclassifiertraining",
+    repo_id="okeowo1014/imgclassifiera",
     commit_message="Optional commit message",
     tags=["image-classifier", "some_other_tag"],
     include_optimizer=True, token=sac