Datasets

BatchAugmenter

Bases: abc.ABC

Abstract base class for batch-level augmentations.

Batch augmenters run after DeepVisionDataset.__getitem__, after normal sample-level augmentations/preprocessing, and after deepvision_collate_fn has padded images and targets to a common shape.

Subclasses must implement get_new_batch and return a new (images, targets, names) tuple. Keeping the batch length unchanged is recommended for stable training and logging.

Custom batch augmenter

from deepvisiontools.datasets.additional_augmentations import BatchAugmenter

class MyBatchAugmenter(BatchAugmenter):
    def get_new_batch(self, images, targets, names=None):
        return images, targets, names

get_new_batch(images, targets, names=None) abstractmethod

Return an augmented batch.

Source code in src/deepvisiontools/datasets/additional_augmentations.py

@abstractmethod
def get_new_batch(
    self,
    images: Tensor,
    targets: BatchData,
    names: list[str] | None = None,
) -> BatchOutput:
    """Return an augmented batch."""

MosaicBatchAugmenter(num_patches=4, p=0.5, *, shuffle=True)

Bases: deepvisiontools.datasets.additional_augmentations.BatchAugmenter

Create same-size mosaic samples by exchanging grid patches inside a batch.

A full group of num_patches samples produces num_patches new mosaic samples, so the batch size is preserved. If the batch size is not divisible by num_patches, the remaining samples are left unchanged. If the runtime batch is smaller than num_patches (for example a last partial DDP batch), the batch is returned unchanged.

Parameters:

Name	Type	Description	Default
num_patches	`Literal[1, 2, 4, 6, 8, 9, 12]`	Number of patches used per mosaic. Defaults to 4.	4
p	`float`	Probability of applying mosaic to a batch. Defaults to 0.5.	0.5
shuffle	`bool`	Randomize sample grouping and patch order. Defaults to True.	True

Supported target types

Works with BboxData, InstanceMaskData, and SemanticMaskData through the existing crop, pad, _new_like, empty, and BatchData helpers.

Usage

from deepvisiontools.datasets import DeepVisionLoader
from deepvisiontools.datasets.additional_augmentations import (
    MosaicBatchAugmenter,
)

trainloader = DeepVisionLoader(
    trainset,
    batch_size=8,
    batch_augmenter=MosaicBatchAugmenter(num_patches=4, p=0.5),
)

Source code in src/deepvisiontools/datasets/additional_augmentations.py

def __init__(
    self,
    num_patches: PatchCount = 4,
    p: float = 0.5,
    *,
    shuffle: bool = True,
) -> None:
    if int(num_patches) not in _MOSAIC_GRID_SHAPES:
        raise ValueError(
            "MosaicBatchAugmenter.num_patches must be one of "
            f"{sorted(_MOSAIC_GRID_SHAPES)}. Got {num_patches}."
        )
    self.num_patches = int(num_patches)
    self.p = _validate_probability(p)
    self.shuffle = bool(shuffle)

RandomCenterCropAndResize(crop, resize, p=0.5, **kwargs)

Bases: torchvision.transforms.v2.Transform

Randomly center-crop an image/target pair, then resize it back to a fixed size.

Parameters:

Name	Type	Description	Default
crop	`int | Sequence[int]`	Crop size forwarded to torchvision.transforms.v2.CenterCrop.	required
resize	`int | Sequence[int]`	Output size forwarded to torchvision.transforms.v2.Resize.	required
p	`float`	Probability of applying the transform. Defaults to 0.5.	0.5

Usage

from deepvisiontools.datasets.additional_augmentations import (
    RandomCenterCropAndResize,
)

dataset.augmentation = [
    RandomCenterCropAndResize(crop=(512, 512), resize=(640, 640), p=0.2)
]

Source code in src/deepvisiontools/datasets/additional_augmentations.py

def __init__(
    self,
    crop: int | Sequence[int],
    resize: int | Sequence[int],
    p: float = 0.5,
    **kwargs: Any,
) -> None:
    super().__init__(**kwargs)
    self.p = _validate_probability(p)
    self.crop = T.CenterCrop(crop)
    self.resize = T.Resize(resize)

RandomChangeBackground(background_dir, p=0.5, *, recursive=False, **kwargs)

Bases: torchvision.transforms.v2.Transform

Replace image background pixels with a random image from a folder.

The foreground is inferred from the first mask-like target received by the transform: pixels where the mask is non-zero are preserved, and all other pixels are replaced by the random background. This is intended for InstanceMaskData and SemanticMaskData. Bounding boxes are intentionally rejected because they do not define the true object silhouette.

Parameters:

Name	Type	Description	Default
background_dir	`str | Path`	Folder containing candidate background images (jpg, jpeg, png, tif, or tiff).	required
p	`float`	Probability of changing the background. Defaults to 0.5.	0.5
recursive	`bool`	If True, scan subfolders too. Defaults to False.	False

Bounding boxes

RandomChangeBackground is not compatible with BboxData, because a bounding box only preserves a rectangular region and would leak background pixels inside the box.

Usage

from deepvisiontools.datasets.additional_augmentations import (
    RandomChangeBackground,
)

trainset.augmentation = [RandomChangeBackground("backgrounds", p=0.4)]

Source code in src/deepvisiontools/datasets/additional_augmentations.py

def __init__(
    self,
    background_dir: str | Path,
    p: float = 0.5,
    *,
    recursive: bool = False,
    **kwargs: Any,
) -> None:
    super().__init__(**kwargs)
    self.p = _validate_probability(p)
    self.background_dir = Path(background_dir)
    self.recursive = bool(recursive)
    self.background_paths = self._scan_backgrounds()

    if not self.background_paths:
        raise FileNotFoundError(
            "RandomChangeBackground did not find any supported background image "
            f"in {self.background_dir}. Supported suffixes: {_IMAGE_SUFFIXES}."
        )

RandomCropAndResize(crop, resize, p=0.5, **kwargs)

Bases: torchvision.transforms.v2.Transform

Randomly crop an image/target pair, then resize it back to a fixed size.

Parameters:

Name	Type	Description	Default
crop	`int | Sequence[int]`	Crop size forwarded to torchvision.transforms.v2.RandomCrop.	required
resize	`int | Sequence[int]`	Output size forwarded to torchvision.transforms.v2.Resize.	required
p	`float`	Probability of applying the transform. Defaults to 0.5.	0.5

Usage

from deepvisiontools.datasets.additional_augmentations import (
    RandomCropAndResize,
)

augs = [RandomCropAndResize(crop=(512, 512), resize=(640, 640), p=0.3)]
dataset.augmentation = augs

Source code in src/deepvisiontools/datasets/additional_augmentations.py

def __init__(
    self,
    crop: int | Sequence[int],
    resize: int | Sequence[int],
    p: float = 0.5,
    **kwargs: Any,
) -> None:
    super().__init__(**kwargs)
    self.p = _validate_probability(p)
    self.crop = T.RandomCrop(crop)
    self.resize = T.Resize(resize)

RandomPadAndResize(maxpad, resize, p=0.5, **kwargs)

Bases: torchvision.transforms.v2.Transform

Randomly pad an image/target pair, then resize it back to a fixed size.

This gives a light zoom-out effect while preserving the target geometry through torchvision v2 transforms.

Parameters:

Name	Type	Description	Default
maxpad	`int | Sequence[int]`	Maximum random padding. An integer applies the same bound to all sides. A sequence must be (left, top, right, bottom).	required
resize	`tuple[int, int]`	Output size (height, width).	required
p	`float`	Probability of applying the transform. Defaults to 0.5.	0.5

Usage

from deepvisiontools.datasets.additional_augmentations import RandomPadAndResize

dataset.augmentation = [RandomPadAndResize(maxpad=64, resize=(640, 640), p=0.2)]

Source code in src/deepvisiontools/datasets/additional_augmentations.py

def __init__(
    self,
    maxpad: int | Sequence[int],
    resize: tuple[int, int],
    p: float = 0.5,
    **kwargs: Any,
) -> None:
    super().__init__(**kwargs)
    self.p = _validate_probability(p)
    self.max_pad = _normalize_padding_bounds(maxpad)
    self.resize = T.Resize(resize)

AnnotationReader

Bases: typing.Protocol

Strategy: enumerate and read (image_name, BaseData).

scan()

Return the ordered list of image filenames (relative to image dir).

Source code in src/deepvisiontools/datasets/base.py

def scan(self) -> list[str]:
    """Return the ordered list of image filenames (relative to image dir)."""
    ...

read_item(index)

Return (image_name, target) for given dataset index.

Source code in src/deepvisiontools/datasets/base.py

def read_item(self, index: int) -> tuple[str, BaseData]:
    """Return (image_name, target) for given dataset index."""
    ...

validation_report(indices=None)

Optional lightweight dataset-validation report based only on annotations/metadata.

Implementations should avoid loading images. Return None when the reader cannot provide such a report cheaply.

Source code in src/deepvisiontools/datasets/base.py

def validation_report(
    self, indices: Sequence[int] | None = None
) -> dict[str, Any] | None:
    """
    Optional lightweight dataset-validation report based only on annotations/metadata.

    Implementations should avoid loading images. Return None when the reader
    cannot provide such a report cheaply.
    """
    ...

AnnotationWriter

Bases: typing.Protocol

Strategy: write individual annotations, then finalize once.

finalize(destination_root)

Called once after all items were written. May group/merge files.

Source code in src/deepvisiontools/datasets/base.py

def finalize(self, destination_root: str | Path) -> None:
    """Called once after all items were written. May group/merge files."""

DeepVisionDataset(dataset_root, reader, writer=None, *, data_type_reader=None, augmentation=None, preprocessing=_DEFAULT_PREPROCESSING, device=None, category_ids=None, cast_to=None, keep_indices=None)

Bases: torch.utils.data.Dataset

Dataset class for deepvisiontools. It handles multiples readers, writers (to export dataset), supports torchvision v2 augmentations, preprocessing etc.

Parameters:

Name	Type	Description	Default
dataset_root	str | pathlib.Path	path to root folder	required
reader	deepvisiontools.datasets.base.AnnotationReader | str	reader type. Either str matching registry (check available_readers() for reader list) or AnnotationReader class for custom.	required
writer	deepvisiontools.datasets.base.AnnotationWriter | str | None	Similar than reader but use for exporting datasets. Defaults to None.	None
data_type_reader	str | None	Used mostly for coco dataset -> enforce BaseData type returned. Defaults to None.	None
augmentation	Optional[List[torchvision.transforms.v2.Transform]]	List of augmentation. You must use torchvision.transforms.v2 Transform (they use tv_tensors). Defaults to None.	None
preprocessing	Optional[collections.abc.Callable]	callable for image tensor preprocessing. Defaults to imagenet_preprocessing().	deepvisiontools.datasets.base._DEFAULT_PREPROCESSING
device	Optional[str]	device to store the BaseData. If None uses current().device Defaults to None.	None
category_ids	Optional[Dict[int, str]]	Dict to name categories. Used only in visualization. Defaults to None.	None
cast_to	Optional[str]	Cast BaseData from reader to another type. casting must be registered in _CASTERS. If None, preserves the reader type. Defaults to None.	None
keep_indices	Optional[collections.abc.Sequence[int]]	preserves only a subset of indices. Defaults to None.	None

Structure of dataset

For registered readers, the dataset folder must respect the following structure

"coco"

• dataset_root/
  • images
  • coco_annotations.json

the json file must have this exact name. The json file is a standard coco format (Dict with "images", "annotations" and "categories" keys, with a list of dict for each of them.)

"semantic_mask_folder":

• dataset_root/
  • images
  • masks

masks are single-channel images with values in [0...Ncls] where 0 is necessarly background.

Tips

• You can call the split method to randomly divide the dataset into train, valid and test subsets
• You can call the keep_indices method to keep only the selected indexes
• You can modify the writer using the property writer so the export method returns the provided format. Note that you may need to modify the cast_to attribute to match writer requirements
• You can export a (sub)dataset to save a split and (optionally) visualizations. Caution to set preprocessing/augmentation to None if you want original images before normalization.
• You can concat multiple DeepVisionDatasets (including subsets etc.) using the concat static method.

Usage

# imports
from deepvisiontools.datasets.base import DeepVisionDataset
from deepvisiontools.config import current, set_current
from torchvision.transforms.v2 import RandomHorizontalFlip, RandomAffine, ColorJitter
# augmentation list
augs = [RandomHorizontalFlip(), RandomAffine(degrees=27, translate=[0.1, 0.2]), ColorJitter(hue=0.3)]
# change configs num classes if needed
mask_cfg = current().data_semantic_mask.update(num_classes=3)
cfg = current().update(data_semantic_mask=mask_cfg)
set_current(cfg)
# Create a dataset
semantic_dataset = DeepVisionDataset("path/to/root/semantic/dataset", "semantic_mask_folder", preprocessing=None)   # preprocessing default to imagenet_preprocessing()
instance_dataset = DeepVisionDataset("path/to/root/semantic/dataset", "coco", data_type_reader="instance_mask", augmentation=augs)  # "coco" needs data_type_reader to be either "bbox" or "instance_mask"
# split dataset
train, valid, test = instance_dataset.split([0.8, 0.1, 0.1])
# export
valid.preprocessing = None  # We want to export original images, not normalized ones.
valid.export("path/to/new/dataset")
train.writer = "semantic_mask_folder"   # prepare exporting to semantic mask dataset
train.cast_to = "semantic_mask"         # enforce InstanceMaskData to be converted to SemanticMaskData
train.export("path/to/converted/train/dataset")
# access item
image, target, img_name = test[0]   # image: torch.Tensor, target: BaseData, img_name: str

Source code in src/deepvisiontools/datasets/base.py

def __init__(
    self,
    dataset_root: str | Path,
    reader: AnnotationReader | str,
    writer: AnnotationWriter | str | None = None,
    *,
    data_type_reader: str | None = None,
    augmentation: list[Transform] | None = None,
    preprocessing: Callable | None = _DEFAULT_PREPROCESSING,
    device: str | None = None,
    category_ids: dict[int, str] | None = None,
    cast_to: str | None = None,
    keep_indices: Sequence[int] | None = None,
) -> None:
    # use default writer same as reader if not provided
    if writer is None and isinstance(reader, str):
        writer = reader

    # Resolve reader
    if isinstance(reader, str):
        if reader in available_readers(include_aliases=True):
            reader = get_reader(reader, dataset_root, data_type=data_type_reader)
        else:
            raise ValueError(
                f"{reader} is not a registered reader type. "
                f"Available: {available_readers(include_aliases=True)}"
            )

    self.root = Path(dataset_root)
    self.image_dir = self.root / "images"

    self.reader = reader
    self._writer: AnnotationWriter | None = None
    self.writer = writer

    self._cat_ids = (
        category_ids if category_ids is not None else self.reader.category_ids
    )
    self.augmentation = augmentation or []
    self.preprocessing = preprocessing
    self.device = device or current().device

    scan = self.reader.scan()
    self._indices = (
        list(range(len(scan))) if keep_indices is None else list(keep_indices)
    )
    self._scan = scan

    self.cast_to = cast_to  # optional conversion
    self._declared_reader_type = data_type_reader  # e.g., 'bbox'

    # Lightweight annotation-only validation (no image loading)
    self._validate()

to_data_type(dst_key)

Return a view that converts every sample's target to dst_key lazily.

Source code in src/deepvisiontools/datasets/base.py

def to_data_type(self, dst_key: str) -> DeepVisionDataset:
    """Return a view that converts every sample's target to `dst_key` lazily."""
    return DeepVisionSubDataset(self, cast_to=dst_key)

split(proportions)

Randomly split the current dataset view into sub-views.

Indices are interpreted relative to the current dataset view, then mapped back to the underlying reader indices by DeepVisionSubDataset. This makes split() safe on both base datasets and already-filtered subdatasets.

Source code in src/deepvisiontools/datasets/base.py

def split(self, proportions: Sequence[float]) -> list[DeepVisionDataset]:
    """
    Randomly split the current dataset view into sub-views.

    Indices are interpreted relative to the current dataset view, then mapped
    back to the underlying reader indices by DeepVisionSubDataset. This makes
    split() safe on both base datasets and already-filtered subdatasets.
    """
    s = sum(proportions)
    assert round(s, 6) == 1.0, f"Split proportions must sum to 1, got {proportions}"

    import random

    positions = list(range(len(self)))
    random.shuffle(positions)

    n = len(positions)
    cuts = [round(p * n) for p in proportions]

    out: list[DeepVisionDataset] = []
    start = 0
    for c in cuts[:-1]:
        out.append(DeepVisionSubDataset(self, indices=positions[start : start + c]))
        start += c

    out.append(DeepVisionSubDataset(self, indices=positions[start:]))
    return out

keep_indices(indices)

Keep a subset of samples from the current dataset view.

The provided indices are positions relative to this dataset/view, not raw reader indices. This matches PyTorch Subset semantics and composes correctly across nested subdatasets.

Source code in src/deepvisiontools/datasets/base.py

def keep_indices(
    self,
    indices: list[int] | torch.Tensor | slice,
) -> DeepVisionDataset:
    """
    Keep a subset of samples from the current dataset view.

    The provided indices are positions relative to this dataset/view, not raw
    reader indices. This matches PyTorch Subset semantics and composes correctly
    across nested subdatasets.
    """
    if isinstance(indices, torch.Tensor):
        if indices.ndim != 1:
            raise ValueError(f"indices must be a 1-D tensor, got {indices.shape}")
        indices = [int(i) for i in indices.tolist()]
    elif isinstance(indices, slice):
        indices = list(range(len(self)))[indices]
    else:
        indices = [int(i) for i in indices]

    return DeepVisionSubDataset(self, indices=indices)

export(destination_root, *, num_visualizations=0)

Write images + annotations using the configured writer.

Source code in src/deepvisiontools/datasets/base.py

def export(
    self,
    destination_root: str | Path,
    *,
    num_visualizations: int | str = 0,
) -> None:
    """
    Write images + annotations using the configured writer.
    """
    if self.writer is None:
        raise RuntimeError(
            "No writer configured on this dataset. Use a FormatPlugin or set_writer(...) first."
        )

    # check actual data type matches writer
    _, t0, _ = self[0]
    actual_key = data_key_from_instance(t0)
    supported = normalize_supported_types(
        getattr(self._writer, "supported_types", None)
    )
    if supported is not None and actual_key not in supported:
        raise TypeError(
            f"Writer {type(self._writer).__name__} supports {sorted(supported)}, "
            f"but got runtime target type '{actual_key}'. "
            f"Set as_data_type to a supported type or switch writer."
        )

    # prepare export
    dest = Path(destination_root)
    img_out = dest / "images"
    visu_out = dest / "visualizations"
    img_out.mkdir(parents=True, exist_ok=True)
    if num_visualizations == "all" or (
        isinstance(num_visualizations, int) and num_visualizations > 0
    ):
        visu_out.mkdir(parents=True, exist_ok=True)

    remaining_visu = (
        len(self) if num_visualizations == "all" else int(num_visualizations)
    )

    # export
    for i in tqdm(
        range(len(self)), total=len(self._indices), desc="Exporting dataset: "
    ):
        image, target, name = self[i]

        # save visualization if requested
        if remaining_visu > 0:
            vis = target.visualize(image)
            save_image(vis.detach().cpu(), (visu_out / f"visu__{name}").as_posix())
            remaining_visu -= 1

        # save image (post-preprocessing), keep original filename
        save_image(image.detach().cpu(), (img_out / name).as_posix())

        # delegate annotation writing
        self.writer.write_item(name, image, target, self.category_ids)

    # allow writer to group files / finalize
    self.writer.finalize(dest)

concat(datasets, writer=None, *, prefix_names=True) staticmethod

Concatenate several DeepVisionDataset views into one DeepVision-compatible view.

The returned dataset can be indexed, split, loaded with DeepVisionLoader, and exported with the configured writer. Input datasets are deep-copied so later edits on the concatenated view do not mutate the original datasets.

Source code in src/deepvisiontools/datasets/base.py

@staticmethod
def concat(
    datasets: Sequence[DeepVisionDataset],
    writer: AnnotationWriter | str | None = None,
    *,
    prefix_names: bool = True,
) -> DeepVisionConcatDataset:
    """
    Concatenate several DeepVisionDataset views into one DeepVision-compatible view.

    The returned dataset can be indexed, split, loaded with DeepVisionLoader,
    and exported with the configured writer. Input datasets are deep-copied so
    later edits on the concatenated view do not mutate the original datasets.
    """
    return DeepVisionConcatDataset(
        datasets,
        writer=writer,
        prefix_names=prefix_names,
    )

DeepVisionLoader(dataset, batch_size=1, shuffle=True, sampler=None, batch_sampler=None, num_workers=0, collate_fn=deepvision_collate_fn, pin_memory=False, drop_last=False, timeout=0, worker_init_fn=None, multiprocessing_context=None, generator=None, *, prefetch_factor=None, persistent_workers=False, pin_memory_device='', in_order=True, batch_augmenter=None)

Bases: torch.utils.data.DataLoader

DeepVisionTools dataloader with BatchData-aware collation.

DeepVisionLoader is a thin wrapper around torch.utils.data.DataLoader. It uses deepvision_collate_fn by default so targets are returned as BatchData, and it can optionally apply a batch-level augmentation after collation.

Batch augmentations are applied on already-collated (images, targets, names) batches. They are useful for operations that need several samples at once, such as mosaic-style augmentations.

Parameters:

Name	Type	Description	Default
dataset	`DeepVisionDataset`	Dataset to wrap.	required
batch_size	`int`	Number of samples per batch. Defaults to 1.	1
shuffle	`bool`	Whether to shuffle samples. Defaults to True.	True
sampler	`Any | None`	Optional PyTorch sampler. Defaults to None.	None
batch_sampler	`Any | None`	Optional PyTorch batch sampler. Defaults to None.	None
num_workers	`int`	Number of worker processes. Defaults to 0.	0
collate_fn	`Any`	Function used to collate samples. Defaults to deepvision_collate_fn.	deepvisiontools.datasets.base.deepvision_collate_fn
pin_memory	`bool`	Whether to pin CPU memory before transfer to CUDA. Defaults to False.	False
drop_last	`bool`	Whether to drop the last incomplete batch. Defaults to False.	False
timeout	`float`	DataLoader worker timeout in seconds. Defaults to 0.	0
worker_init_fn	`Any | None`	Optional worker initialization function. Defaults to None.	None
multiprocessing_context	`Any | None`	Optional multiprocessing context. Defaults to None.	None
generator	`torch.Generator | None`	Optional random generator used by the loader. Defaults to None.	None
prefetch_factor	`int | None`	Number of batches prefetched by each worker. Defaults to None.	None
persistent_workers	`bool`	Whether to keep workers alive between epochs. Defaults to False.	False
pin_memory_device	`str`	Device used for pinned memory. Defaults to "".	''
in_order	`bool`	Whether workers must return batches in order. Defaults to True.	True
batch_augmenter	`BatchAugmenter | None`	Optional batch-level augmenter applied after collation. Defaults to None.	None

Usage

from deepvisiontools.datasets import DeepVisionDataset, DeepVisionLoader

trainset = DeepVisionDataset(
    "path/to/dataset",
    reader="coco",
    data_type_reader="instance_mask",
)

trainloader = DeepVisionLoader(
    trainset,
    batch_size=4,
    shuffle=True,
    num_workers=4,
)

images, targets, names = next(iter(trainloader))

MosaicBatchAugmenter

from deepvisiontools.datasets import DeepVisionDataset, DeepVisionLoader
from deepvisiontools.datasets.additional_augmentations import MosaicBatchAugmenter

trainset = DeepVisionDataset(
    "path/to/dataset",
    reader="coco",
    data_type_reader="instance_mask",
)

trainloader = DeepVisionLoader(
    trainset,
    batch_size=4,
    shuffle=True,
    batch_augmenter=MosaicBatchAugmenter(
        num_patches=4,
        p=0.5,
        randomize_order=True,
    ),
)

Source code in src/deepvisiontools/datasets/base.py

def __init__(
    self,
    dataset,
    batch_size=1,
    shuffle=True,
    sampler=None,
    batch_sampler=None,
    num_workers=0,
    collate_fn=deepvision_collate_fn,
    pin_memory=False,
    drop_last=False,
    timeout=0,
    worker_init_fn=None,
    multiprocessing_context=None,
    generator=None,
    *,
    prefetch_factor=None,
    persistent_workers=False,
    pin_memory_device="",
    in_order=True,
    batch_augmenter: None | BatchAugmenter = None,
) -> None:
    self.batch_augmenter = batch_augmenter
    self.base_collate_fn = collate_fn
    if batch_augmenter is not None and not getattr(
        collate_fn,
        "_deepvision_batch_augmented_collate",
        False,
    ):
        collate_fn = build_batch_augmented_collate_fn(
            collate_fn,
            batch_augmenter,
        )

    super().__init__(
        dataset,
        batch_size,
        shuffle,
        sampler,
        batch_sampler,
        num_workers,
        collate_fn,
        pin_memory,
        drop_last,
        timeout,
        worker_init_fn,
        multiprocessing_context,
        generator,
        prefetch_factor=prefetch_factor,
        persistent_workers=persistent_workers,
        pin_memory_device=pin_memory_device,
        in_order=in_order,
    )

visualize(out_folder, num_visu='all')

Generate batch visualizations.

This is useful to inspect the final batched output, including optional batch-level augmentations such as mosaic.

Parameters:

Name	Type	Description	Default
out_folder	`str | Path`	Output folder where visualization images are saved.	required
num_visu	`int | Literal["all"]`	Number of batches to visualize. Defaults to "all".	'all'

Source code in src/deepvisiontools/datasets/base.py

def visualize(self, out_folder: str | Path, num_visu: int | Literal["all"] = "all"):
    """Generate batch visualizations.

    This is useful to inspect the final batched output, including optional
    batch-level augmentations such as mosaic.

    Args:
        out_folder (`str | Path`): Output folder where visualization images are saved.
        num_visu (`int | Literal["all"]`, optional): Number of batches to visualize. Defaults to "all".
    """
    if num_visu == "all":
        num_visu = len(self)
    elif not isinstance(num_visu, int):
        raise ValueError(
            f"{num_visu} is an invalid number of visu argument (can be int or 'all')"
        )

    out_folder = out_folder if isinstance(out_folder, Path) else Path(out_folder)
    out_folder.mkdir(exist_ok=True, parents=True)

    for b, (ims, targs, _) in tqdm(enumerate(self), total=num_visu):
        if b >= num_visu:
            break

        image_items: list[torch.Tensor] = list(ims.unbind(0))
        for i, im in enumerate(image_items):
            targ: BaseData = targs.datas[i]
            visu = targ.visualize(im)
            save_image(
                visu.detach().cpu(),
                (out_folder / f"batch_{b}_item_{i}.png").as_posix(),
            )

DeepVisionSubDataset(base, *, indices=None, cast_to=None)

Bases: deepvisiontools.datasets.base.DeepVisionDataset

A thin view over a base dataset with its own index list and optional cast override.

Source code in src/deepvisiontools/datasets/base.py

def __init__(
    self,
    base: DeepVisionDataset,
    *,
    indices: Sequence[int] | None = None,
    cast_to: str | None = None,
) -> None:
    self.__dict__ = base.__dict__.copy()

    parent_indices = list(base._indices)

    if indices is None:
        # Preserve the parent view exactly. This fixes subset.to_data_type(...)
        # and any other operation that creates a view without new indices.
        self._indices = parent_indices
    else:
        mapped_indices: list[int] = []
        for index in indices:
            ii = int(index)
            if ii < 0:
                ii += len(parent_indices)
            if ii < 0 or ii >= len(parent_indices):
                raise IndexError(
                    f"Subdataset index {index} is out of range for view "
                    f"of length {len(parent_indices)}."
                )
            mapped_indices.append(parent_indices[ii])
        self._indices = mapped_indices

    if cast_to is not None:
        self.cast_to = cast_to

CocoReader(dataset_root, *, data_type='bbox', category_sort_key='name', annotation_file=DEFAULT_COCO_ANNOT)

Bases: deepvisiontools.datasets.base.AnnotationReader

COCO reader (no pycocotools). Supports bbox and instance mask workflows.

Source code in src/deepvisiontools/datasets/coco_dataset.py

def __init__(
    self,
    dataset_root: str | Path,
    *,
    data_type: Literal["bbox", "instance_mask"] = "bbox",
    category_sort_key: Literal["name", "id"] = "name",
    annotation_file: str = DEFAULT_COCO_ANNOT,
) -> None:
    self.root = Path(dataset_root)
    with open(self.root / annotation_file) as f:
        d = json.load(f)

    # Normalize lists
    self._images: list[dict] = (
        list(d["images"])
        if isinstance(d["images"], list)
        else list(d["images"].values())
    )
    _anns: list[dict] = (
        list(d["annotations"])
        if isinstance(d["annotations"], list)
        else list(d["annotations"].values())
    )
    _cats: list[dict] = (
        list(d["categories"])
        if isinstance(d["categories"], list)
        else list(d["categories"].values())
    )

    # Keep raw declarations exactly as present in the dataset for validation.
    self._raw_cat_ids: dict[int, str] = {
        int(c["id"]): str(c["name"]) for c in _cats if "id" in c and "name" in c
    }

    # Runtime behavior kept unchanged:
    # map raw category ids -> contiguous internal indices.
    cats_sorted = sorted(
        _cats,
        key=(
            (lambda c: c["name"])
            if category_sort_key == "name"
            else (lambda c: int(c["id"]))
        ),
    )
    self._cat_ids: dict[int, str] = {
        i: str(c["name"]) for i, c in enumerate(cats_sorted)
    }
    self._origCatId_to_idx: dict[int, int] = {
        int(c["id"]): i for i, c in enumerate(cats_sorted)
    }

    # Index annotations by image_id.
    anns_by_img: dict[int, list[dict]] = {}
    for a in _anns:
        anns_by_img.setdefault(int(a["image_id"]), []).append(a)
    self._anns_by_img = anns_by_img

    self._data_type = data_type

validation_report(indices=None)

Lightweight annotation-only validation report.

This inspects only the already-loaded COCO JSON structures and never touches the image files. When indices is provided, the report is restricted to the corresponding dataset subset.

Source code in src/deepvisiontools/datasets/coco_dataset.py

def validation_report(
    self, indices: Sequence[int] | None = None
) -> dict[str, object] | None:
    """
    Lightweight annotation-only validation report.

    This inspects only the already-loaded COCO JSON structures and never touches
    the image files. When `indices` is provided, the report is restricted to the
    corresponding dataset subset.
    """
    if indices is None:
        selected_images = self._images
    else:
        selected_images = []
        for i in indices:
            ii = int(i)
            if 0 <= ii < len(self._images):
                selected_images.append(self._images[ii])

    used_category_ids: set[int] = set()
    num_annotations = 0

    for im in selected_images:
        anns = self._anns_by_img.get(im["id"], [])
        num_annotations += len(anns)
        for a in anns:
            cid = a.get("category_id", None)
            if cid is None:
                continue
            try:
                used_category_ids.add(int(cid))
            except Exception:
                continue

    declared_category_ids = sorted(self._raw_cat_ids.keys())
    used_category_ids_sorted = sorted(used_category_ids)

    unused_declared_category_ids = [
        cid for cid in declared_category_ids if cid not in used_category_ids
    ]
    undeclared_used_category_ids = [
        cid for cid in used_category_ids_sorted if cid not in self._raw_cat_ids
    ]

    return {
        "reader": "coco",
        "num_images": len(selected_images),
        "num_annotations": int(num_annotations),
        "declared_category_ids": declared_category_ids,
        "declared_category_names": dict(self._raw_cat_ids),
        "used_category_ids": used_category_ids_sorted,
        "unused_declared_category_ids": unused_declared_category_ids,
        "undeclared_used_category_ids": undeclared_used_category_ids,
    }

CocoWriter(annotation_file=DEFAULT_COCO_ANNOT, write_polygons_for_instances=True)

Bases: deepvisiontools.datasets.base.AnnotationWriter

COCO writer that emits a single coco_annotations.json on finalize.

Source code in src/deepvisiontools/datasets/coco_dataset.py

def __init__(
    self,
    annotation_file: str = DEFAULT_COCO_ANNOT,
    write_polygons_for_instances: bool = True,
):
    self.supported_types = {"bbox", "instance_mask"}
    self._annotation_file = annotation_file
    self._write_poly = bool(write_polygons_for_instances)

    self._images: list[dict] = []
    self._annotations: list[dict] = []
    self._categories: list[dict] = []

    self._img_id = 0
    self._ann_id = 0

SemanticFolderReader(dataset_root, *, masks_subdir='masks', mask_ext=None, category_ids=None, **kwargs)

Bases: deepvisiontools.datasets.base.AnnotationReader

Reader for folder-based semantic masks:

Source code in src/deepvisiontools/datasets/semantic_dataset.py

def __init__(
    self,
    dataset_root: str | Path,
    *,
    masks_subdir: str = "masks",
    mask_ext: str | None = None,
    category_ids: dict[int, str] | None = None,
    **kwargs: object,
) -> None:
    self.root = (
        dataset_root if isinstance(dataset_root, Path) else Path(dataset_root)
    )
    self.image_dir = self.root / "images"
    self.mask_dir = self.root / masks_subdir
    self._imgs = sorted(p.name for p in self.image_dir.iterdir() if p.is_file())
    self._mask_ext = mask_ext
    self._cat_ids = category_ids or {}

validation_report(indices=None)

No cheap annotation-only validation is available for folder-based semantic masks without reading the mask files, so we explicitly opt out.

Source code in src/deepvisiontools/datasets/semantic_dataset.py

def validation_report(self, indices=None):
    """
    No cheap annotation-only validation is available for folder-based semantic masks
    without reading the mask files, so we explicitly opt out.
    """
    return None

SemanticFolderWriter(masks_subdir='masks', ext='.tif')

Bases: deepvisiontools.datasets.base.AnnotationWriter

Writer for folder-based semantic masks.

Source code in src/deepvisiontools/datasets/semantic_dataset.py

def __init__(self, masks_subdir: str = "masks", ext: str = ".tif"):
    self.supported_types = {"semantic_mask"}
    self._pending: list[tuple[torch.Tensor, str]] = []
    self._masks_subdir = str(masks_subdir)

    ext = str(ext)
    if not ext.startswith("."):
        ext = "." + ext
    self._ext = ext

available_readers(*, include_aliases=False)

Return a list of registered reader keys.

Source code in src/deepvisiontools/datasets/base.py

def available_readers(*, include_aliases: bool = False) -> list[str]:
    """
    Return a list of registered reader keys.
    """
    keys = sorted(_READERS._factories.keys())
    if include_aliases:
        alias_keys = sorted(_READERS._aliases.keys())
        return keys + alias_keys
    return keys

available_writers(*, include_aliases=False)

Return a list of registered writer keys. Set include_aliases=True to include alias keys as well.

Source code in src/deepvisiontools/datasets/base.py

def available_writers(*, include_aliases: bool = False) -> list[str]:
    """
    Return a list of registered writer keys.
    Set `include_aliases=True` to include alias keys as well.
    """
    keys = sorted(_WRITERS._factories.keys())
    if include_aliases:
        alias_keys = sorted(_WRITERS._aliases.keys())
        return keys + alias_keys
    return keys

register_reader(key)

Decorator to register a reader factory under a string key.

Source code in src/deepvisiontools/datasets/base.py

def register_reader(key: str) -> Callable[[ReaderFactory], ReaderFactory]:
    """
    Decorator to register a reader factory under a string key.
    """

    def _inner(factory: ReaderFactory) -> ReaderFactory:
        _READERS.register(key, factory)
        return factory

    return _inner

register_writer(key)

Decorator to register a writer factory under a string key.

Source code in src/deepvisiontools/datasets/base.py

def register_writer(key: str) -> Callable[[WriterFactory], WriterFactory]:
    """
    Decorator to register a writer factory under a string key.
    """

    def _inner(factory: WriterFactory) -> WriterFactory:
        _WRITERS.register(key, factory)
        return factory

    return _inner