8000 #17: add proper handling of `split` operator · sdk21/ffmpeg-python@8337d34 · GitHub
[go: up one dir, main page]
Skip to content

Commit 8337d34

Browse files
kkroeningkkroening
committed
kkroening#17: add proper handling of split operator
1 parent 662c56e commit 8337d34

File tree

3 files changed

+67
-42
lines changed

3 files changed

+67
-42
lines changed

‎ffmpeg/_run.py

Lines changed: 27 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from __future__ import unicode_literals
22

3-
from .dag import topo_sort
3+
from .dag import get_outgoing_edges, topo_sort
44
from functools import reduce
55
from past.builtins import basestring
66
import copy
@@ -53,16 +53,31 @@ def _get_input_args(input_node):
5353
return args
5454

5555

56-
def _get_filter_spec(i, node, stream_name_map):
57-
stream_name = _get_stream_name('v{}'.format(i))
58-
stream_name_map[node] = stream_name
59-
inputs = [stream_name_map[edge.upstream_node] for edge in node.incoming_edges]
60-
filter_spec = '{}{}{}'.format(''.join(inputs), node._get_filter(), stream_name)
56+
def _get_filter_spec(node, outgoing_edge_map, stream_name_map):
57+
incoming_edges = node.incoming_edges
58+
outgoing_edges = get_outgoing_edges(node, outgoing_edge_map)
59+
inputs = [stream_name_map[edge.upstream_node, edge.upstream_label] for edge in incoming_edges]
60+
outputs = [stream_name_map[edge.upstream_node, edge.upstream_label] for edge in outgoing_edges]
61+
filter_spec = '{}{}{}'.format(''.join(inputs), node._get_filter(), ''.join(outputs))
6162
return filter_spec
6263

6364

64-
def _get_filter_arg(filter_nodes, stream_name_map):
65-
filter_specs = [_get_filter_spec(i, node, stream_name_map) for i, node in enumerate(filter_nodes)]
65+
def _allocate_filter_stream_names(filter_nodes, outgoing_edge_maps, stream_name_map):
66+
stream_count = 0
67+
for upstream_node in filter_nodes:
68+
outgoing_edge_map = outgoing_edge_maps[upstream_node]
69+
for upstream_label, downstreams in outgoing_edge_map.items():
70+
if len(downstreams) > 1:
71+
# TODO: automatically insert `splits` ahead of time via graph transformation.
72+
raise ValueError('Encountered {} with multiple outgoing edges with same upstream label {!r}; a '
73+
'`split` filter is probably required'.format(upstream_node, upstream_label))
74+
stream_name_map[upstream_node, upstream_label] = _get_stream_name('s{}'.format(stream_count))
75+
stream_count += 1
76+
77+
78+
def _get_filter_arg(filter_nodes, outgoing_edge_maps, stream_name_map):
79+
_allocate_filter_stream_names(filter_nodes, outgoing_edge_maps, stream_name_map)
80+
filter_specs = [_get_filter_spec(node, outgoing_edge_maps[node], stream_name_map) for node in filter_nodes]
6681
return ';'.join(filter_specs)
6782

6883

@@ -78,7 +93,8 @@ def _get_output_args(node, stream_name_map):
7893
raise ValueError('Unsupported output node: {}'.format(node))
7994
args = []
8095
assert len(node.incoming_edges) == 1
81-
stream_name = stream_name_map[node.incoming_edges[0].upstream_node]
96+
edge = node.incoming_edges[0]
97+
stream_name = stream_name_map[edge.upstream_node, edge.upstream_label]
8298
if stream_name != '[0]':
8399
args += ['-map', stream_name]
84100
kwargs = copy.copy(node.kwargs)
@@ -104,8 +120,8 @@ def get_args(stream):
104120
isinstance(node, GlobalNode)]
105121
global_nodes = [node for node in sorted_nodes if isinstance(node, GlobalNode)]
106122
filter_nodes = [node for node in sorted_nodes if node not in (input_nodes + output_nodes + global_nodes)]
107-
stream_name_map = {node: _get_stream_name(i) for i, node in enumerate(input_nodes)}
108-
filter_arg = _get_filter_arg(filter_nodes, stream_name_map)
123+
stream_name_map = {(node, None): _get_stream_name(i) for i, node in enumerate(input_nodes)}
124+
filter_arg = _get_filter_arg(filter_nodes, outgoing_edge_maps, stream_name_map)
109125
args += reduce(operator.add, [_get_input_args(node) for node in input_nodes])
110126
if filter_arg:
111127
args += ['-filter_complex', filter_arg]

‎ffmpeg/nodes.py

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,18 @@ def __repr__(self):
3939
return out
4040

4141

42+
def get_stream_map(stream_spec):
43+
if stream_spec is None:
44+
stream_map = {}
45+
elif isinstance(stream_spec, Stream):
46+
stream_map = {None: stream_spec}
47+
elif isinstance(stream_spec, (list, tuple)):
48+
stream_map = dict(enumerate(stream_spec))
49+
elif isinstance(stream_spec, dict):
50+
stream_map = stream_spec
51+
return stream_map
52+
53+
4254
class Node(KwargReprNode):
4355
"""Node base"""
4456
@classmethod
@@ -55,18 +67,6 @@ def __check_input_types(cls, stream_map, incoming_stream_types):
5567
raise TypeError('Expected incoming stream(s) to be of one of the following types: {}; got {}'
5668
.format(_get_types_str(incoming_stream_types), type(stream)))
5769

58-
@classmethod
59-
def __get_stream_map(cls, stream_spec):
60-
if stream_spec is None:
61-
stream_map = {}
62-
elif isinstance(stream_spec, Stream):
63-
stream_map = {None: stream_spec}
64-
elif isinstance(stream_spec, (list, tuple)):
65-
stream_map = dict(enumerate(stream_spec))
66-
elif isinstance(stream_spec, dict):
67-
stream_map = stream_spec
68-
return stream_map
69-
7070
@classmethod
7171
def __get_incoming_edge_map(cls, stream_map):
7272
incoming_edge_map = {}
@@ -76,7 +76,7 @@ def __get_incoming_edge_map(cls, stream_map):
7676

7777
def __init__(self, stream_spec, name, incoming_stream_types, outgoing_stream_type, min_inputs, max_inputs, args,
7878
kwargs):
79-
stream_map = self.__get_stream_map(stream_spec)
79+
stream_map = get_stream_map(stream_spec)
8080
self.__check_input_len(stream_map, min_inputs, max_inputs)
8181
self.__check_input_types(stream_map, incoming_stream_types)
8282
incoming_edge_map = self.__get_incoming_edge_map(stream_map)

‎ffmpeg/tests/test_ffmpeg.py

Lines changed: 27 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -93,12 +93,19 @@ def test_get_args_simple():
9393

9494

9595
def _get_complex_filter_example():
96-
in_file = ffmpeg.input(TEST_INPUT_FILE)
96+
split = (ffmpeg
97+
.input(TEST_INPUT_FILE)
98+
.vflip()
99+
.split()
100+
)
101+
split0 = split[0]
102+
split1 = split[1]
103+
97104
overlay_file = ffmpeg.input(TEST_OVERLAY_FILE)
98105
return (ffmpeg
99106
.concat(
100-
in_file.trim(start_frame=10, end_frame=20),
101-
in_file.trim(start_frame=30, end_frame=40),
107+
split0.trim(start_frame=10, end_frame=20),
108+
split1.trim(start_frame=30, end_frame=40),
102109
)
103110
.overlay(overlay_file.hflip())
104111
.drawbox(50, 50, 120, 120, color='red', thickness=5)
@@ -110,21 +117,23 @@ def _get_complex_filter_example():
110117
def test_get_args_complex_filter():
111118
out = _get_complex_filter_example()
112119
args = ffmpeg.get_args(out)
113-
assert args == [
114-
'-i', TEST_INPUT_FILE,
120+
assert args == ['-i', TEST_INPUT_FILE,
115121
'-i', TEST_OVERLAY_FILE,
116122
'-filter_complex',
117-
'[0]trim=end_frame=20:start_frame=10[v0];' \
118-
'[0]trim=end_frame=40:start_frame=30[v1];' \
119-
'[v0][v1]concat=n=2[v2];' \
120-
'[1]hflip[v3];' \
121-
'[v2][v3]overlay=eof_action=repeat[v4];' \
122-
'[v4]drawbox=50:50:120:120:red:t=5[v5]',
123-
'-map', '[v5]', os.path.join(SAMPLE_DATA_DIR, 'dummy2.mp4'),
123+
'[0]vflip[s0];' \
124+
'[s0]split[s1][s2];' \
125+
'[s1]trim=end_frame=20:start_frame=10[s3];' \
126+
'[s2]trim=end_frame=40:start_frame=30[s4];' \
127+
'[s3][s4]concat=n=2[s5];' \
128+
'[1]hflip[s6];' \
129+
'[s5][s6]overlay=eof_action=repeat[s7];' \
130+
'[s7]drawbox=50:50:120:120:red:t=5[s8]',
131+
'-map', '[s8]', os.path.join(SAMPLE_DATA_DIR, 'dummy2.mp4'),
124132
'-y'
125133
]
126134

127135

136+
128137
#def test_version():
129138
# subprocess.check_call(['ffmpeg', '-version'])
130139

@@ -156,8 +165,8 @@ def test_custom_filter():
156165
node = ffmpeg.output(node, 'dummy2.mp4')
157166
assert node.get_args() == [
158167
'-i', 'dummy.mp4',
159-
'-filter_complex', '[0]custom_filter=a:b:kwarg1=c[v0]',
160-
'-map', '[v0]',
168+
'-filter_complex', '[0]custom_filter=a:b:kwarg1=c[s0]',
169+
'-map', '[s0]',
161170
'dummy2.mp4'
162171
]
163172

@@ -170,8 +179,8 @@ def test_custom_filter_fluent():
170179
)
171180
assert node.get_args() == [
172181
'-i', 'dummy.mp4',
173-
'-filter_complex', '[0]custom_filter=a:b:kwarg1=c[v0]',
174-
'-map', '[v0]',
182+
'-filter_complex', '[0]custom_filter=a:b:kwarg1=c[s0]',
183+
'-map', '[s0]',
175184
'dummy2.mp4'
176185
]
177186

@@ -197,8 +206,8 @@ def test_pipe():
197206
'-pixel_format', 'rgb24',
198207
'-i', 'pipe:0',
199208
'-filter_complex',
200-
'[0]trim=start_frame=2[v0]',
201-
'-map', '[v0]',
209+
'[0]trim=start_frame=2[s0]',
210+
'-map', '[s0]',
202211
'-f', 'rawvideo',
203212
'pipe:1'
204213
]

0 commit comments

Comments
 (0)
0