Yolo图像增广


在一些场景下的目标检测中,训练集数量较少,模型对灯光、位置、角度、相机的依赖较深,会导致检测的效果比较差,这时就需要进行图像增广。本文介绍几种常见的图像增广技术,包括裁剪、平移、改变亮度、高斯噪声、旋转以及镜像。

图像增广(image augmentation)技术通过对训练图像做一系列随机改变,来产生相似但又不同的训练样本,从而扩大训练数据集的规模。
图像增广的另一种解释是,随机改变训练样本可以降低模型对某些属性的依赖,从而提高模型的泛化能力。

例如,我们可以对图像进行不同方式的裁剪,使感兴趣的物体出现在不同位置,从而减轻模型对物体出现位置的依赖性。我们也可以调整亮度、色彩等因素来降低模型对色彩的敏感度。

可以说,在当年AlexNet的成功中,图像增广技术功不可没。

工作流程

  • 单一图片增广处理流程
一组原始的图片及标注
读取原始图片JPG
读取对应标注XML
图像变换
坐标变换
保存变换过的图片
保存变换过的XML
一组新的图片及标注
  • 整个文件夹的图片增广处理流程
重复设定的次数
指定文件夹
选择目标函数及增广数量
遍历文件夹下所有图片名
一组原始的图片及标注
读取原始图片JPG
读取对应标注XML
图像变换
坐标变换
保存变换过的图片
保存变换过的XML
一组新的图片及标注

辅助工作

读取图片

使用OpenCV读取

filename = "image/test.jpg"
image = cv2.imread(filename)

读取标注XML文件

根据XML标注文件得到标注列表[x_min, y_min, x_max, y_max, cat_name]的列表

    def getBoxes(self, image_name):
        """
        根据XML标注文件得到标注列表[x_min, y_min, x_max, y_max, cat_name]的列表
        :param image_name:
        :return:
        """
        tree = ET.parse(image_name + '.xml')
        root = tree.getroot()
        boxes = []
        for object in root.findall('object'):
            temp_list = []
            name = object.find('name').text
            for coordinate in object.find('bndbox'):
                temp_list.append(int(coordinate.text))
            temp_list.append(name)
            boxes.append(temp_list)
        # print(boxes)
        return boxes

保存变换后的图片

使用OpenCV保存

current_time = datetime.datetime.now().strftime("%Y_%m_%d_%H_%M_%S")
current_num = str(random.randint(0, 9999))
save_image_name = folder + '/' + current_time + '_' + str(i) + current_num + '.jpg'
print("save image name: " + save_image_name)
cv2.imwrite(save_image_name, change_img)

保存更新过的标注

根据对图像的变换操作,保存更新过的标注文件,以新图片相同的名字命名此XML文件。

    def saveXML(self, image_name, xml_name, boxes, shape1, shape0):
        print("xml name ====================================" + xml_name)
        print("image name ====================================" + xml_name)
        folder = ET.Element('folder')
        folder.text = 'image'

        filename = ET.Element('filename')
        filename.text = image_name

        path = ET.Element('path')
        curr_path = os.getcwd()

        path.text = curr_path + '\\image\\' + image_name

        source = ET.Element('source')
        database = ET.SubElement(source, 'database')
        database.text = 'Unknown'

        size = ET.Element('size')
        width = ET.SubElement(size, 'width')

        width.text = str(shape1)
        height = ET.SubElement(size, 'height')
        height.text = str(shape0)
        depth = ET.SubElement(size, 'depth')
        depth.text = '3'
        segmented = ET.Element('segmented')
        segmented.text = '0'

        root = ET.Element('annotation')
        root.extend((folder, filename, path))
        root.extend((source, size, segmented))

        for box in boxes:
            object = ET.Element('object')
            name = ET.SubElement(object, 'name')
            name.text = box[4]
            pose = ET.SubElement(object, 'pose')
            pose.text = 'Unspecified'
            truncated = ET.SubElement(object, 'truncated')
            truncated.text = '0'
            difficult = ET.SubElement(object, 'difficult')
            difficult.text = '0'
            bndbox = ET.SubElement(object, 'bndbox')
            xmin = ET.SubElement(bndbox, 'xmin')
            xmin.text = str(box[0])
            ymin = ET.SubElement(bndbox, 'ymin')
            ymin.text = str(box[1])
            xmax = ET.SubElement(bndbox, 'xmax')
            xmax.text = str(box[2])
            ymax = ET.SubElement(bndbox, 'ymax')
            ymax.text = str(box[3])
            root.extend((object,))

        tree = ET.ElementTree(root)
        tree.write(xml_name)

        tree = ET.parse(xml_name)  # 解析movies.xml这个文件
        root = tree.getroot()  # 得到根元素,Element类
        self.pretty_xml(root, '\t', '\n')  # 执行美化方法
        tree.write(xml_name)

美化更新过的XML

如果打开上一步骤输出的XML文件,会发现所有文本会挤在一行上。这时候应该使用递归去对输出的XML文件进行格式美化。

    def pretty_xml(self, element, indent, newline, level=0):  # elemnt为传进来的Elment类,参数indent用于缩进,newline用于换行
        if element:  # 判断element是否有子元素
            if (element.text is None) or element.text.isspace():  # 如果element的text没有内容
                element.text = newline + indent * (level + 1)
            else:
                element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * (level + 1)
                # else:  # 此处两行如果把注释去掉,Element的text也会另起一行
                # element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * level
        temp = list(element)  # 将element转成list
        for subelement in temp:
            if temp.index(subelement) < (len(temp) - 1):  # 如果不是list的最后一个元素,说明下一个行是同级别元素的起始,缩进应一致
                subelement.tail = newline + indent * (level + 1)
            else:  # 如果是list的最后一个元素, 说明下一行是母元素的结束,缩进应该少一个
                subelement.tail = newline + indent * level
            self.pretty_xml(subelement, indent, newline, level=level + 1)  # 对子元素进行递归操作

图像及坐标变换

1-裁切

裁剪后的图片需要包含所有的标注框,否则会对图像的原始标注造成破坏。

    def __cropImage(self, img, boxes):
        """
        裁切
        :param img: 图像
        :param bboxes: 该图像包含的所有boundingboxes,一个list,每个元素为[x_min,y_min,x_max,y_max]
        :return: crop_img:裁剪后的图像;crop_bboxes:裁剪后的boundingbox的坐标,list
        """
        # 裁剪图像
        w = img.shape[1]
        h = img.shape[0]

        x_min = w
        x_max = 0
        y_min = h
        y_max = 0

        # 最小区域
        for bbox in boxes:
            x_min = min(x_min, bbox[0])
            y_min = min(y_min, bbox[1])
            x_max = max(x_max, bbox[2])
            y_max = max(y_max, bbox[3])
            name = bbox[4]

        # 包含所有目标框的最小框到各个边的距离
        d_to_left = x_min
        d_to_right = w - x_max
        d_to_top = y_min
        d_to_bottom = h - y_max

        # 随机扩展这个最小范围
        crop_x_min = int(x_min - random.uniform(0, d_to_left))
        crop_y_min = int(y_min - random.uniform(0, d_to_top))
        crop_x_max = int(x_max + random.uniform(0, d_to_right))
        crop_y_max = int(y_max + random.uniform(0, d_to_bottom))

        # 确保不出界
        crop_x_min = max(0, crop_x_min)
        crop_y_min = max(0, crop_y_min)
        crop_x_max = min(w, crop_x_max)
        crop_y_max = min(h, crop_y_max)

        crop_img = img[crop_y_min:crop_y_max, crop_x_min:crop_x_max]

        # 裁剪bounding boxes
        crop_bboxes = list()
        for bbox in boxes:
            crop_bboxes.append([int(bbox[0] - crop_x_min), int(bbox[1] - crop_y_min),
                                int(bbox[2] - crop_x_min), int(bbox[3] - crop_y_min), bbox[4]])

        return crop_img, crop_bboxes

2-平移

裁剪后的图片需要包含所有的标注框,否则会对图像的原始标注造成破坏。

    def __translationImage(self, img, boxes):
        """
        平移
        :param img: img
        :param bboxes: bboxes:该图像包含的所有boundingboxes,一个list,每个元素为[x_min,y_min,x_max,y_max]
        :return: shift_img:平移后的图像array;shift_bboxes:平移后的boundingbox的坐标,list
        """

        # 平移图像
        w = img.shape[1]
        h = img.shape[0]

        x_min = w
        x_max = 0
        y_min = h
        y_max = 0

        for bbox in boxes:
            x_min = min(x_min, bbox[0])
            y_min = min(y_min, bbox[1])
            x_max = max(x_max, bbox[2])
            y_max = max(x_max, bbox[3])
            name = bbox[4]

        # 包含所有目标框的最小框到各个边的距离,即每个方向的最大移动距离
        d_to_left = x_min
        d_to_right = w - x_max
        d_to_top = y_min
        d_to_bottom = h - y_max

        # 在矩阵第一行中表示的是[1,0,x],其中x表示图像将向左或向右移动的距离,如果x是正值,则表示向右移动,如果是负值的话,则表示向左移动。
        # 在矩阵第二行表示的是[0,1,y],其中y表示图像将向上或向下移动的距离,如果y是正值的话,则向下移动,如果是负值的话,则向上移动。
        x = random.uniform(-(d_to_left / 3), d_to_right / 3)
        y = random.uniform(-(d_to_top / 3), d_to_bottom / 3)
        M = np.float32([[1, 0, x], [0, 1, y]])

        # 仿射变换
        shift_img = cv2.warpAffine(img, M,
                    (img.shape[1], img.shape[0]))  # 第一个参数表示我们希望进行变换的图片,第二个参数是我们的平移矩阵,第三个希望展示的结果图片的大小

        # 平移boundingbox
        shift_bboxes = list()
        for bbox in boxes:
            shift_bboxes.append([int(bbox[0] + x), int(bbox[1] + y), int(bbox[2] + x), int(bbox[3] + y), bbox[4]])

        return shift_img, shift_bboxes

3-改变亮度

不需要处理标注框的变换问题

    def __changeLightofImage(self, img, boxes):
        """
        改变亮度
        :param img: 图像
        :return: img:改变亮度后的图像array
        """
        '''
        adjust_gamma(image, gamma=1, gain=1)函数:
        gamma>1时,输出图像变暗,小于1时,输出图像变亮
        '''
        flag = random.uniform(0.5, 1.5)  ##flag>1为调暗,小于1为调亮
        newBoxes = copy.deepcopy(boxes)
        newImage = exposure.adjust_gamma(img, flag)
        return newImage, newBoxes

4-高斯噪声

不需要处理标注框的变换问题

    def __addNoiseToImage(self, img, boxes):
        """
        加入噪声
        :param img: 图像
        :return: img:加入噪声后的图像array,由于输出的像素是在[0,1]之间,所以得乘以255
        """
        newBoxes = copy.deepcopy(boxes)
        newImage = random_noise(img, mode='gaussian', clip=True) * 255
        return newImage, newBoxes

5-旋转

旋转后的图片需要包含所有的标注框,否则会对图像的原始标注造成破坏。
需要注意的是,旋转时图像的一些边角可能会被切掉,需要避免这种情况。

    def __rotateImage(self, img, boxes):
        """
        旋转
        :param img: 图像
        :param boxes:
        :param angle: 旋转角度
        :param scale: 默认1
        :return: rot_img:旋转后的图像array;rot_bboxes:旋转后的boundingbox坐标list
        """
        '''
        输入:
            img:array,(h,w,c)
            bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
            angle:
            scale:默认1
        输出:
            
        '''
        # 旋转图像
        w = img.shape[1]
        h = img.shape[0]
        angle = random.uniform(-45, 45)
        scale = random.uniform(0.5, 1.5)
        # 角度变弧度
        rangle = np.deg2rad(angle)
        # 计算新图像的宽度和高度,分别为最高点和最低点的垂直距离
        nw = (abs(np.sin(rangle) * h) + abs(np.cos(rangle) * w)) * scale
        nh = (abs(np.cos(rangle) * h) + abs(np.sin(rangle) * w)) * scale
        # 获取图像绕着某一点的旋转矩阵
        # getRotationMatrix2D(Point2f center, double angle, double scale)
        # Point2f center:表示旋转的中心点
        # double angle:表示旋转的角度
        # double scale:图像缩放因子
        rot_mat = cv2.getRotationMatrix2D((nw * 0.5, nh * 0.5), angle, scale)  # 返回 2x3 矩阵
        # 新中心点与旧中心点之间的位置
        rot_move = np.dot(rot_mat, np.array([(nw - w) * 0.5, (nh - h) * 0.5, 0]))
        # the move only affects the translation, so update the translation
        # part of the transform
        rot_mat[0, 2] += rot_move[0]
        rot_mat[1, 2] += rot_move[1]
        # 仿射变换
        rot_img = cv2.warpAffine(img, rot_mat, (int(math.ceil(nw)), int(math.ceil(nh))),
                                 flags=cv2.INTER_LANCZOS4)  # ceil向上取整

        # 矫正boundingbox
        # rot_mat是最终的旋转矩阵
        # 获取原始bbox的四个中点,然后将这四个点转换到旋转后的坐标系下
        rot_bboxes = list()
        for bbox in boxes:
            x_min = bbox[0]
            y_min = bbox[1]
            x_max = bbox[2]
            y_max = bbox[3]
            name = bbox[4]
            point1 = np.dot(rot_mat, np.array([(x_min + x_max) / 2, y_min, 1]))
            point2 = np.dot(rot_mat, np.array([x_max, (y_min + y_max) / 2, 1]))
            point3 = np.dot(rot_mat, np.array([(x_min + x_max) / 2, y_max, 1]))
            point4 = np.dot(rot_mat, np.array([x_min, (y_min + y_max) / 2, 1]))

            # 合并np.array
            concat = np.vstack((point1, point2, point3, point4))  # 在竖直方向上堆叠
            # 改变array类型
            concat = concat.astype(np.int32)
            # 得到旋转后的坐标
            rx, ry, rw, rh = cv2.boundingRect(concat)
            rx_min = rx
            ry_min = ry
            rx_max = rx + rw
            ry_max = ry + rh
            # 加入list中
            rot_bboxes.append([rx_min, ry_min, rx_max, ry_max, name])
        return rot_img, rot_bboxes

6-镜像

旋转后的图片需要包含所有的标注框,否则会对图像的原始标注造成破坏。
这里介绍两种镜像方式:水平翻转和垂直翻转。

def __flipImage(self, img, bboxes):
    """
    镜像
    :param self:
    :param img:
    :param bboxes:
    :return:
    """
    '''
    镜像后的图片要包含所有的框
    输入:
        img:图像array
        bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
    输出:
        flip_img:镜像后的图像array
        flip_bboxes:镜像后的bounding box的坐标list
    '''
    # 镜像图像
    import copy
    flip_img = copy.deepcopy(img)
    if random.random() < 0.5:
        horizon = True
    else:
        horizon = False
    h, w, _ = img.shape
    if horizon:  # 水平翻转
        flip_img = cv2.flip(flip_img, -1)
    else:
        flip_img = cv2.flip(flip_img, 0)
    # ---------------------- 矫正boundingbox ----------------------
    flip_bboxes = list()
    for bbox in bboxes:
        x_min = bbox[0]
        y_min = bbox[1]
        x_max = bbox[2]
        y_max = bbox[3]
        name = bbox[4]
        if horizon:
            flip_bboxes.append([w - x_max, y_min, w - x_min, y_max, name])
        else:
            flip_bboxes.append([x_min, h - y_max, x_max, h - y_min, name])

    return flip_img, flip_bboxes

完整代码

import os
import cv2
import math
import argparse
import numpy as np
import copy
from skimage import exposure
from skimage.util import random_noise
import random
import datetime
import xml.etree.ElementTree as ET
"""
YOLO图像增广
"""


class ImageAugmentation:
    def __init__(self):
        pass

    def getBoxes(self, image_name):
        """
        根据XML标注文件得到标注列表[x_min, y_min, x_max, y_max, cat_name]的列表
        :param image_name:
        :return:
        """
        tree = ET.parse(image_name + '.xml')
        root = tree.getroot()
        boxes = []
        for object in root.findall('object'):
            temp_list = []
            name = object.find('name').text
            for coordinate in object.find('bndbox'):
                temp_list.append(int(coordinate.text))
            temp_list.append(name)
            boxes.append(temp_list)
        # print(boxes)
        return boxes

    def saveXML(self, image_name, xml_name, boxes, shape1, shape0):
        print("xml name ====================================" + xml_name)
        print("image name ====================================" + xml_name)
        folder = ET.Element('folder')
        folder.text = 'image'

        filename = ET.Element('filename')
        filename.text = image_name

        path = ET.Element('path')
        curr_path = os.getcwd()

        path.text = curr_path + '\\image\\' + image_name

        source = ET.Element('source')
        database = ET.SubElement(source, 'database')
        database.text = 'Unknown'

        size = ET.Element('size')
        width = ET.SubElement(size, 'width')

        width.text = str(shape1)
        height = ET.SubElement(size, 'height')
        height.text = str(shape0)
        depth = ET.SubElement(size, 'depth')
        depth.text = '3'
        segmented = ET.Element('segmented')
        segmented.text = '0'

        root = ET.Element('annotation')
        root.extend((folder, filename, path))
        root.extend((source, size, segmented))

        for box in boxes:
            object = ET.Element('object')
            name = ET.SubElement(object, 'name')
            name.text = box[4]
            pose = ET.SubElement(object, 'pose')
            pose.text = 'Unspecified'
            truncated = ET.SubElement(object, 'truncated')
            truncated.text = '0'
            difficult = ET.SubElement(object, 'difficult')
            difficult.text = '0'
            bndbox = ET.SubElement(object, 'bndbox')
            xmin = ET.SubElement(bndbox, 'xmin')
            xmin.text = str(box[0])
            ymin = ET.SubElement(bndbox, 'ymin')
            ymin.text = str(box[1])
            xmax = ET.SubElement(bndbox, 'xmax')
            xmax.text = str(box[2])
            ymax = ET.SubElement(bndbox, 'ymax')
            ymax.text = str(box[3])
            root.extend((object,))

        tree = ET.ElementTree(root)
        tree.write(xml_name)

        tree = ET.parse(xml_name)  # 解析movies.xml这个文件
        root = tree.getroot()  # 得到根元素,Element类
        self.pretty_xml(root, '\t', '\n')  # 执行美化方法
        tree.write(xml_name)

    def pretty_xml(self, element, indent, newline, level=0):  # elemnt为传进来的Elment类,参数indent用于缩进,newline用于换行
        if element:  # 判断element是否有子元素
            if (element.text is None) or element.text.isspace():  # 如果element的text没有内容
                element.text = newline + indent * (level + 1)
            else:
                element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * (level + 1)
                # else:  # 此处两行如果把注释去掉,Element的text也会另起一行
                # element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * level
        temp = list(element)  # 将element转成list
        for subelement in temp:
            if temp.index(subelement) < (len(temp) - 1):  # 如果不是list的最后一个元素,说明下一个行是同级别元素的起始,缩进应一致
                subelement.tail = newline + indent * (level + 1)
            else:  # 如果是list的最后一个元素, 说明下一行是母元素的结束,缩进应该少一个
                subelement.tail = newline + indent * level
            self.pretty_xml(subelement, indent, newline, level=level + 1)  # 对子元素进行递归操作

    def changeImages(self, folder, function_name, image_name, n):
        if function_name == "crop":
            function = self.__cropImage
        elif function_name == "tran":
            function = self.__translationImage
        elif function_name == "light":
            function = self.__changeLightofImage
        elif function_name == "noise":
            function = self.__addNoiseToImage
        elif function_name == "rotate":
            function = self.__rotateImage
        elif function_name == "flip":
            function = self.__flipImage

        image = cv2.imread(image_name + '.jpg')
        boxes = self.getBoxes(image_name)
        for i in range(1, n + 1):
            print(function_name + " image #" + str(i))
            change_img, change_boxes = function(copy.deepcopy(image), copy.deepcopy(boxes))
            print("Old boxes: ", boxes)
            print("New boxes: ", change_boxes)
            current_time = datetime.datetime.now().strftime("%Y_%m_%d_%H_%M_%S")
            current_num = str(random.randint(0, 9999))
            save_image_name = folder + '/' + current_time + '_' + str(i) + current_num + '.jpg'
            save_xml_name = folder + '/' + current_time + '_' + str(i) + current_num + '.xml'
            print(image_name)
            print("save image name: " + save_image_name)
            print("save xml name:   " + save_xml_name)
            cv2.imwrite(save_image_name, change_img)
            self.saveXML(save_image_name, save_xml_name, change_boxes, change_img.shape[1], change_img.shape[0])
            print("Save new image to current path: " + save_image_name)
            print("Save new xml to current path:   " + save_xml_name)
            print("\n")

    # 1 裁切
    def __cropImage(self, img, boxes):
        """
        裁切
        :param img: 图像
        :param bboxes: 该图像包含的所有boundingboxes,一个list,每个元素为[x_min,y_min,x_max,y_max]
        :return: crop_img:裁剪后的图像;crop_bboxes:裁剪后的boundingbox的坐标,list
        """
        # 裁剪图像
        w = img.shape[1]
        h = img.shape[0]

        x_min = w
        x_max = 0
        y_min = h
        y_max = 0

        # 最小区域
        for bbox in boxes:
            x_min = min(x_min, bbox[0])
            y_min = min(y_min, bbox[1])
            x_max = max(x_max, bbox[2])
            y_max = max(y_max, bbox[3])
            name = bbox[4]

        # 包含所有目标框的最小框到各个边的距离
        d_to_left = x_min
        d_to_right = w - x_max
        d_to_top = y_min
        d_to_bottom = h - y_max

        # 随机扩展这个最小范围
        crop_x_min = int(x_min - random.uniform(0, d_to_left))
        crop_y_min = int(y_min - random.uniform(0, d_to_top))
        crop_x_max = int(x_max + random.uniform(0, d_to_right))
        crop_y_max = int(y_max + random.uniform(0, d_to_bottom))

        # 确保不出界
        crop_x_min = max(0, crop_x_min)
        crop_y_min = max(0, crop_y_min)
        crop_x_max = min(w, crop_x_max)
        crop_y_max = min(h, crop_y_max)

        crop_img = img[crop_y_min:crop_y_max, crop_x_min:crop_x_max]

        # 裁剪bounding boxes
        crop_bboxes = list()
        for bbox in boxes:
            crop_bboxes.append([int(bbox[0] - crop_x_min), int(bbox[1] - crop_y_min),
                                int(bbox[2] - crop_x_min), int(bbox[3] - crop_y_min), bbox[4]])

        return crop_img, crop_bboxes

    # 2-平移
    def __translationImage(self, img, boxes):
        """
        平移
        :param img: img
        :param bboxes: bboxes:该图像包含的所有boundingboxes,一个list,每个元素为[x_min,y_min,x_max,y_max]
        :return: shift_img:平移后的图像array;shift_bboxes:平移后的boundingbox的坐标,list
        """

        # 平移图像
        w = img.shape[1]
        h = img.shape[0]

        x_min = w
        x_max = 0
        y_min = h
        y_max = 0

        for bbox in boxes:
            x_min = min(x_min, bbox[0])
            y_min = min(y_min, bbox[1])
            x_max = max(x_max, bbox[2])
            y_max = max(x_max, bbox[3])
            name = bbox[4]

        # 包含所有目标框的最小框到各个边的距离,即每个方向的最大移动距离
        d_to_left = x_min
        d_to_right = w - x_max
        d_to_top = y_min
        d_to_bottom = h - y_max

        # 在矩阵第一行中表示的是[1,0,x],其中x表示图像将向左或向右移动的距离,如果x是正值,则表示向右移动,如果是负值的话,则表示向左移动。
        # 在矩阵第二行表示的是[0,1,y],其中y表示图像将向上或向下移动的距离,如果y是正值的话,则向下移动,如果是负值的话,则向上移动。
        x = random.uniform(-(d_to_left / 3), d_to_right / 3)
        y = random.uniform(-(d_to_top / 3), d_to_bottom / 3)
        M = np.float32([[1, 0, x], [0, 1, y]])

        # 仿射变换
        shift_img = cv2.warpAffine(img, M,
                    (img.shape[1], img.shape[0]))  # 第一个参数表示我们希望进行变换的图片,第二个参数是我们的平移矩阵,第三个希望展示的结果图片的大小

        # 平移boundingbox
        shift_bboxes = list()
        for bbox in boxes:
            shift_bboxes.append([int(bbox[0] + x), int(bbox[1] + y), int(bbox[2] + x), int(bbox[3] + y), bbox[4]])

        return shift_img, shift_bboxes

    # 3-改变亮度
    def __changeLightofImage(self, img, boxes):
        """
        改变亮度
        :param img: 图像
        :return: img:改变亮度后的图像array
        """
        '''
        adjust_gamma(image, gamma=1, gain=1)函数:
        gamma>1时,输出图像变暗,小于1时,输出图像变亮
        '''
        flag = random.uniform(0.5, 1.5)  ##flag>1为调暗,小于1为调亮
        newBoxes = copy.deepcopy(boxes)
        newImage = exposure.adjust_gamma(img, flag)
        return newImage, newBoxes

    # 4-添加高斯噪声
    def __addNoiseToImage(self, img, boxes):
        """
        加入噪声
        :param img: 图像
        :return: img:加入噪声后的图像array,由于输出的像素是在[0,1]之间,所以得乘以255
        """
        newBoxes = copy.deepcopy(boxes)
        newImage = random_noise(img, mode='gaussian', clip=True) * 255
        return newImage, newBoxes

    # 5-旋转
    def __rotateImage(self, img, boxes):
        """
        旋转
        :param img: 图像
        :param boxes:
        :param angle: 旋转角度
        :param scale: 默认1
        :return: rot_img:旋转后的图像array;rot_bboxes:旋转后的boundingbox坐标list
        """
        '''
        输入:
            img:array,(h,w,c)
            bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
            angle:
            scale:默认1
        输出:
            
        '''
        # 旋转图像
        w = img.shape[1]
        h = img.shape[0]
        angle = random.uniform(-45, 45)
        scale = random.uniform(0.5, 1.5)
        # 角度变弧度
        rangle = np.deg2rad(angle)
        # 计算新图像的宽度和高度,分别为最高点和最低点的垂直距离
        nw = (abs(np.sin(rangle) * h) + abs(np.cos(rangle) * w)) * scale
        nh = (abs(np.cos(rangle) * h) + abs(np.sin(rangle) * w)) * scale
        # 获取图像绕着某一点的旋转矩阵
        # getRotationMatrix2D(Point2f center, double angle, double scale)
        # Point2f center:表示旋转的中心点
        # double angle:表示旋转的角度
        # double scale:图像缩放因子
        rot_mat = cv2.getRotationMatrix2D((nw * 0.5, nh * 0.5), angle, scale)  # 返回 2x3 矩阵
        # 新中心点与旧中心点之间的位置
        rot_move = np.dot(rot_mat, np.array([(nw - w) * 0.5, (nh - h) * 0.5, 0]))
        # the move only affects the translation, so update the translation
        # part of the transform
        rot_mat[0, 2] += rot_move[0]
        rot_mat[1, 2] += rot_move[1]
        # 仿射变换
        rot_img = cv2.warpAffine(img, rot_mat, (int(math.ceil(nw)), int(math.ceil(nh))),
                                 flags=cv2.INTER_LANCZOS4)  # ceil向上取整

        # 矫正boundingbox
        # rot_mat是最终的旋转矩阵
        # 获取原始bbox的四个中点,然后将这四个点转换到旋转后的坐标系下
        rot_bboxes = list()
        for bbox in boxes:
            x_min = bbox[0]
            y_min = bbox[1]
            x_max = bbox[2]
            y_max = bbox[3]
            name = bbox[4]
            point1 = np.dot(rot_mat, np.array([(x_min + x_max) / 2, y_min, 1]))
            point2 = np.dot(rot_mat, np.array([x_max, (y_min + y_max) / 2, 1]))
            point3 = np.dot(rot_mat, np.array([(x_min + x_max) / 2, y_max, 1]))
            point4 = np.dot(rot_mat, np.array([x_min, (y_min + y_max) / 2, 1]))

            # 合并np.array
            concat = np.vstack((point1, point2, point3, point4))  # 在竖直方向上堆叠
            # 改变array类型
            concat = concat.astype(np.int32)
            # 得到旋转后的坐标
            rx, ry, rw, rh = cv2.boundingRect(concat)
            rx_min = rx
            ry_min = ry
            rx_max = rx + rw
            ry_max = ry + rh
            # 加入list中
            rot_bboxes.append([rx_min, ry_min, rx_max, ry_max, name])
        return rot_img, rot_bboxes

    # 6-镜像
    def __flipImage(self, img, bboxes):
        """
        镜像
        :param self:
        :param img:
        :param bboxes:
        :return:
        """
        '''
        镜像后的图片要包含所有的框
        输入:
            img:图像array
            bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
        输出:
            flip_img:镜像后的图像array
            flip_bboxes:镜像后的bounding box的坐标list
        '''
        # 镜像图像
        import copy
        flip_img = copy.deepcopy(img)
        if random.random() < 0.5:
            horizon = True
        else:
            horizon = False
        h, w, _ = img.shape
        if horizon:  # 水平翻转
            flip_img = cv2.flip(flip_img, -1)
        else:
            flip_img = cv2.flip(flip_img, 0)
        # ---------------------- 矫正boundingbox ----------------------
        flip_bboxes = list()
        for bbox in bboxes:
            x_min = bbox[0]
            y_min = bbox[1]
            x_max = bbox[2]
            y_max = bbox[3]
            name = bbox[4]
            if horizon:
                flip_bboxes.append([w - x_max, y_min, w - x_min, y_max, name])
            else:
                flip_bboxes.append([x_min, h - y_max, x_max, h - y_min, name])

        return flip_img, flip_bboxes


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Give old image folder.')
    parser.add_argument('folder', help='old image folder')
    args = parser.parse_args()
    demo = ImageAugmentation()
    folder = os.listdir(args.folder)
    for filename in folder:
        if os.path.splitext(filename)[1] == '.jpg':  # 目录下包含.json的文件
            name = str(args.folder) + '/' + os.path.splitext(filename)[0]
            demo.changeImages(str(args.folder), "crop", name, 5)
            demo.changeImages(str(args.folder), "tran", name, 5)
            demo.changeImages(str(args.folder), "light", name, 5)
            demo.changeImages(str(args.folder), "noise", name, 5)
            demo.changeImages(str(args.folder), "rotate", name, 5)

后续

为简洁起见,代码中并没有加入错误处理的内容。如有需要,可以自行添加。
未来会加入更多图像增广的手段。

  • 2
    点赞
  • 15
    收藏
    觉得还不错? 一键收藏
  • 1
    评论
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值