eclise-clips 0.1.0__tar.gz

eclise_clips-0.1.0/PKG-INFO

@@ -0,0 +1,10 @@
+Metadata-Version: 2.4
+Name: eclise-clips
+Version: 0.1.0
+Requires-Python: >=3.9
+Requires-Dist: numpy>=1.24
+Requires-Dist: h5py>=3.10
+Requires-Dist: imageio[ffmpeg]>=2.36
+Requires-Dist: av>=13.1
+Provides-Extra: test
+Requires-Dist: pytest~=8.4.0; extra == "test"

eclise_clips-0.1.0/README.md

@@ -0,0 +1,113 @@
+# eclise.clips
+
+A shared on-disk contract for video clips carrying per-frame state labels.
+
+This package sits between anything that **produces** clips and anything that
+**consumes** them:
+
+```
+data/phyworld  ──┐                        ┌── video_world_models
+                 ├── eclise.clips shards ──┤
+HF downloads   ──┘                        └── (future consumers)
+```
+
+Producers and consumers live in separate conda envs with incompatible pins and
+never import each other. They agree on this package instead, which is what lets
+generated data be consumed the same way real video will be.
+
+## Reading
+
+```python
+from eclise.clips import ClipDataset
+
+data = ClipDataset("data/phyworld/data/collision_30K")
+item = data[0]
+item["pixels"]      # [3, T, H, W] float32 in [-1, 1]
+item["positions"]   # [T, n_objects, 2] world units, y up
+item["velocity"]    # [2] (vx, vy) px/frame, y down (image convention)
+item["index"]       # global clip index, survives the default collate
+data.source_fps     # 1 / frame_dt
+```
+
+`ClipDataset` is a plain torch `Dataset`; hand it to a `DataLoader` as usual.
+`iter_clip_batches` yields `(clips [B,3,T,H,W], velocities [B,2])` for feeding a
+frozen latent encoder.
+
+Upstream phyworld's published HDF5 carries no metadata, so its family has to be
+named (or inferred from the filename):
+
+```python
+from eclise.clips.readers import phyworld_hdf5
+
+shards = phyworld_hdf5.open_shard_set("collision_30K.hdf5", family="collision")
+data = ClipDataset(shards, velocity_object=0)
+```
+
+## Writing
+
+```python
+from eclise.clips import ShardMeta, ShardWriter, position_spec, params_spec
+
+meta = ShardMeta(
+    family="uniform_motion",
+    num_clips=0,          # overwritten with the realized count on close
+    clip_len=32,
+    image_size=(256, 256),
+    frame_dt=0.1,         # seconds of simulated time between stored frames
+    world_scale=10.0,
+    labels={"positions": position_spec(), "params": params_spec(("r1", "v1"))},
+)
+
+with ShardWriter("shard_000.hdf5", meta) as w:
+    for frames, positions, params in scenes:
+        w.append(frames, positions=positions, params=params)
+```
+
+A shard is one HDF5 file; a dataset is a directory of them. `ShardSet` presents a
+directory as one flat ordered clip sequence, so nothing needs a merge pass.
+
+## Design notes
+
+**Storage.** Clips are mp4 bytestreams in a variable-length `uint8` dataset, one
+row per clip. A dense `uint8` array of 30k 32-frame 256px clips would be ~190GB;
+mp4 brings that to a few GB. Inherited from upstream deliberately, so its
+published files stay readable. Encoding is lossy — consumers see H.264
+artifacts, which is why state-recovery checks run after a decode round-trip.
+
+**`frame_dt`, not `fps`.** An mp4 header's frame rate is a playback preference.
+For simulation output the number that matters is the simulated time between
+stored frames, and upstream never records it: it writes clips with
+`iio.imwrite(..., extension='.mp4')` and no `fps` at all, while the real spacing
+is `timestep * stride = 0.01 * 10 = 0.1s`, set in code the reader never sees. So
+`frame_dt` is required and authoritative, and `source_fps` is derived from it.
+There is no `fps` field, because two rates can disagree.
+
+**Labels are self-describing.** `[N, T, 2, 2]` could be `(clip, frame, ball, xy)`
+in world units, pixels, or normalized coordinates — differing by factors that
+silently rescale any derived velocity. Every label array carries a `LabelSpec`
+recording axis names, units, and whether its y axis points up.
+
+**The y-axis flip.** phyworld's stored positions have y increasing *upward* while
+image rows increase downward: its generators render from normalized positions and
+then reverse the row axis. `positions_to_pixels` and `velocity_from_state` read
+`y_axis_up` off the spec and apply the flip, so no caller has to remember.
+
+## Layout
+
+| Module | Role |
+|---|---|
+| `schema.py` | `ShardMeta`, `LabelSpec`, required metadata, validation |
+| `codec.py` | in-memory mp4 encode/decode |
+| `shards.py` | `ShardWriter`, `ShardReader`, `ShardSet` |
+| `labels.py` | world→pixel conversion, velocity, collision timing |
+| `dataset.py` | `ClipDataset` (the only module needing torch) |
+| `readers/` | adapters for layouts we don't control |
+
+`schema`, `codec`, `shards`, `readers`, and `labels` are pure numpy, so a
+generator env need not install torch.
+
+## Tests
+
+```bash
+cd data/clips && python -m pytest
+```

eclise_clips-0.1.0/pyproject.toml

@@ -0,0 +1,28 @@
+[build-system]
+requires = ["setuptools>=61"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "eclise-clips"
+version = "0.1.0"
+# 3.9, not 3.10: a producer sim may be pinned to an old Python by its simulator.
+# data/phyworld's PHYRE-backed family runs in a 3.9 env because PHYRE's only
+# prebuilt wheel is cp39 (see that project's README), and it must be able to
+# write shards through this package's ShardWriter rather than a second
+# implementation of the contract.
+requires-python = ">=3.9"
+dynamic = ["dependencies"]
+
+[project.optional-dependencies]
+test = ["pytest~=8.4.0"]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["eclise.clips*"]
+
+[tool.setuptools.dynamic]
+dependencies = {file = ["requirements.txt"]}
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+pythonpath = ["src"]

eclise_clips-0.1.0/requirements.txt

@@ -0,0 +1,17 @@
+# Runtime dependencies for eclise.clips.
+#
+# The on-disk clip contract (HDF5 shards holding mp4 bytestreams) plus the torch
+# Dataset that decodes them. Deliberately light: a *producer* (a sim) and a
+# *consumer* (a training project) install this into otherwise-incompatible envs,
+# so anything heavier than these belongs in the consuming project.
+
+numpy>=1.24
+h5py>=3.10
+# mp4 decode/encode of in-memory bytestreams. imageio.v3 does the byte-level
+# imread/imwrite; av (PyAV) is the ffmpeg binding it dispatches to.
+imageio[ffmpeg]>=2.36
+av>=13.1
+
+# torch is an *optional* dependency: schema/shards/readers are pure numpy, and
+# only `eclise.clips.dataset` needs torch. Consumers all have their own pinned
+# torch, so pinning one here would fight them. See dataset.py's import guard.

eclise_clips-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

eclise_clips-0.1.0/src/eclise/clips/__init__.py

@@ -0,0 +1,97 @@
+"""A shared on-disk contract for video clips with per-frame state labels.
+
+Sits between anything that *produces* clips (a simulator, a download of a
+published dataset, eventually a real-video ingest) and anything that *consumes*
+them (a training project). Producers and consumers live in separate conda envs
+and never import each other; they agree on this package instead.
+
+    from eclise.clips import ClipDataset
+
+    data = ClipDataset("data/phyworld/data/collision_30K")
+    item = data[0]
+    item["pixels"]    # [3, T, H, W] float32 in [-1, 1]
+    item["positions"] # [T, n_objects, 2] world units, y up
+    item["velocity"]  # [2] px/frame, y down (image convention)
+    data.source_fps   # 1 / frame_dt -- authoritative, not the mp4 header
+
+Upstream phyworld's published HDF5 carries no metadata, so it needs its family
+named:
+
+    from eclise.clips.readers import phyworld_hdf5
+
+    shards = phyworld_hdf5.open_shard_set("collision_30K.hdf5", family="collision")
+    data = ClipDataset(shards, velocity_object=0)
+
+Writing is the mirror image: declare a
+:class:`~eclise.clips.schema.ShardMeta`, then append clips.
+
+    from eclise.clips import ShardMeta, ShardWriter, position_spec
+
+    meta = ShardMeta(
+        family="uniform_motion", num_clips=0, clip_len=32,
+        image_size=(256, 256), frame_dt=0.1, world_scale=10.0,
+        labels={"positions": position_spec()},
+    )
+    with ShardWriter("shard_000.hdf5", meta) as w:
+        w.append(frames, positions=positions)
+
+Torch is needed only for :class:`~eclise.clips.dataset.ClipDataset`; the schema,
+shard, and reader layers are pure numpy.
+"""
+
+from __future__ import annotations
+
+from .codec import decode_mp4, encode_mp4
+from .labels import (
+    collision_frame,
+    positions_to_pixels,
+    velocities_per_frame,
+    velocity_from_state,
+)
+from .schema import (
+    LABEL_GROUP,
+    META_ATTR,
+    SCHEMA_VERSION,
+    VIDEO_DATASET,
+    LabelSpec,
+    SchemaError,
+    ShardMeta,
+    object_attrs_spec,
+    params_spec,
+    position_spec,
+    velocity_field_spec,
+)
+from .shards import ShardReader, ShardSet, ShardWriter, write_shard
+
+__all__ = [
+    "LABEL_GROUP",
+    "META_ATTR",
+    "SCHEMA_VERSION",
+    "VIDEO_DATASET",
+    "LabelSpec",
+    "SchemaError",
+    "ShardMeta",
+    "ShardReader",
+    "ShardSet",
+    "ShardWriter",
+    "collision_frame",
+    "decode_mp4",
+    "encode_mp4",
+    "object_attrs_spec",
+    "params_spec",
+    "position_spec",
+    "positions_to_pixels",
+    "velocities_per_frame",
+    "velocity_field_spec",
+    "velocity_from_state",
+    "write_shard",
+]
+
+
+def __getattr__(name: str):
+    """Expose :class:`~eclise.clips.dataset.ClipDataset` without importing torch eagerly."""
+    if name in ("ClipDataset", "iter_clip_batches", "PIXEL_RANGES"):
+        from . import dataset
+
+        return getattr(dataset, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

eclise_clips-0.1.0/src/eclise/clips/codec.py

@@ -0,0 +1,85 @@
+"""In-memory mp4 encode/decode for clip bytestreams.
+
+Upstream phyworld carries a copy of this logic in every generator script (four
+copies of ``convert_frames_to_mp4_bytestream_wo_disk``, plus disk-roundtripping
+variants that were left behind after the in-memory version replaced them). This
+is the one home for it.
+
+Encoding is lossy. H.264 on synthetic content -- flat backgrounds with a few
+hard-edged coloured circles -- is close to visually lossless at sane bitrates,
+but "close" is not "exact", and every consumer of a decoded clip is looking at
+compression artifacts. That is a property of the storage format we inherit, not
+something this module can fix; it is why
+``data/phyworld/scripts/verify_dataset.py`` measures state recovery *after* a
+decode round-trip rather than trusting the pre-encode frames.
+"""
+
+from __future__ import annotations
+
+from io import BytesIO
+
+import numpy as np
+
+#: Encoder quality knob. imageio/ffmpeg's default CRF is 23, visibly soft on
+#: hard synthetic edges; 18 is near-transparent for this content at a modest
+#: size increase. Applies only to clips this package writes.
+DEFAULT_CRF = 18
+
+
+def encode_mp4(frames: np.ndarray, *, crf: int = DEFAULT_CRF, fps: float = 24.0) -> np.ndarray:
+    """Encode ``[T, H, W, 3]`` uint8 frames to an mp4 bytestream as ``[n] uint8``.
+
+    ``fps`` lands in the container header only. It is *not* the physics frame
+    rate -- that is :attr:`eclise.clips.schema.ShardMeta.frame_dt`, recorded
+    separately. A value is passed anyway so the header is at least a valid,
+    deliberate number rather than an encoder default (upstream passes nothing,
+    which is how its stored clips ended up with a meaningless rate).
+
+    H.264 requires even frame dimensions; odd ones are rejected here rather than
+    silently rescaled, since a resized clip would invalidate every pixel-unit
+    label in the shard.
+    """
+    import imageio.v3 as iio
+
+    arr = np.asarray(frames)
+    if arr.ndim != 4 or arr.shape[-1] != 3:
+        raise ValueError(f"expected frames [T, H, W, 3], got {arr.shape}")
+    if arr.dtype != np.uint8:
+        raise ValueError(f"expected uint8 frames, got {arr.dtype}")
+    h, w = arr.shape[1:3]
+    if h % 2 or w % 2:
+        raise ValueError(
+            f"H.264 needs even frame dimensions; got {h}x{w}. Render at an even "
+            "size rather than resizing here, which would invalidate pixel labels."
+        )
+    with BytesIO() as buffer:
+        iio.imwrite(
+            buffer,
+            arr,
+            extension=".mp4",
+            fps=fps,
+            codec="libx264",
+            # The pixel format goes through imageio's own keyword, not
+            # output_params: the ffmpeg plugin always emits a -pix_fmt of its own,
+            # so passing a second one makes ffmpeg warn once per encoded clip --
+            # millions of stderr lines on a full generation run, for a flag that
+            # was already set to the same value.
+            pixelformat="yuv420p",
+            output_params=["-crf", str(crf)],
+        )
+        payload = buffer.getvalue()
+    return np.frombuffer(payload, dtype=np.uint8)
+
+
+def decode_mp4(stream: np.ndarray | bytes) -> np.ndarray:
+    """Decode an mp4 bytestream to ``[T, H, W, 3]`` uint8 frames."""
+    import imageio.v3 as iio
+
+    payload = stream.tobytes() if isinstance(stream, np.ndarray) else stream
+    frames = iio.imread(payload, index=None, extension=".mp4")
+    arr = np.asarray(frames)
+    if arr.ndim == 3:  # a single-frame clip decodes without the leading axis
+        arr = arr[None]
+    if arr.shape[-1] == 4:  # some decoders hand back RGBA
+        arr = arr[..., :3]
+    return arr

eclise_clips-0.1.0/src/eclise/clips/dataset.py

@@ -0,0 +1,235 @@
+"""Torch ``Dataset`` over clip shards.
+
+This is the consumer-facing end of the contract. A training project points
+:class:`ClipDataset` at a directory of shards -- generated by a sim, downloaded
+from upstream, or (later) holding real video -- and gets pixel clips plus labels,
+without knowing which.
+
+``torch`` is imported lazily: :mod:`eclise.clips.schema`, :mod:`.shards`, and
+:mod:`.readers` are pure numpy so a generator-side env need not install torch at
+all, while consumers bring their own pinned version.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from .labels import velocity_from_state
+from .schema import SchemaError, ShardMeta
+from .shards import ShardSet
+
+#: How decoded uint8 pixels are mapped to floats.
+#:
+#: - ``"tanh"``: ``[-1, 1]``. The default, because it is what latent video
+#:   encoders (Wan VAE and friends) consume, and what
+#:   ``video_world_models``' own procedural stimuli already produce.
+#: - ``"unit"``: ``[0, 1]``.
+#: - ``"uint8"``: untouched ``0..255``, dtype preserved.
+PIXEL_RANGES = ("tanh", "unit", "uint8")
+
+
+def _to_pixels(frames: np.ndarray, pixel_range: str) -> Any:
+    import torch
+
+    if pixel_range not in PIXEL_RANGES:
+        raise ValueError(
+            f"pixel_range must be one of {PIXEL_RANGES}; got {pixel_range!r}"
+        )
+    arr = np.ascontiguousarray(frames)
+    tensor = torch.from_numpy(arr)  # [T, H, W, 3] uint8
+    tensor = tensor.permute(3, 0, 1, 2)  # -> [3, T, H, W]
+    if pixel_range == "uint8":
+        return tensor.contiguous()
+    tensor = tensor.to(torch.float32) / 255.0
+    if pixel_range == "tanh":
+        tensor = tensor * 2.0 - 1.0
+    return tensor.contiguous()
+
+
+class ClipDataset:
+    """Clips from one or more shards, as ``[3, T, H, W]`` tensors plus labels.
+
+    Each item is a dict:
+
+    - ``"pixels"``: ``[3, T, H, W]``, float32 in the configured range (or uint8).
+    - ``"index"``: global clip index, so a batched metric can be bucketed by any
+      per-clip property after the fact. Labels that are strings or ragged cannot
+      survive the default collate; an index always can.
+    - ``"velocity"``: ``[2]`` ``(vx, vy)`` px/frame, present when the shard has
+      positions and the family has one object (see ``velocity_object``).
+    - one entry per stored label, e.g. ``"positions"`` ``[T, n_obj, 2]``,
+      ``"params"`` ``[K]``.
+
+    ``meta`` (family, ``frame_dt``, ``world_scale``, ...) is a dataset-level
+    attribute, not per item -- it is identical for every clip, and putting it in
+    the item would multiply it by the batch size through the collate.
+
+    Args:
+        source: a shard directory, a single shard file, a list of files, or an
+            already-built :class:`~eclise.clips.shards.ShardSet`.
+        pixel_range: see :data:`PIXEL_RANGES`.
+        labels: which stored labels to include; ``None`` includes all.
+        velocity: emit a ``"velocity"`` entry derived from positions.
+        velocity_object: which object's velocity to use for multi-object
+            families. Required there, since a mean over a head-on collision is
+            ~0 (see :func:`eclise.clips.labels.velocity_from_state`).
+        opener: reader class for non-canonical shards, e.g.
+            ``functools.partial(phyworld_hdf5.open_shard, family="collision")``.
+    """
+
+    def __init__(
+        self,
+        source: str | Path | Sequence[str | Path] | ShardSet,
+        *,
+        pixel_range: str = "tanh",
+        labels: Sequence[str] | None = None,
+        velocity: bool = True,
+        velocity_object: int | None = None,
+        opener: Any = None,
+    ) -> None:
+        if pixel_range not in PIXEL_RANGES:
+            raise ValueError(
+                f"pixel_range must be one of {PIXEL_RANGES}; got {pixel_range!r}"
+            )
+        self.shards = _resolve_source(source, opener=opener)
+        self.pixel_range = pixel_range
+        self.meta: ShardMeta = self.shards.meta
+
+        available = set(self.shards.label_names())
+        if labels is None:
+            self.labels = tuple(sorted(available))
+        else:
+            missing = set(labels) - available
+            if missing:
+                raise SchemaError(
+                    f"requested labels {sorted(missing)} are not in these shards; "
+                    f"available: {sorted(available)}"
+                )
+            self.labels = tuple(labels)
+
+        self._velocity_object = velocity_object
+        self.emits_velocity = bool(velocity and "positions" in available)
+        if velocity and not self.emits_velocity:
+            raise SchemaError(
+                "velocity=True needs a 'positions' label, which these shards do "
+                f"not have (available: {sorted(available)}). Pass velocity=False, "
+                "or use a family that records positions."
+            )
+        if self.emits_velocity:
+            num_objects = int(self.meta.extra.get("num_objects", 1) or 1)
+            if num_objects > 1 and velocity_object is None:
+                raise SchemaError(
+                    f"family {self.meta.family!r} has {num_objects} objects, so a "
+                    "single per-clip velocity is ambiguous. Pass velocity_object=, "
+                    "or velocity=False and derive from the 'positions' label."
+                )
+
+    @property
+    def source_fps(self) -> float:
+        """``1 / frame_dt``, for a latent encoder's FPS-aware path."""
+        return self.meta.source_fps
+
+    def manifest_key(self, *, length: int = 16) -> str:
+        """Stable hash of the underlying shard set, for addressing derived caches."""
+        return self.shards.manifest_key(length=length)
+
+    def manifest(self) -> dict[str, Any]:
+        return self.shards.manifest()
+
+    def __len__(self) -> int:
+        return len(self.shards)
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        import torch
+
+        item: dict[str, Any] = {
+            "pixels": _to_pixels(self.shards.frames(index), self.pixel_range),
+            "index": index,
+        }
+        for name in self.labels:
+            value = np.asarray(self.shards.label(name, index))
+            item[name] = torch.from_numpy(np.ascontiguousarray(value)).to(torch.float32)
+        if self.emits_velocity:
+            positions = np.asarray(self.shards.label("positions", index))
+            item["velocity"] = torch.from_numpy(
+                velocity_from_state(
+                    positions,
+                    self.shards.spec("positions"),
+                    image_size=self.meta.image_size,
+                    world_scale=self.meta.world_scale,
+                    object_index=self._velocity_object,
+                )
+            )
+        return item
+
+    def close(self) -> None:
+        self.shards.close()
+
+    def __enter__(self) -> ClipDataset:
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self.close()
+
+    def __repr__(self) -> str:
+        return (
+            f"ClipDataset(family={self.meta.family!r}, split={self.meta.split!r}, "
+            f"clips={len(self)}, clip_len={self.meta.clip_len}, "
+            f"image_size={self.meta.image_size}, frame_dt={self.meta.frame_dt}, "
+            f"pixel_range={self.pixel_range!r})"
+        )
+
+
+def _resolve_source(
+    source: str | Path | Sequence[str | Path] | ShardSet, *, opener: Any
+) -> ShardSet:
+    if isinstance(source, ShardSet):
+        return source
+    if isinstance(source, (str, Path)):
+        path = Path(source)
+        if path.is_dir():
+            return ShardSet.from_dir(path, opener=opener)
+        return ShardSet([path], opener=opener)
+    return ShardSet(list(source), opener=opener)
+
+
+def iter_clip_batches(
+    dataset: ClipDataset,
+    *,
+    batch_size: int = 8,
+    velocity_fallback: bool = True,
+) -> Any:
+    """Yield ``(clips [B, 3, T, H, W], velocities [B, 2])`` batches.
+
+    The shape a frozen latent encoder wants. Kept here rather than in the
+    consumer so that "render a clip grid" and "read clips off disk" present the
+    same interface to an encode loop.
+
+    With ``velocity_fallback`` and no velocity available, zeros are emitted --
+    appropriate when velocity is only a diagnostic label, not a training target.
+    """
+    import torch
+
+    clips: list[Any] = []
+    vels: list[Any] = []
+    for index in range(len(dataset)):
+        item = dataset[index]
+        clips.append(item["pixels"])
+        if "velocity" in item:
+            vels.append(item["velocity"])
+        elif velocity_fallback:
+            vels.append(torch.zeros(2, dtype=torch.float32))
+        else:
+            raise SchemaError(
+                "dataset emits no velocity; pass velocity_fallback=True to get "
+                "zeros, or construct the ClipDataset with velocity=True"
+            )
+        if len(clips) == batch_size:
+            yield torch.stack(clips), torch.stack(vels)
+            clips, vels = [], []
+    if clips:
+        yield torch.stack(clips), torch.stack(vels)