STM32Cube.AI

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from tensorflow import keras
from tensorflow.keras import layers
from sklearn import metrics

batch_size = 128
epochs = 15
num_classes = 10
input_shape = (28, 28, 1)

def prepare_data():
    (x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

    x_train = x_train.astype("float32") / 255
    x_test = x_test.astype("float32") / 255

    x_train = np.expand_dims(x_train, -1)
    x_test = np.expand_dims(x_test, -1)

    print("x_train shape:", x_train.shape)
    print(x_train.shape[0], "train samples")
    print(x_test.shape[0], "test samples")

    y_train = keras.utils.to_categorical(y_train, num_classes)
    y_test = keras.utils.to_categorical(y_test, num_classes)

    return x_train, x_test, y_train, y_test

def build_model():
    model = keras.Sequential(
        [
            layers.Conv2D(32, kernel_size=3, activation="relu", padding='same', input_shape=input_shape),
            layers.MaxPooling2D(pool_size=2),
            layers.Conv2D(64, kernel_size=3, activation="relu", padding='same'),
            layers.MaxPooling2D(pool_size=2),
            layers.Flatten(),
            layers.Dropout(0.5),
            layers.Dense(10, activation="softmax"),
        ]
    )
    model.summary()
    model.compile(loss="sparse_categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
    return model

def train_model(model, x_train, x_test, y_train, y_test):
    model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
    history = model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_split=0.3)

    test_loss, test_acc = model.evaluate(x_test, y_test)
    print('Test accuracy:', test_acc, " Test loss:", test_loss, "\n")
    model.save('model.h5')

    plt.plot(history.history['acc'])
    plt.plot(history.history['val_acc'])
    plt.title('Accuracy Trend')
    plt.ylabel('accuracy')
    plt.xlabel('epoch')
    plt.legend(['train', 'validation'], loc='best')
    plt.grid()
    plt.savefig('_accuracy_function.png')
    plt.show()

    plt.plot(history.history['loss'])
    plt.plot(history.history['val_loss'])
    plt.title('Loss Trend')
    plt.ylabel('loss')
    plt.xlabel('epoch')
    plt.legend(['train', 'validation'], loc='best')
    plt.grid()
    plt.savefig('_loss_function.png')
    plt.show()

def make_confusion_matrix(model, dataX, dataY):
    predictions = model.predict(dataX)

    predictedClassCode = np.argmax(predictions, axis=1)
    groundTruthClassCode = np.argmax(dataY, axis=1)

    cM = metrics.confusion_matrix(groundTruthClassCode, predictedClassCode)

    cmNormalized = np.around((cM / cM.sum(axis=1)[:, None]) * 100, 2)
    title = 'MNIST confusion matrix'

    fig, ax = plt.subplots(figsize=(7, 7))

    im = ax.imshow(cmNormalized, interpolation='nearest', cmap=plt.cm.Blues)

    plt.xlabel('Predicted Values')
    plt.ylabel('Ground Truth')
    plt.title(title)

    classes = [str(x) for x in range(0, 10)]
    plt.xticks(range(len(classes)), classes)
    plt.yticks(range(len(classes)), classes)
    plt.ylim(bottom=len(classes) - .5)
    plt.ylim(top=-0.5)

    width, height = cM.shape

    print('Accuracy for each class is given below.')
    for predicted in range(width):
        for real in range(height):
            if cmNormalized[predicted, real] > 45:
                color = 'white'
            else:
                color = 'black'
            if (predicted == real):
                print(classes[predicted].ljust(12) + ':', cmNormalized[predicted, real], '%')
            plt.gca().annotate('{}\n{}%'.format(cM[predicted, real], cmNormalized[predicted, real]),
                               xy=(real, predicted),
                               horizontalalignment='center', verticalalignment='center', color=color)
    divider = make_axes_locatable(plt.gca())
    cax = divider.append_axes("right", "5%", pad="3%")
    plt.colorbar(im, cax=cax)

    plt.tight_layout()

    plt.savefig('_confusion_matrix.png')
    plt.show()

x_train, x_test, y_train, y_test = prepare_data()
model = build_model()
train_model(model, x_train, x_test, y_train, y_test)
make_confusion_matrix(model, x_test, y_test)