[recovery mode] Раскрашивание изображений с использованием нейронных сетей11.08.2022 09:32

DeOldify — это проект, основанный на глубоком обучении, для раскрашивания и восстановления изображений. Модель использует архитектуру NoGAN для обучения модели.

Мы будем использовать эту модель, чтобы преобразовать некоторые старые черно-белые фотографии, добавив к ним цвет.

Модель Deoldify

Deoldify использует генеративно-состязательную нейронную сеть (GAN).

Помимо использования GAN в этой модели также используется метод, известный как No-GAN.

Большинство GAN состоят из двух частей — Генератор и Дискриминатор.

Генератор — это часть, которая создает изображение. Дискриминатор пытается выделить настоящие цветные изображения из фейковых перекрашенных изображений. Техника No-GAN работает путем обучения моделей генератора и дискриминатора, присутствующих в GAN, изолированно.

Это похоже на то, как вы обучаете обычную нейронную сеть, но отличается от GAN, поскольку они обычно обучаются параллельно. Затем они настраиваются вместе, как правило, так, как вы тренируете GAN.

Модель работает, беря черно-белое изображение и передавая его модели Deoldify. Затем модель выводит цветное изображение.

Клонирование репозитория GitHub

Мы собираемся использовать репозиторий GitHub, содержащий реальную модель.

 git clone https://github.com/jantic/DeOldify.git DeOldify 

Приведенный выше код клонирует репозиторий DeOldify в папку.

Установка необходимых зависимостей

Чтобы использовать модель, нам нужно установить «пару» зависимостей.

 pip install -r colab_requirements.txt 

Эти зависимости включают в себя:

• fastai==1.0.51

• wandb

• tensorboardX==1.6

• ffmpeg-python

• youtube-dl>=2019.4.17

• jupyterlab

• pillow>=8.0.0

Загрузка модели

Далее нам нужно загрузить предварительно обученную модель.

 mkdir 'models' 
 wget https://data.deepai.org/deoldify/ColorizeArtistic_gen.pth -O ./models/ColorizeArtistic_gen.pth 

Мы создали новую папку с именем models внутри основной папки. С использованием wget мы загружаем предварительно обученную модель в эту только что созданную папку.

Примеры входных изображений

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

Работа с моделью

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

Проведем импорт необходимых библиотек.

from PIL import Image

from deoldify import device
from deoldify.device_id import DeviceId
#choices:  CPU, GPU0...GPU7
device.set(device=DeviceId.GPU0)

import torch

if not torch.cuda.is_available():
    print('GPU not available.')
    
import fastai
from deoldify.visualize import *

torch.backends.cudnn.benchmark = True 

# если версия python >= 3.10 необходимо отдельно переназначить модуль Sized 
import collections
collections.Sized = collections.abc.Sized
# работаю в соседней папке, импортирую папку DeOldify
import sys
sys.path.append('./DeOldify')

Укажем путь к входному файлу

 image_path = '6.jpg' 

Далее создадим переменную colorizer для хранения нашей модели

 colorizer = get_image_colorizer(artistic=True) 

С использованием plot_transformed_imageметод, передадим входное изображения, на выходе колоризованная картинка.

colorizer.plot_transformed_image(image_path, render_factor=35, display_render_factor=True, figsize=(8,8))

Значение по умолчанию 35 для render_factorхорошо работает в большинстве случаев. render_factorопределяет разрешение, при котором отображается цветная часть изображения. Нижний render_factorидеально подходит для изображений с более низким разрешением, в то время как более высокий render_factorдля изображений с высоким разрешением.

Примеры выходных цветных изображений

Мы видим, что модель Deoldify добавила цвета нашим изображениям. Мы достигли этих результатов, написав всего несколько строк кода.

Заключение

Модель Deoldify позволяет раскрашивать черно-белые изображения. Модель с открытым исходным кодом и доступна через GitHub. Вы можете легко поэкспериментировать со своими черно-белыми изображениями и добавить к ним цвет.

© Habrahabr.ru