AI 模型的训练与转换（在 PC 端完成）

Intel RealSense D435i│双目图像 IMU数据采集│▼AI图像质量评估MobileNetV3│┌──────┼──────┐│ │ │清晰度 浑浊度 折射程度│ │ │└──────┼──────┘▼自适应参数优化(Gamma、CLAHE、去雾增强、折射校正、特征提取参数优化)│▼ORB-SLAM3 Stereo-Inertial│特征提取 → 位姿估计 → 回环检测│▼实时定位与三维地图构建一、配置环境conda create -n slam_ai python3.9 conda activate slam_ai # 安装 PyTorch (如果你的电脑没有 Nvidia 显卡请去掉 --index-url 及其后面的内容安装 CPU 版) pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 pip install pandas opencv-python pip install pyrealsense2 pip install onnx二、编写数据采集脚本写一个小工具让你能在电脑前看着摄像头画面按一下键盘就保存一张图片。在你的电脑上新建一个文件叫collect_data.pyimport pyrealsense2 as rs import numpy as np import cv2 import os # 创建保存图片的文件夹 save_dir dataset if not os.path.exists(save_dir): os.makedirs(save_dir) # 配置深度和彩色数据流 pipeline rs.pipeline() config rs.config() # 启用左目红外(用于双目)和彩色图像分辨率设为 640x480 config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30) config.enable_stream(rs.stream.infrared, 1, 640, 480, rs.format.y8, 30) print(正在启动 RealSense 相机...) pipeline.start(config) print(启动成功按下 s 键保存当前帧按下 q 键退出。) img_count 0 try: while True: # 等待一对连贯的帧 frames pipeline.wait_for_frames() color_frame frames.get_color_frame() ir_left_frame frames.get_infrared_frame(1) if not color_frame or not ir_left_frame: continue # 将图像转换为 numpy 数组 color_image np.asanyarray(color_frame.get_data()) ir_left_image np.asanyarray(ir_left_frame.get_data()) # 显示彩色图像 cv2.imshow(RealSense Camera - Press S to Save, Q to Quit, color_image) key cv2.waitKey(1) if key 0xFF ord(s): # 保存图像 img_name os.path.join(save_dir, fframe_{img_count:04d}.jpg) cv2.imwrite(img_name, color_image) print(f已保存: {img_name}) img_count 1 elif key 0xFF ord(q): break finally: pipeline.stop() cv2.destroyAllWindows()运行脚本python3 collect_data.py三、魔改 MobileNetV3 网络结构 (Python 代码)MobileNetV3 原本是用来做“分类”的比如识别猫狗1000类我们需要把它改成“回归”任务输出3个具体的分数。新建一个model.py文件import torch import torch.nn as nn from torchvision.models import mobilenet_v3_small, MobileNet_V3_Small_Weights class ImageQualityNet(nn.Module): def __init__(self): super().__init__() # 1. 加载官方预训练模型 (利用迁移学习加速收敛) self.model mobilenet_v3_small(weightsMobileNet_V3_Small_Weights.DEFAULT) # 2. 找到最后一层分类器并获取其输入维度 in_features self.model.classifier[3].in_features # 3. 将输出维度从 1000 改为 3 (清晰度, 浑浊度, 折射程度) self.model.classifier[3] nn.Linear(in_features, 3) def forward(self, x): # 使用 Sigmoid 激活函数强制让输出结果落在 0.0 到 1.0 之间 return torch.sigmoid(self.model(x)) # 测试一下模型是否正常工作 if __name__ __main__: net ImageQualityNet() dummy_image torch.randn(1, 3, 224, 224) # 模拟一张 224x224 的输入图片 output net(dummy_image) print(模型输出:, output) # 应该输出形如 tensor([[0.51, 0.49, 0.55]]) 的数据四、创建标注文件在项目根目录下也就是和dataset文件夹平级的地方新建一个文本文档命名为labels.csv。用 Excel 或直接用文本编辑器打开它按照下面的格式填入你抓拍的图片名字并凭感觉给它们打个分分值在0.0到1.0之间image_name,clarity,turbidity,refraction frame_0000.jpg,0.9,0.1,0.0 frame_0001.jpg,0.4,0.7,0.2 frame_0002.jpg,0.8,0.2,0.6第一行是表头千万别写错。后面每一行对应一张图的清晰度、浑浊度、折射度。五、编写 PyTorch 数据读取器 (dataset.py)有了图片和 CSV 文件我们需要写一段代码让 PyTorch 能够自动读取它们把图片缩放到 MobileNetV3 要求的224x224分辨率并转换成神经网络认识的张量Tensor。在项目目录下新建dataset.pyimport os import pandas as pd import cv2 import torch from torch.utils.data import Dataset from torchvision import transforms class WaterQualityDataset(Dataset): def __init__(self, csv_file, img_dir): # 读取 CSV 标注文件 self.annotations pd.read_csv(csv_file) self.img_dir img_dir # 定义图像预处理步骤缩放到 224x224并归一化 self.transform transforms.Compose([ transforms.ToPILImage(), transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) def __len__(self): return len(self.annotations) def __getitem__(self, index): # 1. 获取图片路径 img_id self.annotations.iloc[index, 0] img_path os.path.join(self.img_dir, img_id) # 2. 读取图片 (OpenCV 默认是 BGR我们需要转成 RGB) image cv2.imread(img_path) image cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 3. 获取对应的三个标签值 labels self.annotations.iloc[index, 1:].values.astype(float32) labels torch.tensor(labels) # 4. 应用图像预处理 if self.transform: image self.transform(image) return image, labels # 简易测试代码 if __name__ __main__: dataset WaterQualityDataset(csv_filelabels.csv, img_dirdataset) if len(dataset) 0: img, lab dataset[0] print(单张图片张量形状:, img.shape) # 应该是 [3, 224, 224] print(对应的标签数据:, lab) # 应该是三个分数的 tensor else: print(请先在 labels.csv 中添加数据并在 dataset 文件夹中放入图片)六、完善并运行训练脚本 (train.py)现在我们把模型结构之前给你的model.py、数据读取器刚写的dataset.py和训练逻辑全部拼装在一起。在项目目录下新建train.pyimport torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from model import ImageQualityNet from dataset import WaterQualityDataset def main(): # 1. 硬件检测 (有 GPU 用 GPU没有就用 CPU) device torch.device(cuda if torch.cuda.is_available() else cpu) print(f当前使用的训练设备: {device}) # 2. 实例化数据集和数据加载器 (DataLoader) # batch_size 代表一次喂给神经网络几张图因为我们现在是测试先设为 2 dataset WaterQualityDataset(csv_filelabels.csv, img_dirdataset) train_loader DataLoader(dataset, batch_size2, shuffleTrue) # 3. 初始化网络模型并送入设备 model ImageQualityNet().to(device) # 4. 定义损失函数回归任务用 MSE和优化器 criterion nn.MSELoss() optimizer optim.Adam(model.parameters(), lr0.001) # 5. 开始微型训练循环 epochs 10 # 先跑 10 轮试试看 print(开始训练...) for epoch in range(epochs): model.train() running_loss 0.0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) # 梯度清零 optimizer.zero_grad() # 前向传播 outputs model(images) # 计算误差 loss criterion(outputs, labels) # 反向传播与优化 loss.backward() optimizer.step() running_loss loss.item() print(fEpoch [{epoch1}/{epochs}], Loss: {running_loss/len(train_loader):.4f}) # 6. 保存模型权重 torch.save(model.state_dict(), best_quality_net.pth) print(模型训练完成权重已保存为 best_quality_net.pth) if __name__ __main__: main()确保你的项目文件夹里有 3 个.py文件model.py上一轮对话给你的网络结构、dataset.py和train.py。拍十几张图并新建labels.csv填入对应的名字和任意打的分数。尝试在终端运行python3 train.py运行python3 train.py时代码能够正常打印出Epoch [1/10], Loss: ...并在最后生成best_quality_net.pth文件七、把.pth变成.onnx在 PC 端新建一个export.py文件把我们之前提到的导出代码跑一下。这一步会把复杂的 PyTorch 网络结构“压扁”成一个通用的开放神经网络交换格式ONNX。import torch from model import ImageQualityNet # 1. 实例化模型并加载你刚刚训练好的权重 model ImageQualityNet() model.load_state_dict(torch.load(best_quality_net.pth, map_locationcpu)) model.eval() # 切换到推理模式 # 2. 创建一个虚拟的输入张量告诉 ONNX 输入图像的分辨率是 1x3x224x224 dummy_input torch.randn(1, 3, 224, 224) # 3. 导出模型 torch.onnx.export( model, dummy_input, quality_net.onnx, input_names[input], output_names[output], opset_version12 # 稳定的算子版本 ) print( ONNX 模型导出成功请检查目录下是否生成了 quality_net.onnx)运行python3 export.py只要看到打印出成功你就会在项目目录下看到一个quality_net.onnx文件。请把这个文件好好保存它就是我们要带去树莓派5的“核心资产”。八、迁移到树莓派上1.安装基础库# 更新系统包 sudo apt update sudo apt upgrade -y # 安装树莓派上的 OpenCV 和 ONNX Runtime pip install opencv-python numpy pyrealsense2 pip install onnxruntime pip3 install onnxruntime --break-system-packages2.模型送到了环境装好了现在我们要写一个树莓派上的实时测试程序让 RealSense D435i 实时捕获画面喂给quality_net.onnx并在窗口上实时把“清晰度、浑浊度、折射度”的分数打印出来把电脑上的quality_net.onnx文件通过 U盘 或网络如 SCP拷贝到树莓派 5 上。然后在树莓派上新建一个pi_run.py# -*- coding: utf-8 -*- import cv2 import numpy as np import onnxruntime as ort import time # 1. Load ONNX model print(Loading AI Quality Assessment Model...) session ort.InferenceSession(quality_net.onnx) input_name session.get_inputs()[0].name # 2. Start Camera using standard OpenCV (Bypassing pyrealsense2) # RealSense D435i usually maps RGB to index 2 or 4 on Linux. # If it shows black or infrared, change 2 to 0 or 4. camera_index 2 cap cv2.VideoCapture(camera_index) # Set resolution to 640x480 cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640) cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480) if not cap.isOpened(): print(fError: Could not open camera at index {camera_index}. Trying index 0...) cap cv2.VideoCapture(0) if not cap.isOpened(): print(Error: Fully failed to open RealSense via OpenCV. Check USB connection!) exit() print(System started successfully! Real-time evaluating...) try: while True: t_start time.time() # Get frame ret, frame cap.read() if not ret: print(Failed to grab frame) continue # 3. AI Preprocessing img cv2.resize(frame, (224, 224)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img img.astype(np.float32) / 255.0 # 显式指定为 float32防止 NumPy 自动将其转换为 double mean np.array([0.485, 0.456, 0.406], dtypenp.float32) std np.array([0.229, 0.224, 0.225], dtypenp.float32) img (img - mean) / std img np.transpose(img, (2, 0, 1)) img np.expand_dims(img, axis0) # 最终保险确保图像绝对是 float32 类型 img img.astype(np.float32) # 4. Inference outputs session.run(None, {input_name: img}) scores outputs[0][0] clarity, turbidity, refraction scores[0], scores[1], scores[2] fps 1.0 / (time.time() - t_start) # 5. Display results on frame cv2.putText(frame, fFPS: {fps:.1f}, (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) cv2.putText(frame, fClarity: {clarity:.2f}, (20, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2) cv2.putText(frame, fTurbidity: {turbidity:.2f}, (20, 120), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2) cv2.putText(frame, fRefraction: {refraction:.2f}, (20, 160), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2) cv2.imshow(Raspberry Pi 5 - AI Quality Assessment, frame) if cv2.waitKey(1) 0xFF ord(q): break finally: cap.release() cv2.destroyAllWindows()运行代码python3 pi_run.py可以看到3.核心任务把pi_run.py升级为 ROS2 节点我们需要写一个 Python 节点让它扮演一个“中间商”订阅RealSense 发布的彩色图像话题。利用cv_bridge把 ROS2 的图像消息转换成 OpenCV 矩阵。跑 AI 推理算出三个分数。发布评估结果为了简单省事我们先用 ROS2 自带的std_msgs/Float32MultiArray话题把三个分数打包发出去。请在你的树莓派上在orbslam3_ros2源码附近或者任意 ROS2 包里新建water_quality_node.py#!/usr/bin/env python3 # -*- coding: utf-8 -*- import rclpy from rclpy.node import Node from sensor_msgs.msg import Image from std_msgs.msg import Float32MultiArray from cv_bridge import CvBridge import cv2 import numpy as np import onnxruntime as ort class WaterQualityNode(Node): def __init__(self): super().__init__(water_quality_node) # 1. 初始化 ONNX 模型 self.get_logger().info(Loading ONNX Model...) self.session ort.InferenceSession(quality_net.onnx) self.input_name self.session.get_inputs()[0].name # 2. 初始化工具 self.bridge CvBridge() # 3. 创建订阅者订阅 RealSense 官方彩色图像 # 注意如果你的话题名不同请修改下方的 /camera/camera/color/image_raw self.image_sub self.create_subscription( Image, /camera/camera/color/image_raw, self.image_callback, 10 ) # 4. 创建发布者发布水质得分 [Clarity, Turbidity, Refraction] self.quality_pub self.create_publisher( Float32MultiArray, /camera/water_quality, 10 ) self.get_logger().info(AI Quality Node has started successfully!) def image_callback(self, msg): try: # 将 ROS2 Image 转换为 OpenCV BGR 图像 frame self.bridge.imgmsg_to_cv2(msg, desired_encodingbgr8) # --- AI 预处理 --- img cv2.resize(frame, (224, 224)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img img.astype(np.float32) / 255.0 mean np.array([0.485, 0.456, 0.406], dtypenp.float32) std np.array([0.229, 0.224, 0.225], dtypenp.float32) img (img - mean) / std img np.transpose(img, (2, 0, 1)) img np.expand_dims(img, axis0).astype(np.float32) # --- 运行推理 --- outputs self.session.run(None, {self.input_name: img}) scores outputs[0][0] # [clarity, turbidity, refraction] # --- 发布结果 --- quality_msg Float32MultiArray() quality_msg.data [float(scores[0]), float(scores[1]), float(scores[2])] self.quality_pub.publish(quality_msg) # 在终端打印一下方便肉眼观测 self.get_logger().info(fPublished Scores - C: {scores[0]:.2f}, T: {scores[1]:.2f}, R: {scores[2]:.2f}) except Exception as e: self.get_logger().error(fFailed to process image: {str(e)}) def main(argsNone): rclpy.init(argsargs) node WaterQualityNode() try: rclpy.spin(node) except KeyboardInterrupt: pass finally: node.destroy_node() rclpy.shutdown() if __name__ __main__: main()运行相机ros2 launch realsense2_camera rs_launch.py检查话题在另一个终端输入ros2 topic list确认驱动是否在发布/camera/camera/color/image_raw运行 AI 节点运行上面的 Python 脚本python3 water_quality_node.py。观察输出此时打开第三个终端输入ros2 topic echo /camera/water_quality能看到高频刷新的三个水质得分数据流了