Как меняется ландшафт Земли сквозь призму генеративных моделей

от автора

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

Суть задачи — увидеть то, как будет выглядеть ландшафт в будущем по снимкам настоящего

Суть задачи — увидеть то, как будет выглядеть ландшафт в будущем по снимкам настоящего

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

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

В статье есть раздел про сбор данных для датасета, и описание нескольких подходов: одного совсем неудачного и нескольких более успешных. Ну и, забегая вперед, за грибами на соседнее с дачей поле я пойду, если верить обученному черному ящику, не ранее 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, состоящий из зеркального набора блоков со слоями сверток и внимания:

Архитектура диффузионной unet-модели

Архитектура диффузионной unet-модели

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

Пример работы обученной ESRGAN-модели

Пример работы обученной ESRGAN-модели

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

Пример предсказания изменения ландшафта диффузионной моделью

Пример предсказания изменения ландшафта диффузионной моделью

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

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

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

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

Архитектура генератора

Архитектура генератора

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

Архитектура RRDB-блока

Архитектура RRDB-блока

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

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

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

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

Результат предсказания изменения ландшафта с помощью ESRGAN

Результат предсказания изменения ландшафта с помощью ESRGAN

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

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

https://paperswithcode.com/sota/image-super-resolution-on-set14-4x-upscaling

https://paperswithcode.com/sota/image-super-resolution-on-set14-4x-upscaling

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

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

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

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

Архитектура SwinIR

Архитектура SwinIR

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

Архитектура Residual Swin Transformer Block

Архитектура Residual Swin Transformer Block

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

Конфигурация для обучения plain-модели
{   "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.

Результат предсказания изменения ландшафта с помощью SwinIR

Результат предсказания изменения ландшафта с помощью SwinIR

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

Конфигурация для обучения GAN-модели
{   "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 + Discriminator

Результат предсказания изменения ландшафта с помощью SwinIR + Discriminator

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

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

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

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

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

Исчезновение дороги на протяжении ряда предсказаний

Исчезновение дороги на протяжении ряда предсказаний
Исчезновение деревни на протяжении ряда предсказаний

Исчезновение деревни на протяжении ряда предсказаний

Итого

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

Результат предсказания изменения ландшафта с помощью SwinIR + Discriminator

Результат предсказания изменения ландшафта с помощью SwinIR + Discriminator

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

Примеры масштабов входных изображений

Примеры масштабов входных изображений

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

Кривые обучения при различных значениях масштаба входных изображений

Кривые обучения при различных значениях масштаба входных изображений

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

Зависимость точности предсказания от масштаба входных изображений

Зависимость точности предсказания от масштаба входных изображений

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


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


Комментарии

Добавить комментарий

Ваш адрес email не будет опубликован. Обязательные поля помечены *