yamle.data.dataset_wrappers module

class yamle.data.dataset_wrappers.SurrogateDataset(dataset, transform=None, target_transform=None, joint_transform=None)

Bases: Dataset

This class is a dataset wrapper, ensuring that the transforms are applied to the data and targets after splitting the dataset into training and validation.

Parameters:

• dataset (Dataset) – Dataset to wrap.

• transform (Optional[Callable]) – Transformations to apply to the data.

• target_transform (Optional[Callable]) – Transformations to apply to the targets.

• joint_transform (Optional[Callable]) – Transformations to apply to the input as well as the target.

class yamle.data.dataset_wrappers.InputImagePaddingDataset(dataset, padding)

Bases: Dataset

This class is a dataset wrapper, which can pad the input image to a given size.

Parameters:

• dataset (Dataset) – Dataset to wrap.

• padding (int) – Padding to apply to the input image on all sides.

class yamle.data.dataset_wrappers.ImageRotationDataset(dataset, max_angle=90, min_angle=0, seed=42)

Bases: Dataset

This class is a dataset wrapper for image rotation.

It discards the target and replaces it with the rotation angle which should be predicted. This changes the task from anything to regression.

Parameters:

• dataset (Dataset) – Dataset to wrap.

• max_angle (float) – Maximum angle to rotate the image by. Defaults to 90 degrees.

• min_angle (float) – Minimum angle to rotate the image by. Defaults to 0 degrees.

• seed (int) – Seed for the random number generator. Defaults to 42.