wrap_dataset_for_transforms_v2

torchaug.data.dataset.wrap_dataset_for_transforms_v2(dataset, target_keys=None)[source]

Wrap a torchvision.dataset for usage with torchaug.transforms.

Example

>>> dataset = torchvision.datasets.CocoDetection(...)
>>> dataset = wrap_dataset_for_transforms_v2(dataset)

Note

For now, only the most popular datasets are supported. Furthermore, the wrapper only supports dataset configurations that are fully supported by torchaug.transforms. If you encounter an error prompting you to raise an issue to torchvision for a dataset or configuration that you need, please do so.

The dataset samples are wrapped according to the description below.

Special cases:

  • torchvision.datasets.CocoDetection: Instead of returning the target as list of dicts, the wrapper returns a dict of lists. In addition, the key-value-pairs "boxes" (in XYXY coordinate format), "masks" and "labels" are added and wrap the data in the corresponding torchaug.ta_tensors. The original keys are preserved. If target_keys is omitted, returns only the values for the "image_id", "boxes", and "labels".

  • torchvision.datasets.VOCDetection: The key-value-pairs "boxes" and "labels" are added to the target and wrap the data in the corresponding torchaug.ta_tensors. The original keys are preserved. If target_keys is omitted, returns only the values for the "boxes" and "labels".

  • torchvision.datasets.CelebA: The target for target_type="bbox" is converted to the XYXY coordinate format and wrapped into a BoundingBoxes ta_tensors.

  • torchvision.datasets.Kitti: Instead returning the target as list of dicts, the wrapper returns a dict of lists. In addition, the key-value-pairs "boxes" and "labels" are added and wrap the data in the corresponding torchaug.ta_tensors. The original keys are preserved. If target_keys is omitted, returns only the values for the "boxes" and "labels".

  • torchvision.datasets.OxfordIIITPet: The target for target_type="segmentation" is wrapped into a torchaug.ta_tensors._mask.Mask ta_tensors.

  • torchvision.datasets.Cityscapes: The target for target_type="semantic" is wrapped into a torchaug.ta_tensors._mask.Mask ta_tensors. The target for target_type="instance" is replaced by a dictionary with the key-value-pairs "masks" (as torchaug.ta_tensors._mask.Mask ta_tensors) and "labels".

  • torchvision.datasets.WIDERFace: The value for key "bbox" in the target is converted to XYXY coordinate format and wrapped into a BoundingBoxes ta_tensors.

Image classification datasets

This wrapper is a no-op for image classification datasets, since they were already fully supported by torchaug.transforms and thus no change is needed for torchaug.transforms.

Segmentation datasets

Segmentation datasets, e.g. torchvision.datasets.VOCSegmentation, return a two-tuple of PIL.Image.Image’s. This wrapper leaves the image as is (first item), while wrapping the segmentation mask into a torchaug.ta_tensors._mask.Mask (second item).

Video classification datasets

Video classification datasets, e.g. torchvision.datasets.Kinetics, return a three-tuple containing a torch.Tensor for the video and audio and a int as label. This wrapper wraps the video into a Video while leaving the other items as is.

Note

Only datasets constructed with output_format="TCHW" are supported, since the alternative output_format="THWC" is not supported by torchaug.transforms.

Parameters: