Using Mask2Former Segmentor with Custom Dataset (Without gt_masks and gt_labels)

I'm trying to fine-tune a Mask2Former segmentor with a custom dataset by using MMSEGMENTATION plus the code made available in this repository. However, I'm running into some trouble because of the dataset format I'm using.

The code and config file are expecting a dataset in the ADE20K format.

Because of this (as shown below) when collecting the dataset, the train pipeline looks for the keys: 'img', 'gt_semantic_seg', 'gt_masks', 'gt_labels'.

train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations', reduce_zero_label=True),
    dict(type='Resize', img_scale=(2048, 512), ratio_range=(0.5, 2.0)),
    dict(type='RandomCrop', crop_size=(512, 512), cat_max_ratio=0.75),
    dict(type='RandomFlip', prob=0.5),
    dict(type='PhotoMetricDistortion'),
    dict(
        type='Normalize',
        mean=[123.675, 116.28, 103.53],
        std=[58.395, 57.12, 57.375],
        to_rgb=True),
    dict(type='Pad', size=(512, 512), pad_val=0, seg_pad_val=255),
    dict(type='DefaultFormatBundle'),
    dict(type='Collect', keys=['img', 'gt_semantic_seg', 'gt_masks', 'gt_labels'])
]

My dataset, though, is simpler. It only has images in a folder and semantic segmentation masks in another: that is, the first two keys. Becaue of that, when I try to fine-tune the Mask2Former Segmentor using my dataset, I get the error below.

Original Traceback (most recent call last):
  File "/usr/local/lib/python3.8/dist-packages/torch/utils/data/_utils/worker.py", line 308, in _worker_loop
    data = fetcher.fetch(index)
  File "/usr/local/lib/python3.8/dist-packages/torch/utils/data/_utils/fetch.py", line 51, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/usr/local/lib/python3.8/dist-packages/torch/utils/data/_utils/fetch.py", line 51, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/mmsegmentation/mmseg/datasets/custom.py", line 215, in __getitem__
    return self.prepare_train_img(idx)
  File "/mmsegmentation/mmseg/datasets/custom.py", line 232, in prepare_train_img
    return self.pipeline(results)
  File "/mmsegmentation/mmseg/datasets/pipelines/compose.py", line 41, in __call__
    data = t(data)
  File "/mmsegmentation/mmseg/datasets/pipelines/formatting.py", line 284, in __call__
    data[key] = results[key]
KeyError: 'gt_masks'

So the question is: has anyone tried adapting the Mask2Former segmentor to a dataset of this kind?

facebookresearch / dinov2

Using Mask2Former Segmentor with Custom Dataset (Without gt_masks and gt_labels) #259