YOLOv8教程系列:四、使用yolov8仓库训练自己的图像分类数据集(含推理预测)

YOLOv8教程系列:四、使用yolov8仓库训练自己的图像分类数据集(含推理预测)

0.引言

Yolov8是最新一代的You Only Look Once目标检测模型,它由Ultralytics研究团队在2022年开发。相比于之前的Yolo版本,Yolov8在速度和精度上都有很大的提升。

在图像分类任务上,Yolov8使用了以下几个特点:

  1. 网络backbone采用了CSPResNeXt作为主干网络,此网络结合了ResNet和ResNeXt的优点,能够提取更丰富的特征。
  2. 使用了Cross Stage Partial Network作为网络的连接方式,将不同阶段的特征图连接起来,增强特征的表达能力。
  3. 引入了注意力机制,可以学习到目标的更加 discriminative 的特征。
  4. 使用了更大的输入分辨率,即640x640,可以学习到更丰富的信息。
  5. 损失函数上,融合了标签调制、IoU预测以及新颖的Grid Sensitive损失,可以更好地改善模型的定位能力。 在推理上使用了更高效的矩阵运算,可以加速模型的推理速度。

总的来说,Yolov8通过网络设计的改进,损失函数的优化以及高效的推理实现,相比之前的版本取得了显著的进步,在图像分类任务上能够达到更高的精度。
yolov8官方文档
**加粗样式**

1.数据准备

先把自己的数据按类别准备好,格式如下,imagenet数据集格式(文件夹名为datasets):

.
├── ./datasets
│ ├── ./datasets/OK
│ │ ├── ./datasets/OK/1.jpg
│ │ ├── ./datasets/OK/2.jpg
│ │ ├── ./datasets/OK/3.jpg
│ │ ├── …
│ ├── ./datasets/NG
│ │ ├── ./datasets/NG/1.jpg
│ │ ├── ./datasets/NG/1.jpg
│ │ ├── ./datasets/NG/1.jpg
│ │ ├── …

2.数据划分

在datasets上层目录下新建一个split.py文件,运行下面脚本:

# 工具类
import os
import random
from shutil import copy2

def data_set_split(src_data_folder, target_data_folder, train_scale=0.8, val_scale=0.2):
    '''
    读取源数据文件夹,生成划分好的文件夹,分为train、val两个文件夹进行
    :param src_data_folder: 源文件夹
    :param target_data_folder: 目标文件夹
    :param train_scale: 训练集比例
    :param val_scale: 验证集比例
    :return:
    '''
    print("开始数据集划分")
    class_names = os.listdir(src_data_folder)
    # 在目标目录下创建文件夹
    split_names = ['train', 'val']
    for split_name in split_names:
        split_path = os.path.join(target_data_folder, split_name)
        if os.path.isdir(split_path):
            pass
        else:
            os.makedirs(split_path)
        # 然后在split_path的目录下创建类别文件夹
        for class_name in class_names:
            class_split_path = os.path.join(split_path, class_name)
            if os.path.isdir(class_split_path):
                pass
            else:
                os.makedirs(class_split_path)

    # 按照比例划分数据集,并进行数据图片的复制
    # 首先进行分类遍历
    for class_name in class_names:
        current_class_data_path = os.path.join(src_data_folder, class_name)
        current_all_data = os.listdir(current_class_data_path)
        current_data_length = len(current_all_data)
        current_data_index_list = list(range(current_data_length))
        random.shuffle(current_data_index_list)

        train_folder = os.path.join(os.path.join(target_data_folder, 'train'), class_name)
        val_folder = os.path.join(os.path.join(target_data_folder, 'val'), class_name)
        train_stop_flag = current_data_length * train_scale
        current_idx = 0
        train_num = 0
        val_num = 0
        for i in current_data_index_list:
            src_img_path = os.path.join(current_class_data_path, current_all_data[i])
            if current_idx <= train_stop_flag:
                copy2(src_img_path, train_folder)
                train_num = train_num + 1
            else:
                copy2(src_img_path, val_folder)
                val_num = val_num + 1

            current_idx = current_idx + 1

        print("*********************************{}*************************************".format(class_name))
        print("{}类按照{}:{}的比例划分完成,一共{}张图片".format(class_name, train_scale, val_scale, current_data_length))
        print("训练集{}:{}张".format(train_folder, train_num))
        print("验证集{}:{}张".format(val_folder, val_num))


if __name__ == '__main__':
    src_data_folder = "datasets"
    target_data_folder = "dataset/"
    data_set_split(src_data_folder, target_data_folder)

运行结束,会按训练集和验证集8:2的比例生成一个划分后的数据集,名为dataset

3.开始训练

新建一个classify_train.py文件,根据自己情况,调整相关参数即可

from ultralytics import YOLO

model = YOLO("yolo-cls/yolov8s-cls.pt")
model.train(data='/home/lzj/03.AlgoDemo/yolov8/dataset/', epochs=100, batch=2, imgsz=1280)

4.推理预测

新建一个classify_infer.py的脚步,注意修改下面的路径和名称列表,运行结束后,会在指定目录下生成预测的图片

import cv2
import os
from ultralytics import YOLO
from tqdm import tqdm

def read_path(file_pathname, model, name_dict, save_folder):
    file_dir = os.listdir(file_pathname)
    for k,v in name_dict.items():
        name_foler = os.path.join(save_folder, v)
        os.makedirs(name_foler)
    #遍历该目录下的所有图片文件
    for filename in tqdm(file_dir):
        print(filename)
        img = cv2.imread(file_pathname+'/'+filename)
        results = model.predict(source=img)

        for result in results:
            # print(result.names)
            name_dict = result.names
            print(name_dict)
            probs = result.probs.cpu().numpy()
            top1_index = result.probs.top1
            class_name = name_dict[top1_index]
            print(class_name)
            save_img_path = os.path.join(save_folder, class_name, filename)
            cv2.imwrite(save_img_path, img)
        print('---------------------------')



if __name__ == '__main__':
    name_dict = {0: 'NG', 1: 'OK'}
    save_folder = 'classify_infer_folder'
    load_img_folder = '/home/lzj/Downloads/pb'
    model = YOLO('runs/classify/train46/weights/best.pt')
    read_path(load_img_folder, model, name_dict, save_folder)


  • 15
    点赞
  • 113
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 19
    评论
在PyTorch下使用Ultralytics YOLOv5训练自己的数据集,可以按照以下步骤进行操作: 1. 安装依赖库: ```shell pip install torch torchvision ``` 2. 克隆YOLOv5仓库: ```shell git clone https://github.com/ultralytics/yolov5.git ``` 3. 进入yolov5目录: ```shell cd yolov5 ``` 4. 准备数据集: - 将自己的数据集放置在`yolov5/data`目录下,包括图像和对应的标注文件(如YOLO格式的txt文件)。 - 在`yolov5/data`目录下创建一个新的文件夹,用于存放自己的数据集,例如`my_dataset`。 5. 配置数据集: - 在`yolov5/data`目录下创建一个新的文件,命名为`my_dataset.yaml`,用于配置自己的数据集。 - 在`my_dataset.yaml`中,按照以下格式填写配置信息: ```yaml train: ../my_dataset/train/images val: ../my_dataset/val/images nc: 1 # 类别数 names: ['class1'] # 类别名称 ``` 6. 划分训练集和验证集: - 在`yolov5/data`目录下创建`my_dataset/train`和`my_dataset/val`两个文件夹,分别用于存放训练集和验证集的图像和标注文件。 - 将数据集中的图像和对应的标注文件按照一定比例划分到`train`和`val`文件夹中。 7. 开始训练: ```shell python train.py --img 640 --batch 16 --epochs 50 --data my_dataset.yaml --weights yolov5s.pt ``` - `--img`:输入图像的大小。 - `--batch`:每个批次的图像数量。 - `--epochs`:训练的轮数。 - `--data`:数据集的配置文件。 - `--weights`:预训练模型的权重文件。 8. 查看训练结果: - 训练过程中的日志和权重文件保存在`yolov5/runs/train/exp`目录下。 - 可以使用TensorBoard查看训练过程中的损失曲线: ```shell tensorboard --logdir=runs/train/exp ``` 9. 进行推理: - 使用训练好的模型进行目标检测: ```shell python detect.py --source path/to/images --weights runs/train/exp/weights/best.pt --conf 0.4 ``` - `--source`:输入图像或视频的路径。 - `--weights`:训练好的权重文件。 - `--conf`:置信度阈值,用于过滤检测结果。
评论 19
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

Zhijun.li@Studio

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

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

余额充值