Python批量增强图像

Python批量增强图像

基于深度学习的计算机视觉任务，训练模型需丰富的数据集，以提高训练模型的泛化性。当训练数据集图像数量较少时，使用此增强程序批量处理image和mask

图像增强Python程序

import cv2
import glob
import numpy as np

def random_crop_flip(inp_img, gt_img):
    """
    随机裁剪和翻转图像
    :param inp_img:一个 HxWxC 输入图像
    :param gt_img:一个 HxW 输入图像
    :return:随机裁剪和翻转的图像
    """
    h, w = gt_img.shape

    rand_h = np.random.randint(h/8)  # / 除
    rand_w = np.random.randint(w/8)
    offset_h = 0 if rand_h == 0 else np.random.randint(rand_h)
    offset_w = 0 if rand_w == 0 else np.random.randint(rand_w)
    p0, p1, p2, p3 = offset_h, h+offset_h-rand_h, offset_w, w+offset_w-rand_w

    rand_flip = np.random.randint(10)  # 返回一个随机整数
    if rand_flip >= 5:
        inp_img = inp_img[::, ::-1, ::]  # [::-1] 反向排序 [:-1] 从位置0到-1之前的数
        out_img = gt_img[::, ::-1]

    return inp_img[p0:p1, p2:p3], gt_img[p0:p1, p2:p3]
    
def random_rotate(inp_img, gt_img, max_angle=25):
    """
    在+max_angle到-max_angle角度内随机旋转图像
    旋转max_angle角度
    :param inp_img:一个 HxWxC 输入图像
    :param gt_img:一个 HxW 输入GT图像
    :param max_angle:图像可以在任意方向旋转的最大角度
    :return:随机旋转的图像
    """
    # angle = np.random.randint(-max_angle, max_angle) # 返回一个随机角度
    angle = max_angle # 旋转max_angle角度
    h, w = gt_img.shape
    center = (w / 2, h / 2)

    M = cv2.getRotationMatrix2D(center, angle, 1.0)  # 给定一个旋转中心点的坐标、旋转角度和缩放因子，返回一个仿射变换矩阵 M
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])

    #  计算图像的新尺寸并调整旋转矩阵
    new_w = int((h * sin) + (w * cos))
    new_h = int((h * cos) + (w * sin))
    M[0, 2] += (new_w / 2) - center[0]
    M[1, 2] += (new_h / 2) - center[1]

    return cv2.warpAffine(inp_img, M, (new_w, new_h)), cv2.warpAffine(gt_img, M, (new_w, new_h))
    
def random_rotate_lossy(inp_img, gt_img, max_angle):
    """
    在+max_angle到-max_angle角度内随机旋转图像
    旋转max_angle角度
    :param inp_img:一个 HxWxC 输入图像
    :param gt_img:一个 HxW 输入图像
    :param max_angle:图像可以在任意方向旋转的最大角度
    :return:随机旋转的图像
    """
    # angle = np.random.randint(-max_angle, max_angle) # 在+max_angle到-max_angle角度内随机旋转图像
    angle = max_angle # 旋转max_angle角度
    h, w = gt_img.shape
    center = (w / 2, h / 2)
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    return cv2.warpAffine(inp_img, M, (w, h)), cv2.warpAffine(gt_img, M, (w, h))

def random_brightness(inp_img, gt_img):
    """
    随机扰动输入图像的亮度
    :param inp_img:一个 HxWxC 输入图像
    :return:亮度随机扰动的图像
    """
    contrast = np.random.rand(1) + 0.5  # 返回一个随机数，随机数位于[0, 1)中
    light = np.random.randint(-20, 20)
    inp_img = contrast * inp_img + light

    return np.clip(inp_img, 0, 255), gt_img  # np.clip 将一个nd.array的值限制在给定的上下界, 如果元素值小于下界则将值改为下界值a_min, 同理如果大于上界，则将值改为上界值a_max

def aug(image, mask, image_names, mask_names, k):
    image = cv2.imread(image)
    mask = cv2.imread(mask,0)
    image_rotate, mask_rotate = random_rotate_lossy(image, mask, max_angle=90)
    # image_crop_flip, mask_crop_flip = random_crop_flip(image, mask)
    # image_brigh, mask_brigh = random_brightness(image, mask)
    for image_name in image_names: 
        image_name = image_names[k]
        save_image_path = image_name.split('.')[0]+ '_rotate90'+ '.png'
        cv2.imwrite(save_image_path, image_rotate)
    for mask_name in mask_names: 
        mask_name = mask_names[k]
        save_mask_path = mask_name.split('.')[0]+ '_rotate90'+ '.png'
        cv2.imwrite(save_mask_path, mask_rotate)
    k += 1
        
    
if __name__ == '__main__':
    
    image_path = 'H:/image/train/'  
    mask_path =  'H:/gt/train/' 
    
    image_names = sorted(glob.glob(image_path  + '*.jpg'))
    mask_names = sorted(glob.glob(mask_path  + '*.png'))
    idx = 0
    for picture_name in image_names:
        print(mask_names[idx])
        aug(image_names[idx], mask_names[idx], image_names, mask_names, idx) 
        idx += 1

希望本文对大家有帮助，上文若有不妥之处，欢迎指正。

分享决定高度，学习拉开差距