Classification augmentation

aug.py

import cv2 
import numpy as np
import random
import os


def saveImg(img, name,location):
    img = cv2.resize(img,(416, 416))
    cv2.imwrite(location + name + '.jpg', img)


def brightness(img, alpha, beta,location):
    new_img = np.zeros(img.shape, img.dtype)
    alpha = alpha  # Simple contrast control
    beta = beta    # Simple brightness control
    
    for y in range(img.shape[0]):
        for x in range(img.shape[1]):
                new_img[y,x] = np.clip(alpha*img[y,x] + beta, 0, 255)
    saveImg(new_img, "brightness",location)


def gamma(img, correction,location):
    img = img/255.0
    img = cv2.pow(img, correction)
    img = np.uint8(img*255)
    saveImg(img, "gamma",location)


def noise(img ,prob,location):
    output = np.zeros(img.shape,np.uint8)
    thres = 1 - prob 
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            rdn = random.random()
            if rdn < prob:
                output[i][j] = 0
            elif rdn > thres:
                output[i][j] = 255
            else:
                output[i][j] = img[i][j]
    saveImg(output, "noise",location)


def resize(img, x, y,location):
    img = cv2.resize(img,(x, y))
    saveImg(img, "32_32",location)


#print('Found directory: %s' % dirName)
rootDir = 'training_images'
for dirName, subdirList, fileList in os.walk(rootDir):
	for fname in fileList:
		location=dirName+'/'+fname
		print("Augmentation of "+location+" started.")
		img = cv2.imread(location, 0)
		brightness(img, 2.0, 80,location)
		gamma(img, 3.0,location)
		resize(img, 32, 32,location)
		noise(img, 0.05,location)
		print("Augmentation of "+location+" finished.\n\n")