使用自定义数据集

本文档介绍了数据集 API (DatasetCatalogMetadataCatalog) 的工作原理,以及如何使用它们添加自定义数据集。

Detectron2 中具有内置支持的数据集列在 内置数据集 中。如果您想在重用 detectron2 的数据加载器时使用自定义数据集,您需要

  1. 注册您的数据集(即告诉 detectron2 如何获取您的数据集)。

  2. 可选地,注册数据集的元数据

接下来,我们将详细解释上述两个概念。

Colab 教程 提供了如何注册和训练自定义格式数据集的实时示例。

注册数据集

为了让 detectron2 知道如何获取名为“my_dataset”的数据集,用户需要实现一个函数来返回数据集中的项目,然后告诉 detectron2 关于此函数

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")

在这里,代码片段将名为“my_dataset”的数据集与一个返回数据的函数相关联。如果多次调用该函数,它必须返回相同的数据(顺序相同)。注册一直有效,直到进程退出。

该函数可以执行任意操作,并且应该以 list[dict] 格式返回数据,每个字典采用以下两种格式之一

  1. Detectron2 的标准数据集字典,如下所述。这将使其与 detectron2 中许多其他内置功能一起使用,因此建议在足够的情况下使用它。

  2. 任何自定义格式。您也可以返回您自己的格式的任意字典,例如为新任务添加额外的键。然后,您还需要在下游正确处理它们。有关更多详细信息,请参见下文。

标准数据集字典

对于标准任务(实例检测、实例/语义/全景分割、关键点检测),我们将原始数据集加载到 list[dict] 中,其规范类似于 COCO 的注释。这是我们对数据集的标准表示。

每个字典包含有关一张图像的信息。该字典可能包含以下字段,而必需字段会根据数据加载器或任务的需要而有所不同(请参见下文)。

任务

字段

常见

file_name、height、width、image_id

实例检测/分割

annotations

语义分割

sem_seg_file_name

全景分割

pan_seg_file_name、segments_info

  • file_name:图像文件的完整路径。

  • heightwidth:整数。图像的形状。

  • image_id(字符串或整数):标识此图像的唯一 ID。许多评估器需要此 ID 来识别图像,但数据集可能会将其用于其他目的。

  • annotations(list[dict]):实例检测/分割或关键点检测任务需要此字段。每个字典对应于此图像中一个实例的注释,并且可能包含以下键

    • bbox(list[float],必需):包含 4 个数字的列表,表示实例的边界框。

    • bbox_mode(int,必需):边界框的格式。它必须是 structures.BoxMode 的成员。目前支持:BoxMode.XYXY_ABSBoxMode.XYWH_ABS

    • category_id(int,必需):一个介于 [0, num_categories-1] 范围内的整数,表示类别标签。值 num_categories 保留用于表示“背景”类别(如果适用)。

    • segmentation(list[list[float]] 或 dict):实例的分割掩码。

      • 如果为 list[list[float]],它表示一组多边形列表,每个多边形代表对象的每个连通分量。每个 list[float] 是一个简单的多边形,格式为 [x1, y1, ..., xn, yn](n≥3)。X 和 Y 是以像素为单位的绝对坐标。

      • 如果为 dict,它表示 COCO 压缩 RLE 格式的逐像素分割掩码。该字典应包含键“size”和“counts”。您可以通过 pycocotools.mask.encode(np.asarray(mask, order="F")) 将 0 和 1 的 uint8 分割掩码转换为此类字典。如果使用具有此类格式的默认数据加载器,则 cfg.INPUT.MASK_FORMAT 必须设置为 bitmask

    • keypoints(list[float]):格式为 [x1, y1, v1,…, xn, yn, vn]。v[i] 表示此关键点的 可见性n 必须等于关键点类别的数量。X 和 Y 是范围在 [0, W 或 H] 内的绝对实数值坐标。

      (请注意,COCO 格式的关键点坐标是范围在 [0, W-1 或 H-1] 内的整数,这与我们的标准格式不同。Detectron2 将 0.5 添加到 COCO 关键点坐标,以将其从离散像素索引转换为浮点坐标。)

    • iscrowd:0(默认)或 1。此实例是否被标记为 COCO 的“人群区域”。如果您不知道它的含义,请勿包含此字段。

    如果 annotations 为空列表,则表示该图像被标记为没有对象。默认情况下,此类图像将从训练中删除,但可以使用 DATALOADER.FILTER_EMPTY_ANNOTATIONS 将其包含在内。

  • sem_seg_file_name(字符串):语义分割地面实况文件的完整路径。它应该是一张灰度图像,其像素值为整数标签。

  • pan_seg_file_name(字符串):全景分割地面实况文件的完整路径。它应该是一张 RGB 图像,其像素值为使用 panopticapi.utils.id2rgb 函数编码的整数 ID。这些 ID 由 segments_info 定义。如果某个 ID 未出现在 segments_info 中,则该像素被认为是未标记的,通常在训练和评估中被忽略。

  • segments_info(list[dict]):定义全景分割地面实况中每个 ID 的含义。每个字典都包含以下键

    • id(int):出现在地面实况图像中的整数。

    • category_id(int):一个介于 [0, num_categories-1] 范围内的整数,表示类别标签。

    • iscrowd:0(默认)或 1。此实例是否被标记为 COCO 的“人群区域”。

注意

PanopticFPN 模型不使用此处定义的全景分割格式,而是使用实例分割和语义分割数据格式的组合。有关 COCO 的说明,请参见 使用内置数据集

如今很少使用 Fast R-CNN(使用预先计算的建议)。要训练 Fast R-CNN,需要以下额外的键

  • proposal_boxes(数组):形状为 (K, 4) 的二维 numpy 数组,代表此图像的 K 个预先计算的建议框。

  • proposal_objectness_logits(数组):形状为 (K, ) 的 numpy 数组,对应于“proposal_boxes” 中建议的对象性 logits。

  • proposal_bbox_mode(int):预先计算的建议边界框的格式。它必须是 structures.BoxMode 的成员。默认值为 BoxMode.XYXY_ABS

用于新任务的自定义数据集字典

在您的数据集函数返回的 list[dict] 中,字典还可以包含任意自定义数据。这将有助于新任务,这些任务需要标准数据集字典中未涵盖的额外信息。在这种情况下,您需要确保下游代码可以正确处理您的数据。通常,这需要为数据加载器编写新的 mapper(请参见 使用自定义数据加载器)。

设计自定义格式时,请注意所有字典都存储在内存中(有时是序列化和多份副本)。为了节省内存,每个字典都旨在包含少量但足以描述每个样本的信息,例如文件名和注释。通常在数据加载器中加载完整样本。

对于整个数据集共有的属性,请使用 Metadata(请参见下文)。为了避免额外的内存占用,请不要在每个样本中保存此类信息。

数据集的“元数据”

每个数据集都与一些元数据相关联,可以通过 MetadataCatalog.get(dataset_name).some_metadata 访问。元数据是包含整个数据集共享的信息的键值映射,通常用于解释数据集中的内容,例如类的名称、类的颜色、文件的根目录等。此信息对于增强、评估、可视化、日志记录等很有用。元数据的结构取决于相应下游代码的需要。

如果您通过 DatasetCatalog.register 注册新数据集,您可能还想通过 MetadataCatalog.get(dataset_name).some_key = some_value 添加其相应的元数据,以启用任何需要元数据的功能。您可以像这样操作(以元数据键“thing_classes”为例)

from detectron2.data import MetadataCatalog
MetadataCatalog.get("my_dataset").thing_classes = ["person", "dog"]

以下是 Detectron2 中内置功能使用的元数据键列表。如果您添加自己的数据集但没有这些元数据,则某些功能可能不可用。

  • thing_classes (list[str]): 用于所有实例检测/分割任务。每个实例/事物类别的名称列表。如果您加载 COCO 格式的数据集,它将由函数 load_coco_json 自动设置。

  • thing_colors (list[tuple(r, g, b)]): 每个事物类别的预定义颜色(在 [0, 255] 中)。用于可视化。如果没有给出,将使用随机颜色。

  • stuff_classes (list[str]): 用于语义和全景分割任务。每个物质类别的名称列表。

  • stuff_colors (list[tuple(r, g, b)]): 每个物质类别的预定义颜色(在 [0, 255] 中)。用于可视化。如果没有给出,将使用随机颜色。

  • ignore_label (int): 用于语义和全景分割任务。在评估中应忽略具有此类别标签的真实标注中的像素。通常这些是“未标记”的像素。

  • keypoint_names (list[str]): 用于关键点检测。每个关键点的名称列表。

  • keypoint_flip_map (list[tuple[str]]): 用于关键点检测。名称对的列表,其中每对都是两个关键点,如果图像在增强期间水平翻转,则应翻转这两个关键点。

  • keypoint_connection_rules: list[tuple(str, str, (r, g, b))]。每个元组指定一对连接的关键点以及在可视化时用于它们之间线的颜色(在 [0, 255] 中)。

一些特定于某些数据集(例如 COCO)评估的附加元数据

  • thing_dataset_id_to_contiguous_id (dict[int->int]): 用于 COCO 格式中所有实例检测/分割任务。从数据集中实例类 ID 到 [0, #class) 范围内的连续 ID 的映射。将由函数 load_coco_json 自动设置。

  • stuff_dataset_id_to_contiguous_id (dict[int->int]): 用于生成语义/全景分割的预测 json 文件。从数据集中语义分割类 ID 到 [0, num_categories) 中连续 ID 的映射。它仅对评估有用。

  • json_file: COCO 标注 json 文件。用于 COCO 格式数据集的 COCO 评估。

  • panoptic_root, panoptic_json: 用于 COCO 格式的全景评估。

  • evaluator_type: 用于内置主训练脚本选择评估器。不要在新的训练脚本中使用它。您只需在主脚本中直接提供 DatasetEvaluator 即可。

注意

在识别中,有时我们使用“事物”一词来表示实例级任务,而使用“物质”来表示语义分割任务。两者都用于全景分割任务。有关“事物”和“物质”概念的背景,请参阅 On Seeing Stuff: The Perception of Materials by Humans and Machines.

注册 COCO 格式数据集

如果您的实例级(检测、分割、关键点)数据集已经是 COCO 格式的 json 文件,则可以使用以下方法轻松注册数据集及其关联的元数据。

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

如果您的数据集是 COCO 格式,但需要进一步处理,或者具有额外的自定义每个实例注释,则函数 load_coco_json 可能会有用。

更新新数据集的配置

注册数据集后,您可以在 cfg.DATASETS.{TRAIN,TEST} 中使用数据集的名称(例如,上面的示例中的“my_dataset”)。您可能需要更改其他配置以在新数据集上进行训练或评估。

  • MODEL.ROI_HEADS.NUM_CLASSESMODEL.RETINANET.NUM_CLASSES 分别是 R-CNN 和 RetinaNet 模型的事物类数。

  • MODEL.ROI_KEYPOINT_HEAD.NUM_KEYPOINTS 设置关键点 R-CNN 的关键点数。您还需要使用 TEST.KEYPOINT_OKS_SIGMAS 设置 Keypoint OKS 以进行评估。

  • MODEL.SEM_SEG_HEAD.NUM_CLASSES 设置语义 FPN 和全景 FPN 的物质类数。

  • TEST.DETECTIONS_PER_IMAGE 控制要检测的对象的最大数量。如果测试图像可能包含 >100 个对象,请将其设置为更大的数字。

  • 如果您正在训练 Fast R-CNN(使用预先计算的 proposals),则 DATASETS.PROPOSAL_FILES_{TRAIN,TEST} 需要与数据集匹配。proposal 文件的格式记录 here.

新模型(例如 TensorMaskPointRend)通常有类似的自身配置,也需要更改。

提示

更改类数后,预训练模型中的某些层将变得不兼容,因此无法加载到新模型中。这是预期的,加载此类预训练模型将产生有关此类层的警告。