#17: fix merge_outputs; allow stream_spec in get_args+run

This commit is contained in:
Karl Kroening 2017-07-09 15:50:51 -06:00
parent c6e2f05e5b
commit 5d78a2595d
5 changed files with 96 additions and 37 deletions

2
.gitignore vendored
View File

@ -2,6 +2,6 @@
.eggs .eggs
.tox/ .tox/
dist/ dist/
ffmpeg/tests/sample_data/dummy2.mp4 ffmpeg/tests/sample_data/out*.mp4
ffmpeg_python.egg-info/ ffmpeg_python.egg-info/
venv* venv*

View File

@ -13,11 +13,12 @@ from ._ffmpeg import (
overwrite_output, overwrite_output,
) )
from .nodes import ( from .nodes import (
get_stream_spec_nodes,
FilterNode,
GlobalNode, GlobalNode,
InputNode, InputNode,
OutputNode, OutputNode,
output_operator, output_operator,
Stream,
) )
@ -108,18 +109,16 @@ def _get_output_args(node, stream_name_map):
@output_operator() @output_operator()
def get_args(stream): def get_args(stream_spec, overwrite_output=False):
"""Get command-line arguments for ffmpeg.""" """Get command-line arguments for ffmpeg."""
if not isinstance(stream, Stream): nodes = get_stream_spec_nodes(stream_spec)
raise TypeError('Expected Stream; got {}'.format(type(stream)))
args = [] args = []
# TODO: group nodes together, e.g. `-i somefile -r somerate`. # TODO: group nodes together, e.g. `-i somefile -r somerate`.
sorted_nodes, outgoing_edge_maps = topo_sort([stream.node]) sorted_nodes, outgoing_edge_maps = topo_sort(nodes)
input_nodes = [node for node in sorted_nodes if isinstance(node, InputNode)] input_nodes = [node for node in sorted_nodes if isinstance(node, InputNode)]
output_nodes = [node for node in sorted_nodes if isinstance(node, OutputNode) and not output_nodes = [node for node in sorted_nodes if isinstance(node, OutputNode)]
isinstance(node, GlobalNode)]
global_nodes = [node for node in sorted_nodes if isinstance(node, GlobalNode)] global_nodes = [node for node in sorted_nodes if isinstance(node, GlobalNode)]
filter_nodes = [node for node in sorted_nodes if node not in (input_nodes + output_nodes + global_nodes)] filter_nodes = [node for node in sorted_nodes if isinstance(node, FilterNode)]
stream_name_map = {(node, None): _get_stream_name(i) for i, node in enumerate(input_nodes)} stream_name_map = {(node, None): _get_stream_name(i) for i, node in enumerate(input_nodes)}
filter_arg = _get_filter_arg(filter_nodes, outgoing_edge_maps, stream_name_map) filter_arg = _get_filter_arg(filter_nodes, outgoing_edge_maps, stream_name_map)
args += reduce(operator.add, [_get_input_args(node) for node in input_nodes]) args += reduce(operator.add, [_get_input_args(node) for node in input_nodes])
@ -127,17 +126,23 @@ def get_args(stream):
args += ['-filter_complex', filter_arg] args += ['-filter_complex', filter_arg]
args += reduce(operator.add, [_get_output_args(node, stream_name_map) for node in output_nodes]) args += reduce(operator.add, [_get_output_args(node, stream_name_map) for node in output_nodes])
args += reduce(operator.add, [_get_global_args(node) for node in global_nodes], []) args += reduce(operator.add, [_get_global_args(node) for node in global_nodes], [])
if overwrite_output:
args += ['-y']
return args return args
@output_operator() @output_operator()
def run(node, cmd='ffmpeg'): def run(stream_spec, cmd='ffmpeg', **kwargs):
"""Run ffmpeg on node graph.""" """Run ffmpeg on node graph.
Args:
**kwargs: keyword-arguments passed to ``get_args()`` (e.g. ``overwrite_output=True``).
"""
if isinstance(cmd, basestring): if isinstance(cmd, basestring):
cmd = [cmd] cmd = [cmd]
elif type(cmd) != list: elif type(cmd) != list:
cmd = list(cmd) cmd = list(cmd)
args = cmd + node.get_args() args = cmd + get_args(stream_spec, **kwargs)
_subprocess.check_call(args) _subprocess.check_call(args)

View File

@ -52,6 +52,20 @@ def get_stream_map(stream_spec):
return stream_map return stream_map
def get_stream_map_nodes(stream_map):
nodes = []
for stream in stream_map.values():
if not isinstance(stream, Stream):
raise TypeError('Expected Stream; got {}'.format(type(stream)))
nodes.append(stream.node)
return nodes
def get_stream_spec_nodes(stream_spec):
stream_map = get_stream_map(stream_spec)
return get_stream_map_nodes(stream_map)
class Node(KwargReprNode): class Node(KwargReprNode):
"""Node base""" """Node base"""
@classmethod @classmethod
@ -75,8 +89,8 @@ class Node(KwargReprNode):
incoming_edge_map[downstream_label] = (upstream.node, upstream.label) incoming_edge_map[downstream_label] = (upstream.node, upstream.label)
return incoming_edge_map return incoming_edge_map
def __init__(self, stream_spec, name, incoming_stream_types, outgoing_stream_type, min_inputs, max_inputs, args, def __init__(self, stream_spec, name, incoming_stream_types, outgoing_stream_type, min_inputs, max_inputs, args=[],
kwargs): kwargs={}):
stream_map = get_stream_map(stream_spec) stream_map = get_stream_map(stream_spec)
self.__check_input_len(stream_map, min_inputs, max_inputs) self.__check_input_len(stream_map, min_inputs, max_inputs)
self.__check_input_types(stream_map, incoming_stream_types) self.__check_input_types(stream_map, incoming_stream_types)
@ -164,13 +178,13 @@ class OutputNode(Node):
class OutputStream(Stream): class OutputStream(Stream):
def __init__(self, upstream_node, upstream_label): def __init__(self, upstream_node, upstream_label):
super(OutputStream, self).__init__(upstream_node, upstream_label, {OutputNode, GlobalNode}) super(OutputStream, self).__init__(upstream_node, upstream_label, {OutputNode, GlobalNode, MergeOutputsNode})
class MergeOutputsNode(Node): class MergeOutputsNode(Node):
def __init__(self, stream, name): def __init__(self, streams, name):
super(MergeOutputsNode, self).__init__( super(MergeOutputsNode, self).__init__(
stream_spec=None, stream_spec=streams,
name=name, name=name,
incoming_stream_types={OutputStream}, incoming_stream_types={OutputStream},
outgoing_stream_type=OutputStream, outgoing_stream_type=OutputStream,

View File

@ -8,9 +8,10 @@ import random
TEST_DIR = os.path.dirname(__file__) TEST_DIR = os.path.dirname(__file__)
SAMPLE_DATA_DIR = os.path.join(TEST_DIR, 'sample_data') SAMPLE_DATA_DIR = os.path.join(TEST_DIR, 'sample_data')
TEST_INPUT_FILE = os.path.join(SAMPLE_DATA_DIR, 'dummy.mp4') TEST_INPUT_FILE1 = os.path.join(SAMPLE_DATA_DIR, 'in1.mp4')
TEST_OVERLAY_FILE = os.path.join(SAMPLE_DATA_DIR, 'overlay.png') TEST_OVERLAY_FILE = os.path.join(SAMPLE_DATA_DIR, 'overlay.png')
TEST_OUTPUT_FILE = os.path.join(SAMPLE_DATA_DIR, 'dummy2.mp4') TEST_OUTPUT_FILE1 = os.path.join(SAMPLE_DATA_DIR, 'out1.mp4')
TEST_OUTPUT_FILE2 = os.path.join(SAMPLE_DATA_DIR, 'out2.mp4')
subprocess.check_call(['ffmpeg', '-version']) subprocess.check_call(['ffmpeg', '-version'])
@ -94,7 +95,7 @@ def test_get_args_simple():
def _get_complex_filter_example(): def _get_complex_filter_example():
split = (ffmpeg split = (ffmpeg
.input(TEST_INPUT_FILE) .input(TEST_INPUT_FILE1)
.vflip() .vflip()
.split() .split()
) )
@ -109,7 +110,7 @@ def _get_complex_filter_example():
) )
.overlay(overlay_file.hflip()) .overlay(overlay_file.hflip())
.drawbox(50, 50, 120, 120, color='red', thickness=5) .drawbox(50, 50, 120, 120, color='red', thickness=5)
.output(TEST_OUTPUT_FILE) .output(TEST_OUTPUT_FILE1)
.overwrite_output() .overwrite_output()
) )
@ -117,7 +118,7 @@ def _get_complex_filter_example():
def test_get_args_complex_filter(): def test_get_args_complex_filter():
out = _get_complex_filter_example() out = _get_complex_filter_example()
args = ffmpeg.get_args(out) args = ffmpeg.get_args(out)
assert args == ['-i', TEST_INPUT_FILE, assert args == ['-i', TEST_INPUT_FILE1,
'-i', TEST_OVERLAY_FILE, '-i', TEST_OVERLAY_FILE,
'-filter_complex', '-filter_complex',
'[0]vflip[s0];' \ '[0]vflip[s0];' \
@ -128,7 +129,7 @@ def test_get_args_complex_filter():
'[1]hflip[s6];' \ '[1]hflip[s6];' \
'[s5][s6]overlay=eof_action=repeat[s7];' \ '[s5][s6]overlay=eof_action=repeat[s7];' \
'[s7]drawbox=50:50:120:120:red:t=5[s8]', '[s7]drawbox=50:50:120:120:red:t=5[s8]',
'-map', '[s8]', os.path.join(SAMPLE_DATA_DIR, 'dummy2.mp4'), '-map', '[s8]', TEST_OUTPUT_FILE1,
'-y' '-y'
] ]
@ -139,31 +140,38 @@ def test_get_args_complex_filter():
def test_run(): def test_run():
node = _get_complex_filter_example() stream = _get_complex_filter_example()
ffmpeg.run(node) ffmpeg.run(stream)
def test_run_multi_output():
in_ = ffmpeg.input(TEST_INPUT_FILE1)
out1 = in_.output(TEST_OUTPUT_FILE1)
out2 = in_.output(TEST_OUTPUT_FILE2)
ffmpeg.run([out1, out2], overwrite_output=True)
def test_run_dummy_cmd(): def test_run_dummy_cmd():
node = _get_complex_filter_example() stream = _get_complex_filter_example()
ffmpeg.run(node, cmd='true') ffmpeg.run(stream, cmd='true')
def test_run_dummy_cmd_list(): def test_run_dummy_cmd_list():
node = _get_complex_filter_example() stream = _get_complex_filter_example()
ffmpeg.run(node, cmd=['true', 'ignored']) ffmpeg.run(stream, cmd=['true', 'ignored'])
def test_run_failing_cmd(): def test_run_failing_cmd():
node = _get_complex_filter_example() stream = _get_complex_filter_example()
with pytest.raises(subprocess.CalledProcessError): with pytest.raises(subprocess.CalledProcessError):
ffmpeg.run(node, cmd='false') ffmpeg.run(stream, cmd='false')
def test_custom_filter(): def test_custom_filter():
node = ffmpeg.input('dummy.mp4') stream = ffmpeg.input('dummy.mp4')
node = ffmpeg.filter_(node, 'custom_filter', 'a', 'b', kwarg1='c') stream = ffmpeg.filter_(stream, 'custom_filter', 'a', 'b', kwarg1='c')
node = ffmpeg.output(node, 'dummy2.mp4') stream = ffmpeg.output(stream, 'dummy2.mp4')
assert node.get_args() == [ assert stream.get_args() == [
'-i', 'dummy.mp4', '-i', 'dummy.mp4',
'-filter_complex', '[0]custom_filter=a:b:kwarg1=c[s0]', '-filter_complex', '[0]custom_filter=a:b:kwarg1=c[s0]',
'-map', '[s0]', '-map', '[s0]',
@ -172,12 +180,12 @@ def test_custom_filter():
def test_custom_filter_fluent(): def test_custom_filter_fluent():
node = (ffmpeg stream = (ffmpeg
.input('dummy.mp4') .input('dummy.mp4')
.filter_('custom_filter', 'a', 'b', kwarg1='c') .filter_('custom_filter', 'a', 'b', kwarg1='c')
.output('dummy2.mp4') .output('dummy2.mp4')
) )
assert node.get_args() == [ assert stream.get_args() == [
'-i', 'dummy.mp4', '-i', 'dummy.mp4',
'-filter_complex', '[0]custom_filter=a:b:kwarg1=c[s0]', '-filter_complex', '[0]custom_filter=a:b:kwarg1=c[s0]',
'-map', '[s0]', '-map', '[s0]',
@ -185,6 +193,38 @@ def test_custom_filter_fluent():
] ]
def test_merge_outputs():
in_ = ffmpeg.input('in.mp4')
out1 = in_.output('out1.mp4')
out2 = in_.output('out2.mp4')
assert ffmpeg.merge_outputs(out1, out2).get_args() == [
'-i', 'in.mp4', 'out1.mp4', 'out2.mp4'
]
assert ffmpeg.get_args([out1, out2]) == [
'-i', 'in.mp4', 'out2.mp4', 'out1.mp4'
]
def test_multi_passthrough():
out1 = ffmpeg.input('in1.mp4').output('out1.mp4')
out2 = ffmpeg.input('in2.mp4').output('out2.mp4')
out = ffmpeg.merge_outputs(out1, out2)
assert ffmpeg.get_args(out) == [
'-i', 'in1.mp4',
'-i', 'in2.mp4',
'out1.mp4',
'-map', '[1]', # FIXME: this should not be here (see #23)
'out2.mp4'
]
assert ffmpeg.get_args([out1, out2]) == [
'-i', 'in2.mp4',
'-i', 'in1.mp4',
'out2.mp4',
'-map', '[1]', # FIXME: this should not be here (see #23)
'out1.mp4'
]
def test_pipe(): def test_pipe():
width = 32 width = 32
height = 32 height = 32