📚 课程目录第4章使用STM32Cube.AI部署图像分类模型

4. 使用STM32Cube.AI部署图像分类模型

任务目标

本案例将完成一个完整的图像分类部署闭环：使用Keras训练一个MobileNetV2模型（基于ImageNet预训练权重），通过STM32Cube.AI将其转换为STM32可用的C代码，集成到STM32CubeIDE工程中，最终在开发板上对输入图像进行分类并输出Top-1类别索引。

环境与依赖

软件：STM32CubeIDE（≥1.10）、STM32Cube.AI（作为IDE插件，≥8.1）、Python 3.8+（用于模型转换脚本）
硬件：STM32H743系列开发板（推荐Nucleo-H743ZI2，2MB Flash/1MB RAM）
Python库：tensorflow==2.12.0、numpy、Pillow

模型准备与转换

由于完整ImageNet训练在MCU上不可行，本案例采用预训练MobileNetV2（输入224×224×3），通过迁移学习在10类自定义数据集上微调。转换时使用STM32Cube.AI的Python API生成量化后的C代码，避免在IDE中手动配置。

# prepare_model.py — 在PC上运行，生成STM32Cube.AI所需的.h5模型
import tensorflow as tf
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
from tensorflow.keras.models import Model

# 1. 构建迁移学习模型（10类输出）
base = MobileNetV2(weights='imagenet', include_top=False, input_shape=(224,224,3))
x = base.output
x = GlobalAveragePooling2D()(x)
output = Dense(10, activation='softmax')(x)
model = Model(inputs=base.input, outputs=output)

# 2. 冻结预训练层，仅训练分类头（演示用，实际需加载训练好的权重）
model.compile(optimizer='adam', loss='categorical_crossentropy')
model.save('mobilenetv2_10class.h5')
print("模型已保存，输入尺寸: 224x224x3, 输出类别数: 10")

# 3. 调用STM32Cube.AI CLI生成C代码（需先安装stm32ai命令行工具）
import subprocess
cmd = [
    'stm32ai', 'generate',
    '-m', 'mobilenetv2_10class.h5',
    '--name', 'image_classifier',
    '--output', 'stm32_ai_output',
    '--workspace', '.',
    '--compress', '1.0'  # 不压缩，保持精度
]
subprocess.run(cmd, check=True)
print("C代码已生成至 stm32_ai_output/ 目录")

STM32CubeIDE工程集成

在STM32CubeIDE中新建H743工程，将上述生成的stm32_ai_output文件夹复制到工程根目录。在main.c中调用AI运行时API完成推理。核心代码如下：

/* main.c — 在STM32CubeIDE中编译运行 */
#include "app_x-cube-ai.h"
#include "image_classifier.h"

/* 输入输出缓冲区（由STM32Cube.AI生成的头文件定义） */
static ai_handle network = AI_HANDLE_NULL;
static ai_float in_data[1*224*224*3];
static ai_float out_data[1*10];

/* 将摄像头或SD卡中的RGB888图像缩放至224x224并归一化到[0,1] */
void preprocess_image(uint8_t *src_rgb, uint32_t src_w, uint32_t src_h,
                      ai_float *dst, uint32_t dst_size) {
    for (uint32_t y = 0; y < 224; y++) {
        for (uint32_t x = 0; x < 224; x++) {
            uint32_t sx = x * src_w / 224;
            uint32_t sy = y * src_h / 224;
            uint32_t si = (sy * src_w + sx) * 3;
            uint32_t di = (y * 224 + x) * 3;
            dst[di+0] = src_rgb[si+0] / 255.0f;
            dst[di+1] = src_rgb[si+1] / 255.0f;
            dst[di+2] = src_rgb[si+2] / 255.0f;
        }
    }
}

int main(void) {
    HAL_Init();
    SystemClock_Config();  /* 配置为H743最高主频480MHz */

    /* 1. 初始化AI网络 */
    ai_network_create(&network, NULL);
    ai_network_init(network, NULL);

    /* 2. 准备输入图像（此处用模拟数据代替真实摄像头输入） */
    static uint8_t dummy_image[224*224*3];
    for (int i = 0; i < 224*224*3; i++) dummy_image[i] = (uint8_t)(i % 255);
    preprocess_image(dummy_image, 224, 224, in_data, 224*224*3);

    /* 3. 执行推理 */
    ai_input in = {0};
    ai_output out = {0};
    in.data = in_data;
    out.data = out_data;
    ai_run(network, &in, &out);

    /* 4. 获取Top-1类别 */
    uint32_t max_idx = 0;
    ai_float max_val = out_data[0];
    for (int i = 1; i < 10; i++) {
        if (out_data[i] > max_val) {
            max_val = out_data[i];
            max_idx = i;
        }
    }

    /* 5. 通过UART输出结果（115200 8N1） */
    printf("Predicted class: %lu, confidence: %.3f\r\n", max_idx, max_val);

    while (1) { /* 空循环，实际可接LED显示类别 */ }
}

运行与验证

编译：在STM32CubeIDE中编译工程，确认无错误。需在工程属性中开启FPU（Cortex-M7双精度）和优化等级-O2。
烧录：通过ST-Link下载至Nucleo-H743ZI2开发板。
验证方法：打开串口终端（115200），复位开发板。观察串口输出是否包含Predicted class: x行。
通过条件：程序正常输出类别索引（0-9之间），且无HardFault或卡死。由于使用模拟输入数据，类别值本身无实际意义，但输出格式正确且置信度在0到1之间即证明部署链路完整。

常见问题

内存不足：若编译报错Flash/RAM超限，在STM32Cube.AI生成时使用--compress 0.5或改用MobileNetV1（更小）。
推理结果全为0：检查输入归一化是否正确，以及ai_run返回值是否为0（成功）。
性能不足：确认已使能指令缓存（ICache）和数据缓存（DCache），并开启FPU。

本案例验证了从Keras模型到STM32部署的完整工具链。实际项目中需将dummy_image替换为摄像头或SD卡读取的真实图像，并加载训练好的权重文件。串口输出的类别索引与置信度可作为功能正确性的基本判据。