pandora_render/ffmpeg.py

#!/usr/bin/env python3
from argparse import ArgumentParser
import math
import json
import os
import subprocess
import sys
import ox

def run(cmd):
    #print(' '.join('"%s"' % c for c in cmd))
    subprocess.call(cmd)

def add_subtitles(subtitles, clip, position, duration):
    if clip.get('subtitles'):
        if isinstance(clip['subtitles'], list):
            for sub in clip['subtitles']:
                subtitles.append({
                    'in': position + sub['in'],
                    'out': position + sub['out'],
                    'value': sub['value']
                })
        else:
            subtitles.append({
                'in': position,
                'out': position + duration,
                'value': clip['subtitles']
            })


usage = "usage: %(prog)s [options] edit.json"
parser = ArgumentParser(usage=usage)
parser.add_argument('-r', '--resolution', dest='height', type=int,
                    help="output height, default 480", default=480)
parser.add_argument('-a', '--aspect', dest='aspect', type=str,
                    help="aspect ratio, default 16/9 (float or /)", default='16/9')
parser.add_argument('-c', '--config', dest='config',
                    help='config.json containing config',
                    default='~/.ox/client.json')
parser.add_argument('path', metavar='path', type=str,
                    help='edit.json path')
opts = parser.parse_args()


edit_json = opts.path
edit = json.load(open(edit_json))

render = './cache'
output = os.path.splitext(edit_json)[0] + '.mp4'
height = opts.height
if '/' in opts.aspect:
    aspect = [int(p) for p in opts.aspect.split('/')]
    aspect = aspect[0] / aspect[1]
else:
    aspect = float(opts.aspect)

width = int(height * aspect)
width -= width % 2

if not os.path.exists(render):
    os.makedirs(render)

files = []
edit_duration = 0
subtitles = []
position = 0
for clip in edit:
    out = render + '/%s_%0.3f-%0.3f.ts' % (clip['oshash'], clip['in'], clip['out'])
    edit_duration += (clip['out']-clip['in'])
    if os.path.exists(out):
        try:
            duration = ox.avinfo(out)['duration']
        except KeyError:
            os.unlink(out)
            duration = None
        if duration:
            files.append(out)
            src_duration = clip['out']-clip['in']
            if abs(src_duration-duration) > 1:
                print(clip.get('annotation', clip['item']), 'expected', src_duration, 'got', duration, out)
            add_subtitles(subtitles, clip, position, duration)
            position += duration
            continue
    src_info = ox.avinfo(clip['path'])
    clip_aspect = src_info['video'][0]['width'] / src_info['video'][0]['height']
    if 'display_aspect_ratio' in src_info['video'][0]:
        ratio = [int(p) for p in src_info['video'][0]['display_aspect_ratio'].split(':')]
        clip_aspect = ratio[0] / ratio[1]
        #print(clip['path'], ratio, clip_aspect)
    if clip_aspect < aspect:
        x = width
        y = int(x / clip_aspect)
        y -= y % 2
    else:
        y = height
        x = int(y * clip_aspect)
        x -= x % 2
    vf = 'scale=%s:%s' % (x, y)
    if x != width:
        vf += ',crop=%s:%s' % (width, height)  # crop center
    elif y != height:
        offset = int(((y - height) / 3))
        vf += ',crop=w=%s:h=%s:x=0:y=%s' % (width, height, offset)
    options = [
        '-map_metadata', '-1',
        '-vf', vf,
        '-aspect', str(aspect),
        '-c:v', 'libx264',
        '-b:v', '2M',
        '-preset:v', 'medium', '-profile:v', 'high', '-level:v', '4.0',
        '-r', '25',
        '-c:a', 'aac',
        '-ar', '48000',
        '-ac', '2',
        '-b:a', '192k',
    ]
    clip_duration = math.ceil((clip['out'] - clip['in']) / (1/25)) * 1/25
    if not clip_duration:
        print('skip empty clip', clip)
    else:
        files.append(out)
        vid = src_info['video'][0]['id']
        if not src_info['audio']:
            audio = ['-f', 'lavfi', '-i', 'anullsrc=channel_layout=stereo:sample_rate=48000']
            audio_map = ['-map', '0,0', '-map', '1,1']
        else:
            aid = src_info['audio'][0]['id']
            audio = []
            if clip.get('volume', 1) != 1:
                audio += [
                    '-filter:a', 'volume=%s' % clip['volume']
                ]
            audio_map = ['-map', '0:%s,0:0' % aid, '-map', '0:%s,0:1' % vid]
        cmd = [
            'ffmpeg',
            '-hide_banner',
            '-nostats', '-loglevel', 'error',
        ] + audio + [
            '-ss', str(clip['in']),
            '-i', clip['path']
        ] + audio_map + options + [
            '-t', str(clip_duration),
            out
        ]
        run(cmd)
        try:
            duration = ox.avinfo(out)['duration']
        except:
            print('invalid file:', out)
            print('try again?:')
            print(' '.join(cmd).replace('-nostats -loglevel error', ''))
            sys.exit(1)
        src_duration = clip['out']-clip['in']
        if abs(src_duration-duration) > 1:
            print(clip.get('annotation', clip['item']), 'expected', src_duration, 'got', duration, out)
        add_subtitles(subtitles, clip, position, duration)
        position += duration

txt = output + '.txt'
with open(txt, 'w') as fd:
    fd.write('file ' + '\nfile '.join(files))
cmd = [
    'ffmpeg',
    '-hide_banner',
    '-nostats', '-loglevel', 'error',
    '-y', '-f', 'concat', '-safe', '0', '-i', txt, '-c', 'copy', output
]

run(cmd)
os.unlink(txt)

if subtitles:
    srt = output.replace('.mp4', '.srt')
    with open(srt, 'wb') as fd:
        fd.write(ox.srt.encode(subtitles))

duration = ox.avinfo(output)['duration']
if abs(duration - edit_duration) > 1:
    print('file duration is %d, edit was expected to be %d' % (duration, edit_duration))
print('created:', output)