Машинное зрение на Python. Учим нейросеть отличать медведей от слонов

Ты наверняка слышал, что нейросети в последнее время стали чертовски хорошо справляться с распознаванием объектов на картинках. Наша же задача — научиться пользоваться этими нейросетями, ведь их мощь может пригодиться в самых разных случаях. В этой статье я расскажу, как задействовать ее при помощи самых распространенных инструментов: Python и библиотек Tensorflow и Keras.

Бинарная классификация

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

Подключим необходимые библиотеки.

from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout, Activation, BatchNormalization, AveragePooling2D
from tensorflow.keras.optimizers import SGD, RMSprop, Adam
## pip install tensorflow-datasets
import tensorflow_datasets as tfds
import tensorflow as tf
import logging
import numpy as np
import time

Создадим модель нейросети.

def dog_cat_model():
  model = Sequential()
  model.add(Conv2D(32, (3, 3), input_shape=(128, 128, 3), activation='relu'))
  model.add(MaxPooling2D(pool_size=(2, 2)))
  model.add(Conv2D(32, (3, 3), activation='relu'))
  model.add(MaxPooling2D(pool_size=(2, 2)))
  model.add(Flatten())
  model.add(Dense(units=128, activation='relu'))
  model.add(Dense(units=1, activation='sigmoid'))
  model.compile(optimizer=Adam(),
    loss='binary_crossentropy',
    metrics=['accuracy'])
  return model

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

Теперь нейросеть следует обучить. В исходном датасете — 25 тысяч изображений. Мы разделим его на две части для обучения и верификации, а каждое изображение обработаем. Размер изображения должен быть 128 × 128 пикселей, а цвет в диапазоне от нуля до единицы.

def dog_cat_train(model):
  splits = tfds.Split.TRAIN.subsplit(weighted=(80, 10, 10))
  (cat_train, cat_valid, cat_test), info = tfds.load('cats_vs_dogs',
    split=list(splits), with_info=True, as_supervised=True)

  def pre_process_image(image, label):
    image = tf.cast(image, tf.float32)
    image = image/255.0
    image = tf.image.resize(image, (128, 128))
    return image, label

  BATCH_SIZE = 32
  SHUFFLE_BUFFER_SIZE = 1000
  train_batch = cat_train.map(pre_process_image)
    .shuffle(SHUFFLE_BUFFER_SIZE)
    .repeat().batch(BATCH_SIZE)
  validation_batch = cat_valid.map(pre_process_image)
    .repeat().batch(BATCH_SIZE)

  t_start = time.time()
  model.fit(train_batch, steps_per_epoch=4000, epochs=2,
    validation_data=validation_batch,
    validation_steps=10,
    callbacks=None)
  print("Training done, dT:", time.time() - t_start)

Теперь мы запускаем обучение модели и сохраняем результаты в файл:

model = dog_cat_model()
dog_cat_train(model)
model.save('dogs_cats.h5')

На моем компьютере с видеокартой GeForce 1060 обучение сети заняло примерно восемь минут.

Теперь мы можем использовать обученную сеть. Напишем функцию распознавания изображения.

def dog_cat_predict(model, image_file):
  label_names = ["cat", "dog"]

  img = keras.preprocessing.image.load_img(image_file,
    target_size=(128, 128))
  img_arr = np.expand_dims(img, axis=0) / 255.0
  result = model.predict_classes(img_arr)
  print("Result: %s" % label_names[result[0][0]])

Попробуем распознать эти изображения

model = tf.keras.models.load_model('dogs_cats.h5')
dog_cat_predict(model, "cat.png#26759185")
dog_cat_predict(model, "dog1.png#26759185")
dog_cat_predict(model, "dog2.png#26759185")