GAN для преобразования Google -> Yandex

Модуль реализует pix2pix-подобный GAN для парных изображений карт. Генератор получает изображение Google из пары и пытается сгенерировать изображение в стиле Yandex. Сгенерированная картинка сравнивается со вторым изображением пары, то есть с оригинальным *_yandex.png.

Структура

models/GAN/
├── build.py        # генератор notebook.gen.ipynb по схеме
├── _schema.py      # структура генерируемого ноутбука
├── _schema.md      # описание формата схемы
├── notebook.gen.ipynb
├── src/
│   ├── config.py
│   ├── dataloader.py
│   ├── model.py
│   ├── trainer.py
│   ├── analyze.py
│   └── main.py
└── README.md

Архитектура

GeneratorUNet принимает google_img с 3 RGB-каналами и возвращает fake_yandex.

DiscriminatorPatchGAN получает пару (google_img, yandex_img) и отличает настоящую пару от (google_img, fake_yandex).

ImageGAN.generator_step() считает:

adversarial loss: дискриминатор должен принять (google_img, fake_yandex) за реальную пару;
L1 loss: fake_yandex сравнивается с оригинальным yandex_img из той же пары.
SSIM loss: штрафует структурные отличия, чтобы карта не расплывалась;
Sobel edge loss: сохраняет контуры дорог и объектов для последующего поиска ключевых точек.

Итоговая функция потерь генератора:

G_loss =
  lambda_GAN * GAN_loss(D(google_img, fake_yandex), real)
  + lambda_L1 * L1(fake_yandex, yandex_img)
  + lambda_SSIM * SSIMLoss(fake_yandex, yandex_img)
  + lambda_edge * SobelEdgeLoss(fake_yandex, yandex_img)

По умолчанию используется lsgan, усиленная реконструкция и более медленный дискриминатор. Это менее “креативно”, зато обычно даёт более чистые контуры.

Запуск

Из папки models/GAN:

python src/main.py

src/main.py специально написан без функции main(): при генерации ноутбука все переменные остаются доступны после запуска ячейки, как в SiaN.

Чтобы пересобрать ноутбук из схемы:

python build.py

По умолчанию датасет берется из:

C:\Users\admin\Projects\autopilot\datasets\ya_go_maps\images

Путь, размер изображений, batch size и число эпох меняются в src/config.py.

Если CUDA видна, но установленный PyTorch не поддерживает архитектуру GPU например Tesla P100 sm_60, код автоматически переключится на CPU. Это управляется параметром prefer_cuda в src/config.py.

Ожидаемый формат данных

В директории датасета должны лежать пары:

0000_google.png
0000_yandex.png
0001_google.png
0001_yandex.png
...

Даталоадер возвращает:

google_img: вход генератора;
yandex_img: целевое изображение для сравнения;
idx: номер пары.

Чекпоинты

Тренер сохраняет чекпоинты в:

models/GAN/runs/checkpoints

Сохраняются best.pth, периодические epoch_N.pth и final.pth.

После обучения src/main.py также строит:

графики G/D/L1/SSIM/edge loss;
сетку Google input -> Generated Yandex -> Yandex target.

Файлы сохраняются в models/GAN/runs/images.

4.3 KiB Raw Blame History Unescape Escape