feat: gan integration
This commit is contained in:
@@ -8,6 +8,7 @@ import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
import gan
|
||||
import sian_similarity
|
||||
from timer import Timer
|
||||
|
||||
@@ -60,6 +61,7 @@ class AutoPilot(Pilot):
|
||||
reserved_pos: Position | None
|
||||
proccessing_time: float
|
||||
use_sian_similarity: bool
|
||||
use_gan: bool
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -68,6 +70,7 @@ class AutoPilot(Pilot):
|
||||
viz_manager=None,
|
||||
pixel_ratio: float = 1.,
|
||||
use_sian_similarity: bool = False,
|
||||
use_gan: bool = False,
|
||||
):
|
||||
self.prev_chunk = None
|
||||
self.pos = Position(0, 0, 1, 0, 0, 0)
|
||||
@@ -77,6 +80,7 @@ class AutoPilot(Pilot):
|
||||
self.reserved_pos = None
|
||||
self.pixel_ratio = pixel_ratio
|
||||
self.use_sian_similarity = use_sian_similarity
|
||||
self.use_gan = use_gan
|
||||
|
||||
# Пороговые значения качества сопоставления/гомографии
|
||||
self.min_inliers: int = 12
|
||||
@@ -191,6 +195,9 @@ class AutoPilot(Pilot):
|
||||
# Краевой случай: отсутствие чанков
|
||||
if landmark_chunk is None:
|
||||
return None
|
||||
|
||||
if self.use_gan:
|
||||
landmark_chunk = gan.transform_chunk(landmark_chunk)
|
||||
|
||||
landmark_timer.start()
|
||||
src_pts, dst_pts, matches, kp1, kp2 = landmark_chunk.detect_and_match_keypoints(current_chunk)
|
||||
|
||||
@@ -1,2 +1,52 @@
|
||||
2.5 Обучение моделей глубокого обучения
|
||||
|
||||
\section{2.5.4 Обучение «GAN» на датасете «YaGoMaps V2»}
|
||||
|
||||
Для проверки применимости генеративно-состязательной сети к задаче приведения картографических изображений к единому домену было выполнено обучение модели GAN на датасете `YaGoMaps V2`. Датасет содержит парные изображения одних и тех же географических участков, полученные из Google Maps и Яндекс.Карт. В рамках эксперимента изображение Google использовалось как вход генератора, а соответствующее изображение Яндекс.Карт -- как целевой домен.
|
||||
|
||||
Обучение выполнялось на устройстве `cuda`, что позволило проводить полный цикл обучения модели без перехода на CPU. После разделения данных было сформировано 316 обучающих и 79 валидационных примеров. С учетом размера пакета 32 это соответствовало 10 итерациям на эпоху для обучающей выборки и 3 итерациям для валидационной выборки. Генератор содержал 39 157 763 обучаемых параметра, дискриминатор -- 2 769 601 параметр.
|
||||
|
||||
| Параметр | Значение |
|
||||
|---|---|
|
||||
| Датасет | `YaGoMaps V2` |
|
||||
| Устройство обучения | `cuda` |
|
||||
| Обучающая выборка | 316 изображений |
|
||||
| Валидационная выборка | 79 изображений |
|
||||
| Количество эпох | 300 |
|
||||
| Размер пакета | 32 |
|
||||
| Параметры генератора | 39 157 763 |
|
||||
| Параметры дискриминатора | 2 769 601 |
|
||||
| Лучшая валидационная ошибка реконструкции | 35,3304 |
|
||||
|
||||
Таблица X - Основные параметры обучения GAN на датасете `YaGoMaps V2`
|
||||
|
||||
Функция потерь генератора включала состязательную компоненту, L1-ошибку, структурную ошибку SSIM и ошибку по границам, вычисляемую по оператору Собеля. Такая комбинация была выбрана из-за специфики картографических изображений: модель должна не только приблизить цветовую и стилевую палитру Google Maps к Яндекс.Картам, но и сохранить геометрию дорог, перекрестков, кварталов и других ориентиров. Сохранение контуров особенно важно, так как результат генератора далее может использоваться в классическом пайплайне сопоставления ключевых точек и оценки гомографии.
|
||||
|
||||
В начале обучения наблюдалось быстрое снижение ошибки генератора и реконструкционных компонент. На первой эпохе средняя ошибка генератора на обучающей выборке составила 50,3793, L1-компонента -- 32,1376, SSIM-компонента -- 13,7808, edge-компонента -- 4,2270. На валидации после первой эпохи значение ошибки реконструкции составило 43,9219. Уже к 20-й эпохе обучающая ошибка генератора снизилась до 32,6649, а ошибка реконструкции на валидации -- до 35,4286.
|
||||
|
||||
Минимальное значение валидационной ошибки реконструкции было достигнуто на раннем этапе обучения и составило 35,3304. Данный результат соответствует эпохе, на которой модель еще сохраняла баланс между реконструкционным качеством и обобщающей способностью. После этого обучающая ошибка продолжала снижаться, однако валидационная ошибка начала колебаться и постепенно увеличиваться. Например, к 100-й эпохе обучающая ошибка генератора снизилась до 25,0163, но валидационная ошибка реконструкции составила 36,1303. К 300-й эпохе обучающая ошибка генератора достигла 19,5036, тогда как валидационная ошибка реконструкции увеличилась до 38,4682.
|
||||
|
||||
Такое поведение указывает на частичное переобучение генератора под обучающую выборку. Модель постепенно улучшала восстановление обучающих пар, но это не приводило к дальнейшему улучшению качества на отложенных примерах. Дополнительным признаком насыщения состязательного процесса является снижение ошибки дискриминатора на обучающей выборке до малых значений. На 300-й эпохе ошибка дискриминатора на обучении составила 0,0016, тогда как на валидации -- 0,3348. Это означает, что дискриминатор уверенно различал обучающие примеры, а генератор продолжал оптимизироваться преимущественно под обучающие данные.
|
||||
|
||||
Наиболее показательные значения метрик приведены в таблице.
|
||||
|
||||
| Эпоха | train G | train D | val G | val D | val rec |
|
||||
|---:|---:|---:|---:|---:|---:|
|
||||
| 1 | 50,3793 | 0,4947 | 44,1364 | 0,2827 | 43,9219 |
|
||||
| 20 | 32,6649 | 0,1742 | 35,5899 | 0,2482 | 35,4286 |
|
||||
| 23 | 33,0033 | 0,0721 | 35,5553 | 0,2352 | 35,3304 |
|
||||
| 100 | 25,0163 | 0,0066 | 36,4437 | 0,3317 | 36,1303 |
|
||||
| 200 | 21,3662 | 0,0034 | 37,9867 | 0,3280 | 37,6665 |
|
||||
| 300 | 19,5036 | 0,0016 | 38,7948 | 0,3348 | 38,4682 |
|
||||
|
||||
Таблица X - Динамика основных метрик обучения GAN
|
||||
|
||||
По результатам эксперимента для дальнейшего использования целесообразно выбирать не финальное состояние модели после 300 эпох, а чекпоинт с минимальной валидационной ошибкой реконструкции. В данном запуске таким чекпоинтом является модель с `val_rec = 35,3304`. Она обеспечивает лучший компромисс между переносом визуального стиля Яндекс.Карт и сохранением структуры изображения на данных, не участвовавших в обучении.
|
||||
|
||||
Рисунок X - График изменения ошибки генератора при обучении GAN
|
||||
|
||||
Рисунок X - График изменения ошибки дискриминатора при обучении GAN
|
||||
|
||||
Рисунок X - График изменения реконструкционных компонент L1, SSIM и edge loss
|
||||
|
||||
Рисунок X - Пример преобразования изображения Google Maps в стиль Яндекс.Карт
|
||||
|
||||
21
dissertation/chapter_3_11_conclusions.md
Normal file
21
dissertation/chapter_3_11_conclusions.md
Normal file
@@ -0,0 +1,21 @@
|
||||
# Выводы по подразделу 3.11
|
||||
|
||||
По результатам таблиц 7-9 видно, что добавление моделей глубокого обучения действительно снижает ошибку позиционирования, однако приводит к заметному падению скорости обработки кадров.
|
||||
|
||||
## Вывод по ошибке
|
||||
|
||||
По средней RMSE лучший результат показала модель «SiaN-Similarity»: ошибка снизилась с 0.53489 до 0.36168, то есть примерно на 32.38% относительно базового алгоритма. По MSE улучшение еще заметнее: средняя ошибка уменьшилась с 0.28611 до 0.13081, то есть на 54.28%.
|
||||
|
||||
GAN также улучшает точность, но слабее: средняя RMSE снизилась на 14.96%, а MSE -- на 27.69%. Комбинация «GAN» и «SiaN-Similarity» дала промежуточный результат: RMSE улучшилась на 24.89%, MSE -- на 43.59%, но средняя точность оказалась хуже, чем при использовании одной «SiaN-Similarity».
|
||||
|
||||
По отдельным маршрутам картина различается: на маршруте 1 минимальную ошибку дала комбинация GAN и SiaN-Similarity, на маршруте 2 лучше всего сработал GAN, а на маршруте 3 -- SiaN-Similarity. Поэтому нейросетевые методы повышают устойчивость алгоритма, но их эффективность зависит от конкретной траектории и качества совпадения текущих кадров с эталонами.
|
||||
|
||||
## Вывод по FPS
|
||||
|
||||
Базовый алгоритм имеет наибольшую скорость: в среднем 27.24432 FPS. Использование «SiaN-Similarity» снижает скорость до 3.79121 FPS, то есть примерно в 7.19 раза. Это связано с необходимостью дополнительного нейросетевого сравнения кадров и ориентиров.
|
||||
|
||||
GAN работает быстрее, чем «SiaN-Similarity»: средняя скорость составляет 13.12060 FPS, что всего в 2.08 раза ниже базового алгоритма. Комбинация GAN и SiaN-Similarity дает худший результат по скорости -- 3.32183 FPS, то есть примерно в 8.20 раза медленнее базового варианта.
|
||||
|
||||
## Общий вывод
|
||||
|
||||
Если главным критерием является минимальная ошибка, наиболее эффективным вариантом является применение «SiaN-Similarity». Если важен баланс между точностью и скоростью, более предпочтительным выглядит GAN: он улучшает точность слабее, но сохраняет значительно более высокий FPS. Совместное использование GAN и SiaN-Similarity не является оптимальным: оно сильно снижает скорость, но не дает лучшей средней ошибки по сравнению с одной «SiaN-Similarity».
|
||||
39
dissertation/intro_conclusion.md
Normal file
39
dissertation/intro_conclusion.md
Normal file
@@ -0,0 +1,39 @@
|
||||
# Введение и заключение для диссертации
|
||||
|
||||
## Введение
|
||||
|
||||
Беспилотные летательные аппараты в последние годы стали одним из наиболее активно развивающихся направлений робототехники и интеллектуальных транспортных систем. Они применяются для мониторинга инфраструктуры, аэрофотосъемки, доставки грузов, обследования труднодоступных территорий, поисково-спасательных операций и решения специальных задач, в которых требуется быстрое получение информации о местности. Расширение области применения БПЛА приводит к росту требований к их автономности, устойчивости и способности продолжать выполнение полетного задания при ухудшении качества внешних навигационных сигналов.
|
||||
|
||||
В большинстве существующих систем навигация БПЛА основана на совместном использовании глобальных навигационных спутниковых систем, таких как GPS и ГЛОНАСС, инерциальных датчиков, барометрических измерителей и дополнительных радиоэлектронных средств. Такой подход хорошо зарекомендовал себя в штатных условиях, однако имеет существенные ограничения. Спутниковый сигнал может быть недоступен в помещениях, под навесами, в условиях плотной городской застройки, в районах со сложным рельефом, а также в зонах радиоэлектронного подавления или преднамеренного искажения навигационных данных. В подобных ситуациях инерциальная система без внешней коррекции быстро накапливает ошибку, что делает задачу точного возврата аппарата в точку старта особенно сложной.
|
||||
|
||||
Одним из перспективных способов повышения автономности БПЛА является использование оптико-электронных средств навигации, в частности бортовой видеокамеры, направленной вертикально вниз. Такая камера позволяет получать последовательность изображений подстилающей поверхности и использовать их для оценки смещения, ориентации и текущего положения аппарата. Визуальная одометрия, визуально-инерциальная одометрия, корреляционно-экстремальная навигация и алгоритмы сопоставления изображений дают возможность корректировать траекторию даже при отсутствии спутниковой связи. При этом для малых БПЛА особенно важны вычислительная эффективность, устойчивость к шумам, ограниченный объем памяти и способность работать на бортовом вычислителе в режиме, близком к реальному времени.
|
||||
|
||||
Наибольший практический интерес в данной работе представляет задача возврата БПЛА в точку старта. В типичном сценарии аппарат на начальном этапе полета движется по маршруту и формирует эталонное представление местности: сохраняет кадры, выделяет ориентиры или строит карту участка. При потере внешней навигации или при необходимости завершить полет система должна использовать накопленную визуальную информацию, текущие кадры с камеры и модель движения для корректировки курса и последовательного приближения к исходной позиции. В отличие от общего случая SLAM, где построение карты и локализация часто выполняются одновременно, задача возврата допускает более специализированную постановку: карта или набор эталонов могут быть сформированы заранее, а затем использоваться для локализации и коррекции траектории на обратном участке.
|
||||
|
||||
Сложность задачи состоит не только в накоплении ошибки движения, но и в различии между текущими наблюдениями и эталонными изображениями. Даже изображения одного и того же участка местности могут отличаться из-за масштаба, угла поворота, освещения, сезонных изменений, качества съемки, шумов, а также из-за различия источников картографических данных. В работе это рассматривается как рассогласование между доменами изображений. Например, фрагменты Google Maps и Яндекс.Карт для одной территории могут иметь разные цветовые схемы, детализацию дорог, подписи, контуры объектов и визуальное представление ориентиров. Такое рассогласование ухудшает работу классических методов поиска ключевых точек и повышает вероятность ложного сопоставления.
|
||||
|
||||
Целью данной работы является разработка и исследование алгоритма навигации БПЛА для возврата в точку старта на основе обработки изображений местности с бортовой видеокамеры, направленной вертикально вниз, с учетом ограниченных вычислительных ресурсов и возможного рассогласования между эталонными изображениями и текущими наблюдениями. Предлагаемый подход основан на сопоставлении текущих кадров с эталонными фрагментами карты или ранее сохраненными ориентирами, оценке геометрического преобразования между изображениями и последующей коррекции положения и курса аппарата.
|
||||
|
||||
Для достижения поставленной цели в работе решаются следующие задачи. Во-первых, выполняется анализ существующих методов навигации БПЛА, включая интегрированные инерциально-спутниковые системы, визуальную и визуально-инерциальную одометрию, SLAM-подходы и корреляционные методы сопоставления изображений. Во-вторых, формулируется постановка задачи возврата в точку старта как задачи коррекции траектории по визуальным наблюдениям с минимизацией ошибки конечного положения. В-третьих, разрабатывается программная модель движения БПЛА и алгоритм определения текущего положения на основе сопоставления кадров, выделения ключевых точек, вычисления дескрипторов и оценки матрицы гомографии. В-четвертых, исследуются различные методы детекции и описания признаков, включая ORB, SIFT, BRISK и AKAZE, а также оцениваются их точность и скорость. В-пятых, рассматривается возможность повышения устойчивости алгоритма с помощью моделей глубокого обучения: сиамской сети для оценки сходства изображений и генеративно-состязательной сети для приведения картографических изображений к единому визуальному домену.
|
||||
|
||||
Объектом исследования является автономная навигация беспилотного летательного аппарата в условиях отсутствия или ненадежности спутниковых навигационных сигналов. Предметом исследования являются алгоритмы визуального сопоставления изображений и коррекции траектории БПЛА при возврате в точку старта. В качестве основного источника информации рассматривается видеопоток с камеры, направленной вертикально вниз, а в качестве эталонных данных используются фрагменты местности, полученные из картографических источников или сохраненные на этапе прямого полета.
|
||||
|
||||
Методическую основу работы составляют методы компьютерного зрения, цифровой обработки изображений, геометрического моделирования, оценки матрицы гомографии, сопоставления локальных признаков и экспериментального моделирования движения. Для проверки разработанного подхода реализована система симуляции полета, позволяющая задавать маршруты, моделировать перемещение БПЛА, получать кадры подстилающей поверхности, сравнивать текущие изображения с эталонами и визуализировать динамику ошибки. Такой стенд позволяет проводить повторяемые эксперименты при различных настройках алгоритма, разных маршрутах и разных методах обработки изображений.
|
||||
|
||||
Практическая значимость работы заключается в том, что предложенный алгоритм и программная система могут быть использованы как основа для дальнейшей разработки автономного режима возврата БПЛА при недоступности GPS/ГЛОНАСС. Результаты сравнения методов сопоставления изображений позволяют выбрать компромисс между точностью и вычислительной скоростью, что особенно важно для малых летательных аппаратов с ограниченными ресурсами. Дополнительно исследование нейросетевых модулей показывает возможные направления повышения устойчивости системы при междоменном различии изображений и ошибках выбора эталонных ориентиров.
|
||||
|
||||
Структура работы соответствует поставленным задачам. В первой главе рассматриваются современные подходы к навигации БПЛА и анализируются их ограничения применительно к задаче автономного возврата. Во второй главе формулируется математическая и алгоритмическая постановка задачи, описываются методы сопоставления изображений, модели глубокого обучения и используемые наборы данных. В третьей главе представлена программная реализация системы симуляции, описаны основные компоненты алгоритма, проведены экспериментальные запуски и выполнено сравнение классических методов компьютерного зрения и нейросетевых расширений.
|
||||
|
||||
## Заключение
|
||||
|
||||
В ходе выполнения работы была рассмотрена задача автономного возврата беспилотного летательного аппарата в точку старта при отсутствии надежной спутниковой навигации. Актуальность данной задачи связана с тем, что традиционные инерциально-спутниковые системы обеспечивают высокую точность только при наличии устойчивого внешнего сигнала, тогда как в городских каньонах, помещениях, под перекрытиями или в условиях радиоэлектронного подавления БПЛА должен опираться на собственные бортовые средства восприятия. В качестве основного источника информации в работе рассматривалась камера, направленная вертикально вниз, а возврат строился на основе сопоставления текущих кадров с эталонными изображениями местности.
|
||||
|
||||
В первой части исследования был проведен аналитический обзор существующих методов навигации БПЛА. Рассмотрены интегрированные ИНС/ГНСС-системы, визуальная и визуально-инерциальная одометрия, SLAM-подходы, корреляционно-экстремальные методы и специализированные алгоритмы возврата по видеоданным. Показано, что для малых БПЛА важен не только уровень точности, но и вычислительная сложность, возможность работы в реальном времени и устойчивость к ошибкам сопоставления. На основе обзора была обоснована целесообразность подхода, в котором построенная или сохраненная на прямом участке визуальная информация используется для коррекции положения на этапе возврата.
|
||||
|
||||
В практической части был разработан алгоритм визуальной навигации, основанный на выделении ключевых точек, вычислении дескрипторов, сопоставлении текущего кадра с ближайшими эталонными ориентирами и оценке матрицы гомографии. Полученная геометрическая связь между изображениями используется для уточнения положения БПЛА и коррекции траектории. Также была реализована программная система симуляции полета, включающая модель движения, модуль геопозиционирования, обработку маршрутов, кэширование эталонов и визуализатор эксперимента. Это позволило проводить повторяемые запуски и сравнивать методы в одинаковых условиях.
|
||||
|
||||
Экспериментальная проверка была выполнена на серии маршрутов с использованием классических методов детекции и описания признаков ORB, SIFT, BRISK и AKAZE. Всего для классических методов было проведено 28 запусков. Результаты показали, что ORB обладает наибольшей скоростью обработки, в среднем около 23,26 FPS, однако уступает другим методам по точности. SIFT и BRISK обеспечивают приемлемую погрешность, но работают медленнее. Наиболее сбалансированным вариантом оказался AKAZE: он показал высокую точность при сохранении достаточно высокой скорости, в среднем около 19,54 FPS. Отдельно было показано, что использование эталонной коррекции важно не только для снижения текущей ошибки, но и для предотвращения ее накопления на длинных маршрутах.
|
||||
|
||||
Дополнительно были исследованы нейросетевые модули, направленные на повышение устойчивости алгоритма. Модель SiaN-Similarity использовалась для выбора наиболее похожего эталонного ориентира, а GAN-модель рассматривалась как средство приведения изображений из разных картографических доменов к более согласованному виду. Интеграция SiaN-Similarity улучшила точность алгоритма, однако существенно снизила скорость работы, что показывает необходимость дальнейшей оптимизации перед применением на бортовом вычислителе. GAN-модуль подтвердил перспективность доменной адаптации для случаев, когда эталонные и текущие изображения получены из разных источников или имеют заметные визуальные различия.
|
||||
|
||||
Таким образом, поставленная цель работы достигнута: разработан и протестирован алгоритм навигации БПЛА для возврата в точку старта на основе визуального сопоставления изображений. Полученные результаты подтверждают, что коррекция по эталонным ориентирам и оценка гомографии позволяют ограничивать рост ошибки и поддерживать работоспособность навигационного решения без опоры на спутниковый сигнал. Наиболее перспективным направлением дальнейшего развития является оптимизация нейросетевых компонентов, расширение набора экспериментальных сценариев, проверка устойчивости к сезонным и погодным изменениям, а также перенос разработанного алгоритма из симуляционной среды на реальную бортовую платформу.
|
||||
BIN
dissertation/Диссертация.docx
Normal file
BIN
dissertation/Диссертация.docx
Normal file
Binary file not shown.
316
gan.py
Normal file
316
gan.py
Normal file
@@ -0,0 +1,316 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from PIL import Image
|
||||
|
||||
from vision_chunk import VisionChunk
|
||||
|
||||
|
||||
ROOT_DIR = Path(__file__).resolve().parent
|
||||
MODEL_FILE = ROOT_DIR / "models" / "GAN" / "src" / "model.py"
|
||||
DEFAULT_CHECKPOINT_PATH = ROOT_DIR / "models" / "GAN" / "runs" / "checkpoints" / "best.pth"
|
||||
|
||||
IMAGE_SIZE = (256, 256)
|
||||
CHECKPOINT_ENV = "GAN_CHECKPOINT"
|
||||
|
||||
_generator: Optional[torch.nn.Module] = None
|
||||
_device: Optional[torch.device] = None
|
||||
_translated_chunks: dict[int, VisionChunk] = {}
|
||||
|
||||
|
||||
class _LegacyUNetUpBlock(nn.Module):
|
||||
"""Upsampling block used by earlier GAN checkpoints in models/GAN/runs."""
|
||||
|
||||
def __init__(self, in_channels: int, out_channels: int, dropout: float = 0.0):
|
||||
super().__init__()
|
||||
layers = [
|
||||
nn.ConvTranspose2d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=4,
|
||||
stride=2,
|
||||
padding=1,
|
||||
bias=False,
|
||||
),
|
||||
nn.BatchNorm2d(out_channels),
|
||||
nn.ReLU(inplace=True),
|
||||
]
|
||||
if dropout > 0:
|
||||
layers.append(nn.Dropout2d(dropout))
|
||||
self.model = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x: torch.Tensor, skip_input: torch.Tensor) -> torch.Tensor:
|
||||
x = self.model(x)
|
||||
if x.shape != skip_input.shape:
|
||||
diff_h = skip_input.size(2) - x.size(2)
|
||||
diff_w = skip_input.size(3) - x.size(3)
|
||||
x = F.pad(x, [diff_w // 2, diff_w - diff_w // 2, diff_h // 2, diff_h - diff_h // 2])
|
||||
return torch.cat([x, skip_input], dim=1)
|
||||
|
||||
|
||||
class _LegacyGeneratorUNet(nn.Module):
|
||||
"""Generator architecture matching old ConvTranspose2d checkpoints."""
|
||||
|
||||
def __init__(self, down_block_cls, in_channels: int = 3, out_channels: int = 3):
|
||||
super().__init__()
|
||||
self.down1 = down_block_cls(in_channels, 64, normalize=False)
|
||||
self.down2 = down_block_cls(64, 128)
|
||||
self.down3 = down_block_cls(128, 256)
|
||||
self.down4 = down_block_cls(256, 512)
|
||||
self.down5 = down_block_cls(512, 512)
|
||||
self.down6 = down_block_cls(512, 512)
|
||||
self.down7 = down_block_cls(512, 512)
|
||||
|
||||
self.bottleneck = nn.Sequential(
|
||||
nn.Conv2d(512, 512, kernel_size=4, stride=2, padding=1, bias=False),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
self.up1 = _LegacyUNetUpBlock(512, 512, dropout=0.5)
|
||||
self.up2 = _LegacyUNetUpBlock(1024, 512, dropout=0.5)
|
||||
self.up3 = _LegacyUNetUpBlock(512, 512, dropout=0.5)
|
||||
self.up4 = _LegacyUNetUpBlock(1024, 512)
|
||||
self.up5 = _LegacyUNetUpBlock(1024, 256)
|
||||
self.up6 = _LegacyUNetUpBlock(512, 128)
|
||||
self.up7 = _LegacyUNetUpBlock(256, 64)
|
||||
|
||||
self.final = nn.Sequential(
|
||||
nn.ConvTranspose2d(128, out_channels, kernel_size=4, stride=2, padding=1),
|
||||
nn.Tanh(),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
d1 = self.down1(x)
|
||||
d2 = self.down2(d1)
|
||||
d3 = self.down3(d2)
|
||||
d4 = self.down4(d3)
|
||||
d5 = self.down5(d4)
|
||||
|
||||
u = self.bottleneck(d5)
|
||||
u = self.up3(u, d5)
|
||||
u = self.up4(u, d4)
|
||||
u = self.up5(u, d3)
|
||||
u = self.up6(u, d2)
|
||||
u = self.up7(u, d1)
|
||||
return self.final(u)
|
||||
|
||||
|
||||
class _NamedTransposeUNetUpBlock(nn.Module):
|
||||
"""ConvTranspose block with parameter names used by best.pth."""
|
||||
|
||||
def __init__(self, in_channels: int, out_channels: int, dropout: float = 0.0):
|
||||
super().__init__()
|
||||
self.upconv = nn.ConvTranspose2d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=4,
|
||||
stride=2,
|
||||
padding=1,
|
||||
bias=False,
|
||||
)
|
||||
self.norm = nn.BatchNorm2d(out_channels)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.dropout = nn.Dropout2d(dropout) if dropout > 0 else None
|
||||
|
||||
def forward(self, x: torch.Tensor, skip_input: torch.Tensor) -> torch.Tensor:
|
||||
x = self.upconv(x)
|
||||
if x.shape != skip_input.shape:
|
||||
diff_h = skip_input.size(2) - x.size(2)
|
||||
diff_w = skip_input.size(3) - x.size(3)
|
||||
x = F.pad(x, [diff_w // 2, diff_w - diff_w // 2, diff_h // 2, diff_h - diff_h // 2])
|
||||
x = self.norm(x)
|
||||
x = self.relu(x)
|
||||
if self.dropout is not None:
|
||||
x = self.dropout(x)
|
||||
return torch.cat([x, skip_input], dim=1)
|
||||
|
||||
|
||||
class _NamedTransposeGeneratorUNet(nn.Module):
|
||||
"""Full U-Net architecture matching checkpoints with upN.upconv.weight."""
|
||||
|
||||
def __init__(self, down_block_cls, in_channels: int = 3, out_channels: int = 3):
|
||||
super().__init__()
|
||||
self.down1 = down_block_cls(in_channels, 64, normalize=False)
|
||||
self.down2 = down_block_cls(64, 128)
|
||||
self.down3 = down_block_cls(128, 256)
|
||||
self.down4 = down_block_cls(256, 512)
|
||||
self.down5 = down_block_cls(512, 512)
|
||||
self.down6 = down_block_cls(512, 512)
|
||||
self.down7 = down_block_cls(512, 512)
|
||||
|
||||
self.bottleneck = nn.Sequential(
|
||||
nn.Conv2d(512, 512, kernel_size=4, stride=2, padding=1, bias=False),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
self.up1 = _NamedTransposeUNetUpBlock(512, 512, dropout=0.5)
|
||||
self.up2 = _NamedTransposeUNetUpBlock(1024, 512, dropout=0.5)
|
||||
self.up3 = _NamedTransposeUNetUpBlock(1024, 512, dropout=0.5)
|
||||
self.up4 = _NamedTransposeUNetUpBlock(1024, 512)
|
||||
self.up5 = _NamedTransposeUNetUpBlock(1024, 256)
|
||||
self.up6 = _NamedTransposeUNetUpBlock(512, 128)
|
||||
self.up7 = _NamedTransposeUNetUpBlock(256, 64)
|
||||
|
||||
self.final = nn.Sequential(
|
||||
nn.ConvTranspose2d(128, out_channels, kernel_size=4, stride=2, padding=1),
|
||||
nn.Tanh(),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
d1 = self.down1(x)
|
||||
d2 = self.down2(d1)
|
||||
d3 = self.down3(d2)
|
||||
d4 = self.down4(d3)
|
||||
d5 = self.down5(d4)
|
||||
d6 = self.down6(d5)
|
||||
d7 = self.down7(d6)
|
||||
|
||||
u = self.bottleneck(d7)
|
||||
u = self.up1(u, d7)
|
||||
u = self.up2(u, d6)
|
||||
u = self.up3(u, d5)
|
||||
u = self.up4(u, d4)
|
||||
u = self.up5(u, d3)
|
||||
u = self.up6(u, d2)
|
||||
u = self.up7(u, d1)
|
||||
return self.final(u)
|
||||
|
||||
|
||||
def _load_gan_module():
|
||||
spec = importlib.util.spec_from_file_location("gan_model", MODEL_FILE)
|
||||
if spec is None or spec.loader is None:
|
||||
raise ImportError(f"Cannot load GAN model from {MODEL_FILE}")
|
||||
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def _get_checkpoint_path() -> Path:
|
||||
checkpoint_path = os.getenv(CHECKPOINT_ENV)
|
||||
if checkpoint_path:
|
||||
return Path(checkpoint_path).expanduser().resolve()
|
||||
return DEFAULT_CHECKPOINT_PATH
|
||||
|
||||
|
||||
def _get_device() -> torch.device:
|
||||
global _device
|
||||
|
||||
if _device is None:
|
||||
gan_module = _load_gan_module()
|
||||
if hasattr(gan_module, "get_compatible_device"):
|
||||
_device = gan_module.get_compatible_device(prefer_cuda=True)
|
||||
else:
|
||||
_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
return _device
|
||||
|
||||
|
||||
def _extract_generator_state_dict(checkpoint) -> dict:
|
||||
if not isinstance(checkpoint, dict):
|
||||
return checkpoint
|
||||
|
||||
if "generator" in checkpoint:
|
||||
return checkpoint["generator"]
|
||||
|
||||
if "generator_state_dict" in checkpoint:
|
||||
return checkpoint["generator_state_dict"]
|
||||
|
||||
state_dict = checkpoint.get("model_state_dict", checkpoint)
|
||||
if any(key.startswith("generator.") for key in state_dict):
|
||||
return {
|
||||
key.removeprefix("generator."): value
|
||||
for key, value in state_dict.items()
|
||||
if key.startswith("generator.")
|
||||
}
|
||||
|
||||
return state_dict
|
||||
|
||||
|
||||
def _get_generator() -> torch.nn.Module:
|
||||
global _generator
|
||||
|
||||
if _generator is not None:
|
||||
return _generator
|
||||
|
||||
checkpoint_path = _get_checkpoint_path()
|
||||
if not checkpoint_path.exists():
|
||||
raise FileNotFoundError(
|
||||
f"GAN checkpoint not found: {checkpoint_path}. "
|
||||
f"Set {CHECKPOINT_ENV} to another .pth file if needed."
|
||||
)
|
||||
|
||||
gan_module = _load_gan_module()
|
||||
device = _get_device()
|
||||
checkpoint = torch.load(checkpoint_path, map_location=device)
|
||||
state_dict = _extract_generator_state_dict(checkpoint)
|
||||
|
||||
if any(key.endswith(".upconv.weight") for key in state_dict):
|
||||
generator = _NamedTransposeGeneratorUNet(gan_module.UNetDownBlock, in_channels=3, out_channels=3).to(device)
|
||||
elif "final.0.weight" in state_dict:
|
||||
generator = _LegacyGeneratorUNet(gan_module.UNetDownBlock, in_channels=3, out_channels=3).to(device)
|
||||
else:
|
||||
generator = gan_module.GeneratorUNet(in_channels=3, out_channels=3).to(device)
|
||||
|
||||
generator.load_state_dict(state_dict)
|
||||
generator.eval()
|
||||
|
||||
_generator = generator
|
||||
return _generator
|
||||
|
||||
|
||||
def _chunk_to_tensor(chunk: VisionChunk) -> torch.Tensor:
|
||||
image = chunk.image.convert("RGB").resize(IMAGE_SIZE, Image.BILINEAR)
|
||||
array = np.asarray(image, dtype=np.float32) / 127.5 - 1.0
|
||||
tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)
|
||||
return tensor.to(_get_device())
|
||||
|
||||
|
||||
def _tensor_to_image(tensor: torch.Tensor, size: tuple[int, int]) -> Image.Image:
|
||||
array = tensor.squeeze(0).detach().cpu().permute(1, 2, 0).numpy()
|
||||
array = ((array + 1.0) * 127.5).clip(0, 255).astype(np.uint8)
|
||||
image = Image.fromarray(array, mode="RGB")
|
||||
if image.size != size:
|
||||
image = image.resize(size, Image.BILINEAR)
|
||||
return image
|
||||
|
||||
|
||||
def transform_image(image: Image.Image) -> Image.Image:
|
||||
"""Translate a Google-style reference image into the trained GAN target style."""
|
||||
generator = _get_generator()
|
||||
source = VisionChunk(image=image)
|
||||
tensor = _chunk_to_tensor(source)
|
||||
|
||||
with torch.inference_mode():
|
||||
translated = generator(tensor)
|
||||
|
||||
return _tensor_to_image(translated, image.size)
|
||||
|
||||
|
||||
def transform_chunk(chunk: VisionChunk, force: bool = False) -> VisionChunk:
|
||||
"""Return a cached GAN-transformed copy of the reference chunk."""
|
||||
if chunk is None:
|
||||
return chunk
|
||||
|
||||
cache_key = id(chunk)
|
||||
if not force and cache_key in _translated_chunks:
|
||||
return _translated_chunks[cache_key]
|
||||
|
||||
translated = VisionChunk(
|
||||
image=transform_image(chunk.image),
|
||||
feature_method=chunk.feature_method,
|
||||
)
|
||||
translated.pos = chunk.pos
|
||||
_translated_chunks[cache_key] = translated
|
||||
return translated
|
||||
|
||||
|
||||
def clear_cache() -> None:
|
||||
_translated_chunks.clear()
|
||||
17
main.py
17
main.py
@@ -120,7 +120,13 @@ def build(name: str, map_name: str, lat: float, lon: float):
|
||||
sleep(15)
|
||||
online_map.destroy()
|
||||
|
||||
def run(name: str, map_name: str, ref_min_distance: float, use_sian_similarity: bool = False):
|
||||
def run(
|
||||
name: str,
|
||||
map_name: str,
|
||||
ref_min_distance: float,
|
||||
use_sian_similarity: bool = False,
|
||||
use_gan: bool = False,
|
||||
):
|
||||
dir = Path('trajectories')
|
||||
assert dir.exists()
|
||||
dir /= name
|
||||
@@ -164,6 +170,7 @@ def run(name: str, map_name: str, ref_min_distance: float, use_sian_similarity:
|
||||
vis_manager,
|
||||
online_map.pixel_ratio,
|
||||
use_sian_similarity=use_sian_similarity,
|
||||
use_gan=use_gan,
|
||||
)
|
||||
simulator = Simulator(online_map)
|
||||
pilot.target_idx = 0
|
||||
@@ -312,6 +319,12 @@ def parse_args():
|
||||
help='Выбирать ориентир через SiaN similarity вместо ближайшего по текущей позиции'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--use-gan',
|
||||
action='store_true',
|
||||
help='Преобразовывать эталонный vision_chunk через GAN перед поиском ключевых точек'
|
||||
)
|
||||
|
||||
# Парсим аргументы
|
||||
args = parser.parse_args()
|
||||
|
||||
@@ -339,4 +352,4 @@ if __name__ == "__main__":
|
||||
build(name, ref, lat, lon)
|
||||
|
||||
if mode == 'run' or mode == 'standalone':
|
||||
run(name, sim, rmd, args.use_sian_similarity)
|
||||
run(name, sim, rmd, args.use_sian_similarity, args.use_gan)
|
||||
|
||||
Reference in New Issue
Block a user