Время прочтения: 5 мин.

Компьютерное зрение большая область машинного обучения. Если вы решили попробовать свои силы в детектирование объектов, то сегодня я покажу, как, не прибегая к обучению GAN-моделей для генерации изображений быстро подготовить генератор используя библиотеки Numpy и Pillow.

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

Спойлер: в конце поста покажу как его использовать при обучении моделей для детекции объектов.

Суть алгоритма:

*Объект меняет свое местоположение на изображении, размер и угол от 0 до 360.

Перед началом работы я подготовил две папки:

1 – для фоновых изображений.

2 – для объектов.

Импортирую необходимые библиотеки:

import numpy as np 
import os 
import matplotlib.pyplot as plt 
from PIL import Image,ImageEnhance

Перейду к более детальному рассмотрению.

Вспомогательная функция для случайного распределения выбора класса объекта:

def flip_coin():
  if np.random.uniform()>0.5:
    return 'class1'
  return 'class2'

Еще одна вспомогательная функция для добавления шума на объекты.

def salt(image):
    koef = np.random.randint(5)
    koef = koef/10
    img = np.copy(np.array(image))
    salt = np.ceil(koef * img.size * 0.5)
    coords = [np.random.randint(0, i - 1, int(salt)) for i in img.shape]
    img[coords] = 1
    pepper = np.ceil(koef* img.size * 0.5)
    coords = [np.random.randint(0, i - 1, int(pepper)) for i in img.shape]
    img[coords] = 0
    return Image.fromarray(img)

Основная функция.

Создам массив из 0. X-основа для изображения, Y-основа для объекта.

Аргументы функции:

            batch_size – сколько изображений будет сгенерировано за один вызов функции,

            size_img – размер изображений.

get_image(batch_size=5,size_img=128):
X = np.zeros((batch_size, size_img, size_img, 3))
Y = np.zeros((batch_size, 6))

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

size = np.random.randint(int(size_img/10),int(size_img/5))
        file_obj = os.listdir('/content/obj')
        rand = np.random.randint(len(file_obj))
        cat_pil = Image.open('/content/obj'+'/'+file_obj[rand-1])
        temp_cat = cat_pil.thumbnail((size,size))
        temp_cat = temp_cat.rotate(np.random.randint(360), expand=True)
        cat = np.asarray(temp_cat) / 255.
        try:
          cat_x, cat_y,_ = cat.shape
        except ValueError:
          cat_x, cat_y = cat.shape

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

Также не забудьте исправить функцию flip_coin.

      size = np.random.randint(int(size_img/10),int(size_img/5)) #Размер объекта задается случайным образом 
      class_ = flip_coin() #Случайный выбор класса
# Обработка объекта 
      if (class_ == 'class1'):
        file_obj = os.listdir('/content/obj/1')
        rand = np.random.randint(len(file_obj))
        pil = Image.open('/content/obj/1'+'/'+file_obj[rand-1])  
#Случайным образом выбираем объект из папки
        class1 = 1.0
      else:
        file_obj = os.listdir('/content/obj/2')
        rand = np.random.randint(len(file_obj))
        pil = Image.open('/content/obj/2'+'/'+file_obj[rand-1])
        class2 = 1.0
      pil = salt(pil)
      temp = pil.thumbnail((size,size)) #Меняем размер объекта на случайно выбранный 
	size = np.random.randint(8)
	size_ = np.random.randint(8)
      temp = temp.rotate(np.random.randint(360),translate=(size,size_),           expand=True) #Поворачиваем объект на случайное значение 
	color = np.random.randint(5,25)
	enhancer = ImageEnhance.Brightness(temp)
	temp = enhancer.enhancer(color/10) #изменяем яркость объекта
	
	flip = np.random.randint(2)
	if flip==1: #Условие для зеркального отражения 
		temp = np.fliplr(temp)
		temp = Image.fromarray(temp)

      obj = np.asarray(temp) / 255. #Конвертируем входные данные в массив

      try:
        animal_x, animal_y, _ = obj.shape
      except ValueError:
        animal_x, animal_y = obj.shape

Создам пустое изображение, которое нам пригодится для конкатенации изображения и объекта:

bg = Image.new('RGB', (size_img, size_img))

Обработаю фон:

# Обработка фона изображения
      path = '/content/Fon'
      files = os.listdir(path)
      rand = np.random.randint(0, len(files))
      img = Image.open(path+'/'+files[rand-1])
      img = img.resize((size_img,size_img))
      bg.paste(img)

На данный момент у меня подготовлен объект и фон.

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

      x1 = np.random.randint(1,size_img - animal_x) #Задаем координаты для объекта
      y1 = np.random.randint(1,size_img - animal_y)
      h = animal_x
      w = animal_y
      animal_appear = 1.0
      bg.paste(temp, (x1, y1), mask=temp)
      pic = np.asarray(bg) / 255.
      X[i] = pic   #Вместо раннее созданного нулевого массива, подставляем собранное изображение 
      Y[i,0] = x1/float(size_img) #Определяем положение объекта относительно размеров изображения
      Y[i,1] = y1/float(size_img)
      Y[i,2] = animal_x / float(size_img)
      Y[i,3] = animal_appear * class1 // 1 # is_class1
      Y[i,4] = animal_appear * class2 // 1 # is_class2
  yield X,Y
x,y = next(gen_image(size_img=360)) 
plt.imshow(x[0])#Вывод нового изображения

Аннотация: первые 3 числа координаты объекта, остальные — кодировка класса.

Изображения, которые мы загрузили:

Объекты, которые мы загрузили:

Изображения, которые получили. Объект на изображении поворачивается от 0 до 360, зеркально отражается, обрезает углы, меняет свой размер, местоположение и яркость, а также добавляет шум. У второго объекта был фон, но на итоговом изображении его не будет. В этом вы можете убедиться, посмотрев ниже.

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

model.fit(gen_image(),validation_data=gen_image(),steps_per_epoch=60,validation_steps=10 ,epochs=10)

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