Pipes Counting — технический разбор (мобильная разработка + ML)

Приветствуем вас, дорогие читатели! В этой статье мы, ученики 10 класса из школы «Летово», хотим поделиться своим опытом разработки мобильного приложения PipesCounting, созданного для автоматизированного подсчета труб в пачке. За 6 месяцев мы прошли путь от идеи и до публикации приложения, в том числе сбор и расширение датасета, обучение моделей и сборки приложения. Здесь мы хотим поделиться нашим опытом и трудностями с которыми мы столкнулись, чтобы облегчить процесс для тех, кто будет заниматься чем‑то похожим, рассказать о возможных подходах и предупредить о сложностях.

С приложением PipesCounting мы заняли второе место в престижном международном конкурсе AI Challenge в категории «Металлургия», а также посетили с ним международную конференцию AI Journey.

Знакомьтесь с нашей командой: 

• Александр Поваров

• Иван Фенстер 

• Игнат Сухановский 

• Константин Спирин 

• Андрей Бушин 

  

Что мы хотели реализовать

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

С учётом этих требований мы выбрали YOLOv8: модель хорошо подходит для распознавания разных объектов, её просто дообучить под конкретный класс, а скорость позволяет держать хороший баланс между производительностью и качеством.

Данные

Проблема с данными заключалась в том, что размеченных фотографий труб в открытом доступе мало, а в существующих датасетах часто встречаются изображения, которые нам не подходят. Модель должна была распознавать четыре типа труб (круглые, квадратные, уголок, швеллер), причём для «уголка» и «швеллера» (см. Изображение 1) вообще не было размеченных данных в открытых датасетах.

Решение оказалось банальным, но трудоёмким: силами пяти человек с помощью CVAT мы самостоятельно разметили примерно 12 тыс. труб для «уголка» и «швеллера». Тем не менее этого всё равно не хватало для полноценного обучения, поэтому мы прибегли к расширению датасета.

1. Аугментация

Из-за ограниченного объёма данных, а также чтобы контролировать количество размеченных изображений, мы использовали аугментацию — искусственное расширение датасета с помощью простых преобразований исходных изображений. В нашем случае мы «раздували» датасет за счет небольших изменений картинки — например, поворота, отражения, изменения яркости/экспозиции или размытия. 

Для аугментации труб мы выбрали следующие методы: 

• Flip (Отражение) — Удваивает датасет, сохраняя логику укладки труб. 

• Rotation (±15°) — Учит модель работать с небольшими наклонами камеры.

• Brightness/Exposure — Имитирует тени и блики в цеху.

• Blur — Делает модель устойчивой к размытым кадрам с мобильных камер. 

• Mosaic — Собирает 4 изображения в одно, создавая сложные сцены.

Синтетические данные

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

1. Без паттерна

Трубы случайно повернуты и разбросаны по изображению. Помогает с распознаванием одиночных труб или труб вне пачек.

2. С простым паттерном

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

Трубы вложены друг в друга, как в большинстве реальных пачек.

Мы также применили CutOut и Blur к изображениям труб, чтобы снизить риск переобучения модели на конкретной синтетике. Большое преимущество синтетических данных — моментальная разметка (bounding box), что снимает с команды ручную аннотацию.

Что получилось?

После сбора данных и «раздувания» датасета у нас сформировалось несколько версий для экспериментов с обучением моделей.

Примеры фото:

Обучение YOLO:

Чтобы сделать приложение эффективным на производстве, одна из целей наше команды было ускорение инференса модели. Первая идея — уменьшить разрешение входного изображения. Экспериментально мы установили, что при снижении разрешения с 640×640 до 480×480 время инференса сокращается примерно вдвое (см Изображение 2). При этом качество детекции, очевидно, падает, поэтому важно выбрать баланс между скоростью и точностью.

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

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

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

Деление картинки на 4 части и постпроцессинг:

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

Наша идея была в разбиении изображения на четыре части с небольшим “нахлестом” и запуске YOLO отдельно на каждой из частей (см Изображение 3). После этого мы находим все накладывающиеся трубы, путем определения расстояния от центра одной трубы до края другой, из накладывающихся труб оставляем только большую.

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

Результаты методов постобработки: 

Разработка приложения

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

Наша команда решила сделать приложение полностью офлайн из-за возможного отсутствия интернета на производстве. Такое решение — серьезное ограничение: нельзя залить модель на Hugging Face Spaces и обращаться к ней по API; вместо этого модель должна быть непосредственно «вшита» в приложение. 

Так как в качестве модели мы использовали дообученную YOLO с постпроцессингом на Python, мы решили для разработки приложения использовать фреймворк Kivy. Он не типичен для мобильной разработки, зато упрощает интеграцию ML-части нашего продукта. Для сборки самого приложения мы выбрали идущую «в комплекте» утилиту Buildozer. 

Первый месяц выбранная стратегия казалась удачной: язык разметки Kivy мало отличается от других языков разметки (напоминает смесь HTML и CSS), а возможностей хватало — доступ к камере, к фотогалерее, фиксация кадра и последующая обработка реализовывались без особых проблем. Так последовательно был продуман и реализован весь функционал приложения (макет и интерфейс представлены ниже).

Оставался последний этап — собрать приложение под Android. В первую очередь, сложности возникли с прописанием спецификаций не самых популярных библиотек,  например Ultralytics (библиотека с YOLO), для которой не было готового «рецепта» по установке в Buildozer. Важной деталью было также то, что библиотеку OpenCV необходимо указывать в спецификации именно как opencv, а не как привычные cv2 или opencv-python.

Для сохранения размера приложения в пределах 100 MB, мы быстро отказались от установки PyTorch в приложение (он увеличивал размер до ~1,5 ГБ). Вместо этого мы использовали фреймворк TensorFlow Lite (он же LiteRT), созданный специально для деплоя, благо веса YOLO можно импортировать и в этом формате. С ним размер APK-файла обратно уменьшается до 70 MB. Тем не менее появляется новая задача: использование весов YOLO вне Ultralytics требует вручную реализовать pre-/post-processing (в том числе Non-Maximum Suppression).  Мы использовали для этого OpenCV: сборка очередной библиотеки через Buildozer требует много времени и траблшутинга, поэтому хочется использовать уже установленные библиотеки. Благо, нужный код на OpenCV лежит в туториалах на репозитории Ultralytics.

Конечно, от момента законченного проекта на Python до готовых APK и AAB файлов, выложенных в Google Play было ещё много сложностей: и с версиями Android NDK, и с необходимостью «подписать» получившийся файл, и с различающимся для разных версий Android списком разрешений для доступа к галерее пользователя. Все это усугублялось тем, что каждое действие (rebuild) требует часа-двух времени. Тем не менее, все эти сложности предсказуемы и специфичны, поэтому подробнее останавливаться на них не будем. Остановимся, лучше, на более позитивном моменте: работающий APK-файл, наконец, был получен.

Выводы

В результате мы получили полностью автономное мобильное приложение с вшитой YOLO-моделью, способное в офлайн-режиме определять четыре вида труб прямо на производственной площадке.

Главные достижения разработки:

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

2. Постобработка предсказаний — разбиение изображения, фильтрация выбросов и устранение дублей

3. Интеграция ML в Android-приложение через Kivy и TensorFlow Lite, что позволило обойтись без серверных вычислений

4. Создание удобного интерфейса для загрузки или моментального снимка и получения результата в пару секунд.

Мы прошли все этапы — от сбора и разметки данных до публикации приложения в Google Play, решив по дороге десятки «неочевидных» проблем с библиотеками, форматами весов, сборкой APK и оптимизацией размера.

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