【Deep Learning】画像水増し自動化

今回はKeras-yolo3の学習で使用する画像の水増しを自動化するツールの備忘録を載せておきます。フォルダ構成等は合わせてもらう必要ありますが、基本的な動作としては指定したフォルダ内の画像に対して指定した水増し処理を行って、新たにフォルダを作成して画像を生成していきます。
コードは以下のようになります。

import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import os
from PIL import Image, ImageOps, ImageEnhance
import glob
import shutil
import sys

#インプットの画像フォルダ
inputImageFolderPath = '../学習データ/混合データ/品種/image'

#フルーツ固有ですが、全ての品種に対して処理する場合はTrue
allFruits = True

# 0:彩度変更, 1:コントラスト変更, 2:明度変更, 3:シャープネス変更, 4:拡大・縮小, 5:上下反転, 6:左右反転, 7
processList = [8]


# 出力のフォルダを指定。いじらなくていいはずです。
outputRoot = '../学習データ/水増しデータ/混合データ'

# 入力テキストフォルダ。いじらない。
inputTextFolderPath = inputImageFolderPath.replace('image', 'annotation')

# 彩度調整パラメータ。1.0以上で明るく、1.0未満で暗く
saturationRatio = 1.5

# コントラスト調整パラメータ。1.0以上で明るく、1.0未満で暗く
contrastRatio = 1.5

# 明度調整パラメータ。1.0以上で明るく、1.0未満で暗く
brightnessRatio = 1.5

# シャープネス調整パラメータ。1.0以上でくっきり、1.0未満でぼやける
sharpnessRatio = 1.5

# 大きさ調整パラメータ。縦横がこの倍率で大きくなります。
expansionRatio = 1.2

def makeFolder(root, folderName, imageFile):
    outputPath = os.path.join(root, folderName)
    if not os.path.exists(outputPath):
        os.mkdir(outputPath)
    
    
    outputPath = os.path.join(outputPath, imageFile.split('\\')[-5])
    if not os.path.exists(outputPath):
        os.mkdir(outputPath)
    
    outputAnnotationPath = os.path.join(outputPath, "annotation")
    if not os.path.exists(outputAnnotationPath):
        os.mkdir(outputAnnotationPath)
    outputImagePath = os.path.join(outputPath, "image")
    if not os.path.exists(outputImagePath):
        os.mkdir(outputImagePath)
    
    imagePathList = imageFile.split('\\')
    
    outputAnnotationPath = os.path.join(outputAnnotationPath, imagePathList[-3])
    if not os.path.exists(outputAnnotationPath):
        os.mkdir(outputAnnotationPath)
    outputImagePath = os.path.join(outputImagePath, imagePathList[-3])
    if not os.path.exists(outputImagePath):
        os.mkdir(outputImagePath)
    
    outputAnnotationPath = os.path.join(outputAnnotationPath, imagePathList[-2])
    if not os.path.exists(outputAnnotationPath):
        os.mkdir(outputAnnotationPath)
        outputAnnotationPath = os.path.join(outputAnnotationPath, 'classes.txt')
        with open(outputAnnotationPath, mode = 'w') as f:
            f.write("")
    outputImagePath = os.path.join(outputImagePath, imagePathList[-2])
    if not os.path.exists(outputImagePath):
        os.mkdir(outputImagePath)

        
def decideOutputFolderName():
    outputFolderName = ""
    for process in processList:
        if process == 0:
            outputFolderName += "彩度" + str(saturationRatio)
        elif process == 1:
            outputFolderName += "コントラスト" + str(contrastRatio)
        elif process == 2:
            outputFolderName += "明度" + str(brightnessRatio)
        elif process == 3:
            outputFolderName += "シャープネス" + str(sharpnessRatio)
        elif process == 4:
            outputFolderName += "大きさ" + str(expansionRatio)
        elif process == 5:
            outputFolderName += "上下反転"
        elif process == 6:
            outputFolderName += "左右反転"
        elif process == 7:
            outputFolderName += "90度回転"
        elif process == 8:
            outputFolderName += "270度回転"
    return outputFolderName

def copy(inputObject, outputObject):
    if os.path.exists(inputObject):
        shutil.copy(inputObject, outputObject)

def processing(image, coorList):
    for process in processList:
        if process == 0:
            image = changeSaturation(image, saturationRatio)
        elif process == 1:
            image= changeContrast(image, contrastRatio)
        elif process == 2:
            image = changeBrightness(image, brightnessRatio)
        elif process == 3:
            image = changeSharpness(image, sharpnessRatio)
        elif process == 4:
            image, coorList = changeSize(image, coorList, expansionRatio)
        elif process == 5:
            image, coorList = flipVertical(image, coorList)
        elif process == 6:
            image, coorList = flipHorizontal(image, coorList)
        elif process == 7:
            image, coorList = rotation90(image, coorList)
        elif process == 8:
            image, coorList = rotation270(image, coorList)
    return image, coorList

def changeSaturation(image, ratio):
    saturation = ImageEnhance.Color(image)
    return saturation.enhance(ratio)

def changeContrast(image, ratio):
    contrast = ImageEnhance.Contrast(image)
    return contrast.enhance(ratio)

def changeBrightness(image, ratio):
    brightness = ImageEnhance.Brightness(image)
    return brightness.enhance(ratio)

def changeSharpness(image, ratio):
    sharpness = ImageEnhance.Sharpness(image)
    return sharpness.enhance(ratio)

def changeSize(image, coorList, ratio):
    w, h = image.size
    resizedImage = image.resize((int(w * ratio), int(h * ratio)))
    
    resizedCoorList = []
    for oneCoorList in coorList:
        xmin = int(oneCoorList[0])
        ymin = int(oneCoorList[1])
        xmax = int(oneCoorList[2])
        ymax = int(oneCoorList[3])
        resizedCoorList.append([str(int(xmin * ratio)), str(int(ymin * ratio)), str(int(xmax * ratio)), str(int(ymax * ratio)), oneCoorList[4]])
    return resizedImage, resizedCoorList

def flipVertical(image, coorList):
    flippedImage = ImageOps.flip(image)
    
    flippedCoorList = []
    w, h = image.size
    for oneCoorList in coorList:
        xmin = int(oneCoorList[0])
        ymin = int(oneCoorList[1])
        xmax = int(oneCoorList[2])
        ymax = int(oneCoorList[3])
        flippedCoorList.append([str(xmin), str(h - ymax), str(xmax), str(h - ymin), oneCoorList[4]])
    return flippedImage, flippedCoorList

def flipHorizontal(image, coorList):
    mirrorImage = ImageOps.mirror(image)
    
    mirrorCoorList = []
    w, h = image.size
    for oneCoorList in coorList:
        xmin = int(oneCoorList[0])
        ymin = int(oneCoorList[1])
        xmax = int(oneCoorList[2])
        ymax = int(oneCoorList[3])
        mirrorCoorList.append([str(w - xmax), str(ymin), str(w - xmin), str(ymax), oneCoorList[4]])
    return mirrorImage, mirrorCoorList

def rotation90(image, coorList):
    rotatedImage = image.rotate(90, expand = True)
    
    rotatedCoorList = []
    w, h = image.size
    for oneCoorList in coorList:
        xmin = int(oneCoorList[0])
        ymin = int(oneCoorList[1])
        xmax = int(oneCoorList[2])
        ymax = int(oneCoorList[3])
        rotatedCoorList.append([str(ymin), str(w - xmax), str(ymax), str(w - xmin), oneCoorList[4]])
    return rotatedImage, rotatedCoorList

def rotation270(image, coorList):
    rotatedImage = image.rotate(270, expand = True)
    
    rotatedCoorList = []
    w, h = image.size
    for oneCoorList in coorList:
        xmin = int(oneCoorList[0])
        ymin = int(oneCoorList[1])
        xmax = int(oneCoorList[2])
        ymax = int(oneCoorList[3])
        rotatedCoorList.append([str(h - ymax), str(xmin), str(h - ymin), str(xmax), oneCoorList[4]])
    return rotatedImage, rotatedCoorList
    
    
outputFolderName = decideOutputFolderName()

imageName = r'\*\*.jpg'
if allFruits:
    imageName = r'\*\*\*.jpg'

### folder作成
for inputImageFile in glob.glob(inputImageFolderPath + imageName):
    splitInputImageFile = inputImageFile.split('/')
    if 'Thumbs.jpg' in splitInputImageFile[-1]:
        continue
    makeFolder(outputRoot, outputFolderName, inputImageFile)

### 画像ファイル取得
splitinputImageFolderPath = inputImageFolderPath.split('/')
outputImageFolderPath = os.path.join(outputRoot, outputFolderName, splitinputImageFolderPath[-3], splitinputImageFolderPath[-2], splitinputImageFolderPath[-1])
print(inputImageFolderPath)
for image in glob.glob(os.path.join(inputImageFolderPath, "**", "*.jpg"), recursive=True):
    
    splitImage = image.split('/')
    if 'Thumbs.jpg' in splitImage[-1]:
        continue
    originalImage = Image.open(image).convert("RGB")
    originalText = image.replace("image", "annotation", 1).replace(".jpg", ".txt")
    coorList = []
    if os.path.exists(originalText):
        with open(originalText) as f:
            for s_line in f:
                if len(s_line.split(',')) == 1:
                    pass
                else:
                    coorList.append(s_line.split(','))
    
    ### 画像水増し処理
    processedImage, processedCoorList = processing(originalImage, coorList)
    
    ### 出力
    outputImageFolder = os.path.join(outputRoot, outputFolderName, splitImage[-5], splitImage[-4], splitImage[-3], splitImage[-2])
    outputImageFile = os.path.join(outputImageFolder, splitImage[-1])
    if not os.path.isdir(outputImageFolder):
        print("folder made:", outputImageFolder)
        os.makedirs(outputImageFolder)
    processedImage.save(outputImageFile)
    
    outputTextFolder = outputImageFolder.replace("image", "annotation", 1)
    if not os.path.isdir(outputTextFolder):
        print("folder made:", outputTextFolder)
        os.makedirs(outputTextFolder)
    outputTextFile = outputImageFile.replace("image", "annotation").replace(".jpg", ".txt")
    with open(outputTextFile, mode = 'w') as f:
        for oneCoorList in processedCoorList:
            f.write(oneCoorList[0] + "," + oneCoorList[1] + "," + oneCoorList[2] + "," + oneCoorList[3] + "," + oneCoorList[4])

機会があれば、ぞれぞれの水増し結果も見てみましょう。