Привет, Хабр!

Меня зовут Андрей, и я data scientist. В этой статье расскажу о том, как я занял второе место в конкурсе «Цифровой прорыв» с решением по автоматизации привязки фотографии к географическому положению. Главный инсайт — базовые решения не всегда хороши и проверены. Спойлер — самодеятельность и использование современных подходов помогают победить :) Расскажу на своем опыте, как не сделать свое решение хуже базового, и проанализирую подходы других участников.

Информация о задаче, проблема

Я выбрал задачу от института МФТИ, посвященную аэрофотоснимкам. В геодезии ряд задач можно решить с помощью спутниковых снимков. Однако в экстренных ситуациях, связанных с поиском очагов пожара или пропавших людей, сопоставление спутникового снимка с масштабами реальной местности занимает слишком много времени и не дает достоверной информации в режиме онлайн. Для точной обработки информации и обеспечения безопасности внедряется искусственный интеллект. Я занял второе место, предложив одно из лучших решений по автоматизации привязки фотографии к географическому положению. Эта разработка призвана помочь и ускорить геодезические работы, например, при выявлении чрезвычайных ситуаций или вырубки лесов.

Для лучшего понимания задачи стоит ознакомиться с основными терминами. 

• Подложка — крайне большое изображение по высоте и ширине с географической привязкой к местности, т.е. координаты каждого пикселя известны или их можно вычислить. Как правило, на изображении размещена большая площадь земли (квадратные километры и более).

• Аэрофотоснимок — изображение со спутника или беспилотного летательного аппарата, направление камеры при фотографировании смотрело вертикально вниз. Имеет существенно меньшее разрешение в сравнении с подложкой. По сути фотография, сделанная на обычный фотоаппарат. Главная особенность в том, что аэрофотоснимок сделан в отличное от подложки время, время года, или даже в совершенно другой год или на разной высоте.

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

Цель решения — найти местоположение и ориентацию снимка на подложке. 

Исходные данные и метрика качества

В качестве данных выступают 800 аэрофотоснимков фиксированного размера 1024х1024, аннотации к ним, в которых указаны координаты снимка относительно подложки, и угол поворота и изображение самой подложки с расширением 10496x10496. Примечательно, что снимки сделаны в разные временные промежутки и при различных погодных условиях. Например, часть поверхности может быть скрыта облаками или могут присутствовать объекты, которых нет на подложке.

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

Базовое решение

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

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

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

Так я упал на 7 место.. 

Улучшение базового решения

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

• Аугментация данных с добавлением облаков с помощью библиотеки imgaug. Работает очень долго: около 0.2 сек на одну фотографию размера 1024x1024, при том, что другие аугментации отрабатывают за тысячные доли секунды.

• Test time augmentation (TTA) - агрегация предсказаний модели из исходного и аугментированных изображения тестового датасета (пример). Такой подход повышает точность предсказаний, но замедляет общее время работы в несколько раз.

• Обучение нескольких моделей на разных размерах изображения и усреднение их предсказаний.

• Передача предсказаний разных моделей в еще одну модель регрессии (stacking).

Финальное решение

Позже я решил отойти от кода базового решения и написать все самостоятельно, хоть и по похожему сценарию (по-прежнему обучал нейронную сеть, предсказывающую координаты изображения на подложке и угол поворота). Валидацию я проводил на части тренировочного сета, не применяя к нему аугментаций. Т.к. подложка и фотографии тренировочного и тестового набора были взяты из более поздних периодов, при применении аугментаций со старой подложкой для валидации локальные метрики и метрики на платформе могли бы сильно отличаться. В моем случае локальная метрика была 0.99, а на платформе 0.98. 

Аугментации

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

Примеры аугментаций:

На локальной валидациионной выборке такой подход уже дал существенный прирост в метриках, однако я по-прежнему был неудовлетворен. 

Метрика

Во всех предыдущих экспериментах для обучение нейронной сети я использовал либо MAE, либо MSE в качестве функций потерь. Теперь я решил попробовать оптимизировать метрику соревнования напрямую, и данный подход также дал прирост в качестве.

class RegressionLoss(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
 
    def forward(self, y_pred, y_true):
        # input shape (batch_size, 5)
        batch_size = y_pred.shape[0]
 
        x_center_true = (y_true[..., 0] + y_true[..., 2]) / 2
        y_center_true = (y_true[..., 1] + y_true[..., 3]) / 2
        x_center_pred = (y_pred[..., 0] + y_pred[..., 2]) / 2
        y_center_pred = (y_pred[..., 1] + y_pred[..., 3]) / 2
 
        x_metr = torch.abs(x_center_true - x_center_pred)
        y_metr = torch.abs(y_center_true - y_center_pred)
        angle_metr = torch.abs(y_true[..., 4] - y_pred[..., 4])
 
        metric = 1 - (
            0.7 * 0.5 * (x_metr + y_metr)
            + 0.3 * torch.min(angle_metr, torch.abs(angle_metr - 360))
        )
 
        return -(metric.sum() / (batch_size + 1))

Модель

После работы над данными и метриками было решено перебрать разные модели, предобученные на датасете ImageNet. Смотрел как CNN, так и различные вариации Vision transformer, в итоге лучше всего себя показали модели семейства RegNet. Основная идея заключается в том, что мы можем выбирать множество моделей из пространства проектирования, создавая распределение моделей и используя инструменты классической статистики для анализа пространства проектирования. Данный подход отличается от нейросетевого поиска архитектуры (NAS), где целью является нахождение единственной лучшей модели из всего пространства. Подробнее про данный подход можно прочитать в официальной статье авторов. Веса и реализацию моделей на PyTorch можно посмотреть по ссылке.

Дополнительные техники

Также в процесс обучения и тестирования были добавлены техники, представленные в статье torchvision. Каждая из них позволяет немного повысить качество предсказаний модели. Техника, дающая самый большой прирост в качестве заключается в том, что на этапе обучения модели используется разрешение входных фотографий 224x224, а на этапе тестирования сначала делают resize изображения до 256x256 и затем вырезают центр изображения того же размера, что и при обучении (224x224). Подробнее о том, как изменение размера входных изображений во время инференса влияет на качество  описано в статье.

Ансамбль

Выложив свое новое решение, я оказался на 2 месте с тестовым качеством около 0.985. Было радостно, однако результаты 3 и 4 места были сильно близки к моим. Хотелось немного оторваться, чтобы уменьшить шанс проиграть им на закрытой части тестовой выборки, которая открывается уже после окончания соревнования. Соревнование уже подходило к концу и не было времени внедрять новые подходы, поэтому я просто объединил свое улучшения базового решения с новыми. Добавление стекинга моделей и TTA дало прирост метрики на 0.25, сильно замедлив скорость работы алгоритма. Однако в соревновании не было условий про скорость работы, поэтому я оставил такой подход =)

Полученные результаты

В итоге за 7 дней было разработано решение на основе нейронной сети с метрикой выше 0.98 на private set и 2 местом в таблице результатов. 

Подходы других участников

Уже после соревнования я узнал, что лучше всего (почти 1.0 на private set) с данной задачей справляются алгоритмы, основанные на классическом подходе с использованием SIFT, который заключается в поиске ключевых точек и сопоставлении их векторных описаний. Однако в своем базовом варианте такой алгоритм плохо работает при сильных поворотах изображений, а в нашей задаче поворот мог достигать 360 градусов. 

Автор решения, занявшего 1 место, использовал афинно-инвариантный A-SIFT, что позволило решить данную проблему. Также в данном решении представлено улучшение, позволяющее эффективно находить ключевые точки даже при условиях высокой облачности, которая периодически встречается в тестовом датасете.  

Также мне показалось интересным и нестандартным  решение, занявшее 4 место. Оно базируется на свойстве сверточных блоков реагировать на определенные паттерны входного изображения (если паттерн входного изображения совпадает с ядром, то свертка дает максимальный результат).

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

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

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

Такое решение будет работать сильно дольше остальных, однако в нем интересно, то, что вообще отсутствует стадия обучения. 

Подходы остальных участников соревнования можно посмотреть на странице с результатами. 

Итог

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

По ссылкам можно ознакомиться с кодом и презентацией моего решения:

• https://github.com/traptrip/ai-areal-photo

• Привязка аэроснимков к местности. Андрей Попов