diff --git a/ffmpeg/__init__.py b/ffmpeg/__init__.py index 9f1bed7..154d8b6 100644 --- a/ffmpeg/__init__.py +++ b/ffmpeg/__init__.py @@ -1,7 +1,8 @@ from __future__ import unicode_literals + from . import _filters, _ffmpeg, _run from ._filters import * from ._ffmpeg import * from ._run import * -from ._view import * -__all__ = _filters.__all__ + _ffmpeg.__all__ + _run.__all__ + _view.__all__ + +__all__ = _filters.__all__ + _ffmpeg.__all__ + _run.__all__ diff --git a/ffmpeg/_ffmpeg.py b/ffmpeg/_ffmpeg.py index 8dc6c87..968e793 100644 --- a/ffmpeg/_ffmpeg.py +++ b/ffmpeg/_ffmpeg.py @@ -1,4 +1,5 @@ from __future__ import unicode_literals + from .nodes import ( filter_operator, GlobalNode, diff --git a/ffmpeg/_filters.py b/ffmpeg/_filters.py index 79949ab..a422f00 100644 --- a/ffmpeg/_filters.py +++ b/ffmpeg/_filters.py @@ -1,4 +1,5 @@ from __future__ import unicode_literals + from .nodes import ( FilterNode, filter_operator, 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/_utils.py b/ffmpeg/_utils.py index 04f2add..06c5765 100644 --- a/ffmpeg/_utils.py +++ b/ffmpeg/_utils.py @@ -1,4 +1,8 @@ -import hashlib +from __future__ import unicode_literals + +from builtins import str +from past.builtins import basestring +import hashlib def _recursive_repr(item): diff --git a/ffmpeg/dag.py b/ffmpeg/dag.py index c8a634e..3ce3891 100644 --- a/ffmpeg/dag.py +++ b/ffmpeg/dag.py @@ -1,3 +1,5 @@ +from __future__ import unicode_literals + from ._utils import get_hash, get_hash_int from builtins import object from collections import namedtuple @@ -72,14 +74,14 @@ DagEdge = namedtuple('DagEdge', ['downstream_node', 'downstream_label', 'upstrea def get_incoming_edges(downstream_node, incoming_edge_map): edges = [] - for downstream_label, (upstream_node, upstream_label) in incoming_edge_map.items(): + for downstream_label, (upstream_node, upstream_label) in list(incoming_edge_map.items()): edges += [DagEdge(downstream_node, downstream_label, upstream_node, upstream_label)] return edges def get_outgoing_edges(upstream_node, outgoing_edge_map): edges = [] - for upstream_label, downstream_infos in outgoing_edge_map.items(): + for upstream_label, downstream_infos in list(outgoing_edge_map.items()): for (downstream_node, downstream_label) in downstream_infos: edges += [DagEdge(downstream_node, downstream_label, upstream_node, upstream_label)] return edges @@ -91,7 +93,7 @@ class KwargReprNode(DagNode): @property def __upstream_hashes(self): hashes = [] - for downstream_label, (upstream_node, upstream_label) in self.incoming_edge_map.items(): + for downstream_label, (upstream_node, upstream_label) in list(self.incoming_edge_map.items()): hashes += [hash(x) for x in [downstream_label, upstream_node, upstream_label]] return hashes diff --git a/ffmpeg/nodes.py b/ffmpeg/nodes.py index b458fae..17b4bc9 100644 --- a/ffmpeg/nodes.py +++ b/ffmpeg/nodes.py @@ -1,5 +1,6 @@ from __future__ import unicode_literals +from builtins import object from .dag import KwargReprNode from ._utils import get_hash_int @@ -38,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 @@ -49,33 +62,21 @@ class Node(KwargReprNode): @classmethod def __check_input_types(cls, stream_map, incoming_stream_types): - for stream in stream_map.values(): + for stream in list(stream_map.values()): if not _is_of_types(stream, incoming_stream_types): 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 = {} - for downstream_label, upstream in stream_map.items(): + for downstream_label, upstream in list(stream_map.items()): incoming_edge_map[downstream_label] = (upstream.node, upstream.label) return incoming_edge_map 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' ]