playground/ffmpeg.py

#!/usr/bin/env python3.9

import dataclasses
import json
import logging
import subprocess
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: dict[str, str]

    @property
    def resolution(self) -> str:
        return f'{self.width}x{self.height}'

    @property
    def has_metadata(self):
        return any(k in self.tags for k in ['title', 'comment'])

    @property
    def hd_mode(self) -> str:
        if 700 < self.height < 800:
            return '720p'
        if 900 < self.height < 1200:
            return '1080p'
        if 2000 < self.height < 2200:
            return '4K'
        raise ValueError('unknown HD mode')


def ffprobe(video_path: Path) -> FfprobeResult:
    proc = subprocess.run(
        args=[
            "ffprobe",
            "-v",
            "error",
            "-show_streams",
            "-show_format",
            "-print_format",
            "json",
            str(video_path),
        ],
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    try:
        proc.check_returncode()
    except subprocess.CalledProcessError as e:
        logger.error(f"failed to run ffprobe. stderr={e.stderr}", 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

    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_path.suffix.lstrip(".").lower(),
        ar=width / height,
        sample_ar=sample_ar,
        tags=tags,
    )


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.codec == "hevc":
        logging.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
    logging.info(f"calling ffmpeg with {args=}")
    subprocess.run(args, check=True)
    return save_path


def make_thumbnail_tile(
    video: Path,
    columns: int = 3,
    interval_seconds: int = 60,
    tile_width: int = 540,
    skip_first_sec: int = 10,
) -> bytes:
    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"

    args = [
        # fmt: off
        "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",
        "-c:v",
        "jpeg2000",
        "-q:v", "90",
        "-f", "image2pipe",
        "-",
        # fmt: on
    ]
    logger.debug(f'calling ffmpeg with args={args}')
    proc = subprocess.run(
        args=args,
        stdout=subprocess.PIPE,
        timeout=2000,
        # preexec_fn=os.setpgrp,
    )
    proc.check_returncode()
    return bytes(proc.stdout)


if __name__ == "__main__":
    import sys

    logging.basicConfig(level=logging.DEBUG)
    for it in sys.argv[1:]:
        video = Path(it)
        image_path = video.parent / f"{video.stem}.thumbnail.jpg"

        if not video.exists():
            continue

        logger.info("creating thumbnail for %s", video)
        image_jpeg = make_thumbnail_tile(video)
        image_path.write_bytes(image_jpeg)