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

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

Схема решения будет следующей:

1. Установка и импорт всех библиотек

2. Предобработка данных

3. Создание загрузчиков (DataLoader) для подачи данных в модель

4. Вспомогательные функции для обучения модели

5. Создание и обучение модели

6. Тестирование полученного решения

Какие библиотеки нам нужны

Начнем с импорта всех необходимых библиотек. В качестве фреймворка для обучения нейросети выбран 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 тысяч рублей!

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

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