alexeysrus / imagereconstructionbyfrequencies

Применение вейвлет-преобразований для построения интерпретируемой шумоподавляющей сети

Данный репозиторий содержит код для обучения модифицированной архитектуры нейронной сети IS-Net. На изображении ниже приведена схема архитектуры сети.

Подготовка обучающего набора данных

Для запуска обучения необходимо представить его в следующем формате:

PATH_TO_TRAIN_DATASET_FOLDER/
    ---> noisy/
        ---> train_image_1.png
        ---> train_image_2.png
        ...
        ---> train_image_N.png
    ---> clear/
        ---> train_image_1.png
        ---> train_image_2.png
        ...
        ---> train_image_N.png

PATH_TO_VALIDATION_DATASET_FOLDER/
    ---> noisy/
        ---> val_image_1.png
        ---> val_image_2.png
        ...
        ---> val_image_N.png
    ---> clear/
        ---> val_image_1.png
        ---> val_image_2.png
        ...
        ---> val_image_N.png

Где в директориях clear/ должны находится изображания без шума, а в директориях noisy/ - изображения с шумом и имеющие соответвующее имя, как у чистого файла изображения.

По данной ссылке можно скачать набор данных SIDD, преобразованный в используемый формат.

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

Запуск визуализации

В текущем пайплайне используется визуализация через visom, данный пакет устанавливается через пакетный менеджер pip:

pip3 install visdom

Запуск visdom рекомендуется делать фоновым процессом, с помощью screen или tmux. Пример запуска visom:

visdom -port=9001

Запуск обучения

Для запуска обучения необходимо перейти в соответствующую директорию с кодом и установить необходимые зависимости:

cd research_pipelines/denoising_with_wavelets/
pip3 install -r requirements.txt

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

python3 pytorch_wavelet_train.py \
    --train_data_folder $PATH_TO_TRAIN_DATASET_FOLDER \
    --validation_data_folder $PATH_TO_VALIDATION_DATASET_FOLDER \
    --epochs 150 \
    --lr_milestones 5 \
    --image_size 512 \
    --batch_size 4 \
    --visdom 9001 \
    --njobs 4 \
    --exp $PATH_TO_EXPERIMENT_FOLDER/ \
    --preload_datasets

Запуск валидации

Скрипт запускается из директории research_pipelines/denoising_with_wavelets/. Валидация запускается следующим скриптом:

usage: evaluate_model.py [-h] -m MODEL -f FOLDER [-o OUTPUT]

Plot wavelets

options:
  -h, --help            show this help message and exit
  -m MODEL, --model MODEL
                        Path to model checkpoint file
  -f FOLDER, --folder FOLDER
                        Path to folder with images
  -o OUTPUT, --output OUTPUT
                        Path to folder with output visualizations (optional)

Предобученную модель можно скачать по следующей ссылке.

Визулизация коэффициентов вейвлет-преобразований

Для отрисовки матриц с коэффициентами ДВП подготовлен следующий скрипт: research_pipelines/denoising_with_wavelets/scripts/plot_wavelets_pyramid.py

usage: plot_wavelets_pyramid.py [-h] -m MODEL [-i IMAGE] [-o OUTPUT] [-g GAMMA_CORRECTION]

Plot wavelets

options:
  -h, --help            show this help message and exit
  -m MODEL, --model MODEL
                        Path to model checkpoint file
  -i IMAGE, --image IMAGE
                        Path to input image (optional)
  -o OUTPUT, --output OUTPUT
                        Path to folder with output visualizations (optional)
  -g GAMMA_CORRECTION, --gamma_correction GAMMA_CORRECTION
                        Coefficient to adjust gamma of wavelets coefficients visualization (optional)

Замечание: перед запуском python скриптов может понадобиться добавить следующую переменную окружения:

export PYTHONPATH=./