Datasets

ImageBaseDataset

class oml.datasets.images.ImageBaseDataset(paths: List[Path], dataset_root: Optional[Union[Path, str]] = None, bboxes: Optional[Sequence[Optional[Tuple[int, int, int, int]]]] = None, extra_data: Optional[Dict[str, Any]] = None, transform: Optional[Union[Compose, Compose]] = None, f_imread: Optional[Callable[[Union[Path, str, bytes]], Union[Image, ndarray]]] = None, cache_size: Optional[int] = 0, input_tensors_key: str = 'input_tensors', index_key: str = 'idx', paths_key: str = 'paths', x1_key: str = 'x1', x2_key: str = 'x2', y1_key: str = 'y1', y2_key: str = 'y2')[source]

Bases: IBaseDataset, IVisualizableDataset

The base class that handles image specific logic.

__init__(paths: List[Path], dataset_root: Optional[Union[Path, str]] = None, bboxes: Optional[Sequence[Optional[Tuple[int, int, int, int]]]] = None, extra_data: Optional[Dict[str, Any]] = None, transform: Optional[Union[Compose, Compose]] = None, f_imread: Optional[Callable[[Union[Path, str, bytes]], Union[Image, ndarray]]] = None, cache_size: Optional[int] = 0, input_tensors_key: str = 'input_tensors', index_key: str = 'idx', paths_key: str = 'paths', x1_key: str = 'x1', x2_key: str = 'x2', y1_key: str = 'y1', y2_key: str = 'y2')[source]
Parameters

  • paths – Paths to images. Will be concatenated with dataset_root if provided.

  • dataset_root – Path to the images’ dir, set None if you provided the absolute paths in your dataframe

  • bboxes – Bounding boxes of images. Some of the images may not have bounding bboxes.

  • extra_data – Dictionary containing records of some additional information.

  • transform – Augmentations for the images, set None to perform only normalisation and casting to tensor

  • f_imread – Function to read the images, pass None to pick it automatically based on provided transforms

  • cache_size – Size of the dataset’s cache

  • input_tensors_key – Key to put tensors into the batches

  • index_key – Key to put samples’ ids into the batches

  • paths_key – Key put paths into the batches # todo 522: remove

  • x1_key – Key to put x1 into the batches # todo 522: remove

  • x2_key – Key to put x2 into the batches # todo 522: remove

  • y1_key – Key to put y1 into the batches # todo 522: remove

  • y2_key – Key to put y2 into the batches # todo 522: remove

__getitem__(item: int) → Dict[str, Union[FloatTensor, int]][source]
Parameters

item – Idx of the sample

Returns

self.input_tensors_key self.index_key: int = item

Return type

Dictionary including the following keys

visualize(item: int, color: Tuple[int, int, int] = (0, 0, 0)) → ndarray[source]

ImageLabeledDataset

class oml.datasets.images.ImageLabeledDataset(df: DataFrame, extra_data: Optional[Dict[str, Any]] = None, dataset_root: Optional[Union[Path, str]] = None, transform: Optional[Compose] = None, f_imread: Optional[Callable[[Union[Path, str, bytes]], Union[Image, ndarray]]] = None, cache_size: Optional[int] = 0, input_tensors_key: str = 'input_tensors', labels_key: str = 'labels', index_key: str = 'idx', paths_key: str = 'paths', categories_key: Optional[str] = 'categories', sequence_key: Optional[str] = 'sequence', x1_key: str = 'x1', x2_key: str = 'x2', y1_key: str = 'y1', y2_key: str = 'y2')[source]

Bases: ImageBaseDataset, ILabeledDataset

The dataset of images having their ground truth labels.

__init__(df: DataFrame, extra_data: Optional[Dict[str, Any]] = None, dataset_root: Optional[Union[Path, str]] = None, transform: Optional[Compose] = None, f_imread: Optional[Callable[[Union[Path, str, bytes]], Union[Image, ndarray]]] = None, cache_size: Optional[int] = 0, input_tensors_key: str = 'input_tensors', labels_key: str = 'labels', index_key: str = 'idx', paths_key: str = 'paths', categories_key: Optional[str] = 'categories', sequence_key: Optional[str] = 'sequence', x1_key: str = 'x1', x2_key: str = 'x2', y1_key: str = 'y1', y2_key: str = 'y2')[source]
Parameters

  • paths – Paths to images. Will be concatenated with dataset_root if provided.

  • dataset_root – Path to the images’ dir, set None if you provided the absolute paths in your dataframe

  • bboxes – Bounding boxes of images. Some of the images may not have bounding bboxes.

  • extra_data – Dictionary containing records of some additional information.

  • transform – Augmentations for the images, set None to perform only normalisation and casting to tensor

  • f_imread – Function to read the images, pass None to pick it automatically based on provided transforms

  • cache_size – Size of the dataset’s cache

  • input_tensors_key – Key to put tensors into the batches

  • index_key – Key to put samples’ ids into the batches

  • paths_key – Key put paths into the batches # todo 522: remove

  • x1_key – Key to put x1 into the batches # todo 522: remove

  • x2_key – Key to put x2 into the batches # todo 522: remove

  • y1_key – Key to put y1 into the batches # todo 522: remove

  • y2_key – Key to put y2 into the batches # todo 522: remove

__getitem__(item: int) → Dict[str, Any][source]
Parameters

item – Idx of the sample

Return type

Dictionary including the following keys

self.labels_key

get_labels() → ndarray[source]
visualize(item: int, color: Tuple[int, int, int] = (0, 0, 0)) → ndarray

ImageQueryGalleryLabeledDataset

class oml.datasets.images.ImageQueryGalleryLabeledDataset(df: DataFrame, extra_data: Optional[Dict[str, Any]] = None, dataset_root: Optional[Union[Path, str]] = None, transform: Optional[Compose] = None, f_imread: Optional[Callable[[Union[Path, str, bytes]], Union[Image, ndarray]]] = None, cache_size: Optional[int] = 0, input_tensors_key: str = 'input_tensors', labels_key: str = 'labels', paths_key: str = 'paths', categories_key: Optional[str] = 'categories', sequence_key: Optional[str] = 'sequence', x1_key: str = 'x1', x2_key: str = 'x2', y1_key: str = 'y1', y2_key: str = 'y2', is_query_key: str = 'is_query', is_gallery_key: str = 'is_gallery')[source]

Bases: ImageLabeledDataset, IQueryGalleryLabeledDataset

The annotated dataset of images having query/gallery split.

Note, that some datasets used as benchmarks in Metric Learning explicitly provide the splitting information (for example, DeepFashion InShop dataset), but some of them don’t (for example, CARS196 or CUB200). The validation idea for the latter is to perform 1 vs rest validation, where every query is evaluated versus the whole validation dataset (except for this exact query).

So, if you want an item participate in validation as both: query and gallery, you should mark this item as is_query == True and is_gallery == True, as it’s done in the CARS196 or CUB200 dataset.

__init__(df: DataFrame, extra_data: Optional[Dict[str, Any]] = None, dataset_root: Optional[Union[Path, str]] = None, transform: Optional[Compose] = None, f_imread: Optional[Callable[[Union[Path, str, bytes]], Union[Image, ndarray]]] = None, cache_size: Optional[int] = 0, input_tensors_key: str = 'input_tensors', labels_key: str = 'labels', paths_key: str = 'paths', categories_key: Optional[str] = 'categories', sequence_key: Optional[str] = 'sequence', x1_key: str = 'x1', x2_key: str = 'x2', y1_key: str = 'y1', y2_key: str = 'y2', is_query_key: str = 'is_query', is_gallery_key: str = 'is_gallery')[source]
Parameters

  • paths – Paths to images. Will be concatenated with dataset_root if provided.

  • dataset_root – Path to the images’ dir, set None if you provided the absolute paths in your dataframe

  • bboxes – Bounding boxes of images. Some of the images may not have bounding bboxes.

  • extra_data – Dictionary containing records of some additional information.

  • transform – Augmentations for the images, set None to perform only normalisation and casting to tensor

  • f_imread – Function to read the images, pass None to pick it automatically based on provided transforms

  • cache_size – Size of the dataset’s cache

  • input_tensors_key – Key to put tensors into the batches

  • index_key – Key to put samples’ ids into the batches

  • paths_key – Key put paths into the batches # todo 522: remove

  • x1_key – Key to put x1 into the batches # todo 522: remove

  • x2_key – Key to put x2 into the batches # todo 522: remove

  • y1_key – Key to put y1 into the batches # todo 522: remove

  • y2_key – Key to put y2 into the batches # todo 522: remove

__getitem__(idx: int) → Dict[str, Any][source]
Parameters

item – Idx of the sample

Return type

Dictionary including the following keys

self.labels_key

get_query_ids() → LongTensor[source]
get_labels() → ndarray
visualize(item: int, color: Tuple[int, int, int] = (0, 0, 0)) → ndarray

ImageQueryGalleryDataset

class oml.datasets.images.ImageQueryGalleryDataset(df: DataFrame, extra_data: Optional[Dict[str, Any]] = None, dataset_root: Optional[Union[Path, str]] = None, transform: Optional[Compose] = None, f_imread: Optional[Callable[[Union[Path, str, bytes]], Union[Image, ndarray]]] = None, cache_size: Optional[int] = 0, input_tensors_key: str = 'input_tensors', paths_key: str = 'paths', categories_key: Optional[str] = 'categories', sequence_key: Optional[str] = 'sequence', x1_key: str = 'x1', x2_key: str = 'x2', y1_key: str = 'y1', y2_key: str = 'y2', is_query_key: str = 'is_query', is_gallery_key: str = 'is_gallery')[source]

Bases: IVisualizableDataset, IQueryGalleryDataset

The NOT annotated dataset of images having query/gallery split.

__init__(df: DataFrame, extra_data: Optional[Dict[str, Any]] = None, dataset_root: Optional[Union[Path, str]] = None, transform: Optional[Compose] = None, f_imread: Optional[Callable[[Union[Path, str, bytes]], Union[Image, ndarray]]] = None, cache_size: Optional[int] = 0, input_tensors_key: str = 'input_tensors', paths_key: str = 'paths', categories_key: Optional[str] = 'categories', sequence_key: Optional[str] = 'sequence', x1_key: str = 'x1', x2_key: str = 'x2', y1_key: str = 'y1', y2_key: str = 'y2', is_query_key: str = 'is_query', is_gallery_key: str = 'is_gallery')[source]
__getitem__(item: int) → Dict[str, Any][source]
Parameters

item – Idx of the sample

Returns

self.input_tensors_key self.index_key: int = item

Return type

Dictionary including the following keys

get_query_ids() → LongTensor[source]
visualize(item: int, color: Tuple[int, int, int] = (0, 0, 0)) → ndarray[source]

EmbeddingPairsDataset

class oml.datasets.pairs.EmbeddingPairsDataset(embeddings1: Tensor, embeddings2: Tensor, pair_1st_key: str = 'input_tensors_1', pair_2nd_key: str = 'input_tensors_2', index_key: str = 'idx')[source]

Bases: IPairsDataset

Dataset to iterate over pairs of embeddings.

__init__(embeddings1: Tensor, embeddings2: Tensor, pair_1st_key: str = 'input_tensors_1', pair_2nd_key: str = 'input_tensors_2', index_key: str = 'idx')[source]
Parameters

  • embeddings1 – The first input embeddings

  • embeddings2 – The second input embeddings

  • pair_1st_key – Key to put embeddings1 into the batches

  • pair_2nd_key – Key to put embeddings2 into the batches

  • index_key – Key to put samples’ ids into the batches

__getitem__(idx: int) → Dict[str, Tensor][source]
Parameters

item – Idx of the sample

Return type

Dictionary with the following keys

self.pairs_1st_key self.pairs_2nd_key self.index_key

ImagePairsDataset

class oml.datasets.pairs.ImagePairsDataset(paths1: ~typing.List[~pathlib.Path], paths2: ~typing.List[~pathlib.Path], bboxes1: ~typing.Optional[~typing.Sequence[~typing.Optional[~typing.Tuple[int, int, int, int]]]] = None, bboxes2: ~typing.Optional[~typing.Sequence[~typing.Optional[~typing.Tuple[int, int, int, int]]]] = None, transform: ~typing.Optional[~typing.Union[~albumentations.core.composition.Compose, ~torchvision.transforms.transforms.Compose]] = None, f_imread: ~typing.Callable[[~typing.Union[~pathlib.Path, str, bytes]], ~typing.Union[~PIL.Image.Image, ~numpy.ndarray]] = <function imread_pillow>, pair_1st_key: str = 'input_tensors_1', pair_2nd_key: str = 'input_tensors_2', index_key: str = 'idx', cache_size: ~typing.Optional[int] = 0)[source]

Bases: IPairsDataset

Dataset to iterate over pairs of images.

__init__(paths1: ~typing.List[~pathlib.Path], paths2: ~typing.List[~pathlib.Path], bboxes1: ~typing.Optional[~typing.Sequence[~typing.Optional[~typing.Tuple[int, int, int, int]]]] = None, bboxes2: ~typing.Optional[~typing.Sequence[~typing.Optional[~typing.Tuple[int, int, int, int]]]] = None, transform: ~typing.Optional[~typing.Union[~albumentations.core.composition.Compose, ~torchvision.transforms.transforms.Compose]] = None, f_imread: ~typing.Callable[[~typing.Union[~pathlib.Path, str, bytes]], ~typing.Union[~PIL.Image.Image, ~numpy.ndarray]] = <function imread_pillow>, pair_1st_key: str = 'input_tensors_1', pair_2nd_key: str = 'input_tensors_2', index_key: str = 'idx', cache_size: ~typing.Optional[int] = 0)[source]
Parameters

  • paths1 – Paths to the 1st input images

  • paths2 – Paths to the 2nd input images

  • bboxes1 – Should be either None or a sequence of bboxes. If an image has N boxes, duplicate its path N times and provide bounding box for each of them. If you want to get an embedding for the whole image, set bbox to None for this particular image path. The format is x1, y1, x2, y2.

  • bboxes2 – The same as bboxes2, but for the second inputs.

  • transform – Augmentations for the images, set None to perform only normalisation and casting to tensor

  • f_imread – Function to read the images

  • pair_1st_key – Key to put the 1st images into the batches

  • pair_2nd_key – Key to put the 2nd images into the batches

  • index_key – Key to put samples’ ids into the batches

  • cache_size – Size of the dataset’s cache

__getitem__(idx: int) → Dict[str, Union[int, Dict[str, Any]]][source]
Parameters

item – Idx of the sample

Return type

Dictionary with the following keys

self.pairs_1st_key self.pairs_2nd_key self.index_key