Deep learning 数据增强代码

import os,cv2,shutil
import numpy as np
import random

import numpy as np
from PIL import Image
from PIL import ImageDraw
import cv2 as cv
import os.path
import shutil
import math

MAX_VALUE = 100

#对比度和亮度
def Contrast_and_Brightness(alpha, beta, img):
    blank = np.zeros(img.shape, img.dtype)
    # dst = alpha * img + beta * blank
    dst = cv2.addWeighted(img, alpha, blank, 1-alpha, beta)
    return dst


def darker(image,percetage):
    image_copy = image.copy()
    w = image.shape[1]
    h = image.shape[0]
    #get darker
    for xi in range(0,w):
        for xj in range(0,h):
            image_copy[xj,xi,0] = int(image[xj,xi,0]*percetage)
            image_copy[xj,xi,1] = int(image[xj,xi,1]*percetage)
            image_copy[xj,xi,2] = int(image[xj,xi,2]*percetage)
    return image_copy

def brighter(image, percetage):
    w = image.shape[1]
    h = image.shape[0]
    #get brighter
    for xi in range(0,w):
        for xj in range(0,h):
            image[xj,xi,0] = np.clip(int(image[xj,xi,0]*percetage),a_max=255,a_min=0)
            image[xj,xi,1] = np.clip(int(image[xj,xi,1]*percetage),a_max=255,a_min=0)
            image[xj,xi,2] = np.clip(int(image[xj,xi,2]*percetage),a_max=255,a_min=0)
    return image

#旋转
def rotate(image, angle, scale=1):
    w = image.shape[1]
    h = image.shape[0]
    #rotate matrix
    M = cv2.getRotationMatrix2D((w/2,h/2), angle, scale)
    #rotate
    image = cv2.warpAffine(image,M,(w,h))
    return image

#高斯噪声
def addGaussianNoise(image,percetage):
    G_Noiseimg = image.copy()
    w = image.shape[1]
    h = image.shape[0]
    G_NoiseNum=int(percetage*image.shape[0]*image.shape[1])
    for i in range(G_NoiseNum):
        temp_x = np.random.randint(0,h)
        temp_y = np.random.randint(0,w)
        G_Noiseimg[temp_x][temp_y][np.random.randint(3)] = np.random.randn(1)[0]
    return G_Noiseimg

#翻转
def flip(img,angle):
    fliped=cv2.flip(img,angle)
    return fliped

def img_augmentation(path, label_path):
    if path[-3:]=="PNG":
        print("ok")
        img=cv2.imread(path)
        c=random.uniform(0.7,1.3)
        b=random.uniform(-30,30)

        c_img=Contrast_and_Brightness(c,0,img)
        b_img=Contrast_and_Brightness(1,b,img)

        #b=random.uniform(1,1.5)
        #img_brighter=brighter(img,b)
        #d=random.uniform(0.6,0.95)
        #img_dark=darker(img,d)

        # angele_list=[-180,-90,-45,45,90,180]
        # a=random.choice(angele_list)
        a= random.uniform(-5,5)
        print("""""""""""""""""""""""""""""""""""", a)
        rotate_img=rotate(img,a)
        flip_list=[-1,0,1]
        f_p=random.choice(flip_list)
        flip_img=flip(img,f_p)

        (file_path,pic_name)=os.path.split(path)
        b_name="bright_"+pic_name #亮度
        c_name="contrast_"+pic_name #对比度
        r_name="rotate_"+pic_name #旋转
        f_name="flip_"+pic_name #翻转
        label_name = pic_name.split(".")[0] + ".txt"
        b_label_name = b_name.split(".")[0] + ".txt"
        c_label_name = c_name.split(".")[0] + ".txt"
        print(os.path.join(file_path,b_name))
        print(os.path.join(file_path,c_name))
        shutil.copyfile(os.path.join(label_path, label_name), os.path.join(label_path, b_label_name))
        cv2.imwrite(os.path.join(file_path,b_name),b_img)
        shutil.copyfile(os.path.join(label_path, label_name), os.path.join(label_path, c_label_name))
        cv2.imwrite(os.path.join(file_path,c_name),c_img)
        # cv2.imwrite(os.path.join(file_path,r_name),rotate_img)
        # cv2.imwrite(os.path.join(file_path,f_name),flip_img)

# 顺时针旋转
def Nrotation_angle_get_coor_coordinates(point, center, angle):
    src_x, src_y = point
    center_x, center_y = center
    radian = math.radians(angle)

    dest_x = round((src_x - center_x) * math.cos(radian) + (src_y - center_y) * math.sin(radian) + center_x)
    dest_y = round((src_y - center_y) * math.cos(radian) - (src_x - center_x) * math.sin(radian) + center_y)

    return (int(dest_x), int(dest_y))

#逆时针旋转
def Srotation_angle_get_coor_coordinates(point, center, angle):
    src_x, src_y = point
    center_x, center_y = center
    radian = math.radians(angle)

    dest_x = round((src_x - center_x) * math.cos(radian) - (src_y - center_y) * math.sin(radian) + center_x)
    dest_y = round((src_x - center_x) * math.sin(radian) + (src_y - center_y) * math.cos(radian) + center_y)

    return [int(dest_x), int(dest_y)]

def roted_img_and_label(imgpath, labelpath, angle):
    label_dir = os.listdir(label_path)
    for label_name in label_dir:
        img_name = label_name.split(".txt")[0] + ".png"
        im = Image.open(os.path.join(imgpath, img_name))
        f = open(os.path.join(label_path, label_name), "r")
        labels = f.readlines()

        # 指定逆时针旋转的角度
        a = random.uniform(-angle,angle)
        b = random.uniform(-angle,angle)
        for engle in [a, b]:
            im_rotate = im.rotate(engle)
            # a = ImageDraw.ImageDraw(im_rotate)  # 用a来表示
            # im_rotate.show()
            center = [im_rotate.width / 2, im_rotate.height / 2]
            with open(os.path.join(labelpath, "rotated_" + str(engle)[:3] + "_" + label_name), "w") as fw:
                for line in labels:
                    box = line.split()

                    xc = float(box[1]) * im.width
                    yc = float(box[2]) * im.height
                    bw = float(box[3]) * im.width
                    bh = float(box[4]) * im.height
                    point1 = [xc - bw / 2, yc - bh / 2]
                    point2 = [xc + bw / 2, yc + bh / 2]
                    point3 = [xc - bw / 2, yc + bh / 2]
                    point4 = [xc + bw / 2, yc - bh / 2]
                    rotated_point1 = Nrotation_angle_get_coor_coordinates(point1, center, engle)
                    rotated_point2 = Nrotation_angle_get_coor_coordinates(point2, center, engle)
                    rotated_point3 = Nrotation_angle_get_coor_coordinates(point3, center, engle)
                    rotated_point4 = Nrotation_angle_get_coor_coordinates(point4, center, engle)
                    new_box_point1 = [max(5, min(rotated_point1[0], rotated_point2[0], rotated_point3[0],
                                                 rotated_point4[0])),
                                      max(5, min(rotated_point1[1], rotated_point2[1], rotated_point3[1],
                                                 rotated_point4[1]))]
                    new_box_point2 = [min(im_rotate.width - 5, max(rotated_point1[0], rotated_point2[0],
                                                                   rotated_point3[0], rotated_point4[0])),
                                      min(im_rotate.height - 5, max( rotated_point1[1], rotated_point2[1],
                                                                     rotated_point3[1], rotated_point4[1]))]
                    rotated_center = Nrotation_angle_get_coor_coordinates([xc, yc], center, engle)
                    # a.rectangle(((new_box_point1[0], new_box_point1[1]),(new_box_point2[0], new_box_point2[1])),
                    # fill=None, outline='red', width=5)
                    rotated_w = abs(new_box_point2[0] - new_box_point1[0]) / im_rotate.width
                    rotated_h = abs(new_box_point2[1] - new_box_point1[1]) / im_rotate.height
                    new_box = box[0] + " " + str(rotated_center[0] / im_rotate.width) + " " + \
                              str(rotated_center[1] / im_rotate.height) + " " + str(rotated_w) + " " + str(rotated_h)
                    fw.write(new_box)
                    fw.write('\n')
            im_rotate.save(os.path.join(imgpath, "rotated_" + str(engle)[:3] + "_" + img_name))
if __name__ == "__main__":
    path=r'D:\project\cuicui\darknet\build\darknet\x64\data\voc\VOCks\JPEGImages'#改成自己的path
    label_path = r'D:\project\cuicui\darknet\build\darknet\x64\data\voc\VOCks\labels'
    for image_name in os.listdir(path):
        image_path = os.path.join(path, image_name)
        print(image_path)
        img_augmentation(image_path, label_path)
    roted_img_and_label(path, label_path, 5)