Модель восстановления лиц GPEN

1. Введение

GPEN — это модель восстановления лиц без разметки. Автор встраивает декодер StyleGAN V2, предложенный предыдущей моделью, в качестве декодера GPEN; реконструирует простой энкодер с DNN для предоставления входных данных для декодера. Таким образом, сохраняя превосходную производительность декодера StyleGAN V2, функция модели меняется с преобразования стиля изображения на восстановление лиц без разметки. Общая структура модели показана на следующем рисунке:

Для более подробного ознакомления с моделью и обратитесь к репозиторию, вы можете просмотреть следующий проект AI Studio [ссылка]([GPEN Blind Face Repair Model Reproduction - Paddle AI Studio (baidu.com)](https://aistudio.baidu.com/ The latest version of aistudio/projectdetail/3936241?contributionType=1)).

2. Подготовка к работе

2.1 Подготовка набора данных

Обучающий набор модели GPEN представляет собой классический набор данных FFHQ с лицами, всего 70 000 изображений лиц высокого разрешения 1024 x 1024, а тестовый набор представляет собой набор данных CELEBA-HQ, всего 2 000 изображений лиц высокого разрешения. Для получения подробной информации см. URL набора данных: FFHQ, CELEBA-HQ. Конкретные ссылки для скачивания приведены ниже:

Адрес загрузки исходного набора данных:

• FFHQ: https://drive.google.com/drive/folders/1tZUcXDBeOibC6jcMCtgRRz67pzrAHeHL?usp=drive_open

• CELEBA-HQ: https://drive.google.com/drive/folders/0B4qLcYyJmiz0TXY1NG02bzZVRGs?resourcekey=0-arAVTUfW9KRhN-irJchVKQ&usp=sharing

Поскольку исходный набор данных FFHQ слишком велик, вы также можете загрузить набор данных FFHQ 256 разрешений по следующей ссылке:

https://paddlegan.bj.bcebos.com/datasets/images256x256.tar

После загрузки организация файлов следующая:

|-- data/GPEN
    |-- ffhq/images256x256/
        |-- 00000
            |-- 00000.png
            |-- 00001.png
            |-- ......
            |-- 00999.png
        |-- 01000
            |-- ......
        |-- ......
        |-- 69000
            |-- ......
                |-- 69999.png
    |-- test
        |-- 2000张png图片

Пожалуйста, измените параметры dataroot наборов данных train и test в файле конфигурации configs/gpen_256_ffhq.yaml на путь к вашему обучающему и тестовому набору.

2.2 Подготовка модели

Файл параметров модели и адрес загрузки журнала обучения:

link: https://paddlegan.bj.bcebos.com/models/gpen.zip

Загрузите параметры модели и тестовые изображения по ссылке и поместите их в папку data/ в корневом каталоге проекта. Конкретная структура файлов следующая:

data/gpen/weights
    |-- model_ir_se50.pdparams
    |-- weight_pretrain.pdparams  
data/gpen/lite_data

3. Начало использования

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

Введите следующий код в консоли, чтобы начать обучение:

python tools/main.py -c configs/gpen_256_ffhq.yaml

Модель поддерживает только обучение на одной карте.

Обучение модели требует использования paddle2.3 и выше, и дождитесь реализации paddle функций, связанных с операторами второго порядка elementwise_pow. Версия paddle2.2.2 может нормально работать, но модель не может быть успешно обучена, потому что некоторые функции потерь будут вычислять неправильный градиент. Если во время обучения сообщается об ошибке, обучение в настоящее время не поддерживается. Вы можете пропустить часть обучения и напрямую использовать предоставленные параметры модели для тестирования. Оценка модели и тестирование могут использовать paddle2.2.2 и выше.

3.2 Оценка модели

При оценке модели введите следующий код в консоль, используя загруженные параметры модели, упомянутые выше:

python tools/main.py -c configs/gpen_256_ffhq.yaml -o dataset.test.amount=2000 --load data/gpen/weights/weight_pretrain.pdparams --evaluate-only
``` **3.3 Модельное предсказание**

#### 3.3.1 Экспорт весов генератора

После обучения необходимо использовать программу tools/extract_weight.py для извлечения весов генератора из обученной модели (включая генератор и дискриминатор) для логического вывода в applications/tools/gpen.py, чтобы реализовать различные приложения модели GPEN. Введите следующую команду, чтобы извлечь веса генератора:

```bash
python tools/extract_weight.py data/gpen/weights/weight_pretrain.pdparams --net-name g_ema --output data/gpen/weights/g_ema.pdparams

3.3.2 Обработка одного изображения

После извлечения весов генератора введите следующую команду для тестирования изображений по пути --test_img. Изменение параметра --seed может генерировать разные деградированные изображения для демонстрации более разнообразных эффектов. Вы можете изменить путь после --test_img на любое изображение, которое хотите протестировать. Если после параметра --weight_path не указаны веса, автоматически загрузятся веса обученной модели для тестирования.

python applications/tools/gpen.py --test_img data/gpen/lite_data/15006.png --seed=100 --weight_path data/gpen/weights/g_ema.pdparams --model_type gpen-ffhq-256

Ниже приведены примеры изображений и соответствующие восстановленные изображения, слева направо: деградированное изображение, сгенерированное изображение и исходное чёткое изображение:

Пример результата:

result saved in : output_dir/gpen_predict.png
        FID: 92.11730631094356
        PSNR:19.014782083825743

4. Tipc

4.1 Экспортируйте модель логического вывода

python tools/export_model.py -c configs/gpen_256_ffhq.yaml --inputs_size=1,3,256,256 --load data/gpen/weights/weight_pretrain.pdparams

Приведённая выше команда создаст файлы структуры модели gpenmodel_g_ema.pdmodel и файлы весов модели gpenmodel_g_ema.pdiparams и gpenmodel_g_ema.pdiparams.info, необходимые для прогнозирования, которые будут сохранены в каталоге inference_model/. Вы также можете изменить параметры после --load на файл параметров модели, который вы хотите протестировать.

4.2 Логический вывод с помощью механизма прогнозирования

python tools/inference.py --model_type GPEN --seed 100 -c configs/gpen_256_ffhq.yaml -o dataset.test.dataroot="./data/gpen/lite_data/" --output_path test_tipc/output/ --model_path inference_model/gpenmodel_g_ema

В конце логического вывода восстановленное изображение, созданное моделью, будет сохранено в test_tipc/output/GPEN по умолчанию, а значение FID, полученное в результате теста, будет выведено в test_tipc/output/GPEN/metric.txt.

Результат по умолчанию:

Metric fid: 187.0158

Примечание: поскольку операция ухудшения качества изображений высокой чёткости имеет определённую степень случайности, результаты каждого теста будут разными. Чтобы обеспечить согласованность результатов тестов, здесь я зафиксировал случайное начальное число, так что при каждом тесте выполняется одна и та же операция ухудшения качества изображения.

4.3 Вызов сценария для завершения обучения и выполнения теста в два этапа

Чтобы вызвать режим lite_train_lite_infer тестовой базы обучения функции прогнозирования, запустите:

# Исправленный формат файла sh
sed -i 's/\r//' test_tipc/prepare.sh
sed -i 's/\r//' test_tipc/test_train_inference_python.sh
sed -i 's/\r//' test_tipc/common_func.sh
# подготовка данных
bash test_tipc/prepare.sh ./test_tipc/configs/GPEN/train_infer_python.txt 'lite_train_lite_infer'
# запуск теста
bash test_tipc/test_train_inference_python.sh ./test_tipc/configs/GPEN/train_infer_python.txt 'lite_train_lite_infer'

5. Ссылки

@misc{2021GAN,
      title={GAN Prior Embedded Network for Blind Face Restoration in the Wild},
      author={ Yang, T.  and  Ren, P.  and  Xie, X.  and  Zhang, L. },
      year={2021},
      archivePrefix={CVPR},
      primaryClass={cs.CV}
}