Упрощение процесса создания специально обученных моделей обнаружения объектов

Detecto - это библиотека Python, построенная на основе PyTorch, которая упрощает процесс построения моделей обнаружения объектов. Библиотека действует как облегченный пакет, который сокращает объем кода, необходимый для инициализации моделей, применяет переносное обучение к пользовательским наборам данных и выполняет логический вывод для изображений и видео.



Начиная

Чтобы увидеть, насколько просто начать работу с Detecto, давайте загрузим предварительно обученную модель из зоопарка моделей torchvision и запустим логический вывод на следующем изображении:

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

Приведенный выше код считывается с сохраненного изображения (в моем случае с именем «fruit.jpg»), генерирует для него прогнозы из предварительно обученной модели и отображает результаты:

Класс Detecto Model построен на архитектуре Faster R-CNN ResNet-50 FPN из подпакета моделей torchvision, который предварительно обучен на наборе данных COCO 2017. По умолчанию он может обнаруживать около 80 различных объектов, таких как фрукты, животные, автомобили, кухонная техника и многое другое.

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

Трансферное обучение

Существует пара руководств, в которых вы узнаете, как использовать предварительно обученную модель и применить переносное обучение к пользовательскому набору данных. Однако во многих из этих сценариев разработчикам приходится определять настраиваемые классы для своих наборов данных, вносить изменения в предварительно обученную модель или писать свои собственные методы обучения и визуализации с нуля. Иногда все, что вам нужно, - это быстро добиться хороших результатов. К счастью, сделать это с помощью Detecto очень просто.

Для начала, в Detecto есть класс Dataset (расширяющий класс PyTorch), который принимает любые данные в формате PASCAL VOC; то есть с каждым изображением связан файл аннотации XML (здесь - отличный инструмент для маркировки этого формата). Чтобы увидеть, как это будет выглядеть, вы можете иметь свой набор данных в любом из следующих форматов:

# All images and XML files in the same folder:
images/ 
| image0.jpg
| image0.xml
| image1.jpg
| image1.xml
| ...

# Images and XML files in separate folders:
images/ 
| image0.jpg
| image1.jpg
| ...
labels/
| image0.xml
| image1.xml
| ...

В обоих случаях чтение вашего набора данных так же просто, как следующее:

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

В приведенном выше примере после загрузки нашего набора данных из папки «images» мы инициализируем Model списком классов ['alien', 'bat', 'witch'], сообщающих ему, что мы хотим предсказать. Затем мы вызываем fit, который настроит предварительно обученную модель, чтобы узнать, как обнаруживать наши настраиваемые объекты.

Теперь давайте запустим модель на изображении и распечатаем результаты:

Выход:

['alien', 'bat', 'witch']
tensor([[ 569.2125,  203.6702, 1003.4383,  658.1044],
        [ 276.2478,  144.0074,  579.6044,  508.7444],
        [ 277.2929,  162.6719,  627.9399,  511.9841]])
tensor([0.9952, 0.9837, 0.5153])

Здесь наш главный прогноз был инопланетянином с координатами [569, 204, 1003, 658] и достоверностью 99,5%. Давайте также построим наши прогнозы:

Модуль visualize в Detecto поставляется со многими другими методами визуализации, включая обнаружение видео и прямую трансляцию с камеры. Вот как выглядит вывод на видео:

Когда вы закончите работу, вы можете сохранить и загрузить свои модели в файл .pth стандартным способом PyTorch:

Расширенное использование

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

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

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

В целом, Detecto по-прежнему является легкой библиотекой, поэтому после обучения модели вам могут потребоваться возможности точной настройки, которые еще не поддерживаются. К счастью, вам не нужно ограничиваться API Detecto: просто используйте метод get_internal_model для доступа к базовой модели PyTorch, которую вы затем можете интегрировать в свой код, как если бы это была любая другая модель PyTorch.

Заключение

В этой статье я представляю Detecto и показываю, как его можно использовать, чтобы значительно упростить обнаружение объектов с помощью PyTorch. Чтобы узнать больше, ознакомьтесь с этими ресурсами:

Не стесняйтесь обращаться к нам с любыми вопросами или сообщать о проблеме!