Фиксация различных нарушений, контроль доступа, розыск и отслеживание автомобилей – лишь часть задач, для которых требуется по фотографии определить номер автомобиля (государственный регистрационный знак или ГРЗ).  В этой статье мы рассмотрим создание модели для распознавания с помощью Catalyst – одного из самых популярных высокоуровневых фреймворков для Pytorch. Он позволяет избавиться от большого количества повторяющегося из проекта в проект кода – цикла обучения, расчёта метрик, создания чек-поинтов моделей и другого – и сосредоточиться непосредственно на эксперименте. Сделать модель для распознавания можно с помощью разных подходов, например, путем поиска и определения отдельных символов, или в виде задачи image-to-text. Мы рассмотрим модель с несколькими выходами (multihead-модель). В качестве датасета возьмём датасет с российскими номерами от проекта Nomeroff Net. Примеры изображений из датасета представлены на рис. 1. Рис. 1. Примеры изображений из датасета   Необходимо разработать модель, которая на входе будет принимать изображение ГРЗ, а на выходе отдавать строку распознанных символов. Модель будет состоять из экстрактора фичей и нескольких классификационных “голов”. В датасете представлены ГРЗ из 8 и 9 символов, поэтому голов будет девять. Каждая голова будет предсказывать один символ из алфавита “1234567890ABEKMHOPCTYX”, плюс специальный символ “-” (дефис) для обозначения отсутствия девятого символа в восьмизначных ГРЗ. Архитектура схематично представлена на рис. 2. Рис. 2. Архитектура модели   В качестве loss-функции возьмём стандартную кросс-энтропию. Будем применять её к каждой голове в отдельности, а затем просуммируем полученные значения для получения общего лосса модели. Оптимизатор – Adam. Используем также OneCycleLRWithWarmup как планировщик leraning rate. Размер батча – 128. Длительность обучения установим в 10 эпох.  В качестве предобработки входных изображений будем выполнять нормализацию и преобразование к единому размеру. Далее рассмотрим основные моменты кода. Класс датасета (листинг 1) в общем обычный для CV-задач на Pytorch. Обратить внимание стоит лишь на то, как мы возвращаем список кодов символов в качестве таргета. В параметре label_encoder передаётся служебный класс, который умеет преобразовывать символы алфавита в их коды и обратно.  Листинг 1. Класс датасета   В классе модели (листинг 2) мы используем библиотеку PyTorch Image Models для создания экстрактора фичей. Каждую из классификационных голов модели мы добавляем в ModuleList, чтобы их параметры были доступны оптимизатору. Логиты с выхода каждой из голов возвращаются списком. Листинг 2. Класс модели   Центральным звеном, связывающим все компоненты и обеспечивающим обучение модели, является Runner. Он представляет абстракцию над циклом обучения-валидации модели и отдельными его компонентами. В случае обучения multihead-модели нас будет интересовать реализация метода handle_batch и набор колбэков. Метод handle_batch, как следует из названия, отвечает за обработку батча данных. Мы в нём будем только вызывать модель с данными батча, а обработку полученных результатов – расчёт лосса, метрик и т.д. – мы реализуем с помощью колбэков. Код метода представлен в листинге 3. Листинг 3. Реализация runner’а   Колбэки мы будем использовать следующие:   CriterionCallback – для расчёта лосса. Нам потребуется по отдельному экземпляру для каждой из голов модели. MetricAggregationCallback – для агрегации лоссов отдельных голов в единый лосс модели. OptimizerCallback – чтобы запускать оптимизатор и обновлять веса модели. SchedulerCallback – для запуска LR Scheduler’а. AccuracyCallback – чтобы иметь представление о точности классификации каждой из голов в ходе обучения модели. CheckpointCallback – чтобы сохранять лучшие веса модели. Код, формирующий список колбэков, представлен в листинге 4.   Листинг 4. Код получения колбэков Остальные части кода являются тривиальными для Pytorch и Catalyst, поэтому мы не станем приводить их здесь. Полный код к статье доступен на GitHub. Рис. 3. График лосс-функции модели в процессе обучения. Оранжевая линия – train loss, синяя – valid loss   В списке ниже перечислены некоторые ошибки, которые модель допустила на тест-сете:   Incorrect prediction: T970XT23- instead of T970XO123 Incorrect prediction: X399KT161 instead of X359KT163 Incorrect prediction: E166EP133 instead of E166EP123 Incorrect prediction: X225YY96- instead of X222BY96- Incorrect prediction: X125KX11- instead of X125KX14- Incorrect prediction: X365PC17- instead of X365PC178 Здесь присутствуют все возможные типы: некорректно распознанные буквы и цифры основной части ГРЗ, некорректно распознанные цифры кода региона, лишняя цифра в коде региона, а также неверно предсказанное отсутствие последней цифры. В статье мы рассмотрели способ реализации multihead-модели для распознавания ГРЗ автомобилей с помощью фреймворка Catalyst. Основными компонентами явились собственно модель, а также раннер и набор колбэков для него. Модель успешно обучилась и показала высокую точность на тестовой выборке. Спасибо за внимание! Надеемся, что наш опыт был вам полезен.  Больше наших статей по машинному обучению и обработке изображений: Data Science: предсказание бизнес-событий для улучшения сервиса Как мы используем алгоритмы компьютерного зрения: обработка видео в мобильном браузере с помощью OpenCV.js Тестируем комплементарную кросс-энтропию в задачах классификации текста