如何结合TensorFlow与YOLOv7开发AOI SMT PCB缺陷检测系统？

一、两者在PCB缺陷检测场景的核心作用 #

• YOLOv7：是专门做目标检测的算法框架，在PCB缺陷检测里，它的核心任务是快速定位缺陷位置（画 bounding box）+ 识别缺陷粗类别，比如能一眼找出PCB上焊锡短路、漏焊、元件偏移这类缺陷的位置，同时判断属于“焊锡缺陷”还是“元件缺陷”这类大类。它的优势是速度快、检测精度高，完全适配工业AOI场景的实时检测需求，单独用就能完成基础的缺陷定位+粗分类任务。
• TensorFlow：是通用的深度学习开发框架，它能搭建、训练各种深度学习模型，不止是分类，也能实现检测、分割任务。在这个场景里，它的作用可以是：
  1. 训练细粒度缺陷分类模型：比如YOLO检测出某块区域是缺陷后，用TensorFlow模型进一步区分是“锡珠”“锡桥”还是“虚焊”这类更细分的缺陷类型；
  2. 做数据预处理、特征提取，或者对YOLO的检测结果做后处理优化；
  3. 搭建系统部署框架，比如用TensorFlow Lite把模型压缩后部署到工业边缘设备上。

二、YOLOv7 + TensorFlow 结合做缺陷分类检测的实现思路 #

最实用的结合方式是YOLO负责定位粗分，TensorFlow负责细粒度分类，或者统一用TensorFlow框架实现YOLOv7并扩展分类分支，以下是具体方案：

方案1：YOLOv7定位 + TensorFlow细分类（工业场景常用） #

1. 第一步：用YOLOv7完成缺陷检测训练

   • 准备带标注的PCB缺陷数据集：标注每个缺陷的边界框（bbox）和粗类别（比如“缺陷”/“正常”，或“焊锡缺陷”/“元件缺陷”）；
   • 用YOLOv7官方代码训练模型，得到能识别并定位PCB缺陷的检测器；
   • 推理时，YOLOv7会输出图像中所有缺陷的位置坐标和粗类别。

2. 第二步：用TensorFlow训练细粒度分类模型

   • 从数据集中裁剪出YOLOv7检测到的缺陷区域图像，作为分类模型的数据集，标注更细分的缺陷类型（比如“锡珠”“锡桥”“漏焊”“元件偏移”）；
   • 用TensorFlow搭建分类模型，推荐用ResNet、MobileNet这类预训练模型做迁移学习，减少训练成本，提升分类精度；
   • 推理时，把YOLO输出的缺陷区域裁剪出来，输入到TensorFlow分类模型，得到精准的缺陷细分类型。

3. 代码整合示例

   # 导入YOLOv7检测相关模块（需提前配置YOLOv7环境）
   from yolov7.utils.datasets import LoadImages
   from yolov7.models.experimental import attempt_load
   from yolov7.utils.general import non_max_suppression, scale_coords
   # 导入TensorFlow和CV模块
   import tensorflow as tf
   import cv2
   
   # 加载YOLOv7预训练模型
   yolo_weights = "yolov7_pcb_defect.pt"
   yolo_model = attempt_load(yolo_weights, map_location="cpu")
   # 加载TensorFlow细分类模型
   tf_classifier = tf.keras.models.load_model("pcb_fine_classifier.h5")
   # 分类标签映射
   cls_labels = ["锡珠", "锡桥", "漏焊", "元件偏移"]
   
   # 处理单张PCB图像
   img_path = "pcb_test.jpg"
   img = cv2.imread(img_path)
   img_tensor = tf.convert_to_tensor(img, dtype=tf.float32) / 255.0
   img_tensor = tf.expand_dims(img_tensor, axis=0)
   
   # YOLOv7推理检测缺陷
   pred = yolo_model(img_tensor)[0]
   pred = non_max_suppression(pred, conf_thres=0.5, iou_thres=0.5)
   
   # 遍历检测结果
   for det in pred:
       if len(det):
           # 缩放边界框到原图像尺寸
           det[:, :4] = scale_coords(img_tensor.shape[2:], det[:, :4], img.shape).round()
           for *xyxy, conf, cls in det:
               x1, y1, x2, y2 = map(int, xyxy)
               # 裁剪缺陷区域
               defect_roi = img[y1:y2, x1:x2]
               # 预处理适配TensorFlow模型
               defect_roi = cv2.resize(defect_roi, (224, 224)) / 255.0
               defect_roi = tf.expand_dims(defect_roi, axis=0)
               # TensorFlow细分类推理
               pred_cls_idx = tf.argmax(tf_classifier.predict(defect_roi), axis=1)[0]
               pred_cls = cls_labels[pred_cls_idx]
               # 输出结果
               print(f"缺陷位置：({x1},{y1})-({x2},{y2})，置信度：{conf:.2f}，缺陷类型：{pred_cls}")
   

方案2：基于TensorFlow重写YOLOv7并扩展分类分支 #

如果想统一用TensorFlow框架开发，可以把YOLOv7的网络结构用TensorFlow的tf.keras层重写，然后在检测头后新增一个分类分支，专门处理缺陷的细粒度分类：

• 参考YOLOv7的官方网络结构，用TensorFlow搭建包含特征提取、检测头、分类头的多任务模型；
• 训练时使用多任务损失函数（检测损失+分类损失），同时优化缺陷定位和细分类任务；
• 这种方式适合需要统一框架部署的场景，方便后续用TensorFlow生态的工具做模型优化和部署。

三、关键注意事项 #

• 数据集对齐：YOLO的标注和分类模型的标注要一一对应，确保裁剪的缺陷区域标注准确；
• 部署优化：如果要部署到工业设备，YOLOv7可以转成ONNX格式后导入TensorFlow，或者直接用TensorFlow Lite优化分类模型，保证实时检测速度；
• 迁移学习优先：不管是YOLOv7还是TensorFlow分类模型，都用预训练模型做迁移学习，能大幅减少训练数据量，加快训练速度。

内容的提问来源于stack exchange，提问作者Candy