小白必备！YOLOv8实例分割训练自己的数据集指南

后端

2023-08-04 07:25:46

实例分割：使用 YOLOv8 训练自己的模型

实例分割是一种计算机视觉任务，用于识别图像中不同对象的实例。与目标检测（仅识别对象类别）不同，实例分割的任务是识别每个对象的像素，从而实现更精细的分割。

1. 数据准备

数据转换和划分

将数据转换为 YOLOv8 格式 ：将图片和 JSON 文件放在不同的文件夹中。运行 json2txt.py 脚本将 JSON 标注转换为 YOLOv8 支持的 TXT 格式。
划分数据集 ：运行 split.py 脚本将图片随机分为训练集和测试集。

2. 模型训练

安装 YOLOv8

克隆 YOLOv8 仓库并安装依赖项：

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

修改配置文件

打开 yolov5/data/coco128.yaml 文件并修改 num_classes 为您数据集中的类别数。

开始训练

运行以下命令开始训练模型：

python train.py --data your_dataset.yaml --weights yolov5s.pt --img 640 --batch 16 --epochs 300 --name your_model_name

3. 模型评估

训练后评估

训练完成后，使用以下命令评估模型性能：

python val.py --data your_dataset.yaml --weights your_model_name.pt --img 640

4. 实例分割

进行实例分割

训练完成后，使用以下命令进行实例分割：

python detect.py --weights your_model_name.pt --img 640 --conf 0.5 --iou 0.5 --source your_image.jpg

代码示例

# json2txt.py
import json
import os

def main():
    # 文件路径
    json_dir = "path/to/json_dir"
    image_dir = "path/to/image_dir"
    txt_dir = "path/to/txt_dir"

    # 遍历 JSON 文件
    for json_file in os.listdir(json_dir):
        # 读取 JSON 文件
        with open(os.path.join(json_dir, json_file), "r") as f:
            data = json.load(f)

        # 将标注信息写入 TXT 文件
        with open(os.path.join(txt_dir, json_file.replace(".json", ".txt")), "w") as f:
            for obj in data["objects"]:
                class_id = obj["category"]
                x = (obj["bbox"]["x"] + obj["bbox"]["width"] / 2) / obj["image_width"]
                y = (obj["bbox"]["y"] + obj["bbox"]["height"] / 2) / obj["image_height"]
                width = obj["bbox"]["width"] / obj["image_width"]
                height = obj["bbox"]["height"] / obj["image_height"]
                f.write(f"{class_id} {x} {y} {width} {height}\n")

if __name__ == "__main__":
    main()

# split.py
import os
import random

def main():
    # 文件路径
    image_dir = "path/to/image_dir"
    train_txt = "path/to/train.txt"
    test_txt = "path/to/test.txt"

    # 获取图片路径
    images = [os.path.join(image_dir, f) for f in os.listdir(image_dir)]

    # 随机划分数据集
    random.shuffle(images)
    train_size = int(len(images) * 0.8)
    train_images = images[:train_size]
    test_images = images[train_size:]

    # 将路径写入文件
    with open(train_txt, "w") as f:
        for image in train_images:
            f.write(image + "\n")
    with open(test_txt, "w") as f:
        for image in test_images:
            f.write(image + "\n")

if __name__ == "__main__":
    main()