Как мы 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
Импортируем основные библиотеки
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
- Colab Demo for Real-ESRGAN
- Colab Demo for Real-ESRGAN (anime videos)
А еще совсем недавно выложили portable версии
Играйтесь, улучшайте аниме и кидайте свои результаты в комментарии. Подписывайтесь на наш телеграмм канал