yolov8-detect 训练自己的数据集(Linux)--详细步骤

1、拉取代码

GitHub - ultralytics/ultralytics: NEW - YOLOv8 🚀 in PyTorch > ONNX > OpenVINO > CoreML > TFLiteNEW - YOLOv8 🚀 in PyTorch > ONNX > OpenVINO > CoreML > TFLite - GitHub - ultralytics/ultralytics: NEW - YOLOv8 🚀 in PyTorch > ONNX > OpenVINO > CoreML > TFLiteicon-default.png?t=N7T8https://github.com/ultralytics/ultralytics.git

2、搭建环境

        docker pull ultralytics/ultralytics:latest

3、准备数据集(labelme标注,json格式)

  •   将json转化为txt,运行labelme2txt.py,按照自己的数据集修改label
import json
import os


def read_json(json_file):
    with open(json_file, 'r') as f:
        load_dict = json.load(f)
    f.close()
    return load_dict

def mk_file(file_path: str):
    if os.path.exists(file_path):
        # 如果文件夹存在,则先删除原文件夹在重新创建
        rmtree(file_path)
    os.makedirs(file_path)

def json2txt(json_path, txt_path):
    mk_file(txt_path)
    for json_file in os.listdir(json_path):
        txt_name = os.path.join(txt_path , json_file[0:-5] + '.txt')
        txt_file = open(txt_name, 'w')
        json_file_path = os.path.join(json_path, json_file)
        json_data = read_json(json_file_path)
        imageWidth = json_data['imageWidth']
        imageHeight = json_data['imageHeight']

        for i in range(len(json_data['shapes'])):
            label = json_data['shapes'][i]['label']

            if label == 'car':
                index = 0
            else:
                index = 1

            x1 = json_data['shapes'][i]['points'][0][0]
            x2 = json_data['shapes'][i]['points'][1][0]
            y1 = json_data['shapes'][i]['points'][0][1]
            y2 = json_data['shapes'][i]['points'][1][1]
            # 将标注框按照图像大小压缩
            x_center = (x1 + x2) / 2 / imageWidth
            y_center = (y1 + y2) / 2 / imageHeight
            bbox_w = (x2 - x1) / imageWidth
            bbox_h = (y2 - y1) / imageHeight
            bbox = (x_center, y_center, bbox_w, bbox_h)
            txt_file.write(str(index) + " " + " ".join([str(a) for a in bbox]) + '\n')

            #print(label)


if __name__ == "__main__":
    json_path = './data/json/'
    txt_path = './data/txt/'

    json2txt(json_path, txt_path)

  • 划分数据集,运行split.py,修改 train_percent 调整数据集划分比例。
# -*- coding: utf-8 -*-
"""
将数据集划分为训练集,验证集,测试集
"""
import os
import random
import shutil


# 创建保存数据的文件夹
def makedir(new_dir):
    if not os.path.exists(new_dir):
        os.makedirs(new_dir)


def split_data(img_dir, label_dir,save_dir):
    random.seed(1)  # 随机种子
    # 1.确定原图片数据集路径
    datasetimg_dir = img_dir
    # 确定原label数据集路径
    datasetlabel_dir = label_dir

    # 2.确定数据集划分后保存的路径
    split_dir = save_dir
    train_dir = os.path.join(split_dir, "train")
    valid_dir = os.path.join(split_dir, "val")
    test_dir = os.path.join(split_dir, "test")
    dir_list = [train_dir, valid_dir, test_dir]
    image_label = ['images', 'labels']

    for i in range(len(dir_list)):
        for j in range(len(image_label)):
            makedir(os.path.join(dir_list[i], image_label[j]))

    # 3.确定将数据集划分为训练集,验证集,测试集的比例
    train_pct = 0.8
    valid_pct = 0.2
    #test_pct = 0.0
    # 4.划分
    imgs = os.listdir(datasetimg_dir)  # 展示目标文件夹下所有的文件名
    imgs = list(filter(lambda x: x.endswith('.jpg'), imgs))  # 取到所有以.png结尾的文件,如果改了图片格式,这里需要修改
    random.shuffle(imgs)  # 乱序路径
    img_count = len(imgs)  # 计算图片数量
    train_point = int(img_count * train_pct)  # 0:train_pct
    valid_point = int(img_count * (train_pct + valid_pct))  # train_pct:valid_pct
    for i in range(img_count):
        if i < train_point:  # 保存0-train_point的图片到训练集
            out_dir = os.path.join(train_dir, 'images')
            label_out_dir = os.path.join(train_dir, 'labels')

        elif i < valid_point:  # 保存train_point-valid_point的图片到验证集
            out_dir = os.path.join(valid_dir, 'images')
            label_out_dir = os.path.join(valid_dir, 'labels')
        else:  # 保存test_point-结束的图片到测试集
            out_dir = os.path.join(test_dir, 'images')
            label_out_dir = os.path.join(test_dir, 'labels')

        target_path = os.path.join(out_dir, imgs[i])  # 指定目标保存路径
        src_path = os.path.join(datasetimg_dir, imgs[i])  # 指定目标原图像路径
        label_target_path = os.path.join(label_out_dir, imgs[i][0:-4] + '.txt')
        label_src_path = os.path.join(datasetlabel_dir, imgs[i][0:-4] + '.txt')
        shutil.copy(src_path, target_path)  # 复制图片
        shutil.copy(label_src_path, label_target_path)  # 复制txt

    print('train:{}, valid:{}, test:{}'.format(train_point, valid_point - train_point,
                                               img_count - valid_point))


if __name__ == "__main__":
    img_dir = './data/jpg/'
    label_dir = './data/txt/'
    save_dir='./data/spilt/'
    split_data(img_dir, label_dir,save_dir)

4、修改配置文件

  • ultralytics\cfg\datasets\coco128.yaml,修改自己数据集图片所在路径和类别名字

  • ultralytics/cfg/models/v8/yolov8.yaml,修改类别数
  • ultralytics/cfg/default.yaml 修改配置文件,预训练权重的路径,训练次数等

5、运行ultralytics/models/yolo/detect/train.py,最新版的不能直接跑,没有main函数,加一个main函数

def train(cfg=DEFAULT_CFG, use_python=False):
    """Train and optimize YOLO model given training data and device."""

    cfg.model='./yolov8n.pt'
    model=cfg.model
    #model = cfg.model or 'yolov8n.pt'
    data = cfg.data or 'coco128.yaml'  # or yolo.ClassificationDataset("mnist")
    device = cfg.device if cfg.device is not None else ''

    args = dict(model=model, data=data, device=device)
    if use_python:
        from ultralytics import YOLO
        YOLO(model).train(**args)
    else:
        trainer = DetectionTrainer(overrides=args)
        trainer.train()


if __name__ == '__main__':
    train()
  • 0
    点赞
  • 4
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值