diff --git a/README.md b/README.md index d65b8c3..86976c1 100644 --- a/README.md +++ b/README.md @@ -64,7 +64,7 @@ overlay_file = ffmpeg.input('overlay.png') ) .overlay(overlay_file.hflip()) .drawbox(50, 50, 120, 120, color='red', thickness=5) - .output(TEST_OUTPUT_FILE) + .output('out.mp4') .run() ) ``` diff --git a/ffmpeg/__init__.py b/ffmpeg/__init__.py index 154d8b6..d4f48d8 100644 --- a/ffmpeg/__init__.py +++ b/ffmpeg/__init__.py @@ -4,5 +4,6 @@ from . import _filters, _ffmpeg, _run from ._filters import * from ._ffmpeg import * from ._run import * +from ._view import * -__all__ = _filters.__all__ + _ffmpeg.__all__ + _run.__all__ +__all__ = _filters.__all__ + _ffmpeg.__all__ + _run.__all__ + _view.__all__ diff --git a/ffmpeg/_view.py b/ffmpeg/_view.py new file mode 100644 index 0000000..cdb41b0 --- /dev/null +++ b/ffmpeg/_view.py @@ -0,0 +1,86 @@ +from __future__ import unicode_literals + +from builtins import str +from .dag import get_outgoing_edges +from ._run import topo_sort +import os +import tempfile + +from ffmpeg.nodes import ( + FilterNode, + get_stream_spec_nodes, + InputNode, + OutputNode, + Stream, + stream_operator, +) + + +_RIGHT_ARROW = '\u2192' + + +def _get_node_color(node): + if isinstance(node, InputNode): + color = '#99cc00' + elif isinstance(node, OutputNode): + color = '#99ccff' + elif isinstance(node, FilterNode): + color = '#ffcc00' + else: + color = None + return color + + +@stream_operator() +def view(stream_spec, **kwargs): + try: + import graphviz + except ImportError: + raise ImportError('failed to import graphviz; please make sure graphviz is installed (e.g. `pip install ' + 'graphviz`)') + + filename = kwargs.pop('filename', None) + show_labels = kwargs.pop('show_labels', True) + if filename is None: + filename = tempfile.mktemp() + + nodes = get_stream_spec_nodes(stream_spec) + + sorted_nodes, outgoing_edge_maps = topo_sort(nodes) + graph = graphviz.Digraph() + graph.attr(rankdir='LR') + if len(list(kwargs.keys())) != 0: + raise ValueError('Invalid kwargs key(s): {}'.format(', '.join(list(kwargs.keys())))) + + for node in sorted_nodes: + color = _get_node_color(node) + + graph.node(str(hash(node)), node.short_repr, shape='box', style='filled', fillcolor=color) + outgoing_edge_map = outgoing_edge_maps.get(node, {}) + + for edge in get_outgoing_edges(node, outgoing_edge_map): + kwargs = {} + up_label = edge.upstream_label + down_label = edge.downstream_label + if show_labels and (up_label is not None or down_label is not None): + if up_label is None: + up_label = '' + if down_label is None: + down_label = '' + if up_label != '' and down_label != '': + middle = ' {} '.format(_RIGHT_ARROW) + else: + middle = '' + kwargs['label'] = '{} {} {}'.format(up_label, middle, down_label) + upstream_node_id = str(hash(edge.upstream_node)) + downstream_node_id = str(hash(edge.downstream_node)) + graph.edge(upstream_node_id, downstream_node_id, **kwargs) + + graph.view(filename) + + return stream_spec + + +__all__ = [ + 'view', +]