Method _split_generators

tensorflow_datasets/object_detection/coco.py:231–277

Returns SplitGenerators.

(self, dl_manager)

    )

  def _split_generators(self, dl_manager):
    """Returns SplitGenerators."""

    # Merge urls from all splits together
    urls = {}
    for split in self.builder_config.splits:
      urls['{}_images'.format(split.name)] = 'zips/{}.zip'.format(split.images)
      urls['{}_annotations'.format(split.name)] = 'annotations/{}.zip'.format(
          split.annotations
      )

    # DownloadManager memoize the url, so duplicate urls will only be downloaded
    # once.
    root_url = 'http://images.cocodataset.org/'
    extracted_paths = dl_manager.download_and_extract(
        {key: root_url + url for key, url in urls.items()}
    )

    splits = []
    for split in self.builder_config.splits:
      image_dir = extracted_paths['{}_images'.format(split.name)]
      annotations_dir = extracted_paths['{}_annotations'.format(split.name)]
      if self.builder_config.has_panoptic:
        panoptic_image_zip_path = os.path.join(
            annotations_dir,
            'annotations',
            'panoptic_{}.zip'.format(split.images),
        )
        panoptic_dir = dl_manager.extract(panoptic_image_zip_path)
        panoptic_dir = os.path.join(
            panoptic_dir, 'panoptic_{}'.format(split.images)
        )
      else:
        panoptic_dir = None
      splits.append(
          tfds.core.SplitGenerator(
              name=split.name,
              gen_kwargs=dict(
                  image_dir=image_dir,
                  annotation_dir=annotations_dir,
                  split_name=split.images,
                  annotation_type=split.annotation_type,
                  panoptic_dir=panoptic_dir,
              ),
          )
      )
    return splits

  def _generate_examples(
      self, image_dir, annotation_dir, split_name, annotation_type, panoptic_dir