293 lines
12 KiB
Python
293 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
from queue import SimpleQueue, Empty
|
|
from subprocess import Popen, DEVNULL, PIPE, TimeoutExpired
|
|
import select
|
|
from shutil import rmtree
|
|
from glob import glob
|
|
from os import mkdir, unlink, listdir, environ
|
|
from os.path import basename, abspath, expanduser, isfile, isdir, getsize
|
|
from threading import Thread
|
|
from time import sleep
|
|
from typing import List, NamedTuple, Callable, Union
|
|
from signal import signal, SIGINT
|
|
from sys import exit, stderr
|
|
from shlex import split, join
|
|
from hashlib import md5
|
|
from time import strptime
|
|
from tqdm import tqdm
|
|
import re
|
|
|
|
DEBUG = 'DEBUG' in environ
|
|
|
|
def dprint(*args, **kwargs):
|
|
if DEBUG:
|
|
print(*args, *kwargs)
|
|
|
|
def _popen(args, **kwargs):
|
|
dprint(f'calling subprocess: {args}')
|
|
kwargs['stderr'] = PIPE
|
|
kwargs['stdout'] = DEVNULL
|
|
kwargs['universal_newlines'] = True
|
|
return Popen(args, **kwargs)
|
|
|
|
class Task(NamedTuple):
|
|
input_file: str
|
|
output_file: str
|
|
ffmpeg_args: List[str] = []
|
|
|
|
class FFMPEGProc:
|
|
_duration_re = re.compile(r'.*Duration:\s*-?(?P<time_h>[0-9]+):(?P<time_m>[0-9]+):(?P<time_s>[0-9.]+),')
|
|
_progress_re = re.compile(r'frame=\s*(?P<frame>[0-9]+)\s+fps=\s*(?P<fps>[0-9]+).*time=-?(?P<time_h>[0-9]+):(?P<time_m>[0-9]+):(?P<time_s>[0-9,.]+)\s+.*speed=(?P<speed>[0-9\.]+)x')
|
|
|
|
@staticmethod
|
|
def _match_to_sec(match):
|
|
return int(match.group('time_h'))*3600+int(match.group('time_m'))*60+float(match.group('time_s'))
|
|
|
|
def __init__(self, cmd: Union[list, str], shell=False, stdin=DEVNULL, stdout=DEVNULL, update_callback: Callable[[int,int,float,float,float], None] = None):
|
|
self._cmd = cmd
|
|
self._update_callback = update_callback
|
|
self._should_stop = False
|
|
self._shell = shell
|
|
self._duration = None
|
|
self._stdin = stdin
|
|
self._stdout = stdout
|
|
self.stderr = ''
|
|
|
|
def stop(self):
|
|
self._should_stop = True
|
|
|
|
def _read_stderr_loop(self, stderr_lines: list):
|
|
"""Read stderr in a loop (used on Windows where select.poll is unavailable)."""
|
|
while True:
|
|
line = self._proc.stderr.readline()
|
|
if not line and self._proc.poll() is not None:
|
|
break
|
|
if line:
|
|
stderr_lines.append(line)
|
|
match = self._progress_re.match(line)
|
|
if match and self._update_callback:
|
|
self._update_callback(
|
|
int(match.group('frame')),
|
|
int(match.group('fps')),
|
|
self._match_to_sec(match),
|
|
self._duration,
|
|
float(match.group('speed')),
|
|
)
|
|
elif self._duration is None:
|
|
dm = self._duration_re.match(line)
|
|
if dm:
|
|
self._duration = self._match_to_sec(dm)
|
|
|
|
def run(self):
|
|
self._proc = Popen(self._cmd, shell=self._shell, stderr=PIPE, stdin=self._stdin, stdout=self._stdout, universal_newlines=True)
|
|
if getattr(select, 'poll', None):
|
|
poll = select.poll()
|
|
poll.register(self._proc.stderr)
|
|
while self._proc.poll() is None and not self._should_stop:
|
|
if not poll.poll(1):
|
|
sleep(0.1)
|
|
continue
|
|
sleep(0.001)
|
|
line = self._proc.stderr.readline()
|
|
match = self._progress_re.match(line)
|
|
if not match:
|
|
self.stderr += line
|
|
if match and self._update_callback:
|
|
self._update_callback(
|
|
int(match.group('frame')),
|
|
int(match.group('fps')),
|
|
self._match_to_sec(match),
|
|
self._duration,
|
|
float(match.group('speed')),
|
|
)
|
|
elif self._duration is None:
|
|
match = self._duration_re.match(line)
|
|
if match:
|
|
self._duration = self._match_to_sec(match)
|
|
try:
|
|
_, err = self._proc.communicate(timeout=1)
|
|
self.stderr += err or ''
|
|
except TimeoutExpired:
|
|
pass
|
|
else:
|
|
stderr_lines = []
|
|
reader = Thread(target=self._read_stderr_loop, args=(stderr_lines,), daemon=True)
|
|
reader.start()
|
|
while self._proc.poll() is None and not self._should_stop:
|
|
sleep(0.2)
|
|
reader.join(timeout=2)
|
|
self.stderr = ''.join(stderr_lines)
|
|
return self._proc.returncode
|
|
|
|
class TqdmAbsolute(tqdm):
|
|
def __init__(self, *args, **kwargs):
|
|
kwargs['bar_format'] = '{l_bar}{bar}|{n:.1f}/{total:.1f} [{elapsed}<{remaining}]'
|
|
kwargs['dynamic_ncols'] = True
|
|
if not 'total' in kwargs:
|
|
kwargs['total'] = 99999999
|
|
if not 'leave' in kwargs:
|
|
kwargs['leave'] = False
|
|
super().__init__(*args, **kwargs)
|
|
def update(self, to):
|
|
super().update(to - self.n) # will also set self.n = b * bsize
|
|
|
|
class TaskThread(Thread):
|
|
def __init__(self, host: str, task_queue: SimpleQueue, bar_pos):
|
|
super().__init__()
|
|
self._should_stop = False
|
|
self._host = host
|
|
self._task_queue = task_queue
|
|
self._ffmpeg = None
|
|
self._bar = TqdmAbsolute(desc=host, position=bar_pos)
|
|
self._current_file = None
|
|
|
|
def stop(self):
|
|
self._should_stop = True
|
|
if self._ffmpeg:
|
|
self._ffmpeg.stop()
|
|
|
|
def run(self):
|
|
def upd(frames, fps, time, duration, speed):
|
|
self._bar.total = duration or 999
|
|
self._bar.desc = self._host + ': ' + self._current_file
|
|
self._bar.update(time)
|
|
try:
|
|
while not self._should_stop:
|
|
task = self._task_queue.get(False)
|
|
|
|
self._current_file = basename(task.input_file)
|
|
with open(task.input_file, 'r') as infile, open(task.output_file, 'w') as outfile:
|
|
ffmpeg_cmd = [
|
|
'nice', '-n10', 'ionice', '-c3',
|
|
'ffmpeg', '-f', 'matroska', '-i', 'pipe:',
|
|
*task.ffmpeg_args,
|
|
'-f', 'matroska', 'pipe:'
|
|
]
|
|
if self._host != 'localhost':
|
|
ffmpeg_cmd = ['ssh', self._host, join(ffmpeg_cmd)]
|
|
self._ffmpeg = FFMPEGProc(ffmpeg_cmd, stdin=infile, stdout=outfile, update_callback=upd)
|
|
|
|
ret = self._ffmpeg.run()
|
|
if ret != 0:
|
|
tqdm.write(f'task for {self._current_file} failed on host {self._host}', file=stderr)
|
|
tqdm.write(self._ffmpeg.stderr, file=stderr)
|
|
self._task_queue.put(task)
|
|
except Empty:
|
|
pass
|
|
self._bar.close()
|
|
|
|
def encode(hosts: List[str], input_file: str, output_file: str, segment_seconds: float = 60, remote_args: str = '', concat_args: str = '', tmp_dir: str = None, keep_tmp=False, resume=False, copy_input=False):
|
|
input_file = abspath(expanduser(input_file))
|
|
output_file = abspath(expanduser(output_file))
|
|
tmp_dir = tmp_dir or 'ffmpeg_segments_'+md5(input_file.encode()).hexdigest()
|
|
tmp_in = f'{tmp_dir}/in'
|
|
tmp_out = f'{tmp_dir}/out'
|
|
try:
|
|
mkdir(tmp_dir)
|
|
mkdir(tmp_in)
|
|
mkdir(tmp_out)
|
|
except FileExistsError:
|
|
if not resume:
|
|
raise
|
|
|
|
# skip splitting on resume
|
|
if len(listdir(tmp_in)) == 0 or not resume:
|
|
cv = ['copy'] if copy_input else ['libx264', '-crf', '0', '-preset', 'ultrafast', '-bf', '0']
|
|
with TqdmAbsolute(desc="splitting input file") as bar:
|
|
def upd(frames, fps, time, duration, speed):
|
|
bar.total = duration
|
|
bar.update(time)
|
|
ffmpeg = FFMPEGProc([
|
|
'ffmpeg', '-i', expanduser(input_file),
|
|
'-an', '-sn',
|
|
'-c:v', *cv,
|
|
'-f', 'segment', '-reset_timestamps', '1', '-segment_time', str(segment_seconds) + 's',
|
|
tmp_in + '/%08d.mkv'
|
|
],
|
|
update_callback=upd
|
|
)
|
|
ret = ffmpeg.run()
|
|
if ret != 0:
|
|
tqdm.write(ffmpeg.stderr, file=stderr)
|
|
return
|
|
|
|
task_queue = SimpleQueue()
|
|
for f in sorted(glob(tmp_in+'/*')):
|
|
output_segment = tmp_out+f'/{basename(f)}'
|
|
# skip already encoded segments
|
|
if not isfile(output_segment):
|
|
task_queue.put(Task(f, output_segment, split(remote_args)))
|
|
|
|
threads = [TaskThread(host, task_queue, pos) for pos,host in enumerate(hosts,0)]
|
|
|
|
def sigint(sig, stack):
|
|
print('Got SIGINT, stopping...')
|
|
for thread in threads:
|
|
thread.stop()
|
|
for thread in threads:
|
|
thread.join()
|
|
exit(1)
|
|
|
|
signal(SIGINT, sigint)
|
|
|
|
for thread in threads:
|
|
thread.start()
|
|
for thread in threads:
|
|
thread.join()
|
|
|
|
with open('output_segments.txt', 'w') as f:
|
|
f.write('\n'.join([f"file '{file}'" for file in sorted(glob(tmp_out+'/*'))]))
|
|
|
|
with TqdmAbsolute(desc='concatenating output segments') as bar:
|
|
def upd(frames, fps, time, duration, speed):
|
|
bar.total = duration
|
|
bar.update(time)
|
|
ffmpeg = FFMPEGProc([
|
|
'ffmpeg', '-i', input_file,
|
|
'-f', 'concat', '-safe', '0', '-i', 'output_segments.txt',
|
|
'-map_metadata', '0:g',
|
|
'-map', '1:v',
|
|
'-map', '0:a?',
|
|
'-map', '0:s?',
|
|
'-c:v', 'copy',
|
|
'-c:s', 'copy',
|
|
*split(concat_args),
|
|
'-y', output_file
|
|
],
|
|
update_callback=upd
|
|
)
|
|
if ffmpeg.run() != 0:
|
|
tqdm.write(ffmpeg.stderr, file=stderr)
|
|
return
|
|
unlink('output_segments.txt')
|
|
|
|
if not keep_tmp:
|
|
rmtree(tmp_dir)
|
|
|
|
if __name__ == '__main__':
|
|
import argparse
|
|
parser = argparse.ArgumentParser(description='Splits a file into segments and processes them on multiple hosts in parallel using ffmpeg over SSH.')
|
|
parser.add_argument('input_file', help='File to encode.')
|
|
parser.add_argument('output_file', help='Path to encoded output file.')
|
|
parser.add_argument('remote_args', help='Arguments to pass to the remote ffmpeg instances. For example: "-c:v libx264 -crf 23 -preset fast"')
|
|
parser.add_argument('concat_args', default='', help='Arguments to pass to the local ffmpeg concatenating the processed video segments and muxing it with the original audio/subs/metadata. Mainly useful for audio encoding options, or "-an" to get rid of it.')
|
|
parser.add_argument('-s', '--segment-length', type=float, default=10, help='Segment length in seconds.')
|
|
parser.add_argument('-H', '--host', action='append', help='SSH hostname(s) to encode on. Use "localhost" to include the machine you\'re running this from. Can include username.', required=True)
|
|
parser.add_argument('-k', '--keep-tmp', action='store_true', help='Keep temporary segment files instead of deleting them on successful exit.')
|
|
parser.add_argument('-r', '--resume', action='store_true', help='Don\'t split the input file again, keep existing segments and only process the missing ones.')
|
|
parser.add_argument('-t', '--tmp-dir', default=None, help='Directory to use for temporary files. Should not already exist and will be deleted afterwards.')
|
|
parser.add_argument('-c', '--copy-input', action='store_true', help='Don\'t (losslessly) re-encode input while segmenting. Only use this if your input segments frame-perfectly with "-c:v copy" (i.e. it has no B-frames)')
|
|
args = parser.parse_args()
|
|
encode(
|
|
args.host,
|
|
args.input_file,
|
|
args.output_file,
|
|
segment_seconds=args.segment_length,
|
|
remote_args=args.remote_args,
|
|
concat_args=args.concat_args,
|
|
tmp_dir=args.tmp_dir,
|
|
keep_tmp=args.keep_tmp,
|
|
resume=args.resume,
|
|
copy_input=args.copy_input
|
|
)
|