From 8337d34b8248134da99164e8a294f66796106e24 Mon Sep 17 00:00:00 2001 From: Karl Kroening Date: Thu, 6 Jul 2017 03:35:03 -0600 Subject: [PATCH] #17: add proper handling of `split` operator --- ffmpeg/_run.py | 38 ++++++++++++++++++++++--------- ffmpeg/nodes.py | 26 ++++++++++----------- ffmpeg/tests/test_ffmpeg.py | 45 ++++++++++++++++++++++--------------- 3 files changed, 67 insertions(+), 42 deletions(-) diff --git a/ffmpeg/_run.py b/ffmpeg/_run.py index 81acb47..eb00c90 100644 --- a/ffmpeg/_run.py +++ b/ffmpeg/_run.py @@ -1,6 +1,6 @@ from __future__ import unicode_literals -from .dag import topo_sort +from .dag import get_outgoing_edges, topo_sort from functools import reduce from past.builtins import basestring import copy @@ -53,16 +53,31 @@ def _get_input_args(input_node): return args -def _get_filter_spec(i, node, stream_name_map): - stream_name = _get_stream_name('v{}'.format(i)) - stream_name_map[node] = stream_name - inputs = [stream_name_map[edge.upstream_node] for edge in node.incoming_edges] - filter_spec = '{}{}{}'.format(''.join(inputs), node._get_filter(), stream_name) +def _get_filter_spec(node, outgoing_edge_map, stream_name_map): + incoming_edges = node.incoming_edges + outgoing_edges = get_outgoing_edges(node, outgoing_edge_map) + inputs = [stream_name_map[edge.upstream_node, edge.upstream_label] for edge in incoming_edges] + outputs = [stream_name_map[edge.upstream_node, edge.upstream_label] for edge in outgoing_edges] + filter_spec = '{}{}{}'.format(''.join(inputs), node._get_filter(), ''.join(outputs)) return filter_spec -def _get_filter_arg(filter_nodes, stream_name_map): - filter_specs = [_get_filter_spec(i, node, stream_name_map) for i, node in enumerate(filter_nodes)] +def _allocate_filter_stream_names(filter_nodes, outgoing_edge_maps, stream_name_map): + stream_count = 0 + for upstream_node in filter_nodes: + outgoing_edge_map = outgoing_edge_maps[upstream_node] + for upstream_label, downstreams in outgoing_edge_map.items(): + if len(downstreams) > 1: + # TODO: automatically insert `splits` ahead of time via graph transformation. + raise ValueError('Encountered {} with multiple outgoing edges with same upstream label {!r}; a ' + '`split` filter is probably required'.format(upstream_node, upstream_label)) + stream_name_map[upstream_node, upstream_label] = _get_stream_name('s{}'.format(stream_count)) + stream_count += 1 + + +def _get_filter_arg(filter_nodes, outgoing_edge_maps, stream_name_map): + _allocate_filter_stream_names(filter_nodes, outgoing_edge_maps, stream_name_map) + filter_specs = [_get_filter_spec(node, outgoing_edge_maps[node], stream_name_map) for node in filter_nodes] return ';'.join(filter_specs) @@ -78,7 +93,8 @@ def _get_output_args(node, stream_name_map): raise ValueError('Unsupported output node: {}'.format(node)) args = [] assert len(node.incoming_edges) == 1 - stream_name = stream_name_map[node.incoming_edges[0].upstream_node] + edge = node.incoming_edges[0] + stream_name = stream_name_map[edge.upstream_node, edge.upstream_label] if stream_name != '[0]': args += ['-map', stream_name] kwargs = copy.copy(node.kwargs) @@ -104,8 +120,8 @@ def get_args(stream): 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)] - stream_name_map = {node: _get_stream_name(i) for i, node in enumerate(input_nodes)} - filter_arg = _get_filter_arg(filter_nodes, stream_name_map) + 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) args += reduce(operator.add, [_get_input_args(node) for node in input_nodes]) if filter_arg: args += ['-filter_complex', filter_arg] diff --git a/ffmpeg/nodes.py b/ffmpeg/nodes.py index bc1d55d..17b4bc9 100644 --- a/ffmpeg/nodes.py +++ b/ffmpeg/nodes.py @@ -39,6 +39,18 @@ class Stream(object): return out +def get_stream_map(stream_spec): + if stream_spec is None: + stream_map = {} + elif isinstance(stream_spec, Stream): + stream_map = {None: stream_spec} + elif isinstance(stream_spec, (list, tuple)): + stream_map = dict(enumerate(stream_spec)) + elif isinstance(stream_spec, dict): + stream_map = stream_spec + return stream_map + + class Node(KwargReprNode): """Node base""" @classmethod @@ -55,18 +67,6 @@ class Node(KwargReprNode): raise TypeError('Expected incoming stream(s) to be of one of the following types: {}; got {}' .format(_get_types_str(incoming_stream_types), type(stream))) - @classmethod - def __get_stream_map(cls, stream_spec): - if stream_spec is None: - stream_map = {} - elif isinstance(stream_spec, Stream): - stream_map = {None: stream_spec} - elif isinstance(stream_spec, (list, tuple)): - stream_map = dict(enumerate(stream_spec)) - elif isinstance(stream_spec, dict): - stream_map = stream_spec - return stream_map - @classmethod def __get_incoming_edge_map(cls, stream_map): incoming_edge_map = {} @@ -76,7 +76,7 @@ class Node(KwargReprNode): def __init__(self, stream_spec, name, incoming_stream_types, outgoing_stream_type, min_inputs, max_inputs, args, kwargs): - stream_map = self.__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_types(stream_map, incoming_stream_types) incoming_edge_map = self.__get_incoming_edge_map(stream_map) diff --git a/ffmpeg/tests/test_ffmpeg.py b/ffmpeg/tests/test_ffmpeg.py index cd52837..32b8be5 100644 --- a/ffmpeg/tests/test_ffmpeg.py +++ b/ffmpeg/tests/test_ffmpeg.py @@ -93,12 +93,19 @@ def test_get_args_simple(): def _get_complex_filter_example(): - in_file = ffmpeg.input(TEST_INPUT_FILE) + split = (ffmpeg + .input(TEST_INPUT_FILE) + .vflip() + .split() + ) + split0 = split[0] + split1 = split[1] + overlay_file = ffmpeg.input(TEST_OVERLAY_FILE) return (ffmpeg .concat( - in_file.trim(start_frame=10, end_frame=20), - in_file.trim(start_frame=30, end_frame=40), + split0.trim(start_frame=10, end_frame=20), + split1.trim(start_frame=30, end_frame=40), ) .overlay(overlay_file.hflip()) .drawbox(50, 50, 120, 120, color='red', thickness=5) @@ -110,21 +117,23 @@ def _get_complex_filter_example(): def test_get_args_complex_filter(): out = _get_complex_filter_example() args = ffmpeg.get_args(out) - assert args == [ - '-i', TEST_INPUT_FILE, + assert args == ['-i', TEST_INPUT_FILE, '-i', TEST_OVERLAY_FILE, '-filter_complex', - '[0]trim=end_frame=20:start_frame=10[v0];' \ - '[0]trim=end_frame=40:start_frame=30[v1];' \ - '[v0][v1]concat=n=2[v2];' \ - '[1]hflip[v3];' \ - '[v2][v3]overlay=eof_action=repeat[v4];' \ - '[v4]drawbox=50:50:120:120:red:t=5[v5]', - '-map', '[v5]', os.path.join(SAMPLE_DATA_DIR, 'dummy2.mp4'), + '[0]vflip[s0];' \ + '[s0]split[s1][s2];' \ + '[s1]trim=end_frame=20:start_frame=10[s3];' \ + '[s2]trim=end_frame=40:start_frame=30[s4];' \ + '[s3][s4]concat=n=2[s5];' \ + '[1]hflip[s6];' \ + '[s5][s6]overlay=eof_action=repeat[s7];' \ + '[s7]drawbox=50:50:120:120:red:t=5[s8]', + '-map', '[s8]', os.path.join(SAMPLE_DATA_DIR, 'dummy2.mp4'), '-y' ] + #def test_version(): # subprocess.check_call(['ffmpeg', '-version']) @@ -156,8 +165,8 @@ def test_custom_filter(): node = ffmpeg.output(node, 'dummy2.mp4') assert node.get_args() == [ '-i', 'dummy.mp4', - '-filter_complex', '[0]custom_filter=a:b:kwarg1=c[v0]', - '-map', '[v0]', + '-filter_complex', '[0]custom_filter=a:b:kwarg1=c[s0]', + '-map', '[s0]', 'dummy2.mp4' ] @@ -170,8 +179,8 @@ def test_custom_filter_fluent(): ) assert node.get_args() == [ '-i', 'dummy.mp4', - '-filter_complex', '[0]custom_filter=a:b:kwarg1=c[v0]', - '-map', '[v0]', + '-filter_complex', '[0]custom_filter=a:b:kwarg1=c[s0]', + '-map', '[s0]', 'dummy2.mp4' ] @@ -197,8 +206,8 @@ def test_pipe(): '-pixel_format', 'rgb24', '-i', 'pipe:0', '-filter_complex', - '[0]trim=start_frame=2[v0]', - '-map', '[v0]', + '[0]trim=start_frame=2[s0]', + '-map', '[s0]', '-f', 'rawvideo', 'pipe:1' ]