基于Albumentations的图像增强，对yolo数据，生成增强后的图片以及标签。

sguangxiang

已于 2024-02-04 09:52:47 修改

阅读量1.8k

点赞数 17

文章标签： YOLO python opencv

于 2024-01-31 11:24:45 首次发布

本文链接：https://blog.csdn.net/weixin_67977644/article/details/135949095

版权

本文介绍了一位开发者如何利用Albumentations库对YOLO数据集进行增强，包括图像裁剪、旋转、翻转等操作，同时自定义了封装类，实现在指定路径下批量处理图片和标签，以提升机器学习模型的鲁棒性。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

Albumentations是一个用于图像增强的Python库，它提供了多种增强技术，包括随机裁剪、旋转、缩放、翻转、变形、颜色变换、模糊等操作。使用Albumentations库可以快速、高效地对图像数据进行增强，从而提升机器学习模型的鲁棒性。

本人根据非常棒的Albumentations数据增强库进行二次封装，将yolo数据生成增强后的标签跟图片，代码更改路径可直接调用。

from albumentations import *
import os
import cv2
from tqdm import tqdm


class enhancement:
    def __init__(self, picture_path, label_path, save_img_path, save_lable_path):
        self.picture_name = sorted(os.listdir(picture_path))
        self.label_name = sorted(os.listdir(label_path))
        self.picture_path = [picture_path + i for i in self.picture_name]
        self.label_path = [label_path + i for i in self.label_name]
        self.save_img_path = save_img_path
        self.save_lable_path = save_lable_path

    def iter(self):
        batch_size = 10
        for index_bin in tqdm(range(0, len(self.picture_path), batch_size), desc='批次进度'):
            # print(index_bin)
            picture_batch = self.picture_path[index_bin:index_bin + batch_size]
            label_batch = self.label_path[index_bin:index_bin + batch_size]
            yield picture_batch, label_batch, [index_bin, index_bin + batch_size]

    def get_transform(self):
        '''

最低0.47元/天解锁文章