Как сгенерировать датасет изображений с помощью алгоритмов?

Сгенерируйте собственный датасет изображений для машинного обучения с помощью Python и его библиотек Pillow, Pandas и Numpy.

Эта статья вдохновлена такими недавно выпущенными NFT-коллекциями, как Cryptopunks, Sollamas, ON1 Force, Bored Ape Yacht Club. Среди этих проектов есть такие, где изображения нарисованы художниками, в то время как в других они сгенерированы алгоритмически. Для нас, как для специалистов в Data Science, интереснее второй вариант.

С помощью Python мы создадим коллекцию случайных изображений со множеством общих признаков. Конфигурировать финальные версии изображений будем путём добавления дополнительных элементов (персонажей и объектов) к базовым изображениям (фону).

В итоге у нас должен получиться датасет, каждое изображение в котором будет иметь особый атрибут: например, пират и пончик в пустыне или монстр и ключ в лесу.

Поскольку у нас нет цели создать произведение искусства, в качестве основы мы будем использовать свободно распространяемый пиксельный арт. С помощью этих базовых изображений мы сгенерируем сотни их производных. При желании вы можете выбрать другие изображения и создать датасет, отличающийся от датасета в статье.

Мы будем использовать атрибуты трёх типов: фон, персонаж и объект. Чтобы определить, насколько распространённым будет тот или иной атрибут, для каждого из них установим вероятностное распределение в процентах. Не бойтесь экспериментировать с параметрами — так вы сможете создать различные датасеты (например, датасет, в котором в 99% изображений используется один атрибут, а в оставшемся 1% — другой). Проводя такие эксперименты и создавая на их основе модели предсказания, можно многому научиться.

Вероятностное распределение свойств

Чтобы сгенерировать наш датасет, будем использовать следующее распределение свойств:

Приступаем к работе

Прежде всего, необходимо найти изображения, которые будут выступать в качестве базы. Главный критерий выбора — лицензионное разрешение на бесплатное скачивание и модификацию изображений. Все изображения, которые используются в данной статье, вы можете найти в этом репозитории.

Как только изображения выбраны, можем приступать к программированию алгоритмического генератора. Мы используем Python с библиотекой Pillow. Ниже представлен базовый скрипт для генерации изображений:

from os import path, mkdir

from PIL import Image

output_folder = "generated"

if not path.exists(output_folder):

mkdir(output_folder)background_file = path.join("backgrounds", "forest.png")

background_image = Image.open(background_file)

output_file = path.join(output_folder, "generated1.png")

background_image.save(output_file)

Папка generated создаётся для хранения сгенерированных файлов. Если такая папка уже существует, создавать её вторично не нужно. Далее загружаем одно из наших фоновых изображений и сохраняем его в папку generated. Итак, мы выполнили первую операцию с изображениями! Теперь дело за добавлением дополнительных элементов.

Зафиксируем размер фоновых изображений, чтобы у итоговых сгенерированных изображений он также был одинаковым. Чтобы облегчить себе работу, приведём масштаб самых крупных изображений к масштабу самых мелких. Масштаб персонажей менять не будем, поскольку цель данного проекта — функциональность, а не эстетика. Что касается объектов, изменим масштаб лишь тех из них, размеры которых сильно отличаются от других, а также тех, которые выступают за пределы фоновых изображений.

Итоговое разрешение фоновых изображений — 1920x1078 px.

Процесс несколько усложняется, когда дело доходит до добавления персонажей и объектов на фон. Нам необходимо определить, как именно мы расположим персонажей и насколько крупными они будут. Есть несколько вариантов определения положения персонажей:

• фиксированное для всех изображений;
• фиксированное для каждого типа персонажей/фонов;
• случайное.

Если мы выберем последний вариант, нужно будет соблюдать ограничения:

• избегать границ фона;
• избегать расположения, когда видна только малая часть изображения.

То же самое касается и объектов: необходимо определить их расположение, исходя либо из положения персонажа, либо из заранее определённого места на фоне. При создании данного датасета мы решили располагать объекты справа от персонажей.

В библиотеке Pillow используется система координат с началом в верхнем левом углу изображения:

Чтобы упростить себе задачу, поместим всех персонажей в одну и ту же фиксированную область на фоновых изображениях — она определяется отцентрированной по горизонтали точкой внизу. Размеры h и w прямоугольника — это высота и ширина персонажа. Чтобы указать Pillow, где располагать персонажа, нам необходимо вычислить положение жирной чёрной точки.

Эту точку можно вычислить по формуле ниже, где 1920 — ширина фона (мы делим её на 2, чтобы найти середину), а 1000 — результат округления 1078, чтобы оставить отступ между нижней частью персонажа и нижним краем изображения.

Чтобы расположить объекты справа от персонажей, воспользуемся той же самой формулой с единственным отличием — добавим отступ m, чтобы объект не затрагивал персонажа. В данном случае рабочее значение m — 30 px.

Если перевести все эти вычисления в код, получится следующее:

def generate_image(background, character, object, file_name):

  background_file = path.join("backgrounds", f"{background}.png")
  
  background_image = Image.open(background_file)
  
  #Создаём персонажа
  
  character_file = path.join("characters", f"{character}.png")
  
  character_image = Image.open(character_file)
  
  coordinates = (int(1920/2-character_image.width/2), int(1000-character_image.height)) #x, y
  
  background_image.paste(character_image, coordinates, mask=character_image)
  
  #Создаём объект
  
  if object != "none":
  
  object_file = path.join("objects", f"{object}.png")
  
  object_image = Image.open(object_file)
  
  coordinates = (int(1920/2+character_image.width/2+30), int(1000-object_image.height)) #x, y
  
  background_image.paste(object_image, coordinates, mask=object_image)
  
  output_file = path.join(output_folder, f"{file_name}.png")
  
  background_image.save(output_file)

А что насчёт данных?

Одновременно с генерацией изображений мы будем регистрировать их в базе данных вместе с их характеристиками. Для регистрации воспользуемся библиотекой pandas, данные будем сохранять в CSV-файл.

Создадим пустой DataFrame с названиями колонок. С каждой генерацией изображения в DataFrame будет добавляться новая строка. Когда сгенерируются все изображения, CSV-файл будет сохранён. Переменная num — последовательностный счётчик для наименования изображений: например, сгенерированные изображения будут называться generated1, generated2, generated2894, …

df = pd.DataFrame(columns = ["background", "character", "object", "generated image"])

for … …. ….

  #Код для генерации изображения  
  
  data = [background, character, object, f"generated{num}"]  
  
  s = pd.Series(data, index=df.columns)  
  
  df = df.append(s, ignore_index=True)
  
df.to_csv('data.csv', index=False)

Сколько изображений в итоге получится?

Вычислить итоговое количество изображений можно с помощью простой формулы — необходимо перемножить число возможных атрибутов в каждой категории. В нашем случае это 4 фона * 5 персонажей * 11 объектов = 220 изображений.

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

backgrounds = ["countryside", "desert", "forest", "glacial"]

characters = ["mage", "warrior", "pirate", "monster", "toadking"]

objects = ["none", "barrel", "anchor", "axe", "bomb", "key", "chest", "bananas", "cupcake", "donut", "heart",]

def generate_all_imgs():

  num = 0
  
  df = pd.DataFrame(columns = ["background", "character", "object", "generated image"])
  
  for background in backgrounds:
  
      for character in characters:
      
          for object in objects:
            
                generate_image(background, character, object, f"generated{num}")
                
                data = [background, character, object, f"generated{num}"]
                
                s = pd.Series(data, index=df.columns)
                
                df = df.append(s, ignore_index=True)
                
                num += 1  
                
   df.to_csv('data.csv', index=False)

Рандомизация

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

С помощью кода ниже сгенерируем заданное количество (total_imgs) изображений. То, какие свойства будут использованы в каждом изображении, будем определять случайным образом.

Мы будем использовать вероятности для каждого атрибута, определённые в таблице выше, а с помощью функции random.choice из Numpy будем выбирать свойство с учётом наиболее вероятного.

def generate_random_imgs(total_imgs):

    df = pd.DataFrame(columns = ["background", "character", "object", "generated image"])
  
    for num in range(total_imgs):
  
        background = np.random.choice(np.arange(0,len(backgrounds)), p=[0.3, 0.3, 0.3, 0.1])
        
        background = backgrounds[background]
        
        character = np.random.choice(np.arange(0,len(characters)), p=[0.4, 0.3, 0.2, 0.095, 0.005])
        
        character = characters[character]
        
        object = np.random.choice(np.arange(0,len(objects)), p=[0.3, 0.2, 0.1, 0.1, 0.1, 0.05, 0.05, 0.04, 0.03, 0.025, 0.005])
        
        object = objects[object]
        
        generate_image(background, character, object, f"generated{num}")
        
        data = [background, character, object, f"generated{num}"]
        
        s = pd.Series(data, index=df.columns)
        
        df = df.append(s, ignore_index=True)
        
    df.to_csv('data.csv', index=False)

Не используйте слишком большое количество параметров или их значений, иначе процесс генерации может длиться часами и занимать большие объёмы памяти.

Заключение

Итак, мы написали небольшую простую программу для генерации составных изображений из простых. Вооружившись итоговым датасетом, мы можем обучать нейронную сеть для классификации изображений по фону или персонажу.

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

Кроме того, наш генератор можно использовать как песочницу для аугментаций данных и пробовать создать новые изображения, где положение персонажей не будет статичным, где они будут повёрнуты определённым образом и т. п.

Весь упомянутый в статье код, а также использованные изображения вы можете найти в этом репозитории на GitHub.

Источники изображений:

• персонажи и объекты;
• фон «Ледник»;
• фон «Пустыня»;
• фон «Деревня»;
• фон «Лес».