Меня давно интересовала тема апскейла изображений, отдельно - апскейла старых видео. Одно из первых решений, которое попалось в руки несколько лет назад - waifu2x (https://github.com/nagadomi/waifu2x). Но эта нейронка больше подходила для апскейла аниме (насколько я помню на них она и тренировалась). То есть, waifu2x подходила для довольно простых изображений без избытка деталей и сложности текстур.
Затем я поизучал ESRGAN (https://github.com/xinntao/ESRGAN) и Real-ESRGAN (https://github.com/xinntao/Real-ESRGAN). Довольно неплохие модельки, вполне годятся для апскейла изображений, но очень часто заметна синтетичность, особенно в сложных сценах, например когда на изображении есть деревья. Я даже попробовал дотренировать Real-ESRGAN, к слову это делать не сложно, на их гитхабе есть скрипты и инструкции (https://github.com/xinntao/Real-ESRGAN/blob/master/docs/Training.md), но пока дособирал свой датасет для тренировки на глаза попалась другая модель - SwinIR (https://github.com/JingyunLiang/SwinIR), потестировав которую понял - она покрывает мои текущие потребности, если не полностью, то по меньшей мере процентов на 80%. А потребности были - заапскейлить несколько старых фильмов, и чтобы после апскейла фильм смотрелся как фильм, а не как пластилиновый театр. В целом все получилось. Именно об этом эта статья.
Апскейлить будем фильм "Пираты Силиконовой долины" (1999г, США). Он повествует о появлении домашнего ПК и становлении компаний Apple и Microsoft. Довольно интересный фильм с бунтарским духом той эпохи. Главные герои - молодые Джобс, Возняк, Гейтс и другие участники "революции домашних ПК". Кстати, апскейлить фильм будем конечно же на домашнем ПК.
Пример изображения, что можно получить (рекомендуется смотреть в увеличенном):
Или в виде анимированного гифа:
Характеристики диска:
- DVD5, MPEG-2, разрешение 720x480 (NTSC)
Моя конфигурация для этой задачи:
- Gigabyte A520M, AMD Ryzen 5 PRO 3600, 32Gb DDR4 3200 MT/s (16+16)
- Gigabyte GeForce RTX 3060 12GB, CUDA Version: 12.5
- Ubuntu 22.04
Нам понадобятся:
- ПК и видеокарта с поддержкой CUDA
- Python
- Библиотека Spandrel для Python
- Одна из моделей SwinIR
- ffmpeg и ffprobe
- Порядка 80Гб дискового пространства
- Несколько дней работы ПК (если непрерывно, то ~5 дней для апскейла в 2x на конфигурации вроде моей)
Кратко суть алгоритма:
- С помощью ffmpeg распаковываем фильм по кадрам в формате PNG
- Апскейлим каждый кадр
- С помощью ffmpeg кодируем полученные изображения в HD-версию фильма с присоединением аудио-дорожек из исходного материала
Приступим.
Установка ПО
Установка ffmpeg
Под Ubuntu:
apt update
apt install ffmpeg
(при установке ffmpeg одновременно подтягивается ffprobe)
Под Windows:
https://www.ffmpeg.org/download.html
Качаем один из последних build, распаковываем архив, например в c:\ffmpeg. Из архива нам нужны две утилиты: ffmpeg и ffprobe.
Можно прописать путь к папке ffmpeg в PATH, чтобы вызывать ffmpeg из командной строки в любой директории.
Установка Spandrel
pip install spandrel
Скачивание модели
Список доступных моделей SwinIR:
https://github.com/JingyunLiang/SwinIR/releases
Я тестировал все модели из списка, наиболее понравились эти 2:
Для апскейла 2x: 003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x2_GAN.pth
Для апскейла 4x: 003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth
Качаем модель, запоминаем путь к модели.
Я делал апскейл до 2x, соответственно модель: 003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x2_GAN.pth
О требуемом времени и ресурсах будет ниже, замечу лишь, что апскейл в 4x будет существенно дольше, а качество в сравнении с 2x скорее незначительное, а в чем-то хуже.
Этап1: Распаковка видео
Склейка VOB'ов
В моем DVD экземпляре 4 основных vob'а:
VTS_01_1.VOB
VTS_01_2.VOB
VTS_01_3.VOB
VTS_01_4.VOB
Склеим их.
Команда для Linux:cat VTS_01_1.VOB VTS_01_2.VOB VTS_01_3.VOB VTS_01_4.VOB > video.vob
Команда для Windows:copy /b VTS_01_1.VOB+VTS_01_2.VOB+VTS_01_3.VOB+VTS_01_4.VOB video.vob
Дальше будем работать с объединенным video.vob.
Сразу можно выгрузить аудио-дорожки. Посмотрим какие вообще есть:ffprobe -i video.vob
Видим 2 аудиодорожки:
Stream #0:2[0x80]: Audio: ac3, 48000 Hz, stereo, fltp, 192 kb/s, Start-Time 0.281s
Stream #0:3[0x81]: Audio: ac3, 48000 Hz, stereo, fltp, 192 kb/s, Start-Time 0.281s
Здесь первая дорожка английская, вторая русская, хоть это и не обозначено явно. Выгружаем их:
ffmpeg -i video.vob -map 0:a:0 -c copy audio_track_eng.ac3 -map 0:a:1 -c copy audio_track_rus.ac3
Параметр "-c copy" указывает выгружать аудио как есть, без пережатия.
Теперь можно приступить к распаковке фреймов.
Раскадровка
Ремарка:
Здесь и далее все команды приводятся под Linux. Под Windows они обычно идентичные, если не задано другого варианта явно, за исключением указания путей к директориям. В Linux путь к директории идет через прямой слеш "/", в Windows через обратный "\".
В идеальном сценарии мы могли бы выполнить простую команду:ffmpeg -i video.vob "video_in_png/file_%06d.png"
После чего в папке "video_in_png" появились бы все фреймы в исходном количестве и в исходном разрешении. Но... Из двух DVD, которые я успел обработать, в обоих случаях были "особенности". Данный диск оказался самым сложным из них.
Во-первых, реальная частота кадров. Я смотрел данные по ней несколькими утилитами и разными способами, вот краткий список FPS, которые определялись для фильма:
60000/1001 = 59,94...
29970/1000 = 29,97
119/4 = 29,75
24000/1001 = 23,98...
Мне даже пришлось смотреть VOB файлы в hex-редакторе и по определенным сигнатурам расшифровывать значения байтов (большое спасибо ChatGPT).
Там я узнал "точно", что:
фреймрейт 29.97 fps - не правда (спойлер: либо диск скомпонован криво, либо я просто не до конца разбираюсь в структуре DVD)
разрешение 720x480 - не правда, фактическое разрешение оказалось 640x480
UPD:
В комментариях поправили, возможно отображаемое разрешение диска было 720x540 (исходное 720x480, где пиксели растягиваются по вертикали до 540). Сложно понять, как на самом деле, я совсем не доверяю служебной информации этого диска. В любом случае растянутые пиксели исходной картинки нам были не нужны, ведь мы и так собрались их "растягивать" через апскейл, поэтому 640x480 вероятно оптимальный выбор в любом случае.
Кстати, утилита ffprobe тоже выводит эти данные:Stream #0:1[0x1e0]: Video: mpeg2video (Main), yuv420p(tv, smpte170m, progressive), 720x480 [SAR 32:27 DAR 16:9], 29.75 fps
Еще и уверяет, что видео формата 16:9, тогда как оно 4:3.
В общем методом перебора я пришел к тому, что реальная частота кадров:
24000/1001, то есть, ~23,98 fps. Не хочу сильно задерживаться на этом, это история размером с отдельную статью, и надеюсь данный случай был редким исключением, поэтому нет смысла его описывать.
Отдельно - история с разрешением. Все утилиты, а также биты в исходных VOB'ах, говорят, что разрешение видео стандартное - 720x480, но если выгружать фреймы в этом разрешении, картинка становится растянута по горизонтали. Как уже писал выше, фактическое разрешение фреймов оказалось 640x480. Это я вывел эмпирически, и позже подтвердил расчетами (на всякий случай). Здесь тоже не хотелось бы задерживаться, надеюсь этот диск просто исключение, видеоряд формата 640x480 записали как стандартное для DVD разрешение 720x480, и вероятно просто не корректно внесли разметку в структуру данных.
Если бы вместо:
720x480 [SAR 32:27 DAR 16:9]
было указано:
720x480 [SAR 8:9 DAR 4:3]
Тогда бы все сходилось. SAR (Sample Aspect Ratios) если простыми словами, указывает, как растянуть или сузить видео. На всякий случай формула расчета:
Ширина фактическая = Ширина исходная (в файле) x (SAR_width / SAR_height) = 720 x (8/9) = 640, высота не меняется, итого 640x480.
В общем, пора заканчивать с этими дебрями, это был самый муторный этап, продолжаем.
Итоговая команда для выгрузки фреймов:ffmpeg -i video.vob -r 24000/1001 -vf "scale=640:480:flags=lanczos" "/home/user/frames_orig/file_%06d.png"
Здесь:
"-i video.vob" - исходный объединенный файл
"-r 24000/1001" - частота кадров равная ~23,98 кадров/сек
"-vf "scale=640:480:flags=lanczos"" - фильтр выходного потока: меняем разрешение на 640x480 и добавляем метод интерполяции Ланцоша, последний хорошо сглаживает неровности на изображении после любой его деформации, но сохраняет четкость изображения
"/home/user/frames_orig" - директория куда будут выгружаться фреймы; нужно предварительно создать
"file_%06d.png" - формат PNG, маска файлов %06d - 6-значный счетчик начиная с 000000, файлы будут вида "file_000000.png", "file_000001.png" и тд.
Вариант этой же команды, но с использованием GPU (CUDA), работает обычно быстрее:ffmpeg -hwaccel cuda -i video.vob -r 24000/1001 -vf "scale=640:480:flags=lanczos" "/home/user/frames_orig/file_%06d.png"
Еще раз напомню, если бы не было проблем с конкретным диском, команда была бы проще и очевидней:ffmpeg -i video.vob -vf "scale=640:480:flags=lanczos" "/home/user/frames_orig/file_%06d.png"
То есть просто выгружаем все фреймы, но тут пришлось явно задавать входящий фреймрейт (который сначала еще нужно было точно определить).
Если бы мы выгружали фреймы этой командой, тогда каждый 3-5 фрейм был бы дублем предыдущего. Всего бы выгрузилось 174000 фреймов, тогда как фактически их 139202, или 139235 по другой оценке, но не будем про это, и так достаточно :)
Кстати, насчет другого DVD который я апскейлил.
Параметры видеоряда там были следующие:Stream #0:0: Video: mpeg2video (Main), yuv420p(tv, top first), 720x576 [SAR 64:45 DAR 16:9], 25 fps, 25 tbr, 1k tbn
Они соответствовали действительности. Исходя из исходного разрешения 720x576 и SAR 64:45, выходило, что кадр будет выведен на экран в разрешении 1024x576 16:9 (720 x (64/45) = 1024). Я хотел увеличить разрешение видео 2 раза, до FullHD, соответственно разрешение исходного кадра нужно было 960x540 (умножаем на 2 получается 1920x1080). Здесь главное, что пропорции сохранены, 1024x576 и 960x540 оба 16:9. Соответственно выгружал фреймы командой:ffmpeg -i video_in.mkv -vf scale=960:540:flags=lanczos "/home/user/frames_orig/file_%06d.png"
Этап2: Апскейл фреймов
Теперь можно приступить к апскейлу фреймов. Это процесс длительный, в зависимости от количества фреймов и их разрешения занимает до нескольких суток. Я обычно делаю итерациями, пачками по 10000 фреймов. Для этого можно, например, разбить общую папку с фреймами на несколько по 10тыс файлов в каждой, либо доработать скрипт ниже, чтобы он работал по диапазону файлов (главное не запутаться в их порядке, тогда можно словить рассинхронизацию).
Пример команды в Linux для разбивки общей папки на несколько по 10тыс файлов:i=1; for file in all_frames/*; do mkdir -p "frames_
file" "frames_$((i/10000+1))"; ((i++)); done
Команда для Windows (PowerShell):i=1; Get-ChildItem all_frames | ForEach-Object { $d=([math]::Floor(($i-1)/10000)+1)"; if (!(Test-Path $d)) {New-Item -ItemType Directory -Path $d | Out-Null}; Move-Item $_.FullName $d; $i++ }
Не буду останавливаться здесь на долго, как именно разбить процесс на части - дело вкуса.
Основной скрипт:
Код:
import os
import torch
from PIL import Image
import torchvision.transforms as transforms
from spandrel import ImageModelDescriptor, ModelLoader
# ПАРАМЕТРЫ
images_path = "/home/user/frames_in"
output_dir = "/home/user/frames_upscaled"
os.makedirs(output_dir, exist_ok=True) # Создаем папку если не найдена
model_path = "/home/user/models/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x2_GAN.pth"
batch_size = 2 # Размер батча
batch_images = [] # Очередь изображений в батче
output_format = "JPG" # PNG, JPG
# Список изображений
all_files = sorted([f for f in os.listdir(images_path) if os.path.isfile(os.path.join(images_path, f))])
all_files_count = len(all_files)
# Очищать torch.cuda.empty_cache(), True или False, иногда помогает вместить изображения в батч
clean_cache = 0
# Загрузка модели
model = ModelLoader().load_from_file(model_path)
assert isinstance(model, ImageModelDescriptor)
model.cuda().eval()
def save_image(image_name, output_tensor):
''' Преобразование изображения из тензора и сохранение в зависимости от заданного формата '''
output_image = transforms.ToPILImage()(output_tensor.cpu().clamp(0, 1))
output_path = os.path.join(output_dir, f"{image_name}.{output_format.lower()}")
fmt = output_format.upper()
if fmt == "PNG":
output_image.save(output_path, format="PNG")
elif fmt == "JPG":
output_image.save(output_path, format="JPEG", quality=100)
else:
raise ValueError(f"Ошибка формата: {output_format!r}")
# СТАРТ ОБРАБОТКИ
for idx, image_in in enumerate(all_files):
image_path = os.path.join(images_path, image_in)
image_name = os.path.splitext(image_in)[0]
image = Image.open(image_path).convert("RGB")
input_tensor = transforms.ToTensor()(image).unsqueeze(0).cuda()
batch_images.append({'image_name': image_name, 'input_tensor': input_tensor})
# Если накопился полный батч или это последнее изображение - обрабатываем батч
if len(batch_images) == batch_size or idx == all_files_count - 1:
try:
# Очищаем память Cuda если clean_cache = True (иногда помогает вместить изображения в батч)
if clean_cache: torch.cuda.empty_cache()
# Проверяем доступную память перед объединением изображений в батч
required_memory = sum(item['input_tensor'].element_size() * item['input_tensor'].nelement() for item in batch_images)
free_memory = torch.cuda.memory_reserved(0) - torch.cuda.memory_allocated(0)
if free_memory < required_memory:
raise RuntimeError("CUDA out of memory")
batch_tensor = torch.cat([item['input_tensor'] for item in batch_images], dim=0)
# Отправляем модели батч с изображениями
with torch.no_grad():
output_tensor = model(batch_tensor)
# Сохраняем обработанные изображения
for i, item in enumerate(batch_images):
image_name = item['image_name']
save_image(image_name, output_tensor[i])
except RuntimeError as e:
if "CUDA out of memory" in str(e):
print("Ошибка наличия свободной памяти, обрабатываем изображения по одному...")
for item in batch_images:
image_name = item['image_name']
single_tensor = item['input_tensor']
with torch.no_grad():
output_tensor = model(single_tensor)
save_image(image_name, output_tensor[0])
else:
raise
except Exception as e: # Другая ошибка
print(f"Ошибка: {e}")
finally: # Очищаем батч
batch_images.clear()
print("ГОТОВО.")
# Выгружаем модель
del model
torch.cuda.empty_cache()Пояснение по некоторым параметрам.
batch_size = 2
Размер батча который мы передаем модели, сколько изображений обрабатывать одновременно. Я обычно использую 2, т.к., во-первых - это существенно ускоряет обработку в сравнении с 1 изображением, во-вторых - оптимальное расходование памяти, в третьих - дальнейшее увеличение батча, в моем случае, почти не дает прироста скорости, зато потребление памяти возрастает существенно.
К слову, для фреймов в разрешении 640x480 на моей RTX 3600 12Gb batch_size максимум 3, больше не хватит памяти; для 960x540 максимум 2 фрейма можно отправить видеокарте единовременно.
А вот замер обработки фреймов в разрешении 320x240 с разными батчами, апскейл 2x:
Открыть
1 batch:
Время выполнения: 1 минута 40 секунд
Максимум памяти использовалось: 849.29 MB
Максимум памяти включая резерв: 982.00 MB
2 batches:
Время выполнения: 1 минута 14 секунд
Максимум памяти использовалось: 1612.66 MB
Максимум памяти включая резерв: 1840.00 MB
3 batches:
Время выполнения: 1 минута 13 секунд
Максимум памяти использовалось: 2365.52 MB
Максимум памяти включая резерв: 2712.00 MB
4 batches:
Время выполнения: 1 минута 12 секунд
Максимум памяти использовалось: 3122.71 MB
Максимум памяти включая резерв: 3570.00 MB
6 batches:
Время выполнения: 1 минута 12 секунд
Максимум памяти использовалось: 4640.16 MB
Максимум памяти включая резерв: 5328.00 MB
8 batches:
Время выполнения: 1 минута 12 секунд
Максимум памяти использовалось: 6153.86 MB
Максимум памяти включая резерв: 7046.00 MB
10 batches:
Время выполнения: 1 минута 12 секунд
Максимум памяти использовалось: 7674.44 MB
Максимум памяти включая резерв: 8778.00 MB
12 batches:
Время выполнения: 1 минута 12 секунд
Максимум памяти использовалось: 9183.08 MB
Максимум памяти включая резерв: 10538.00 MB
Как видно прирост в скорости заметен лишь в сравнении 2 batches и 1 batch. Возможно на других конфигурациях и видеокартах увеличение батча скажется существеннее.
output_format = "JPG" # PNG, JPG
В каком формате сохранить увеличенное изображение. Итоговые файлы я обычно сохраняю в JPG, т.к. увеличенные фреймы в PNG будут занимать слишком много места, а вот исходные фреймы выгружаю через ffmpeg в PNG.
clean_cache = 0 # True или False
Нужно ли очищать кеш CUDA перед проверкой доступной GPU памяти. Честно признаться - это костыль, чтобы впихнуть невпихуемое вместить фреймы в батчи в отдельных случаях. Только с помощью этого костыля мне удалось вместить 2 фрейма в батч в случае обработки другого диска, где разрешение входящих фреймов было 960x540.
По другим параметрам должно быть все очевидно.
Запускаем скрипт и уходим на несколько дней ждать обработки.
Примеры изображений ДО и ПОСЛЕ
(рекомендуется смотреть в увеличенном)
Или к примеру анимированные гифки в большем разрешении:
Сравнение в гифах:
Этап3: Финал, кодирование видео
После того как мы запскейлили все фреймы, остается закодировать итоговое видео.
Команда:ffmpeg -r 24000/1001 -i "/home/user/frames_upscaled/file_%06d.jpg" -i "/home/user/audio_track_rus.ac3" -i "/home/user/audio_track_eng.ac3" -c:v hevc_nvenc -b:v 10M -minrate 5M -maxrate 15M -bufsize 30M -preset p7 -map 0:v -map 1:a -map 2:a -metadata:s:a:0 title="Russian" -metadata:s:a:0 language=rus -metadata:s:a:1 title="English" -metadata:s:a:1 language=eng -c:a copy -pix_fmt yuv420p video_hd.mkv
Здесь:
"-r 24000/1001" - частота кадров равная ~23,98 кадров/сек
"-i /home/user/frames_upscaled/file_%06d.jpg" - директория с увеличенными фреймами
"-i "/home/user/audio_track_rus.ac3" -i "/home/user/audio_track_eng.ac3"" - подключаем аудиодорожки
"-c:v hevc_nvenc" - кодек
"-b:v 10M -minrate 5M -maxrate 15M" - переменный битрейт, среднее значение 10Мбит/сек, минимальное 5Мбит/сек, максимальное 15Мбит/сек
"-bufsize 30M" - размер буфера для переменного битрейта, рекомендуется использовать 2x от maxrate (2x15M=30M), либо можно не указывать, оставить на усмотрение ffmpeg
"-preset p7" - пресет 7 для кодека hevc_nvenc, высокое качество
"-map 0:v -map 1:a -map 2:a" - порядок потоков которые подключаем в итоговый файл: видео из фреймов, аудио1, аудио2 - которые указали до этого
"-metadata:s:a:0 title="Russian" -metadata:s:a:0 language=rus -metadata:s:a:1 title="English" -metadata:s:a:1 language=eng" - прописываем мета-информацию о языках аудиодорожек
"-c:a copy" - копировать аудио без пережатия
"-pix_fmt yuv420p" - цветовой формат пикселей, для обычных видео обычно используется yuv420p
"video_hd.mkv" - имя выходного видео в контейнере MKV
Ждем компиляции и все готово.
Заключение
Мы заапскейлили фильм с помощью модели SwinIR. В принципе можно попробовать любую другую модель, библиотека Spandrel позволяет работать со многими моделями. Есть сайт https://openmodeldb.info/ где сотни моделей на разных архитектурах, в основном там тюны базовых моделей.
Внимание:
Если у вас PyTorch ниже версии 2.6, тогда настоятельно рекомендуется запускать модели .pth/.bin от неизвестных авторов с флагом weights_only=True. Это связано с тем, что в бинарные файлы моделей может быть встроен вредоносный код, который может произвольно выполниться при десериализации модели (при загрузке). Если мы задаем флаг weights_only=True, тогда PyTorch загружает только веса модели.
Начиная с PyTorch 2.6 флаг по умолчанию стал weights_only=True, если значение параметра не передается явно.
Итого, из исходных 640x480 мы получили видео 1280x960 (4:3). Это не стандартное разрешение формата HD (1280x720 16:9) или FullHD (1920x1080 16:9), но и исходный материал оказался не стандартным.
Ниже будут ссылки на примеры кадров ДО/ПОСЛЕ, а также фрагмент видео ДО/ПОСЛЕ. В целом качество вполне приличное, как-будто действительно полноценное HD.
Какие замечены недостатки? Не очень реалистично получаются трава и деревья, часто заметно, что нейронка их перерисовала включив фантазию. Также иногда можно заметить синтетичность в отдельных деталях или объектах, но не часто и если всматриваться. Вообще чем более размытое исходное изображение или отдельные объекты на нем, тем хуже они заапскейлятся, будет заметна синтетичность.
В целом, нормальные обычные DVD апскейлятся хорошо, там где изначальная детализация хорошая. Плохо апскейлятся сильно сжатые видео, где картинка замыленная и мало деталей. Но, еще раз, если исходный материл не был сильно сжат, тогда апскейл скорее всего отработает хорошо.
Дополнительные материалы
Примеры кадров ДО и ПОСЛЕ можно посмотреть и скачать на моем гугл-диске. Там же есть примеры изображений увеличенных в 4 раза. Некоторые примеры опубликую у себя в Тг-канале.
Фрагмент видео: оригинал и улучшенный.
После того как из DVD видео сделали HD, его можно переделать в 3D. Об этом я писал в другой статье:
Как сделать 3D версию любого фильма на примере StarWars4 (DepthAnythingV2 + Parallax)
Спасибо за внимание.