Классификация изображений свёрточной нейросетью

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

Подготовим окружение

Установим необходимые пакеты и настроим среду на вывод статичных графиков. В случае проблем с библиотеками выполните команду EngeePkg.purge(), которая устранит из вашей системы все пакеты кроме системных и позволит установить необходимые пакеты без проблем с совместимостью.

# Установка необходимых пакетов
Pkg.add(["Flux", "BSON", "ImageTransformations"])
gr();

Обучающие данные расположены в нескольких папках в каталоге "учебные данные". Названия папок считываются из файловой системы и становятся названиями классов. Мы также указываем, где находятся примеры для классификации (с неизвестными метками).

DATA_DIR = "$(@__DIR__)/учебные данные";
UNKNOWN_DIR = "$(@__DIR__)/неизвестно";

Обучение и проверка модели

В этом проекте с нуля мы обучим сверточную модель следующего вида:

В скрипте train.jl собраны функции, которые осуществляют:

Загрузку изображений из папок классов, ресайз до 128×128, нормализация
Подсчёт дисбаланса классов (сколько вилок, сколько ложек)
Разделение данных на train/test (стратифицированно, 75%/25%)
Сборка модели — свёрточная сеть с BatchNorm и Dropout
Сбалансированные батчи — чтобы в каждом батче были оба класса поровну
Аугментация — удвоение образцов в датасете при помощи отражения и высветления
Цикл обучения — forward pass, подсчёт потерь, backward pass, обновление весов
Оценка качества — точность, precision, recall на тесте
Early stopping — остановка если нет улучшений 8 эпох
Сохранение лучшей модели

include("$(@__DIR__)/_scripts/train.jl")
model, classes = train_model(DATA_DIR; epochs=25, batch_size=16, lr=0.0005, test_split=0.25);

Размер батча: 16, Скорость обучения: 0.0005
Доля тестовой выборки: 25.0%
Найдено классов: 2: ["вилка", "ложка"]

=== Распределение классов ===

Всего изображений: 335 (128×128)
  вилка: 189 изображений (56.4%)
  ложка: 146 изображений (43.6%)

=== Разделение данных ===
  Тренировочных: 252 (75.2%)
  Тестовых: 83 (24.8%)
Параметров модели: 607810

=== Обучение ===
  Эпоха  1/25, Train Loss: 1.3937, Train Acc: 44.4%, Test Acc: 44.6% ★  (precision/recall по классам: вилка: 65.2%/31.9%, ложка: 46.7%/77.8%)
  Эпоха  2/25, Train Loss: 1.3289, Train Acc: 48.0%, Test Acc: 42.2%    (precision/recall по классам: вилка: 75.0%/12.8%, ложка: 45.3%/94.4%)
  Эпоха  3/25, Train Loss: 1.3427, Train Acc: 66.3%, Test Acc: 61.4% ★  (precision/recall по классам: вилка: 63.8%/78.7%, ложка: 60.0%/41.7%)
  Эпоха  4/25, Train Loss: 1.1816, Train Acc: 57.1%, Test Acc: 49.4%    (precision/recall по классам: вилка: 72.7%/17.0%, ложка: 45.8%/91.7%)
  Эпоха  5/25, Train Loss: 1.1303, Train Acc: 79.0%, Test Acc: 75.9% ★  (precision/recall по классам: вилка: 71.7%/80.9%, ложка: 70.0%/58.3%)
  Эпоха  6/25, Train Loss: 1.0433, Train Acc: 73.4%, Test Acc: 68.7%    (precision/recall по классам: вилка: 77.8%/74.5%, ложка: 68.4%/72.2%)
  Эпоха  7/25, Train Loss: 0.9447, Train Acc: 78.2%, Test Acc: 67.5%    (precision/recall по классам: вилка: 78.0%/68.1%, ложка: 64.3%/75.0%)
  Эпоха  8/25, Train Loss: 0.9592, Train Acc: 79.4%, Test Acc: 73.5%    (precision/recall по классам: вилка: 90.6%/61.7%, ложка: 64.7%/91.7%)
  Эпоха  9/25, Train Loss: 1.0558, Train Acc: 84.1%, Test Acc: 75.9%    (precision/recall по классам: вилка: 93.8%/63.8%, ложка: 66.7%/94.4%)
  Эпоха 10/25, Train Loss: 0.8999, Train Acc: 81.3%, Test Acc: 67.5%    (precision/recall по классам: вилка: 78.9%/63.8%, ложка: 62.2%/77.8%)
  Эпоха 11/25, Train Loss: 0.97, Train Acc: 86.5%, Test Acc: 75.9%    (precision/recall по классам: вилка: 78.3%/76.6%, ложка: 70.3%/72.2%)
  Эпоха 12/25, Train Loss: 0.8364, Train Acc: 89.3%, Test Acc: 75.9%    (precision/recall по классам: вилка: 82.9%/72.3%, ложка: 69.0%/80.6%)
  Эпоха 13/25, Train Loss: 0.7675, Train Acc: 79.4%, Test Acc: 66.3%    (precision/recall по классам: вилка: 67.7%/89.4%, ложка: 76.2%/44.4%)

  ⏹ Ранний останов после 13 эпох (нет улучшения 8 эпох)

Загружена лучшая модель (Test Acc: 75.9%)

=== Результаты ===
  Лучшая точность на тесте: 75.9%

  Точность train/test: 78.2% / 72.3%
  ✓ Нет переобучения (разрыв 5.9%)
Обучение завершено! 🚀
Модель сохранена в model.bson

После обучения мы преврим, как ведет себя модель. Скрипт выполняет следующие шаги:

Загрузка модели, метаданных и неизвестных изображений — из BSON-файла извлекается архитектура сети, веса, названия классов и количество примеров в каждом классе обучающего набора, и переключение модели в режим оценки. Сканируется указанная папка, каждое изображение приводится к размеру 128×128, нормализуется в диапазон [-1, 1] и переводится в формат тензора W×H×C
Классификация каждого изображения — тензор подаётся на вход модели, выходные логиты преобразуются в вероятности через softmax, определяется класс с максимальной вероятностью и соответствующая уверенность и группировка по классам
Сортировка по уверенности и вывод статистики — внутри каждого класса изображения сортируются от самых уверенных предсказаний к наименее уверенным. Количество предсказанных изображений, процент от общего числа, средняя, максимальная и минимальная уверенность
Анализ смещения — сравнение процента предсказаний каждого класса с процентом этого класса в обучающем наборе, вывод индикатора

include("$(@__DIR__)/_scripts/visualize.jl")
results = classify_and_visualize(UNKNOWN_DIR);

=== Статистика обучающего набора ===
  вилка: 189 изображений
  ложка: 146 изображений

=== Результаты классификации ===
  вилка:
    Количество: 5 (50.0%)
    Средняя уверенность: 0.902
    Макс/мин уверенность: 1.0 / 0.752
  ложка:
    Количество: 5 (50.0%)
    Средняя уверенность: 0.626
    Макс/мин уверенность: 0.692 / 0.578

=== Анализ смещения ===
  вилка: 50.0% предсказано vs 56.4% обучено ✓
  ложка: 50.0% предсказано vs 43.6% обучено ✓

Классифицируем все изображения из папки "неизвестно" и выведем до 10 изображений для каждого класса.

include("$(@__DIR__)/_scripts/simple_mosaic.jl")
plot(create_simple_mosaic(UNKNOWN_DIR))

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

include("$(@__DIR__)/_scripts/predict_to_csv.jl")
predict_to_csv(UNKNOWN_DIR, confidence_threshold=0.6, output_csv="$(@__DIR__)/predictions.csv")

Обработано файлов: 10
  вилка: 5
  ложка: 4
  неизвестно: 1

Сохранено в /user/_retrain_resnet/predictions.csv

Заключение

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

Результаты, представленные в примере, получены при помощи нейросети, прошедшей десятки итераций процесса обучения. Чтобы получить хороший алгоритм классификации, нужно вариировать гиперпараметры, изучать датасет, менять длительность обучения, топологию нейросети (больше слоёв, включать или выключать батч-нормализацию и т.д.), размер батчей или оптимизатор, снять барьер преждевременной остановки или же просто поменять настройки генератора случайных чисел.

Row	Файл	Предсказанный_класс	Уверенность	Вероятность_вилка	Вероятность_ложка
	String	String	Float32	Float32	Float32
1	00000495.jpg	вилка	0.999757	0.999757	0.000242674
2	00000496.png	вилка	0.995816	0.995816	0.00418395
3	00000494.jpg	вилка	0.936342	0.936342	0.0636576
4	00000497.jpg	вилка	0.824374	0.824374	0.175625
5	00000498.jpg	вилка	0.752225	0.752225	0.247775
6	00000194.jpg	ложка	0.692343	0.307657	0.692343
7	00000193.png	ложка	0.62094	0.37906	0.62094
8	00000196.jpg	ложка	0.618481	0.381519	0.618481
9	00000195.jpg	ложка	0.618284	0.381716	0.618284
10	00000192.jpg	неизвестно	0.577928	0.422072	0.577928

Классификатор изображений для вашего датасета

Классификация изображений свёрточной нейросетью

Подготовим окружение

Обучение и проверка модели

Заключение

Тип

Краткое описание

Связанные материалы

Теги

Категории

Языки

Форматы

Уровень

Статус

Опубликовано

Обновлено

Источник