
Привет, Хабр!
Это последняя из трех статей, в которых я (автор канала Зайцем по ХаХатонам) рассказываю о задачах Всеросийского чемпионата Цифрового Прорыва, объясняю базовые решения (baseline) и даю советы, которые помогут подняться выше по рейтингу. В данной статье будет рассмотрен кейс от МФТИ по привязке аэроснимков к местности.
Данная статья является особенной, так как она содержит исправленный бейзлайн, который изначально не работал. Сейчас же приведенное ниже решение дает результат на 9 место в лидерборде!
Спойлер: в конце статьи есть советы для улучшения базового решения.
Цифровой Прорыв
Думаю, все и так знают, что такое Цифровой Прорыв. Однако, напомню, что в этом году основной тематикой стал искусственный интеллект. И сезон этого года в самом разгаре!
Хоть часть мероприятий уже прошла, впереди участников ждет ещё 19 региональных чемпионатов, 5 окружных хакатонов и 3 всероссийских чемпионата. Советую присоединиться ко мне и другим участникам, чтобы не упустить возможность выиграть денежные призы и крутые путешествия, а также набраться опыта на самых разных задачах.
Введение
В современном мире огромное количество задач решается с помощью спутниковых фотографий и аэрофотоснимков. Зачастую от скорости и качества интерпретации этих данных зависит то, как быстро выявляются пожары, наводнения и другие чрезвычайные ситуации. Технологии машинного зрения только начинают применяться в решении такого рода задач, однако потребность в их использовании постоянно растет.
Решение данной задачи позволит оперативно привязывать изображения к географическим координатам, что в дальнейшем может ускорить геодезические работы, поможет оперативно искать пропавших людей, контролировать вырубку лесов. И это только краткий список того, где требуется привязка аэрофотоснимков к местности.
Участникам чемпионата будет предложено найти местоположение и ориентацию снимка на крайне большом изображении по высоте и ширине с географической привязкой к местности.
Условие задачи
Цель задачи — необходимо найти местоположение и ориентацию снимка на подложке.
Для лучшего понимания контекста задачи участникам стоит ознакомиться со следующими терминами:
Подложка — крайне большое изображение по высоте и ширине с географической привязкой к местности, т.е. координаты каждого пикселя известны или их можно вычислить. Как правило, на изображении размещена большая площадь земли (квадратные километры и более)
Аэрофотоснимок — изображение со спутника или беспилотного летательного аппарата, направление камеры при фотографировании смотрело вертикально вниз. Имеет существенно меньшее разрешение в сравнении с подложкой. По сути фотография, сделанная на обычный фотоаппарат. Главная особенность в том, что аэрофотоснимок сделан в отличное от подложки время, время года, или даже в совершенно другой год или на разной высоте.
Перекрытие — положение изображений, при котором одна и та же площадь местности видна на двух и более аэрофотоснимках. Взаимное ориентирование разновременных снимков разного разрешения подразумевает под собой сопоставление снимков и получение их географической привязки за счет ручного сопоставления оператором с картой.
Данные
В качестве данных выступают аэрофотоснимки фиксированного размера:
train/img — папка, содержащая 800 фотографий тренировочного набора;
train/json — папка с данными в формате json со следующими значениями
left top — координата левого верхнего угла фотографии относительно подложки;
right top — координата правого верхнего угла;
left bottom — координата левого нижнего угла;
right bottom — координата правого нижнего угла;
angle — угол поворота.
test/ — папка, содержащая 400 фотографий для предсказания;
original.tiff — подложка с расширением 10496 x 10496:
На что стоит обратить внимание
Снимки сделаны в разные временные промежутки и при различных погодных условиях. Например, часть поверхности может быть скрыта за облаками. Стоит также заметить, что фотографий для обучения мало, расширить набор для обучения можно за счет самостоятельной нарезки фотографий с подложки.
Метрика
Для такой специфичной задачи разработана своя метрика, которая определяет разницу между предсказанным центром, углом поворота фотографии и их оригинальными значениями.

Подробно о решении
Методология решения
В базовом решении предложено решить задачу, как регрессии. Таргетными значениями данном случае выступают 4 координаты и угол поворота изображения относительно оригинального.
Схема решения будет следующей:
Установка и импорт всех библиотек
Предобработка данных
Создание загрузчиков (DataLoader) для подачи данных в модель
Вспомогательные функции для обучения модели
Создание и обучение модели
Тестирование полученного решения
Какие библиотеки нам нужны
Начнем с импорта всех необходимых библиотек. В качестве фреймворка для обучения нейросети выбран torch.
# Общие библиотеки import pandas as pd import numpy as np import glob from tqdm import tqdm import os from sklearn.model_selection import train_test_split import json from math import sin, cos # Для создания и обучения модели import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import Dataset from torchvision import datasets, models, transforms # Для работы с изображениями import cv2 from PIL import Image # Для визуализации import matplotlib.pyplot as plt from IPython.display import clear_output
Преобразование начального датасета
На данном этапе данные, хранящиеся в json файлах, преобразовываются в pandas-датафрейм.
json_dir = "/content/json/" data_df = pd.DataFrame({'id': [], "left_top_x": [], 'left_top_y': [], "right_bottom_x": [], 'right_bottom_y': [], 'angle': []}) json_true = [] for _, _, files in os.walk(json_dir): for x in files: if x.endswith(".json"): data = json.load(open(json_dir + x)) new_row = {'id':x.split(".")[0]+".img", 'left_top_x':data["left_top"][0], 'left_top_y':data["left_top"][1], 'right_bottom_x': data["right_bottom"][0], "right_bottom_y": data["right_bottom"][1], 'angle': data["angle"]} data_df = data_df.append(new_row, ignore_index=True) data_df.head(5)

Загрузчик данных
Для эффективного обучен��я нейросети данные нужно подавать в виде батчей. Батч хранит в себе несколько экземпляров данных. К примеру, в базовом решении параметр batch_size равен 16, то есть каждый батч, подаваемый в модель содержит в себе 16 экземпляра данных. Ниже представлен один из вариантов реализации подобного загрузчика данных.
Сначала напишем класс, в котором данные непосредственно загружаются и приводятся в нужный формат.
class ImageDataset(Dataset): def __init__(self, data_df, transform=None): self.data_df = data_df self.transform = transform def __getitem__(self, idx): # достаем имя изображения и ее лейбл image_name, labels = self.data_df.iloc[idx]['id'], [self.data_df.iloc[idx]['left_top_x']/10496, self.data_df.iloc[idx]['left_top_y']/10496, self.data_df.iloc[idx]['right_bottom_x']/10496, self.data_df.iloc[idx]['right_bottom_y']/10496, self.data_df.iloc[idx]['angle']/360] # читаем картинку. read the image image = cv2.imread(f"/content/train/{image_name}") image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image = Image.fromarray(image) # преобразуем, если нужно. transform it, if necessary if self.transform: image = self.transform(image) return image, torch.tensor(labels).float() def __len__(self): return len(self.data_df)
Далее зададим аугментации, которые будут использоваться при обучении модели.
train_transform = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) valid_transform = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ])
Посмотрим на количество данных и разделим их на тренировочную и валидационную части.
from os import listdir print("Обучающей выборки " ,len(listdir("/content/train"))) print("Тестовой выборки " ,len(listdir("/content/test"))) Обучающей выборки 800 Тестовой выборки 400
# разделим датасет на трейн и валидацию, чтобы смотреть на качество train_df, valid_df = train_test_split(data_df, test_size=0.2, random_state=43)
Каждую из выборок подадим в ранее созданный класс. После чего обернем это в другой класс, уже существующий в библиотеке torch - DataLoader.
train_dataset = ImageDataset(train_df, train_transform) valid_dataset = ImageDataset(valid_df, valid_transform)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=16, shuffle=True, pin_memory=True, num_workers=2) valid_loader = torch.utils.data.DataLoader(dataset=valid_dataset, batch_size=16, # shuffle=True, pin_memory=True, num_workers=2)
Вспомогательные функции
Для обучения модели нам понадобятся функции расчета метрики, построения графика обучения и непосредственно обучения.
Ниже приведена функция рассчета метрики. Важное замечание - метрика для лидерборда рассчитывается исправленной метрикой, что именно исправлено можно прочитать в разделе "Метрика".
def compute_metric(data_true, data_pred, outImageW = 10496, outImageH = 10496): x_center_true = np.array((data_true[0] + data_true[2])/2).astype(int) y_center_true = np.array((data_true[1] + data_true[3])/2).astype(int) x_metr = x_center_true - np.array((data_pred[0] + data_pred[2])/2).astype(int) y_metr = y_center_true - np.array((data_pred[1] + data_pred[3])/2).astype(int) metr = 1 - 0.7 * (abs(x_metr)/outImageH + abs(y_metr)/outImageW)/2 + 0.3 *abs(data_pred[4] - data_true[4])/359 return metr
Функция визуализации графиков обучения.
def plot_history(train_history, val_history, title='loss'): plt.figure() plt.title('{}'.format(title)) plt.plot(train_history, label='train', zorder=1) points = np.array(val_history) steps = list(range(0, len(train_history) + 1, int(len(train_history) / len(val_history))))[1:] plt.scatter(steps, val_history, marker='+', s=180, c='orange', label='val', zorder=2) plt.xlabel('train steps') plt.legend(loc='best') plt.grid() plt.show()
Процесс обучения модели написан вручную без готовых функций (таких как TrainEpoch). Это дает более четкий контроль процесса обучения и возможность его кастомизировать.
def train(res_model, criterion, optimizer, train_dataloader, test_dataloader, NUM_EPOCH=15): train_loss_log = [] val_loss_log = [] train_acc_log = [] val_acc_log = [] for epoch in tqdm(range(NUM_EPOCH)): model.train() train_loss = 0. train_size = 0 train_pred = [] for imgs, labels in train_dataloader: optimizer.zero_grad() imgs = imgs.cuda() labels = labels.cuda() y_pred = model(imgs) loss = criterion(y_pred, labels) loss.backward() train_loss += loss.item() train_size += y_pred.size(0) train_loss_log.append((loss.data.cpu().detach().numpy() / y_pred.size(0)) * 100) y_pred[:, :4] = y_pred[:, :4] * 10496 y_pred[:, -1] = y_pred[:, -1] * 360 labels[:, :4] = labels[:, :4] * 10496 labels[:, -1] = labels[:, -1] * 360 for label, pr in zip(labels, y_pred): train_pred.append(compute_metric(label.cpu().detach().numpy(), pr.cpu().detach().numpy())) optimizer.step() train_acc_log.append(train_pred) val_loss = 0. val_size = 0 val_pred = [] model.eval() with torch.no_grad(): for imgs, labels in test_dataloader: imgs = imgs.cuda() labels = labels.cuda() pred = model(imgs) loss = criterion(pred, labels) pred[:, :4] = pred[:, :4] * 10496 pred[:, -1] = pred[:, -1] * 360 labels[:, :4] = labels[:, :4] * 10496 labels[:, -1] = labels[:, -1] * 360 val_loss += loss.item() val_size += pred.size(0) for label, pr in zip(labels, pred): val_pred.append(compute_metric(label.cpu().detach().numpy(), pr.cpu().detach().numpy())) val_loss_log.append((val_loss / val_size)*100) val_acc_log.append(val_pred) clear_output() plot_history(train_loss_log, val_loss_log, 'loss') print('Train loss:', (train_loss / train_size)*100) print('Val loss:', (val_loss / val_size)*100) print('Train metric:', (np.mean(train_pred))) print('Val metric:', (np.mean(val_pred))) return train_loss_log, train_acc_log, val_loss_log, val_acc_log
Обучение модели
В качестве модели используем resnet50, предобученный на датасете imagenet с выходным слоем размером 5, так как мы хотим предсказывать 5 параметров. Функцией потерь же будет MSELoss, используемый в задачах регрессии.
torch.cuda.empty_cache() # Подргружаем модель model = models.resnet50(pretrained=True) model.fc = nn.Linear(2048, 5) model = model.cuda() criterion = torch.nn.MSELoss() optimizer = torch.optim.Adam(model.fc.parameters(), lr=0.001)
Запустим обучение и понаблюдаем за изменением лоссов.
train_loss_log, train_acc_log, val_loss_log, val_acc_log = train(model, criterion, optimizer, train_loader, valid_loader, 15)

Проверка модели
Посчитаем метрику на датасете для валидации. Получаем метрику 0.98. Однако, не стоит забывать о том, что обучались мы на небольшом наборе данных и итоговая метрика на лидерборде может отличаться.
total_metric = [] for imgs, labels in valid_loader: imgs = imgs.cuda() labels = labels.cpu().detach().numpy() pred = model(imgs) pred = pred.cpu().detach().numpy() pred[:, :4] = pred[:, :4] * 10496 pred[:, -1] = pred[:, -1] * 360 labels[:, :4] = labels[:, :4] * 10496 labels[:, -1] = labels[:, -1] * 360 for label, pr in zip(labels, pred): total_metric.append(compute_metric(label, pr)) total_metric = np.mean(total_metric) print('Valid metric:', total_metric) Valid metric: 0.9801261833663446
Создадим предсказания на тестовом наборе данных
Для начала нужно написать класс для загрузки тестовых данных, аналогичный тому, что написан для тренировочных.
class TestImageDataset(Dataset): def __init__(self, files, transform=None): self.files = files self.transform = transform def __getitem__(self, idx): image_name = self.files[idx] # читаем картинку. read the image image = cv2.imread(f"/content/test/{image_name}") image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image = Image.fromarray(image) # преобразуем, если нужно. transform it, if necessary if self.transform: image = self.transform(image) return image def __len__(self): return len(self.files)
Далее собираем названия всех тестовых файлов и объявляем даталоадер с размером батча 16, так как тестовых картинок 400, а 400%16==0.
test_images_dir = '/content/test/' for _, _, test_files in os.walk(test_images_dir): break test_dataset = TestImageDataset(test_files, valid_transform) test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=16, # shuffle=True, pin_memory=True, num_workers=2 )
Собираем предсказания в список.
indexes = [x.split('.')[0] for x in test_files] preds = [] for imgs in test_loader: imgs = imgs.cuda() pred = model(imgs) pred = pred.cpu().detach().numpy() pred[:, :4] = np.clip(pred[:, :4] * 10496, 0, 10496) pred[:, -1] = np.clip(pred[:, -1] * 360, 0, 360) preds.extend(list(pred))
Записываем полученные предсказания в нужный для сабмита формат. После чего можно сжать все .json-файлы в архив и загрузить его на платформу.
sub_dir = "/content/submission/" if not os.path.exists(sub_dir): os.makedirs(sub_dir) json_true = [] for indx, pred in zip(indexes, preds): pred = [int(x) for x in pred] left_top = [pred[0], pred[1]] right_top = [pred[2], pred[1]] left_bottom = [pred[0], pred[3]] right_bottom = [pred[2], pred[3]] res = { 'left_top': left_top, 'right_top': right_top, 'left_bottom': left_bottom, 'right_bottom': right_bottom, 'angle': pred[4] } with open(sub_dir+indx+'.json', 'w') as f: json.dump(res, f)
Пример того, что содержит в себе каждый .json-файл.
{ "left_top": [7000, 4000], "right_top": [6000, 4000], "left_bottom": [7000, 3000], "right_bottom": [6000, 3000], "angle": 178 }
Рекомендации по улучшению решения
Первым вариантом улучшения решения является увеличения количества эпох и дообучение модели.
Также можно попробовать изменить архитектуру модели на более сложную.
Попробовать создать ансамбль моделей и применить метод TTA (Test Time Augmentation).
Поменять размер входных изображений и изучить возможность улучшить используемые аугментации.
Расширить датасет из предоставленной подложки или из сторонних ресурсов.
Поразмышлять над иными подходами к решению задачи - не регрессии.
Итоги
Кейс весьма интересен за счет нестандартной постановки задачи. Есть простор для экспериментов над подходами. А если вам удастся подобрать эффективный подход, то еще и сможете получить денежный приз до 250 тысяч рублей!
Все интересующие вас вопросе вы можете задать в канале Зайцем по ХаХатонам.
Всем удачи на чемпионатах и хакатонах!
