import math import os from pathlib import Path import shutil import subprocess import tempfile def with_fifo(func): def f(*args, **kwargs): with tempfile.TemporaryDirectory() as tmpdir: fifo_path = str(Path(tmpdir) / 'rip.fifo') try: os.mkfifo(fifo_path) except FileExistsError: pass kwargs['fifo_path'] = fifo_path return func(*args, **kwargs) return f def convert_chapters_to_mkvmerge_format(chapters): """Convert a list of chapter timestamps into a string (of multiple lines) as accepted by mkvmerge""" # determine how many zeroes we have as a prefix number_of_zeros = 0 if len(chapters) > 1: number_of_zeros = int(math.log10(len(chapters) - 1)) chapter_template = 'CHAPTER{{:0{}d}}'.format(number_of_zeros) lines = [] for i, timestamp in enumerate(chapters): chapter_str = chapter_template.format(i) # if we start with e.g. 0:21, we need another 0 to be compatible if len(timestamp.split(':')[0]) == 1: timestamp = '0' + timestamp # if we have a higher resolution than 1 millisecond, we need to cut # some chars if len(timestamp.split('.')[-1]) > 3: strip_length = len(timestamp.split('.')[-1]) - 3 timestamp = timestamp[:-strip_length] lines.append(f"{chapter_str}={timestamp}") lines.append(f"{chapter_str}NAME=") return os.linesep.join(lines) def set_chapters(mkv_file, chapters): """Set the chapters in the given MKV file `chapters` is a list of chapter start times. """ if not Path(mkv_file).exists(): raise RuntimeError('MKV file does not exist.') if not shutil.which('mkvpropedit'): raise RuntimeError('Cannot find "mkvpropedit". Please install the "mkvtoolnix" package.') with tempfile.NamedTemporaryFile() as f: mkvmerge_chapters = convert_chapters_to_mkvmerge_format(chapters) f.write(mkvmerge_chapters.encode('utf-8')) f.flush() mkvpropedit_cmd = ['mkvpropedit', mkv_file, '--chapters', f.name] subprocess.check_call(mkvpropedit_cmd) def get_parameterized_ffmpeg_cmd(media_type, quality, extra_opts, use_gpu, audio_map, subtitle_map, pre_read=1024, video_filters=None, scale=None): """Build an ffmpeg command based on the given parameters Returns an function taking input file and output file as arguments to build the final command. """ if scale is not None: if video_filters is None: video_filters = [] video_filters.append(f"scale={scale}") ffmpeg_cmd = [] ffmpeg_cmd.append(['ffmpeg', '-probesize', '{}M'.format(pre_read), '-analyzeduration', '{}M'.format(pre_read), '-hwaccel', 'auto', '-i']) ffmpeg_cmd.append(extra_opts) ffmpeg_cmd[-1].extend(['-map', '0:v:0']) if '-c:a' not in ffmpeg_cmd[-1]: ffmpeg_cmd[-1].extend(['-c:a', 'copy']) ffmpeg_cmd[-1].extend(['-c:s', 'copy']) if use_gpu: ffmpeg_cmd[-1] += [ '-c:v', 'h264_nvenc', '-rc', 'vbr_hq', '-rc-lookahead', '20', '-cq', '{}'.format(quality), # '-maxrate:v', '3M', # '-b:v', '8M', '-maxrate:v', '10M', '-preset:v', 'slow'] else: ffmpeg_cmd[-1] += [ '-c:v', 'h264', '-crf', '{}'.format(quality), '-preset:v', 'slower', '-x264-params', 'opencl=true', # '-x264-params', 'keyint=240:min-keyint=20', ] ffmpeg_cmd[-1] += [ '-pix_fmt', 'yuv420p', '-movflags', '+faststart', # '-profile:v', 'baseline', # '-level', '3.0', ] if video_filters: ffmpeg_cmd[-1] += [ '-vf', ','.join(video_filters) ] if audio_map: for i, (channel, lang) in enumerate(audio_map): if channel < 0x80: map_target = '0:a:{}'.format(channel) else: map_target = 'i:0x{:x}'.format(channel) ffmpeg_cmd[-1].extend([ '-map', map_target, '-metadata:s:a:{}'.format(i), 'language={}'.format(lang) ]) else: ffmpeg_cmd[-1].extend(['-map', '0:a:0', '-metadata:s:a:0', 'language=de']) if subtitle_map: for i, (channel, lang) in enumerate(subtitle_map): if channel < 0x20: map_target = '0:s:{}'.format(channel) else: map_target = 'i:0x{:x}'.format(channel) ffmpeg_cmd[-1].extend([ '-map', map_target, '-metadata:s:s:{}'.format(i), 'language={}'.format(lang) ]) # TODO we are probably able to do forced subtitles. we probably want that # we can set a default stream. make this available somehow # -disposition:a:1 default if media_type and not use_gpu: ffmpeg_cmd[-1].extend(['-tune:v', media_type]) def parameterized(output_file, input_file): return ffmpeg_cmd[0] + [input_file] + ffmpeg_cmd[1] + [output_file] return parameterized def get_ffmpeg_cmd(output_file, input_path, media_type, quality, extra_opts, use_gpu, audio_map, subtitle_map, video_filters=None, scale=None): cmd = get_parameterized_ffmpeg_cmd(media_type, quality, extra_opts, use_gpu, audio_map, subtitle_map, video_filters=video_filters, scale=scale) cmd = cmd(output_file=output_file, input_file=input_path) print(cmd) return cmd @with_fifo def rip_title(title, output_dir, mpv_cmd, ffmpeg_cmd, fifo_path, output_file_suffix='.mkv'): filename = f"{title}{output_file_suffix}" out_file = output_dir / filename mpv_cmd = mpv_cmd(title, fifo_path) # we need to redirect stdout, because otherwise we endup botching the terminal input somehow mpv_proc = subprocess.Popen(mpv_cmd, stdout=subprocess.DEVNULL) ffmpeg_cmd = ffmpeg_cmd(output_file=out_file, input_file=fifo_path) print(ffmpeg_cmd) subprocess.check_call(ffmpeg_cmd) # mpv should exit if ffmpeg is gone, because the FIFO got closed mpv_proc.communicate() return out_file