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

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()

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)