Data API¶

API reference for data modules and datasets.

LitData Streaming¶

LitDataDataModule¶

class objdet.data.formats.litdata.LitDataDataModule(data_dir, train_subdir='train', val_subdir='val', test_subdir='test', class_names=None, class_index_mode=ClassIndexMode.TORCHVISION, **kwargs)[source]¶

Bases: BaseDataModule

Lightning DataModule for LitData optimized datasets.

This DataModule uses LitData’s StreamingDataset and StreamingDataLoader for efficient data loading from preprocessed datasets.

Features: - Uses native LitData streaming for optimal performance - Custom collate function for variable-length detection targets - Automatic transform chain for numpy-to-tensor conversion

Parameters:

data_dir (str | Path) – Root directory containing train/val/test subdirectories.
train_subdir (str) – Subdirectory name for training data.
val_subdir (str) – Subdirectory name for validation data.
test_subdir (str) – Subdirectory name for test data.
**kwargs (Any) – Additional arguments for BaseDataModule.

Example

>>> datamodule = LitDataDataModule(
...     data_dir="/data/coco_litdata",
...     batch_size=16,
...     num_workers=4,
... )
>>> datamodule.setup("fit")
>>> train_loader = datamodule.train_dataloader()

setup(stage=None)[source]¶

Set up datasets for each stage.

Parameters:: stage (str | None) – One of “fit”, “validate”, “test”, or “predict”.
Return type:: None

train_dataloader()[source]¶

Create training DataLoader using LitData’s StreamingDataLoader.

Return type:: Any
Returns:: StreamingDataLoader for training data with shuffling enabled.

val_dataloader()[source]¶

Create validation DataLoader using LitData’s StreamingDataLoader.

Return type:: Any
Returns:: StreamingDataLoader for validation data without shuffling.

test_dataloader()[source]¶

Create test DataLoader using LitData’s StreamingDataLoader.

Return type:: Any
Returns:: StreamingDataLoader for test data without shuffling.

DetectionStreamingDataset¶

Factory function that creates a StreamingDataset with detection transforms applied.

from objdet.data.formats.litdata import DetectionStreamingDataset

dataset = DetectionStreamingDataset(
    input_dir="/data/coco_litdata/train",
    shuffle=True,
    drop_last=False,
    transforms=None,  # Optional additional transforms
)

Parameters:

input_dir (str | Path): Directory containing LitData optimized chunks
shuffle (bool): Whether to shuffle the data (default: False)
drop_last (bool): Whether to drop the last incomplete batch (default: False)
transforms (Callable | None): Optional transform to apply after detection format conversion

Returns: A litdata.StreamingDataset configured for object detection.

create_streaming_dataloader¶

from objdet.data.formats.litdata import create_streaming_dataloader

loader = create_streaming_dataloader(
    dataset=streaming_dataset,
    batch_size=16,
    num_workers=4,
    pin_memory=True,
    drop_last=False,
)

Parameters:

dataset: A StreamingDataset instance
batch_size (int): Batch size for the dataloader
num_workers (int): Number of worker processes (default: 4)
pin_memory (bool): Whether to pin memory for faster GPU transfer (default: True)
drop_last (bool): Whether to drop the last incomplete batch (default: False)

Returns: A litdata.StreamingDataLoader with detection-specific collation.

Standard Formats¶

COCODataModule¶

class objdet.data.formats.coco.COCODataModule(data_dir, train_ann_file='annotations/instances_train2017.json', val_ann_file='annotations/instances_val2017.json', test_ann_file=None, train_img_dir='train2017', val_img_dir='val2017', test_img_dir='test2017', **kwargs)[source]¶

Bases: BaseDataModule

Lightning DataModule for COCO format datasets.

Parameters:

data_dir (str | Path) – Root directory containing images.
train_ann_file (str) – Path to training annotations JSON.
val_ann_file (str) – Path to validation annotations JSON.
test_ann_file (str | None) – Optional path to test annotations JSON.
train_img_dir (str) – Subdirectory for training images (relative to data_dir).
val_img_dir (str) – Subdirectory for validation images.
test_img_dir (str) – Subdirectory for test images.
**kwargs (Any) – Additional arguments for BaseDataModule.

Example

>>> datamodule = COCODataModule(
...     data_dir="/data/coco",
...     train_ann_file="annotations/instances_train2017.json",
...     val_ann_file="annotations/instances_val2017.json",
...     train_img_dir="train2017",
...     val_img_dir="val2017",
...     batch_size=16,
... )

setup(stage=None)[source]¶

Set up datasets for each stage.

Parameters:: stage (str | None) – One of “fit”, “validate”, “test”, or “predict”.
Return type:: None

COCODataset¶

class objdet.data.formats.coco.COCODataset(data_dir, ann_file, transforms=None, class_index_mode=ClassIndexMode.TORCHVISION)[source]¶

Bases: Dataset

PyTorch Dataset for COCO format annotations.

Parameters:

data_dir (str | Path) – Root directory containing images.
ann_file (str | Path) – Path to annotation JSON file.
transforms (Callable | None) – Optional transform to apply to images and targets.
class_index_mode (ClassIndexMode) – How to handle class indices.

images¶: List of image metadata dicts.

annotations¶: Dict mapping image_id to list of annotations.

categories¶: Dict mapping category_id to category info.

class_names¶: List of class names in order.

__len__()[source]¶

Return number of images in dataset.

Return type:: int

__getitem__(index)[source]¶

Get image and annotations by index.

Parameters:: index (int) – Dataset index.
Return type:: tuple[Tensor, DetectionTarget]
Returns:: Tuple of (image_tensor, target_dict).

VOCDataModule¶

class objdet.data.formats.voc.VOCDataModule(data_dir, train_split='trainval', val_split='val', test_split='test', **kwargs)[source]¶

Bases: BaseDataModule

Lightning DataModule for Pascal VOC format datasets.

Parameters:

data_dir (str | Path) – Root VOC directory.
train_split (str) – Training split name (default: “trainval”).
val_split (str) – Validation split name (default: “val”).
test_split (str) – Test split name (default: “test”).
**kwargs (Any) – Additional arguments for BaseDataModule.

setup(stage=None)[source]¶

Set up datasets for each stage.

Return type:: None

YOLODataModule¶

class objdet.data.formats.yolo.YOLODataModule(data_dir, images_subdir='images', labels_subdir='labels', train_split='train', val_split='val', test_split='test', **kwargs)[source]¶

Bases: BaseDataModule

Lightning DataModule for YOLO format datasets.

Expected directory structure:

data_dir/
├── images/
│   ├── train/
│   └── val/
└── labels/
    ├── train/
    └── val/

Parameters:

data_dir (str | Path) – Root dataset directory.
images_subdir (str) – Subdirectory name for images (default: “images”).
labels_subdir (str) – Subdirectory name for labels (default: “labels”).
train_split (str) – Training split name (default: “train”).
val_split (str) – Validation split name (default: “val”).
**kwargs (Any) – Additional arguments for BaseDataModule.

setup(stage=None)[source]¶

Set up datasets for each stage.

Return type:: None

Base Classes¶

BaseDataModule¶

class objdet.data.base.BaseDataModule(data_dir, dataset_format=DatasetFormat.COCO, class_index_mode=ClassIndexMode.TORCHVISION, class_names=None, batch_size=8, num_workers=4, pin_memory=True, persistent_workers=True, train_transforms=None, val_transforms=None, test_transforms=None, predict_path=None)[source]¶

Bases: LightningDataModule

Abstract base class for detection data modules.

This class provides common functionality for all detection datasets: - Standard DataLoader configuration - Train/val/test/predict dataset management - Augmentation pipeline integration - Class index mode handling

Subclasses must implement: - setup(): Create datasets for each stage

Parameters:

data_dir (str | Path) – Root directory containing the dataset.
dataset_format (DatasetFormat | str) – Format of the dataset annotations.
class_index_mode (ClassIndexMode | str) – How class indices should be handled.
class_names (list[str] | None) – List of class names (excluding background for YOLO, including all classes for TorchVision).
batch_size (int) – Batch size for DataLoaders.
num_workers (int) – Number of workers for DataLoaders.
pin_memory (bool) – Whether to pin memory in DataLoaders.
train_transforms (Callable | None) – Transform pipeline for training data.
val_transforms (Callable | None) – Transform pipeline for validation data.
test_transforms (Callable | None) – Transform pipeline for test data.

data_dir¶: Dataset root directory.

dataset_format¶: Dataset format enum.

class_index_mode¶: Class index handling mode.

class_names¶: List of class names.

batch_size¶: DataLoader batch size.

num_workers¶: DataLoader workers.

Example

>>> datamodule = COCODataModule(
...     data_dir="/path/to/coco",
...     class_index_mode=ClassIndexMode.TORCHVISION,
...     batch_size=16,
... )
>>> datamodule.setup("fit")
>>> train_loader = datamodule.train_dataloader()

abstractmethod setup(stage=None)[source]¶

Set up datasets for training, validation, testing, or prediction.

This method must be implemented by subclasses to create the appropriate datasets for each stage.

Parameters:: stage (str | None) – One of “fit”, “validate”, “test”, or “predict”. If None, set up all datasets.
Return type:: None

train_dataloader()[source]¶

Create training DataLoader.

Return type:: DataLoader
Returns:: DataLoader for training data with shuffling enabled.

val_dataloader()[source]¶

Create validation DataLoader.

Return type:: DataLoader
Returns:: DataLoader for validation data without shuffling.

test_dataloader()[source]¶

Create test DataLoader.

Return type:: DataLoader
Returns:: DataLoader for test data without shuffling.

predict_dataloader()[source]¶

Create prediction DataLoader.

Return type:: DataLoader
Returns:: DataLoader for prediction data without shuffling.

property num_classes: int¶

Get the number of classes in the dataset.

Returns:: Number of classes (not including background for YOLO mode).

get_class_name(class_id)[source]¶

Get class name by ID.

Parameters:: class_id (int) – Class index (accounting for class_index_mode).
Return type:: str
Returns:: Class name string.

get_dataset_info()[source]¶

Get dataset information dictionary.

Return type:: dict[str, Any]
Returns:: Dictionary with dataset metadata.

detection_collate_fn¶

objdet.data.base.detection_collate_fn(batch)[source]¶

Collate function for detection batches.

Unlike classification where we can stack images, detection requires lists because targets have variable numbers of boxes.

Parameters:: batch (list[tuple[Tensor, DetectionTarget]]) – List of (image, target) tuples.
Return type:: tuple[list[Tensor], list[DetectionTarget]]
Returns:: Tuple of (list of images, list of targets).

Data Preprocessing¶

convert_to_litdata¶

objdet.data.preprocessing.litdata_converter.convert_to_litdata(output_dir, format_name, input_dir=None, num_workers=4, class_names=None, splits=None, coco_ann_file=None, coco_images_dir=None, voc_images_dir=None, voc_annotations_dir=None, voc_imagesets_dir=None, yolo_images_dir=None, yolo_labels_dir=None)[source]¶

Convert dataset to LitData format.

This is the main entry point for the CLI preprocess command.

Either provide input_dir to use standard directory conventions, or provide the format-specific path arguments directly for custom dataset structures.

Parameters:

output_dir (str | Path) – Output directory for LitData format.
format_name (str) – Source format (“coco”, “voc”, “yolo”).
input_dir (str | Path | None) – Source dataset directory (optional if format-specific paths provided).
num_workers (int) – Number of parallel workers.
class_names (list[str] | None) – Class names (required for YOLO).
splits (list[str] | None) – List of splits to convert (default: [“train”, “val”]).
coco_ann_file (str | Path | None) – COCO annotation file path.
coco_images_dir (str | Path | None) – COCO images directory.
voc_images_dir (str | Path | None) – VOC JPEGImages directory.
voc_annotations_dir (str | Path | None) – VOC Annotations directory.
voc_imagesets_dir (str | Path | None) – VOC ImageSets/Main directory.
yolo_images_dir (str | Path | None) – YOLO images directory.
yolo_labels_dir (str | Path | None) – YOLO labels directory.

Return type:

None

Example

>>> # Using standard directory structure
>>> convert_to_litdata(
...     input_dir="/data/coco",
...     output_dir="/data/coco_litdata",
...     format_name="coco",
... )

>>> # Using custom paths (input_dir not needed)
>>> convert_to_litdata(
...     output_dir="/data/litdata_output",
...     format_name="coco",
...     coco_ann_file="/data/my_coco/train_annotations.json",
...     coco_images_dir="/data/my_coco/images",
...     splits=["train"],
... )

LitDataConverter¶

class objdet.data.preprocessing.litdata_converter.LitDataConverter(output_dir, dataset_format, input_dir=None, class_names=None, num_workers=4, coco_ann_file=None, coco_images_dir=None, voc_images_dir=None, voc_annotations_dir=None, voc_imagesets_dir=None, yolo_images_dir=None, yolo_labels_dir=None)[source]¶

Bases: object

Converter for creating LitData optimized datasets.

Either provide input_dir to use standard directory conventions, or provide the format-specific path arguments directly for custom dataset structures.

Parameters:

output_dir (str | Path) – Output directory for LitData format.
dataset_format (DatasetFormat | str) – Source dataset format (coco, voc, yolo).
input_dir (str | Path | None) – Source dataset directory (optional if format-specific paths provided).
class_names (list[str] | None) – List of class names (required for YOLO format).
num_workers (int) – Number of parallel workers.
coco_ann_file (str | Path | None) – COCO annotation file path.
coco_images_dir (str | Path | None) – COCO images directory.
voc_images_dir (str | Path | None) – VOC JPEGImages directory.
voc_annotations_dir (str | Path | None) – VOC Annotations directory.
voc_imagesets_dir (str | Path | None) – VOC ImageSets/Main directory.
yolo_images_dir (str | Path | None) – YOLO images directory.
yolo_labels_dir (str | Path | None) – YOLO labels directory.

Example

>>> # Using standard COCO directory structure
>>> converter = LitDataConverter(
...     input_dir="/data/coco",
...     output_dir="/data/coco_litdata",
...     dataset_format=DatasetFormat.COCO,
... )
>>> converter.convert()

>>> # Using custom paths (input_dir not needed)
>>> converter = LitDataConverter(
...     output_dir="/data/litdata_output",
...     dataset_format=DatasetFormat.COCO,
...     coco_ann_file="/data/my_dataset/labels.json",
...     coco_images_dir="/data/my_dataset/photos",
... )

convert(split='train')[source]¶

Convert dataset to LitData format.

Parameters:: split (str) – Dataset split to convert.
Return type:: None