bdn_render.py

#!/usr/local/bin/python3

# -*- coding: utf-8 -*-

'''Render ASS to PGS/Blu-ray Subtitles'''

import os
import sys
from colorama import init, Fore, Style
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
import shutil
import tempfile
import subprocess
import pysubs2

ALLOWED_VF  = ('1080p', '1080i', '720p', '576i', '480p', '480i')
ALLOWED_FPS = ('23.976', '24', '25', '29.97', '50', '59.94')

def winepath(path):
    return subprocess.check_output(['winepath', '-w', path]).decode().replace('\n', '')

def main():
    init(autoreset=True)
    parser = ArgumentParser(description=__doc__, formatter_class=ArgumentDefaultsHelpFormatter)
    parser.add_argument('input', help='Input ass')
    parser.add_argument('-o', '--output', required=True, help='Output dir')
    parser.add_argument('--fps', default='23.976', help='Target framerate')
    parser.add_argument('--video-format', default='1080p', help='Output subtitles format')
    parser.add_argument('--lang', default='rus', help='Output language code for subtitles')
    parser.add_argument('--sup', action='store_true', help='Addionatinly output PGS/SUP file')
    args = parser.parse_args()

    if os.path.isfile(args.input):
        filename = os.path.basename(args.input)
        if os.path.splitext(filename)[-1].lower() == '.ass':
            ass = args.input
        else:
            sys.exit(f'{Fore.RED}{Style.BRIGHT}Not an ass file')
    else:
        sys.exit(f'{Fore.RED}{Style.BRIGHT}ASS file not found')

    if not os.path.isdir(args.output):
        sys.exit(f'{Fore.RED}{Style.BRIGHT}Output dir not specified')

    if args.video_format not in ALLOWED_VF:
        sys.exit(f'{Fore.RED}{Style.BRIGHT}{args.video_format} not allowed')

    if args.fps not in ALLOWED_FPS:
        sys.exit(f'{Fore.RED}{Style.BRIGHT}{args.fps} not allowed')

    subs = pysubs2.load(ass)
    subs.sort()
    duration = pysubs2.time.ms_to_frames(subs[-1].end, float(args.fps)) + 20

    if duration < 10:
        sys.exit(f'{Fore.RED}{Style.BRIGHT}Duration of subtitles is lower than 10 frames. Abort')

    basename   = os.path.splitext(os.path.basename(ass))[0]
    output_dir = os.path.join(args.output, os.path.splitext(os.path.basename(ass))[0])
    output_xml = os.path.join(output_dir, basename) + '.xml'
    output_sup = os.path.join(args.output, basename) + '.sup'

    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    if os.path.exists(output_sup):
            os.remove(output_sup)
    os.mkdir(output_dir)

    avs_temp = tempfile.mkstemp(suffix='.avs', dir=output_dir)[-1]

    with open(avs_temp, 'w') as avs:
        print('Loadplugin("C:\\Program Files (x86)\\AviSynth+\\plugins+\\VSFilter.dll")', file=avs)
        print(f'MaskSub(\"{winepath(args.input)}\", 1920, 1080, {args.fps}, {duration})', file=avs)

    render = ['wine', 'avs2bdnxml', '--trackname', basename,
            '--language', f'{args.lang}', '--video-format', f'{args.video_format}', '--fps', f'{args.fps}',
            '--autocrop', '1', '--palette', '1', '--buffer-opt', '0', '--min-split', '3',
            '--ugly', '0', '--even-y', '0', '--null-xml', '0', '--stricter', '0',
            '--output', output_xml
        ]

    if args.sup:
        if os.path.exists(output_sup):
            os.remove(output_sup)
        render.extend(('--output', output_sup))
    render.append(avs_temp)

    try:
        subprocess.run(render, check=True)
    except subprocess.CalledProcessError:
        try:
            shutil.rmtree(output_dir)
            os.remove(output_sup)
        except:
            pass
        sys.exit(f'{Fore.RED}{Style.BRIGHT}Error occurred during rendering')
    else:
        print(f'{Fore.GREEN}{Style.BRIGHT}Done')

    try:
        os.remove(avs_temp)
    except:
        pass

if __name__ == "__main__":
    main()