簡易画像データ水増しツール

image-padder.py

import glob
import cv2
import numpy as np
import os
import tqdm
from PIL import Image, ImageChops

def effect_gamma(image):
    gamma = 1.5
    gamma_cvt = np.zeros((256,1),dtype = 'uint8')
    for i in range(256):
        gamma_cvt[i][0] = 255 * (float(i)/255) ** (1.0/gamma)
    return  cv2.LUT(image,gamma_cvt)

def extract_filename(path):
    basename = os.path.basename(path)
    return os.path.splitext(basename)[0]

def copy_image(image):
    image = cv2.imread(file)
    file_name = extract_filename(file)
    cv2.imwrite('./output/' + file_name + '_original.png', image)        

def save_gamma_image(file):
    image = cv2.imread(file)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = effect_gamma(image)
    file_name = extract_filename(file)
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    cv2.imwrite('./output/' + file_name + '_gamma.png', image)

def save_gaussian_image(file):
    illust_im = Image.open(file).convert('RGB')
    noise_im = Image.effect_noise((illust_im.width, illust_im.height), 50).convert('RGB')
    result_im = ImageChops.multiply(illust_im, noise_im)
    file_name = extract_filename(file)
    result_im.save('./output/' + file_name + '_gaussian.png')

def save_salt_and_papper_image(file):
    image = cv2.imread(file)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    row,col,ch = image.shape
    # 白
    pts_x = np.random.randint(0, col-1 , 10000) #0から(col-1)までの乱数を千個作る
    pts_y = np.random.randint(0, row-1 , 10000)
    image[(pts_y,pts_x)] = (255,255,255) #y,xの順番になることに注意

    # 黒
    pts_x = np.random.randint(0, col-1 , 10000)
    pts_y = np.random.randint(0, row-1 , 10000)
    image[(pts_y,pts_x)] = (0,0,0)
    
    file_name = extract_filename(file)
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    cv2.imwrite('./output/' + file_name + '_salt.png', image)    

files = glob.glob("./input/*")

for file in tqdm.tqdm(files):
    print(file)
    copy_image(file)
    save_gamma_image(file)
    save_gaussian_image(file)
    save_salt_and_papper_image(file)