Ищем потерявшихся питомцев с помощью нейросетей
Представьте себе, ваш питомец вдруг сбежал отправился в увлекательное путешествие по соседским дворам. Паника, нервы, поиски, бумажные объявления… И хорошо, если ваш пушистик чипирован, но ведь о таком заботятся не все. Можно пойти шерстить порталы с потеряшками, но где искать? Как они работают? Тот еще квест! Хотелось бы автоматизировать этот поиск и здесь как нельзя кстати подойдут нейросети. Мы обучили сеточки для детектирования и распознавания мордочек, которые могут стать основой для удобного сервиса поиска потерявшихся животных.
Мы запилили удобного демо-бота с этими сеточками, он называется FindPet. И теперь с удовольствием представляем его вам и рассказываем, как мы его создавали.
Позвольте представиться, меня зовут Трофимова Оля, я специалист по машинному обучению в команде распознавания лиц компании NtechLab. Мы в компании давно и успешно занимаемся исследованиями в области искусственного интеллекта, преимущественно связанными с компьютерным зрением. Особенно хорошо компания известна своей технологией распознавания лиц и мы даже можем утверждать, что собаку съели на этой задаче. Мы отлично умеем работать с лицами людей и не раз занимали первое место в авторитетном международном соревновании NIST (National Institute of Standards and Technology). И почему бы не применить наши знания к животным, у них ведь тоже лица!
А теперь о том, как работает наш демо-бот. Он поддерживает два действия: «загрузить мордочку в базу» (если вы потеряли пушистика) и «поискать мордочку в базе» (если вы нашли пушистика). В первом случае вместе с фотографией питомца вы можете оставить имя и контакт (используйте любые, мы их все равно не сохраняем — ПД), чтобы нашедший вашего «искателя приключений» смог с вами связаться. Во втором случае бот сравнивает загруженное фото с базой пропавших питомцев и показывает карточку самого похожего питомца или сообщает о том, что такого питомца не нашел.
Порог сходства обученной модели составляет 70% (т. е. при сходстве выше 70% можно утверждать, что две фотографии принадлежат одному и тому же животному), но для бота мы снизили его до 60%, чтобы уж точно не пропустить потеряшку. Поэтому иногда, если вашего животного нет в базе, в выдаче может быть совсем другое животное, которое просто максимально похоже на ваше.
Что у бота под капотом?
Мы очень подробно рассказывали о том, как работает распознавание лиц, в нашей статье. В целом задача поиска пропавших питомцев сводится к двум основным нейросеткам: детектору и идентификатору мордочек. Сначала в своих экспериментах мы делили детектор на кошачий и собачий, но задумались: зачем усложнять? И сделали один общий класс «кошка или собака». Мы думали добавить нормализатор, который бы выравнивал мордочки и облегчал их распознавание, но победила лень, и мы решили, что можно сначала проверить, справится ли сетка без нормализатора, а потом уже по необходимости добавлять его (спойлер: добавлять не пришлось).
Обучение модели распознавания для мордочек питомцев ничем не отличается от обучения такой модели для человеческих мордочек лиц. Нам надо из входящей картинки получить вектор — закодированное представление мордочки, в котором будет содержаться вся необходимая для распознавания информация. Для получения таких векторов мы немного модифицировали архитектуру Res2Net. Вообще, большинство подобных задач решают с помощью ResNet‑подобных моделей. Они отлично справляются с кодированием информации и легко масштабируются, так как residual‑слои позволяют обучать очень глубокие модели и не терять важную информацию во время передачи признаков между слоями сети. Мы делали подходы к использованию других видов архитектур, например, трансформеров, но не смогли получить такие же хорошие результаты.
В качестве лосс‑функции мы использовали классический для задач распознавания CosFace, а метрикой качества выбрали TMR@FMR, про которую мы уже подробно рассказывали в нашей статье.
Мы провели несколько экспериментов: пробовали разные варианты препроцессинга данных, подбирали наиболее удачные параметры лосса, варьировали некоторые блоки и глубину нейросети (легкий вариант, соизмеримый с Res2Net50, и вариант потяжелее, соизмеримый с Res2Net100). Мы часто обучаем разные по «тяжести» сетки, т.к. иногда нам важна высокая точность, пусть даже ценой большего количества вычислительных ресурсов, а иногда важнее быстро получать результаты, и тогда высокой точностью можно пренебречь. На нашем тестовом датасете метрика TMR@FMR составила 92.77%. Но выбрать хорошую архитектуру — не самая важная часть процесса обучения нейросетей. На первом месте здесь, конечно же, данные.
Данные для обучения
Детектирование
Датасет для обучения детектора животных — это набор картинок и разметка к ним в виде координат области, содержащей в себе питомца. В интернете в свободном доступе можно найти очень большое количество датасетов с животными. Мы нашли несколько датасетов (например, раз, два) с разметкой именно мордочек, но основная часть датасетов содержит разметку животного целиком, что нам не подходило.
Решали мы эту проблему с помощью псевдоразметки. Сначала мы обучили детектор на имеющемся у нас датасете с «правильными» координатами мордочек и с его помощью разметили неразмеченные фотографии. Из этих фотографий мы выбрали те, где детектор был наиболее «уверен» в своих ответах. Мы объединили эти новые фотографии с первоначальным набором данных и на нем переобучили детектор. Во время такого итеративного процесса модель детектора учится на собственной разметке и со временем улучшается.
Тесты для оценки качества работы детектора должны быть очень сложными, максимально сложными! Именно такие мы подготовили. Вот вы можете найти здесь собак?
Первая версия нашего детектора с задачей справилась плохо:
И нам пришлось добавлять в качестве отрицательных примеров (фотографий без какой‑либо разметки) в обучающий датасет маффины, а заодно еще людей и мартышек.
Конечно же, помимо пары прикольных «чихуахуа vs маффины» квестов, мы использовали полноценные тестовые датасеты, составленные из изображений, не участвовавших в обучении. Всего тестовая выборка состояла из ~ 2000 фотографий кошек и собак.
Для оценки работы детектора в качестве метрики мы использовали mAP@0.5 (mean Average Precision при Intersection over Union = 0.5) и mAP@0.5:0.95 (mean Average Precision при Intersection over Union из отрезка [0.5:0.95]). Финальными метриками нашего детектора кошек и собак были mAP@0.5:0.95 = 0.8888, mAP@0.5 = 0.9948, и этот результат нас вполне устраивал.
Распознавание
Чтобы обучить сетку для идентификации, нужны не просто фотографии питомцев, а несколько размеченных фотографий одного и того же питомца, желательно с разным ракурсом, светом, может, даже возрастом, что‑то вроде того:
Здесь чем больше фотографий одной и той же мордочки, тем лучше. Но даже котик с единственной фотографией будет полезен в обучении, т. к. научит сетку не путать с ним другие мордочки из датасета. По количеству уникальных животных в датасете ограничений сверху также нет. Мы смогли найти несколько датасетов в интернете, но для хорошей точности этого было недостаточно, и мы стали думать, где бы взять фотки животных, сгруппированных по ID.
Нам нужно было два набора данных — большой набор для обучения и небольшой для тестирования нейросети. Найти достаточно большой датасет, не содержащий в себе ошибок разметки, — задача очень сложная, но главное, что она необязательная. Очевидные ошибки довольно просто отфильтровать (об этом подробнее будет ниже), а небольшое наличие оставшихся ошибок компенсируется большим размером датасета.
А вот для тестов нужен точно‑точно не содержащий ошибок набор данных, и автоматическая фильтрация тестового датасета крайне нежелательна.
Мы попробовали краудсорсинг. Разместили задание, в котором просили людей пофотографировать своих животных с разных сторон, люди такое делают бесплатно и с удовольствием каждый день, почему бы им немного не заработать на этом? Фотографий было много, но там, помимо кошек и собак, было что угодно, начиная с фоток стен и заканчивая мемами на турецком языке. Мы даже сделали отдельно небольшую сетку‑классификатор, чтобы отсеивать очевидный мусор, но в целом быстро поняли, что много данных мы так не соберем.
Параллельно мы запустили локальный краудсорсинг — создали чат на всех сотрудников компании, в котором попросили покидать фотки своих питомцев (мой самый любимый рабочий чат, честное слово).
Локально мы собрали пару тысяч фотографий животных, для обучения это не годилось, но для тестов — идеальный вариант! Чем более приближены фотографии в тестовом датасете к реальному домену, тем лучше, а здесь как раз люди фотографировали своих питомцев без четких инструкций на разные устройства в разных условиях.
А еще в сборе тестового датасета нам любезно согласился помочь приют для животных Дубовая роща. Поездка в приют — это прекрасный способ хорошо провести время быстро собрать данные, а что самое главное, качество такого датасета будет зависеть только от нас самих.
Так мы собрали отличный тестовый датасет из 10 тысяч фотографий полутора тысяч пушистиков.
Параллельно мы искали данные для обучения модели, и данных нужно было очень много. Мы нашли золото — сайт-агрегатор приютов и центров спасения животных Urgent Need for Pet Adoption — Find Dogs & Cats & More | Petfinder. Датасет прямо такой, как нужно: фотографии питомцев сгруппированы по отдельным карточкам, все отснято в разных условиях на разные устройства, даже есть данные о породе, поле и возрасте на случай, если мы решим когда‑то использовать это в поиске.
Наверное, вы скажете, что парсить чужие сайты без разрешения нельзя. И даже если там нет никакой защиты от скачивания, все равно нельзя. И мы с вами, конечно же, полностью согласны!
Так или иначе, мы смогли собрать достаточно большой датасет со структурой, как у Petfinder.
Фильтрация
Чтобы из набора фотографий составить датасет, мы прогнали их через наш детектор мордочек, и оставалось только «почистить» его. Самая распространенная проблема нашего датасета — это наличие нескольких питомцев на одном фото:
В датасете это выглядит как подписанные именем «Jensen and Hackles» фотографии мордочек разных кошек. Самый простой вариант решения такой проблемы — просто удалять ID, содержащие несколько животных на одной фотографии. Но это ленивый вариант на случай, если у вас огромный датасет и не страшно пожертвовать несколькими ID. В идеальном случае мы хотим как‑то понять, что это два разных животных, и сохранить их по отдельности как «Jensen and Hackles 1» и «Jensen and Hackles 2», как минимум мы хотим понять, чьих фотографий в этой папке больше, и оставить только его мордочку. Глазами это сделать несложно (почти всегда), но нам, конечно, надо автоматизировать этот процесс.
Для старта нам нужен хороший датасет, пусть даже и небольшой. Например, такой. На нем мы обучим нейросеть получать из фотографий мордочек векторы‑эмбеддинги — закодированные представления мордочек.
На следующем этапе прогоняем фотографии из нашего «грязного» датасета этой сеткой и для каждой мордочки получаем описывающий ее вектор‑эмбеддинг. На основании схожести (cosine similarity) этих векторов внутри каждого ID мы выделяем несколько кластеров с помощью DBSCAN с некоторым порогом схожести.
Поскольку мы не слишком доверяем сетке, обученной на малом количестве данных, то и порог задаем нестрогий. Можно попробовать подобрать хороший порог «в бою»: учим сетку на датасетах, отфильтрованных с разными порогами, а потом по лучшему бенчмарку выбираем лучший порог.
Второй этап фильтрации — поиск повторяющихся ID. Может случиться такое, что в датасете есть набор «Jensen and Hackles», а потом еще отдельно «Jensen» и «Hackles». Если оставить это как есть, то сетка будет учиться разделять «Jensen and Hackles» и «Jensen», что приведет к ошибкам в тестах. Чтобы этого избежать, надо как‑то сравнить между собой похожесть разных ID. Для этого мы будем сравнивать между собой центроиды — средний вектор‑эмбеддинг для каждого ID. Здесь, опять же, важно хорошо подобрать порог. Если схожесть двух центроидов выше некоторого порога — мы считаем, что там одно и то же животное, и тогда либо объединяем их, либо оставляем только тот ID, который содержит больше фотографий.
Конечно, это менее надежно, чем «ручная» разметка и фильтрация, зато это быстро и почти бесплатно! В датасете, скорее всего, останется некоторое количество ошибок, но если датасет достаточно большой, то сетка с этим справится.
В результате фильтрации у нас получился датасет из ~ 150 тысяч пушистиков, всего 500 тысяч фотографий. Мы решили, что для первой версии модели этого вполне достаточно.
Про проблемы
Самая большая проблема нашего демо-бота для тестирования — это малая наполненность базы. Когда вы ищете пушистика, бот попытается выдать вам карточку максимально похожего животного. Но если этого животного в базе нет, то бот вернет ближайшую по схожести мордочку, которая может принадлежать вообще другой породе (иногда и виду).
Мы могли бы поднять порог схожести для выдачи, но пока намеренно этого не делаем, чтобы точно‑точно не пропустить потеряшку.
У детектора иногда случаются проблемы с некоторыми мордочками в профиль: он их не видит, потому что таких фотографий в обучении было очень мало. В целом очевидно, что это «лечится» добавлением нужного домена в обучающий датасет, но мы пока ограничились просьбой к пользователям бота использовать фотографии, на которых животное смотрит прямо.
Поскольку выравнивать мордочки мы поленились, то при загрузке перевернутой фотографии сетка может отрабатывать некорректно. На скриншотах пример как раз такого поведения:
Если искать перевернутого котика, то находится такой же перевернутый, но другой котик (пусть даже и со схожестью ниже порога). А если котика правильно повернуть, то в базе тут же находится нужная карточка. Чтобы такого поведения избежать, мы все же обучим выравниватель мордочек бот при загрузке в базу просит пользователей использовать хорошие фотки:
P.S. Лучше всего подходят фото, где крупно и отчетливо видна мордочка животного.
А иногда животных реально очень сложно отличить, особенно часто это случается с черными котами. Ниже пример некорректной работы сетки (и мы ее понимаем).
Про планы
В целом наша модель уже показывает отличные результаты, но пространство для улучшений все еще есть. Сетку точно можно дообучить на большем количестве фотографий, чтобы повысить точность предсказаний. Еще один вариант улучшения — поиск в базе по нескольким фотографиям вместо одной.
Дополнительно мы планируем обучить несколько сеток атрибутов — для определения породы, возраста, а может быть, даже пола животных. Учитывание таких параметров в сравнении поможет реже ошибаться.
Если у вас есть идеи, пожелания или вопросы, то не стесняйтесь и смело делитесь ими!
