TinyML Lab: Simple MLP for STM32 TinyML Deployment with X-Cube-AI

				
					# The following lines are installation commands, they are not part of the core script logic
%pip install micromlgen
%pip install skl2onnx
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
# PCA is removed from imports
from sklearn.neural_network import MLPClassifier
from sklearn.metrics import accuracy_score, confusion_matrix, ConfusionMatrixDisplay
# --- External Libraries for Embedded Export ---
import onnx
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
from micromlgen import port # Kept for potential C export of the MLP itself
# 1. Load Data
iris = load_iris()
X = iris.data
y = iris.target
feature_names = iris.feature_names
target_names = iris.target_names

# 2. Standardize the data (Crucial for Neural Networks) - left out
#scaler = StandardScaler()
#X_scaled = scaler.fit_transform(X) # X_scaled now has 4 features

# --- VISUALIZATION: 3D Plot of Original Data (using 3 of 4 features) ---
# This plot remains useful to visualize the complexity of the 4D space.
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

for i, target_name in enumerate(target_names):
    # Using the first three features for 3D plot
    ax.scatter(X[y == i, 0], X[y == i, 1], X[y == i, 2], label=target_name)

ax.set_xlabel(feature_names[0])
ax.set_ylabel(feature_names[1])
ax.set_zlabel(feature_names[2])
ax.set_title(f'Iris Dataset (3 of 4 features) - Original Data')
ax.legend()
plt.tight_layout()
plt.show() # Display the plot
# -------------------------------------------------------------

# 4. Split Data (using the full 4 standardized features)
# NOTE: X and not X_scaled is used here
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, random_state=42, stratify=y
)

# 5. Train Simple Neural Network (MLP)
# The input layer is now implicitly 4 neurons (matching X_scaled)
mlp = MLPClassifier(
    hidden_layer_sizes=(10,),
    max_iter=1000,
    activation='relu',
    solver='adam',
    random_state=42
)
mlp.fit(X_train, y_train)

# 6. Evaluate
y_pred = mlp.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)

print("--- Results with Simple Neural Network (MLP) on 4 Features ---")
print(f"Classifier Accuracy (on 4 features): {accuracy:.4f}")
# PCA related print statements are removed
# -----------------------------------------

# --- VISUALIZATION: Confusion Matrix ---
cm = confusion_matrix(y_test, y_pred)
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=target_names)
fig, ax = plt.subplots(figsize=(8, 8))
disp.plot(ax=ax, cmap=plt.cm.Blues)
plt.title("Confusion Matrix for MLP (using 4 original features)")
plt.show() # Display the plot
# -----------------------------------------

# --- EXPORT 1: Neural Network Model to ONNX ---
try:
    # Define the input type: 4 input features (Sepal/Petal) with float type
    # IMPORTANT: Changed input size from 2 to 4
    initial_type = [('float_input', FloatTensorType([None, 4]))]

    # Export the MLP model
    onnx_model = convert_sklearn(mlp, initial_types=initial_type, target_opset=13)

    # Save to file
    with open("iris_mlp_4f.onnx", "wb") as f: # Renaming the output file
        f.write(onnx_model.SerializeToString())
    print("\n✅ Exported MLP model to ONNX format: 'iris_mlp_4f.onnx' (4 features)")

except Exception as e:
    print(f"\n⚠️ Skipping ONNX export: Error during conversion: {e}")

# --- EXPORT 2: MLP Model to C Code using micromlgen ---
# The original request was to export PCA, but since PCA is removed,
# we can use micromlgen to export the MLP model itself (which includes the Standard Scaler).
try:
    # Export the MLP model (it will include the Standard Scaler logic if provided)
    # Note: micromlgen is best suited for simpler models like Decision Trees,
    # but it can port the scaler and the MLP parameters.
    c_code = port(mlp, classmap={
        0: 'SETOSA',
        1: 'VERSICOLOR',
        2: 'VIRGINICA'
    }, scaler=scaler) # Passing the scaler ensures input data is standardized in C code

    # Save to C file
    with open('iris_mlp_model.c','w') as f:
        f.write(c_code)

    print("✅ Exported MLP model + Standard Scaler to C format: 'iris_mlp_model.c'")

except Exception as e:
    print(f"\n⚠️ Skipping C export: Error during conversion or import: {e}")