Repository: open-mmlab/mmsegmentation

Issue:
Troubleshooting LoadAnnotations Error with Channel Masks
I have three channel masks, each representing one of three classes. The masks have binary values of 0 and 1 for the background and target. Unfortunately, there are some conflicting areas in the masks that prevent me from using multi-class labeling. I am currently using the `LoadAnnotations ` function to load the mask PNGs, but I am encountering an error. How can I resolve this issue?

Traceback (most recent call last):
  File "/content/mmsegmentation/tools/train.py", line 106, in <module>
    main()
  File "/content/mmsegmentation/tools/train.py", line 102, in main
    runner.train()
  File "/usr/local/lib/python3.8/dist-packages/mmengine/runner/runner.py", line 1686, in train
    model = self.train_loop.run()  # type: ignore
  File "/usr/local/lib/python3.8/dist-packages/mmengine/runner/loops.py", line 263, in run
    data_batch = next(self.dataloader_iterator)
  File "/usr/local/lib/python3.8/dist-packages/mmengine/runner/loops.py", line 158, in __next__
    data = next(self._iterator)
  File "/usr/local/lib/python3.8/dist-packages/torch/utils/data/dataloader.py", line 628, in __next__
    data = self._next_data()
  File "/usr/local/lib/python3.8/dist-packages/torch/utils/data/dataloader.py", line 1333, in _next_data
    return self._process_data(data)
  File "/usr/local/lib/python3.8/dist-packages/torch/utils/data/dataloader.py", line 1359, in _process_data
    data.reraise()
  File "/usr/local/lib/python3.8/dist-packages/torch/_utils.py", line 543, in reraise
    raise exception
AssertionError: Caught AssertionError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/usr/local/lib/python3.8/dist-packages/torch/utils/data/_utils/worker.py", line 302, in _worker_loop
    data = fetcher.fetch(index)
  File "/usr/local/lib/python3.8/dist-packages/torch/utils/data/_utils/fetch.py", line 58, 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 58, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/usr/local/lib/python3.8/dist-packages/mmengine/dataset/base_dataset.py", line 408, in __getitem__
    data = self.prepare_data(idx)
  File "/usr/local/lib/python3.8/dist-packages/mmengine/dataset/base_dataset.py", line 789, in prepare_data
    return self.pipeline(data_info)
  File "/usr/local/lib/python3.8/dist-packages/mmengine/dataset/base_dataset.py", line 58, in __call__
    data = t(data)
  File "/usr/local/lib/python3.8/dist-packages/mmcv/transforms/base.py", line 12, in __call__
    return self.transform(results)
  File "/content/mmsegmentation/mmseg/datasets/transforms/formatting.py", line 74, in transform
    data_sample.gt_sem_seg = PixelData(**gt_sem_seg_data)
  File "/usr/local/lib/python3.8/dist-packages/mmengine/structures/base_data_element.py", line 219, in __init__
    self.set_data(kwargs)
  File "/usr/local/lib/python3.8/dist-packages/mmengine/structures/base_data_element.py", line 249, in set_data
    setattr(self, k, v)
  File "/usr/local/lib/python3.8/dist-packages/mmengine/structures/pixel_data.py", line 78, in __setattr__
    assert value.ndim in [
AssertionError: The dim of value must be 2 or 3, but got 4





Discussion:
I think you can convert the segmentation masks of your dataset into single channel mask. 
here is an example:
https://github.com/open-mmlab/mmsegmentation/blob/b2fdae7ce794483e590e8ec8dfb63e86877ab115/tools/convert_datasets/isaid.py#L37
Actually, single channel mask solution is not my case. As I mentioned before, some classes have joint masks. For example, a pixel target can be [0,1,1] as a result of three labels. But, I just wonder that, can I use `LoadBiomedicalAnnotation` to load multichannel masks for a PNG image with `LoadImage`?

Resolving patch:
diff --git a/mmseg/datasets/transforms/formatting.py b/mmseg/datasets/transforms/formatting.py
--- a/mmseg/datasets/transforms/formatting.py
+++ b/mmseg/datasets/transforms/formatting.py
@@ -1,4 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import warnings
+
 import numpy as np
 from mmcv.transforms import to_tensor
 from mmcv.transforms.base import BaseTransform
@@ -72,9 +74,16 @@ def transform(self, results: dict) -> dict:
 
         data_sample = SegDataSample()
         if 'gt_seg_map' in results:
-            gt_sem_seg_data = dict(
-                data=to_tensor(results['gt_seg_map'][None,
-                                                     ...].astype(np.int64)))
+            if results['gt_seg_map'].shape == 2:
+                data = to_tensor(results['gt_seg_map'][None,
+                                                       ...].astype(np.int64))
+            else:
+                warnings.warn('Please pay attention your ground truth '
+                              'segmentation map, usually the segentation '
+                              'map is 2D, but got '
+                              f'{results["gt_seg_map"].shape}')
+                data = to_tensor(results['gt_seg_map'].astype(np.int64))
+            gt_sem_seg_data = dict(data=data)
             data_sample.gt_sem_seg = PixelData(**gt_sem_seg_data)
 
         if 'gt_edge_map' in results:


Test patch:
diff --git a/tests/test_datasets/test_formatting.py b/tests/test_datasets/test_formatting.py
--- a/tests/test_datasets/test_formatting.py
+++ b/tests/test_datasets/test_formatting.py
@@ -45,6 +45,11 @@ def test_transform(self):
                               BaseDataElement)
         self.assertEqual(results['data_samples'].ori_shape,
                          results['data_samples'].gt_sem_seg.shape)
+        results = copy.deepcopy(self.results)
+        results['gt_seg_map'] = np.random.rand(3, 300, 400)
+        results = transform(results)
+        self.assertEqual(results['data_samples'].ori_shape,
+                         results['data_samples'].gt_sem_seg.shape)
 
     def test_repr(self):
         transform = PackSegInputs(meta_keys=self.meta_keys)