Простой инструмент аугментации изображений на Python с использованием Google Диска

Выполняет добавление к набору изображений следующих типов:
1) 3 дополнительных с разными углами поворота;
2) 16 с разными кадрами;
3) И 8 с разным контрастом.

В общей сложности одно изображение преобразуется в набор из 28 для более полного обучения распознаванию изображений в сверточных нейронных сетях.

Ссылка на блокнот Gooogle Colab: https://colab.research.google.com/drive/1mtOahe5_Fum-2agZhMmzU1cBo9PWqDc5?usp=sharing

Сначала подключим необходимые модули:

# import libraries

#!/usr/bin/python
from PIL import Image
import numpy as np
import os, sys
import pathlib
import time
from google.colab import drive

Затем мы подключаем Google Диск и указываем каталог, в котором будут находиться исходные изображения, а также куда поместить дополненные изображения.

Важный! Убедитесь, что указанные каталоги существуют на вашем диске:

# mount drive and set input and output folders (create manually if not exist)

INPUT_FOLDER = '/content/drive/MyDrive' + '/test/';
OUTPUT_FOLDER = '/content/drive/MyDrive' + '/test/output/';

drive.mount('/content/drive')

Далее мы опишем основные функции, которые изменят качества изображений (поворот, масштабирование и контрастность):

# augment functions

def aRotate(img, angle):
  imRotate = img.rotate(angle, expand=True)
  imRotate.save(f + '_rotated_' + str(angle + 90) + '.jpg', 'JPEG', quality=100)

def aZoom(img, scale):
  width, height = img.size
  imCrop = img.crop((width*scale, height*scale, width*(1-scale), height*(1-scale)))
  imCrop.save(f + '_cropped_' + str(scale) + '.jpg', 'JPEG', quality=100)

def aContrast(img, level):
  factor = (259 * (level + 255)) / (255 * (259 - level))
  def contrast(c):
    return 128 + factor * (c - 128)
  imContrast = img.point(contrast)
  imContrast.save(f + '_contrasted_' + str(level) + '.jpg', 'JPEG', quality=100)

Наконец, код, который выполняет процесс аугментации. Сначала мы создаем 4 копии изображения под разными углами, сохраняем их, а затем применяем к ним функции обрезки и контраста:

# rotate first

inputDir = os.listdir(INPUT_FOLDER)
for item in inputDir:
  if os.path.isfile(INPUT_FOLDER + item):
    im = Image.open(INPUT_FOLDER + item)
    f, e = os.path.splitext(OUTPUT_FOLDER + item)
    for i in range(0, 360, 90):
      aRotate(im, i)

# then crop and contrast all files with rotated

outputDir = os.listdir(OUTPUT_FOLDER)
for item in outputDir:
  if os.path.isfile(OUTPUT_FOLDER + item):
    im = Image.open(OUTPUT_FOLDER + item)
    f, e = os.path.splitext(OUTPUT_FOLDER + item)
    for i in np.arange(0.02, 0.1, 0.02):
      aZoom(im, i)
    aContrast(im, -50)
    aContrast(im, 50)

В результате мы должны получить 28 дополненных изображений из одного исходного изображения.