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

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

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

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

Подготовка датасета

Как уже сказано выше, существует ряд геосервисов, позволяющих получить доступ к api для запроса спутниковых снимков местности в заданных координатах и определенный момент времени (пример: Sentinel Hub). Источники данных, их характер и качество при этом сильно разнятся. Доступное при этом через api разрешение составляет порядка единиц метров на пиксель, чего недостаточно для задачи.

В то же время обнаружилось, что Google Earth Pro умеет выгружать 46Мп склейки за разные годы с разрешением уже почти 2 пикселя на метр, что позволяет брать во внимание значительно меньшие детали местности. А нарезая изображения на квадраты 128 на 128, из каждой такой склейки можно получить около 2800 патчей-изображений для обучения модели, что уже совсем хорошо.

Стоит отметить, что у снимков за разные годы могут отличаться:

• источники (группировки спутников)

• настройки

• погодные условия

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

# image1 и image2 - изображения для выравнивания

image2_equalized = np.zeros_like(image2)

for i in range(3):  # Проход по каждому каналу RGB
    hist1 = cv2.calcHist([image1], [i], None, [256], [0, 256])
    hist2 = cv2.calcHist([image2], [i], None, [256], [0, 256])

    cdf1 = hist1.cumsum()
    cdf1 = (cdf1 / cdf1[-1]) * 255

    cdf2 = hist2.cumsum()
    cdf2 = (cdf2 / cdf2[-1]) * 255

    lut = np.interp(cdf2, cdf1, range(256)).astype(np.uint8)
    image2_equalized[:,:,i] = cv2.LUT(image2[:,:,i], lut)

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

• имеют слишком мало деталей

• имеют слишком много деталей

• отличаются слишком сильно

Для поиска первых двух проблем используем пороговые значения для суммы краев полученных с помощью детектора Кэнни.

factor = 1 / (PATSH_SIZE * PATSH_SIZE)
edges = kornia.filters.canny(img.unsqueeze(0).float(), kernel_size=KERNEL_SIZE)
edges_sum = edges[0].sum() * factor

Примеры изображений со слишком малым и слишком большим количеством деталей:

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

factor = 1 / (CHANNELS * PATSH_SIZE * PATSH_SIZE)
diff_sum = (torch.abs(img_1 - img_2) > DIFF_THRESHOLD).sum().item() * factor

Примеры изображений с количеством отличий выше заданного порога:

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

Первый подход: diffusers

Первая мысль использовать диффузионную unet архитектуру скорее была продиктована желанием познакомиться с этим классом моделей, но тем не менее поделюсь опытом. За основу было взято решение из библиотеки diffusers а-ля conditioned unet, где в качестве условия выступают нарезанные патчи из исходного изображения. Чтобы уместиться в память пришлось ограничиться генерируемым размером 32 на 32 пикселя. Соответственно из исходных 128-размерных изображений набиралось 48-канальный (4 * 4 * 3) тензор с условием:

Итого на вход подаем 51 канал размером 32 на 32 пикселя:

• первые три канала - белый шум

• оставшиеся 48 - нарезанное на патчи исходное изображение размера 128 на 128 пикселей

Учим модель расшумлять разницу между снимками.

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

Архитектура модели — обычный unet, состоящий из зеркального набора блоков со слоями сверток и внимания:

• Downsample-блок: diffusers на github

• Upsample-блок: diffusers на github

Полученное на выходе изображение все еще небольшого размера 32 на 32. Увеличиваем его с помощью отдельной модели класса super resolution. Для этого обучаем на том же наборе данных ESRGAN в режиме 4x.

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

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

Второй подход: ESRGAN

Используем каноническую архитектуру из оригинальной статьи: arxiv.

Генераторный модуль использует метод рекурсивной группы RRDB-блоков (Residual-in-Residual Dense Block), который представляет собой серию повторяющихся преобразований, включающих в себя остаточные соединения и плотные связи, что позволяет более эффективно передавать информацию через слои.

В отличие от оригинальной модели опускаем лишь модуль собственно самого увеличения разрешения и отыскиваем подходящее число RRDB-блоков (остановился на 16).

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

На поверку вылезти из заманчивого локального минимума — «просто отдавать исходное изображение» оказалось сложно. Пришлось подмешивать в функцию потерь компонент штрафа за слишком малые изменения относительно входного изображения.

l1_loss = l1_coeff * l1(fake, high_res)
adversarial_loss = adv_coeff * -torch.mean(disc(fake))
loss_for_vgg = vgg_coeff * vgg_loss(fake, high_res)
diff_loss = diff_coeff * -torch.abs(low_res - fake).sum()

gen_loss = l1_loss + adversarial_loss + loss_for_vgg + diff_loss

Результат получился уже повеселее, нежели диффузионка с апскейлом, но тем не менее визуально метод "угадывал" не часто.

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

Идем дальше, смотрим на SOTA в сфере super resolution:

Останавливаемся на SwinIR — решении на основе трансформеров.

Третий подход: SwinIR

Код реализации и скрипты для обучения живут здесь: github.

Это решение также предназначено для восстановления изображений, но на основе Swin Transformer. SwinIR состоит из трех частей: поверхностное извлечение признаков, глубокое извлечение признаков и финальное восстановление изображения на их основе.

В частности, модуль глубокого извлечения признаков состоит из нескольких блоков Residual Swin Transformer Block (RSTB), каждый из которых включает в себя несколько слоев трансформеров Swin с остаточной связью.

Обучение сети конфигурируется json-файликом с настройками модели и гиперпараметров обучения. Создаем свой:

{
  "task": "swinir_geo_predict"
  , "model": "plain"
  , "gpu_ids": [0,1,2,3,4,5,6,7]
  , "dist": true

  , "n_channels": 3

  , "path": {
    "root": "geo_predict"
    , "pretrained_netG": null
    , "pretrained_netE": null
  }

  , "datasets": {
    "train": {
      "name": "train_dataset"
      , "dataset_type": "sr"
      , "dataroot_H": "../train_2021"
      , "dataroot_L": "../train_2014"

      , "H_size": 128
      , "sigma": 15
      , "sigma_test": 15

      , "dataloader_shuffle": true
      , "dataloader_num_workers": 16
      , "dataloader_batch_size": 8
    }
    , "test": {
      "name": "test_dataset" 
      , "dataset_type": "sr"
      , "dataroot_H": "../test_2021"
      , "dataroot_L": "../test_2014"

      , "sigma": 15
      , "sigma_test": 15

    }
  }

  , "netG": {
    "net_type": "swinir" 
    , "upscale": 1 // не меняем разрешение, у нас другая задача
    , "in_chans": 3 
    , "img_size": 128 
    , "window_size": 8  
    , "img_range": 1.0 
    , "depths": [6, 6, 6, 6, 6] 
    , "embed_dim": 120 
    , "num_heads": [6, 6, 6, 6, 6]
    , "mlp_ratio": 2 
    , "upsampler": null
    , "resi_connection": "1conv"

    , "init_type": "default"
  }

  , "train": {
    "G_lossfn_type": "vgg"
    , "G_lossfn_weight": 1.0
    , "G_charbonnier_eps": 1e-9

    , "E_decay": 0.999

    , "G_optimizer_type": "adam"
    , "G_optimizer_lr": 2e-4
    , "G_optimizer_wd": 0 
    , "G_optimizer_clipgrad": null
    , "G_optimizer_reuse": true

    , "G_scheduler_type": "MultiStepLR"
    , "G_scheduler_milestones": [800000, 1200000, 1400000, 1500000, 1600000]
    , "G_scheduler_gamma": 0.5

    , "G_regularizer_orthstep": null
    , "G_regularizer_clipstep": null

    , "G_param_strict": true
    , "E_param_strict": true

    , "checkpoint_test": 5000
    , "checkpoint_save": 5000
    , "checkpoint_print": 200
  }
}

Параметры обучения оставляем прежними, функцию потерь также не трогаем, это стандартная сумма l1 и perceptual loss на основе vgg.

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

{
  "task": "swinir_geo_predict_gan" 
  , "model": "gan"        
  , "gpu_ids": [0]        

  , "scale": 1       
  , "n_channels": 3  
  , "path": {
    "root": "geo_predict"  
    , "pretrained_netG": null  
    , "pretrained_netD": null  
	, "pretrained_netE": null  
	}

, "datasets": {
    "train": {
      "name": "train_dataset" 
      , "dataset_type": "sr" 
      , "dataroot_H": "../train_2021"
      , "dataroot_L": "../train_2014" 
      , "H_size": 128  
      , "sigma": 15       
      , "sigma_test": 15 

      , "dataloader_shuffle": true
      , "dataloader_num_workers": 16
      , "dataloader_batch_size": 8 
    }
    , "test": {
      "name": "test_dataset" 
      , "dataset_type": "sr" 
      , "dataroot_H": "../test_2021"
      , "dataroot_L": "../test_2014" 

      , "sigma": 15    
      , "sigma_test": 15 

    }
  }


  , "netG": {
    "net_type": "swinir" 
    , "upscale": 1 
    , "in_chans": 3 
    , "img_size": 128 
    , "window_size": 8  
    , "img_range": 1.0 
    , "depths": [6, 6, 6, 6, 6] 
    , "embed_dim": 120 
    , "num_heads": [6, 6, 6, 6, 6]
    , "mlp_ratio": 2 
    , "upsampler": null        
    , "resi_connection": "1conv"

    , "init_type": "default"
  }

  , "netD": {
    "net_type": "discriminator_unet"
    , "in_nc": 3
    , "base_nc": 64
    , "n_layers": 3     
    , "norm_type": "spectral" 

    , "init_type": "orthogonal" 
    , "init_bn_type": "uniform" 
    , "init_gain": 0.2
  }

  , "train": {
    "G_lossfn_type": "l1"
    , "G_lossfn_weight": 1

    , "F_lossfn_type": "l1" 
    , "F_lossfn_weight": 1
    , "F_feature_layer": [2,7,16,25,34]  
    , "F_weights": [0.1,0.1,1.0,1.0,1.0]
    , "F_use_input_norm": true
    , "F_use_range_norm": false

    , "gan_type": "lsgan" 
    , "D_lossfn_weight": 1

    , "E_decay": 0.999 

    , "D_init_iters": 0

    , "G_optimizer_type": "adam"
    , "G_optimizer_lr": 5e-5 
    , "G_optimizer_wd": 0

    , "D_optimizer_type": "adam"
    , "D_optimizer_lr": 5e-5 
    , "D_optimizer_wd": 0

    , "G_scheduler_type": "MultiStepLR"
    , "G_scheduler_milestones": [800000, 1600000]
    , "G_scheduler_gamma": 0.5
    , "G_optimizer_reuse": true

    , "D_scheduler_type": "MultiStepLR"
    , "D_scheduler_milestones": [800000, 1600000]
    , "D_scheduler_gamma": 0.5
    , "D_optimizer_reuse": false

    , "G_param_strict": true
    , "D_param_strict": true
    , "E_param_strict": true

    , "checkpoint_test": 50000000000 
    , "checkpoint_save": 5000
    , "checkpoint_print": 200
  }
}

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

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

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

Итого

Подход с использованием SwinIR в составе GAN действительно позволяет предугадывать, как изменится ландшафт в той или иной точке с течением времени. Отобрав более интересный датасет, включающий в себя не только одно исходное изображение, а несколько на протяжении ряда наблюдений, вероятно, может позволить сильно улучшить точность предсказаний. Тем более, что к 2024 году наблюдений накопилось достаточно, особенно в окрестностях крупных городов. А пока ограничиваемся лишь единичным снимком как источником контекста. Еще несколько примеров предсказания изменений ландшафта местности спустя 7 лет (исходное изображение, предсказание, действительность):

И напоследок про вопрос — начиная с какого размера площади поверхности Земли модель умудряется находить тот самый контекст, который позволяет хоть как-то предсказать как же изменится эта самая поверхность. Для наглядности, вот примеры размеров входных изображений:

Можно обучить модель на одинаковых датасетах (и по размеру изображений и по их количеству), но на разном масштабе от 40м до 130м на изображение. Затем результат сравнить с реальными снимками с помощью perceptual loss. Логично, что чем больше площадь, тем для модели понятнее характер этой местности, она обучается быстрее и качественнее. Кривые обучения для разных масштабов следующие:

А если взглянуть на зависимость perceptual loss метрики от масштаба изображений, то можно заметить, что качество стремительно растет при увеличении масштаба и достигает насыщения в районе 80-120м.

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