added true iframe deletion datamoshing, fixed composite argument order

Author: readrawhex

dithertools/__init__.py

@@ -2,15 +2,16 @@
 
 Core classes: :class:`Image`, :class:`Video`, :class:`FrameGenerator`.
 Compositing: :func:`composite`, :func:`composite_color`.
-Datamoshing: :func:`iframe_delete`, :func:`pframe_dup`.
+Datamoshing: :func:`iframe_delete`, :func:`iframe_delete_optflow`,
+:func:`pframe_dup`.
 I/O shortcut: :func:`load` (returns ``Image`` or ``Video`` based on file type).
 Mask generation: :func:`random_matrix`, :func:`random_matrix_color`.
 Animated masks: :class:`MaskGenerator`.
 """
 
 __version__ = "0.1.0"
 
-from dithertools.datamosh import iframe_delete, pframe_dup
+from dithertools.datamosh import iframe_delete, iframe_delete_optflow, pframe_dup
 from dithertools.dither import composite, composite_color
 from dithertools.generators import FrameGenerator, MaskGenerator
 from dithertools.image import Image
@@ -26,6 +27,7 @@
     "composite",
     "composite_color",
     "iframe_delete",
+    "iframe_delete_optflow",
     "load",
     "pframe_dup",
     "random_matrix",

dithertools/datamosh.py

@@ -1,22 +1,27 @@
-"""Datamoshing effects via optical-flow-driven frame warping.
+"""Datamoshing effects.
 
-:func:`iframe_delete` simulates keyframe deletion at a scene cut: the last
-frame of scene A is warped forward by scene B's motion vectors, producing the
-characteristic "melting between scenes" look.
+:func:`iframe_delete` performs true MPEG-4 bitstream datamoshing by binary-
+patching encoded I-frames so the decoder treats them as P-frames, producing
+authentic block artefacts and colour smearing at the codec level.  Requires
+``ffmpeg`` and ``ffprobe`` in ``PATH``.
+
+:func:`iframe_delete_optflow` simulates I-frame deletion in decoded pixel
+space: the last frame of scene A is warped forward by scene B's motion
+vectors using OpenCV dense optical flow, no compressed video files required.
 
 :func:`pframe_dup` simulates P-frame duplication: a frozen reference frame is
 warped by the accumulated motion of subsequent frames, producing pixel
 streaking and blooming in the direction of movement.
-
-Both functions work entirely in decoded pixel space (``uint8`` RGB NumPy
-arrays) using OpenCV dense optical flow — no compressed video files are
-required.
 """
 
 from __future__ import annotations
 
 import itertools
+import json
+import os
+import subprocess
 import sys
+import tempfile
 from typing import TYPE_CHECKING, Iterator
 
 import cv2
@@ -138,7 +143,7 @@ def _frame_count(source: "Image | Video | FrameGenerator") -> int | None:
     return getattr(source, "frame_count", None)
 
 
-def iframe_delete(
+def iframe_delete_optflow(
     scene_a: "Image | Video | FrameGenerator",
     scene_b: "Image | Video | FrameGenerator | None" = None,
     *,
@@ -148,7 +153,7 @@ def iframe_delete(
     dis: bool = False,
     progress: bool = True,
 ) -> "Image | Video":
-    """Simulate I-frame deletion at a scene cut.
+    """Simulate I-frame deletion at a scene cut using optical flow.
 
     The last frame of ``scene_a`` (or the frame just before
     ``transition_frame`` in single-source mode) acts as the decoder's frozen
@@ -157,6 +162,9 @@ def iframe_delete(
     from the new scene's own motion — the visual equivalent of removing the
     keyframe that would otherwise reset the image to the new scene content.
 
+    Works entirely in decoded pixel space (``uint8`` RGB NumPy arrays) using
+    OpenCV dense optical flow — no compressed video files are required.
+
     **Output frames** — only the warped frames derived from ``scene_a``'s
     reference are returned.  ``scene_a``'s original frames are consumed solely
     to extract the reference and are *not* included in the output.  In
@@ -212,14 +220,14 @@ def iframe_delete(
     Example::
 
         from dithertools import Video
-        from dithertools.datamosh import iframe_delete
+        from dithertools.datamosh import iframe_delete_optflow
 
         # Two-source: apply scene_b's motion to scene_a's last frame
-        result = iframe_delete(Video("scene_a.mp4"), Video("scene_b.mp4"))
+        result = iframe_delete_optflow(Video("scene_a.mp4"), Video("scene_b.mp4"))
         result.save("moshed.mp4")
 
         # Single-source: delete the keyframe at frame 60
-        result = iframe_delete(Video("clip.mp4"), transition_frame=60)
+        result = iframe_delete_optflow(Video("clip.mp4"), transition_frame=60)
         result.save("moshed.mp4")
     """
     from dithertools.image import Image
@@ -301,6 +309,240 @@ def _single_source() -> Iterator[np.ndarray]:
     return Video.from_frames(_progress_iter(_single_source(), total, progress), fps=fps)
 
 
+def _aviglitch(data: bytearray, kf_positions: list[tuple[int, int]]) -> None:
+    """Zero non-first keyframe payloads and clear AVIIF_KEYFRAME in idx1.
+
+    kf_positions: list of (abs_file_pos, payload_size) for each video keyframe
+    in stream order, as returned by ffprobe show_packets.  The first entry is
+    kept intact; all subsequent ones are zeroed (triggering decoder error
+    concealment — the decoder copies the last decoded frame instead of
+    resetting to the intra content) and their AVIIF_KEYFRAME flag in the AVI
+    idx1 index is cleared (preventing decoders from flushing their reference
+    buffer at the transition point).
+    """
+    AVIIF_KEYFRAME = 0x10
+
+    for pos, size in kf_positions[1:]:
+        data[pos : pos + size] = bytes(size)
+
+    idx1_off = bytes(data).find(b"idx1")
+    if idx1_off == -1:
+        return
+
+    idx1_size = int.from_bytes(data[idx1_off + 4 : idx1_off + 8], "little")
+    idx1_start = idx1_off + 8
+    n_entries = idx1_size // 16
+    kf_count = 0
+
+    for i in range(n_entries):
+        e = idx1_start + i * 16
+        if e + 16 > len(data):
+            break
+        fcc = bytes(data[e : e + 4])
+        flags = int.from_bytes(data[e + 4 : e + 8], "little")
+        if fcc[2:4] not in (b"dc", b"db") or not (flags & AVIIF_KEYFRAME):
+            continue
+        kf_count += 1
+        if kf_count > 1:
+            data[e + 4 : e + 8] = (flags & ~AVIIF_KEYFRAME).to_bytes(4, "little")
+
+
+def iframe_delete(
+    scene_a: "Image | Video | FrameGenerator",
+    scene_b: "Image | Video | FrameGenerator | None" = None,
+    *,
+    transition_frame: int | None = None,
+    quality: int = 3,
+    progress: bool = True,
+) -> "Video":
+    """True MPEG-4 bitstream datamosh at a scene cut.
+
+    Encodes all input frames to a temporary MPEG-4 AVI, forces I-frames only
+    at frame 0 and the transition point, then applies the aviglitch-style
+    effect to the transition I-frame: its payload is zeroed (the decoder
+    applies error concealment — copying scene A's last frame — instead of
+    resetting to scene B's intra content) and its ``AVIIF_KEYFRAME`` flag in
+    the AVI ``idx1`` index is cleared (preventing decoders from flushing
+    their reference buffer at the cut).  Scene B's subsequent P-frames then
+    apply their motion vectors to scene A's pixel state, producing the
+    characteristic datamosh: block artefacts, colour smearing, and
+    macroblocks at the codec level.
+
+    Requires ``ffmpeg`` and ``ffprobe`` in ``PATH``.
+
+    **Mode selection** — supply exactly one of ``scene_b`` or
+    ``transition_frame``:
+
+    - *Two-source mode* (``scene_b`` provided): all frames from ``scene_a``
+      followed by all frames from ``scene_b`` are returned.  The transition
+      artefact appears at the scene cut; ``scene_a``'s frames provide the
+      visual anchor the effect needs.
+    - *Single-source mode* (``transition_frame`` provided): frames before the
+      cut pass through unchanged; frames from ``transition_frame`` onward
+      carry the datamosh artefacts.
+
+    Args:
+        scene_a: The donor scene.  In two-source mode, its last frame becomes
+            the motion reference for ``scene_b``.  In single-source mode,
+            this is the full source.  May be an
+            :class:`~dithertools.image.Image`,
+            :class:`~dithertools.video.Video`, or
+            :class:`~dithertools.generators.FrameGenerator`.
+        scene_b: The motion-donor scene (two-source mode).  ``None`` when
+            using single-source mode.
+        transition_frame: Index at which the I-frame deletion occurs
+            (single-source mode).  ``None`` when ``scene_b`` is provided.
+        quality: MPEG-4 ``qscale:v`` value.  ``1`` = best quality / least
+            artefacts; ``31`` = worst quality / most macroblocks.  Defaults
+            to ``3``.
+        progress: Display a progress spinner on stderr while encoding frames.
+            Only active when stderr is a TTY.  Defaults to ``True``.
+
+    Returns:
+        A memory-backed :class:`~dithertools.video.Video` containing all
+        output frames — scene_a + moshed scene_b in two-source mode, or the
+        full source with post-transition artefacts in single-source mode.
+
+    Raises:
+        ValueError: If both ``scene_b`` and ``transition_frame`` are given,
+            or if neither is given.
+        RuntimeError: If ``ffmpeg`` or ``ffprobe`` is not found in ``PATH``,
+            or if encoding or probing fails.
+
+    Example::
+
+        from dithertools import Video
+        from dithertools.datamosh import iframe_delete
+
+        # Two-source: authentic datamosh at the cut between two scenes
+        result = iframe_delete(Video("scene_a.mp4"), Video("scene_b.mp4"))
+        result.save("moshed.mp4")
+
+        # Single-source: delete the keyframe at frame 60
+        result = iframe_delete(Video("clip.mp4"), transition_frame=60)
+        result.save("moshed.mp4")
+    """
+    from dithertools.image import Image
+    from dithertools.video import Video
+
+    if scene_b is not None and transition_frame is not None:
+        raise ValueError(
+            "Provide either scene_b (two-source mode) or transition_frame "
+            "(single-source mode), not both."
+        )
+    if scene_b is None and transition_frame is None:
+        raise ValueError(
+            "Provide either scene_b (two-source mode) or transition_frame "
+            "(single-source mode)."
+        )
+
+    def _collect(src: "Image | Video | FrameGenerator") -> list[np.ndarray]:
+        if isinstance(src, Image):
+            return [src.data]
+        return list(src)
+
+    if scene_b is not None:
+        a_frames = _collect(scene_a)
+        b_frames = _collect(scene_b)
+        all_frames = a_frames + b_frames
+        transition = len(a_frames)
+        fps = _fps(scene_b) or _fps(scene_a) or 24.0
+    else:
+        all_frames = _collect(scene_a)
+        transition = transition_frame  # type: ignore[assignment]
+        fps = _fps(scene_a) or 24.0
+
+    if not all_frames:
+        raise ValueError("source yielded no frames")
+
+    h, w = all_frames[0].shape[:2]
+    total = len(all_frames)
+
+    tmp_path: str | None = None
+    try:
+        fd, tmp_path = tempfile.mkstemp(suffix=".avi")
+        os.close(fd)
+
+        ffmpeg_cmd = [
+            "ffmpeg", "-y",
+            "-f", "rawvideo",
+            "-pix_fmt", "rgb24",
+            "-s", f"{w}x{h}",
+            "-r", str(fps),
+            "-i", "pipe:0",
+            "-c:v", "mpeg4",
+            "-qscale:v", str(quality),
+            "-g", "99999",
+            "-sc_threshold", "0",
+            "-force_key_frames", f"expr:eq(n,0)+eq(n,{transition})",
+            tmp_path,
+        ]
+        proc = subprocess.Popen(
+            ffmpeg_cmd,
+            stdin=subprocess.PIPE,
+            stderr=subprocess.DEVNULL,
+        )
+        assert proc.stdin is not None
+        for frame in _progress_iter(iter(all_frames), total, progress):
+            proc.stdin.write(frame.tobytes())
+        proc.stdin.close()
+        proc.wait()
+        if proc.returncode != 0:
+            raise RuntimeError(f"ffmpeg failed with exit code {proc.returncode}")
+
+        probe_cmd = [
+            "ffprobe", "-v", "error",
+            "-select_streams", "v:0",
+            "-show_packets",
+            "-show_entries", "packet=flags,pos,size",
+            "-of", "json",
+            tmp_path,
+        ]
+        probe = subprocess.run(probe_cmd, capture_output=True, text=True, check=True)
+        pkts = json.loads(probe.stdout).get("packets", [])
+
+        kf_positions = [
+            (int(p["pos"]), int(p["size"]))
+            for p in pkts
+            if "K" in p.get("flags", "")
+        ]
+
+        with open(tmp_path, "r+b") as f:
+            avi_data = bytearray(f.read())
+
+        _aviglitch(avi_data, kf_positions)
+
+        with open(tmp_path, "wb") as f:
+            f.write(avi_data)
+
+        # Decode via ffmpeg, which applies error concealment on the zeroed
+        # I-frame (copying scene A's last decoded frame into the reference
+        # buffer) before processing scene B's P-frames.  OpenCV's
+        # VideoCapture treats the decode error as fatal and stops reading,
+        # which would truncate the output.
+        decode_cmd = [
+            "ffmpeg", "-v", "error",
+            "-i", tmp_path,
+            "-f", "rawvideo", "-pix_fmt", "rgb24",
+            "pipe:1",
+        ]
+        decode_proc = subprocess.run(decode_cmd, capture_output=True)
+        raw = decode_proc.stdout
+        frame_size = h * w * 3
+        n_frames = len(raw) // frame_size
+        frames = [
+            np.frombuffer(
+                raw[i * frame_size : (i + 1) * frame_size], dtype=np.uint8
+            ).reshape(h, w, 3).copy()
+            for i in range(n_frames)
+        ]
+        return Video.from_frames(iter(frames), fps=fps)
+
+    finally:
+        if tmp_path and os.path.exists(tmp_path):
+            os.unlink(tmp_path)
+
+
 def pframe_dup(
     source: "Image | Video | FrameGenerator",
     transition_frame: int | None = None,