mmsegmentation 数据集

CustomDataset 类

def __getitem__(self, idx):
    """Get training/test data after pipeline.
    
    Args:
    idx (int): Index of data.
    
    Returns:
    dict: Training/test data (with annotation if `test_mode` is set
    False).
    """
    
    if self.test_mode:
        return self.prepare_test_img(idx)
    else:
        return self.prepare_train_img(idx)
    
def prepare_train_img(self, idx):
        """Get training data and annotations after pipeline.
    Args:
        idx (int): Index of data.

    Returns:
        dict: Training data and annotation after pipeline with new keys
            introduced by pipeline.
    """

    img_info = self.img_infos[idx]
    ann_info = self.get_ann_info(idx)
    results = dict(img_info=img_info, ann_info=ann_info)
    self.pre_pipeline(results)
    return self.pipeline(results)

def prepare_test_img(self, idx):
    """Get testing data after pipeline.

    Args:
        idx (int): Index of data.

    Returns:
        dict: Testing data after pipeline with new keys introduced by
            pipeline.
    """

    img_info = self.img_infos[idx]
    results = dict(img_info=img_info)
    self.pre_pipeline(results)
    return self.pipeline(results)