imaginaire.datasets package

Submodules

imaginaire.datasets.base module

All datasets are inherited from this class.

class imaginaire.datasets.base.BaseDataset(cfg, is_inference, is_test)[source]

Bases: torch.utils.data.dataset.Dataset

Base class for image/video datasets.

Parameters

  • cfg (Config object) – Input config.

  • is_inference (bool) – Training if False, else validation.

  • is_test (bool) – Final test set after training and validation.

apply_ops(data, op_dict, full_data=False)[source]

Apply any ops from op_dict to data types.

Parameters

  • data (dict) – Dict containing data_type as key, with each value as a list of numpy.ndarrays.

  • op_dict (dict) – Dict containing data_type as key, with each value containing string of operations to apply.

  • full_data (bool) – Do these ops require access to the full data?

Returns

Dict containing data_type as key, with each value modified by the op if any.

Return type

data (dict)

flip_hdr(data, is_flipped=False)[source]

Flip hdr images.

Parameters

  • data (dict) – Keys are from data types. Values can be numpy.ndarray or list of numpy.ndarray (image or list of images).

  • is_flipped (bool) – Applying left-right flip to the hdr images

Returns

  • data (dict): Augmented data, with same keys as input data.

Return type

(tuple)

get_op(op)[source]

Get function to apply for specific op.

Parameters

op (str) – Name of the op.

Returns

function handle.

perform_augmentation(data, paired, augment_ops=None)[source]

Perform data augmentation on images only.

Parameters

  • data (dict) – Keys are from data types. Values can be numpy.ndarray or list of numpy.ndarray (image or list of images).

  • paired (bool) – Apply same augmentation to all input keys?

  • augment_ops (list) – The augmentation operations.

Returns

  • data (dict): Augmented data, with same keys as input data.

  • is_flipped (bool): Flag which tells if images have been left-right flipped.

Return type

(tuple)

perform_individual_video_frame(data, augment_ops)[source]

Perform data augmentation on images only.

Parameters

  • data (dict) – Keys are from data types. Values can be numpy.ndarray or list of numpy.ndarray (image or list of images).

  • augment_ops (list) – The augmentation operations for individual frames.

Returns

  • data (dict): Augmented data, with same keys as input data.

  • is_flipped (bool): Flag which tells if images have been left-right flipped.

Return type

(tuple)

to_tensor(data)[source]

Convert all images to tensor.

Parameters

data (dict) – Dict containing data_type as key, with each value as a list of numpy.ndarrays.

Returns

Dict containing data_type as key, with each value as a list of torch.Tensors.

Return type

data (dict)

imaginaire.datasets.cache module

class imaginaire.datasets.cache.Cache(root, size_GB)[source]

Bases: object

This creates an on disk cache, which saves files as bytes. :param root: Path to the cache dir. :type root: str :param size_MB: Size of cache in MB. :type size_MB: float

read(key)[source]
write(key, value)[source]

imaginaire.datasets.dummy module

class imaginaire.datasets.dummy.Dataset(cfg, is_inference=False, is_test=False)[source]

Bases: torch.utils.data.dataset.Dataset

Dummy dataset, returns nothing.

imaginaire.datasets.folder module

class imaginaire.datasets.folder.FolderDataset(root, metadata)[source]

Bases: torch.utils.data.dataset.Dataset

This deals with opening, and reading from an Folder dataset.

Parameters

  • root (str) – Path to the folder.

  • metadata (dict) – Containing extensions.

getitem_by_path(path, data_type)[source]

Load data item stored for key = path.

Parameters

  • path (str) – Key into Folder dataset.

  • data_type (str) – Key into self.extensions e.g. data/data_segmaps/…

Returns

Contents of file for this key.

Return type

img (PIL.Image) or buf (str)

imaginaire.datasets.images module

class imaginaire.datasets.images.Dataset(cfg, is_inference=False, is_test=False)[source]

Bases: imaginaire.datasets.base.BaseDataset

Image dataset for use in class conditional GAN.

Parameters

  • cfg (Config) – Loaded config object.

  • is_inference (bool) – In train or inference mode?

set_sample_class_idx(class_idx)[source]

Set sample class idx. This is not used in this class…

Parameters

class_idx (int) – Which class idx to sample from.

imaginaire.datasets.lmdb module

class imaginaire.datasets.lmdb.LMDBDataset(root)[source]

Bases: torch.utils.data.dataset.Dataset

This deals with opening, and reading from an LMDB dataset. :param root: Path to the LMDB file. :type root: str

getitem_by_path(path, data_type)[source]

Load data item stored for key = path.

Parameters

  • path (str) – Key into LMDB dataset.

  • data_type (str) – Key into self.extensions e.g. data/data_segmaps/…

Returns

Contents of LMDB value for this key.

Return type

img (PIL.Image) or buf (str)

imaginaire.datasets.object_store module

class imaginaire.datasets.object_store.ObjectStoreDataset(root, aws_credentials_file, data_type='', cache=None)[source]

Bases: torch.utils.data.dataset.Dataset

This deals with opening, and reading from an AWS S3 bucket. :param root: Path to the AWS S3 bucket. :type root: str :param aws_credentials_file: Path to file containing AWS credentials. :type aws_credentials_file: str :param data_type: Which data type should this dataset load? :type data_type: str

getitem_by_path(path, data_type)[source]

Load data item stored for key = path.

Parameters

  • path (str) – Path into AWS S3 bucket, without data_type prefix.

  • data_type (str) – Key into self.extensions e.g. data/data_segmaps/…

Returns

Contents of LMDB value for this key.

Return type

img (PIL.Image) or buf (str)

imaginaire.datasets.paired_few_shot_videos module

class imaginaire.datasets.paired_few_shot_videos.Dataset(cfg, is_inference=False, sequence_length=None, few_shot_K=None, is_test=False)[source]

Bases: imaginaire.datasets.paired_videos.Dataset

Paired video dataset for use in few-shot vid2vid.

Parameters

  • cfg (Config) – Loaded config object.

  • is_inference (bool) – In train or inference mode?

  • sequence_length (int) – What sequence of images to provide?

  • few_shot_K (int) – How many images to provide for few-shot?

set_inference_sequence_idx(index, k_shot_index, k_shot_frame_index)[source]

Get frames from this sequence during inference.

Parameters

  • index (int) – Index of inference sequence.

  • k_shot_index (int) – Index of sequence from which k_shot is sampled.

  • k_shot_frame_index (int) – Index of frame to sample.

set_sequence_length(sequence_length, few_shot_K=None)[source]

Set the length of sequence you want as output from dataloader.

Parameters

  • sequence_length (int) – Length of output sequences.

  • few_shot_K (int) – Number of few-shot frames.

imaginaire.datasets.paired_few_shot_videos_native module

class imaginaire.datasets.paired_few_shot_videos_native.Dataset(cfg, is_inference=False, is_test=False)[source]

Bases: imaginaire.datasets.base.BaseDataset

Dataset for paired few shot videos.

Parameters

  • cfg (Config) – Loaded config object.

  • is_inference (bool) – In train or inference mode?

get_label_lengths()[source]

Get num channels of all labels to be concated.

Returns

Dict mapping image data_type to num channels.

Return type

label_lengths (OrderedDict)

num_inference_sequences()[source]

Number of sequences available for inference.

Returns

(int)

imaginaire.datasets.paired_images module

class imaginaire.datasets.paired_images.Dataset(cfg, is_inference=False, is_test=False)[source]

Bases: imaginaire.datasets.paired_videos.Dataset

Paired image dataset for use in pix2pixHD, SPADE.

Parameters

  • cfg (Config) – Loaded config object.

  • is_inference (bool) – In train or inference mode?

num_inference_sequences()[source]

Number of sequences available for inference. Overriden from super as this is not applicable for images.

Returns

(int)

set_inference_sequence_idx(index)[source]

Get frames from this sequence during inference. Overriden from super as this is not applicable for images.

Parameters

index (int) – Index of inference sequence.

set_sequence_length(sequence_length)[source]

Set the length of sequence you want as output from dataloader. Ignore this as this is an image loader.

Parameters

sequence_length (int) – Length of output sequences.

imaginaire.datasets.paired_videos module

class imaginaire.datasets.paired_videos.Dataset(cfg, is_inference=False, sequence_length=None, is_test=False)[source]

Bases: imaginaire.datasets.base.BaseDataset

Paired video dataset for use in vid2vid, wc_vid2vid.

Parameters

  • cfg (Config) – Loaded config object.

  • is_inference (bool) – In train or inference mode?

  • sequence_length (int) – What sequence of images to provide?

get_label_lengths()[source]

Get num channels of all labels to be concated.

Returns

Dict mapping image data_type to num channels.

Return type

label_lengths (OrderedDict)

num_inference_sequences()[source]

Number of sequences available for inference.

Returns

(int)

set_inference_sequence_idx(index)[source]

Get frames from this sequence during inference.

Parameters

index (int) – Index of inference sequence.

set_sequence_length(sequence_length)[source]

Set the length of sequence you want as output from dataloader.

Parameters

sequence_length (int) – Length of output sequences.

imaginaire.datasets.unpaired_few_shot_images module

class imaginaire.datasets.unpaired_few_shot_images.Dataset(cfg, is_inference=False, is_test=False)[source]

Bases: imaginaire.datasets.base.BaseDataset

Image dataset for use in FUNIT.

Parameters

  • cfg (Config) – Loaded config object.

  • is_inference (bool) – In train or inference mode?

set_sample_class_idx(class_idx=None)[source]

Set sample class idx.

Parameters

class_idx (int) – Which class idx to sample from.

imaginaire.datasets.unpaired_images module

class imaginaire.datasets.unpaired_images.Dataset(cfg, is_inference=False, is_test=False)[source]

Bases: imaginaire.datasets.base.BaseDataset

Unpaired image dataset for use in MUNIT.

Parameters

  • cfg (Config) – Loaded config object.

  • is_inference (bool) – In train or inference mode?

Module contents