Введение

Недавно я и моя команда участвовали в хакатоне от компании «Норникель». Мы выбрали трек «Грязные дела», где наша задача заключалась в разработке алгоритма компьютерного зрения для решения проблем на производстве.

Вот условия задачи

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

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

Переписывать всё на другой фреймворк или переносить веса модели было рискованно, и времени на это не было.

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

Описание задачи

Датасет состоял из 170 изображений с дефектами и 140 изображений с искусственно наложенными дефектами. Каждое изображение сопровождалось бинарной маской, где дефекты были помечены единицей, а фон — нулем. Это означало, что наша задача сводилась к бинарной сегментации, где для каждого пикселя нужно было классифицировать, принадлежит ли он к дефекту или к фону.

Пример данных

Выбор фреймворка и архитектуры модели

Решил не менять фреймворк, хотя на сервере организатора нельзя было отправлять решение, при котором использовался tensorflow, т.е все писали на pytorch, или переводили в ONNX. У меня разыгрался азарт все равно решить на tensorflow.

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

Unet — это архитектура сверточной нейронной сети, которая широко используется для решения задач сегментации изображений, в частности медицинских снимков. Она состоит из двух основных частей: энкодера (сжимающая часть) и декодера (восстанавливающая часть).

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

• Декодер — это часть сети, которая восстанавливает пространство изображения, увеличивая его разрешение. В качестве декодера я использовал архитектуру pix2pix, которая включает в себя слои upsampling (увеличение размера изображения) и downsampling (уменьшение размера). Эти операции помогают извлекать и восстанавливать пространственные зависимости изображения, такие как углы, контуры и текстуры, и таким образом точно выделять дефекты.

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

Реализация решения

Примерно определились, теперь код
Импортируем нужные библиотеки

import tensorflow as tf
import numpy as np
import glob
import matplotlib.pyplot as plt
import keras
import cv2
from tensorflow.keras import layers, models
from tensorflow_examples.models.pix2pix import pix2pix

Сформируем наш датасет, и разобьем на train и test.

image_dataset = tf.data.Dataset.from_tensor_slices(image_paths)
mask_dataset = tf.data.Dataset.from_tensor_slices(mask_paths)

n_train = int(len(image_dataset)*0.7)
n_test = int(len(image_dataset)*0.3)

dataset = tf.data.Dataset.zip((image_dataset, mask_dataset))
train_dataset = dataset.take(n_train)
test_dataset = dataset.skip(n_train)

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

class Augment(tf.keras.layers.Layer):
    def __init__(self, seed=42):
        super().__init__()
        self.augment_inputs = tf.keras.layers.RandomFlip(mode="horizontal", seed=seed)
        self.augment_labels = tf.keras.layers.RandomFlip(mode="horizontal", seed=seed)
    def call(self, inputs, labels):
        inputs = self.augment_inputs(inputs)
        labels = self.augment_labels(labels)
        return inputs, labels

Реализуем downsampling (информацию взял с документации)

base_model = tf.keras.applications.MobileNetV2(input_shape=[224, 224, 3], include_top=False)
layer_names = [
    'block_1_expand_relu',   # 64x64
    'block_3_expand_relu',   # 32x32
    'block_6_expand_relu',   # 16x16
    'block_13_expand_relu',  # 8x8
    'block_16_project',      # 4x4
]
base_model_outputs = [base_model.get_layer(name).output for name in layer_names]
down_stack = tf.keras.Model(inputs=base_model.input, outputs=base_model_outputs)
down_stack.trainable = False

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

up_stack = [
    pix2pix.upsample(512, 3),  # 4x4 -> 8x8
    pix2pix.upsample(256, 3),  # 8x8 -> 16x16
    pix2pix.upsample(128, 3),  # 16x16 -> 32x32
    pix2pix.upsample(64, 3),   # 32x32 -> 64x64
]

Сама Unet

def unet_model(output_channels:int):
  # Входной слой для модели.
  inputs = tf.keras.layers.Input(shape=[224, 224, 3])
  # Применение down_stack 
  skips = down_stack(inputs)
  # Берем последний элемент из списка skips самый глубокий слой
  x = skips[-1]
  # Переворачиваем список skips кроме последнего слоя
  skips = reversed(skips[:-1])
  for up, skip in zip(up_stack, skips):
    x = up(x)
    # Конкатенация текущего слоя x с соответствующим слоем из down_stack
    concat = tf.keras.layers.Concatenate()
    x = concat([x, skip])
  last = tf.keras.layers.Conv2DTranspose(
      filters=output_channels, kernel_size=3, strides=2,
      padding='same', activation='sigmoid')
      # Применяем последний слой
  x = last(x)
  return tf.keras.Model(inputs=inputs, outputs=x)

Далее запускаем обучение модели, использовать будем классические Adam со скоростью обучения 1e-5, функцию потерь — кросс-энтропию, а метрикой выбрана IoU (Intersection over Union). Количество эпох — 25.

model = unet_model(output_channels=1)
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-5), loss=tf.keras.losses.BinaryCrossentropy(), metrics = ['Accuracy', tf.keras.metrics.BinaryIoU()])

EPOCHS = 25
model.fit(train_batches, validation_data=train_batches, epochs=EPOCHS)

Результаты и выводы

Результаты

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

• Accuracy: 0.8983

• IoU (Intersection over Union): 0.8022

• Loss: 0.1662

• Validation Accuracy: 0.8788

• Validation IoU: 0.7796

• Validation Loss: 0.0572

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

Выводы

• Архитектура UNet с предобученным энкодером на основе MobileNetV2 показала хорошие результаты для задачи сегментации, что подтверждается высокой точностью и значением IoU. Несмотря на ограниченность данных, модель успешно справляется с задачей выделения дефектов на изображениях.

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

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

Таблица 1. Метрики модели на тренировочных и валидационных данных

Заключение

С полным кодом можно ознакомиться на моем GitHub. Буду рад выслушать ваши замечания, предложения и конструктивную критику. Также приглашаю подписаться на мой телеграм-канал, где я делюсь своими проектами в области машинного обучения и спортивными достижениями.

Ваши отзывы помогут мне совершенствоваться и двигаться дальше!