simple image data augmentation using OpenCV. flip, rotation, resize, statuer.

image_data_augmentation.py

import os
import shutil
from math import ceil, floor
import numpy as np
import cv2

"""
Data augmentation. 
"""


class Op:
    def __init__(self, file_path, size, hflip_p=0.5, vflip_p=0.5,
                 resize_range=(0.5, 1.5), rotate_range=(-45, 45),
                 alfa_range=(0.9, 1.2), beta_range=(-10, 10)):
        self.root = file_path
        self.fns = os.listdir(file_path)

        self.hflip_p = hflip_p
        self.vflip_p = vflip_p
        self.resize_range = resize_range
        self.rotate_range = rotate_range
        self.alfa_range = alfa_range
        self.beta_range = beta_range
        self.size = size

    def get_transform_param(self):
        param = {}
        param['hflip'] = True if np.random.rand() < self.hflip_p else False
        param['vflip'] = True if np.random.rand() < self.vflip_p else False
        param['resize'] = np.random.rand() * (
            self.resize_range[1] -
            self.resize_range[0]) + self.resize_range[0]
        param['rotate'] = np.random.rand()*(self.rotate_range[1]-self.rotate_range[0])
        param['stature_alfa'] = np.random.rand() * (
            self.alfa_range[1] -
            self.alfa_range[0]) + self.alfa_range[0]
        param['stature_beta'] = np.random.rand() * (self.beta_range[1] -
                                                    self.beta_range[0]) + self.beta_range[0]
        return param

    def transform(self, img, param):
        hflip = param['hflip']
        vflip = param['vflip']
        if hflip:
            img = np.flip(img, axis=1)
        if vflip:
            img = np.flip(img, axis=0)
        resize = param['resize']
        img = cv2.resize(img, None, None, resize, resize)
        angle = param['rotate']
        img = self.rotate(img, angle)
        img = self.stature(img, param['stature_alfa'], param['stature_beta'])
        return img

    def resize_pad(self, img, size=32):
        img = img[:, :]
        r, c = img.shape
        bigger = max(r, c)

        img = cv2.resize(img, (floor(c/bigger*size), floor(r/bigger*size)))
        out = np.zeros((size, size), np.uint8)
        a = ceil((size-img.shape[0])/2.)
        b = ceil((size-img.shape[1])/2.)
        out[a:a+img.shape[0], b:b+img.shape[1]] = img
        return out

    def stature(self, img, alfa, beta):
        img = img.astype(np.float32)
        img = img*alfa + beta
        img[img > 255] = 255
        img[img < 0] = 0
        return img.astype(np.uint8)

    def rotate(self, img, angle):
        src_h, src_w = img.shape
        cx = src_w/2
        cy = src_h/2
        rm = cv2.getRotationMatrix2D((cx, cy), angle, 1)
        rm_inv = cv2.invertAffineTransform(rm)
        p = np.array([[0,     0, src_w, src_w, cx],
                      [0, src_h,     0, src_h, cy],
                      [1,     1,     1,     1, 1]], np.float32)
        p_new = np.matmul(rm_inv, p)
        new_w = int(np.max(p_new[0, :-1]) - np.min(p_new[0, :-1]))
        new_h = int(np.max(p_new[1, :-1]) - np.min(p_new[1, :-1]))
        p_new_nc = p_new[:, :-1].transpose()

        new_cx = cx - np.min(p_new[0, :-1])
        new_cy = cy - np.min(p_new[1, :-1])
        p_new_nc[:, 0] += (new_cx-cx)
        p_new_nc[:, 1] += (new_cy-cy)
        rm = cv2.getAffineTransform(
            p[0:-1, :-1].transpose()[0:3, :].astype(np.float32), p_new_nc[0:3, :].astype(np.float32))
        return cv2.warpAffine(img, rm, (new_w, new_h), None, cv2.INTER_CUBIC)

    def __getitem__(self, index):
        fn = self.fns[index]
        fn_abs = os.path.join(self.root, fn)

        trans_param = self.get_transform_param()

        img = cv2.imread(fn_abs)
        img = img[:, :, 0]

        img = self.transform(img, trans_param)
        img = self.resize_pad(img, self.size)
        return img, fn_abs, fn

    def __len__(self):
        return len(self.fns)

    def __repr__(self):
        fmt_str = 'diver dataset'
        return fmt_str


op = Op("tmp_pos", 60, hflip_p=0.5, vflip_p=0,
        resize_range=(0.5, 1.5), rotate_range=(-30, 30),
        alfa_range=(0.9, 1.5), beta_range=(-10, 20))
dst = "pos"
times = 50
show = False

if show:
    cv2.namedWindow("img")
for t in range(times):
    for i in range(len(op)):
        img, fn_abs, fn = op[i]

        cv2.imwrite(dst+'/'+str(t)+'-'+str(i)+'-'+fn, img)
        
        if show:
            cv2.imshow("img", img)
            key = cv2.waitKey(0)
            if key == 27:
                break
    if show:
        key = cv2.waitKey(0)
        if key == 27:
break