112 lines
4.3 KiB
Markdown
112 lines
4.3 KiB
Markdown
# GAN для преобразования Google -> Yandex
|
||
|
||
Модуль реализует pix2pix-подобный GAN для парных изображений карт. Генератор получает изображение Google из пары и пытается сгенерировать изображение в стиле Yandex. Сгенерированная картинка сравнивается со вторым изображением пары, то есть с оригинальным `*_yandex.png`.
|
||
|
||
## Структура
|
||
|
||
```text
|
||
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: сохраняет контуры дорог и объектов для последующего поиска ключевых точек.
|
||
|
||
Итоговая функция потерь генератора:
|
||
|
||
```text
|
||
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`:
|
||
|
||
```bash
|
||
python src/main.py
|
||
```
|
||
|
||
`src/main.py` специально написан без функции `main()`: при генерации ноутбука
|
||
все переменные остаются доступны после запуска ячейки, как в SiaN.
|
||
|
||
Чтобы пересобрать ноутбук из схемы:
|
||
|
||
```bash
|
||
python build.py
|
||
```
|
||
|
||
По умолчанию датасет берется из:
|
||
|
||
```text
|
||
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`.
|
||
|
||
## Ожидаемый формат данных
|
||
|
||
В директории датасета должны лежать пары:
|
||
|
||
```text
|
||
0000_google.png
|
||
0000_yandex.png
|
||
0001_google.png
|
||
0001_yandex.png
|
||
...
|
||
```
|
||
|
||
Даталоадер возвращает:
|
||
|
||
- `google_img`: вход генератора;
|
||
- `yandex_img`: целевое изображение для сравнения;
|
||
- `idx`: номер пары.
|
||
|
||
## Чекпоинты
|
||
|
||
Тренер сохраняет чекпоинты в:
|
||
|
||
```text
|
||
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`.
|