Как стать автором
Обновить

Конвертация видео из 2D в 3D через нейросети и параллакс (скрипт)

Уровень сложностиСредний
Время на прочтение13 мин
Количество просмотров561

Эта статья продолжение основной статьи:
Как сделать 3D версию любого фильма на примере StarWars4 (DepthAnythingV2 + Parallax) (https://habr.com/ru/articles/897860/)

Сначала рекомендуется ознакомиться с первой статьей, там все основные детали: суть алгоритма, необходимые библиотеки, первоначальные скрипты и описание параметров в них. Также там приведены примеры обработанных изображений и есть ссылки на готовые 3D видео (отрывок StarWars4), в том числе для VR. Эта статья продолжение, здесь приводится доработанный скрипт и комментарии к нему. Также ниже будут обозначены другие решения, которые можно использовать для конвертации видео из 2D в 3D.

По традиции будут приложены несколько изображений, в том числе анимированные 3D-гифы, примеры того, что можно получить через DepthAnythingV2 + Parallax.

Новый скрипт

Скрипт:
import os
import subprocess
from threading import Lock
from concurrent.futures import ThreadPoolExecutor
from multiprocessing import Value
import cv2
import torch
import numpy as np

from depth_anything_v2.dpt import DepthAnythingV2


# ПАРАМЕТРЫ ОБЩИЕ
# Исходный файл
video_file = "/home/user/video.mkv"
video_name = os.path.splitext(os.path.basename(video_file))[0]

# Папка для выгрузки фреймов и папка для итоговых 3D фреймов
frames_dir = os.path.join(os.path.dirname(video_file), f"{video_name}_frames")
images3d_dir = os.path.join(os.path.dirname(video_file), f"{video_name}_3d")
os.makedirs(frames_dir, exist_ok=True)
os.makedirs(images3d_dir, exist_ok=True)

frame_counter = Value('i', 0) # Счетчик для именования кадров
threads_count = Value('i', 0) # Счетчик текущих потоков, чтобы не выходить за пределы max_threads

chunk_size = 5000  # Количество файлов на один поток
max_threads = 3 # Максимальное количество потоков

# Устройство для вычислений
device = torch.device('cuda')


# ПАРАМЕТРЫ 3D
PARALLAX_SCALE = 15  # Максимальное значение параллакса в пикселях, рекомендуется от 10 до 20
PARALLAX_METHOD = 2  # 1 или 2
INPAINT_RADIUS = 2  # Рекомендуется от 2 до 5, оптимальное значение 2-3
INTERPOLATION_TYPE = cv2.INTER_LINEAR
TYPE3D = "FSBS"  # HSBS, FSBS, HOU, FOU
LEFT_RIGHT = "LEFT"  # LEFT or RIGHT

# 0 - если не нужно менять размеры полученного изображения
new_width = 1920
new_height = 1080

# Путь к папке с моделями, указывать без слеша на конце, например: "/home/user/DepthAnythingV2/models"
depth_model_dir = "/home/user/DepthAnythingV2/models"

model_depth_configs = {
        'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
        'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
        'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]}
}

encoder = 'vitl' # 'vitl', 'vitb', 'vits'

model_depth = DepthAnythingV2(**model_depth_configs[encoder])
model_depth.load_state_dict(torch.load(f'{depth_model_dir}/depth_anything_v2_{encoder}.pth', weights_only=True, map_location=device))
model_depth = model_depth.to(device).eval()
 

def image_size_correction(current_height, current_width, left_image, right_image):
    ''' Коррекция размеров изображений если заданы new_width и new_height '''
    
    # Вычисляем смещения для центрирования
    top = (new_height - current_height) // 2
    left = (new_width - current_width) // 2
    
    # Создаем черный холст нужного размера
    new_left_image = np.zeros((new_height, new_width, 3), dtype=np.uint8)
    new_right_image = np.zeros((new_height, new_width, 3), dtype=np.uint8)
    
    # Размещаем изображение на черном фоне
    new_left_image[top:top + current_height, left:left + current_width] = left_image
    new_right_image[top:top + current_height, left:left + current_width] = right_image
    
    return new_left_image, new_right_image
            
def depth_processing(image):
    ''' Создание карты глубины для изображения '''
    
    # Вычисление глубины
    with torch.no_grad():
        depth = model_depth.infer_image(image)
        
    # Нормализация глубины
    depth_normalized = (depth - depth.min()) / (depth.max() - depth.min())

    return depth_normalized

def image3d_processing_method1(image, depth, height, width):
    ''' Функция создания стереопары на основе исходного изображения и карты глубины.
        Метод1: более быстрый, контуры более сглаженные, но может быть менее точным
    '''
    
    # Вычисление значения для параллакса
    parallax = depth * PARALLAX_SCALE
    
    # Сетка координат
    y, x = np.indices((height, width), dtype=np.float32)

    # Вычисление смещений
    shift_left = np.clip(x - parallax, 0, width - 1)
    shift_right = np.clip(x + parallax, 0, width - 1)

    # Применение смещений с cv2.remap
    left_image = cv2.remap(image, shift_left, y, interpolation=INTERPOLATION_TYPE)
    right_image = cv2.remap(image, shift_right, y, interpolation=INTERPOLATION_TYPE)

    return left_image, right_image
    
def image3d_processing_method2(image, depth, height, width):
    ''' Функция создания стереопары на основе исходного изображения и карты глубины.
        Метод2: немного медленнее первого метода, но может быть точнее.
    '''
    
    # Вычисление значения для параллакса
    parallax = depth * PARALLAX_SCALE
    
    # Округление параллакса и преобразование в int32
    shift = np.round(parallax).astype(np.int32)

    # Сетка координат
    y, x = np.indices((height, width))

    # Подготовка изображений
    left_image  = np.zeros_like(image)
    right_image = np.zeros_like(image)

    # Формирование левого изображения по смещенным координатам
    x_src_left = x - shift
    valid_left = (x_src_left >= 0) & (x_src_left < width)
    left_image[y[valid_left], x[valid_left]] = image[y[valid_left], x_src_left[valid_left]]

    # Формирование правого изображения по смещенным координатам
    x_src_right = x + shift
    valid_right = (x_src_right >= 0) & (x_src_right < width)
    right_image[y[valid_right], x[valid_right]] = image[y[valid_right], x_src_right[valid_right]]
    
    # Маски пропущенных пикселей для инпейнтинга
    mask_left  = (~valid_left).astype(np.uint8) * 255
    mask_right = (~valid_right).astype(np.uint8) * 255

    # Заполнение пустот через инпейнтинг
    left_image  = cv2.inpaint(left_image,  mask_left,  INPAINT_RADIUS,  cv2.INPAINT_TELEA)
    right_image = cv2.inpaint(right_image, mask_right, INPAINT_RADIUS, cv2.INPAINT_TELEA)

    return left_image, right_image
    
def image3d_combining(left_image, right_image, height, width):   
    ''' Объединение изображений стереопары в единое 3D изображение '''
    
    # Корректировка размеров изображений, если заданы new_width и new_height
    if new_width and new_height:
        left_image, right_image = image_size_correction(height, width, left_image, right_image)
        # Меняем значения исходных размеров изображений на new_height и new_width для корректного склеивания ниже
        height = new_height
        width = new_width
        
    # Порядок изображений, сначала левое или сначала правое
    img1, img2 = (left_image, right_image) if LEFT_RIGHT == "LEFT" else (right_image, left_image)
    
    # Объединение левого и правого изображений в единое 3D изображение
    if TYPE3D == "HSBS":  # Сужение и склейка изображений по горизонтали
        combined_image = np.hstack((cv2.resize(img1, (width // 2, height), interpolation=cv2.INTER_AREA),
                          cv2.resize(img2, (width // 2, height), interpolation=cv2.INTER_AREA)))
                          
    elif TYPE3D == "HOU":  # Сужение и склейка изображений по вертикали
        combined_image = np.vstack((cv2.resize(img1, (width, height // 2), interpolation=cv2.INTER_AREA),
                          cv2.resize(img2, (width, height // 2), interpolation=cv2.INTER_AREA)))
                          
    elif TYPE3D == "FSBS":  # Склейка изображений по горизонтали
        combined_image = np.hstack((img1, img2))
    
    elif TYPE3D == "FOU":  # Склейка изображений по вертикали
        combined_image = np.vstack((img1, img2))
    
    return combined_image

def get_total_frames():
    ''' Определение точного количества фреймов в видео.
        Сначала пробуется первый вариант, он быстрее, но срабатывает редко.
        Если не сработал первый вариант, пробуется второй, он долгий, но обычно отрабатывает хорошо.
    '''
    
    cmd1 = ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=nb_frames",
            "-of", "default=nokey=1:noprint_wrappers=1", video_file]
    cmd2 = ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=nb_read_frames", "-count_frames",
            "-of", "default=nokey=1:noprint_wrappers=1", video_file]
    
    try:
        result = subprocess.check_output(cmd1).splitlines()[0].decode().strip()
        print(f"Вариант1: {result}")
        if result != "N/A":
            return int(result)
    except Exception:
        pass

    try:
        result = subprocess.check_output(cmd2).splitlines()[0].decode().strip()
        print(f"Вариант2: {result}")
        if result != "N/A":
            return int(result)
    except Exception:
        pass
        
    raise RuntimeError("Ошибка, не удалось определить исходное количество фреймов.")


def extract_frames(start_frame, end_frame):
    ''' Извлечение фреймов и распределение их по чанкам исходя из chunk_size '''
    
    frames_to_process = end_frame - start_frame + 1
    extracted_frames = []

    with frame_counter.get_lock():
        start_counter = frame_counter.value
        frame_counter.value += frames_to_process

    for chunk_start in range(start_frame, end_frame + 1, chunk_size):
        chunk_end = min(chunk_start + chunk_size - 1, end_frame)
        extract_frames_dir = os.path.join(frames_dir, f"file_%06d.png")

        cmd = [
            "ffmpeg", "-hwaccel", "cuda", "-i", video_file,
            "-vf", f"select='between(n,{chunk_start},{chunk_end})'",
            "-vsync", "0", "-start_number", str(chunk_start), extract_frames_dir
        ]
        subprocess.run(cmd, check=True)
        print(cmd)

        for i in range(chunk_end - chunk_start + 1):
            frame_number = start_counter + i + (chunk_start - start_frame)
            frame_path = extract_frames_dir % frame_number
            extracted_frames.append(frame_path)
                
    return extracted_frames
    
def chunk_processing(extracted_frames):
    ''' Старт обработки каждого заполненного чанка '''
    
    for frame_path in extracted_frames:
    
        # Извлекаем имя изображения для последующего сохранения 3D изображения
        frame_name = os.path.splitext(os.path.basename(frame_path))[0]
        
        # Загрузка изображения
        image = cv2.imread(frame_path)
        
        # Размеры изображения
        height, width = image.shape[:2]

        # Запуск depth_processing и получение карты глубины
        depth = depth_processing(image)

        # Запуск image3d_processing и получение двух изображений стереопары
        if PARALLAX_METHOD == 1:
            left_image, right_image = image3d_processing_method1(image, depth, height, width)
        elif PARALLAX_METHOD == 2:
            left_image, right_image = image3d_processing_method2(image, depth, height, width)
        else:
            print(f"Задайте корректный {PARALLAX_METHOD}.")

        # Объединение стереопары в общее 3D изображение
        image3d = image3d_combining(left_image, right_image, height, width)
        
        # Сохранение 3D изображения
        output_image3d_path = os.path.join(images3d_dir, f'{frame_name}.jpg')
        cv2.imwrite(output_image3d_path, image3d, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
        #cv2.imwrite(output_image3d_path, image3d) # Если PNG

        # Удаление исходного файла
        os.remove(frame_path)
        
    with threads_count.get_lock():
        threads_count.value = max(1, threads_count.value - 1) # Уменьшение счетчика после завершения текущего потока


def run_processing():
    ''' Глобальная функция старта обработки с учетом многопоточности'''
    
    # Получение количества фреймов в видео
    total_frames = get_total_frames()
                        
    # Управление потоками
    if total_frames:
        with ThreadPoolExecutor(max_workers=max_threads) as executor:
            futures = []
            for start_frame in range(0, total_frames, chunk_size):
                while True:
                    with threads_count.get_lock():
                        if threads_count.value < max_threads:
                            threads_count.value += 1
                            break
                            
                    time.sleep(5) # Пауза перед повторной проверкой на количество работающих потоков

                end_frame = min(start_frame + chunk_size - 1, total_frames - 1)
                extracted_frames = extract_frames(start_frame, end_frame)
                future = executor.submit(chunk_processing, extracted_frames)
                futures.append(future)
            
            # Ожидаем завершения задач
            for future in futures:
                future.result()


# ЗАПУСК ОБРАБОТКИ
run_processing()

print("ГОТОВО.")


# Очищаем Cuda
del model_depth
torch.cuda.empty_cache()

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

За выгрузку отвечает ffmpeg. Задаем chunk_size (количество фреймов на поток) и max_threads (количество потоков), скрипт последовательно обрабатывает все фреймы до последнего. Предварительно мы получаем значение общего количества фреймов с помощью ffprobe. Все подробности по настройке параметров в предыдущей (основной) статье. Могу лишь заметить, что на моей конфигурации (AMD Ryzen 5 PRO 3600, 32Gb DDR4, RTX 3060 12GB) в среднем достаточно 3-5 потоков и примерно по 5000 фреймов на поток.

Для чего вообще возникла идея с многопоточной обработкой (псевдо-многопоточной)? Во-первых, медленная выгрузка фреймов. Мы выгружаем по диапазону, например:
ffmpeg -hwaccel cuda -i video.mkv -vf "select='between(n,5000,10000)'" -vsync 0 -start_number 5000 "extracted_frames/file_%06d.png"

затем:
ffmpeg -hwaccel cuda -i "video.mkv" -vf "select='between(n,10001,15000)'" -vsync 0 -start_number 10001 "extracted_frames/file_%06d.png"

и тд.

Соответственно ffmpeg должен пересчитать все кадры перед выгрузкой (а возможно вообще все кадры), чтобы выгрузить корректно. Это также зависит от конкретного кодека и алгоритма кодирования. У меня пока не получилось ускорить этот процесс, чтобы соблюсти точную синхронизацию (чтобы не было пропуска или дублирования кадров), если кто-то знает, как это сделать лучше, поделитесь пожалуйста в комментариях, я протестирую и обновлю скрипт.

C-3PO в объеме
C-3PO в объеме

Вкратце по команде:
ffmpeg -hwaccel cuda -i video.mkv -vf "select='between(n,5000,10000)'" -vsync 0 -start_number 5000 "extracted_frames/file_%06d.png"

"-hwaccel cuda" - используем CUDA для выгрузки, обычно быстрее чем на CPU
"-i video.mkv" - исходный видео-файл
"-vf "select='between(n,5000,10000)'"" - фильтр диапазона, от 5000 до 10000 фрейма
"-vsync 0" - отключаем синхронизацию по временным меткам, выгружаем фреймы как есть
"-start_number" - счетчик для названия, начинаем от 5000
""extracted_frames/file_%06d.png"" - путь куда будут выгружаться фреймы и маска файлов, где %06d - 6-значный счетчик, файлы будут вида "file_005000.png", "file_005001.png" и тд.

После полной обработки из полученных фреймов нужно будет "вручную" скомпилировать фильм, не забыв подключить звуковые дорожки из исходного файла.
Пример команды:
ffmpeg -r 24000/1001 -i "frames_3d/file_%06d.jpg" -i video.mkv -c:v hevc_nvenc -b:v 20M -minrate 10M -maxrate 30M -bufsize 60M -preset p7 -map 0:v -map 1:a -c:a copy -pix_fmt yuv420p video_3d.mkv

Здесь:
"-r 24000/1001" - частота кадров (как было в исходнике), 24000/1001=23,976 кадров в секунду
"-i "frames_3d/file_%06d.jpg"" - папка с итоговыми фреймами
"-i video.mkv" - исходный видео-файл для экспорта из него аудио-дорожек
"-c:v hevc_nvenc" - кодек, в данном случае для CUDA - HEVC H265(быстро и качественно)
"-b:v 20M -minrate 10M -maxrate 30M" - переменный битрейт, среднее значение 20Мбит/сек, минимальное 10Мбит/сек, максимальное 30Мбит/сек
"-bufsize 60M" - размер буфера, рекомендуется использовать 2x от maxrate (2x30M=60M), либо можно не указывать, будет на усмотрение ffmpeg
"-preset p7" - пресет 7 для кодека hevc_nvenc, высокое качество
"-map 0:v" - указываем использовать папку с фреймами для основного видеоряда
"-map 1:a -c:a copy" - указываем использовать аудио-дорожки из "-i video.mkv" без перекодирования, "-c:a copy" - прямое копирование
"-pix_fmt yuv420p" - цветовой формат пикселей, для выходных видео рекомендуется использовать yuv420p
"video_3d.mkv" - имя выходного файла

Скрипт можно доработать и включить эту команду на автовыполнение после завершения обработки фреймов, например так:

Код:
compile_video3d = [
	"ffmpeg",
	"-r", "24000/1001",
	"-i", "frames_3d/file_%06d.jpg",
	"-i", "video.mkv",
	"-c:v", "hevc_nvenc",
	"-b:v", "20M",
	"-minrate", "10M",
	"-maxrate", "30M",
	"-bufsize", "60M",
	"-preset", "p7",
	"-map", "0:v",
	"-map", "1:a",
	"-c:a", "copy",
	"-pix_fmt", "yuv420p",
	"video_3d.mkv"
]

subprocess.run(compile_video3d, check=True)

Лично я предпочитаю делать вручную, т.к. постоянно приходится что-то корректировать, например убирать часть аудио-дорожек, или экспериментировать с кодеками, фреймрейтом и чем угодно еще.

После компиляции нужно не забыть удалить каталог с полученными фреймами.

Пример карты глубины для фрейма
Пример карты глубины для фрейма

Новая функция параллакса

В скрипте появилась новая функция создания параллакса, теперь их две:
image3d_processing_method1
image3d_processing_method2

Первая более быстрая и с более мягким сглаживанием при смещениях. Вторая делает смещения по другому принципу, 3D получается немного другим. Объекты через эти функции могут смещаться по разному, рекомендую проверить оба варианта и выбрать предпочтительный. Глобально 3D-объем хорошо получается через оба метода, и в целом картинки очень похожи, но в то же время есть заметные различия. Выбор метода задается через параметр PARALLAX_METHOD, доступны 2 варианта: 1 и 2 соответственно.

Появился еще один новый параметр - INPAINT_RADIUS. Это радиус заполнения смещений в пикселях для второго метода (image3d_processing_method2). В данном случае это заполнение соседними пикселями на краях изображений при их смещении. То есть, это заполнение черной рамки вокруг смещенных изображений, чтобы вместо черного цвета использовались соседние пиксели. Рекомендуется от 2 до 5, в большинстве случаев достаточно 2-3. Если значение больше, например INPAINT_RADIUS = 15, тогда края будут слишком размыты и время обработки существенно увеличится. Если же наоборот, выставить 0 или 1, тогда рамка будет выглядеть слишком резкой и не точной. В общем, в большинстве случаев достаточно указать 2 или 3 пикселя.

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

Пример карты глубины для фрейма
Пример карты глубины для фрейма

Другие решения

VapourSynth

В комментариях к предыдущим статьям мне подсказали пару интересных решений.

Во-первых, вместо цепочки: выгрузка фреймов -> их обработка -> компиляция из готовых фреймов итогового видео, можно использовать промежуточный сервер обработки видео на лету, например VapourSynth. Схема там примерно в следующем: ffmpeg выгружает фрейм, сразу (без сохранения) передает его в функцию обработки (в данном случае это генерация 3D-версии кадра), полученный фрейм кодируется в выходной видео-файл (точнее наверное встает в очередь на кодирование). Все это происходит в RAM / VRAM, минуя промежуточные этапы сохранения кадров на жесткий диск.

Я пока не экспериментировал с этим. Попробовал поставить на Убунту, но VapourSynth стал требовать самых последних версий ffmpeg и некоторых других библиотек (apt update не помог, стабильных версий оказалось не достаточно). Пришлось компилировать вручную последний ffmpeg (хотя лично меня вполне устраивал последний стабильный), и еще несколько других библиотек, но VapourSynth так и не получилось завести. Позже я обязательно вернусь к этому, когда будет больше свободного времени. Возможно под Windows это дело заводится проще.

Важный момент по обработке на лету. С одной стороны это удобно, с другой - есть нюансы. Обработка одного фильма на модели Depth-Anything-V2 Large может длиться больше суток, а то и несколько. Я приводил приблизительный расчет для фильма StarWars4 формата FullHD продолжительностью 2 часа 4 минуты в основной статье. На моей конфигурации (AMD Ryzen 5 PRO 3600, 32Gb DDR4, RTX 3060 12GB) с моделью Large потребовалось бы порядка 32 часов на обработку данного фильма, и если в процессе произойдет сбой или случайно выключится компьютер - придется начинать все с начала.

Дарт Вейдер просто душка
Дарт Вейдер не стерпел бы этого

Другой момент. Надо быть точно уверенным в исходниках. В статье про апскейл старых видео (https://habr.com/ru/articles/904784/) я подробно описал, какие проблемы могут возникнуть при работе с некоторыми исходниками и форматами, особенно если это DVD-MPEG2 или что-то другое из юрского периода той эпохи. Там могут быть проблемы с точным определением частоты кадров, с выходным форматом изображения и чем угодно еще. Это нужно учитывать и предварительно проверять исходники и то, что из них получается на выходе.

А так, в целом, реализация с сервером обработки без сохранения фреймов - замечательная идея, ведь здесь совсем не требуется место под фреймы на диске, а если мы работаем с форматом 4K и PNG, это очень критичный момент.

Другая библиотека для конвертации 2D -> 3D

К комментариях также подсказали другое возможное решение:
https://github.com/nagadomi/nunif/tree/master/iw3

Я его не пробовал, лишь бегло посмотрел. Там есть GUI и очень много настроек. Можно выбрать модель глубины, метод обработки, есть поддержка анаглифа и много чего еще. Пожалуй в этой реализации будет сложнее разобраться, но решение определенно заслуживает внимания.

Заключение

За сим я наверное буду заканчивать с серией публикаций по теме конвертации видео из 2D в 3D. Еще раз напоминаю основную статью:
Как сделать 3D версию любого фильма на примере StarWars4 (DepthAnythingV2 + Parallax)
Там более подробно - как это работает, описание параметров, примеры изображений и тд.

Скрипты я периодически обновляю, за актуальными версиями можно следить тут:
https://github.com/peterplv/MakeAnythingStereo3D, либо в Тг-канале (https://t.me/peter_touch_ai).

На эксперименты я потратил несколько недель, просмотрел около десятка фильмов в 3D (это была исключительно приятная часть).

Вообще, изначально появилось желание "отдохнуть", отвлечься от основной деятельности (LLM, RAGи, агенты и их применение в различных задачах), переключиться на что-то более "легкое" и приятное. В итоге отдых перешел в практическое русло и получилось реализовать 2 pet-проекта - конвертацию видео 2D -> 3D, и апскейл старых видео (статья об этом). Параллельно пришлось глубже изучить незаменимый ffmpeg, погрузиться в некоторые форматы видео, немного в структуру используемых нейронок, познакомиться с принципом эффекта параллакса и тд тд. В общем, полезный получился отдых :)

Буду признателен, если поставите лайк за старания.

До прибудет с вами сила.

Чуви и Хан благодарят за внимание
Чуви и Хан благодарят за внимание
Только зарегистрированные пользователи могут участвовать в опросе. Войдите, пожалуйста.
Смотрите ли вы 3D?
50% Да, люблю 3D5
10% Редко, но люблю 3D1
10% Мне безразлично 3D, но могу посмотреть1
30% Не смотрю3
Проголосовали 10 пользователей. Воздержавшихся нет.
Теги:
Хабы:
+4
Комментарии3

Публикации

Ближайшие события