Chapter 1.2 - Adapt and move the Jupyter Notebook to Python scripts¶

import json
import sys
from pathlib import Path
from typing import List

import keras
import matplotlib.pyplot as plt
import tensorflow as tf
import yaml
from utils.seed import set_seed


def get_preview_plot(ds: tf.data.Dataset, labels: List[str]) -> plt.Figure:
    """Plot a preview of the prepared dataset"""
    fig, axes = plt.subplots(2, 5, figsize=(10, 5), tight_layout=True)
    for images, label_idxs in ds.take(1):
        for ax, image, label_idx in zip(axes.ravel(), images, label_idxs):
            ax.imshow(image.numpy().astype("uint8"), cmap="gray")
            ax.set_title(labels[label_idx.numpy()])
            ax.set_xticks([])
            ax.set_yticks([])

    return fig


def main() -> None:
    if len(sys.argv) != 3:
        print("Arguments error. Usage:\n")
        print("\tpython3 prepare.py <raw-dataset-folder> <prepared-dataset-folder>\n")
        exit(1)

    # Load parameters
    prepare_params = yaml.safe_load(open("params.yaml"))["prepare"]

    raw_dataset_folder = Path(sys.argv[1])
    prepared_dataset_folder = Path(sys.argv[2])
    seed = prepare_params["seed"]
    split = prepare_params["split"]
    image_size = prepare_params["image_size"]
    grayscale = prepare_params["grayscale"]
    batch_size = prepare_params["batch_size"]

    # Set seed for reproducibility
    set_seed(seed)

    # Read data
    ds_train, ds_val = keras.utils.image_dataset_from_directory(
        raw_dataset_folder,
        labels="inferred",
        label_mode="int",
        color_mode="grayscale" if grayscale else "rgb",
        batch_size=batch_size,
        image_size=image_size,
        shuffle=True,
        seed=seed,
        validation_split=split,
        subset="both",
    )
    labels = ds_train.class_names

    prepared_dataset_folder.mkdir(parents=True, exist_ok=True)

    # Save the preview plot
    preview_plot = get_preview_plot(ds_train, labels)
    preview_plot.savefig(prepared_dataset_folder / "preview.png")

    # Normalize the data
    normalization_layer = keras.layers.Rescaling(1.0 / 255)
    ds_train = ds_train.map(lambda x, y: (normalization_layer(x), y))
    ds_val = ds_val.map(lambda x, y: (normalization_layer(x), y))

    # Save the prepared dataset
    with open(prepared_dataset_folder / "labels.json", "w") as f:
        json.dump(labels, f)
    tf.data.Dataset.save(ds_train, str(prepared_dataset_folder / "train"))
    tf.data.Dataset.save(ds_val, str(prepared_dataset_folder / "val"))

    print(f"\nDataset saved at {prepared_dataset_folder.absolute()}")


if __name__ == "__main__":
    main()

import sys
from pathlib import Path
from typing import Tuple

import keras
import numpy as np
import tensorflow as tf
import yaml

from utils.seed import set_seed


def get_model(
    image_shape: Tuple[int, int, int],
    conv_size: int,
    dense_size: int,
    output_classes: int,
) -> keras.Model:
    """Create a simple CNN model"""
    model = keras.models.Sequential(
        [
            keras.layers.Input(shape=image_shape),
            keras.layers.Conv2D(conv_size, (3, 3), activation="relu"),
            keras.layers.MaxPooling2D((3, 3)),
            keras.layers.Flatten(),
            keras.layers.Dense(dense_size, activation="relu"),
            keras.layers.Dense(output_classes),
        ]
    )
    return model


def main() -> None:
    if len(sys.argv) != 3:
        print("Arguments error. Usage:\n")
        print("\tpython3 train.py <prepared-dataset-folder> <model-folder>\n")
        exit(1)

    # Load parameters
    params = yaml.safe_load(open("params.yaml"))
    prepare_params = params["prepare"]
    train_params = params["train"]

    prepared_dataset_folder = Path(sys.argv[1])
    model_folder = Path(sys.argv[2])

    image_size = prepare_params["image_size"]
    grayscale = prepare_params["grayscale"]
    image_shape = (*image_size, 1 if grayscale else 3)

    seed = train_params["seed"]
    lr = train_params["lr"]
    epochs = train_params["epochs"]
    conv_size = train_params["conv_size"]
    dense_size = train_params["dense_size"]
    output_classes = train_params["output_classes"]

    # Set seed for reproducibility
    set_seed(seed)

    # Load the prepared datasets and shuffle the training set at each epoch
    ds_train = tf.data.Dataset.load(str(prepared_dataset_folder / "train"))
    ds_train = ds_train.shuffle(
        buffer_size=ds_train.cardinality(), seed=seed, reshuffle_each_iteration=True
    )
    ds_val = tf.data.Dataset.load(str(prepared_dataset_folder / "val"))

    # Define the model
    model = get_model(image_shape, conv_size, dense_size, output_classes)
    model.compile(
        optimizer=keras.optimizers.Adam(lr),
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=[keras.metrics.SparseCategoricalAccuracy()],
    )
    model.summary()

    # Train the model
    model.fit(
        ds_train,
        epochs=epochs,
        validation_data=ds_val,
    )

    # Save the model
    model_folder.mkdir(parents=True, exist_ok=True)
    model_path = model_folder.absolute() / "model.keras"
    model.save(model_path)

    # Save the model history
    np.save(model_folder.absolute() / "history.npy", model.history.history)

    print(f"\nModel saved at {model_folder.absolute()}")


if __name__ == "__main__":
    main()

import json
import sys
from pathlib import Path
from typing import List

import keras
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from sklearn.metrics import (
    ConfusionMatrixDisplay,
    f1_score,
    precision_score,
    recall_score,
)


def get_training_plot(model_history: dict) -> plt.Figure:
    """Plot the training and validation loss"""
    epochs = range(1, len(model_history["loss"]) + 1)

    fig = plt.figure(figsize=(10, 4))
    plt.plot(epochs, model_history["loss"], label="Training loss")
    plt.plot(epochs, model_history["val_loss"], label="Validation loss")
    plt.xticks(epochs)
    plt.title("Training and validation loss")
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.legend()
    plt.grid(True)

    return fig


def get_pred_preview_plot(
    model: keras.Model, ds_val: tf.data.Dataset, labels: List[str]
) -> plt.Figure:
    """Plot a preview of the predictions"""
    fig, axes = plt.subplots(2, 5, figsize=(10, 5), tight_layout=True)
    for images, label_idxs in ds_val.take(1):
        pred_idxs = np.argmax(model.predict(images, verbose=0), axis=1)
        for ax, image, true_idx, pred_idx in zip(
            axes.ravel(), images, label_idxs, pred_idxs
        ):
            true_label = labels[true_idx.numpy()]
            pred_label = labels[pred_idx]
            ax.imshow(image.numpy().squeeze(), cmap="gray")
            ax.set_title(f"True: {true_label}\nPred: {pred_label}")
            ax.set_xticks([])
            ax.set_yticks([])

            border_color = "lime" if true_idx.numpy() == pred_idx else "red"
            for spine in ax.spines.values():
                spine.set_edgecolor(border_color)
                spine.set_linewidth(4)

    return fig


def get_confusion_matrix_plot(
    y_true: np.ndarray, y_pred: np.ndarray, labels: List[str]
) -> plt.Figure:
    """Plot the confusion matrix"""
    fig, ax = plt.subplots(figsize=(6, 6), tight_layout=True)
    display = ConfusionMatrixDisplay.from_predictions(
        y_true,
        y_pred,
        display_labels=labels,
        normalize="true",
        cmap="Blues",
        values_format=".2f",
        ax=ax,
        colorbar=True,
    )

    for value, text in zip(display.confusion_matrix.ravel(), display.text_.ravel()):
        text.set_fontsize(7)
        if np.isclose(value, 0.0):
            text.set_color("lightgray")

    ax.set_xticklabels(labels, rotation=90)
    ax.set_title("Validation confusion matrix")

    return fig


def main() -> None:
    if len(sys.argv) != 3:
        print("Arguments error. Usage:\n")
        print("\tpython3 evaluate.py <model-folder> <prepared-dataset-folder>\n")
        exit(1)

    model_folder = Path(sys.argv[1])
    prepared_dataset_folder = Path(sys.argv[2])
    evaluation_folder = Path("evaluation")
    plots_folder = Path("plots")

    # Create folders
    (evaluation_folder / plots_folder).mkdir(parents=True, exist_ok=True)

    # Load files
    ds_val = tf.data.Dataset.load(str(prepared_dataset_folder / "val"))
    with open(prepared_dataset_folder / "labels.json") as f:
        labels = json.load(f)

    # Load model
    model_path = model_folder.absolute() / "model.keras"
    model = keras.models.load_model(model_path)
    model_history = np.load(
        model_folder.absolute() / "history.npy", allow_pickle=True
    ).item()

    # Log metrics
    val_loss, val_acc = model.evaluate(ds_val)
    preds = model.predict(ds_val)
    y_true = tf.concat([y for _, y in ds_val], axis=0).numpy()
    y_pred = np.argmax(preds, axis=1)

    metrics = {
        "val_loss": val_loss,
        "val_acc": val_acc,
        "precision": precision_score(y_true, y_pred, average="macro", zero_division=0),
        "recall": recall_score(y_true, y_pred, average="macro", zero_division=0),
        "f1_score": f1_score(y_true, y_pred, average="macro", zero_division=0),
    }

    print(f"Validation loss: {metrics['val_loss']:.2f}")
    print(f"Validation accuracy: {metrics['val_acc'] * 100:.2f}%")
    print(f"Precision: {metrics['precision']:.2f}")
    print(f"Recall:    {metrics['recall']:.2f}")
    print(f"F1 score:  {metrics['f1_score']:.2f}")

    with open(evaluation_folder / "metrics.json", "w") as f:
        json.dump(metrics, f)

    # Save training history plot
    fig = get_training_plot(model_history)
    fig.savefig(evaluation_folder / plots_folder / "training_history.png")

    # Save predictions preview plot
    fig = get_pred_preview_plot(model, ds_val, labels)
    fig.savefig(evaluation_folder / plots_folder / "pred_preview.png")

    # Save confusion matrix plot
    fig = get_confusion_matrix_plot(y_true, y_pred, labels)
    fig.savefig(evaluation_folder / plots_folder / "confusion_matrix.png")

    print(
        f"\nEvaluation metrics and plot files saved at {evaluation_folder.absolute()}"
    )


if __name__ == "__main__":
    main()

import os

import keras
import tensorflow as tf


def set_seed(seed: int) -> None:
    """Set all random seeds and enable deterministic operations"""
    os.environ["PYTHONHASHSEED"] = str(seed)

    keras.utils.set_random_seed(seed)
    tf.config.experimental.enable_op_determinism()

    tf.config.threading.set_inter_op_parallelism_threads(1)
    tf.config.threading.set_intra_op_parallelism_threads(1)