Тестирование нейросети ESRGAN для детализации изображений

Небольшой тест нейронной сети ESRGAN на время и качество обработки нескольких образцов изображений при определенной конфигурации железа. Установка Python и нужных библиотек.

Установка

Для работы нам понадобится среда исполнения Python, а также набор библиотек PyTorch, OpenCV и сам скрипт ESRGAN с натренированной моделью. Скачиваем с python.org дистрибутив (в моём случае Python версии 3.8.5 x64 под ОС Windows) и производим установку. Далее запускаем (Win+R) консоль Сmd. Надо добраться до папки менеджера пакетов Pip питона. В моём случае это C:\Python38\Scripts, меняем директорию на нужную:

cd C:\Python38\Scripts

Теперь можно запустить установку PyTorch. На pytorch.org есть конфигуратор командной строки для инсталляции, выбираем версию, ОС, менеджер пакетов, язык и поддержку CUDA, которой в моём случае не будет, в виду отсутствия видеокарты от компании Nvidia.

Запускаем предложенную команду в консоли, менеджер пакетов скачает и установит всё необходимое:

pip install torch==1.6.0+cpu torchvision==0.7.0+cpu -f https://download.pytorch.org/whl/torch_stable.html

 И ещё одна установка библиотеки OpenCV:

pip install numpy opencv-python

Осталась пара шагов. На гитхабе забираем проект ESRGAN (Code -> Download ZIP) и скачиваем натренированные модели сети. Проект извлекаем в удобное место. Модели сети распаковываем в папку models проекта. Папка LR предназначена для исходников изображений, результаты соответственно получим в папке results. Для исполнения задачи нужно запустить файл test.py, который я немного модифицировал.

model_path = 'models/RRDB_ESRGAN_x4.pth'  # models/RRDB_ESRGAN_x4.pth OR models/RRDB_PSNR_x4.pth
device = torch.device('cpu')  # if you want to run on CPU, change 'cuda' -> cpu

В строчке с переменной model_path указан путь до используемой модели сети, а в device включен параметр cpu, расчёты будут производиться силами центрального процессора. Ещё я добавил метки времени старта и конца выполнения скрипта с подсчётом результата в секундах. Полный листинг:

import os.path as osp
import glob
import cv2
import numpy as np
import torch
import RRDBNet_arch as arch
import time

sec_start = time.time()
print('Start: ', time.ctime(sec_start))

model_path = 'models/RRDB_ESRGAN_x4.pth'  # models/RRDB_ESRGAN_x4.pth OR models/RRDB_PSNR_x4.pth
device = torch.device('cpu')  # if you want to run on CPU, change 'cuda' -> cpu

test_img_folder = 'LR/*'

model = arch.RRDBNet(3, 3, 64, 23, gc=32)
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)

print('Model path {:s}. \nTesting...'.format(model_path))

idx = 0

for path in glob.glob(test_img_folder):
    idx += 1
    base = osp.splitext(osp.basename(path))[0]
    print(idx, base)
    # read images
    img = cv2.imread(path, cv2.IMREAD_COLOR)
    img = img * 1.0 / 255
    img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
    img_LR = img.unsqueeze(0)
    img_LR = img_LR.to(device)

    with torch.no_grad():
        output = model(img_LR).data.squeeze().float().cpu().clamp_(0, 1).numpy()
    output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
    output = (output * 255.0).round()
    cv2.imwrite('results/{:s}_rlt.png'.format(base), output)

sec_end = time.time()
print('End: ', time.ctime(sec_end))

sec_final = sec_end - sec_start
print('Total time in seconds: ', round(sec_final))

Тестирование

Для тестирования возьмём три экземпляра изображений разных разрешений: 640x360, 1280x720 и 1920x1080 пикселей. Замерим время выполнения задачи увеличения разрешения в 4 раза и сравним участки картин с стандартной бикубической интерполяцией против обработки нейронной сетью. В конфигурации для расчётов использован 16 поточный Xeon E5-2650 v2 и 32 GB оперативной памяти в четырёх канальном режиме.

Примерные выводы по времени обработки кадра из графика сделать можно. На выходе получили изображения разрешением: 2560 x 1440, 5120 x 2880 и 7680 x 4320 пикселей. Опытным путём было установлено, что для исходника более 2.6 мегапикселя 32 GB памяти становится недостаточно. Данные сбрасываются в файл подкачки на диск и всё резко тормозит на долгие часы с постоянной передачей данных между памятью и диском.

Посмотрим относительное сравнение качества участков изображений из эксперимента. Слева привычная бикубическая интерполяция, справа работа алгоритма ESRGAN с текущей натренированной моделью. Для наглядности брались экземпляры 640 x 360 пикселей.

Работа Rob Garlington - Bag End

Работа Derrick Chew - Ghost in the shell Fanart

Работа Patrick Sutton - Portable Cognition Unit

Заключение

Результат получается не идеальный, но вполне интересный. Это хороший инструмент, который, с ростом вычислительных мощностей, становится доступен многим. На рынке есть коммерческие продукты с встроенными алгоритмами нейросетей, которые более удобны в использовании, работают быстрее (используя GPU) и имеют различные настройки. В данном же случае просто знакомство и тест на CPU.