193 lines
6.3 KiB
Python
193 lines
6.3 KiB
Python
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
|