【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])
機会があれば、ぞれぞれの水増し結果も見てみましょう。