During the model training and testing, there will be a large amount of data passed through different components, and different algorithms usually have different kinds of data. For example, single-stage detectors may only need ground truth bounding boxes and ground truth box labels, whereas Mask R-CNN also requires the instance masks.
The training codes can be shown as:
```python
for img, img_metas, gt_bboxes, gt_labels in data_loader:
loss = retinanet(img, img_metas, gt_bboxes, gt_labels)
```
```python
for img, img_metas, gt_bboxes, gt_masks, gt_labels in data_loader:
loss = mask_rcnn(img, img_metas, gt_bboxes, gt_masks, gt_labels)
```
We can see that without encapsulation, the inconsistency of data required by different algorithms leads to the inconsistency of interfaces among different algorithm modules, which affects the extensibility of the whole algorithm library. Moreover, the modules within one algorithm library often need redundant interfaces in order to maintain compatibility.
These disadvantages are more obvious among different algorithm libraries, which makes it difficult to reuse modules and expand interfaces when implementing multi-task perception models (multiple tasks such as semantic segmentation, detection, key point detection, etc.).
To solve the above problems, MMEngine defines a set of abstract data interfaces to encapsulate various data during the implementation of the model. Suppose the above different data are encapsulated into `data_sample`, the training of different algorithms can be abstracted and unified into the following code:
```python
for img, data_sample in dataloader:
loss = model(img, data_sample)
```
The abstracted interface unifies and simplifies the interface between modules in the algorithm library, and can be used to pass data between datasets, models, visualizers, evaluates, or even within different modules in one model.
Besides the basic add, delete, update, and query functions, this interface also supports transferring data between different devices and the operation of `dict` and `torch.Tensor`, which can fully satisfy the requirements of the algorithm library.
Those algorithm libraries based on MMEngine can inherit from this design and implement their own interfaces to meet the characteristics and custom needs of data in different algorithms, improving the expandability while maintaining a unified interface.
During the implementation, there are two types of data interfaces for the algorithm libraries:
- A collection of all annotation information and prediction information for a training or testing sample, such as the output of a dataset, the inputs of model and visualizer, typically constitutes all the information of an individual training or testing sample. MMEngine defines this as a `DataSample`.
- A single type of prediction or annotation, typically the output of a sub-module in an algorithm model, such as the output of the RPN in two-stage detection, the output of a semantic segmentation model, the output of a keypoint branch, or the output of the generator in GANs, is defined by MMEngine as a data element (`XXXData`).
The following section first introduces the base class [BaseDataElement](mmengine.structures.BaseDataElement) for `DataSample` and `XXXData`.
## BaseDataElement
There are two types of data in `BaseDataElement`. One is `data` such as the bounding box, label, and the instance mask, etc., the other is `metainfo` which contains the meta information of the data to ensure the integrity of the data, including `img_shape`, `img_id`, and some other basic information of the images. These information facilitate the recovery and the use of the data in visualization and other cases. Therefore, users need to explicitly distinguish and declare the data of these two types of attributes while creating the `BaseDataElement`.
To make it easier to use `BaseDataElement`, the data in both `data` and `metainfo` are attributes of `BaseDataElement`. We can directly access the data and metainfo by accessing the class attributes. In addition, `BaseDataElement` provides several methods for manipulating the data in `data`.
- Add, delete, update, and query data in different fields of `data`.
- Copy `data` to target devices.
- Support accessing data in the same way as a dictionary or a tensor to fully satisfy the algorithm's requirements.
### 1. Create BaseDataElement
The data parameter of `BaseDataElement` can be freely added by means of `key=value`. The fields of `metainfo`, however, need to be explicitly specified using the keyword `metainfo`.
```python
import torch
from mmengine.structures import BaseDataElement
# declare an empty object
data_element = BaseDataElement()
bboxes = torch.rand((5, 4)) # suppose bboxes is a tensor in the shape of Nx4. N represents the number of the boxes
scores = torch.rand((5,)) # suppose scores is a tensor with N dimensions. N represents the number of the noxes.
img_id = 0 # image ID
H = 800 # image height
W = 1333 # image width
# Set the data parameter directly in BaseDataElement
# Explicitly declare the metainfo in BaseDataElement
data_element = BaseDataElement(
bboxes=bboxes,
scores=scores,
metainfo=dict(img_id=img_id, img_shape=(H, W)))
```
### 2. `new` and `clone`
Users can use the `new()` method to create an abstract data interface with the same state and data from an existing data interface. You can set `metainfo` and `data` while creating a new `BaseDataElement` to create an abstract interface with the same state and data as `data` or `metainfo`. For example, `new(metainfo=xx)` makes the new `BaseDataElement` has the same content as the cloned `BaseDataElement`, but `metainfo` is set to the newly specified content. You can also use `clone()` directly to get a deep copy. The behavior of the `clone()` is the same as the `clone()` in PyTorch Tensor operation.
```python
data_element = BaseDataElement(
bboxes=torch.rand((5, 4)),
scores=torch.rand((5,)),
metainfo=dict(img_id=1, img_shape=(640, 640)))
# set metainfo and data while creating BaseDataElement
print('bboxes is not in data_element2', 'bboxes' not in data_element2) # True
print('img_id in data_element2 is same as img_id in data_element', data_element2.img_id == data_element.img_id)
print('label in data_element2 is', 'label' in data_element2)
# create a new object using `clone`, which makes the new object has the same data, same metainfo, and the same status as the data_element
data_element2 = data_element1.clone()
```
```
bboxes is in data_element1: True
bboxes in data_element1 is same as bbox in data_element tensor(True)
img_id in data_element1 is True
bboxes is not in data_element2 True
img_id in data_element2 is same as img_id in data_element True
label in data_element2 is True
```
### 3. Add and query attributes
When it comes to adding attributes, users can add attributes to the `data` in the same way they add class attributes. For `metainfo`, it generally stores metadata about images and is not usually modified. If there is a need to add attributes to `metainfo`, users should use the `set_metainfo` interface to explicitly modify it.
For querying, users can access the key-value pairs that exist only in `data` using `keys`, `values`, and `items`. Similarly, they can access the key-value pairs that exist only in `metainfo` using `metainfo_keys`, `metainfo_values`, and `metainfo_items`. Users can also access all attributes of the BaseDataElement, regardless of their type, using `all_keys`, `all_values`, and `all_items`.
To facilitate usage, users can access the data within `data` and `metainfo` in the same way they access class attributes. Alternatively, they can use the `get()` interface in a dictionary-like manner to access the data.
**Note:**
1.`BaseDataElement` does not support having the same field names in both `metainfo` and `data` attributes. Therefore, users should avoid setting the same field names in them, as it would result in an error in `BaseDataElement`.
2. Considering that `InstanceData` and `PixelData` support slicing operations on the data, in order to maintain consistency with the use of `[]` and reduce the number of different methods for the same need, BaseDataElement does not support accessing and setting its attributes like a dictionary. Therefore, operations like `BaseDataElement[name]` for value assignment and retrieval are not supported.
```python
data_element = BaseDataElement()
# Set the `metainfo` field of the data_element using `set_metainfo`,
# with img_id and img_shape becoming attributes of the data_element.
Users can modify the `data` attribute of `BaseDataElement` in the same way they modify instance attributes. As for `metainfo`, it generally stores metadata about images and is not usually modified. If there is a need to modify `metainfo`, users should use the `set_metainfo` interface to make explicit modifications.
For convenience in operations, `data` and `metainfo` can be directly deleted using del. Additionally, the pop method is supported to delete attributes after accessing them.
print('The keys in metainfo is ', data_element.metainfo_keys())
```
```
(1280, 1280)
img_id: 10
img_shape: (1280, 1280)
img_id: 10
The keys in metainfo is []
```
### 5. Tensor-like operations
Users can transform the data status in `BaseDataElement` like the operations in tensor.Tensor. Currently, we support `cuda`, `cpu`, `to`, and `numpy`, etc. `to` has the same interface as `torch.Tensor.to()`, which allows users to change the status of the encapsulted tensor freely.
**Note:** These interfaces only handle sequences types in `np.array`, `torch.Tensor`, and numbers. Data in other types will be skipped, such as strings.
print('The type of bboxes in fp16_instances is', fp16_instances.bboxes.dtype) # torch.float16
# detach all data gradients
cuda_element_3 = cuda_element_2.detach()
print('The data in cuda_element_3 requires grad: ', cuda_element_3.bboxes.requires_grad)
# transform data to numpy array
np_instances = cpu_element_1.numpy()
print('The type of cpu_element_1 is convert to', type(np_instances.bboxes))
```
```
cuda_element_1 is on the device of cuda:0
cuda_element_1 is on the device of cuda:0
cpu_element_1 is on the device of cpu
cpu_element_2 is on the device of cpu
The type of bboxes in fp16_instances is torch.float16
The data in cuda_element_3 requires grad: False
The type of cpu_element_1 is convert to <class'numpy.ndarray'>
```
### 6. Show properties
`BaseDataElement` also implements `__repr__` which allows users to get all the data information through `print`. Meanwhile, to facilitate debugging, all attributes in `BaseDataElement` are added to `__dict__`. Users can visualize the contents directly in their IDEs. A complete property display is as follows:
MMEngine categorizes the data elements into three categories:
- InstanceData: mainly for high-level tasks that encapsulated all instance-related data in the image, such as bounding boxes, labels, instance masks, key points, polygons, tracking ids, etc. All instance-related data has the same **length**, which is the number of instances in the image.
- PixelData: mainly for low-level tasks and some high-level tasks that require pixel-level labels. It encapsulates pixel-level data such as segmentation map for semantic segmentations, flow map for optical flow tasks, panoptic segmentation map for panoramic segmentations, and various images generated by bottom-level tasks like super-resolution maps, denoising maps, and other various style maps generated. These data typically have three or four dimensions, with the last two dimensions representing the height and width of the data, which are consistent across the dataset.
- LabelData: mainly for encapsulating label-level data, such as class labels in image classification or multi-class classification, content categories for generated images in image generation, text in text recognition tasks, and more.
### InstanceData
[`InstanceData`](mmengine.structures.InstanceData) builds upon `BaseDataElement` and introduces restrictions on the data stored in `data`, requiring that the length of the data is consistent. For example, in object detection, assuming an image has N objects (instances), you can store all the bounding boxes and labels in InstanceData, where the lengths of bounding boxes and label in InstanceData are the same. Based on this assumption, InstanceData is extended to include the following features:
- length validation of the data stored in InstanceData's data.
- support for dictionary-like access and assignment of attributes in the `data`.
- support for basic indexing, slicing, and advanced indexing capabilities.
- support for concatenation of InstanceData with the same keys but different instances.
These extended features support basic data structures such as `torch.tensor`, `numpy.ndarray`, list, str, and tuple, as well as custom data structures, as long as the custom data structure implements `__len__`, `__getitem__`, and `cat` methods.
#### Data verification
All data stored in `InstanceData` must have the same length.
`InstanceData` supports the list indexing and slicing operations similar to Python, meanwhile, it also supports advanced indexing operations like numpy.
Users can concatenate two `InstanceData` with the same key into one new `InstanceData`. For two different `InstanceData` with different length as N and M, the length of the output `InstanceData` is N + M.
# AssertionError: The dim of value must be 2 or 3, but got 4
```
```
AssertionError: The dim of value must be 2 or 3, but got 4
```
#### Querying in spatial dimension
`PixelData` supports indexing and slicing in spatial dimension on part of the data instances. Users only need to pass in the index of the length and width.
print('The shape of pixel_data is: ', pixel_data.shape)
```
```
The shape of pixel_data is (20, 40)
```
- Indexing
```python
index_data = pixel_data[10, 20]
print('The shape of index_data is: ', index_data.shape)
```
```
The shape of index_data is (1, 1)
```
- Slicing
```python
slice_data = pixel_data[10:20, 20:40]
print('The shape of slice_data is: ', slice_data.shape)
```
```
The shape of slice_data is (10, 20)
```
### LabelData
[`LabelData`](mmengine.structures.LabelData) is mainly used to encapsulate label data such as classiciation labels, predicted text labels, etc. `LabelData` has no limitations to `data`, and it provides two extra features: `onehot` transformation and `index` transformation.
10 is convert to tensor([0, 1, 0, 0, 0, 0, 0, 0, 0, 0])
tensor([0, 1, 0, 0, 0, 0, 0, 0, 0, 0]) is convert to tensor([1])
```
## xxxDataSample
There may be different types of labels in one sample, for example, there may be both instance-level labels (Box) and pixel-level labels (SegMap) in one image. Therefore, we need to have a higher-level encapsulation on top of PixelData, InstanceData, and PixelData to represent the image-level labels. This layer is named `XXXDataSample` across the OpenMMLab series algorithms. In MMDet we have `DetDataSample`. All the labels are encapsulated in `XXXDataSample` during the training process, so different deep learning tasks can maintain a uniform data flow and data processing method.
### Downstream library usage
We take MMDet as an example to illustrate the use of the `DataSample` in downstream libraries and its constraints and naming styles. MMDet defined `DetDataSample` and seven fields, which are:
- Annotation Information
- gt_instance (InstanceData): Instance annotation information includes the instance class, bounding box, etc. The type constraint is `InstanceData`.
- gt_panoptic_seg (PixelData): For panoptic segmentation annotation information, the required type is PixelData.
- gt_semantic_seg (PixelData): Semantic segmentation annotation information. The type constraint is `PixelData`.
- Prediction Results
- pred_instance (InstanceData): Instance prediction results include the instance class, bounding boxes, etc. The type constraint is `InstanceData`.
- pred_panoptic_seg (PixelData): Panoptic segmentation prediction results. The type constraint is `PixelData`.
- pred_semantic_seg (PixelData): Semantic segmentation prediction results. The type constraint is `PixelData`.
- Intermediate Results
- proposal (InstanceData): Mostly used for the RPN results in the two-stage algorithms. The type constraint is `InstanceData`.
`DetDataSample` is used in the following way. It will throw an error when the data type is invalid, for example, using `torch.Tensor` to define `proposals` instead of `InstanceData`.
[0.4524, 0.8265, 0.4262, 0.2215]]) should be a <class'mmengine.data.instance_data.InstanceData'> but got <class'torch.Tensor'>
```
## Simpify the interfaces
In this section, we use MMDetection to demonstrate how to migrate the abstract data interfaces to simplify the module and component interfaces. We suppose both `DetDataSample` and `InstanceData` have been implemented in MMDetection and MMEngine.
### 1. Simplify the module interface
Detector's external interfaces can be significantly simplified and unified. In the training process of a single-stage detection and segmentation algorithm in MMDet 2.X, `SingleStageDetector` requires `img`, `img_metas`, `gt_bboxes`，`gt_labels` and `gt_bboxes_ignore` as the inputs, but `SingleStageInstanceSegmentor` requires `gt_masks` as well. This causes inconsistency in the training interface and affects flexibility.
```python
class SingleStageDetector(BaseDetector):
...
def forward_train(self,
img,
img_metas,
gt_bboxes,
gt_labels,
gt_bboxes_ignore=None):
class SingleStageInstanceSegmentor(BaseDetector):
...
def forward_train(self,
img,
img_metas,
gt_masks,
gt_labels,
gt_bboxes=None,
gt_bboxes_ignore=None,
**kwargs):
```
In MMDet 3.X, the training interfaces of all the detectors can be unified as `img` and `data_samples` using `DetDataSample`. Different modules can use `data_samples` to encapsulate their own attributes.
```python
class SingleStageDetector(BaseDetector):
...
def forward_train(self,
img,
data_samples):
class SingleStageInstanceSegmentor(BaseDetector):
...
def forward_train(self,
img,
data_samples):
```
### 2. Simplify the model interfaces
In MMDet 2.X, `HungarianAssigner` and `MaskHungarianAssigner` will be used to assign bboxes and instance segment information with annotated instances, respectively. The assignment logics of these two modules are the same, and the only differences are the interface and the calculation of the loss functions. However, this difference makes the code of `HungarianAssigner` cannot be directly used in `MaskHungarianAssigner`, which caused the redundancy.
```python
class HungarianAssigner(BaseAssigner):
def assign(self,
bbox_pred,
cls_pred,
gt_bboxes,
gt_labels,
img_meta,
gt_bboxes_ignore=None,
eps=1e-7):
class MaskHungarianAssigner(BaseAssigner):
def assign(self,
cls_pred,
mask_pred,
gt_labels,
gt_mask,
img_meta,
gt_bboxes_ignore=None,
eps=1e-7):
```
In MMDet 3.X, `InstanceData` can encapsulate the bounding boxes, scores, and masks. With this, we can simplify the core parameters of `HungarianAssigner` to `pred_instances`, `gt_instances`, and `gt_instances_ignore`. This unifies the two assigners into one `HungarianAssianger`.
