Для написания бота использовалась библиотека aiogram для построения асинхронной архитектуры. Весь код в файле bot.py. В файле keyboards.py находятся используемые клавиатуры, а также длинные сообщения, чтобы не засорять основной файл.
Бот умеет делать следующие действия:
- Переносить стиль
- Для работы необходимо отправить сначала боту фотографию стиля, а затем фотографию, на которую перенести стиль.
- Есть возможность выбрать картинку стиля у бота. Предлагается на выбор 6 всемирно известных картин с краткой справкой (автор - название - год написания)
- Использовать CycleGAN
- Превратить летнюю фотографию в зимнюю
- Превратить зимнюю фотографию в летнюю
- Бот может отправить примеры, получаемые в результате работы 3 вышеназванных алгоритмов
- Перенаправить на данный репозиторий, чтобы ознакомиться с проектом.
Особенности:
- Запоминание пользователей и отправляемых ими фотографий реализовано через словарь, в котором сохраняются объекты класса для каждого пользователя, начавшего общение с ботом. В объекте класса хранится id пользователя, название выбранного алгоритма и в случае использования style transfer сохраняется фотография стиля. Фотография контента и фото для CycleGAN не сохраняются, а передаются сразу в пайплайн выбранного алгоритма.
- Бот отдает фотографию в том же разрешении, что и получил от пользователя. Так как для моделей необходимы квадратные фотографии, то используются паддинги
- StyleTranfer: сжатие до 150 пикселей (причина в разделе Heroku)
- CycleGAN: сжатие до 512 пикселей
- Алгоритмы запускаются в отдельном потоке с помощью threading.Thread
- Длительность работы:
- StyleTrasfer: в среднем 4 минуты
- CycleGAN: в среднем 30 секунд
Для переноса стилей были выполнены следующие действия:
- Взята предобученная модель VGG16
- Модель была дополнительно доучена на датасете Caltech 101, содержащим 101 категорию данных и 8677 картинок. Код обучения представлен в файле model_nst/train_vgg16.ipynb
- Чтобы сэкономить память, были сохранены веса только нужных слоев (первые 11) в файл models_wts/nst_features.pth
- Для переноса стилей использован классический алгоритм, который подразумевает итеративное изменение входной картинки на основе разницы с style и content картинками
- При формировании модели для Style Transfer не используются слои модели VGG16 напрямую, объявляются только необходимые 11 слоев через nn.Sequential с целью сэкономить память при дальнейшем использовании Heroku
Для использования CycleGAN был взят код junyanz pytorch-CycleGAN-and-pix2pix Исходный код дает возможность использовать CycleGAN и pix2pix моделей в различных режимах на различных датасетах
Для использования такого большого проекта в собственных нуждах были выполнены следующие действия:
- Взята архитектура только CycleGAN в режиме для теста, чтобы генерировать изображения только в одну сторону (папка models)
- Убраны лишние функции
- Переделан вход модели так, чтобы она принимала не датасет, а сразу картинки
- Модель сразу отдает результат генерации, а не пускает его дальше по пайплайну
- Пред и пост обработка фотографий делается собственными функциями
- Взяты параметры для запуска скрипта в тестовом режиме (папка options)
- Код переработан так, чтобы не использовать датасеты, также удалены любые функции связанные с ними.
- В папке scripts находится 1 скрипт для загрузки весов натренированного CycleGAN. Для загрузки необходимо перейти в папку с CycleGAN и запустить скрипт с выбранными весами. Я использовал summer2winter_yosemite и winter2summer_yosemite. Веса загружаются в папку checkpoints
$ cd model_cyclegan
$ bash ./scripts/download_cyclegan_model.sh summer2winter_yosemite
- Был переделан скрипт test.py таким образом, чтобы создать модель со стандартными параметрами и сохранить ее полностью в файлы pth для дальнейшего использования ботом. Сделано это для того, чтобы отпала необходимость сложного создания модели при работе бота. Модели сохраняются в папку models_wts. Скрипт test.py запускается следующим образом:
$ cd model_cyclegan
$ python test.py --name summer2winter_yosemite_pretrained --model test --no_dropout --gpu_ids -1
gpu_ids равняется -1, так как подразумевается использование модели только на CPU
Бот был задеплоен на heroku.
Для этого создан пустой init.py и Procfile, в котором указана команда, выполняемая при запуске heroku
web: python bot.py
Через Heroku были заполнены config vars, в которых хранятся токен телеграм бота и настройки для webhook:
- WEBHOOK_HOST=https://serene-beach-54092.herokuapp.com
- WEBAPP_HOST=0.0.0.0
- WEBAPP_PORT=5000
Возникшая проблема при работе с Heroku: Весь проект с установленными зависимостями, указанными в requirements.txt, занимает на сервере 351.4 Мб из 500 Мб. Создаваемая модель для style transfer при работе с разрешением картинок уже как минимум 160x160 пикселей, уже не помещается на сервере, несмотря на то, что все слои VGG16 не загружаются, а создаются только необходимые 11 слоев и веса к ним. По этой причине картинки для style transfer сжимаются до разрешения 150х150 пикселей, а затем получаемая картинка растягивается до исходного разрешения картинки, полученной от пользователя.
Чтобы проект поместился на сервере в 500 Мб и по той причине, что все вычисления происходят на CPU, torch и torchvision загружаются только с поддержкой CPU. Heroku автоматически не сможет установить данные пакеты, поэтому в файле указывается ссылка откуда их загружать.
-f https://download.pytorch.org/whl/torch_stable.html
aiogram==2.13
Pillow==8.2.0
torch==1.9.0+cpu
torchvision==0.10.0+cpu
| Описание файлов и папок | ||
|---|---|---|
| images | ||
| examples | Примеры работы модели | |
| styles | Варианты картинок для стиля | |
| model_cyclegan | ||
| checkpoints | Веса, используемые для создания СycleGAN | |
| models | Архитектура СycleGAN | |
| options | Параметры для создания СycleGAN | |
| scripts | Внутри скрипт для загрузки весов CycleGAN | |
| CycleGAN.py | Пайплайн для СycleGAN | |
| test.py | Скрипт для создания моделей с загруженными весами | |
| model_nst | ||
| StyleTransfer.py | Модель и пайплайн для StyleTransfer | |
| train_vgg16.py | Дообучение модели VGG16 | |
| models_wts | Веса для всех моделей | |
| bot.py | Код телеграм бота | |
| keyboards.py | Клавиатуры и сообщения для бота | |
| image_processing.py | Запуск пайплайнов для NST и CycleGAN | |
| requirements.txt | Необходимые пакеты | |
| Procfile | Команда при запуске Heroku | |
| init.py | Пустой файл для работы Heroku |