Update train_tensorflow.py

train_tensorflow.py

@@ -1,9 +1,3 @@
-# This module trains the CNN based on the labels provided in ./data/CNN
-# Note that data must be first split into train, validation, and test data
-# by running split_data.py.
-# Reference:
-# https://towardsdatascience.com/a-single-function-to-streamline-image-classification-with-keras-bd04f5cfe6df
-
 from matplotlib import pyplot as plt
 from tensorflow.keras.preprocessing.image import ImageDataGenerator
 from tensorflow.keras.models import Sequential
@@ -148,10 +142,9 @@ def test_model(model):
     Does not return anything.'''
     testdir = DATA_FOLDER + 'test'
 
-    # pieces = ['Empty', 'Rook', 'Knight', 'Bishop', 'Queen', 'Pawn', 'King']
-    pieces = ['Empty', 'Rook_White', 'Rook_Black', 'Knight_White', 'Knight_Black', 'Bishop_White',
-              'Bishop_Black', 'Queen_White', 'Queen_Black', 'King_White', 'King_Black', 'Pawn_White', 'Pawn_Black']
-    pieces.sort()
+    # IMPORTANT: Ordre EXACT comme dans le générateur (alphabétique)
+    pieces = ['Bishop_Black', 'Bishop_White', 'Empty', 'King_Black', 'King_White', 'Knight_Black',
+              'Knight_White', 'Pawn_Black', 'Pawn_White', 'Queen_Black', 'Queen_White', 'Rook_Black', 'Rook_White']
     score = 0
     total_size = 0
     for subdir, dirs, files in os.walk(testdir):
@@ -160,7 +153,8 @@ def test_model(model):
                 continue
             piece = subdir.split('/')[-1]
             path = os.path.join(subdir, file)
-            y_prob = model.predict(cv2.imread(path).reshape(1, 300, 150, 3))
+            img = cv2.imread(path).reshape(1, 300, 150, 3) / 255.0
+            y_prob = model.predict(img, verbose=0)
             y_pred = y_prob.argmax()
             if y_pred < 0 or y_pred >= len(pieces):
                 print(y_pred, y_prob)
@@ -176,7 +170,7 @@ if __name__ == '__main__':
     history = fit_model(model, train_generator,
                         validation_generator, save=False)
     save_history(history, "./history.json")
-    plot_accuracy(history)
-    plot_loss(history)
+    # plot_accuracy(history)
+    # plot_loss(history)
     test_model(model)
-    model.save_weights('./model_weights.h5')
+    model.save_weights('./model_weights.weights.h5')