目标检测带标签数据增强代码

注意：我的数据集名字从1开始的，所以是for i in range（1，num），所以倒数第三句
str(i + num-1) 如果是从零开始就不用-1。另外文件名以及图片大小（1280，1024）也要根据自己的数据集做修改
图片平移：

import cv2 as cv
import numpy as np
import random
import xml.etree.ElementTree as et

# 图片平移 i + num
def photo_move(name, num):

    for i in range(1, num):
        a = random.choice([-1, 1])
        b = random.choice([-1, 1])
        c = random.randint(30, 50)
        d = random.randint(30, 50)
        e1 = random.randint(0, 255)
        f1 = random.randint(0, 255)
        g1 = random.randint(0, 255)

        # 定义位移矩阵，x方向移动a*c，y方向移动b*d,并且使用random库实现随机偏移
        matshift = np.float32([[1, 0, a*c], [0, 1, b*d]])
        # C:\Users\lx\Desktop\dataset\ban\photo_raw
        # 生成图片并保存
        img = cv.imread('C:/Users/86188/Desktop/dataset/' + name + '/photo_raw/' + name + str(i) + '.jpg')
        dst = cv.warpAffine(img, matshift, (1280,1024), borderValue = (e1, f1, g1))
        cv.imwrite('C:/Users/86188/Desktop/dataset/' + name + '/photo_enhancement/' + name + str(i + num-1) + '.jpg', dst)


        # 获得二维矩阵
        xml_path = 'C:/Users/86188/Desktop/dataset/' + name + '/label_raw/' + name + str(i) + '.xml'
        xml_file = et.parse(xml_path)
        root = xml_file.getroot()
        objs = root.findall('object')
        for obj in objs:
            bnd = obj.find('bndbox')
            # if obj.find('name').text != name:
            #      print('wrong')
            #     break
            xmin = float(bnd.find('xmin').text)
            ymin = float(bnd.find('ymin').text)
            xmax = float(bnd.find('xmax').text)
            ymax = float(bnd.find('ymax').text)
            ju_zhen = np.array([[xmin, xmax],
                                [ymin, ymax],
                                [1, 1]])

            # 保存变换后坐标点文件
            juzhen_new = np.dot(matshift, ju_zhen)      # 矩阵乘法
            if juzhen_new[0][0] + juzhen_new[1][0] > juzhen_new[0][1] + juzhen_new[1][1]:
                break

            bnd.find('xmin').text = str(int(juzhen_new[0][0]))
            bnd.find('ymin').text = str(int(juzhen_new[1][0]))
            bnd.find('xmax').text = str(int(juzhen_new[0][1]))
            bnd.find('ymax').text = str(int(juzhen_new[1][1]))

        xml_path_new = 'C:/Users/86188/Desktop/dataset/' + name + '/label_enhancement/' + name + str(i + num-1) + '.xml'
        xml_file.write(xml_path_new)  # 更新xml文件

        print('平移', i+num)

图片缩放 ：

def photo_zoom_in_and_out(name, num):

    for i in range(1, num):

        a = random.choice([0.8, 0.9, 0.95, 1.05, 1.08])
        b = random.choice([0.8, 0.9, 0.98, 1.05, 1.08])
        d = random.randint(0, 255)
        e = random.randint(0, 255)
        f = random.randint(0, 255)


        # 定义缩放矩阵
        matshift = np.float32([[a, 0, 0], [0, b, 0]])  # 定义缩放矩阵，长宽比为随机数

        # 生成并保存图片
        img = cv.imread('C:/Users/86188/Desktop/dataset/' + name + '/photo_raw/' + name + str(i) + '.jpg')
        dst = cv.warpAffine(img, matshift, (1280,1024), borderValue=(d, e, f))  # 画布尺寸是和原来一样
        cv.imwrite('C:/Users/86188/Desktop/dataset/' + name + '/photo_enhancement/' + name + str(i + 2* num-1) + '.jpg', dst)

        # 获得二维矩阵
        xml_path = 'C:/Users/86188/Desktop/dataset/' + name + '/label_raw/' + name + str(i) + '.xml'
        xml_file = et.parse(xml_path)
        root = xml_file.getroot()
        objs = root.findall('object')
        for obj in objs:
            bnd = obj.find('bndbox')
            xmin = float(bnd.find('xmin').text)
            ymin = float(bnd.find('ymin').text)
            xmax = float(bnd.find('xmax').text)
            ymax = float(bnd.find('ymax').text)
            ju_zhen = np.array([[xmin, xmax],
                                [ymin, ymax],
                                [1, 1]])


            # 保存变换后坐标点文件
            juzhen_new = np.dot(matshift, ju_zhen)
            if juzhen_new[0][0] + juzhen_new[1][0] > juzhen_new[0][1] + juzhen_new[1][1]:
                break

            bnd.find('xmin').text = str(int(juzhen_new[0][0]))
            bnd.find('ymin').text = str(int(juzhen_new[1][0]))
            bnd.find('xmax').text = str(int(juzhen_new[0][1]))
            bnd.find('ymax').text = str(int(juzhen_new[1][1]))

        xml_path_new = 'C:/Users/86188/Desktop/dataset/' + name + '/label_enhancement/' + name + str(i + 2* num-1) + '.xml'
        xml_file.write(xml_path_new)  # 更新xml文件

        print('缩放', i + 2 * num)

水平翻转：

def shuiping_photo_fan_zhuan(name, num):
    # 参数2 必选参数。用于指定镜像翻转的类型，0表示绕×轴正直翻转，即垂直镜像翻转；1表示绕y轴翻转，即水平镜像翻转；-1表示绕×轴、y轴两个轴翻转，即对角镜像翻转。
    for i in range(1, num):

        img = cv.imread('C:/Users/86188/Desktop/dataset/' + name + '/photo_raw/' + name + str(i) + '.jpg')
        dst = cv.flip(img, 1)
        cv.imwrite('C:/Users/86188/Desktop/dataset/' + name + '/photo_enhancement/' + name + str(i + 3* num-1) + '.jpg', dst)


        # 获得坐标值
        xml_path = 'C:/Users/86188/Desktop/dataset/' + name + '/label_raw/' + name + str(i) + '.xml'
        xml_file = et.parse(xml_path)
        root = xml_file.getroot()
        objs = root.findall('object')
        for obj in objs:
            bnd = obj.find('bndbox')
            x1 = float(bnd.find('xmin').text)
            # y1 = float(bnd.find('ymin').text)
            x2 = float(bnd.find('xmax').text)
            # y2 = float(bnd.find('ymax').text)
            # x 值 变化
            xmin = 1280 - x2
            xmax = 1280 - x1

            bnd.find('xmin').text = str(int(xmin))
            bnd.find('xmax').text = str(int(xmax))

        xml_path_new = 'C:/Users/86188/Desktop/dataset/' + name + '/label_enhancement/' + name + str(i +3*  num-1) + '.xml'
        xml_file.write(xml_path_new)  # 更新xml文件

        print('水平翻转', i + 3 * num)

垂直翻转：

def chuizhi_photo_fan_zhuan(name, num):
    # 参数2 必选参数。用于指定镜像翻转的类型，0表示绕×轴正直翻转，即垂直镜像翻转；1表示绕y轴翻转，即水平镜像翻转；-1表示绕×轴、y轴两个轴翻转，即对角镜像翻转。
    for i in range(1, num):
        img = cv.imread('C:/Users/86188/Desktop/dataset/' + name + '/photo_raw/' + name + str(i) + '.jpg')
        dst = cv.flip(img, 0)
        cv.imwrite('C:/Users/86188/Desktop/dataset/' + name + '/photo_enhancement/' + name + str(i + 4 * num-1) + '.jpg', dst)

        # 获得坐标值
        xml_path = 'C:/Users/86188/Desktop/dataset/' + name + '/label_raw/' + name + str(i) + '.xml'
        xml_file = et.parse(xml_path)
        root = xml_file.getroot()
        objs = root.findall('object')
        for obj in objs:
            bnd = obj.find('bndbox')
            x1 = float(bnd.find('xmin').text)
            y1 = float(bnd.find('ymin').text)
            x2 = float(bnd.find('xmax').text)
            y2 = float(bnd.find('ymax').text)
            # y 值 变化
            ymin = 1024 - y2
            ymax = 1024 - y1

            bnd.find('ymin').text = str(int(ymin))
            bnd.find('ymax').text = str(int(ymax))

            bnd.find('ymin').text = str(int(ymin))
            bnd.find('ymax').text = str(int(ymax))

        xml_path_new = 'C:/Users/86188/Desktop/dataset/' + name + '/label_enhancement/' + name + str(i + 4 * num-1) + '.xml'
        xml_file.write(xml_path_new)  # 更新xml文件

        print('垂直翻转', i + 4 * num)

垂直水平翻转：

def chuizhi_shuiping_photo_fan_zhuan(name, num):
    # 参数2 必选参数。用于指定镜像翻转的类型，0表示绕×轴正直翻转，即垂直镜像翻转；1表示绕y轴翻转，即水平镜像翻转；-1表示绕×轴、y轴两个轴翻转，即对角镜像翻转。
    for i in range(1, num):
        img = cv.imread('C:/Users/86188/Desktop/dataset/' + name + '/photo_raw/' + name + str(i) + '.jpg')
        dst = cv.flip(img, -1)
        cv.imwrite('C:/Users/86188/Desktop/dataset/' + name + '/photo_enhancement/' + name + str(i + 5 * num-1) + '.jpg', dst)

        # 获得坐标值
        xml_path = 'C:/Users/86188/Desktop/dataset/' + name + '/label_raw/' + name + str(i) + '.xml'
        xml_file = et.parse(xml_path)
        root = xml_file.getroot()
        obj = root.find('object')
        bnd = obj.find('bndbox')
        x1 = float(bnd.find('xmin').text)
        y1 = float(bnd.find('ymin').text)
        x2 = float(bnd.find('xmax').text)
        y2 = float(bnd.find('ymax').text)
        # y 值 变化
        xmin = 1280 - x2
        xmax = 1280 - x1
        ymin = 1024 - y2
        ymax = 1024 - y1

        bnd.find('xmin').text = str(int(xmin))
        bnd.find('xmax').text = str(int(xmax))
        bnd.find('ymin').text = str(int(ymin))
        bnd.find('ymax').text = str(int(ymax))

        xml_path_new = 'C:/Users/86188/Desktop/dataset/' + name + '/label_enhancement/' + name + str(i + 5 * num-1) + '.xml'
        xml_file.write(xml_path_new)  # 更新xml文件

        print('垂直水平翻转', i + 5 * num)

亮度和对比度改变：

def adjust_brightness_and_contrast(name, num):

    for i in range(1, num):

        a = random.choice([0.8,0.9, 1.05, 1.08])
        b = random.randint(-30, 30)
        img = cv.imread('C:/Users/86188/Desktop/dataset/' + name + '/photo_raw/' + name + str(i) + '.jpg')
        h, w, ch = img.shape
        blank = np.zeros([h, w, ch], img.dtype)
        dst = cv.addWeighted(img, a, blank, 1-a, b)
        cv.imwrite('C:/Users/86188/Desktop/dataset/' + name + '/photo_enhancement/' + name + str(i + 6*num-1) + '.jpg', dst)

        xml_path = 'C:/Users/86188/Desktop/dataset/' + name + '/label_raw/' + name + str(i) + '.xml'
        xml_file = et.parse(xml_path)
        xml_path_new = 'C:/Users/86188/Desktop/dataset/' + name + '/label_enhancement/' + name + str(i + 6*num-1) + '.xml'
        xml_file.write(xml_path_new)  # 更新xml文件

        print('亮度对比度',i + num)

最后：

num_表示的是你数据集每个类别的个数

list_lingjian = ['TB', 'D71', 'TB20K']
num_ = 45
k=0
photo_move(name=list_lingjian[k], num=num_)
photo_zoom_in_and_out(name=list_lingjian[k], num=num_)
#fang_she_change(name=list_lingjian[k], num=num_)
shuiping_photo_fan_zhuan(name=list_lingjian[k], num=num_)
chuizhi_photo_fan_zhuan(name=list_lingjian[k], num=num_)
chuizhi_shuiping_photo_fan_zhuan(name=list_lingjian[k], num=num_)
adjust_brightness_and_contrast(list_lingjian[k], num=num_ )