4х повышение разрешения изображения с использованием ESRGAN

В данной статье хочу разобрать применение предобученной нейронной сети ESRGAN для увеличения разрешения изображения в четыре раза c использованием tensorflow hub.

Инструменты

Буду пользоваться языком программирования python 3.10, среда разработки - jupyter notebook.

Проведем импорт необходимых библиотек. Для работы с нейронными сетями воспользуемся tensorflow. Для работы с уже обученной моделью ESRGAN будем пользовать tensorflow hub. Также нам потребуется библиотека для работы с изображениями, возьмем PILLOW.

from PIL import Image
import tensorflow as tf
import tensorflow_hub as hub

Предподготовка изображения

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

• Картинка должна быть с плавающей запятой, преобразованной с использованием tf.cast(image, tf.float32).

• 4-х мерные входные параметры, [batch_size, height, width, 3]. Чтобы применить сверхвысокое разрешение к одному изображению, воспользуемся tf.expand_dims(image, 0) чтобы добавить пакетное измерение.

• Для отображения изображения, необходимо преобразовать его обратно в uint8с использованием tf.cast(tf.clip_by_value(image[index_of_image_to_display], 0, 255), tf.uint8)

Входное изображение. Размер 640х480 пикселей. (изображено ультразвуковое исследование сердца )

# определим путь к входному изображению
image_path = './result_images/frame6.jpg'
преобразуем входное изображения согласно условиям
hr_image = tf.image.decode_image(tf.io.read_file(image_path))
image = tf.expand_dims(hr_image, 0)
image = tf.cast(image, tf.float32)

Работа с ESRGAN

Модель располагается в хранилище tensorflow hub.

# проведем загрузку модели
model = hub.load("https://tfhub.dev/captain-pool/esrgan-tf2/1")
выполним 4x увеличения изображения
super_resolution = model(image)

Постобработка изображения

Преобразуем 4х картинку обратно в uint8 и сохраним в папку.

image = np.asarray(super_resolution)
image = tf.clip_by_value(image, 0, 255)
image = Image.fromarray(tf.cast(image[0], tf.uint8).numpy())
image.save('./SR.jpg'% count)

Выходное изображение. Разрешение 2560 * 1920 пикселей.

Заключение

В данной статье рассмотрено как с использованием предобученной нейронной сети ESRGAN можно увеличить изображение в четыре раза с использованием минимального количества кода.