From 6b875035cbf0d57e1de15ba554744d802a76c469 Mon Sep 17 00:00:00 2001 From: SuperOptimizer <156155735+SuperOptimizer@users.noreply.github.com> Date: Wed, 14 Jan 2026 20:17:34 +0000 Subject: [PATCH] h264 changes --- src/zarr/codecs/__init__.py | 7 + src/zarr/codecs/blosc2.py | 318 ++++++++++++++++++++++++++ src/zarr/codecs/numcodecs/__init__.py | 2 + src/zarr/codecs/numcodecs/_codecs.py | 4 + 4 files changed, 331 insertions(+) create mode 100644 src/zarr/codecs/blosc2.py diff --git a/src/zarr/codecs/__init__.py b/src/zarr/codecs/__init__.py index 4c621290e7..aa79e3e1b3 100644 --- a/src/zarr/codecs/__init__.py +++ b/src/zarr/codecs/__init__.py @@ -1,6 +1,7 @@ from __future__ import annotations from zarr.codecs.blosc import BloscCname, BloscCodec, BloscShuffle +from zarr.codecs.blosc2 import Blosc2Cname, Blosc2Codec, Blosc2Shuffle from zarr.codecs.bytes import BytesCodec, Endian from zarr.codecs.crc32c_ import Crc32cCodec from zarr.codecs.gzip import GzipCodec @@ -15,6 +16,7 @@ AsType, BitRound, Blosc, + Blosc2 as NumcodecsBlosc2, Delta, FixedScaleOffset, Fletcher32, @@ -37,6 +39,9 @@ "BloscCname", "BloscCodec", "BloscShuffle", + "Blosc2Cname", + "Blosc2Codec", + "Blosc2Shuffle", "BytesCodec", "Crc32cCodec", "Endian", @@ -50,6 +55,7 @@ ] register_codec("blosc", BloscCodec) +register_codec("blosc2", Blosc2Codec) register_codec("bytes", BytesCodec) # compatibility with earlier versions of ZEP1 @@ -74,6 +80,7 @@ register_codec("numcodecs.astype", AsType, qualname="zarr.codecs.numcodecs.AsType") register_codec("numcodecs.bitround", BitRound, qualname="zarr.codecs.numcodecs.BitRound") register_codec("numcodecs.blosc", Blosc, qualname="zarr.codecs.numcodecs.Blosc") +register_codec("numcodecs.blosc2", NumcodecsBlosc2, qualname="zarr.codecs.numcodecs.Blosc2") register_codec("numcodecs.delta", Delta, qualname="zarr.codecs.numcodecs.Delta") register_codec( "numcodecs.fixedscaleoffset", diff --git a/src/zarr/codecs/blosc2.py b/src/zarr/codecs/blosc2.py new file mode 100644 index 0000000000..f423c8a8f5 --- /dev/null +++ b/src/zarr/codecs/blosc2.py @@ -0,0 +1,318 @@ +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field, replace +from enum import Enum +from functools import cached_property +from typing import TYPE_CHECKING, Final, Literal, NotRequired, TypedDict + +from zarr.abc.codec import BytesBytesCodec +from zarr.core.buffer.cpu import as_numpy_array_wrapper +from zarr.core.common import JSON, NamedRequiredConfig, parse_enum, parse_named_configuration +from zarr.core.dtype.common import HasItemSize + +if TYPE_CHECKING: + from typing import Self + + from zarr.core.array_spec import ArraySpec + from zarr.core.buffer import Buffer + +Shuffle = Literal["noshuffle", "shuffle", "bitshuffle"] +"""The shuffle values permitted for the blosc2 codec""" + +SHUFFLE: Final = ("noshuffle", "shuffle", "bitshuffle") + +# Standard blosc2 compressors plus openh264 +CName = Literal["lz4", "lz4hc", "blosclz", "zlib", "zstd", "openh264"] +"""The codec identifiers used in the blosc2 codec""" + + +class Blosc2ConfigV3(TypedDict): + """Configuration for the V3 Blosc2 codec""" + + cname: CName + clevel: int + shuffle: Shuffle + blocksize: int + typesize: int + qp: NotRequired[int] # H.264 quantization parameter (0-51), only for openh264 + + +class Blosc2JSON_V3(NamedRequiredConfig[Literal["blosc2"], Blosc2ConfigV3]): + """ + The JSON form of the Blosc2 codec in Zarr V3. + """ + + +class Blosc2Shuffle(Enum): + """ + Enum for shuffle filter used by blosc2. + """ + + noshuffle = "noshuffle" + shuffle = "shuffle" + bitshuffle = "bitshuffle" + + @classmethod + def from_int(cls, num: int) -> Blosc2Shuffle: + blosc2_shuffle_int_to_str = { + 0: "noshuffle", + 1: "shuffle", + 2: "bitshuffle", + } + if num not in blosc2_shuffle_int_to_str: + raise ValueError(f"Value must be between 0 and 2. Got {num}.") + return Blosc2Shuffle[blosc2_shuffle_int_to_str[num]] + + +class Blosc2Cname(Enum): + """ + Enum for compression library used by blosc2. + """ + + lz4 = "lz4" + lz4hc = "lz4hc" + blosclz = "blosclz" + zstd = "zstd" + zlib = "zlib" + openh264 = "openh264" + + +def parse_typesize(data: JSON) -> int: + if isinstance(data, int): + if data > 0: + return data + else: + raise ValueError( + f"Value must be greater than 0. Got {data}, which is less or equal to 0." + ) + raise TypeError(f"Value must be an int. Got {type(data)} instead.") + + +def parse_clevel(data: JSON) -> int: + if isinstance(data, int): + return data + raise TypeError(f"Value should be an int. Got {type(data)} instead.") + + +def parse_blocksize(data: JSON) -> int: + if isinstance(data, int): + return data + raise TypeError(f"Value should be an int. Got {type(data)} instead.") + + +@dataclass(frozen=True) +class Blosc2Codec(BytesBytesCodec): + """ + Blosc2 compression codec for zarr. + + Blosc2 is the next-generation version of Blosc, providing improved + compression algorithms and support for user-defined codecs like OpenH264. + + Attributes + ---------- + is_fixed_size : bool + Always False for Blosc2 codec, as compression produces variable-sized output. + typesize : int + The data type size in bytes used for shuffle filtering. + cname : Blosc2Cname + The compression algorithm being used (lz4, lz4hc, blosclz, zlib, zstd, or openh264). + clevel : int + The compression level (0-9). + shuffle : Blosc2Shuffle + The shuffle filter mode (noshuffle, shuffle, or bitshuffle). + Note: Ignored for openh264 codec. + blocksize : int + The size of compressed blocks in bytes (0 for automatic). + Note: For openh264, automatically set to chunk size. + + Parameters + ---------- + typesize : int, optional + The data type size in bytes. Default: 1. + cname : Blosc2Cname or str, optional + The compression algorithm to use. Default: 'zstd'. + Use 'openh264' for H.264 video compression of 3D cubic chunks. + clevel : int, optional + The compression level, from 0 (no compression) to 9 (maximum compression). + Default: 5. + shuffle : Blosc2Shuffle or str, optional + The shuffle filter to apply before compression. Default: 'noshuffle'. + blocksize : int, optional + The requested size of compressed blocks in bytes. Default: 0. + + qp : int, optional + Quantization parameter for H.264 codec (0-51). Only used with cname='openh264'. + 0 = highest quality/largest files, 51 = lowest quality/smallest files. + Default: 26. + + Notes + ----- + When using cname='openh264': + - Data must be uint8 (typesize=1) + - Chunks must be cubic (NxNxN where N is even) + - Compression is lossy (H.264 video codec) + - blocksize is automatically set to match chunk size + - qp controls quality/compression tradeoff (0-51) + - Provides excellent compression for volumetric data (CT/MRI scans, etc.) + + Examples + -------- + Create a Blosc2 codec with default settings: + + >>> codec = Blosc2Codec() + + Create a codec with H.264 compression for 3D volumes: + + >>> codec = Blosc2Codec(cname='openh264') + + Create H.264 codec with high compression (QP=39): + + >>> codec = Blosc2Codec(cname='openh264', qp=39) + + See Also + -------- + BloscCodec : Original Blosc codec (v1) + """ + + _tunable_attrs: set[Literal["typesize", "shuffle"]] = field(init=False) + is_fixed_size = False + + typesize: int + cname: Blosc2Cname + clevel: int + shuffle: Blosc2Shuffle + blocksize: int + qp: int + + def __init__( + self, + *, + typesize: int | None = None, + cname: Blosc2Cname | CName = Blosc2Cname.zstd, + clevel: int = 5, + shuffle: Blosc2Shuffle | Shuffle | None = None, + blocksize: int = 0, + qp: int = 26, + ) -> None: + object.__setattr__(self, "_tunable_attrs", set()) + + # For openh264, force typesize=1 and noshuffle + cname_parsed = parse_enum(cname, Blosc2Cname) + is_openh264 = cname_parsed == Blosc2Cname.openh264 + + if typesize is None: + typesize = 1 + if not is_openh264: + self._tunable_attrs.update({"typesize"}) + + if shuffle is None: + shuffle = Blosc2Shuffle.noshuffle if is_openh264 else Blosc2Shuffle.bitshuffle + if not is_openh264: + self._tunable_attrs.update({"shuffle"}) + + # Validate openh264 constraints + if is_openh264: + if typesize != 1: + raise ValueError("openh264 codec requires typesize=1 (uint8 data)") + if not (0 <= qp <= 51): + raise ValueError(f"QP must be 0-51 for openh264. Got {qp}") + + typesize_parsed = parse_typesize(typesize) + clevel_parsed = parse_clevel(clevel) + shuffle_parsed = parse_enum(shuffle, Blosc2Shuffle) + blocksize_parsed = parse_blocksize(blocksize) + + object.__setattr__(self, "typesize", typesize_parsed) + object.__setattr__(self, "cname", cname_parsed) + object.__setattr__(self, "clevel", clevel_parsed) + object.__setattr__(self, "shuffle", shuffle_parsed) + object.__setattr__(self, "blocksize", blocksize_parsed) + object.__setattr__(self, "qp", qp) + + @classmethod + def from_dict(cls, data: dict[str, JSON]) -> Self: + _, configuration_parsed = parse_named_configuration(data, "blosc2") + return cls(**configuration_parsed) # type: ignore[arg-type] + + def to_dict(self) -> dict[str, JSON]: + config: Blosc2ConfigV3 = { + "typesize": self.typesize, + "cname": self.cname.value, + "clevel": self.clevel, + "shuffle": self.shuffle.value, + "blocksize": self.blocksize, + } + # Include QP for openh264 codec + if self.cname == Blosc2Cname.openh264: + config["qp"] = self.qp + result: Blosc2JSON_V3 = { + "name": "blosc2", + "configuration": config, + } + return result # type: ignore[return-value] + + def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self: + """ + Create a new codec with typesize and shuffle parameters adjusted + according to the size of each element in the data type. + """ + # For openh264, don't evolve - keep typesize=1 + if self.cname == Blosc2Cname.openh264: + return self + + item_size = 1 + if isinstance(array_spec.dtype, HasItemSize): + item_size = array_spec.dtype.item_size + new_codec = self + if "typesize" in self._tunable_attrs: + new_codec = replace(new_codec, typesize=item_size) + if "shuffle" in self._tunable_attrs: + new_codec = replace( + new_codec, + shuffle=(Blosc2Shuffle.bitshuffle if item_size == 1 else Blosc2Shuffle.shuffle), + ) + + return new_codec + + @cached_property + def _blosc2_codec(self): + from numcodecs.blosc2 import Blosc2 + + map_shuffle_str_to_int = { + Blosc2Shuffle.noshuffle: 0, + Blosc2Shuffle.shuffle: 1, + Blosc2Shuffle.bitshuffle: 2, + } + return Blosc2( + cname=self.cname.value, + clevel=self.clevel, + shuffle=map_shuffle_str_to_int[self.shuffle], + blocksize=self.blocksize, + typesize=self.typesize, + qp=self.qp, + ) + + async def _decode_single( + self, + chunk_bytes: Buffer, + chunk_spec: ArraySpec, + ) -> Buffer: + return await asyncio.to_thread( + as_numpy_array_wrapper, self._blosc2_codec.decode, chunk_bytes, chunk_spec.prototype + ) + + async def _encode_single( + self, + chunk_bytes: Buffer, + chunk_spec: ArraySpec, + ) -> Buffer | None: + return await asyncio.to_thread( + lambda chunk: chunk_spec.prototype.buffer.from_bytes( + self._blosc2_codec.encode(chunk.as_numpy_array()) + ), + chunk_bytes, + ) + + def compute_encoded_size(self, _input_byte_length: int, _chunk_spec: ArraySpec) -> int: + raise NotImplementedError diff --git a/src/zarr/codecs/numcodecs/__init__.py b/src/zarr/codecs/numcodecs/__init__.py index d68ad3fba6..ed63e13ccf 100644 --- a/src/zarr/codecs/numcodecs/__init__.py +++ b/src/zarr/codecs/numcodecs/__init__.py @@ -11,6 +11,7 @@ AsType, BitRound, Blosc, + Blosc2, Delta, FixedScaleOffset, Fletcher32, @@ -39,6 +40,7 @@ "AsType", "BitRound", "Blosc", + "Blosc2", "Delta", "FixedScaleOffset", "Fletcher32", diff --git a/src/zarr/codecs/numcodecs/_codecs.py b/src/zarr/codecs/numcodecs/_codecs.py index 4a3d88a84f..ae176d2d90 100644 --- a/src/zarr/codecs/numcodecs/_codecs.py +++ b/src/zarr/codecs/numcodecs/_codecs.py @@ -193,6 +193,10 @@ class Blosc(_NumcodecsBytesBytesCodec, codec_name="blosc"): pass +class Blosc2(_NumcodecsBytesBytesCodec, codec_name="blosc2"): + pass + + class LZ4(_NumcodecsBytesBytesCodec, codec_name="lz4"): pass