datasets.enrico package

Submodules

datasets.enrico.get_data module

Implements dataloaders for ENRICO dataset.

class datasets.enrico.get_data.EnricoDataset(data_dir, mode='train', noise_level=0, img_noise=False, wireframe_noise=False, img_dim_x=128, img_dim_y=256, random_seed=42, train_split=0.65, val_split=0.15, test_split=0.2, normalize_image=False, seq_len=64)

Bases: Dataset

Implements torch dataset class for ENRICO dataset.

__init__(data_dir, mode='train', noise_level=0, img_noise=False, wireframe_noise=False, img_dim_x=128, img_dim_y=256, random_seed=42, train_split=0.65, val_split=0.15, test_split=0.2, normalize_image=False, seq_len=64)

Instantiate ENRICO dataset.

Parameters:

  • data_dir (str) – Data directory.

  • mode (str, optional) – What data to extract. Defaults to “train”.

  • noise_level (int, optional) – Noise level, as defined in robustness. Defaults to 0.

  • img_noise (bool, optional) – Whether to apply noise to images or not. Defaults to False.

  • wireframe_noise (bool, optional) – Whether to apply noise to wireframes or not. Defaults to False.

  • img_dim_x (int, optional) – Image width. Defaults to 128.

  • img_dim_y (int, optional) – Image height. Defaults to 256.

  • random_seed (int, optional) – Seed to split dataset on and shuffle data on. Defaults to 42.

  • train_split (float, optional) – Percentage of training data split. Defaults to 0.65.

  • val_split (float, optional) – Percentage of validation data split. Defaults to 0.15.

  • test_split (float, optional) – Percentage of test data split. Defaults to 0.2.

  • normalize_image (bool, optional) – Whether to normalize image or not Defaults to False.

  • seq_len (int, optional) – Length of sequence. Defaults to 64.

featurizeElement(element)

Convert element into tuple of (bounds, one-hot-label).

datasets.enrico.get_data.get_dataloader(data_dir, batch_size=32, num_workers=0, train_shuffle=True, return_class_weights=True)

Get dataloaders for this dataset.

Parameters:

  • data_dir (str) – Data directory.

  • batch_size (int, optional) – Batch size. Defaults to 32.

  • num_workers (int, optional) – Number of workers. Defaults to 0.

  • train_shuffle (bool, optional) – Whether to shuffle training data or not. Defaults to True.

  • return_class_weights (bool, optional) – Whether to return class weights or not. Defaults to True.

Returns:

Tuple of ((train dataloader, validation dataloader, test dataloader), class_weights) if return_class_weights, otherwise just the dataloaders

Return type:

tuple

Module contents