diff --git a/docs/utils/iterables.md b/docs/utils/iterables.md new file mode 100644 index 00000000..cde710f1 --- /dev/null +++ b/docs/utils/iterables.md @@ -0,0 +1,18 @@ +--- +comments: true +status: new +--- + +# Iterables Utils + +
+

create_batches

+
+ +:::supervision.utils.iterables.create_batches + +
+

fill

+
+ +:::supervision.utils.iterables.fill \ No newline at end of file diff --git a/mkdocs.yml b/mkdocs.yml index 5444290b..2afadfa5 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -65,6 +65,7 @@ nav: - Utils: - Video: utils/video.md - Image: utils/image.md + - Iterables: utils/iterables.md - Notebook: utils/notebook.md - File: utils/file.md - Assets: assets.md diff --git a/supervision/utils/iterables.py b/supervision/utils/iterables.py index 529ceadb..70b658b0 100644 --- a/supervision/utils/iterables.py +++ b/supervision/utils/iterables.py @@ -1,25 +1,33 @@ from typing import Generator, Iterable, List, TypeVar -SequenceElement = TypeVar("SequenceElement") +V = TypeVar("V") def create_batches( - sequence: Iterable[SequenceElement], batch_size: int -) -> Generator[List[SequenceElement], None, None]: + sequence: Iterable[V], batch_size: int +) -> Generator[List[V], None, None]: """ - Provides a generator that yields chunks of input sequence - of size specified by `batch_size` parameter. Last - chunk may be smaller batch. + Provides a generator that yields chunks of the input sequence + of the size specified by the `batch_size` parameter. The last + chunk may be a smaller batch. Args: - sequence (Iterable[SequenceElement]): Sequence to be - split into batches. - batch_size (int): Expected size of a batch + sequence (Iterable[V]): The sequence to be split into batches. + batch_size (int): The expected size of a batch. Returns: - Generator[List[SequenceElement], None, None]: Generator - to yield chinks of `sequence` of size `batch_size`, - up to the length of input `sequence`. + Generator[List[V], None, None]: A generator that yields chunks + of `sequence` of size `batch_size`, up to the length of + the input `sequence`. + + Examples: + ```python + list(create_batches([1, 2, 3, 4, 5], 2)) + # [[1, 2], [3, 4], [5]] + + list(create_batches("abcde", 3)) + # [['a', 'b', 'c'], ['d', 'e']] + ``` """ batch_size = max(batch_size, 1) current_batch = [] @@ -28,31 +36,35 @@ def create_batches( yield current_batch current_batch = [] current_batch.append(element) - if len(current_batch) > 0: + if current_batch: yield current_batch -def fill( - sequence: List[SequenceElement], - desired_size: int, - content: SequenceElement, -) -> List[SequenceElement]: +def fill(sequence: List[V], desired_size: int, content: V) -> List[V]: """ - Fill the sequence with padding elements until sequence reaches - desired size. + Fill the sequence with padding elements until the sequence reaches + the desired size. Args: - sequence (List[SequenceElement]): Input sequence. - desired_size (int): Expected size of output list - difference - between this value and actual `sequence` length (if positive) - dictates how many elements will be added as padding. - content (SequenceElement): Element to be placed at the end of - input `sequence` as padding. + sequence (List[V]): The input sequence. + desired_size (int): The expected size of the output list. The + difference between this value and the actual length of `sequence` + (if positive) dictates how many elements will be added as padding. + content (V): The element to be placed at the end of the input + `sequence` as padding. Returns: - List[SequenceElement]: Padded version of input `sequence` (if needed) + List[V]: A padded version of the input `sequence` (if needed). + + Examples: + ```python + fill([1, 2], 4, 0) + # [1, 2, 0, 0] + + fill(['a', 'b'], 3, 'c') + # ['a', 'b', 'c'] + ``` """ missing_size = max(0, desired_size - len(sequence)) - required_padding = [content] * missing_size - sequence.extend(required_padding) + sequence.extend([content] * missing_size) return sequence