一、环境准备

1. 安装必要的环境

确保安装以下依赖项：

• Python：建议使用 Python 3.8 以上版本。
• PyTorch：根据你的 GPU 支持安装合适的版本。
• YOLO：我们使用 ultralytics 提供的 YOLOv8。

2. 安装步骤

1. 安装 PyTorch：
   根据 CUDA 版本选择适合的安装命令（以 CUDA 11.8 为例）：

   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
   

2. 安装 ultralytics：

   pip install ultralytics
   

3. 安装 OpenCV（用于绘制结果）：

   pip install opencv-python
   

4. 验证安装：
   在 Python 环境中运行以下代码：

   import torch
   import cv2
   from ultralytics import YOLO
   
   print("PyTorch Version:", torch.__version__)
   print("OpenCV Version:", cv2.__version__)
   print("CUDA Available:", torch.cuda.is_available())
   

---

二、实现 YOLO 缺陷检测

1. 检测逻辑说明

以下代码实现了对一个文件夹中的所有图片进行检测，并将检测到缺陷的图片保存到指定文件夹。核心功能包括：

• 加载 YOLO 模型。
• 遍历图片文件夹，对每张图片进行检测。
• 如果检测到缺陷，绘制缺陷框并保存结果。

2. 完整代码

以下是实现 YOLO 缺陷检测的代码：

import os
import sys
import cv2
import time
import torch
from ultralytics import YOLO

# 验证 PyTorch 和 OpenCV 是否正常安装
print("PyTorch Version:", torch.__version__)
print("OpenCV Version:", cv2.__version__)
print("CUDA Available:", torch.cuda.is_available())

if torch.cuda.is_available():
    print("CUDA Device Name:", torch.cuda.get_device_name(0))

    # 输入和输出路径
    input_folder = r'D:\BaiduNetdiskDownload\TEST'  # 替换为你的输入文件夹路径
    output_folder = r'D:\BaiduNetdiskDownload\TEST_RESULTS'  # 替换为你的输出文件夹路径
    os.makedirs(output_folder, exist_ok=True)

    # 加载 YOLO 模型并设置置信度阈值
    model = YOLO('yolov8n.pt')  # 使用 YOLOv8 的轻量化预训练模型
    model.conf = 0.25  # 设置置信度阈值

    # 遍历文件夹中的图片
    for filename in os.listdir(input_folder):
        if filename.endswith('.jpg'):  # 仅处理 .jpg 文件
            input_path = os.path.join(input_folder, filename)

            # 开始计时
            start_time = time.time()
            results = model(input_path)
            end_time = time.time()
            elapsed_time_ms = (end_time - start_time) * 1000

            # 加载原始图片
            img = cv2.imread(input_path)
            if img is None:
                print(f"无法加载图片: {input_path}")
                continue

            # 检测结果处理
            has_defect = False  # 是否检测到缺陷
            for result in results:
                if result.boxes:
                    has_defect = True  # 标记为有缺陷
                    for box in result.boxes:
                        cls = result.names[int(box.cls)]  # 缺陷类型
                        conf = box.conf.item()  # 置信度
                        x1, y1, x2, y2 = map(int, box.xyxy[0])
                        h, w, _ = img.shape
                        x1, y1, x2, y2 = max(0, x1), max(0, y1), min(w, x2), min(h, y2)

                        # 绘制矩形框和标签
                        color = (0, 0, 255)  # 红色
                        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
                        label = f"{cls} ({x1},{y1},{x2 - x1},{y2 - y1})"
                        cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)

            # 仅保存有缺陷的图片
            if has_defect:
                save_path = os.path.join(output_folder, filename)
                cv2.imwrite(save_path, img)
                print(f"图片: {filename} | 耗时: {elapsed_time_ms:.2f} 毫秒 | 已保存检测结果。")
            else:
                print(f"图片: {filename} 未检测到缺陷，跳过保存。")

    print(f"检测完成！检测结果已保存到: {output_folder}")

---

三、运行步骤

1. 文件准备

• 输入文件夹：将待检测的图片放入 D:\BaiduNetdiskDownload\TEST 文件夹。
• 输出文件夹：检测到有缺陷的图片将保存到 D:\BaiduNetdiskDownload\TEST_RESULTS。

2. 运行代码

确保代码中的路径正确，运行代码后，控制台会输出检测结果，例如：

图片: image1.jpg | 耗时: 350.12 毫秒 | 已保存检测结果。
图片: image2.jpg 未检测到缺陷，跳过保存。

---

四、结果查看

1. 检测结果

• 检测到有缺陷的图片会保存到 TEST_RESULTS 文件夹。
• 在保存的图片上可以看到：
  • 红色矩形框标记缺陷区域。
  • 矩形框上方显示缺陷类型及坐标信息。

2. 控制台日志

每张图片的检测时间和检测状态都会在控制台输出。

---

五、常见问题

1. YOLO 模型未检测到缺陷

• 确保输入图片内容清晰。
• 调整置信度阈值 model.conf：

  model.conf = 0.2  # 降低阈值
  

2. OpenCV 显示或保存失败

• 确保 opencv-python 安装正确。
• 确保图片文件路径有效。

3. CUDA 未启用

• 如果控制台输出 CUDA Available: False，请检查 GPU 驱动和 CUDA 工具包是否正确安装。

---

六、总结

通过上述流程，你可以快速实现基于 YOLO 的缺陷检测。该代码灵活易用，可根据实际需求调整检测逻辑，例如修改绘制样式或保存策略。如果有进一步需求（如扩展功能），可以随时联系我们！