diff options
Diffstat (limited to 'train.py')
| -rw-r--r-- | train.py | 270 |
1 files changed, 270 insertions, 0 deletions
diff --git a/train.py b/train.py new file mode 100644 index 0000000..8ed60c6 --- /dev/null +++ b/train.py @@ -0,0 +1,270 @@ +import sys + +if len(sys.argv) < 3: + print("Run script as: train.py [DATASET PATH] [MODEL SAVE NAME]") + exit() + +import tensorflow as tf +from tensorflow import keras +from tensorflow.keras import layers +import matplotlib.pyplot as plt +import pathlib + +# --- Ange din dataset-mapp här --- +data_dir = pathlib.Path(f"{sys.argv[1]}/").with_suffix('') + +# --- Bildparametrar --- +img_height = 192 +img_width = 50 +batch_size = 32 + +# --- Ladda dataset från mappar --- +train_ds = tf.keras.preprocessing.image_dataset_from_directory( + data_dir, + validation_split=0.2, + subset="training", + seed=123, + color_mode="grayscale", + image_size=(img_height, img_width), + batch_size=batch_size, + labels="inferred", + label_mode="int", + verbose=True +) + +val_ds = tf.keras.preprocessing.image_dataset_from_directory( + data_dir, + validation_split=0.2, + subset="validation", + seed=123, + color_mode="grayscale", + image_size=(img_height, img_width), + batch_size=batch_size, + labels="inferred", + label_mode="int", + verbose=True +) + +image_count = len(list(data_dir.glob('*/*.png'))) + len(list(data_dir.glob('p/*/*.png'))) +print(f"Number of images: {image_count}") + +class_names = train_ds.class_names +num_classes = len(class_names) +print(f"Class names: {class_names}") + +plt.figure(figsize=(20, 10)) +for images, labels in train_ds.take(1): + for i in range(9): + ax = plt.subplot(3, 3, i + 1) + plt.imshow(images[i].numpy().astype("uint8"), interpolation='nearest', aspect='auto', cmap='gray') + plt.title(class_names[labels[i]]) + plt.axis("off") +plt.suptitle("Training images", fontsize=30) +plt.show() + +for image_batch, labels_batch in train_ds: + print(image_batch.shape) + print(labels_batch.shape) + break + + +# --- Pipeline-optimering --- +# AUTOTUNE = tf.data.AUTOTUNE +# train_ds = train_ds.cache().shuffle(1000).prefetch(AUTOTUNE) +# val_ds = val_ds.cache().prefetch(AUTOTUNE) + +# --- Data augmentation --- +data_augmentation = tf.keras.Sequential([ + layers.RandomFlip("horizontal"), + #layers.RandomZoom(0.2), + layers.RandomContrast(0.4) +]) + +plt.figure(figsize=(10, 10)) +for images, labels in train_ds.take(1): + for i in range(9): + # Add the image to a batch. + image = tf.cast(tf.expand_dims(images[i], 0), tf.float32) + + augmented_image = data_augmentation(image) + ax = plt.subplot(3, 3, i + 1) + plt.imshow(augmented_image[0], interpolation='nearest', aspect='auto', cmap='gray') + plt.title(class_names[labels[i]]) + plt.axis("off") +plt.suptitle("Augmented training images", fontsize=30) +plt.show() + +train_ds = train_ds.repeat(5).shuffle(1000) +train_ds = train_ds.map(lambda x, y: (data_augmentation(x, training=True), y)) + +# --- Modell --- +model = keras.Sequential([ + layers.Rescaling(1./255, input_shape=(img_height, img_width, 1), name="Input_image"), + # layers.Input(shape=(img_height, img_width, 1), name="Input_image"), + + layers.Conv2D(32, (3, 3), activation="relu"), + layers.MaxPooling2D(), + + layers.Conv2D(64, (3, 3), activation="relu"), + layers.MaxPooling2D(), + + layers.Conv2D(128, (3, 3), activation="relu"), + layers.MaxPooling2D(), + + layers.Flatten(), + layers.Dense(64, activation="relu"), + layers.Dense(num_classes, activation="softmax", name="Prediction") +]) + +earlystop = tf.keras.callbacks.EarlyStopping(monitor="val_loss", min_delta=0.1, patience=10, mode="min") + +model.compile( + optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001), + loss=tf.keras.losses.SparseCategoricalCrossentropy(reduction="sum"), + metrics=[tf.keras.metrics.SparseCategoricalAccuracy()], +) + +model.summary() + +keras.utils.plot_model(model, "model.png", show_shapes=True) + +# --- Train presence/no presence model --- +history = model.fit( + train_ds, + validation_data=val_ds, + epochs=500, + callbacks=[earlystop] +) + +# --- Plotta träning/validering --- +plt.figure() +plt.style.use('default') +plt.plot(history.history["sparse_categorical_accuracy"], label="Training") +plt.plot(history.history["val_sparse_categorical_accuracy"], label="Validation") +plt.xlabel("Epochs") +plt.ylabel("Accuracy") +plt.legend() +plt.show() + +model.save(f"{sys.argv[2]}.keras") + +# --- Activities --- + +activities_dir = pathlib.Path(f"{sys.argv[1]}/p/").with_suffix('') + +train_activities_ds = tf.keras.preprocessing.image_dataset_from_directory( + activities_dir, + validation_split=0.2, + subset="training", + seed=123, + color_mode="grayscale", + image_size=(img_height, img_width), + batch_size=batch_size, + labels="inferred", + label_mode="int", + verbose=True +) + +val_activities_ds = tf.keras.preprocessing.image_dataset_from_directory( + activities_dir, + validation_split=0.2, + subset="validation", + seed=123, + color_mode="grayscale", + image_size=(img_height, img_width), + batch_size=batch_size, + labels="inferred", + label_mode="int", + verbose=True +) + +image_count = len(list(data_dir.glob('p/*/*.png'))) +print(f"Number of images: {image_count}") + +class_names = train_activities_ds.class_names +num_classes = len(class_names) +print(f"Class names: {class_names}") + +plt.figure(figsize=(10, 10)) +for images, labels in train_activities_ds.take(1): + for i in range(9): + ax = plt.subplot(3, 3, i + 1) + plt.imshow(images[i].numpy().astype("uint8"), interpolation='nearest', aspect='auto', cmap='gray') + plt.title(class_names[labels[i]]) + plt.axis("off") +plt.suptitle("Training images", fontsize=30) +plt.show() + +for image_batch, labels_batch in train_activities_ds: + print(image_batch.shape) + print(labels_batch.shape) + break + +# --- Pipeline-optimering --- +# AUTOTUNE = tf.data.AUTOTUNE +# train_ds = train_ds.cache().shuffle(1000).prefetch(AUTOTUNE) +# val_ds = val_ds.cache().prefetch(AUTOTUNE) + +plt.figure(figsize=(10, 10)) +for images, labels in train_activities_ds.take(1): + for i in range(9): + # Add the image to a batch. + image = tf.cast(tf.expand_dims(images[i], 0), tf.float32) + + augmented_image = data_augmentation(image) + ax = plt.subplot(3, 3, i + 1) + plt.imshow(augmented_image[0], interpolation='nearest', aspect='auto', cmap='gray') + plt.title(class_names[labels[i]]) + plt.axis("off") +plt.suptitle("Augmented training images", fontsize=30) +plt.show() + +train_activities_ds = train_activities_ds.repeat(20).shuffle(1000) +train_activities_ds = train_activities_ds.map(lambda x, y: (data_augmentation(x, training=True), y)) + +# --- Modell --- +activity_model = keras.Sequential([ + layers.Rescaling(1./255, input_shape=(img_height, img_width, 1), name="Input_image"), + # layers.Input(shape=(img_height, img_width, 1), name="Input_image"), + + layers.Conv2D(32, (3, 3), activation="relu"), + layers.MaxPooling2D(), + + layers.Conv2D(64, (3, 3), activation="relu"), + layers.MaxPooling2D(), + + layers.Conv2D(128, (3, 3), activation="relu"), + layers.MaxPooling2D(), + + layers.Flatten(), + layers.Dense(64, activation="relu"), + layers.Dense(num_classes, activation="softmax", name="Prediction") +]) + +earlystop = tf.keras.callbacks.EarlyStopping(monitor="val_loss", min_delta=0.1, patience=10, mode="min") + +activity_model.compile( + optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001), + loss=tf.keras.losses.SparseCategoricalCrossentropy(reduction="sum"), + metrics=[tf.keras.metrics.SparseCategoricalAccuracy()], +) + +# --- Train activities model --- +history = activity_model.fit( + train_activities_ds, + validation_data=val_activities_ds, + epochs=500, + callbacks=[earlystop] +) + +# --- Plotta träning/validering --- +plt.figure() +plt.style.use('default') +plt.plot(history.history["sparse_categorical_accuracy"], label="Training") +plt.plot(history.history["val_sparse_categorical_accuracy"], label="Validation") +plt.xlabel("Epochs") +plt.ylabel("Accuracy") +plt.legend() +plt.show() + +activity_model.save(f"{sys.argv[2]}_activities.keras") |