changed composite functions to allow video and image mask values

Author: readrawhex

dithertools/dither.py

@@ -85,10 +85,73 @@ def _blend_color(
     return np.where(top > mask, top, bottom).astype(np.uint8)
 
 
+def _rgb_to_gray(arr: np.ndarray) -> np.ndarray:
+    if arr.ndim == 2:
+        return arr
+    return (
+        0.299 * arr[:, :, 0].astype(np.float32)
+        + 0.587 * arr[:, :, 1].astype(np.float32)
+        + 0.114 * arr[:, :, 2].astype(np.float32)
+    ).astype(np.uint8)
+
+
+def _build_mask_iter(
+    source: object,
+    w: int,
+    h: int,
+    fc: int | None,
+    to_gray: bool,
+) -> Iterator[np.ndarray]:
+    """Build a per-frame mask iterator from an Image, Video, or FrameGenerator."""
+    from dithertools.image import Image
+    from dithertools.video import Video
+
+    if isinstance(source, Image):
+        m = source.resize(w, h).data
+        if to_gray:
+            m = _rgb_to_gray(m)
+        return itertools.repeat(m)
+
+    if isinstance(source, Video):
+        if fc is not None:
+            source.resample(fc)
+        source.resize(w, h)
+
+        def _video_mask() -> Iterator[np.ndarray]:
+            for frame in source:
+                yield _rgb_to_gray(frame) if to_gray else frame
+
+        return _video_mask()
+
+    if isinstance(source, FrameGenerator):
+        if source._source_type == "video":
+            if fc is not None:
+                source._video.resample(fc)  # type: ignore[union-attr]
+            source._video.resize(w, h)  # type: ignore[union-attr]
+            iter(source._video)  # type: ignore[union-attr]
+            source._frame_idx = 0
+            source._buffer = None
+
+            def _fg_video_mask() -> Iterator[np.ndarray]:
+                for frame in source:
+                    yield _rgb_to_gray(frame) if to_gray else frame
+
+            return _fg_video_mask()
+
+        def _fg_image_mask() -> Iterator[np.ndarray]:
+            while True:
+                frame = _resize_frame(next(source), w, h)  # type: ignore[arg-type]
+                yield _rgb_to_gray(frame) if to_gray else frame
+
+        return _fg_image_mask()
+
+    raise TypeError(f"unsupported mask source type {type(source).__name__!r}")
+
+
 def composite(
     top: "Image | Video | FrameGenerator",
     bottom: "Image | Video | FrameGenerator",
-    mask: np.ndarray | None = None,
+    mask: "np.ndarray | Image | Video | FrameGenerator | MaskGenerator | None" = None,
     tile_every_x: int | None = 8,
     tile_every_y: int | None = 8,
     repeat_every_x: int | None = 1,
@@ -132,18 +195,28 @@ def composite(
             :class:`~dithertools.image.Image`,
             :class:`~dithertools.video.Video`, or
             :class:`~dithertools.generators.FrameGenerator`.
-        mask: A ``(height, width)`` ``uint8`` array, or a
-            :class:`~dithertools.generators.MaskGenerator`, against which the
-            top pixel brightness is compared.  Static arrays such as those
-            returned by :func:`~dithertools.matrix.random_matrix` are reused
-            for every frame.  A :class:`~dithertools.generators.MaskGenerator`
-            yields a new mask per frame, enabling animated dithering patterns;
-            ``color=True`` generators are rejected.  If ``None``, a fresh
-            random grayscale mask is generated at the bottom layer's dimensions.
-            If a static array does not match the bottom layer dimensions it is
-            resized with nearest-neighbour interpolation.  A generator with no
-            ``initial`` is lazily initialized to a random mask at the bottom
-            layer's dimensions.
+        mask: Grayscale mask source.  Accepted types:
+
+            - ``None`` — a random grayscale mask is generated at the bottom
+              layer's dimensions.
+            - ``np.ndarray`` ``(H, W)`` ``uint8`` — static mask, reused every
+              frame; resized with nearest-neighbour interpolation if dimensions
+              differ.
+            - :class:`~dithertools.image.Image` — converted to grayscale (RGB
+              images use ``0.299R + 0.587G + 0.114B``), resized to bottom
+              dimensions, and reused every frame.
+            - :class:`~dithertools.video.Video` or video-backed
+              :class:`~dithertools.generators.FrameGenerator` — frames are
+              used one-per-output-frame; resampled to match the output frame
+              count and resized to bottom dimensions.  RGB frames are
+              auto-converted to grayscale.
+            - Image-backed :class:`~dithertools.generators.FrameGenerator` —
+              infinite; ``fn`` is called each output frame; frames resized and
+              converted to grayscale.
+            - :class:`~dithertools.generators.MaskGenerator` — yields a new
+              mask per frame; ``color=True`` generators are rejected.  A
+              generator with no ``initial`` is lazily initialized at the bottom
+              layer's dimensions.
         tile_every_x: Tile width passed to :func:`~dithertools.matrix.random_matrix`
             when generating an auto mask (i.e. when ``mask`` is ``None``).
             The mask pattern repeats horizontally every this many columns.
@@ -199,7 +272,7 @@ def composite(
             repeat_every_x=repeat_every_x,
             repeat_every_y=repeat_every_y,
         )
-        mask_iter: Iterator[np.ndarray] = itertools.repeat(_arr)
+        mask_iter: Iterator[np.ndarray] | None = itertools.repeat(_arr)
     elif isinstance(mask, MaskGenerator):
         if mask.color:
             raise ValueError(
@@ -219,10 +292,25 @@ def composite(
                 )
             )
         mask_iter = mask
-    else:
+    elif isinstance(mask, np.ndarray):
         if mask.shape[:2] != (h, w):
             mask = np.array(_PILImage.fromarray(mask).resize((w, h), _PILImage.NEAREST))
         mask_iter = itertools.repeat(mask)
+    elif isinstance(mask, Image):
+        _m = mask.resize(w, h).data
+        mask_iter = itertools.repeat(_rgb_to_gray(_m))
+    elif isinstance(mask, (Video, FrameGenerator)):
+        mask_iter = None  # resolved per-branch once frame count is known
+    else:
+        raise TypeError(
+            f"mask must be a numpy array, Image, Video, FrameGenerator, "
+            f"MaskGenerator, or None; got {type(mask).__name__!r}"
+        )
+
+    def _resolve_mask(fc: int | None = None) -> Iterator[np.ndarray]:
+        if mask_iter is not None:
+            return mask_iter  # type: ignore[return-value]
+        return _build_mask_iter(mask, w, h, fc, to_gray=True)
 
     def _is_fg_video(x: object) -> bool:
         return isinstance(x, FrameGenerator) and x._source_type == "video"
@@ -233,13 +321,14 @@ def _is_fg_image(x: object) -> bool:
     # Image × Image → Image
     if isinstance(bottom, Image) and isinstance(top, Image):
         return Image.from_array(
-            _blend(top.resize(w, h).data, bottom.data, next(mask_iter))
+            _blend(top.resize(w, h).data, bottom.data, next(_resolve_mask()))
         )
 
     # --- Bottom is finite (Video or fg_video): bottom drives frame count ---
     if isinstance(bottom, Video) or _is_fg_video(bottom):
         fc = bottom.frame_count
         fps = bottom.fps
+        mask_iter = _resolve_mask(fc)
 
         if isinstance(top, Image):
             top_frame = top.resize(w, h).data
@@ -293,6 +382,7 @@ def _fg_video_over_video() -> Iterator[np.ndarray]:
         if isinstance(top, Video):
             top.resize(w, h)
             total = top.frame_count
+            mask_iter = _resolve_mask(total)
 
             def _video_over_image() -> Iterator[np.ndarray]:
                 for top_frame, m in zip(top, mask_iter):
@@ -307,6 +397,7 @@ def _video_over_image() -> Iterator[np.ndarray]:
             iter(top._video)  # type: ignore[union-attr]
             top._frame_idx = 0  # type: ignore[union-attr]
             total = top.frame_count
+            mask_iter = _resolve_mask(total)
 
             def _fg_video_over_image() -> Iterator[np.ndarray]:
                 for top_frame, m in zip(top, mask_iter):
@@ -321,6 +412,7 @@ def _fg_video_over_image() -> Iterator[np.ndarray]:
         if isinstance(top, Video):
             top.resize(w, h)
             total = top.frame_count
+            mask_iter = _resolve_mask(total)
 
             def _video_over_fg_image() -> Iterator[np.ndarray]:
                 for top_frame, m in zip(top, mask_iter):
@@ -335,6 +427,7 @@ def _video_over_fg_image() -> Iterator[np.ndarray]:
             iter(top._video)  # type: ignore[union-attr]
             top._frame_idx = 0  # type: ignore[union-attr]
             total = top.frame_count
+            mask_iter = _resolve_mask(total)
 
             def _fg_video_over_fg_image() -> Iterator[np.ndarray]:
                 for top_frame, m in zip(top, mask_iter):
@@ -353,7 +446,7 @@ def _fg_video_over_fg_image() -> Iterator[np.ndarray]:
 def composite_color(
     top: "Image | Video | FrameGenerator",
     bottom: "Image | Video | FrameGenerator",
-    mask: np.ndarray | None = None,
+    mask: "np.ndarray | Image | Video | FrameGenerator | MaskGenerator | None" = None,
     tile_every_x: int | None = 8,
     tile_every_y: int | None = 8,
     repeat_every_x: int | None = 1,
@@ -397,21 +490,26 @@ def composite_color(
             :class:`~dithertools.image.Image`,
             :class:`~dithertools.video.Video`, or
             :class:`~dithertools.generators.FrameGenerator`.
-        mask: A ``(height, width, 3)`` or ``(height, width)`` ``uint8`` array,
-            or a :class:`~dithertools.generators.MaskGenerator`, against which
-            the top channels are compared.  A three-channel array or
-            ``color=True`` generator applies independent per-channel thresholds;
-            a grayscale array or ``color=False`` generator applies the same
-            threshold to all three channels.  Static arrays such as those
-            returned by :func:`~dithertools.matrix.random_matrix_color` are
-            reused for every frame.  A :class:`~dithertools.generators.MaskGenerator`
-            yields a new mask per frame.  If ``None``, a fresh random color mask
-            is generated at the bottom layer's dimensions.  If a static array
-            does not match the bottom layer dimensions it is resized with
-            nearest-neighbour interpolation.  A generator with no ``initial``
-            is lazily initialized using
-            :func:`~dithertools.matrix.random_matrix_color` (for ``color=True``)
-            or :func:`~dithertools.matrix.random_matrix` (for ``color=False``).
+        mask: Per-channel mask source.  Accepted types:
+
+            - ``None`` — a random color mask is generated at the bottom layer's
+              dimensions.
+            - ``np.ndarray`` ``(H, W)`` or ``(H, W, 3)`` ``uint8`` — static
+              mask, reused every frame; resized if dimensions differ.  A
+              non-3-channel 3-D array raises ``ValueError``.
+            - :class:`~dithertools.image.Image` — used as-is (``(H, W)``
+              grayscale or ``(H, W, 3)`` color); resized to bottom dimensions.
+              A non-3-channel color image raises ``ValueError``.
+            - :class:`~dithertools.video.Video` or video-backed
+              :class:`~dithertools.generators.FrameGenerator` — frames are used
+              one-per-output-frame; resampled to match output frame count and
+              resized to bottom dimensions.
+            - Image-backed :class:`~dithertools.generators.FrameGenerator` —
+              infinite; ``fn`` is called each output frame; frames resized.
+            - :class:`~dithertools.generators.MaskGenerator` — yields a new
+              mask per frame; ``color=True`` applies per-channel thresholds,
+              ``color=False`` broadcasts the same threshold to all channels.
+              A generator with no ``initial`` is lazily initialized.
         tile_every_x: Tile width passed to
             :func:`~dithertools.matrix.random_matrix_color` when generating an
             auto mask (i.e. when ``mask`` is ``None``). The mask pattern repeats
@@ -468,7 +566,7 @@ def composite_color(
             repeat_every_x=repeat_every_x,
             repeat_every_y=repeat_every_y,
         )
-        mask_iter: Iterator[np.ndarray] = itertools.repeat(_arr)
+        mask_iter: Iterator[np.ndarray] | None = itertools.repeat(_arr)
     elif isinstance(mask, MaskGenerator):
         if not mask.initialized:
             if mask.color:
@@ -491,7 +589,7 @@ def composite_color(
                 )
             mask._initialize(_arr)
         mask_iter = mask
-    else:
+    elif isinstance(mask, np.ndarray):
         if mask.ndim == 3 and mask.shape[2] != 3:
             raise ValueError(
                 f"composite_color requires a (H, W) or (H, W, 3) mask, "
@@ -505,6 +603,31 @@ def composite_color(
         if mask.shape[:2] != (h, w):
             mask = np.array(_PILImage.fromarray(mask).resize((w, h), _PILImage.NEAREST))
         mask_iter = itertools.repeat(mask)
+    elif isinstance(mask, Image):
+        _m = mask.resize(w, h).data
+        if _m.ndim == 3 and _m.shape[2] != 3:
+            raise ValueError(
+                f"composite_color requires a (H, W) or (H, W, 3) mask, "
+                f"got Image with shape {_m.shape}"
+            )
+        if _m.ndim not in (2, 3):
+            raise ValueError(
+                f"composite_color requires a (H, W) or (H, W, 3) mask, "
+                f"got Image with shape {_m.shape}"
+            )
+        mask_iter = itertools.repeat(_m)
+    elif isinstance(mask, (Video, FrameGenerator)):
+        mask_iter = None  # resolved per-branch once frame count is known
+    else:
+        raise TypeError(
+            f"mask must be a numpy array, Image, Video, FrameGenerator, "
+            f"MaskGenerator, or None; got {type(mask).__name__!r}"
+        )
+
+    def _resolve_mask(fc: int | None = None) -> Iterator[np.ndarray]:
+        if mask_iter is not None:
+            return mask_iter  # type: ignore[return-value]
+        return _build_mask_iter(mask, w, h, fc, to_gray=False)
 
     def _is_fg_video(x: object) -> bool:
         return isinstance(x, FrameGenerator) and x._source_type == "video"
@@ -515,13 +638,14 @@ def _is_fg_image(x: object) -> bool:
     # Image × Image → Image
     if isinstance(bottom, Image) and isinstance(top, Image):
         return Image.from_array(
-            _blend_color(top.resize(w, h).data, bottom.data, next(mask_iter))
+            _blend_color(top.resize(w, h).data, bottom.data, next(_resolve_mask()))
         )
 
     # --- Bottom is finite (Video or fg_video): bottom drives frame count ---
     if isinstance(bottom, Video) or _is_fg_video(bottom):
         fc = bottom.frame_count
         fps = bottom.fps
+        mask_iter = _resolve_mask(fc)
 
         if isinstance(top, Image):
             top_frame = top.resize(w, h).data
@@ -575,6 +699,7 @@ def _fg_video_over_video() -> Iterator[np.ndarray]:
         if isinstance(top, Video):
             top.resize(w, h)
             total = top.frame_count
+            mask_iter = _resolve_mask(total)
 
             def _video_over_image() -> Iterator[np.ndarray]:
                 for top_frame, m in zip(top, mask_iter):
@@ -589,6 +714,7 @@ def _video_over_image() -> Iterator[np.ndarray]:
             iter(top._video)  # type: ignore[union-attr]
             top._frame_idx = 0  # type: ignore[union-attr]
             total = top.frame_count
+            mask_iter = _resolve_mask(total)
 
             def _fg_video_over_image() -> Iterator[np.ndarray]:
                 for top_frame, m in zip(top, mask_iter):
@@ -603,6 +729,7 @@ def _fg_video_over_image() -> Iterator[np.ndarray]:
         if isinstance(top, Video):
             top.resize(w, h)
             total = top.frame_count
+            mask_iter = _resolve_mask(total)
 
             def _video_over_fg_image() -> Iterator[np.ndarray]:
                 for top_frame, m in zip(top, mask_iter):
@@ -617,6 +744,7 @@ def _video_over_fg_image() -> Iterator[np.ndarray]:
             iter(top._video)  # type: ignore[union-attr]
             top._frame_idx = 0  # type: ignore[union-attr]
             total = top.frame_count
+            mask_iter = _resolve_mask(total)
 
             def _fg_video_over_fg_image() -> Iterator[np.ndarray]:
                 for top_frame, m in zip(top, mask_iter):