dataset module

Armory Dataset Classes

class dataset.ArmoryDataset(dataset, adapter: Callable[[...], Mapping[str, Any]])

Bases: Dataset

Wrapper around a PyTorch dataset to apply an adapter to all samples obtained from the dataset.

Example:

from armory.dataset import ArmoryDataset

def rename_fields(sample):
    # Rename the 'data' field in the sample to 'image'
    sample["image"] = sample.pop("data")
    return sample

# assuming `dataset` has been defined elsewhere
renamed_dataset = ArmoryDataset(dataset, rename_fields)

dataset.DatasetOutputAdapter

An adapter for dataset samples. The output must be a dictionary of column names to values.

alias of Callable[[…], Mapping[str, Any]]

class dataset.ImageClassificationDataLoader(*args, dim: ImageDimensions, image_key: str, label_key: str, scale: Scale, **kwargs)

Bases: ShuffleableDataLoader

Data loader for image classification datasets.

Example:

import armory.data
from armory.dataset import ImageClassificationDataLoader

# assuming `dataset` has been defined elsewhere
dataloader = ImageClassificationDataLoader(
    dataset,
    dim=armory.data.ImageDimensions.CHW,
    image_key="image",
    label_key="label",
    scale=armory.data.Scale(dtype=armory.data.DataType.FLOAT, max=1.0),
)

batch_size: int | None

dataset: Dataset[T_co]

drop_last: bool

num_workers: int

pin_memory: bool

pin_memory_device: str

prefetch_factor: int | None

sampler: Sampler | Iterable

timeout: float

class dataset.ObjectDetectionDataLoader(*args, boxes_key: str, dim: ImageDimensions, format: BBoxFormat, image_key: str, labels_key: str, objects_key: str, scale: Scale, **kwargs)

Bases: ShuffleableDataLoader

Data loader for object detection datasets.

Example:

import armory.data
from armory.dataset import ObjectDetectionDataLoader

# assuming `dataset` has been defined elsewhere
dataloader = ObjectDetectionDataLoader(
    dataset,
    boxes_key="boxes",
    dim=armory.data.ImageDimensions.HWC,
    format=armory.data.BBoxFormat.XYWH,
    image_key="image",
    labels_key="category",
    objects_key="objects",
    scale=armory.data.Scale(dtype=armory.data.DataType.FLOAT, max=1.0),
)

batch_size: int | None

dataset: Dataset[T_co]

drop_last: bool

num_workers: int

pin_memory: bool

pin_memory_device: str

prefetch_factor: int | None

sampler: Sampler | Iterable

timeout: float

class dataset.ShuffleableDataLoader(*args, seed: int | None = None, shuffle: bool = False, **kwargs)

Bases: DataLoader

Extension to the PyTorch data loader that enables determinstic shuffling when the data loader is used multiple times.

Example:

from armory.dataset import ShuffleableDataLoader

# assuming `dataset` has been defined elsewhere
dataloader = ShuffleableDataLoader(
    dataset,
    shuffle=True,
    seed=8675309,
)

batch1 = next(iter(dataloader))
batch2 = next(iter(dataloader))
assert batch1 == batch2

batch_size: int | None

dataset: Dataset[T_co]

drop_last: bool

num_workers: int

pin_memory: bool

pin_memory_device: str

prefetch_factor: int | None

sampler: Sampler | Iterable

timeout: float

class dataset.TupleDataset(dataset, keys: Iterable[str])

Bases: ArmoryDataset

Dataset wrapper with a pre-applied adapter to adapt tuples to map-like samples.

Example:

from armory.dataset import TupleDataset

# assuming `dataset` has been defined elsewhere
print(dataset[0])
# output: [[0, 0, 0], [0, 0, 0]], [5]

tuple_ds = TupleDataset(dataset, ("image", "label"))
print(tuple_ds[0])
# output: {'image': [[0, 0, 0], [0, 0, 0]], 'label': [5]}