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

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

INFO

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

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

Научим нашу нейросеть распознавать, кошка или собака изображены на фото. Это классический пример, и к нему в 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")
Результат

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

Продолжение доступно только участникам

Материалы из последних выпусков становятся доступны по отдельности только через два месяца после публикации. Чтобы продолжить чтение, необходимо стать участником сообщества «Xakep.ru».

Присоединяйся к сообществу «Xakep.ru»!

Членство в сообществе в течение указанного срока откроет тебе доступ ко ВСЕМ материалам «Хакера», увеличит личную накопительную скидку и позволит накапливать профессиональный рейтинг Xakep Score! Подробнее

DmitrySpb79: IoT/Python Developer

Комментарии (2)

  • В статью не вникал! - нет времени.

    Но за одну лишь тематику! и поддержку автора - Спасибо!

    Продолжай!