CityscapesDataModule#
- class torch_uncertainty.datamodules.segmentation.CityscapesDataModule(root, batch_size, eval_batch_size=None, crop_size=1024, eval_size=(1024, 2048), train_transform=None, test_transform=None, basic_augment=True, val_split=None, num_workers=1, pin_memory=True, persistent_workers=True)[source]#
DataModule for the Cityscapes dataset.
- Parameters:
root (str or Path) – Root directory of the datasets.
batch_size (int) – Number of samples per batch during training.
eval_batch_size (int | None) – Number of samples per batch during evaluation (val and test). Set to
batch_size
ifNone
. Defaults toNone
.crop_size (sequence or int, optional) – Desired input image and segmentation mask sizes during training. If
crop_size
is an int instead of sequence like \((H, W)\), a square crop \((\text{size},\text{size})\) is made. If provided a sequence of length \(1\), it will be interpreted as \((\text{size[0]},\text{size[1]})\). Has to be provided iftrain_transform
is not provided. Otherwise has no effect. Defaults to1024
.eval_size (sequence or int, optional) – Desired input image and segmentation mask sizes during evaluation. If size is an int, smaller edge of the images will be matched to this number, i.e., \(\text{height}>\text{width}\), then image will be rescaled to \((\text{size}\times\text{height}/\text{width},\text{size})\). Has to be provided if
test_transform
is not provided. Otherwise has no effect. Defaults to(1024,2048)
.train_transform (nn.Module | None) – Custom training transform. Defaults to
None
. If not provided, a default transform is used.test_transform (nn.Module | None) – Custom test transform. Defaults to
None
. If not provided, a default transform is used.basic_augment (bool) – Whether to apply base augmentations. Defaults to
True
. Only used iftrain_transform
is not provided.val_split (float or None, optional) – Share of training samples to use for validation. Defaults to
None
.num_workers (int, optional) – Number of dataloaders to use. Defaults to
1
.pin_memory (bool, optional) – Whether to pin memory. Defaults to
True
.persistent_workers (bool, optional) – Whether to use persistent workers. Defaults to
True
.
Note
By default this datamodule injects the following transforms into the training and validation/test datasets:
Training transforms:
from torchvision.transforms import v2 v2.Compose( [ v2.ToImage(), RandomRescale(min_scale=0.5, max_scale=2.0, antialias=True), v2.RandomCrop(size=crop_size, pad_if_needed=True), v2.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5), v2.RandomHorizontalFlip(), v2.ToDtype( { tv_tensors.Image: torch.float32, tv_tensors.Mask: torch.int64, "others": None, }, scale=True, ), v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ] )
Validation/Test transforms:
from torchvision.transforms import v2 v2.Compose( [ v2.ToImage(), v2.Resize(size=eval_size, antialias=True), v2.ToDtype( { tv_tensors.Image: torch.float32, tv_tensors.Mask: torch.int64, "others": None, }, scale=True, ), v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ] )
This behavior can be modified by setting up
train_transform
andtest_transform
at initialization.