Taisei replay video recording script

taisei-record-replay.py

#!/usr/bin/env python3

import subprocess
import tempfile
import argparse
import os
import math
from pathlib import Path

FPS = 60

def main(argv):
    ncpu = len(os.sched_getaffinity(0))
    ncpu = min(16, ncpu)

    ffmpeg_args = {
        'h264': {
            'common': ['-c:v', 'libx264', '-preset', 'veryslow'],
            'pass1': ['-pass', '1'],
            'pass2': ['-pass', '2'],
        },
        'h265': {
            'common': [
                '-c:v', 'libx265', 
                '-preset', 'veryslow',
                '-threads', f'{min(16, ncpu)}',
            ],
            'pass1': ['-x265-params', 'pass=1'],
            'pass2': ['-x265-params', 'pass=2'],
        },
        'vp9': {
            'common': [
                '-c:v', 'libvpx-vp9',
                '-deadline', 'good',
                '-cpu-used', '2',
                '-row-mt', '1',
                '-tile-rows', '2',
                '-tile-columns', '2',
            ],
            'pass1': ['-pass', '1'],
            'pass2': ['-pass', '2'],
        }
    }

    p = argparse.ArgumentParser(description='Record a Taisei replay')

    p.add_argument('-v', '--viewport',
        help='Record the viewport instead of the screen',
        action='store_true',
    )

    p.add_argument('-s', '--size',
        help='File size to target',
        type=float,
        default=25,
    )

    p.add_argument('-W', '--width',
        help='Game framebuffer width',
        type=int,
        default=0,
    )

    p.add_argument('-H', '--height',
        help='Game framebuffer height',
        type=int,
        default=0,
    )

    p.add_argument('-c', '--codec',
        help='Video codec to use',
        default='vp9',
        const='vp9',
        nargs='?',
        type=str,
        action='store',
        choices=list(ffmpeg_args),
    )

    p.add_argument('replay',
        help='The replay file',
        type=Path,
    )

    p.add_argument('video',
        help='Path to the output video',
        type=Path,
    )

    args = p.parse_args()
    print(args)
    ffmpeg_args = ffmpeg_args[args.codec]

    if args.width <= 0 and args.height <= 0:
        args.width = 1600

    with tempfile.TemporaryDirectory(prefix='taisei-framedump-') as td:
        prefix = str(Path(td) / 'frame')
        env = os.environ.copy()
        env.update({
            'SDL_VIDEODRIVER': 'offscreen',
            'TAISEI_FRAMEDUMP': prefix,
            'TAISEI_FRAMEDUMP_SOURCE': 'viewport' if args.viewport else 'screen',
            'TAISEI_LOGLVLS_STDOUT': '-a',
            'TAISEI_LOGLVLS_STDERR': '+iwe',
        })

        def run(cmd, *args, **kwargs):
            print("RUN:", cmd)
            return subprocess.run(cmd, *args, **kwargs)

        run([
            'taisei', 
            '-f1', 
            '-r', args.replay, 
            # '--renderer=gles30',
            '--width', str(args.width),
            '--height', str(args.height),
        ], check=True, env=env)

        num_frames = len(tuple(Path(td).glob('frame*.png')))
        duration = num_frames / FPS
        bitrate = math.floor(args.size * 8000 / duration)

        ffmpeg = [
            'ffmpeg',
            '-hide_banner',
            '-y',
            '-framerate', f'{FPS}',
            '-pattern_type', 'glob',
            '-i', prefix + '*.png',
            '-pix_fmt', 'yuv420p',
            '-b:v', f'{bitrate}k',
            '-threads', str(ncpu),
        ] + ffmpeg_args['common']

        run(ffmpeg + ffmpeg_args['pass1'] + [
            '-f', 'null',
            '/dev/null',
        ], check=True)

        run(ffmpeg + ffmpeg_args['pass2'] + [
            args.video
        ], check=True)


if __name__ == '__main__':
    import sys
    exit(main(sys.argv))