首页 > 分享 > 如何基于YOLOv8的训练使用宠物狗数据集

如何基于YOLOv8的训练使用宠物狗数据集

萌宠菠菠乐园
2024-11-19 16:07

在这里插入图片描述

宠物狗数据集，并提供基于YOLOv8的训练代码。这个数据集包含26000多张图像，标注了多个类别，标注格式为YOLO格式，标注文件为JSON格式。我们将首先介绍数据集的结构，然后提供转换标注格式的脚本，最后提供训练和预测的脚本。

1. 数据集介绍数据集目录结构

假设你的数据集已经准备好，并且分为训练集、验证集和测试集。数据集目录结构如下：

dog_detection_dataset/ ├── images/ │ ├── train/ │ ├── val/ │ └── test/ ├── labels_json/ │ ├── train/ │ ├── val/ │ └── test/ ├── labels_yolo/ │ ├── train/ │ ├── val/ │ └── test/ └── data.yaml 1234567891011121314 数据集标注格式

JSON格式的标注文件示例如下：

{ "image_path": "images/train/image_0001.jpg", "annotations": [ { "class_id": 0, "bbox": [0.3, 0.4, 0.2, 0.2] }, { "class_id": 1, "bbox": [0.7, 0.5, 0.1, 0.1] } ] } 12345678910111213

其中：

image_path 是图像的路径。annotations 是一个列表，每个元素是一个字典，包含类别ID和边界框坐标。 2. 数据集配置文件 (data.yaml)

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

path: ./dog_detection_dataset # 数据集路径 train: images/train # 训练集图像路径 val: images/val # 验证集图像路径 test: images/test # 测试集图像路径 nc: 2 # 类别数 names: - dog - other # 如果有其他类别，可以在这里添加 123456789 3. 转换标注格式

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

转换脚本

import json import os def convert_json_to_yolo(json_file, yolo_file, class_names): with open(json_file, 'r') as f: data = json.load(f) image_path = data['image_path'] annotations = data['annotations'] with open(yolo_file, 'w') as f: for annotation in annotations: class_id = annotation['class_id'] if class_id >= len(class_names): continue bbox = annotation['bbox'] x_center, y_center, width, height = bbox f.write(f"{class_id} {x_center} {y_center} {width} {height}n") def convert_all_json_to_yolo(json_dir, yolo_dir, class_names): os.makedirs(yolo_dir, exist_ok=True) for filename in os.listdir(json_dir): if filename.endswith('.json'): json_file = os.path.join(json_dir, filename) yolo_file = os.path.join(yolo_dir, filename.replace('.json', '.txt')) convert_json_to_yolo(json_file, yolo_file, class_names) if __name__ == "__main__": class_names = ['dog', 'other'] # 根据实际情况添加类别 json_train_dir = 'dog_detection_dataset/labels_json/train' yolo_train_dir = 'dog_detection_dataset/labels_yolo/train' convert_all_json_to_yolo(json_train_dir, yolo_train_dir, class_names) json_val_dir = 'dog_detection_dataset/labels_json/val' yolo_val_dir = 'dog_detection_dataset/labels_yolo/val' convert_all_json_to_yolo(json_val_dir, yolo_val_dir, class_names) json_test_dir = 'dog_detection_dataset/labels_json/test' yolo_test_dir = 'dog_detection_dataset/labels_yolo/test' convert_all_json_to_yolo(json_test_dir, yolo_test_dir, class_names)

12345678910111213141516171819202122232425262728293031323334353637383940414243 4. 训练脚本 (train.py)

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/dog_detection/best.pt") if __name__ == "__main__": data_yaml_path = 'dog_detection_dataset/data.yaml' model_config = 'yolov8n.yaml' # 你可以选择不同的YOLOv8模型配置，如yolov8s.yaml, yolov8m.yaml等 epochs = 100 batch_size = 16 img_size = 640 augment = True train_model(data_yaml_path, model_config, epochs, batch_size, img_size, augment)

123456789101112131415161718192021222324252627 5. 预测脚本 (predict.py)

import cv2 import torch from ultralytics import YOLO def predict_image(image_path, model_path, img_size=640): # 加载模型 model = YOLO(model_path) # 读取图像 image = cv2.imread(image_path) image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 进行预测 results = model(image_rgb, size=img_size) # 处理预测结果 for result in results: boxes = result.boxes.xyxy.cpu().numpy() scores = result.boxes.conf.cpu().numpy() labels = result.boxes.cls.cpu().numpy().astype(int) for box, score, label in zip(boxes, scores, labels): x1, y1, x2, y2 = map(int, box) class_name = ['dog', 'other'][label] color = (0, 255, 0) if label == 0 else (0, 0, 255) cv2.rectangle(image, (x1, y1), (x2, y2), color, 2) cv2.putText(image, f'{class_name} {score:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2) # 显示图像 cv2.imshow('Prediction', image) cv2.waitKey(0) cv2.destroyAllWindows() if __name__ == "__main__": image_path = 'path_to_your_image.jpg' model_path = 'runs/train/dog_detection/best.pt' predict_image(image_path, model_path)

12345678910111213141516171819202122232425262728293031323334353637 6. 运行脚本

转换标注格式：

python convert_json_to_yolo.py 1

训练模型：

python train.py 1

进行预测：

python predict.py 1 7. 详细解释转换标注格式脚本 (convert_json_to_yolo.py)

导入依赖项：

import json：导入JSON解析库。import os：导入操作系统接口库。

定义转换函数：

convert_json_to_yolo：将单个JSON格式的标注文件转换为YOLO格式的TXT文件。convert_all_json_to_yolo：将指定目录下的所有JSON格式的标注文件转换为YOLO格式的TXT文件。

主函数：

设置类别名称和各个子集的路径。调用convert_all_json_to_yolo函数进行转换。训练脚本 (train.py)

导入依赖项：

from ultralytics import YOLO：导入YOLOv8模型。

定义训练函数：

train_model：加载模型，设置训练参数，训练模型，并保存最佳模型。

主函数：

设置数据集路径、模型配置、训练参数等。调用train_model函数进行训练。预测脚本 (predict.py)

导入依赖项：

import cv2：导入OpenCV库。import torch：导入PyTorch库。from ultralytics import YOLO：导入YOLOv8模型。

定义预测函数：

predict_image：加载模型，读取图像，进行预测，处理预测结果，并显示带有标注的图像。

主函数：

设置图像路径和模型路径。调用predict_image函数进行预测。 8. 注意事项 数据集路径：确保数据集路径正确，特别是data.yaml文件中的路径。模型配置：可以选择不同的YOLOv8模型配置，如yolov8s.yaml, yolov8m.yaml等，根据你的计算资源和需求选择合适的模型。图像大小：img_size可以根据实际需求调整，通常使用640或1280。数据增强：augment参数控制是否启用数据增强，可以在训练过程中提高模型的泛化能力。 9. 数据增强

为了增加数据集的多样性，可以使用数据增强技术。YOLOv8在训练过程中默认支持多种数据增强方法，如随机裁剪、翻转、颜色抖动等。如果需要自定义数据增强，可以参考YOLOv8的文档进行配置。

总结

通过以上步骤，你可以构建一个基于YOLOv8模型的宠物狗检测系统。convert_json_to_yolo.py用于将JSON格式的标注文件转换为YOLO格式，train.py用于训练模型，predict.py用于加载训练好的模型并进行预测。