Skip to content

writer

sleap_nn.inference.writer

IncrementalLabelsWriter — buffered write of a .slp.

The legacy predictor accumulates every frame's predictions in memory before writing to disk. For long videos (100k frames × hundreds of MB of confmaps) that's an OOM risk. This writer accepts Outputs one batch at a time and slim()s each before converting to LabeledFrames, so the heavy intermediate tensors (confmaps, PAFs) are dropped immediately — the dominant memory sink. It finalizes with an atomic rename so a crash mid-write doesn't corrupt the output.

Two design constraints from the design doc:

  • Atomic rename — write to <path>.tmp, rename on close(). Crash → no corrupted output (just an orphan .tmp the user can rm).
  • Resume-friendly — closed writers can be re-opened to append. Not implemented in this commit (deferred to a follow-up).

Memory note: Outputs are slim()-ed per write so heavy tensors are dropped right away. NOTE: until sleap-io exposes an incremental .slp append surface, the (lightweight) LabeledFrame objects accumulate in memory and the .slp is written once at :meth:close/finalize — peak memory is therefore O(n_frames) of slimmed frames, not O(write_interval). write_interval currently controls only the slim/convert cadence, not a per-flush disk write.

Classes:

Name Description
IncrementalLabelsWriter

Stream-write Outputs → sio.LabeledFrame → .slp.

IncrementalLabelsWriter

Stream-write Outputs → sio.LabeledFrame → .slp.

Parameters:

Name Type Description Default
path

Destination .slp path. Writes go to <path>.tmp and rename on close().

 required
skeleton

sio.Skeleton for instance conversion.

 required
videos

Optional list of sio.Video indexed by video_indices.

 required
write_interval

Convert + slim buffered Outputs to LabeledFrames every write_interval frames; the consolidated .slp is written once at :meth:close (no per-flush disk write until sleap-io supports incremental append).

 required

Usage::

writer = IncrementalLabelsWriter(path="out.slp", skeleton=skel)
with writer:  # context-manager closes + atomic-renames
    for outputs in predictor.predict_streaming(provider):
        writer.write(outputs)

Methods:

Name Description
__enter__

Context manager entry — return self.
__exit__

Context manager exit — call :meth:close (atomic-renames).
close

Flush the remaining buffer + atomic-rename .tmp → final path.
write

Buffer an Outputs for eventual disk flush.

Attributes:

Name Type Description
frame_count int

Number of LabeledFrames buffered + accumulated so far.
tmp_path Path

The intermediate .tmp path written to before atomic rename.
Source code in sleap_nn/inference/writer.py 
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
@attrs.define
class IncrementalLabelsWriter:
    """Stream-write ``Outputs`` → ``sio.LabeledFrame`` → ``.slp``.

    Args:
        path: Destination ``.slp`` path. Writes go to ``<path>.tmp`` and
            rename on ``close()``.
        skeleton: ``sio.Skeleton`` for instance conversion.
        videos: Optional list of ``sio.Video`` indexed by ``video_indices``.
        write_interval: Convert + slim buffered ``Outputs`` to
            ``LabeledFrame``s every ``write_interval`` frames; the
            consolidated ``.slp`` is written once at :meth:`close` (no
            per-flush disk write until sleap-io supports incremental append).

    Usage::

        writer = IncrementalLabelsWriter(path="out.slp", skeleton=skel)
        with writer:  # context-manager closes + atomic-renames
            for outputs in predictor.predict_streaming(provider):
                writer.write(outputs)
    """

    path: str
    skeleton: Any
    videos: Optional[List[Any]] = None
    write_interval: int = 500
    # Provenance dict attached to the finalized Labels. May be set after
    # construction (e.g. once end-of-run timestamps are known). #583.
    provenance: Optional[dict] = None
    # Centroid-only packaging — threaded from
    # ``Predictor._resolve_centroid_packaging`` so the streamed ``.slp`` matches
    # the in-memory ``predict()`` output (collapse + sio.Centroid emission).
    anchor_ind: Optional[int] = None
    collapse_skeleton: Optional[Any] = None
    emit_centroid: str = "instance"
    source: str = "center_of_mass"
    # Segmentation mask output representation — threaded so the streamed ``.slp``
    # matches the in-memory ``predict()`` output (mask / polygon ROI / both).
    mask_output: str = "mask"
    polygon_epsilon: float = 0.01

    _buffer: List[Any] = attrs.field(factory=list, init=False, repr=False)
    _all_frames: List[Any] = attrs.field(factory=list, init=False, repr=False)
    _closed: bool = attrs.field(default=False, init=False, repr=False)
    _resolved_videos: Optional[List[Any]] = attrs.field(
        default=None, init=False, repr=False
    )

    @property
    def tmp_path(self) -> Path:
        """The intermediate ``.tmp`` path written to before atomic rename."""
        return Path(self.path).with_suffix(Path(self.path).suffix + ".tmp")

    @property
    def frame_count(self) -> int:
        """Number of ``LabeledFrame``s buffered + accumulated so far."""
        return len(self._all_frames) + len(self._buffer)

    def __enter__(self) -> "IncrementalLabelsWriter":
        """Context manager entry — return self."""
        return self

    def __exit__(self, *exc) -> None:
        """Context manager exit — call :meth:`close` (atomic-renames)."""
        self.close()

    def write(self, outputs: Any) -> None:
        """Buffer an ``Outputs`` for eventual disk flush.

        Args:
            outputs: An :class:`Outputs` from one inference batch. The
                writer slims it (drops heavy intermediates, moves to CPU,
                detaches autograd) before converting to LabeledFrame.
        """
        if self._closed:
            raise RuntimeError("write() called on a closed writer")

        slim = outputs.slim()
        videos = self._resolve_videos()
        sub = slim.to_labels(
            skeleton=self.skeleton,
            videos=videos,
            anchor_ind=self.anchor_ind,
            collapse_skeleton=self.collapse_skeleton,
            emit_centroid=self.emit_centroid,
            source=self.source,
            mask_output=self.mask_output,
            polygon_epsilon=self.polygon_epsilon,
        )
        self._buffer.extend(sub.labeled_frames)

        if len(self._buffer) >= self.write_interval:
            self._flush()

    def _resolve_videos(self) -> List[Any]:
        """Return a list of Videos for label conversion (memoized).

        ``sio.Labels.save`` cannot serialize ``LabeledFrame`` objects
        whose ``video`` is ``None``, so when the caller didn't supply a
        list (or supplied ``[None]``) we substitute a single backend-less
        ``Video`` placeholder. The placeholder serializes cleanly and
        stays loadable via ``sio.load_slp``; users typically rebind a
        real ``Video`` after load.

        The result is computed once and cached so ``write()`` (per batch)
        and ``_finalize()`` share the **same** ``Video`` object(s). Minting
        a fresh placeholder per call produced N+1 distinct videos, so every
        frame referenced a ``Video`` absent from the saved ``Labels.videos``
        list — a corrupt multi-video ``.slp`` (#582).
        """
        import sleap_io as sio

        if self._resolved_videos is not None:
            return self._resolved_videos

        if self.videos:
            real = [v for v in self.videos if v is not None]
            if real:
                self._resolved_videos = real
                return self._resolved_videos
        self._resolved_videos = [sio.Video(filename="unknown", backend=None)]
        return self._resolved_videos

    def close(self) -> None:
        """Flush the remaining buffer + atomic-rename ``.tmp`` → final path.

        Idempotent: subsequent calls are no-ops.
        """
        if self._closed:
            return
        if self._buffer:
            self._flush()
        self._finalize()
        self._closed = True

    # ──────────────────────────────────────────────────────────────────
    # Internals
    # ──────────────────────────────────────────────────────────────────

    def _flush(self) -> None:
        """Move the in-memory buffer onto the on-disk accumulator.

        We don't actually write per-flush to disk in this implementation
        (sleap-io's ``.slp`` format is HDF5; appending one frame at a
        time is high-overhead). Instead we accumulate in
        ``self._all_frames`` and write once at finalization, but the
        caller-visible memory contract (drop ``Outputs`` after
        ``write_interval``) is preserved because we slim() per write
        and only keep ``LabeledFrame`` references.

        A future refactor could switch to true streaming HDF5 writes
        once sleap-io exposes that surface; the public API here doesn't
        change.
        """
        self._all_frames.extend(self._buffer)
        self._buffer.clear()

    def _finalize(self) -> None:
        """Write the ``.tmp`` file and atomically rename to the final path."""
        import sleap_io as sio

        videos = self._resolve_videos()
        # Mask-only (segmentation) models may carry no skeleton; emit an empty
        # skeleton list rather than ``[None]`` (which makes ``sio.Labels.save``
        # raise ``AttributeError: 'NoneType' has no attribute 'nodes'``). Mirrors
        # the in-memory path (``Outputs.to_labels``).
        pkg_skeleton = self.collapse_skeleton or self.skeleton
        labels = sio.Labels(
            labeled_frames=self._all_frames,
            videos=videos,
            skeletons=[pkg_skeleton] if pkg_skeleton is not None else [],
            provenance=self.provenance or {},
        )
        tmp = self.tmp_path
        tmp.parent.mkdir(parents=True, exist_ok=True)
        # Pass the format explicitly: ``sio.Labels.save`` infers from the
        # filename suffix, but our ``.tmp`` suffix isn't recognized.
        # restore_original_videos=False matches the non-streaming path's
        # default (PRESERVE_SOURCE): a .pkg.slp input backreferences itself,
        # not a pre-embedding source video that's often not on disk.
        labels.save(str(tmp), format="slp", restore_original_videos=False)
        # Atomic rename — no half-written file on the destination path.
        tmp.replace(self.path)

frame_count property

Number of LabeledFrames buffered + accumulated so far.

tmp_path property

The intermediate .tmp path written to before atomic rename.

__enter__()

Context manager entry — return self.

Source code in sleap_nn/inference/writer.py 
93
94
95
def __enter__(self) -> "IncrementalLabelsWriter":
    """Context manager entry — return self."""
    return self

__exit__(*exc)

Context manager exit — call :meth:close (atomic-renames).

Source code in sleap_nn/inference/writer.py 
97
98
99
def __exit__(self, *exc) -> None:
    """Context manager exit — call :meth:`close` (atomic-renames)."""
    self.close()

close()

Flush the remaining buffer + atomic-rename .tmp → final path.

Idempotent: subsequent calls are no-ops.

Source code in sleap_nn/inference/writer.py 
158
159
160
161
162
163
164
165
166
167
168
def close(self) -> None:
    """Flush the remaining buffer + atomic-rename ``.tmp`` → final path.

    Idempotent: subsequent calls are no-ops.
    """
    if self._closed:
        return
    if self._buffer:
        self._flush()
    self._finalize()
    self._closed = True

write(outputs)

Buffer an Outputs for eventual disk flush.

Parameters:

Name Type Description Default
outputs Any

An :class:Outputs from one inference batch. The writer slims it (drops heavy intermediates, moves to CPU, detaches autograd) before converting to LabeledFrame.

 required
Source code in sleap_nn/inference/writer.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def write(self, outputs: Any) -> None:
    """Buffer an ``Outputs`` for eventual disk flush.

    Args:
        outputs: An :class:`Outputs` from one inference batch. The
            writer slims it (drops heavy intermediates, moves to CPU,
            detaches autograd) before converting to LabeledFrame.
    """
    if self._closed:
        raise RuntimeError("write() called on a closed writer")

    slim = outputs.slim()
    videos = self._resolve_videos()
    sub = slim.to_labels(
        skeleton=self.skeleton,
        videos=videos,
        anchor_ind=self.anchor_ind,
        collapse_skeleton=self.collapse_skeleton,
        emit_centroid=self.emit_centroid,
        source=self.source,
        mask_output=self.mask_output,
        polygon_epsilon=self.polygon_epsilon,
    )
    self._buffer.extend(sub.labeled_frames)

    if len(self._buffer) >= self.write_interval:
        self._flush()