diff --git a/av/filter/graph.pxd b/av/filter/graph.pxd index 2e52bd6ec..0e216cf12 100644 --- a/av/filter/graph.pxd +++ b/av/filter/graph.pxd @@ -12,7 +12,6 @@ cdef class Graph: cpdef configure(self, bint auto_buffer=*, bint force=*) cdef dict _name_counts - cdef str _get_unique_name(self, str name) cdef _register_context(self, FilterContext) cdef _auto_register(self) diff --git a/av/filter/graph.py b/av/filter/graph.py index 471b6ebc4..2cca2a712 100644 --- a/av/filter/graph.py +++ b/av/filter/graph.py @@ -17,7 +17,6 @@ class Graph: def __cinit__(self): self.ptr = lib.avfilter_graph_alloc() self.configured = False - self._name_counts = {} self._nb_filters_seen = 0 self._context_by_ptr = {} self._context_by_name = {} @@ -28,14 +27,10 @@ def __dealloc__(self): # This frees the graph, filter contexts, links, etc.. lib.avfilter_graph_free(cython.address(self.ptr)) - @cython.cfunc - def _get_unique_name(self, name: str) -> str: - count = self._name_counts.get(name, 0) - self._name_counts[name] = count + 1 - if count: - return "%s_%s" % (name, count) - else: - return name + def _ctx_from_name_or_die(self, name: str) -> FilterContext: + if (ctx := self._context_by_name.get(name)) is None: + raise ValueError(f"'{name}' is not a valid filter context") + return ctx @cython.ccall def configure(self, auto_buffer: cython.bint = True, force: cython.bint = False): @@ -53,10 +48,10 @@ def link_nodes(self, *nodes): Links nodes together for simple filter graphs. """ for c, n in zip(nodes, nodes[1:]): - c.link_to(n) + self._ctx_from_name_or_die(c).link_to(self._ctx_from_name_or_die(n)) return self - def add(self, filter, args=None, **kwargs): + def add(self, name, filter, args=None, **kwargs): cy_filter: Filter if isinstance(filter, str): cy_filter = Filter(filter) @@ -65,7 +60,9 @@ def add(self, filter, args=None, **kwargs): else: raise TypeError("filter must be a string or Filter") - name: str = self._get_unique_name(kwargs.pop("name", None) or cy_filter.name) + if name in self._context_by_name: + raise ValueError(f"Filter context name '{name}' already taken") + ptr: cython.pointer[lib.AVFilterContext] = lib.avfilter_graph_alloc_filter( self.ptr, cy_filter.ptr, name ) @@ -74,7 +71,7 @@ def add(self, filter, args=None, **kwargs): # Manually construct this context (so we can return it). ctx: FilterContext = wrap_filter_context(self, cy_filter, ptr) - ctx.init(args, **kwargs) + ctx.init(args, name=name, **kwargs) self._register_context(ctx) # There might have been automatic contexts added (e.g. resamplers, @@ -109,11 +106,11 @@ def _auto_register(self): def add_buffer( self, + name, template=None, width=None, height=None, format=None, - name=None, time_base=None, ): if template is not None: @@ -141,8 +138,8 @@ def add_buffer( time_base = Fraction(1, 1000) return self.add( + name, "buffer", - name=name, video_size=f"{width}x{height}", pix_fmt=str(int(VideoFormat(format))), time_base=str(time_base), @@ -151,12 +148,12 @@ def add_buffer( def add_abuffer( self, + name, template=None, sample_rate=None, format=None, layout=None, channels=None, - name=None, time_base=None, ): """ @@ -194,7 +191,7 @@ def add_abuffer( if channels: kwargs["channels"] = f"{channels}" - return self.add("abuffer", name=name, **kwargs) + return self.add(name, "abuffer", **kwargs) def set_audio_frame_size(self, frame_size): """ @@ -211,7 +208,10 @@ def set_audio_frame_size(self, frame_size): cython.cast(FilterContext, sink).ptr, frame_size ) - def push(self, frame): + def filter_push(self, frame, name=None): + if name: + return self._ctx_from_name_or_die(name).push(frame) + if frame is None: contexts = self._context_by_type.get( "buffer", [] @@ -228,13 +228,16 @@ def push(self, frame): for ctx in contexts: ctx.push(frame) - def vpush(self, frame: VideoFrame | None): + def filter_vpush(self, frame: VideoFrame | None): """Like `push`, but only for VideoFrames.""" for ctx in self._context_by_type.get("buffer", []): ctx.push(frame) # TODO: Test complex filter graphs, add `at: int = 0` arg to pull() and vpull(). - def pull(self): + def filter_pull(self, name=None): + if name: + return self._ctx_from_name_or_die(name).pull() + vsinks = self._context_by_type.get("buffersink", []) asinks = self._context_by_type.get("abuffersink", []) @@ -244,7 +247,7 @@ def pull(self): return (vsinks or asinks)[0].pull() - def vpull(self): + def filter_vpull(self): """Like `pull`, but only for VideoFrames.""" vsinks = self._context_by_type.get("buffersink", []) nsinks = len(vsinks) @@ -252,3 +255,20 @@ def vpull(self): raise ValueError(f"can only auto-pull with single sink; found {nsinks}") return vsinks[0].pull() + + def filter_link_to(self, output_name, input_name, output_idx=0, input_idx=0): + self._ctx_from_name_or_die(output_name).link_to( + self._ctx_from_name_or_die(input_name), + output_idx, + input_idx + ) + + def filter_process_command( + self, + name, + cmd, + arg=None, + res_len=1024, + flags=0 + ): + return self._ctx_from_name_or_die(name).process_command(cmd, arg, res_len, flags) \ No newline at end of file diff --git a/av/filter/graph.pyi b/av/filter/graph.pyi index e170c2ce7..3af994741 100644 --- a/av/filter/graph.pyi +++ b/av/filter/graph.pyi @@ -17,21 +17,21 @@ class Graph: def __init__(self) -> None: ... def configure(self, auto_buffer: bool = True, force: bool = False) -> None: ... - def link_nodes(self, *nodes: FilterContext) -> Graph: ... - def add( - self, filter: str | Filter, args: Any = None, **kwargs: str - ) -> FilterContext: ... + def link_nodes(self, *nodes: str) -> Graph: ... + def add(self, node_name: str, filter: str | Filter, args: Any = None, **kwargs: str) -> None: ... def add_buffer( self, + node_name: str, template: VideoStream | None = None, width: int | None = None, height: int | None = None, format: VideoFormat | None = None, name: str | None = None, time_base: Fraction | None = None, - ) -> FilterContext: ... + ) -> None: ... def add_abuffer( self, + node_name: str, template: AudioStream | None = None, sample_rate: int | None = None, format: AudioFormat | str | None = None, @@ -39,9 +39,16 @@ class Graph: channels: int | None = None, name: str | None = None, time_base: Fraction | None = None, - ) -> FilterContext: ... + ) -> None: ... def set_audio_frame_size(self, frame_size: int) -> None: ... - def push(self, frame: None | AudioFrame | VideoFrame) -> None: ... - def pull(self) -> VideoFrame | AudioFrame: ... - def vpush(self, frame: VideoFrame | None) -> None: ... - def vpull(self) -> VideoFrame: ... + def filter_vpush(self, frame: VideoFrame | None) -> None: ... + def filter_vpull(self) -> VideoFrame: ... + def filter_link_to( + self, output_name: str, input_name: str, output_idx: int = 0, input_idx: int = 0 + ) -> None: ... + def filter_push(self, name: str, frame: Frame | None) -> None: ... + def filter_pull(self, name: str) -> Frame: ... + def filter_process_command( + self, name: str, cmd: str, arg: str | None = None, res_len: int = 1024, flags: int = 0 + ) -> str | None: ... + def _ctx_from_name_or_die(self, name: str) -> FilterContext: ... \ No newline at end of file