Segmentation Data Module

The data_module module provides a PyTorch Lightning LightningDataModule for semantic segmentation tasks.

class ucs.data.data_module.SegmentationDataModule(config=None, transform=None, **kwargs)[source]

Bases: LightningDataModule

A PyTorch Lightning DataModule for semantic segmentation tasks. This class handles dataset preparation, transformation, and creation of data loaders for training, validation, and testing.

Variables:

  • dataset_path (str) – Path to the dataset or dataset identifier, initialized from the config.

  • batch_size (int) – Batch size for the data loaders, initialized from the config.

  • num_workers (int) – Number of workers for data loading, initialized from the config.

  • do_reduce_labels (bool) – Whether to reduce label values, initialized from the config.

  • pin_memory (bool) – Whether to use pinned memory for faster data transfer, initialized from the config.

  • transform (callable, optional) – Transformations to apply to the dataset.

  • persistent_workers (bool) – Whether to use persistent workers in data loading.

  • feature_extractor (SegformerImageProcessor) – Pre-trained feature extractor initialized with the model name.

  • raw_dataset (Dataset or None) – The raw dataset loaded from the dataset source.

  • train_dataset (Dataset or None) – The processed training dataset.

  • val_dataset (Dataset or None) – The processed validation dataset.

  • test_dataset (Dataset or None) – The processed test dataset.

Parameters:

  • config (DatasetConfig, optional) – Configuration object containing dataset parameters.

  • transform (callable, optional) – Transformations to apply to the dataset.

  • **kwargs – Additional keyword arguments for overriding dataset configurations.