【Python】手持ちの2D動画を3D化！VRで楽しむ3D動画の作り方（Video Depth Anything）

import os
import cv2
import numpy as np
import argparse
import subprocess
import shutil
import json

# ------------------------------
# フレーム読み込み関数
# ------------------------------
def read_frame_pair(src_video, depth_video=None):
    """
    動画からフレームを読み込み、元画像と深度マップを返す関数

    Parameters
    ----------
    src_video : cv2.VideoCapture
        元動画キャプチャ
    depth_video : cv2.VideoCapture or None, optional
        深度動画キャプチャ。None の場合は src_video の右半分を深度マップとして使用

    Returns
    -------
    tuple
        (src_frame, depth_frame)。フレーム末尾の場合は (None, None)
    """
    ret_src, frame = src_video.read()
    if not ret_src:
        return None, None

    if depth_video is None:
        width = frame.shape[1] // 2
        src_frame = frame[:, :width]
        depth_frame = frame[:, width:]
    else:
        ret_depth, depth_frame = depth_video.read()
        if not ret_depth:
            return None, None
        src_frame = frame

    return src_frame, depth_frame

# ------------------------------
# VR用SBS画像生成関数
# ------------------------------
def create_vr_image(src_frame, depth_frame, max_shift=10, half_sbs=True, depth_min=0, depth_max=255):
    """
    元画像と深度マップから左右視差を持つVR用SBS画像を生成

    Parameters
    ----------
    src_frame : np.ndarray
        元画像
    depth_frame : np.ndarray
        深度マップ画像
    max_shift : int
        最大視差（ピクセル）
    half_sbs : bool
        Trueなら幅半分SBSで出力
    depth_min : int
        深度マップの最小値
    depth_max : int
        深度マップの最大値

    Returns
    -------
    np.ndarray
        左右視差付きSBS画像
    """
    height, width = src_frame.shape[:2]
    depth_gray = cv2.cvtColor(depth_frame, cv2.COLOR_BGR2GRAY).astype(np.float32)
    depth_norm = np.clip((depth_gray - depth_min) / (depth_max - depth_min), 0, 1)
    shift = depth_norm * max_shift
    map_x, map_y = np.meshgrid(np.arange(width), np.arange(height))
    map_x = map_x.astype(np.float32)
    map_y = map_y.astype(np.float32)
    map_x_left  = map_x - shift
    map_x_right = map_x + shift
    left  = cv2.remap(src_frame, map_x_left, map_y, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
    right = cv2.remap(src_frame, map_x_right, map_y, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)

    if half_sbs:
        left_half  = cv2.resize(left,  (width // 2, height))
        right_half = cv2.resize(right, (width // 2, height))
        combined = np.hstack((left_half, right_half))
    else:
        combined = np.hstack((left, right))

    return combined

# ------------------------------
# VR動画生成関数
# ------------------------------
def generate_3d_sbs_video(src_video_path, depth_video_path=None, output_path="output.mp4", max_shift=10, half_sbs=True, depth_min=0, depth_max=255):
    """
    元動画と深度動画からVR動画を生成

    Parameters
    ----------
    src_video_path : str
        元動画ファイルパス
    depth_video_path : str or None
        深度動画ファイルパス。Noneの場合は src_video の右半分を使用
    output_path : str
        出力動画パス
    max_shift : int
        最大視差
    half_sbs : bool
        Trueなら幅半分SBSで出力
    depth_min : int
        深度マップ最小値
    depth_max : int
        深度マップ最大値

    Returns
    -------
    str
        生成したVR動画のパス
    """
    video = cv2.VideoCapture(src_video_path)
    depth = cv2.VideoCapture(depth_video_path) if depth_video_path else None
    width  = int(video.get(cv2.CAP_PROP_FRAME_WIDTH) // (2 if depth_video_path is None else 1))
    height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps    = video.get(cv2.CAP_PROP_FPS)
    frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    out_width  = width if half_sbs else width * 2
    out_height = height
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (out_width, out_height))

    frame_idx = 0
    while True:
        src_frame, depth_frame = read_frame_pair(video, depth_video=depth)
        if src_frame is None:
            break
        vr_frame = create_vr_image(src_frame, depth_frame, max_shift=max_shift, half_sbs=half_sbs, depth_min=depth_min, depth_max=depth_max)
        out.write(vr_frame)
        frame_idx += 1
        progress = (frame_idx / frame_count) * 100
        print(f"\r進捗: {frame_idx}/{frame_count} ({progress:.2f}%)", end="")

    video.release()
    if depth:
        depth.release()
    out.release()
    print(f"\nVR動画生成完了: {output_path}")
    return output_path

# ------------------------------
# 音声有無チェック関数
# ------------------------------
def has_audio_stream(file_path):
    """
    ffprobe を使って音声ストリームの有無を確認

    Parameters
    ----------
    file_path : str
        チェックする動画ファイルパス

    Returns
    -------
    bool
        True なら音声あり、False なら音声なし
    """
    cmd = [
        "ffprobe", "-v", "error",
        "-select_streams", "a",
        "-show_entries", "stream=index",
        "-of", "json",
        file_path
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    info = json.loads(result.stdout or "{}")
    return "streams" in info and len(info["streams"]) > 0

# ------------------------------
# FFmpegで音声上書き（無音トラック対応）
# ------------------------------
def overwrite_audio_ffmpeg(video_path, audio_src_path, output_path):
    """
    VR動画に音声を上書き。音声が無ければ無音トラックを追加

    Parameters
    ----------
    video_path : str
        音声を上書きするVR動画
    audio_src_path : str or None
        音声元の動画。Noneの場合は無音トラックを追加
    output_path : str
        最終出力ファイル
    """
    if not audio_src_path or not has_audio_stream(audio_src_path):
        print(f"⚠ 音声ストリームが見つかりません → 無音トラックを追加")
        cmd = [
            "ffmpeg", "-y",
            "-i", video_path,
            "-f", "lavfi",
            "-i", "anullsrc=channel_layout=stereo:sample_rate=44100",
            "-c:v", "copy",
            "-c:a", "aac",
            "-shortest",
            output_path,
        ]
        subprocess.run(cmd, check=True)
        return

    cmd = [
        "ffmpeg",
        "-i", video_path,
        "-i", audio_src_path,
        "-c:v", "copy",
        "-map", "0:v:0",
        "-map", "1:a:0",
        "-y",
        output_path
    ]
    subprocess.run(cmd, check=True)
    print(f"音声上書き完了: {output_path}")

# ------------------------------
# 1本の動画を処理
# ------------------------------
def process_single_video(src_path, depth_path, audio_dir, output_dir, max_shift=20, half_sbs=False):
    """
    1本の動画をVR化し、音声付きで出力

    Parameters
    ----------
    src_path : str
        元動画パス
    depth_path : str or None
        深度動画パス
    audio_dir : str or None
        音声フォルダ。Noneの場合は無音トラック追加
    output_dir : str
        出力フォルダ
    max_shift : int
        最大視差
    half_sbs : bool
        Trueなら幅半分SBS
    """
    os.makedirs(output_dir, exist_ok=True)
    base_name = os.path.splitext(os.path.basename(src_path))[0]

    # 元動画の末尾 "_src" を削除して "_3D" を付与
    if base_name.endswith("_src"):
        base_name = base_name[:-4]
    base_name += "_3D"

    # 一時出力ファイル
    temp_output = os.path.join(output_dir, f"{base_name}_temp.mp4")
    # 最終出力ファイル（_output を削除）
    final_output = os.path.join(output_dir, f"{base_name}.mp4")

    # VR動画生成
    generate_3d_sbs_video(src_path, depth_video_path=depth_path, output_path=temp_output,
                          max_shift=max_shift, half_sbs=half_sbs)

    # 音声上書き（audio_dir が None でも無音トラックで対応）
    audio_src_path = os.path.join(audio_dir, os.path.basename(src_path).replace("_src", "")) if audio_dir else None
    overwrite_audio_ffmpeg(temp_output, audio_src_path, final_output)

    # 一時ファイル削除
    os.remove(temp_output)


# ------------------------------
# バッチ処理関数
# ------------------------------
def batch_process_videos(input_dir, output_dir, audio_dir=None, max_shift=20, half_sbs=False):
    """
    フォルダ内の動画をすべてVR化して出力

    Parameters
    ----------
    input_dir : str
        元動画フォルダ
    output_dir : str
        出力先フォルダ
    audio_dir : str or None
        音声フォルダ。Noneの場合は無音トラックで出力
    max_shift : int
        最大視差
    half_sbs : bool
        Trueなら幅半分SBSで出力
    """
    os.makedirs(output_dir, exist_ok=True)
    files = os.listdir(input_dir)
    processed = set()

    for file in files:
        if not file.endswith(".mp4"):
            continue

        file_path = os.path.join(input_dir, file)
        base_name = file.rsplit(".mp4", 1)[0]

        if base_name.endswith("_src"):
            core_name = base_name[:-4]
            vis_file = f"{core_name}_vis.mp4"
            vis_path = os.path.join(input_dir, vis_file)
            depth_video_path = vis_path if os.path.exists(vis_path) else None
            print(f"処理中: {file} + {'_visあり' if depth_video_path else '右半分を深度として使用'}")
            process_single_video(file_path, depth_video_path, audio_dir, output_dir, max_shift=max_shift, half_sbs=half_sbs)
            processed.add(file)
            if depth_video_path:
                processed.add(vis_file)
        elif file not in processed:
            print(f"処理中 (右半分を深度として使用): {file}")
            process_single_video(file_path, None, audio_dir, output_dir, max_shift=max_shift, half_sbs=half_sbs)
            processed.add(file)

# ------------------------------
# コマンドライン引数
# ------------------------------
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Batch 3D SBS Video Generation with optional Audio")
    parser.add_argument("--input_dir", required=True, help="入力動画フォルダ")
    parser.add_argument("--output_dir", required=True, help="出力先フォルダ")
    parser.add_argument("--audio_dir", help="オーディオフォルダ（省略可）")
    parser.add_argument("--max_shift", type=int, default=10, help="最大視差")
    parser.add_argument("--half_sbs", action="store_true", help="幅半分SBSで出力")
    args = parser.parse_args()

    batch_process_videos(args.input_dir, args.output_dir, args.audio_dir, max_shift=args.max_shift, half_sbs=args.half_sbs)

import os
import argparse
import torch
import numpy as np
from video_depth_anything.video_depth import VideoDepthAnything
from utils.dc_utils import read_video_frames, save_video

def process_videos(input_dir, output_dir, encoder_type='vitl', metric=False,
                   input_size=518, max_res=1280, fp32=False, grayscale=False):

    os.makedirs(output_dir, exist_ok=True)
    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

    # --- モデルロード（1回だけ） ---
    model_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]},
    }
    checkpoint_name = 'metric_video_depth_anything' if metric else 'video_depth_anything'

    video_depth_anything = VideoDepthAnything(**model_configs[encoder_type], metric=metric)
    video_depth_anything.load_state_dict(
        torch.load(f'./checkpoints/{checkpoint_name}_{encoder_type}.pth', map_location='cpu'),
        strict=True
    )
    video_depth_anything = video_depth_anything.to(DEVICE).eval()

    # --- フォルダ内の動画を処理 ---
    video_extensions = ('.mp4', '.mov', '.avi', '.mkv')

    for file_name in os.listdir(input_dir):
        if not file_name.lower().endswith(video_extensions):
            continue

        input_path = os.path.join(input_dir, file_name)
        print(f"Processing: {input_path}")

        frames, target_fps = read_video_frames(input_path, process_length=-1, max_res=max_res)
        depths, fps = video_depth_anything.infer_video_depth(frames, target_fps, input_size=input_size, device=DEVICE, fp32=fp32)

        # 保存
        base_name = os.path.splitext(file_name)[0]
        processed_video_path = os.path.join(output_dir, base_name + '_src.mp4')
        depth_vis_path = os.path.join(output_dir, base_name + '_vis.mp4')

        save_video(frames, processed_video_path, fps=fps)
        save_video(depths, depth_vis_path, fps=fps, is_depths=True, grayscale=grayscale)

        print(f"Saved: {processed_video_path}, {depth_vis_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Batch Video Depth Conversion")
    parser.add_argument("--input_dir", required=True, help="入力動画フォルダ")
    parser.add_argument("--output_dir", required=True, help="出力先フォルダ")
    parser.add_argument("--encoder", default="vitb", choices=["vits", "vitb", "vitl"], help="使用するエンコーダ")
    parser.add_argument("--metric", action="store_true", help="metric モデルを使用")
    parser.add_argument("--input_size", type=int, default=518, help="入力サイズ")
    parser.add_argument("--max_res", type=int, default=1280, help="最大解像度")
    parser.add_argument("--fp32", action="store_true", help="fp32で推論")
    parser.add_argument("--grayscale", action="store_true", help="グレースケールで保存")
    args = parser.parse_args()

    process_videos(
        input_dir=args.input_dir,
        output_dir=args.output_dir,
        encoder_type=args.encoder,
        metric=args.metric,
        input_size=args.input_size,
        max_res=args.max_res,
        fp32=args.fp32,
        grayscale=args.grayscale
    )