【数据增强】对图片数据进行旋转、镜像、加噪等

'''
注：运行环境为win10，如果要在Linux下运行请修改文件路径即可
opencv3
python3.6
'''
import cv2
import numpy as np
import os.path
import copy
import matplotlib.pyplot as plt

def cv_imread(filePath): #读取中文路径的图片
    cv_img=cv2.imdecode(np.fromfile(filePath,dtype=np.uint8),cv2.IMREAD_UNCHANGED)
    # imdecode读取图像默认会按bgr通道，如果后续需要RGB可以通过如下转换
    #cv_img=cv2.cvtColor(cv_img,cv2.COLOR_BGR2RGB)
    return cv_img
def cv_imwrite(filePathName, img):

    try:
        _, cv_img = cv2.imencode(".jpeg", img)[1].tofile(filePathName)
        return True
    except:
        return False

def rotate(image, angle, center=None, scale=1.0):
    (h, w) = image.shape[:2]
    # If no rotation center is specified, the center of the image is set as the rotation center
    if center is None:
        center = (w / 2, h / 2)
    M = cv2.getRotationMatrix2D(center, angle, scale)
    rotated = cv2.warpAffine(image, M, (w, h))
    return rotated

# 通道数值拆分式高斯加噪
def noiseing(img):
    # img = cv2.cvtColor(rgbimg, cv2.COLOR_BGR2GRAY)
    part1 = 25.5         # noise part
    part2 = 255 - part1   # normal part
    w = img.shape[1]
    h = img.shape[0]
    newimg = np.zeros((h, w, 3), np.uint8)
    # row and col
    for x in range(0, h):
        for y in range(0, w-1, 2):  # Avoid exceeding boundaries
            r1 = np.random.random_sample()
            r2 = np.random.random_sample()
            z1 = part1 * np.cos(2 * np.pi * r2) * np.sqrt((-2) * np.log(r1))
            z2 = part1 * np.sin(2 * np.pi * r2) * np.sqrt((-2) * np.log(r1))

            fxy_0 = img[x, y, 0] + z1
            fxy_1 = img[x, y, 1] + z1
            fxy_2 = img[x, y, 2] + z1
            fxy1_0 = img[x, y + 1, 0] + z2
            fxy1_1 = img[x, y + 1, 1] + z2
            fxy1_2 = img[x, y + 1, 2] + z2
            a = [fxy_0,fxy_1,fxy_2,fxy1_0,fxy1_1,fxy1_2]
            a = sorted(a)
            part1_noise =  part1 / (a[-1]-a[0])
            part2_normal = part2 / 255
            
            newimg[x, y, 0] = (fxy_0 - a[0]) * part1_noise + fxy_0 * part2_normal
            newimg[x, y, 1] = (fxy_1 - a[0]) * part1_noise + fxy_1 * part2_normal
            newimg[x, y, 2] = (fxy_2 - a[0]) * part1_noise + fxy_2 * part2_normal
            newimg[x, y + 1, 0] = (fxy1_0 - a[0]) * part1_noise + fxy1_0 * part2_normal
            newimg[x, y + 1, 1] = (fxy1_1 - a[0]) * part1_noise + fxy1_1 * part2_normal
            newimg[x, y + 1, 2] = (fxy1_2 - a[0]) * part1_noise + fxy1_2 * part2_normal

    return newimg
def mirror(img):
    w = img.shape[1]
    h = img.shape[0]
    mir_img = copy.deepcopy(img)
    for i in range(h):
        for j in range(w):
            mir_img[i, w-1-j] = img[i,j]
    return mir_img



def main():
    file_dir = r"G:\win_linux中转站\人脸属性\Erised\Erised\ricemee-age\Images\25岁" #待读入的图片文件夹
    output_dir = r"G:\test"  #待写入的文件夹
    count = 0

    for num, class_name in enumerate(os.listdir(file_dir)):

        img_path = file_dir +'\\'+ class_name
        #image = cv2.imread(img_path)
        image = cv_imread(img_path)
        # Simple rotation 90 degrees
        rotated = rotate(image, 90)
        cv_imwrite(output_dir + '\\' + class_name[:-4] + '_ro90.jpg', rotated)
        # Rotate 45 degrees and add noise
        rotated = rotate(image, 45)
        cv_imwrite(output_dir + '\\' + class_name[:-4] + '_ro45.jpg', rotated)
        # Rotate 180 degrees and add noise
        rotated = rotate(image, 180)
        newimg = noiseing(rotated)
        cv_imwrite(output_dir + '\\' + class_name[:-4] + '_noise.jpg', newimg)
        # Image processing
        iLR = mirror(image)
        cv_imwrite(output_dir + '\\' + class_name[:-4] + '_mirr.jpg', iLR)
        count += 1
        print("The {}th image is processed successfully".format(count))
        if count == 2:
            print("The work is Done.")
            break
main()