mediapipe_model_maker.object_detector.Dataset

Dataset library for object detector.

Inherits From: ClassificationDataset, Dataset

mediapipe_model_maker.object_detector.Dataset(
    dataset: tf.data.Dataset,
    label_names: List[str],
    size: Optional[int] = None
)

tf_dataset A tf.data.Dataset object that contains a potentially large set of elements, where each element is a pair of (input_data, target). The input_data means the raw input data, like an image, a text etc., while the target means the ground truth of the raw input data, e.g. the classification label of the image etc.
size The size of the dataset. tf.data.Dataset donesn't support a function to get the length directly since it's lazy-loaded and may be infinite.

label_names

num_classes

size Returns the size of the dataset.

Same functionality as calling len. See the len method definition for more information.

Methods

from_cache

View source

@classmethod
from_cache(
    tfrecord_cache_files: cache_files.TFRecordCacheFiles
) -> 'Dataset'

Loads the TFRecord data from cache.

Args
tfrecord_cache_files The TFRecordCacheFiles object containing the already cached TFRecord and metadata files.

Returns
ObjectDetectorDataset object.

Raises
ValueError if tfrecord_cache_files are not already cached.

from_coco_folder

View source

@classmethod
from_coco_folder(
    data_dir: str,
    max_num_images: Optional[int] = None,
    cache_dir: Optional[str] = None
) -> 'Dataset'

Loads images and labels from the given directory in COCO format.

Folder structure should be:

  <data_dir>/
    images/
      <file0>.jpg
      ...
    labels.json

The labels.json annotations file should should have the following format:

{
    "categories": [{"id": 0, "name": "background"}, ...],
    "images": [{"id": 0, "file_name": "<file0>.jpg"}, ...],
    "annotations": [{
       "id": 0,
       "image_id": 0,
       "category_id": 2,
       "bbox": [x-top left, y-top left, width, height],
       }, ...]
}

Note that category id 0 is reserved for the "background" class. It is optional to include, but if included it must be set to "background".

Args
data_dir Name of the directory containing the data files.
max_num_images Max number of images to process.
cache_dir The cache directory to save TFRecord and metadata files. The TFRecord files are a standardized format for training object detection while the metadata file is used to store information like dataset size and label mapping of id to label name. If the cache_dir is not set, a temporary folder will be created and will not be removed automatically after training which means it can be reused later.

Returns
Dataset containing images and labels and other related info.

Raises
ValueError If the input data directory is empty.
ValueError If the label_name for id 0 is set to something other than the 'background' class.

from_pascal_voc_folder

View source

@classmethod
from_pascal_voc_folder(
    data_dir: str,
    max_num_images: Optional[int] = None,
    cache_dir: Optional[str] = None
) -> 'Dataset'

Loads images and labels from the given directory in PASCAL VOC format.

Folder structure should be:

  <data_dir>/
    images/
      <file0>.jpg
      ...
    Annotations/
      <file0>.xml
      ...

Each .xml annotation file should have the following format:

  <annotation>
    <filename>file0.jpg</filename>
    <object>
      <name>kangaroo</name>
      <bndbox>
        <xmin>233</xmin>
        <ymin>89</ymin>
        <xmax>386</xmax>
        <ymax>262</ymax>
      </bndbox>
    </object>
    <object>...</object>
  </annotation>

Args
data_dir Name of the directory containing the data files.
max_num_images Max number of images to process.
cache_dir The cache directory to save TFRecord and metadata files. The TFRecord files are a standardized format for training object detection while the metadata file is used to store information like dataset size and label mapping of id to label name. If the cache_dir is not set, a temporary folder will be created and will not be removed automatically after training which means it can be reused later.

Returns
Dataset containing images and labels and other related info.

Raises
ValueError if the input data directory is empty.

gen_tf_dataset

View source

gen_tf_dataset(
    batch_size: int = 1,
    is_training: bool = False,
    shuffle: bool = False,
    preprocess: Optional[Callable[..., Any]] = None,
    drop_remainder: bool = False
) -> tf.data.Dataset

Generates a batched tf.data.Dataset for training/evaluation.

Args
batch_size An integer, the returned dataset will be batched by this size.
is_training A boolean, when True, the returned dataset will be optionally shuffled and repeated as an endless dataset.
shuffle A boolean, when True, the returned dataset will be shuffled to create randomness during model training.
preprocess A function taking three arguments in order, feature, label and boolean is_training.
drop_remainder boolean, whether the finally batch drops remainder.

Returns
A TF dataset ready to be consumed by Keras model.

split

View source

split(
    fraction: float
) -> Tuple[ds._DatasetT, ds._DatasetT]

Splits dataset into two sub-datasets with the given fraction.

Primarily used for splitting the data set into training and testing sets.

Args
fraction float, demonstrates the fraction of the first returned subdataset in the original data.

Returns
The splitted two sub datasets.

__len__

View source

__len__() -> int

Returns the number of element of the dataset.

If size is not set, this method will fallback to using the len method of the tf.data.Dataset in self._dataset. Calling len on a tf.data.Dataset instance may throw a TypeError because the dataset may be lazy-loaded with an unknown size or have infinite size.

In most cases, however, when an instance of this class is created by helper functions like 'from_folder', the size of the dataset will be preprocessed, and the _size instance variable will be already set.

Raises
TypeError if self._size is not set and the cardinality of self._dataset is INFINITE_CARDINALITY or UNKNOWN_CARDINALITY.