Разбор базового решения для задачи привязки аэроснимков к местности с Цифрового Прорыва

ruDALL-E Kandinsky уже знает как привязывать аэроснимки к местности, а Вы?

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

Это последняя из трех статей, в которых я (автор канала Зайцем по ХаХатонам) рассказываю о задачах Всеросийского чемпионата Цифрового Прорыва, объясняю базовые решения (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:

На что стоит обратить внимание

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

Метрика

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

Формула расчета метрики (позже была скорректирована, так как 10 градусов или 350 градусов это ошибка в 10 градусов, а не в 350)

Подробно о решении

Методология решения

В базовом решении предложено решить задачу, как регрессии. Таргетными значениями данном случае выступают 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']]          # читаем картинку. 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 }

Итоги

Кейс весьма интересен за счет нестандартной постановки задачи. Есть простор для экспериментов над подходами. А если вам удастся подобрать эффективный подход, то еще и сможете получить денежный приз до 250 тысяч рублей!

Все интересующие вас вопросе вы можете задать в канале Зайцем по ХаХатонам.

Всем удачи на чемпионатах и хакатонах!

ссылка на оригинал статьи https://habr.com/ru/post/676080/