深度学习之数据扩充，对应生成扩充的标签文件

最新推荐文章于 2024-07-06 21:27:32 发布

原创

最新推荐文章于 2024-07-06 21:27:32 发布 · 2.4k 阅读

26 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #pytorch #python

该代码实现了一套图像数据增强流程，包括使用imgaug库进行中心坐标的转换、仿射变换、亮度和色调调整等操作。通过 AbsoluteLabels 和 RelativeLabels 将标签坐标转换为适合模型训练的格式，并进行图像的缩放和填充以保持正方形。此外，还定义了DefaultAug类来执行随机像素删除、锐化和图像翻转等数据增强。整个流程适用于目标检测任务，提高模型的泛化能力。

需要用到imgaug库，可自行安装。

直接上代码。transforms.py主要是对标签坐标进行转换及执行扩充操作

# import torch
import torch.nn.functional as F
import numpy as np

import imgaug.augmenters as iaa
from imgaug.augmentables.bbs import BoundingBox, BoundingBoxesOnImage


import torchvision.transforms as transforms

# function: 将中心坐标（x,y），宽高（w, h）转为左上角坐标和右下角坐标的形式（x1, y1, x2, y2）
def xywh2xyxy_np(x):
    y = np.zeros_like(x)
    y[..., 0] = x[..., 0] - x[..., 2] / 2
    y[..., 1] = x[..., 1] - x[..., 3] / 2
    y[..., 2] = x[..., 0] + x[..., 2] / 2
    y[..., 3] = x[..., 1] + x[..., 3] / 2
    return y

class ImgAug(object):
    def __init__(self, augmentations=[]):
        self.augmentations = augmentations

    def __call__(self, data):
        img, boxes = data

        # convert xywh to xyxy
        boxes = np.array(boxes)
        boxes[:, 1:] = xywh2xyxy_np(boxes[:,1:])

        # convert bounding boxes to imgaug
        bounding_boxes = BoundingBoxesOnImage(
            [BoundingBox(*box[1:], label = box[0]) for box in boxes],
            shape=img.shape
        )

        # apply augmentations
        img, bounding_boxes = self.augmentations(
            image = img,
            bounding_boxes = bounding_boxes
        )

        bounding_boxes = bounding_boxes.clip_out_of_image()

        # convert bounding boxes back to numpy
        boxes = np.zeros((len(bounding_boxes), 5))
        for box_idx, box in enumerate(bounding_boxes):
            x1 = box.x1

最低0.47元/天解锁文章