防震锤缺陷检测数据集 1000张 防震锤 带标注 voc yolo分类名: (图片张数， 标注个数)在这里插入图片描述

none_ defect ive:
(636，1 808)
defective: (831， 933)
总数:
(1000，2741)
总类(nc): 2类
在这里插入图片描述
在这里插入图片描述
使用YOLOv8进行训练的完整步骤。我们将确保所有步骤都是一键运行的，以便你可以轻松地进行模型训练和评估。
在这里插入图片描述

数据集介绍

数据集概述
  • 数据集名称：Anti-Vibration Hammer Defect Detection Dataset (AVHDDD)
  • 数据类型：图像
  • 目标任务：防震锤缺陷检测
  • 样本数量：1000张图像
  • 标注格式：VOC (XML) 和 YOLO (TXT)
  • 类别：2类
    • 无缺陷 (none_defective)
    • 缺陷 (defective)
数据集统计
  • 无缺陷 (none_defective)：636张图像，1808个框
  • 缺陷 (defective)：831张图像，933个框
  • 总计：1000张图像，2741个框
数据集目录结构
AVHDDD/
├── images/
│   ├── train/
│   └── val/
├── labels_voc/
│   ├── train/
│   └── val/
├── labels_yolo/
│   ├── train/
│   └── val/
└── data.yaml

数据集配置文件

创建一个data.yaml文件，配置数据集的路径和类别信息：

path: ./AVHDDD  # 数据集路径
train: images/train  # 训练集图像路径
val: images/val  # 验证集图像路径

nc: 2  # 类别数
names: ['none_defective', 'defective']  # 类别名称

转换标注格式

假设标注文件是VOC格式的XML文件，我们需要将它们转换为YOLO格式的TXT文件。

转换脚本
import xml.etree.ElementTree as ET
import os

def convert_voc_to_yolo(voc_file, yolo_file, class_names):
    tree = ET.parse(voc_file)
    root = tree.getroot()

    width = int(root.find('size/width').text)
    height = int(root.find('size/height').text)

    with open(yolo_file, 'w') as f:
        for obj in root.findall('object'):
            class_name = obj.find('name').text
            if class_name not in class_names:
                continue
            class_id = class_names.index(class_name)

            bbox = obj.find('bndbox')
            x_min = float(bbox.find('xmin').text)
            y_min = float(bbox.find('ymin').text)
            x_max = float(bbox.find('xmax').text)
            y_max = float(bbox.find('ymax').text)

            x_center = (x_min + x_max) / 2.0 / width
            y_center = (y_min + y_max) / 2.0 / height
            w = (x_max - x_min) / width
            h = (y_max - y_min) / height

            f.write(f"{class_id} {x_center} {y_center} {w} {h}\n")

def convert_all_voc_to_yolo(voc_dir, yolo_dir, class_names):
    os.makedirs(yolo_dir, exist_ok=True)
    for filename in os.listdir(voc_dir):
        if filename.endswith('.xml'):
            voc_file = os.path.join(voc_dir, filename)
            yolo_file = os.path.join(yolo_dir, filename.replace('.xml', '.txt'))
            convert_voc_to_yolo(voc_file, yolo_file, class_names)

if __name__ == "__main__":
    class_names = ['none_defective', 'defective']
    voc_train_dir = 'AVHDDD/labels_voc/train'
    yolo_train_dir = 'AVHDDD/labels_yolo/train'
    convert_all_voc_to_yolo(voc_train_dir, yolo_train_dir, class_names)

    voc_val_dir = 'AVHDDD/labels_voc/val'
    yolo_val_dir = 'AVHDDD/labels_yolo/val'
    convert_all_voc_to_yolo(voc_val_dir, yolo_val_dir, class_names)

YOLOv8训练代码

  1. 安装YOLOv8库和依赖项：

    git clone https://github.com/ultralytics/ultralytics.git
cd ultralytics
pip install -r requirements.txt

  2. 训练模型：

    from ultralytics import YOLO

def train_model(data_yaml_path, model_config, epochs, batch_size, img_size, augment):
    # 加载模型
    model = YOLO(model_config)

    # 训练模型
    results = model.train(
        data=data_yaml_path,
        epochs=epochs,
        batch=batch_size,
        imgsz=img_size,
        augment=augment
    )

    # 保存模型
    model.save("runs/train/avhddd/best.pt")

if __name__ == "__main__":
    data_yaml_path = 'AVHDDD/data.yaml'
    model_config = 'yolov8n.yaml'
    epochs = 100
    batch_size = 16
    img_size = 640
    augment = True

    train_model(data_yaml_path, model_config, epochs, batch_size, img_size, augment)

详细解释

  1. 安装YOLOv8和依赖项：

    • 克隆YOLOv8仓库并安装所有必要的依赖项。

  2. 训练模型：

    • 导入YOLOv8库。
    • 加载模型配置文件。
    • 调用model.train方法进行训练。
    • 保存训练后的最佳模型。

运行训练脚本

将上述脚本保存为一个Python文件（例如train_yolov8_avhddd.py），然后运行它。

python train_yolov8_avhddd.py

评估模型

  1. 评估模型： 
    from ultralytics import YOLO

def evaluate_model(data_yaml_path, weights_path, img_size, conf_threshold):
    # 加载模型
    model = YOLO(weights_path)

    # 评估模型
    results = model.val(
        data=data_yaml_path,
        imgsz=img_size,
        conf=conf_threshold
    )

    # 打印评估结果
    print(results)

if __name__ == "__main__":
    data_yaml_path = 'AVHDDD/data.yaml'
    weights_path = 'runs/train/avhddd/best.pt'
    img_size = 640
    conf_threshold = 0.4

    evaluate_model(data_yaml_path, weights_path, img_size, conf_threshold)

详细解释

  1. 评估模型：
    • 导入YOLOv8库。
    • 加载训练好的模型权重。
    • 调用model.val方法进行评估。
    • 打印评估结果。

运行评估脚本

将上述脚本保存为一个Python文件（例如evaluate_yolov8_avhddd.py），然后运行它。

python evaluate_yolov8_avhddd.py

一键运行脚本

为了实现一键运行，可以将图像预处理、训练和评估脚本合并到一个主脚本中，并添加命令行参数来控制运行模式。

import argparse
import os
import xml.etree.ElementTree as ET
from ultralytics import YOLO

def convert_voc_to_yolo(voc_file, yolo_file, class_names):
    tree = ET.parse(voc_file)
    root = tree.getroot()

    width = int(root.find('size/width').text)
    height = int(root.find('size/height').text)

    with open(yolo_file, 'w') as f:
        for obj in root.findall('object'):
            class_name = obj.find('name').text
            if class_name not in class_names:
                continue
            class_id = class_names.index(class_name)

            bbox = obj.find('bndbox')
            x_min = float(bbox.find('xmin').text)
            y_min = float(bbox.find('ymin').text)
            x_max = float(bbox.find('xmax').text)
            y_max = float(bbox.find('ymax').text)

            x_center = (x_min + x_max) / 2.0 / width
            y_center = (y_min + y_max) / 2.0 / height
            w = (x_max - x_min) / width
            h = (y_max - y_min) / height

            f.write(f"{class_id} {x_center} {y_center} {w} {h}\n")

def convert_all_voc_to_yolo(voc_dir, yolo_dir, class_names):
    os.makedirs(yolo_dir, exist_ok=True)
    for filename in os.listdir(voc_dir):
        if filename.endswith('.xml'):
            voc_file = os.path.join(voc_dir, filename)
            yolo_file = os.path.join(yolo_dir, filename.replace('.xml', '.txt'))
            convert_voc_to_yolo(voc_file, yolo_file, class_names)

def train_model(data_yaml_path, model_config, epochs, batch_size, img_size, augment):
    # 加载模型
    model = YOLO(model_config)

    # 训练模型
    results = model.train(
        data=data_yaml_path,
        epochs=epochs,
        batch=batch_size,
        imgsz=img_size,
        augment=augment
    )

    # 保存模型
    model.save("runs/train/avhddd/best.pt")

def evaluate_model(data_yaml_path, weights_path, img_size, conf_threshold):
    # 加载模型
    model = YOLO(weights_path)

    # 评估模型
    results = model.val(
        data=data_yaml_path,
        imgsz=img_size,
        conf=conf_threshold
    )

    # 打印评估结果
    print(results)

def main(mode):
    class_names = ['none_defective', 'defective']
    data_yaml_path = 'AVHDDD/data.yaml'
    model_config = 'yolov8n.yaml'
    epochs = 100
    batch_size = 16
    img_size = 640
    conf_threshold = 0.4
    augment = True

    if mode == 'convert':
        voc_train_dir = 'AVHDDD/labels_voc/train'
        yolo_train_dir = 'AVHDDD/labels_yolo/train'
        convert_all_voc_to_yolo(voc_train_dir, yolo_train_dir, class_names)

        voc_val_dir = 'AVHDDD/labels_voc/val'
        yolo_val_dir = 'AVHDDD/labels_yolo/val'
        convert_all_voc_to_yolo(voc_val_dir, yolo_val_dir, class_names)
    elif mode == 'train':
        train_model(data_yaml_path, model_config, epochs, batch_size, img_size, augment)
    elif mode == 'eval':
        weights_path = 'runs/train/avhddd/best.pt'
        evaluate_model(data_yaml_path, weights_path, img_size, conf_threshold)
    else:
        print("Invalid mode. Use 'convert', 'train', or 'eval'.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Convert, train, or evaluate YOLOv8 on the Anti-Vibration Hammer Defect Detection Dataset.")
    parser.add_argument('mode', type=str, choices=['convert', 'train', 'eval'], help="Mode: 'convert', 'train', or 'eval'")
    args = parser.parse_args()

    main(args.mode)

详细解释

  1. 命令行参数：

    • 使用argparse库添加命令行参数，控制脚本的运行模式（转换、训练或评估）。

  2. 主函数：

    • 根据传入的模式参数，调用相应的转换、训练或评估函数。

运行主脚本

将上述脚本保存为一个Python文件（例如main_yolov8_avhddd.py），然后运行它。

转换标注格式
python main_yolov8_avhddd.py convert
训练模型
python main_yolov8_avhddd.py train
评估模型
python main_yolov8_avhddd.py eval

总结

通过以上步骤，你可以准备好防震锤缺陷检测数据集，并使用YOLOv8进行训练和评估。希望这些信息对你有帮助！

# 目标检测 # 深度学习
Logo
魔乐社区

魔乐社区（Modelers.cn) 是一个中立、公益的人工智能社区，提供人工智能工具、模型、数据的托管、展示与应用协同服务，为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作，由全产业链共同建设、共同运营、共同享有，推动国产AI生态繁荣发展。

更多推荐

cover

大数据毕业设计选题推荐-基于大数据的农作物产量数据分析与可视化系统-Hadoop-Spark-数据可视化-BigData

avatar 魔乐社区
cover

大模型推理适配实战：手把手带你完成vLLM Ascend迁移实操

avatar 魔乐社区
cover

基于python大数据的汽车数据分析系统设计与实现

avatar 魔乐社区