playground/ffmpeg.py

#!/usr/bin/env python3.9

import dataclasses
import json
import logging
import subprocess
import typing
from datetime import timedelta
from math import ceil
from pathlib import Path


logger = logging.getLogger(__name__)


@dataclasses.dataclass
class FfprobeResult:
    duration_sec: float
    duration_human: str
    codec: str
    fps: float
    size_bytes: int
    width: int
    height: int
    bitrate: int
    container: str
    ar: float
    sample_ar: float
    tags: typing.Dict[str, str]


def ffprobe(video: Path) -> FfprobeResult:
    logger.debug("executing ffprobe")
    proc = subprocess.run(
        args=[
            "ffprobe",
            "-v",
            "error",
            "-show_streams",
            "-show_format",
            "-print_format",
            "json",
            str(video),
        ],
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    try:
        proc.check_returncode()
    except subprocess.CalledProcessError:
        logger.error("failed to run ffprobe", exc_info=True)
        raise
    output: dict = json.loads(proc.stdout)
    video_stream: dict = [s for s in output.get("streams", []) if s.get("codec_type") == "video"][0]
    video_format: dict = output.get("format", {})
    width, height = video_stream.get("width"), video_stream.get("height")
    profile = video_stream.get("profile")
    codec = video_stream.get("codec_name")
    duration = float(video_format.get("duration"))
    fps_long = eval(video_stream.get("r_frame_rate", ""))
    fps = float(f"{fps_long:.3f}")
    size = int(video_format.get("size", 0))
    duration_time = timedelta(seconds=ceil(duration))
    bitrate = int(video_stream.get("bit_rate", video_format.get("bit_rate", 0)))
    tags = video_format.get("tags", {})

    try:
        sample_w, sample_h = list(map(int, video_stream["sample_aspect_ratio"].split(":")))
        sample_ar = (sample_w / sample_h) or 1
    except:
        sample_ar = 1

    logger.debug(f"dimensions={width}x{height} duration={duration_time}")
    # total_frames = video_stream.get("nb_frames")

    return FfprobeResult(
        duration_sec=duration,
        duration_human=str(duration_time),
        codec=codec,
        fps=fps,
        size_bytes=size,
        width=width,
        bitrate=bitrate,
        height=height,
        container=video.suffix.lstrip(".").lower(),
        ar=width / height,
        sample_ar=sample_ar,
        tags=tags,
        # total_frames=total_frames,
    )


def strip_metadata(video_path: Path, save_path: Path) -> Path:
    media_info = ffprobe(video_path)
    # check if the first video stream has hevc codec
    if media_info["streams"][0]["codec_name"] == "hevc":
        logger.info(f"Video codec is hevc, adding hvc1 tag")
        extra_args = ["-tag:v", "hvc1"]
    else:
        extra_args = []

    # fmt: off
    args = [
        'ffmpeg',
        '-i',
        str(video_path),
        '-c:v', 'copy',
        '-movflags', '+faststart', 
        '-map_metadata', '-1',
        *extra_args,
        '-c:a', 'copy',
        '-y',
        str(save_path),
    ]
    # fmt: on
    logger.info(f"calling ffmpeg with {args=}")
    subprocess.run(args, check=True)
    return save_path


def make_thumbnail_tile(
    video: Path,
    image_path: Path = None,
    columns: int = 3,
    interval_seconds: int = 60,
    tile_width: int = 540,
    skip_first_sec: int = 10,
    skip_if_exists: bool = False,
) -> Path:
    if not image_path:
        image_path = video.parent / f"{video.stem}.thumbnail.jpg"
    if image_path.is_file() and skip_if_exists:
        logger.debug("thumbnail already exists")
        return image_path

    info = ffprobe(video)

    min_frames = columns * 3

    sec_per_frame = min(interval_seconds, ceil(info.duration_sec / min_frames))

    rows = ceil(info.duration_sec // sec_per_frame / columns)
    tile = f"{columns}x{rows}"

    scaled_width, scaled_height = tile_width, ceil(tile_width / info.ar / info.sample_ar)
    scale = f"{scaled_width}:{scaled_height}"

    # font_path = Path(__file__).parent / 'iosevka.ttf'
    # timestamp_filter = rf"drawtext=r=1:timecode='00\:00\:00\:00':fontsize=16:fontcolor=white:x=10:y=10:box=1:boxcolor=black@0.5"
    # timestamp_filter = fr"drawtext=text='%{{(pts\\+{skip_first_sec})\:hms}}':fontsize=16:fontcolor=white:x=10:y=10:box=1:boxcolor=black@0.5"

    proc = subprocess.run(
        # fmt: off
        args=[
            "ffmpeg",
            "-v", "error",
            "-skip_frame", "nokey",
            "-ss", f"{skip_first_sec}",
            "-i", str(video.absolute()),
            # "-vf", f"fps=1/{sec_per_frame},scale={scale},{timestamp_filter},tile={tile}",
            "-vf", f"fps=1/{sec_per_frame},scale={scale},tile={tile}",
            "-frames", "1",
            "-y", str(image_path.absolute()),
        ],
        # fmt: on
        stdout=subprocess.DEVNULL,
        timeout=2000,
    )
    proc.check_returncode()
    return image_path


if __name__ == "__main__":
    import sys

    logging.basicConfig(level=logging.DEBUG)
    for it in sys.argv[1:]:
        video = Path(it)

        if not video.exists():
            continue

        logger.info("creating thumbnail for %s", video)
        make_thumbnail_tile(video)