diff --git a/autopilot.py b/autopilot.py index 8ff8752..8efc298 100644 --- a/autopilot.py +++ b/autopilot.py @@ -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) diff --git a/dissertation/chapter_2/2.5_training/2.5_training.md b/dissertation/chapter_2/2.5_training/2.5_training.md index 7c6a076..3406e68 100644 --- a/dissertation/chapter_2/2.5_training/2.5_training.md +++ b/dissertation/chapter_2/2.5_training/2.5_training.md @@ -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 в стиль Яндекс.Карт diff --git a/dissertation/chapter_3_11_conclusions.md b/dissertation/chapter_3_11_conclusions.md new file mode 100644 index 0000000..1f21bbf --- /dev/null +++ b/dissertation/chapter_3_11_conclusions.md @@ -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». diff --git a/dissertation/intro_conclusion.md b/dissertation/intro_conclusion.md new file mode 100644 index 0000000..6942788 --- /dev/null +++ b/dissertation/intro_conclusion.md @@ -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-модуль подтвердил перспективность доменной адаптации для случаев, когда эталонные и текущие изображения получены из разных источников или имеют заметные визуальные различия. + +Таким образом, поставленная цель работы достигнута: разработан и протестирован алгоритм навигации БПЛА для возврата в точку старта на основе визуального сопоставления изображений. Полученные результаты подтверждают, что коррекция по эталонным ориентирам и оценка гомографии позволяют ограничивать рост ошибки и поддерживать работоспособность навигационного решения без опоры на спутниковый сигнал. Наиболее перспективным направлением дальнейшего развития является оптимизация нейросетевых компонентов, расширение набора экспериментальных сценариев, проверка устойчивости к сезонным и погодным изменениям, а также перенос разработанного алгоритма из симуляционной среды на реальную бортовую платформу. diff --git a/dissertation/Диссертация.docx b/dissertation/Диссертация.docx new file mode 100644 index 0000000..f25584b Binary files /dev/null and b/dissertation/Диссертация.docx differ diff --git a/gan.py b/gan.py new file mode 100644 index 0000000..201051b --- /dev/null +++ b/gan.py @@ -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() diff --git a/main.py b/main.py index 807491c..2c4a8e1 100644 --- a/main.py +++ b/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)