Система генерации изображений людей по тексту и фото

Описание проекта

Данная система реализует генерацию новых изображений человека на основе исходного изображения и текстового описания изменений. Система включает:

• Семантическую сегментацию (U-Net)
• Текстовый энкодер (BiGRU для русского языка)
• Генератор на основе модифицированного StyleGAN2 с SPADE нормализацией
• Многомасштабный дискриминатор
• Метрики качества (FID, LPIPS, mIoU)

Структура проекта

kyrsProject/
├── src/                    # Исходный код
│   ├── data_prep.py       # Подготовка датасета CelebAMask-HQ
│   ├── unet_segmentation.py # Модель сегментации U-Net
│   ├── text_encoder.py    # Текстовый энкодер BiGRU
│   ├── stylegan2_generator.py # Генератор на основе StyleGAN2
│   ├── discriminator.py   # Многомасштабный дискриминатор
│   ├── metrics.py         # Метрики качества (FID, LPIPS, mIoU)
│   ├── pipeline.py        # Основной пайплайн генерации
│   ├── train_all.py       # Скрипт обучения всех моделей
│   └── demo.py            # Демо скрипт
├── data/                  # Датасет CelebAMask-HQ
│   ├── img_align_celeba/  # Изображения лиц
│   └── CelebAMask-HQ-mask-anno/  # Аннотации масок
├── demo_images/           # Тестовые изображения
├── weights/               # Сохраненные веса моделей
├── requirements.txt       # Зависимости Python
└── README.md              # Этот файл

Установка зависимостей

pip install -r requirements.txt

Подготовка данных

Датасет CelebAMask-HQ уже размещен в папке data/. Для подготовки данных:

cd src
python data_prep.py

Обучение моделей

Полное обучение системы

cd src
python train_all.py

Это выполнит поэтапное обучение:

1. Обучение модели сегментации (U-Net)
2. Обучение генератора и дискриминатора (GAN)
3. Совместная тонкая настройка всех компонентов

Модели сохраняются в папку weights/.

Запуск демо

С командной строки

cd src
python demo.py --input ../demo_images/artur_standart.jpg --text "сделать волосы рыжими" --output result.jpg

Использование пайплайна в коде

from src.pipeline import ImageGenerationPipeline

pipeline = ImageGenerationPipeline()
result = pipeline.generate('input.jpg', 'добавить очки')
result.save('output.jpg')

Что делает система

Вход:

• Исходное изображение человека (из demo_images/ или датасета)
• Текстовое описание изменений на русском (например: "сделать волосы длиннее", "осветлить кожу")

Выход:

• Сгенерированное изображение с изменениями в result/generated.jpg
• Графики сравнения методов в result/comparison_graphs.png

Процесс генерации:

1. Сегментация: U-Net разбивает входное изображение на 19 семантических классов
2. Кодирование текста: BiGRU преобразует русскоязычное описание в векторное представление
3. Генерация: Модифицированный StyleGAN2 с SPADE нормализацией создает новое изображение на основе маски и текста

Результаты обучения

• Модели сохраняются в папку weights/
• Метрики качества: FID, LPIPS, mIoU
• Достигнутые результаты: FID = 26.8, mIoU = 0.78

Как работать с системой

1. Установите зависимости: pip install -r requirements.txt
2. Подготовьте данные: python src/data_prep.py
3. Обучите модели: python src/train_all.py (требует GPU)
4. Тестируйте: python src/demo.py --input image.jpg --text "описание"

Примеры текстовых команд

• "сделать волосы рыжими"
• "добавить очки"
• "сделать улыбку"

Метрики качества

• FID (Fréchet Inception Distance): Меньше - лучше (близость распределений реальных и сгенерированных изображений)
• LPIPS (Learned Perceptual Image Patch Similarity): Меньше - лучше (перцептивное сходство)
• mIoU (mean Intersection over Union): Больше - лучше (точность сегментации)