decode_image_bit.py

"""Local copy of XPolicyLab's image-bit codec, for reference.

Source: ``XPolicyLab.utils.process_data``
(https://github.com/XPolicyLab/XPolicyLab) — ``decode_image_bit``,
``encode_image_bit``, and ``images_encoding``. Prefer importing that
module when the package is available; this file is the same pair kept
next to the dataset so collectors and loaders can see how RoboTwin /
XPolicyLab HDF5 camera buffers must be written and read.
"""

import cv2
import numpy as np

# Two image bit stream formats exist in RoboTwin / XPolicyLab data, and both
# decode to RGB.
#
#   legacy    A JPEG written by handing an RGB array straight to cv2.imencode,
#             which reads its input as BGR. The stored bytes are therefore
#             channel-reversed with respect to the JPEG standard: cv2.imdecode
#             reverses them a second time and returns the original RGB, but any
#             conforming decoder (PIL, ffmpeg, a browser) shows red and blue
#             swapped. All data collected before the marker existed is this.
#   standard  A conforming RGB JPEG written by encode_image_bit, stamped with a
#             COM segment holding _RGB_MARKER_PAYLOAD. cv2.imdecode returns BGR
#             for it, so decoding needs exactly one swap.
#
# The marker lives inside the buffer rather than in an HDF5 attribute because
# decode_image_bit is only ever handed one buffer and has no file context to
# consult. COM is a standard segment carrying its own length, so every decoder
# skips it; PIL exposes it as Image.open(...).info["comment"], which lets a
# non-OpenCV reader tell the two formats apart as well.
_RGB_MARKER_PAYLOAD = b"XPL-RGB1"

_JPEG_PAD = 0xFF
_JPEG_STUFF = 0x00
_JPEG_TEM = 0x01
_JPEG_RST_FIRST = 0xD0
_JPEG_RST_LAST = 0xD7
_JPEG_SOI = 0xD8
_JPEG_EOI = 0xD9
_JPEG_SOS = 0xDA
_JPEG_APP0 = 0xE0
_JPEG_COM = 0xFE

_COM_SEGMENT_LENGTH_BYTES = 2
_RGB_MARKER_SEGMENT_LENGTH = len(_RGB_MARKER_PAYLOAD) + _COM_SEGMENT_LENGTH_BYTES


def _normalize_image_buffer(image_bit):
    """Reduce any container holding one encoded buffer to bytes or a contiguous
    uint8 array, so the decode and encode paths agree on what they are looking
    at."""
    if isinstance(image_bit, np.ndarray) and image_bit.dtype.kind in {"S", "U"}:
        image_bit = image_bit.item() if image_bit.ndim == 0 else image_bit.tobytes()

    if isinstance(image_bit, str):
        image_bit = image_bit.encode("utf-8")
    elif isinstance(image_bit, memoryview):
        image_bit = image_bit.tobytes()

    if isinstance(image_bit, (bytes, bytearray)):
        # Fixed-width HDF5 byte columns pad the tail with NUL.
        image_bit = image_bit.rstrip(b"\0")
    elif isinstance(image_bit, np.ndarray):
        image_bit = np.ascontiguousarray(image_bit)

    return image_bit


def _byte_view(image_bit) -> memoryview:
    """A flat, byte-addressable view of one normalized buffer, without copying."""
    if isinstance(image_bit, (bytes, bytearray, memoryview)):
        return memoryview(image_bit).cast("B")

    array = np.ascontiguousarray(image_bit)

    if array.dtype != np.uint8:
        return memoryview(array.tobytes())

    return memoryview(array.reshape(-1))


def _has_rgb_marker(image_bit) -> bool:
    """
    Whether a buffer is a standard RGB JPEG stamped by `encode_image_bit`.

    The header segments are walked rather than probing a fixed offset, so the
    marker is still found if another tool inserted a segment of its own ahead of
    it. Walking stops at SOS, which bounds the work to the header and keeps the
    entropy-coded payload — the bulk of a frame — untouched.
    """
    view = _byte_view(_normalize_image_buffer(image_bit))
    size = len(view)

    if size < 4 or view[0] != _JPEG_PAD or view[1] != _JPEG_SOI:
        return False

    offset = 2

    while offset + 4 <= size:
        if view[offset] != _JPEG_PAD:
            return False

        marker = view[offset + 1]

        if marker == _JPEG_PAD:  # fill byte ahead of the real marker
            offset += 1
            continue

        if marker == _JPEG_STUFF:  # only valid inside entropy-coded data
            return False

        if marker in (_JPEG_TEM, _JPEG_SOI, _JPEG_EOI) or (
            _JPEG_RST_FIRST <= marker <= _JPEG_RST_LAST
        ):
            offset += 2
            continue

        if marker == _JPEG_SOS:  # entropy-coded data starts; no header left
            return False

        length = (view[offset + 2] << 8) | view[offset + 3]

        if length < _COM_SEGMENT_LENGTH_BYTES:
            return False

        if marker == _JPEG_COM and length == _RGB_MARKER_SEGMENT_LENGTH:
            payload = view[offset + 4 : offset + 2 + length]
            if bytes(payload) == _RGB_MARKER_PAYLOAD:
                return True

        offset += 2 + length

    return False


def _insert_rgb_marker(jpeg_bytes: bytes) -> bytes:
    """
    Stamp a JPEG COM segment marking the stream as standard RGB.

    The segment goes after the JFIF APP0 rather than straight after SOI: JFIF
    requires APP0 to follow SOI immediately, and a hand-written parser that
    relies on that layout would otherwise trip over the marker.
    """
    segment = (
        bytes([_JPEG_PAD, _JPEG_COM])
        + _RGB_MARKER_SEGMENT_LENGTH.to_bytes(2, "big")
        + _RGB_MARKER_PAYLOAD
    )

    offset = 2

    if len(jpeg_bytes) >= 6 and jpeg_bytes[2:4] == bytes([_JPEG_PAD, _JPEG_APP0]):
        app0_end = 4 + int.from_bytes(jpeg_bytes[4:6], "big")
        if app0_end <= len(jpeg_bytes):
            offset = app0_end

    return jpeg_bytes[:offset] + segment + jpeg_bytes[offset:]


def _decode_single_image_bit(image_bit):
    """Decode one encoded image buffer into an HWC uint8 RGB array."""
    image_bit = _normalize_image_buffer(image_bit)

    image = cv2.imdecode(np.frombuffer(image_bit, np.uint8), cv2.IMREAD_COLOR)

    if image is None:
        raise ValueError(
            f"Failed to decode image bits (type={type(image_bit).__name__}, "
            f"size={getattr(image_bit, 'size', len(image_bit) if hasattr(image_bit, '__len__') else '?')})."
        )

    # cv2.imdecode returns BGR for a conforming JPEG. A marked buffer is one, so
    # it needs the swap; a legacy buffer stores reversed channels that imdecode
    # has already reversed back, so swapping it is what would break the order.
    # Either way the caller gets RGB, which is why no caller ever swaps.
    if _has_rgb_marker(image_bit):
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    return image


def _decode_image_bit_sequence(image_bits):
    frames = []

    for index, image_bit in enumerate(image_bits):
        try:
            frames.append(decode_image_bit(image_bit))
        except ValueError as exc:
            raise ValueError(f"Frame {index}: {exc}") from exc

    if not frames:
        return np.zeros((0,), dtype=np.uint8)

    return np.stack(frames, axis=0)


def decode_image_bit(image_bits):
    """
    Decode encoded image bit stream(s) into uint8 RGB image array(s).

    Copied from ``XPolicyLab.utils.process_data.decode_image_bit``.

    The output is RGB for both stored byte formats — the legacy
    channel-reversed streams and the standard RGB JPEGs written by
    `encode_image_bit` — because this function reads the marker that tells them
    apart and swaps only where a swap is owed. Never add a COLOR_BGR2RGB after
    this function to "correct" the output: the two formats are
    indistinguishable to the eye and to any single sample, so a caller-side
    swap is right on at most one of them and silently wrong on the other. This
    is also why hand-rolled decoding is unsupported, PIL included: PIL reads the
    standard format correctly and the legacy format reversed.

    Deliberately converting the RGB result to BGR is a different thing and is
    allowed where a checkpoint was trained on BGR data; that must be an opt-in
    documented in the adapter (see Dexora_1B's `input_color_order`), never a
    silent fix applied at the decode site.

    Values that are already decoded are returned unchanged, so this function is
    safe to call on an observation or trajectory field without knowing whether
    the producer encoded it.

    Dispatch is on dtype first, then ndim:
        - bytes / bytearray / memoryview / str  -> one encoded buffer
        - ndarray of dtype kind 'S', 'U', 'O'   -> sequence of encoded buffers,
                                                   or one buffer when 0-d
        - uint8 ndarray, ndim == 1              -> one encoded buffer
        - uint8 ndarray, ndim == 2              -> (T, N) stack of encoded buffers
        - uint8 ndarray, ndim >= 3              -> already decoded, returned as is
        - ndarray of any other dtype            -> already decoded, returned as is
        - list / tuple                          -> element-wise, stacked on axis 0

    Grayscale (H, W) uint8 images are not supported: a 2-D uint8 array is always
    read as a stack of encoded buffers.

    Raises:
        ValueError: if any buffer fails to decode.
    """
    if isinstance(image_bits, (bytes, bytearray, memoryview, str)):
        return _decode_single_image_bit(image_bits)

    if isinstance(image_bits, np.ndarray):
        if image_bits.dtype.kind in {"S", "U", "O"}:
            if image_bits.ndim == 0:
                return _decode_single_image_bit(image_bits.item())
            return _decode_image_bit_sequence(image_bits)

        if image_bits.dtype == np.uint8:
            if image_bits.ndim == 1:
                return _decode_single_image_bit(image_bits)
            if image_bits.ndim == 2:
                return _decode_image_bit_sequence(image_bits)

        return image_bits

    if isinstance(image_bits, (list, tuple)):
        return _decode_image_bit_sequence(image_bits)

    return _decode_single_image_bit(image_bits)


def _encode_single_image_bit(image, quality=None):
    """Encode one uint8 RGB image into a marked standard JPEG buffer."""
    array = np.asarray(image)

    if array.dtype != np.uint8:
        raise TypeError(
            f"Expected a uint8 RGB image, got dtype {array.dtype}. Scale and cast "
            "before encoding — a float array is ambiguous between [0, 1] and "
            "[0, 255], and guessing is how channel and range bugs start."
        )

    if array.ndim != 3 or array.shape[-1] != 3:
        raise ValueError(
            f"Expected one (H, W, 3) RGB image, got shape {array.shape}. "
            "Grayscale is not supported, and an RGBA frame must be sliced to "
            "[..., :3] first."
        )

    params = [] if quality is None else [int(cv2.IMWRITE_JPEG_QUALITY), int(quality)]
    bgr = cv2.cvtColor(np.ascontiguousarray(array), cv2.COLOR_RGB2BGR)
    success, encoded_image = cv2.imencode(".jpg", bgr, params)

    if not success:
        raise ValueError(f"Failed to JPEG-encode an RGB frame of shape {array.shape}.")

    return _insert_rgb_marker(encoded_image.tobytes())


def _passthrough_encoded_buffer(image_bit):
    """
    Return an already-encoded buffer unchanged, refusing legacy bytes.

    Relabelling a legacy buffer as standard RGB is the one way this scheme can
    silently reverse channels, so an unmarked buffer is an error rather than a
    no-op.
    """
    image_bit = _normalize_image_buffer(image_bit)

    if not _has_rgb_marker(image_bit):
        raise ValueError(
            "Refusing to pass through an unmarked image buffer: it is in the "
            "legacy channel-reversed format, and stamping it as standard RGB "
            "would reverse its channels. Decode it with decode_image_bit and "
            "encode the resulting array if a lossy re-encode is intended."
        )

    if isinstance(image_bit, (bytes, bytearray)):
        return bytes(image_bit)

    return np.asarray(image_bit).reshape(-1).tobytes()


def _encode_image_bit_sequence(images, quality):
    buffers = []

    for index, image in enumerate(images):
        try:
            buffers.append(encode_image_bit(image, quality=quality))
        except (ValueError, TypeError) as exc:
            raise type(exc)(f"Frame {index}: {exc}") from exc

    return buffers


def encode_image_bit(images, quality=None):
    """
    Encode uint8 RGB image(s) into the image bit stream(s) that
    `decode_image_bit` reads back, and mark them as standard RGB.

    Copied from ``XPolicyLab.utils.process_data.encode_image_bit``.

    This is the inverse of `decode_image_bit` and the only supported encoder.
    The buffers it produces are conforming RGB JPEGs, so PIL, ffmpeg, a dataset
    viewer or a browser all show the right colours, and they carry a JPEG COM
    marker so `decode_image_bit` knows not to treat them as legacy
    channel-reversed data. Encoding by hand with `cv2.imencode` skips the
    marker and produces a buffer that reads back with red and blue swapped.

    Dispatch mirrors `decode_image_bit`, on dtype first, then ndim:
        - uint8 ndarray, ndim == 3      -> one (H, W, 3) RGB image  -> bytes
        - uint8 ndarray, ndim == 4      -> (T, H, W, 3) stack       -> list[bytes]
        - list / tuple                  -> element-wise             -> list[bytes]
        - bytes / bytearray / memoryview / str, uint8 ndarray with ndim == 1,
          and ndarray of dtype kind 'S', 'U', 'O'
                                        -> already encoded, returned as bytes
        - uint8 ndarray, ndim == 2      -> stack of encoded buffers -> list[bytes]

    Grayscale is not supported, matching `decode_image_bit`: a 2-D uint8 array
    is always read as a stack of already-encoded buffers, never as one image.

    Args:
        images: One RGB image, a sequence of them, or already-encoded buffers.
        quality: Optional JPEG quality; OpenCV's default is used when omitted.

    Returns:
        bytes for a single image, list[bytes] for a sequence.

    Raises:
        TypeError: if image data is not uint8.
        ValueError: if a frame has an unsupported shape, fails to encode, or is
            an already-encoded buffer in the legacy format.
    """
    if isinstance(images, (bytes, bytearray, memoryview, str)):
        return _passthrough_encoded_buffer(images)

    if isinstance(images, np.ndarray):
        if images.dtype.kind in {"S", "U", "O"}:
            if images.ndim == 0:
                return _passthrough_encoded_buffer(images.item())
            return _encode_image_bit_sequence(images, quality)

        if images.dtype == np.uint8:
            if images.ndim == 0:
                raise ValueError(
                    "Expected an RGB image or a sequence of them, got a 0-d uint8 "
                    "array."
                )
            if images.ndim == 1:
                return _passthrough_encoded_buffer(images)
            if images.ndim == 3:
                return _encode_single_image_bit(images, quality)
            if images.ndim == 2:
                # Ambiguous shape, so say which reading was attempted: the most
                # likely way to land here is passing a grayscale frame.
                try:
                    return _encode_image_bit_sequence(images, quality)
                except (ValueError, TypeError) as exc:
                    raise ValueError(
                        "A 2-D uint8 array is read as a stack of already-encoded "
                        "buffers, never as a grayscale image, and this one is not "
                        f"a valid stack: {exc}"
                    ) from exc
            return _encode_image_bit_sequence(images, quality)

        raise TypeError(
            f"Expected uint8 image data, got an array of dtype {images.dtype}. "
            "Scale and cast before encoding."
        )

    if isinstance(images, (list, tuple)):
        return _encode_image_bit_sequence(images, quality)

    return _encode_single_image_bit(images, quality)


def images_encoding(imgs):
    """
    JPEG-encode RGB frames for storage in the XPolicyLab layout and report the
    width of the fixed-size HDF5 byte column they fit into, which pads the
    shorter buffers with NUL on write.

    Copied from ``XPolicyLab.utils.process_data.images_encoding``.

    Feed this RGB, never BGR and never a channel-swapped array. Encoding runs
    through `encode_image_bit`, so the buffers are conforming RGB JPEGs marked
    for `decode_image_bit`.
    """
    encode_data = encode_image_bit(imgs)

    if isinstance(encode_data, bytes):
        encode_data = [encode_data]

    max_len = max((len(jpeg_data) for jpeg_data in encode_data), default=0)

    return encode_data, max_len