简介:本资源为基于YOLOv8-OBB的芯片引脚缺陷检测完整项目,采用TensorRT进行推理加速,面向计算机、人工智能、电子信息等专业的在校学生、教师及企业开发者,可用于毕业设计、课程设计、项目立项或算法进阶学习。压缩包共394个文件,约4.7MB,以C++头文件与源文件为主,辅以CUDA核函数、ONNX解析与TensorRT部署相关代码,并包含少量说明文档、配置文件与预览图,整体结构清晰,便于按模块查阅与二次开发。项目已通过导师评审,答辩成绩达95分,代码经测试可正常运行。内容涵盖旋转框目标检测、模型转换、TensorRT推理加速及缺陷识别流程,读者可据此掌握从训练到部署的完整链路,并在此基础上修改以适配其他检测任务。目前已有63人学习下载,适合具备一定深度学习基础、希望深入理解工业缺陷检测与推理优化的读者参考使用。
1. 芯片引脚缺陷检测为什么值得用 YOLOv8-OBB 重做一遍
芯片引脚缺陷检测这个场景,传统做法是模板匹配加形态学,规则写得越多越脆。引脚歪一点、氧化一点、背景光变一点,阈值就得重调,产线换批次就是一场玄学。YOLOv8-OBB 的价值在于把「框」换成「带角度的旋转框」,引脚这种细长目标用水平框标注会引入大量背景,OBB 直接贴合引脚走向,分类和定位一起学,泛化比手写规则稳得多。再叠上 TensorRT 加速,单张推理从几十毫秒压到几毫秒,才够得上产线节拍。这套「源码+文档+全部资料」的组合,适合两类人:一类是想把检测模型真正推到工控机或边缘盒子上跑起来的算法工程师,另一类是做课程设计、毕业设计需要一份能跑通、能讲清楚全流程的完整工程的人。下面我按自己落地的顺序,把数据、训练、导出、加速、排错一条线讲透。
2. 从标注到 OBB 数据集:引脚检测的数据准备与格式转换
2.1 为什么引脚必须用旋转框而不是水平框
引脚在芯片图像里是细长条,长宽比经常到 8:1 甚至更高。用水平框标注,框里塞进去的绝大部分是背景和相邻引脚,模型学到的特征被稀释,密集引脚还会出现框重叠、NMS 互相压制。旋转框用中心点、宽、高、角度五个量描述目标,框和引脚几乎重合,分类头拿到的特征干净,回归头也不用在角度上瞎猜。
YOLOv8-OBB 的角度定义是重点,踩过坑的都懂。它用的是「角度在 [0, 90) 区间、宽高按长边短边归一」的表示,也就是 DOTA 那套定义。你自己标数据时如果角度从 -90 到 90 随便给,训练时 loss 会震荡,mAP 上不去。常见做法是标注阶段就统一:让宽始终对应较长的那条边,角度取锐角。标注工具用 roLabelImg 或 X-AnyLabeling 都行,导出格式选 DOTA 或 YOLO-OBB。
2.2 DOTA 转 YOLO-OBB 的转换脚本与四个边界坑
DOTA 格式是每行八个坐标加类别名,YOLO-OBB 是class x y w h angle归一化。转换脚本不长,但边界情况特别多。
import os import cv2 import numpy as np def dota_to_yolo_obb(dota_line, img_w, img_h, class_map): parts = dota_line.strip().split() # DOTA: x1 y1 x2 y2 x3 y3 x4 y4 class_name difficult coords = np.array(parts[:8], dtype=np.float32).reshape(4, 2) cls_name = parts[8] cls_id = class_map[cls_name] # 用 minAreaRect 反推中心点、宽高、角度,避免自己算角度出错 rect = cv2.minAreaRect(coords) (cx, cy), (w, h), angle = rect # 关键:统一成长边为宽、角度落在 [0, 90) if w < h: w, h = h, w angle += 90 angle = angle % 180 if angle >= 90: angle -= 180 w, h = h, w angle = angle % 90 # 归一化 cx /= img_w cy /= img_h w /= img_w h /= img_h return f"{cls_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f} {angle:.6f}"逻辑说明:用cv2.minAreaRect而不是手写向量叉积算角度,是因为 OpenCV 内部对角度和宽高的约定和 YOLOv8-OBB 接近,少一层换算就少一层错。参数上class_map是类别名到 id 的字典,必须和训练时的data.yaml顺序一致,顺序错了模型会把「缺脚」认成「弯脚」。
四个边界坑:第一,minAreaRect返回的角度在不同 OpenCV 版本里范围不一样,老版本是 [-90, 0),新版本是 (0, 90],所以上面用取模加判断兜底。第二,坐标越界,标注时框超出图像边界的,转换后 cx、cy 会大于 1,训练直接报错,转换前要 clip。第三,空行和 difficult 标记,DOTA 里 difficult=1 的样本建议直接丢弃,别喂给模型。第四,类别名大小写,pin_bent和Pin_Bent会被当成两类,统一小写最省事。
2.3 data.yaml 与目录结构怎么摆
YOLOv8-OBB 对目录结构不挑,但data.yaml里路径写错是新手第一翻车点。推荐结构:
dataset/ images/ train/ val/ labels/ train/ val/path: /home/user/dataset train: images/train val: images/val names: 0: normal 1: bent 2: missing 3: offset注意names用字典还是列表都行,但 id 必须从 0 连续。标签文件名要和图片同名、只换后缀,img_001.jpg对应img_001.txt,对不上会被静默跳过,训练时 loss 正常但 mAP 是 0,这种黑匣子现象八成是路径或文件名问题。
3. YOLOv8-OBB 训练:参数怎么设、指标怎么看
3.1 从预训练权重起步的最小训练命令
别从零训,OBB 的预训练权重能省掉大量收敛时间。
yolo obb train \ model=yolov8n-obb.pt \ data=dataset/data.yaml \ epochs=200 \ imgsz=1024 \ batch=8 \ device=0 \ workers=4 \ patience=50 \ lr0=0.01 \ cos_lr=True \ project=runs/obb \ name=pin_v1逻辑说明:model选 n 还是 s 看你的算力和精度要求,引脚缺陷这种细粒度任务,n 在 1024 输入下通常够用,s 更稳但慢。imgsz=1024是关键,引脚太细,640 下小目标特征基本丢光,我一般直接上 1024 甚至 1280。batch受显存限制,8G 显存跑 1024 大概只能到 8。patience=50是早停,OBB 后期容易过拟合,早停能救回最优权重。cos_lr配合lr0=0.01比固定学习率收敛更平滑。
3.2 必调的四个参数与它们对引脚检测的影响
| 参数 | 建议值 | 作用与调法 |
|---|---|---|
| imgsz | 1024 / 1280 | 决定小目标可见度,引脚细必须加大,代价是显存和耗时 |
| batch | 显存允许的最大值 | 太小 BN 统计不稳,太大会 OOM,先试 8 再上下调 |
| lr0 | 0.01 | 太大 loss 炸,太小收敛慢,配合 cos_lr 用 |
| degrees | 0(默认) | OBB 自带角度回归,别开 Mosaic 的旋转增强去干扰角度学习 |
degrees这条是血泪经验。有人习惯性开旋转增强,结果 OBB 的角度回归和增强后的角度打架,mAP 反而降。OBB 任务里角度是标签的一部分,增强要谨慎,mosaic、mixup可以留,旋转类增强建议关掉或调很小。
3.3 训练日志里该盯哪几个指标
box_loss、cls_loss、dfl_loss三条曲线,正常是同步下降。如果cls_loss降但box_loss不降,多半是角度标注不统一;如果dfl_loss一直高,是框回归没学好,检查标注框是否贴合引脚。验证集看mAP50和mAP50-95,OBB 的 mAP 普遍比水平框低几个点,别拿水平框的 0.9 去要求 OBB,0.75 以上在引脚场景就算能用。混淆矩阵重点看bent和offset有没有互相混,这两类角度接近,混了说明角度特征没学出来,回去查标注。
4. 导出 ONNX 再转 TensorRT:pt 文件转换 tensorrt 的完整链路
4.1 为什么不能直接从 pt 转 engine
Ultralytics 支持format=engine直接导出,但生产环境我强烈建议走 ONNX 中转。原因有三:ONNX 是中间表示,出问题能单独验证;TensorRT 版本和 CUDA 版本耦合,ONNX 让你换环境时不用重训;OBB 的输出头结构特殊,直接转 engine 有时输出维度对不上,ONNX 能先看清楚。
# 第一步:pt 转 onnx,注意 opset 和 simplify yolo export \ model=runs/obb/pin_v1/weights/best.pt \ format=onnx \ imgsz=1024 \ opset=12 \ simplify=True \ dynamic=False逻辑说明:opset=12是兼容性最好的选择,太高有些 TensorRT 版本不认。simplify=True会跑 onnx-simplifier 去掉冗余节点,对 OBB 的旋转框解码尤其有用。dynamic=False固定输入尺寸,TensorRT 能针对固定 shape 做最优优化,产线上图片尺寸固定,没必要开动态。
4.2 trtexec 转 engine 与精度选择
trtexec \ --onnx=best.onnx \ --saveEngine=best_fp16.engine \ --fp16 \ --workspace=4096 \ --minShapes=images:1x3x1024x1024 \ --optShapes=images:8x3x1024x1024 \ --maxShapes=images:16x3x1024x1024逻辑说明:--fp16在支持 fp16 的卡上几乎无损提速,引脚检测对精度敏感度中等,fp16 通常够。--workspace=4096是 4G 显存给优化器用,太小会 fallback 到慢的 kernel。shape 三个档位是给动态 batch 用的,如果你固定 batch=1,直接--shapes=images:1x3x1024x1024更省事。转完看 trtexec 打印的Throughput和GPU Compute Time,这才是真实性能。
4.3 TensorRT 10.x 在 GTX1070 上到底能不能跑
这是热搜里问得最多的问题。结论:能装能跑,但有限制。GTX1070 是 Pascal 架构,compute capability 6.1,TensorRT 10.x 官方支持列表里 Pascal 还在,但 fp16 的加速收益比 Turing 之后的卡小很多,因为 Pascal 的 fp16 吞吐是 fp32 的 1/64(消费级卡被砍过)。所以 1070 上跑 fp16 engine 可能比 fp32 还慢,建议直接--fp32或者--int8(int8 需要校准,麻烦)。另外 TensorRT 10.x 对 CUDA 版本要求高,1070 能装的驱动版本有限,装之前先确认驱动支持的 CUDA 上限,别硬上最新版。我一般在这种老卡上就用 TensorRT 8.x,稳定且资料多。
5. 推理部署与避坑:从 engine 到产线节拍
5.1 Python 端加载 engine 并做 OBB 后处理
import tensorrt as trt import pycuda.driver as cuda import numpy as np import cv2 class OBBInfer: def __init__(self, engine_path): logger = trt.Logger(trt.Logger.WARNING) with open(engine_path, "rb") as f, trt.Runtime(logger) as runtime: self.engine = runtime.deserialize_cuda_engine(f.read()) self.context = self.engine.create_execution_context() # 绑定输入输出,TensorRT 10 用 get_tensor_name 遍历 self.bindings = [] for i in range(self.engine.num_io_tensors): name = self.engine.get_tensor_name(i) shape = self.engine.get_tensor_shape(name) dtype = trt.nptype(self.engine.get_tensor_dtype(name)) self.bindings.append((name, shape, dtype)) def preprocess(self, img, size=1024): # letterbox 保持比例,OBB 对形变敏感,别直接 resize h, w = img.shape[:2] scale = min(size / h, size / w) nh, nw = int(h * scale), int(w * scale) resized = cv2.resize(img, (nw, nh)) canvas = np.full((size, size, 3), 114, dtype=np.uint8) canvas[:nh, :nw] = resized blob = canvas[:, :, ::-1].transpose(2, 0, 1).astype(np.float32) / 255.0 return np.ascontiguousarray(blob[None]), scale逻辑说明:OBB 对图像形变比水平框更敏感,因为角度会被拉伸改变,所以预处理必须 letterbox 而不是直接 resize。114是 YOLO 系列惯用的填充灰度。后处理部分要把 TensorRT 输出的[1, 4+nc+1, num_anchors]解码成旋转框,再做旋转 NMS,这部分建议直接复用 Ultralytics 的non_max_suppression里 OBB 分支,自己写容易在角度周期上出错。
5.2 避坑与排查:五条产线级踩坑记录
现象:engine 推理结果和 pt 推理对不上,框位置偏移。原因:预处理不一致,pt 推理时 Ultralytics 内部做了 letterbox,你手写的预处理没对齐 padding 比例。解决:把 pt 推理的预处理参数打印出来,逐项对齐,尤其是 scale 和 pad 的计算方式。
现象:TensorRT 转换报 unsupported op。原因:ONNX 里有些算子 TensorRT 版本不支持,OBB 的旋转框解码里常见GridSample或自定义 op。解决:升级 TensorRT 到支持该 op 的版本,或在导出时用simplify把 op 融合掉,实在不行把后处理挪到 CPU 用 numpy 做。
现象:fp16 engine 精度掉得厉害,mAP 掉 10 个点。原因:Pascal 或老架构卡 fp16 吞吐被砍,且部分层 fp16 溢出。解决:改 fp32,或对敏感层用--precisionConstraints强制 fp32,TensorRT 支持逐层精度设置。
现象:batch 加大后显存 OOM。原因:workspace 和 activation 显存叠加,1024 输入下 activation 很大。解决:降 workspace,或改用--optShapes让 TensorRT 按实际 batch 优化,别一上来就 maxShapes 拉满。
现象:产线跑几小时后推理变慢。原因:显存碎片或 context 没复用,每帧新建 context。解决:engine 和 context 全局只建一次,输入输出 buffer 预分配复用,别在循环里反复 allocate。
5.3 节拍估算与硬件选型
1024 输入、YOLOv8n-OBB、fp16,在 RTX 3060 上单张大概 4-6ms,加上预处理和后处理,端到端 10ms 左右,100 FPS 够大多数产线。GTX1070 上 fp32 大概 20-30ms,30-50 FPS,如果产线节拍是每秒 10 个芯片,够用;要更快就得上 Turing 之后的卡。选型原则:先算清产线节拍要求,再倒推需要的 FPS,别盲目堆卡。
6. 把 OBB 检测做稳的一个进阶技巧:角度一致性校验
模型训完、engine 转完,不代表能上产线。引脚检测最隐蔽的问题是角度抖动:同一个引脚,连续几帧检测出来的角度在 0 和 90 附近跳,后处理算偏移量时就会误判。根因是 OBB 的角度回归在边界处不连续,0 度和 90 度物理上是同一个方向,但数值上差 90。
我的做法是加一层角度一致性校验,在推理后处理里做:
def normalize_angle(angle, w, h): # 把角度统一到 [0, 90),并保证宽是长边 if w < h: w, h = h, w angle += 90 angle = angle % 180 if angle >= 90: angle -= 180 w, h = h, w return angle % 90, w, h def angle_stable(prev_angle, cur_angle, threshold=5.0): # 角度差超过阈值就认为抖动,用上一帧平滑 diff = abs(prev_angle - cur_angle) diff = min(diff, 90 - diff) # 角度周期是 90 if diff > threshold: return prev_angle return cur_angle逻辑说明:normalize_angle和训练时的标注约定保持一致,保证推理和训练同分布。angle_stable利用角度 90 度周期性算最小差,超过阈值就用上一帧角度,相当于一个轻量卡尔曼。threshold=5.0是我在引脚场景试出来的,太小会跟不住真实变化,太大抖动滤不掉,按你的引脚尺寸和相机帧率调。
验证这套是否有效,别只看 mAP。我会单独统计连续帧的角度方差,方差降下来才算稳。另外准备一批「边界样本」——引脚刚好水平或垂直的图,专门看这些样本的角度输出有没有跳变,这是 OBB 最容易翻车的地方。
最后说个习惯:每次改完标注或增强策略,我都会把同一批图在 pt 和 engine 上各跑一遍,逐框对比角度和坐标,差超过 1 个像素就查预处理。这个后悔药比上线后返工便宜太多。希望帮到你。
本文还有配套的精品资源,点击获取