January 1, 2022

Как мы Evangelion улучшали

К нам в МАИ (8 институт) привезли новый сервер с A100x8 и 2TB RAM. И первым делом мы решили его потестить, под конец нового года реальными задачами не очень хотелось заниматься, поэтому решили нагрузить его задачей Enhancement.

Enhancement - это улучшений изображения/видео/звука

Для улучшения взяли аниме Neon Genesis Evangelion (1995-1996). Нашли DVD образ, перенесли его в цифровой формате .mkv. Оригинальный размер 704 x 480, кадры аниме сильно зашумлены, нет четкости линий (как в новых аниме), иногда проступают белые точки/линии от процесса съемки кадров.

В сети есть версии в 1080p, но это просто растянутые 480p.

Результаты вышли очень крутые, видео в формате 4К (2560x1080). Хотя и не без минусов, белые линии которые появляются при процессе съемки аниме некуда не исчезли.

Совместно с 8 институтом мы обработали 26 серий + фильм End Of Evangelion. Серии будут выходить постепенно:
Даты выхода серий:
1 Января 1-6 серия
2 Января 6-12 серия
3 Января 12-18 серия
4 Января 18-24 серия
5 Января 24-26 серия
6 Января The End Of Evangelion

Подписывайтесь на наш телеграмм канал

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

Real-ESRGAN

Real-ESRGAN (Real Enhanced Super-Resolution Generative Adversarial Network)[Paper, GitHub] мы уже упоминали про нее у себя в Телеграмм канале

Немного про GAN

Задача GAN сети, генерировать ненастоящие данные которые выглядят очень реалистично. Типовая архитектура GAN сети, содержит две сети, генератор и дискриминатора. Задача генератора создать не настоящие данные, а дискриминатора отличить, где настоящие, а где не настоящие. Тем самым помогаю генератору создавать более реалистичные данные.

В чем отличие от ESRGAN (про него подробнее тут)? В том что обучение происходит на чистых синтетических данных и в качестве дискриминатора используется U-net со спектральной нормализацией.

Синтетические данные

При создании обучающего набора, используются всевозможные варианты размытия, методы изменения размера, шума и методы сжатия JPEG.

После того как собрали данные, сеть ERSGAN приступает обучаться на разные масштаба увеличения изображения (x1, x2, x4)

Использование U-net

U-Net со спектральной нормализацией необходима, для стабилизации динамики обучения, так же она помогает смягчать черезмерно острые артефакты изображения появивщиеся при обучении. С помощью U-Net удалось найти баланс между улучшения локальных мест на изображении и подавлении артефкатов.

Кодес (с пояснением)

Собственно теория это круто, но как добиться результатов на своих данных?

Первым делом надо склонировать репозитрий с Real-ESRGAN.

Откроем терминал, создадим виртуальное окружение, активируем его, поставим все пакеты что требует репозиторий и поставим сам realesrgan

git clone https://github.com/xinntao/Real-ESRGAN.git
cd Real-ESRGAN
virtualenv venv 
source venv/bin/activate
pip install -r requirements.txt
pip install tqdm
python setup.py develop

Для удобной работы стоит добавить еще Jupyter Notebook и запустим его.

pip install jupyter notebook
jupyter notebook 

Создадим ipynb и начнем.

Импортируем основные библиотеки

import torch
from basicsr.archs.rrdbnet_arch import RRDBNet
from basicsr.utils.logger import AvgTimer
from tqdm import tqdm

from realesrgan import IOConsumer, PrefetchReader, RealESRGANer
from realesrgan.archs.srvgg_arch import SRVGGNetCompact 

Загрузим нужную нам обученую сеть, список можно найти тут

!wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.3.0/RealESRGANv2-animevideo-xsx4.pth -P experiments/pretrained_models

Далее указываем параметры для сети (данные параметры только для RealESRGANv2-animevideo-xsx4)

model_name = "RealESRGANv2-animevideo-xsx4"
model_path = os.path.join('experiments/pretrained_models', model_name + '.pth')
model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')
netscale = 4
outscale=4
tile = 0
tile_pad = 10
pre_pad = 0
half= False
upsampler = RealESRGANer(
        scale=netscale,
        model_path=model_path,
        model=model,
        tile=tile,
        tile_pad=tile_pad,
        pre_pad=pre_pad,
        half=half)

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

Поставим библиотеку для работы с видео

!pip install scikit-video
import skvideo.io
siries = "Evangelion.(10.serij.iz.26).1995-1996.x264.DVDRip.mkv"

Видео загрузили, теперь посмотрим что с параметрами видео

metadata = skvideo.io.ffprobe(f"./NeonEvangelion/{siries}")
fps = float(metadata["video"]['@r_frame_rate'].split('/')[0])/1001
print("Частота кадров:", fps)
вот тут может быть проблема, так как не во всех видео единый формат, тут пришлось поделить на 1001 так как указанная частота 23976/1001
frame_width = int(metadata['video']["@width"])
frame_height = int(metadata['video']["@height"])
print("Разрешение видео:", frame_width, "x", frame_height)

Грузим наш видос

videodata = skvideo.io.vread(f"./NeonEvangelion/{siries}")
frame_count = videodata.shape[0]
print('Количестов кадров:', frame_count)  # float `frame_count`
print('Длительность видео (сек):', frame_count/fps)  # float `frame_count`

Посмотрим как сильно апскейлится наше видео

output, _ = upsampler.enhance(videodata[0], outscale=outscale)
height = output.shape[0]
width = output.shape[1]
print("Выходное разрешение видео:", width, "x", height)

Создадим метод для записи видео

crf = 17
writer = skvideo.io.FFmpegWriter('Magmadiver.mp4',
            inputdict={'-r': str(fps), '-s':'{}x{}'.format(width,height)},
            outputdict={'-r': str(fps), '-c:v': 'libx264', '-crf': str(crf), '-preset': 'veryslow', '-pix_fmt': 'yuv444p'}
)

И начнем по кадрово обрабатывать и записывать видео

for frame in tqdm(videodata):
    output, _ = upsampler.enhance(frame, outscale=outscale)
    writer.writeFrame(output)
    torch.cuda.synchronize()
writer.close()

Теперь если был звук в видео, надо его как то вытащить и поможет с этим ffmpeg. Сохраним его в отдельный mp3 файл

!ffmpeg -i "./NeonEvangelion/Evangelion.(10.serij.iz.26).1995-1996.x264.DVDRip.mkv" "Magmadiver.mp3"

И потом возьмем созданное видео и добавим в него звук.

!ffmpeg -i "Magmadiver.mp4" -i "Magmadiver.mp3" -c copy -map 0:v:0 -map 1:a:0 "10.Magmadiver.mp4"

Все спустя 1 час 30 минут очередная серия Evangelion готова.

Colab

Ну или вам впадлу разбираться с кодом и вы хотите просто посмотреть на качество улучшения всегда можете воспользоваться colab

А еще совсем недавно выложили portable версии

Играйтесь, улучшайте аниме и кидайте свои результаты в комментарии. Подписывайтесь на наш телеграмм канал