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

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

Прежде всего нам предстоит разделить этот набор на 3 категории: обучающая, тестовая и проверочная. Первая предназначена для обучения модели. Рекомендуется искусственно увеличивать этот набор, фрагментируя, переворачивая или отзеркаливая изображения.

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

data_augmentation = tf.keras.Sequential([
  tf.keras.layers.experimental.preprocessing.RandomFlip('horizontal'),
  tf.keras.layers.experimental.preprocessing.RandomRotation(0.2),
])
for image, _ in train_dataset.take(1):
  plt.figure(figsize=(10, 10))
  first_image = image[0]
  for i in range(9):
    ax = plt.subplot(3, 3, i + 1)
    augmented_image = data_augmentation(tf.expand_dims(first_image, 0))
    plt.imshow(augmented_image[0] / 255)
    plt.axis('off')

Следующим шагом является подготовка набора данных, для загрузки в нейронную сеть. Все изображения должны иметь одинаковый размер и находиться в одном цветовом пространстве. Как правило, используется пространство RGB. В нем каждый пиксель имеет значения, лежащие в диапазоне от 0 до 255. Для обучения нужно использовать как можно меньшие значения, поэтому значения каждого пикселя делятся на 255, чтобы в итоге они лежали в пределе от 0 до 1.

train_ds = tf.keras.preprocessing.image_dataset_from_directory(
  data_dir,
  validation_split=0.2,
  subset="training",
  seed=123,
  image_size=(img_height, img_width),
  batch_size=batch_size)
normalization_layer = layers.experimental.preprocessing.Rescaling(1./255)

Теперь можно переходить к созданию и обучению модели

num_classes = 2

model = Sequential([
  layers.experimental.preprocessing.Rescaling(1./255, input_shape=(img_height, img_width, 3)),
  layers.Conv2D(16, 3, padding='same', activation='relu'),
  layers.MaxPooling2D(),
  layers.Conv2D(32, 3, padding='same', activation='relu'),
  layers.MaxPooling2D(),
  layers.Conv2D(64, 3, padding='same', activation='relu'),
  layers.MaxPooling2D(),
  layers.Flatten(),
  layers.Dense(128, activation='relu'),
  layers.Dense(num_classes)
])

model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

epochs=10
history = model.fit(
  train_ds,
  validation_data=val_ds,
  epochs=epochs
)

В процессе обучения модель оценивает точность работы после каждой эпохи. Делает это она на проверочных данных. Так как после каждой эпохи этот набор не меняется, нейросеть запоминает их и к последней эпохе оценка становится необъективна. Последующая проверка на тестовых данных показывает реальную точность работы обученной модели. Обычно она отличается на 5-7%. Как видно из графика ниже наш случай не исключение. При этом точность работы все равно довольно высокая.

Теперь необходимо сохранить обученную модель.

history.save('Cats_vs_Dogs')

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

image_batch, label_batch = test_dataset.as_numpy_iterator().next()
predictions = model.predict_on_batch(image_batch).flatten()
predictions = tf.nn.sigmoid(predictions)
predictions = tf.where(predictions < 0.5, 0, 1)
print('Predictions:\n', predictions.numpy())
print('Labels:\n', label_batch)


plt.figure(figsize=(10, 10))
for i in range(9):
  ax = plt.subplot(3, 3, i + 1)
  plt.imshow(image_batch[i].astype("uint8"))
  plt.title(class_names[predictions[i]])
  plt.axis("off")

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

Таким образом, можно подвести итог, что нейросеть по классификации изображений работает довольно точно и ее можно использовать в будущих проектах.