Shiina-Mahiru commited on
Commit
90c3474
·
verified ·
1 Parent(s): f0f994f

Upload ffmpeg\_view.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. ffmpeg//_view.py +108 -0
ffmpeg//_view.py ADDED
@@ -0,0 +1,108 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import unicode_literals
2
+
3
+ from builtins import str
4
+ from .dag import get_outgoing_edges
5
+ from ._run import topo_sort
6
+ import tempfile
7
+
8
+ from ffmpeg.nodes import (
9
+ FilterNode,
10
+ get_stream_spec_nodes,
11
+ InputNode,
12
+ OutputNode,
13
+ stream_operator,
14
+ )
15
+
16
+
17
+ _RIGHT_ARROW = '\u2192'
18
+
19
+
20
+ def _get_node_color(node):
21
+ if isinstance(node, InputNode):
22
+ color = '#99cc00'
23
+ elif isinstance(node, OutputNode):
24
+ color = '#99ccff'
25
+ elif isinstance(node, FilterNode):
26
+ color = '#ffcc00'
27
+ else:
28
+ color = None
29
+ return color
30
+
31
+
32
+ @stream_operator()
33
+ def view(stream_spec, detail=False, filename=None, pipe=False, **kwargs):
34
+ try:
35
+ import graphviz
36
+ except ImportError:
37
+ raise ImportError(
38
+ 'failed to import graphviz; please make sure graphviz is installed (e.g. `pip install '
39
+ 'graphviz`)'
40
+ )
41
+
42
+ show_labels = kwargs.pop('show_labels', True)
43
+ if pipe and filename is not None:
44
+ raise ValueError('Can\'t specify both `filename` and `pipe`')
45
+ elif not pipe and filename is None:
46
+ filename = tempfile.mktemp()
47
+
48
+ nodes = get_stream_spec_nodes(stream_spec)
49
+
50
+ sorted_nodes, outgoing_edge_maps = topo_sort(nodes)
51
+ graph = graphviz.Digraph(format='png')
52
+ graph.attr(rankdir='LR')
53
+ if len(list(kwargs.keys())) != 0:
54
+ raise ValueError(
55
+ 'Invalid kwargs key(s): {}'.format(', '.join(list(kwargs.keys())))
56
+ )
57
+
58
+ for node in sorted_nodes:
59
+ color = _get_node_color(node)
60
+
61
+ if detail:
62
+ lines = [node.short_repr]
63
+ lines += ['{!r}'.format(arg) for arg in node.args]
64
+ lines += [
65
+ '{}={!r}'.format(key, node.kwargs[key]) for key in sorted(node.kwargs)
66
+ ]
67
+ node_text = '\n'.join(lines)
68
+ else:
69
+ node_text = node.short_repr
70
+ graph.node(
71
+ str(hash(node)), node_text, shape='box', style='filled', fillcolor=color
72
+ )
73
+ outgoing_edge_map = outgoing_edge_maps.get(node, {})
74
+
75
+ for edge in get_outgoing_edges(node, outgoing_edge_map):
76
+ kwargs = {}
77
+ up_label = edge.upstream_label
78
+ down_label = edge.downstream_label
79
+ up_selector = edge.upstream_selector
80
+
81
+ if show_labels and (
82
+ up_label is not None
83
+ or down_label is not None
84
+ or up_selector is not None
85
+ ):
86
+ if up_label is None:
87
+ up_label = ''
88
+ if up_selector is not None:
89
+ up_label += ":" + up_selector
90
+ if down_label is None:
91
+ down_label = ''
92
+ if up_label != '' and down_label != '':
93
+ middle = ' {} '.format(_RIGHT_ARROW)
94
+ else:
95
+ middle = ''
96
+ kwargs['label'] = '{} {} {}'.format(up_label, middle, down_label)
97
+ upstream_node_id = str(hash(edge.upstream_node))
98
+ downstream_node_id = str(hash(edge.downstream_node))
99
+ graph.edge(upstream_node_id, downstream_node_id, **kwargs)
100
+
101
+ if pipe:
102
+ return graph.pipe()
103
+ else:
104
+ graph.view(filename, cleanup=True)
105
+ return stream_spec
106
+
107
+
108
+ __all__ = ['view']