caffeinism-utils 0.0.171__tar.gz → 0.0.173__tar.gz

{caffeinism_utils-0.0.171 → caffeinism_utils-0.0.173}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: caffeinism-utils
-Version: 0.0.171
+Version: 0.0.173
 Summary: 
 Author: Kim Minjong
 Author-email: make.dirty.code@gmail.com

{caffeinism_utils-0.0.171 → caffeinism_utils-0.0.173}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "caffeinism-utils"
-version = "0.0.171"
+version = "0.0.173"
 description = ""
 authors = ["Kim Minjong <make.dirty.code@gmail.com>"]
 readme = "README.md"

{caffeinism_utils-0.0.171 → caffeinism_utils-0.0.173}/src/caffeinism_utils/aiteration.py

@@ -26,35 +26,35 @@ async def rate_limit_iterator(aiterator, iters_per_second):
 
 T = TypeVar("T")
 
-__CLOSE = object()
 
+class BaseStreamQueue(Generic[T]):
+    _CLOSE = object()
 
-class StreamQueue(Generic[T]):
     def __init__(self):
-        self.queue = asyncio.Queue[T]()
+        self.queue = asyncio.Queue()
+
+    async def close(self):
+        await self.queue.put(self._CLOSE)
+
+
+class StreamQueue(BaseStreamQueue):
+    queue: asyncio.Queue[T]
 
     def put(self, data: T) -> None:
         return self.queue.put(data)
 
     async def __aiter__(self) -> AsyncGenerator[T]:
-        while (it := await self.queue.get()) is not __CLOSE:
+        while (it := await self.queue.get()) is not self._CLOSE:
             yield it
 
-    async def close(self):
-        await self.queue.put(__CLOSE)
 
-
-class StreamIteratorQueue(Generic[T]):
-    def __init__(self):
-        self.queue = asyncio.Queue[AsyncIterable[T]]()
+class StreamIteratorQueue(BaseStreamQueue):
+    queue: asyncio.Queue[AsyncIterable[T]]
 
     def put(self, data: AsyncIterable[T]) -> None:
         return self.queue.put(data)
 
     async def __aiter__(self) -> AsyncGenerator[T]:
-        while (iterator := await self.queue.get()) is not __CLOSE:
+        while (iterator := await self.queue.get()) is not self._CLOSE:
             async for it in iterator:
                 yield it
-
-    async def close(self):
-        await self.queue.put(__CLOSE)

{caffeinism_utils-0.0.171 → caffeinism_utils-0.0.173}/src/caffeinism_utils/av/io.py

@@ -6,7 +6,7 @@ from abc import ABC, abstractmethod
 from concurrent.futures import Future, ThreadPoolExecutor
 from pathlib import Path
 from types import TracebackType
-from typing import Iterable
+from typing import Generator, Iterable
 
 import av
 import av.container
@@ -100,7 +100,7 @@ class BasePyAVReader(PyAVInterface):
         self._codec_contexts = {}
 
     @abstractmethod
-    def __iter__(self):
+    def __iter__(self) -> Generator[av.VideoFrame | av.AudioFrame, None, None]:
         raise NotImplementedError
 
     def _create_container(self):
@@ -178,7 +178,7 @@ class PyAVReader(BasePyAVReader):
                 raise NotImplementedError
         return self._alpha_merger
 
-    def __iter__(self):
+    def __iter__(self) -> Generator[av.VideoFrame | av.AudioFrame, None, None]:
         with self:
             for packet in self.container.demux(self.streams + self.audio_streams):
                 for frame in self.codec_contexts[packet.stream].decode(packet):

caffeinism_utils-0.0.173/src/caffeinism_utils/av/utils.py

@@ -0,0 +1,292 @@
+import av
+import numba
+
+@numba.jit(
+    (
+        numba.types.MemoryView(numba.uint8, 1, "C"),
+        numba.types.MemoryView(numba.uint8, 1, "C"),
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+    ),
+    nopython=True,
+    nogil=True,
+)
+def _frame_paste_rgba(
+    dst_view: memoryview,
+    src_view: memoryview,
+    x: int,
+    y: int,
+    dw: int,
+    dh: int,
+    sw: int,
+    sh: int,
+    src_stride: int,
+    dst_stride: int,
+):
+    sx = sy = 0
+
+    if x < 0:
+        sx -= x
+        sw += x
+        x = 0
+    if y < 0:
+        sy -= y
+        sh += y
+        y = 0
+
+    sw = min(sw, sw - sx, dw - x)
+    sh = min(sh, sh - sy, dh - y)
+
+    if sw <= 0 or sh <= 0:
+        return
+
+    s_base = sy * src_stride + sx * 4
+    d_base = y * dst_stride + x * 4
+    row_bytes = sw * 4
+
+    for r in range(sh):
+        so = s_base + r * src_stride
+        do = d_base + r * dst_stride
+        dst_view[do : do + row_bytes] = src_view[so : so + row_bytes]
+
+
+def frame_paste_rgba(dst: av.VideoFrame, src: av.VideoFrame, x: int, y: int):
+    assert src.format.name == "rgba" and dst.format.name == "rgba"
+
+    dst.make_writable()
+
+    src_plane = src.planes[0]
+    dst_plane = dst.planes[0]
+
+    _frame_paste_rgba(
+        memoryview(dst_plane),
+        memoryview(src_plane),
+        x,
+        y,
+        dst.width,
+        dst.height,
+        src.width,
+        src.height,
+        src_plane.line_size,
+        dst_plane.line_size,
+    )
+    return dst
+
+
+@numba.jit(
+    (
+        numba.types.MemoryView(numba.uint8, 1, "C"),
+        numba.types.MemoryView(numba.uint8, 1, "C"),
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+    ),
+    nopython=True,
+    nogil=True,
+)
+def _blit_8bit_planar(
+    dst_view: memoryview,
+    src_view: memoryview,
+    dx: int,
+    dy: int,
+    sx: int,
+    sy: int,
+    w: int,
+    h: int,
+    dst_stride: int,
+    src_stride: int,
+):
+    if w <= 0 or h <= 0:
+        return
+    row_bytes = w  # 1 byte per pixel in 8-bit planar
+    s_base = sy * src_stride + sx
+    d_base = dy * dst_stride + dx
+    for r in range(h):
+        so = s_base + r * src_stride
+        do = d_base + r * dst_stride
+        dst_view[do : do + row_bytes] = src_view[so : so + row_bytes]
+
+
+@numba.jit(
+    (
+        numba.types.MemoryView(numba.uint8, 1, "C"),
+        numba.types.MemoryView(numba.uint8, 1, "C"),
+        numba.types.MemoryView(numba.uint8, 1, "C"),
+        numba.types.MemoryView(numba.uint8, 1, "C"),
+        numba.types.MemoryView(numba.uint8, 1, "C"),
+        numba.types.MemoryView(numba.uint8, 1, "C"),
+        numba.types.MemoryView(numba.uint8, 1, "C"),
+        numba.types.MemoryView(numba.uint8, 1, "C"),
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+        numba.int64,
+    ),
+    nopython=True,
+    nogil=True,
+)
+def _frame_paste_yuva420p(
+    dst_view_y: memoryview,
+    dst_view_u: memoryview,
+    dst_view_v: memoryview,
+    dst_view_a: memoryview,
+    src_view_y: memoryview,
+    src_view_u: memoryview,
+    src_view_v: memoryview,
+    src_view_a: memoryview,
+    x: int,
+    y: int,
+    dw: int,
+    dh: int,
+    sw: int,
+    sh: int,
+    src_stride_y: int,
+    src_stride_u: int,
+    src_stride_v: int,
+    src_stride_a: int,
+    dst_stride_y: int,
+    dst_stride_u: int,
+    dst_stride_v: int,
+    dst_stride_a: int,
+):
+    sw0, sh0 = sw, sh
+    dw, dh = dw, dh
+
+    sx = max(0, -x)
+    sy = max(0, -y)
+    dx = max(0, x)
+    dy = max(0, y)
+
+    w = min(sw0 - sx, dw - dx)
+    h = min(sh0 - sy, dh - dy)
+
+    if w <= 0 or h <= 0:
+        return
+
+    _blit_8bit_planar(
+        dst_view_y,
+        src_view_y,
+        dx,
+        dy,
+        sx,
+        sy,
+        w,
+        h,
+        dst_stride_y,
+        src_stride_y,
+    )
+
+    _blit_8bit_planar(
+        dst_view_a,
+        src_view_a,
+        dx,
+        dy,
+        sx,
+        sy,
+        w,
+        h,
+        dst_stride_a,
+        src_stride_a,
+    )
+
+    hsub = 1
+    vsub = 1
+
+    sx_c = sx >> hsub
+    sy_c = sy >> vsub
+    dx_c = dx >> hsub
+    dy_c = dy >> vsub
+
+    src_cw = (sw0 + (1 << hsub) - 1) >> hsub
+    src_ch = (sh0 + (1 << vsub) - 1) >> vsub
+    dst_cw = (dw + (1 << hsub) - 1) >> hsub
+    dst_ch = (dh + (1 << vsub) - 1) >> vsub
+
+    w_c = ((sx + w + (1 << hsub) - 1) >> hsub) - (sx >> hsub)
+    h_c = ((sy + h + (1 << vsub) - 1) >> vsub) - (sy >> vsub)
+
+    w_c = min(w_c, src_cw - sx_c, dst_cw - dx_c)
+    h_c = min(h_c, src_ch - sy_c, dst_ch - dy_c)
+
+    if w_c > 0 and h_c > 0:
+        _blit_8bit_planar(
+            dst_view_u,
+            src_view_u,
+            dx_c,
+            dy_c,
+            sx_c,
+            sy_c,
+            w_c,
+            h_c,
+            dst_stride_u,
+            src_stride_u,
+        )
+        _blit_8bit_planar(
+            dst_view_v,
+            src_view_v,
+            dx_c,
+            dy_c,
+            sx_c,
+            sy_c,
+            w_c,
+            h_c,
+            dst_stride_v,
+            src_stride_v,
+        )
+
+
+def frame_paste_yuva420p(dst: av.VideoFrame, src: av.VideoFrame, x: int, y: int):
+    assert src.format.name == "yuva420p" and dst.format.name == "yuva420p"
+
+    dst.make_writable()
+
+    src_Y, src_U, src_V, src_A = src.planes
+    dst_Y, dst_U, dst_V, dst_A = dst.planes
+
+    _frame_paste_yuva420p(
+        memoryview(dst_Y),
+        memoryview(dst_U),
+        memoryview(dst_V),
+        memoryview(dst_A),
+        memoryview(src_Y),
+        memoryview(src_U),
+        memoryview(src_V),
+        memoryview(src_A),
+        x,
+        y,
+        dst.width,
+        dst.height,
+        src.width,
+        src.height,
+        src_Y.line_size,
+        src_U.line_size,
+        src_V.line_size,
+        src_A.line_size,
+        dst_Y.line_size,
+        dst_U.line_size,
+        dst_V.line_size,
+        dst_A.line_size,
+    )
+
+    return dst