Use Custom Datasets — detectron2 0.6 documentation (original) (raw)

detectron2

This document explains how the dataset APIs (DatasetCatalog, MetadataCatalog) work, and how to use them to add custom datasets.

Datasets that have builtin support in detectron2 are listed in builtin datasets. If you want to use a custom dataset while also reusing detectron2’s data loaders, you will need to:

  1. Register your dataset (i.e., tell detectron2 how to obtain your dataset).
  2. Optionally, register metadata for your dataset.

Next, we explain the above two concepts in detail.

The Colab tutorialhas a live example of how to register and train on a dataset of custom formats.

Register a Dataset

To let detectron2 know how to obtain a dataset named “my_dataset”, users need to implement a function that returns the items in your dataset and then tell detectron2 about this function:

def my_dataset_function(): ... return list[dict] in the following format

from detectron2.data import DatasetCatalog DatasetCatalog.register("my_dataset", my_dataset_function)

later, to access the data:

data: List[Dict] = DatasetCatalog.get("my_dataset")

Here, the snippet associates a dataset named “my_dataset” with a function that returns the data. The function must return the same data (with same order) if called multiple times. The registration stays effective until the process exits.

The function can do arbitrary things and should return the data in list[dict], each dict in either of the following formats:

  1. Detectron2’s standard dataset dict, described below. This will make it work with many other builtin features in detectron2, so it’s recommended to use it when it’s sufficient.
  2. Any custom format. You can also return arbitrary dicts in your own format, such as adding extra keys for new tasks. Then you will need to handle them properly downstream as well. See below for more details.

Standard Dataset Dicts

For standard tasks (instance detection, instance/semantic/panoptic segmentation, keypoint detection), we load the original dataset into list[dict] with a specification similar to COCO’s annotations. This is our standard representation for a dataset.

Each dict contains information about one image. The dict may have the following fields, and the required fields vary based on what the dataloader or the task needs (see more below).

Task Fields
Common file_name, height, width, image_id
Instance detection/segmentation annotations
Semantic segmentation sem_seg_file_name
Panoptic segmentation pan_seg_file_name, segments_info

Note

The PanopticFPN model does not use the panoptic segmentation format defined here, but a combination of both instance segmentation and semantic segmentation data format. See Use Builtin Datasets for instructions on COCO.

Fast R-CNN (with pre-computed proposals) models are rarely used today. To train a Fast R-CNN, the following extra keys are needed:

Custom Dataset Dicts for New Tasks

In the list[dict] that your dataset function returns, the dictionary can also have arbitrary custom data. This will be useful for a new task that needs extra information not covered by the standard dataset dicts. In this case, you need to make sure the downstream code can handle your data correctly. Usually this requires writing a new mapper for the dataloader (see Use Custom Dataloaders).

When designing a custom format, note that all dicts are stored in memory (sometimes serialized and with multiple copies). To save memory, each dict is meant to contain small but sufficient information about each sample, such as file names and annotations. Loading full samples typically happens in the data loader.

For attributes shared among the entire dataset, use Metadata (see below). To avoid extra memory, do not save such information inside each sample.

Register a COCO Format Dataset

If your instance-level (detection, segmentation, keypoint) dataset is already a json file in the COCO format, the dataset and its associated metadata can be registered easily with:

from detectron2.data.datasets import register_coco_instances register_coco_instances("my_dataset", {}, "json_annotation.json", "path/to/image/dir")

If your dataset is in COCO format but need to be further processed, or has extra custom per-instance annotations, the load_coco_jsonfunction might be useful.

Update the Config for New Datasets

Once you’ve registered the dataset, you can use the name of the dataset (e.g., “my_dataset” in example above) in cfg.DATASETS.{TRAIN,TEST}. There are other configs you might want to change to train or evaluate on new datasets:

New models (e.g. TensorMask,PointRend) often have similar configs of their own that need to be changed as well.

Tip

After changing the number of classes, certain layers in a pre-trained model will become incompatible and therefore cannot be loaded to the new model. This is expected, and loading such pre-trained models will produce warnings about such layers.