image_classification

Это консольная утилита, позволяющая классифицировать предметы мусора на .jpg изображениях. Утилита "под капотом" использует классификатор RandomForestClassifier из библиотеки sklearn. Заранее обученная версия классификатора предоставляется вместе с утилитой, однако имеется возможность использования собственной модели, веса которой сохранены в .pkl файле. Для обучения дефолтной модели использовался этот датасет.

Датасет

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

Dataset link

Использование

Запуск классификации:

python src/image_classification predict -m <MODEL_FILE.pkl> -i <FILE.jpg>

поддерживаются только .jpg файлы

Запуск обучения:

python src/image_classification fit -s <SAVE_DIR> -d garbage-dataset/

Архитектура

Для возможности дальнейшего масштабирования был создан интерфейс IClassifier имеющий методы:

• predict() -- классификация изображения;
• fit() -- обучение классификатора;
• restore() -- загрузка предобученного классификатора;
• store() -- сохранение классификатора в .pkl формат.

ImageClassifier реализует интерфейс IClassifier. Предсказание основывется на 12 фичах (более подробно про фичи в hypotheses_report.ipynb).

Каждая фича реализует интерфейс IFeature:

• __call__() -- получение значения фичи из изображения;
• name() -- имя фичи. Создание IFeature позволяет гибко конфигурировать классификатор, разрешая включать/выключать отдельные фичи.

Сама утилита представлена классом ClassifierUtil, именно она отвечает за предсказание и обучение (внутри пользуется ImageClassifier) При инициализации ClassifierUtil используется CFG. Этот класс инкапсулирует в себе пользовательские аргументы, а также парсит их. При создании ClassifierUtil CFG валидируется в зависимости от типа запуска:

• predict -- классификация объекта на картинке;
• fit -- обучение классификатора.

API

Расширение функциональности утилиты происходит через реализацию интерфейсов IFeature и IClassifier. Кроме того наличие CFG позволяет интегрировать существующий классификатор ClassifierUtil в интерфейсы альтернативные консольному.

Тестирование

Тестирование проходило по трем сценариям:

1. Тестирования валидации пользовательской конфигурации
2. Тестирование обученного классификатора (использовался тестовый датасет, который не участвовал в обучении)
   • данные распознаются корректно
   • уровень уверенности не ниже минимального значения
3. Интеграционное тестирование утилиты на корректных сценариях
   • предсказание
   • обучение (использовался минимальный датасет для проверки пайплайна)

процент тестового покрытия доступен в артeфактах CI

Классификатор

В качестве классификатора был выбран Random Forest

Результаты на тестовых изображениях получились следующие:

Класс	Корректные предсказания	Всего изображений	Процент (%)
trash	87	190	45.8%
glass	365	613	59.6%
clothes	1006	1066	94.4%
metal	43	204	21.1%
plastic	171	397	43.0%
cardboard	216	365	59.2%
paper	183	336	54.4%
biological	136	200	68.0%
shoes	201	396	50.8%

Более подробный тест представлен в ноутбуке classifier_test.ipynb

NB: Классификатор и тестирование производилось с зафиксированным сидом 42 для разбиения датасета на тестовую и тренировочную выборку