Coverage for trimesh/voxel/runlength.py: 90%
298 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-24 04:40 +0000
1"""
2Numpy encode/decode/utility implementations for run length encodings.
4# Run Length Encoded Features
6Encoding/decoding functions for run length encoded data.
8We include code for two variations:
10* run length encoding (RLE)
11* binary run length encdoing (BRLE)
13RLE stores sequences of repeated values as the value followed by its count, e.g.
15```python
16dense_to_rle([5, 5, 3, 2, 2, 2, 2, 6]) == [5, 2, 3, 1, 2, 4, 6, 1]
17```
19i.e. the value `5` is repeated `2` times, then `3` is repeated `1` time, `2` is
20repeated `4` times and `6` is repeated `1` time.
22BRLE is an optimized form for when the stored values can only be `0` or `1`.
23This means we only need to save the counts, and assume the values alternate
24(starting at `0`).
26```python
27dense_to_brle([1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0]) == \
28 [0, 2, 4, 7, 2]
29```
31i.e. the value zero occurs `0` times, followed by `2` ones, `4` zeros, `7` ones
32and `2` zeros.
34Sequences with counts exceeding the data type's maximum value have to be
35handled carefully. For example, the `uint8` encoding of 300 zeros
36(`uint8` has a max value of 255) is:
38* RLE: `[0, 255, 0, 45]` (`0` repeated `255` times + `0` repeated `45` times)
39* BRLE: `[255, 0, 45, 0]` (`255` zeros + `0` ones + `45` zeros + `0` ones)
41This module contains implementations of various RLE/BRLE operations.
42"""
44import functools
46import numpy as np
49def brle_length(brle):
50 """Optimized implementation of `len(brle_to_dense(brle))`"""
51 return np.sum(brle)
54def rle_length(rle):
55 """Optimized implementation of `len(rle_to_dense(rle_to_brle(rle)))`"""
56 return np.sum(rle[1::2])
59def rle_to_brle(rle, dtype=None):
60 """
61 Convert run length encoded (RLE) value/counts to BRLE.
63 RLE data is stored in a rank 1 array with each pair giving:
64 (value, count)
66 e.g. the RLE encoding of [4, 4, 4, 1, 1, 6] is [4, 3, 1, 2, 6, 1].
68 Parameters
69 ----------
70 rle : (n,) int
71 Run length encoded data
73 Returns
74 ----------
75 equivalent binary run length encoding. a list if dtype is None,
76 otherwise brle_to_brle is called on that list before returning.
78 Raises
79 ----------
80 ValueError
81 If any of the even counts of `rle` are not zero or 1.
82 """
83 curr_val = 0
84 out = [0]
85 acc = 0
86 for value, count in np.reshape(rle, (-1, 2)):
87 acc += count
88 if value not in (0, 1):
89 raise ValueError("Invalid run length encoding for conversion to BRLE")
90 if value == curr_val:
91 out[-1] += count
92 else:
93 out.append(int(count))
94 curr_val = value
95 if len(out) % 2:
96 out.append(0)
97 if dtype is not None:
98 out = brle_to_brle(out, dtype=dtype)
99 return out
102def brle_logical_not(brle):
103 """
104 Get the BRLE encoding of the `logical_not`ed dense form of `brle`.
106 Equivalent to `dense_to_brle(np.logical_not(brle_to_dense(brle)))` but
107 highly optimized - just pads brle with a 0 on each end (or strips is
108 existing endpoints are both zero).
110 Parameters
111 ----------
112 brle: rank 1 int array of binary run length encoded data
114 Returns
115 ----------
116 rank 1 int array of binary run length encoded data corresponding to
117 element-wise not of the input.
118 """
119 if brle[0] or brle[-1]:
120 return np.pad(brle, [1, 1], mode="constant")
121 else:
122 return brle[1:-1]
125def merge_brle_lengths(lengths):
126 """Inverse of split_long_brle_lengths."""
127 if len(lengths) == 0:
128 return []
130 out = [int(lengths[0])]
131 accumulating = False
132 for length in lengths[1:]:
133 if accumulating:
134 out[-1] += length
135 accumulating = False
136 else:
137 if length == 0:
138 accumulating = True
139 else:
140 out.append(int(length))
141 return out
144def split_long_brle_lengths(lengths, dtype=np.int64):
145 """
146 Split lengths that exceed max dtype value.
148 Lengths `l` are converted into [max_val, 0] * l // max_val + [l % max_val]
150 e.g. for dtype=np.uint8 (max_value == 255)
151 ```
152 split_long_brle_lengths([600, 300, 2, 6], np.uint8) == \
153 [255, 0, 255, 0, 90, 255, 0, 45, 2, 6]
154 ```
155 """
156 lengths = np.asarray(lengths)
157 max_val = np.iinfo(dtype).max
158 bad_length_mask = lengths > max_val
159 if np.any(bad_length_mask):
160 # there are some bad lengths
161 nl = len(lengths)
162 repeats = np.asarray(lengths) // max_val
163 remainders = (lengths % max_val).astype(dtype)
165 lengths = np.concatenate(
166 [
167 np.array([max_val, 0] * repeat + [remainder], dtype=dtype)
168 for repeat, remainder in zip(repeats, remainders)
169 ]
170 )
171 lengths = lengths.reshape((np.sum(repeats) * 2 + nl,)).astype(dtype)
172 return lengths
173 elif lengths.dtype != dtype:
174 return lengths.astype(dtype)
175 else:
176 return lengths
179def dense_to_brle(dense_data, dtype=np.int64):
180 """
181 Get the binary run length encoding of `dense_data`.
183 Parameters
184 ----------
185 dense_data: rank 1 bool array of data to encode.
186 dtype: numpy int type.
188 Returns
189 ----------
190 Binary run length encoded rank 1 array of dtype `dtype`.
192 Raises
193 ----------
194 ValuError if dense_data is not a rank 1 bool array.
195 """
196 if dense_data.dtype != bool:
197 raise ValueError("`dense_data` must be bool")
198 if len(dense_data.shape) != 1:
199 raise ValueError("`dense_data` must be rank 1.")
200 n = len(dense_data)
201 starts = np.r_[0, np.flatnonzero(dense_data[1:] != dense_data[:-1]) + 1]
202 lengths = np.diff(np.r_[starts, n])
203 lengths = split_long_brle_lengths(lengths, dtype=dtype)
204 if dense_data[0]:
205 lengths = np.pad(lengths, [1, 0], mode="constant")
206 return lengths
209_ft = np.array([False, True], dtype=bool)
212def brle_to_dense(brle_data, vals=None):
213 """Decode binary run length encoded data to dense.
215 Parameters
216 ----------
217 brle_data: BRLE counts of False/True values
218 vals: if not `None`, a length 2 array/list/tuple with False/True substitute
219 values, e.g. brle_to_dense([2, 3, 1, 0], [7, 9]) == [7, 7, 9, 9, 9, 7]
221 Returns
222 ----------
223 rank 1 dense data of dtype `bool if vals is None else vals.dtype`
225 Raises
226 ----------
227 ValueError if vals it not None and shape is not (2,)
228 """
229 if vals is None:
230 vals = _ft
231 else:
232 vals = np.asarray(vals)
233 if vals.shape != (2,):
234 raise ValueError(f"vals.shape must be (2,), got {vals.shape}")
235 ft = np.repeat(_ft[np.newaxis, :], (len(brle_data) + 1) // 2, axis=0).flatten()
236 return np.repeat(ft[: len(brle_data)], np.asarray(brle_data, dtype=int)).flatten()
239def rle_to_dense(rle_data, dtype=np.int64):
240 """Get the dense decoding of the associated run length encoded data."""
241 values, counts = np.split(np.reshape(rle_data, (-1, 2)), 2, axis=-1)
242 if dtype is not None:
243 values = np.asanyarray(values, dtype=dtype)
244 try:
245 result = np.repeat(
246 np.squeeze(values, axis=-1), np.squeeze(counts, axis=-1).astype(int)
247 )
248 except TypeError:
249 # on windows it sometimes fails to cast data type
250 result = np.repeat(
251 np.squeeze(values.astype(np.int64), axis=-1),
252 np.squeeze(counts.astype(int), axis=-1),
253 )
254 return result
257def dense_to_rle(dense_data, dtype=np.int64):
258 """Get run length encoding of the provided dense data."""
259 n = len(dense_data)
260 starts = np.r_[0, np.flatnonzero(dense_data[1:] != dense_data[:-1]) + 1]
261 lengths = np.diff(np.r_[starts, n])
262 values = dense_data[starts]
263 values, lengths = split_long_rle_lengths(values, lengths, dtype=dtype)
264 out = np.stack((values, lengths), axis=1)
265 return out.flatten()
268def split_long_rle_lengths(values, lengths, dtype=np.int64):
269 """
270 Split long lengths in the associated run length encoding.
272 e.g.
273 ```python
274 split_long_rle_lengths([5, 300, 2, 12], np.uint8) == [5, 255, 5, 45, 2, 12]
275 ```
277 Parameters
278 ----------
279 values: values column of run length encoding, or `rle[::2]`
280 lengths: counts in run length encoding, or `rle[1::2]`
281 dtype: numpy data type indicating the maximum value.
283 Returns
284 ----------
285 values, lengths associated with the appropriate splits. `lengths` will be
286 of type `dtype`, while `values` will be the same as the value passed in.
287 """
288 max_length = np.iinfo(dtype).max
289 lengths = np.asarray(lengths, dtype=np.int64)
290 repeats = lengths // max_length
291 if repeats.any():
292 repeats = (repeats + 1).astype(int)
293 remainder = lengths % max_length
294 values = np.repeat(values, repeats)
295 lengths = np.zeros(len(repeats), dtype=dtype)
296 lengths.fill(max_length)
297 lengths = np.repeat(lengths, repeats)
298 lengths[np.cumsum(repeats) - 1] = remainder
299 elif lengths.dtype != dtype:
300 lengths = lengths.astype(dtype)
301 return values, lengths
304def merge_rle_lengths(values, lengths):
305 """Inverse of split_long_rle_lengths except returns normal python lists."""
306 ret_values = []
307 ret_lengths = []
308 curr = None
309 for value, length in zip(values, lengths):
310 if length == 0:
311 continue
312 if value == curr:
313 ret_lengths[-1] += length
314 else:
315 curr = value
316 ret_lengths.append(int(length))
317 ret_values.append(value)
318 return ret_values, ret_lengths
321def brle_to_rle(brle, dtype=np.int64):
322 if len(brle) % 2 == 1:
323 brle = np.concatenate([brle, [0]])
324 lengths = brle
325 values = np.tile(_ft, len(brle) // 2)
326 return rle_to_rle(np.stack((values, lengths), axis=1).flatten(), dtype=dtype)
329def brle_to_brle(brle, dtype=np.int64):
330 """
331 Almost the identity function.
333 Checks for possible merges and required splits.
334 """
335 return split_long_brle_lengths(merge_brle_lengths(brle), dtype=dtype)
338def rle_to_rle(rle, dtype=np.int64):
339 """
340 Almost the identity function.
342 Checks for possible merges and required splits.
343 """
344 values, lengths = np.reshape(rle, (-1, 2)).T
345 values, lengths = merge_rle_lengths(values, lengths)
346 values, lengths = split_long_rle_lengths(values, lengths, dtype=dtype)
347 return np.stack((values, lengths), axis=1).flatten()
350def _unsorted_gatherer(indices, sorted_gather_fn):
351 if not isinstance(indices, np.ndarray):
352 indices = np.array(indices, copy=False)
353 order = np.argsort(indices)
354 ordered_indices = indices[order]
356 def f(data, dtype=None):
357 result = np.zeros(len(order), dtype=dtype or getattr(data, "dtype", None))
358 result[order] = tuple(sorted_gather_fn(data, ordered_indices))
359 return result
361 return f
364def sorted_rle_gather_1d(rle_data, ordered_indices):
365 """
366 Gather brle_data at ordered_indices.
368 This is equivalent to `rle_to_dense(brle_data)[ordered_indices]` but avoids
369 the decoding.
371 Parameters
372 ----------
373 brle_data: iterable of run-length-encoded data.
374 ordered_indices: iterable of ints in ascending order.
376 Returns
377 ----------
378 `brle_data` iterable of values at the dense indices, same length as
379 ordered indices.
380 """
381 data_iter = iter(rle_data)
382 index_iter = iter(ordered_indices)
383 try:
384 index = next(index_iter)
385 except StopIteration:
386 return
387 start = 0
388 while True:
389 while start <= index:
390 try:
391 value = next(data_iter)
392 start += next(data_iter)
393 except StopIteration:
394 raise IndexError(
395 "Index %d out of range of raw_values length %d", index, start
396 )
398 try:
399 while index < start:
400 yield value
401 index = next(index_iter)
402 except StopIteration:
403 break
406def rle_mask(rle_data, mask):
407 """
408 Perform masking of the input run-length data.
410 Parameters
411 ----------
412 rle_data: iterable of run length encoded data
413 mask: iterable of bools corresponding to the dense mask.
415 Returns
416 ----------
417 iterable of dense values of rle_data wherever mask is True.
418 """
419 data_iter = iter(rle_data)
420 mask_iter = iter(mask)
421 while True:
422 try:
423 value = next(data_iter)
424 count = next(data_iter)
425 except StopIteration:
426 break
427 for _ in range(count):
428 m = next(mask_iter)
429 if m:
430 yield value
433def brle_mask(rle_data, mask):
434 """
435 Perform masking of the input binary run-length data.
437 Parameters
438 ----------
439 brle_data: iterable of binary run length encoded data
440 mask: iterable of bools corresponding to the dense mask.
442 Returns
443 ----------
444 iterable dense values of brle_data wherever mask is True.
445 """
446 data_iter = iter(rle_data)
447 mask_iter = iter(mask)
448 value = True
449 while True:
450 try:
451 value = not value
452 count = next(data_iter)
453 except StopIteration:
454 break
455 for _ in range(count):
456 m = next(mask_iter)
457 if m:
458 yield value
461def rle_gatherer_1d(indices):
462 """
463 Get a gather function at the given indices.
465 Because gathering on RLE data requires sorting, for instances where
466 gathering at the same indices on different RLE data this can save the
467 sorting process.
469 If only gathering on a single RLE iterable, use `rle_gather_1d`.
471 Parameters
472 ----------
473 indices: iterable of integers
475 Returns
476 ----------
477 gather function, mapping `(rle_data, dtype=None) -> values`.
478 `values` will have the same length as `indices` and dtype provided,
479 or rle_data.dtype if no dtype is provided.
480 """
481 return _unsorted_gatherer(indices, sorted_rle_gather_1d)
484def rle_gather_1d(rle_data, indices, dtype=None):
485 """
486 Gather RLE data values at the provided dense indices.
488 This is equivalent to `rle_to_dense(rle_data)[indices]` but the
489 implementation does not require the construction of the dense array.
491 If indices is known to be in order, use `sorted_gather_1d`.
493 Parameters
494 ----------
495 rle_data: run length encoded data
496 indices: dense indices
497 dtype: numpy dtype. If not provided, uses rle_data.dtype
499 Returns
500 ----------
501 numpy array, dense data at indices, same length as indices and dtype as
502 rle_data
503 """
504 return rle_gatherer_1d(indices)(rle_data, dtype=dtype)
507def sorted_brle_gather_1d(brle_data, ordered_indices):
508 """
509 Gather brle_data at ordered_indices.
511 This is equivalent to `brle_to_dense(brle_data)[ordered_indices]` but
512 avoids the decoding.
514 Parameters
515 ----------
516 raw_data: iterable of run-length-encoded data.
517 ordered_indices: iterable of ints in ascending order.
519 Returns
520 ----------
521 `raw_data` iterable of values at the dense indices, same length as
522 ordered indices.
523 """
524 data_iter = iter(brle_data)
525 index_iter = iter(ordered_indices)
526 try:
527 index = next(index_iter)
528 except StopIteration:
529 return
530 start = 0
531 value = True
532 while True:
533 while start <= index:
534 try:
535 value = not value
536 start += next(data_iter)
537 except StopIteration:
538 raise IndexError(
539 "Index %d out of range of raw_values length %d", index, start
540 )
542 try:
543 while index < start:
544 yield value
545 index = next(index_iter)
546 except StopIteration:
547 break
550def brle_gatherer_1d(indices):
551 """
552 Get a gather function at the given indices.
554 Because gathering on BRLE data requires sorting, for instances where
555 gathering at the same indices on different RLE data this can save the
556 sorting process.
558 If only gathering on a single RLE iterable, use `brle_gather_1d`.
560 Parameters
561 ----------
562 indices: iterable of integers
564 Returns
565 ----------
566 gather function, mapping `(rle_data, dtype=None) -> values`.
567 `values` will have the same length as `indices` and dtype provided,
568 or rle_data.dtype if no dtype is provided.
569 """
570 return functools.partial(
571 _unsorted_gatherer(indices, sorted_brle_gather_1d), dtype=bool
572 )
575def brle_gather_1d(brle_data, indices):
576 """
577 Gather BRLE data values at the provided dense indices.
579 This is equivalent to `rle_to_dense(rle_data)[indices]` but the
580 implementation does not require the construction of the dense array.
582 If indices is known to be in order, use `sorted_brle_gather_1d`.
584 Parameters
585 ----------
586 rle_data: run length encoded data
587 indices: dense indices
589 Returns
590 ----------
591 numpy array, dense data at indices, same length as indices and dtype as
592 rle_data
593 """
594 return brle_gatherer_1d(indices)(brle_data)
597def brle_reverse(brle_data):
598 """Equivalent to dense_to_brle(brle_to_dense(brle_data)[-1::-1])."""
599 if len(brle_data) % 2 == 0:
600 brle_data = np.concatenate([brle_data, [0]], axis=0)
601 end = -1 if brle_data[-1] == 0 else None
602 return brle_data[-1:end:-1]
605def rle_reverse(rle_data):
606 """Get the rle encoding of the reversed dense array."""
607 if not isinstance(rle_data, np.ndarray):
608 rle_data = np.array(rle_data, copy=False)
609 rle_data = np.reshape(rle_data, (-1, 2))
610 rle_data = rle_data[-1::-1]
611 return np.reshape(rle_data, (-1,))
614def rle_to_sparse(rle_data):
615 """Get dense indices associated with non-zeros."""
616 indices = []
617 values = []
618 it = iter(rle_data)
619 index = 0
620 try:
621 while True:
622 value = next(it)
623 counts = int(next(it))
624 end = index + counts
625 if value:
626 indices.append(np.arange(index, end, dtype=np.int64))
627 values.append(np.repeat(value, counts))
628 index = end
629 except StopIteration:
630 pass
631 if len(indices) == 0:
632 assert len(values) == 0
633 return indices, values
635 indices = np.concatenate(indices)
636 values = np.concatenate(values, dtype=rle_data.dtype)
637 return indices, values
640def brle_to_sparse(brle_data, dtype=np.int64):
641 ends = np.cumsum(brle_data)
642 indices = [np.arange(s, e, dtype=dtype) for s, e in zip(ends[::2], ends[1::2])]
643 return np.concatenate(indices)
646def rle_strip(rle_data):
647 """
648 Remove leading and trailing zeros.
650 Parameters
651 ----------
652 rle_data: run length encoded data
654 Returns
655 ----------
656 (stripped_rle_data, padding)
657 stripped_rle_data: rle data without any leading or trailing zeros
658 padding: 2-element dense padding
659 """
660 rle_data = np.reshape(rle_data, (-1, 2))
661 start = 0
662 final_i = len(rle_data)
663 for i, (val, count) in enumerate(rle_data):
664 if val and count > 0:
665 final_i = i
666 break
667 else:
668 start += count
670 end = 0
671 final_j = len(rle_data)
672 for j, (val, count) in enumerate(rle_data[::-1]):
673 if val and count > 0:
674 final_j = j
675 break
676 else:
677 end += count
679 rle_data = rle_data[final_i : None if final_j == 0 else -final_j].reshape((-1,))
680 return rle_data, (start, end)
683def brle_strip(brle_data):
684 """
685 Remove leading and trailing zeros.
687 Parameters
688 ----------
689 brle_data: binary run length encoded data.
691 Returns
692 ----------
693 (stripped_brle_data, padding)
694 stripped_brle_data: rle data without any leading or trailing zeros
695 padding: 2-element dense padding
696 """
697 start = 0
698 val = True
699 final_i = len(brle_data)
700 for i, count in enumerate(brle_data):
701 val = not val
702 if val and count > 0:
703 final_i = i
704 break
705 else:
706 start += count
707 end = 0
708 final_j = len(brle_data)
709 val = bool(len(brle_data) % 2)
710 for j, count in enumerate(brle_data[::-1]):
711 val = not val
712 if val and count > 0:
713 final_j = j
714 break
715 else:
716 end += count
718 brle_data = brle_data[final_i : None if final_j == 0 else -final_j]
719 brle_data = np.concatenate([[0], brle_data])
720 return brle_data, (start, end)