[Fix] Fix binary segmentation

mmseg/models/decode_heads/decode_head.py

@@ -1,4 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import warnings
 from abc import ABCMeta, abstractmethod
 from typing import List, Tuple
 
@@ -44,6 +45,9 @@ class BaseDecodeHead(BaseModule, metaclass=ABCMeta):
  in_channels (int|Sequence[int]): Input channels.
  channels (int): Channels after modules, before conv_seg.
  num_classes (int): Number of classes.
+ out_channels (int): Output channels of conv_seg.
+ threshold (float): Threshold for binary segmentation in the case of
+ `num_classes==1`. Default: None.
  dropout_ratio (float): Ratio of dropout layer. Default: 0.1.
  conv_cfg (dict|None): Config of conv layers. Default: None.
  norm_cfg (dict|None): Config of norm layers. Default: None.
@@ -82,6 +86,8 @@ def __init__(self,
  channels,
  *,
  num_classes,
+ out_channels=None,
+ threshold=None,
  dropout_ratio=0.1,
  conv_cfg=None,
  norm_cfg=None,
@@ -100,7 +106,6 @@ def __init__(self,
  super().__init__(init_cfg)
  self._init_inputs(in_channels, in_index, input_transform)
  self.channels = channels
- self.num_classes = num_classes
  self.dropout_ratio = dropout_ratio
  self.conv_cfg = conv_cfg
  self.norm_cfg = norm_cfg
@@ -110,6 +115,30 @@ def __init__(self,
  self.ignore_index = ignore_index
  self.align_corners = align_corners
 
+ if out_channels is None:
+ if num_classes == 2:
+ warnings.warn('For binary segmentation, we suggest using'
+ '`out_channels = 1` to define the output'
+ 'channels of segmentor, and use `threshold`'
+ 'to convert seg_logist into a prediction'
+ 'applying a threshold')
+ out_channels = num_classes
+
+ if out_channels != num_classes and out_channels != 1:
+ raise ValueError(
+ 'out_channels should be equal to num_classes,'
+ 'except binary segmentation set out_channels == 1 and'
+ f'num_classes == 2, but got out_channels={out_channels}'
+ f'and num_classes={num_classes}')
+
+ if out_channels == 1 and threshold is None:
+ threshold = 0.3
+ warnings.warn('threshold is not defined for binary, and defaults'
+ 'to 0.3')
+ self.num_classes = num_classes
+ self.out_channels = out_channels
+ self.threshold = threshold
+
  if isinstance(loss_decode, dict):
  self.loss_decode = build_loss(loss_decode)
  elif isinstance(loss_decode, (list, tuple)):
@@ -125,7 +154,7 @@ def __init__(self,
  else:
  self.sampler = None
 
- self.conv_seg = nn.Conv2d(channels, num_classes, kernel_size=1)
+ self.conv_seg = nn.Conv2d(channels, self.out_channels, kernel_size=1)
  if dropout_ratio > 0:
  self.dropout = nn.Dropout2d(dropout_ratio)
  else:

mmseg/models/segmentors/base.py

@@ -148,11 +148,9 @@ def postprocess_result(self,
  segmentation before normalization.
  """
  batch_size, C, H, W = seg_logits.shape
- assert C > 1, ('This post processes does not binary segmentation, and '
- f'channels `seg_logtis` must be > 1 but got {C}')
 
  if data_samples is None:
- data_samples = []
+ data_samples = [SegDataSample() for _ in range(batch_size)]
  only_prediction = True
  else:
  only_prediction = False
@@ -174,23 +172,19 @@ def postprocess_result(self,
  mode='bilinear',
  align_corners=self.align_corners,
  warning=False).squeeze(0)
- # i_seg_logits shape is C, H, W with original shape
- i_seg_pred = i_seg_logits.argmax(dim=0, keepdim=True)
- data_samples[i].set_data({
- 'seg_logits':
- PixelData(**{'data': i_seg_logits}),
- 'pred_sem_seg':
- PixelData(**{'data': i_seg_pred})
- })
  else:
  i_seg_logits = seg_logits[i]
+
+ if C > 1:
  i_seg_pred = i_seg_logits.argmax(dim=0, keepdim=True)
- prediction = SegDataSample()
- prediction.set_data({
- 'seg_logits':
- PixelData(**{'data': i_seg_logits}),
- 'pred_sem_seg':
- PixelData(**{'data': i_seg_pred})
- })
- data_samples.append(prediction)
+ else:
+ i_seg_pred = (i_seg_logits >
+ self.decode_head.threshold).to(i_seg_logits)
+ data_samples[i].set_data({
+ 'seg_logits':
+ PixelData(**{'data': i_seg_logits}),
+ 'pred_sem_seg':
+ PixelData(**{'data': i_seg_pred})
+ })
+
  return data_samples

mmseg/models/segmentors/encoder_decoder.py

@@ -100,6 +100,7 @@ def _init_decode_head(self, decode_head: ConfigType) -> None:
  self.decode_head = MODELS.build(decode_head)
  self.align_corners = self.decode_head.align_corners
  self.num_classes = self.decode_head.num_classes
+ self.out_channels = self.decode_head.out_channels
 
  def _init_auxiliary_head(self, auxiliary_head: ConfigType) -> None:
  """Initialize ``auxiliary_head``"""

tests/test_models/test_heads/test_decode_head.py

@@ -0,0 +1,193 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from unittest.mock import patch
+
+import pytest
+import torch
+from mmengine.structures import PixelData
+
+from mmseg.models.decode_heads.decode_head import BaseDecodeHead
+from mmseg.structures import SegDataSample
+from .utils import to_cuda
+
+
+@patch.multiple(BaseDecodeHead, __abstractmethods__=set())
+def test_decode_head():
+
+ with pytest.raises(AssertionError):
+ # default input_transform doesn't accept multiple inputs
+ BaseDecodeHead([32, 16], 16, num_classes=19)
+
+ with pytest.raises(AssertionError):
+ # default input_transform doesn't accept multiple inputs
+ BaseDecodeHead(32, 16, num_classes=19, in_index=[-1, -2])
+
+ with pytest.raises(AssertionError):
+ # supported mode is resize_concat only
+ BaseDecodeHead(32, 16, num_classes=19, input_transform='concat')
+
+ with pytest.raises(AssertionError):
+ # in_channels should be list|tuple
+ BaseDecodeHead(32, 16, num_classes=19, input_transform='resize_concat')
+
+ with pytest.raises(AssertionError):
+ # in_index should be list|tuple
+ BaseDecodeHead([32],
+ 16,
+ in_index=-1,
+ num_classes=19,
+ input_transform='resize_concat')
+
+ with pytest.raises(AssertionError):
+ # len(in_index) should equal len(in_channels)
+ BaseDecodeHead([32, 16],
+ 16,
+ num_classes=19,
+ in_index=[-1],
+ input_transform='resize_concat')
+
+ with pytest.raises(ValueError):
+ # out_channels should be equal to num_classes
+ BaseDecodeHead(32, 16, num_classes=19, out_channels=18)
+
+ # test out_channels
+ head = BaseDecodeHead(32, 16, num_classes=2)
+ assert head.out_channels == 2
+
+ # test out_channels == 1 and num_classes == 2
+ head = BaseDecodeHead(32, 16, num_classes=2, out_channels=1)
+ assert head.out_channels == 1 and head.num_classes == 2
+
+ # test default dropout
+ head = BaseDecodeHead(32, 16, num_classes=19)
+ assert hasattr(head, 'dropout') and head.dropout.p == 0.1
+
+ # test set dropout
+ head = BaseDecodeHead(32, 16, num_classes=19, dropout_ratio=0.2)
+ assert hasattr(head, 'dropout') and head.dropout.p == 0.2
+
+ # test no input_transform
+ inputs = [torch.randn(1, 32, 45, 45)]
+ head = BaseDecodeHead(32, 16, num_classes=19)
+ if torch.cuda.is_available():
+ head, inputs = to_cuda(head, inputs)
+ assert head.in_channels == 32
+ assert head.input_transform is None
+ transformed_inputs = head._transform_inputs(inputs)
+ assert transformed_inputs.shape == (1, 32, 45, 45)
+
+ # test input_transform = resize_concat
+ inputs = [torch.randn(1, 32, 45, 45), torch.randn(1, 16, 21, 21)]
+ head = BaseDecodeHead([32, 16],
+ 16,
+ num_classes=19,
+ in_index=[0, 1],
+ input_transform='resize_concat')
+ if torch.cuda.is_available():
+ head, inputs = to_cuda(head, inputs)
+ assert head.in_channels == 48
+ assert head.input_transform == 'resize_concat'
+ transformed_inputs = head._transform_inputs(inputs)
+ assert transformed_inputs.shape == (1, 48, 45, 45)
+
+ # test multi-loss, loss_decode is dict
+ with pytest.raises(TypeError):
+ # loss_decode must be a dict or sequence of dict.
+ BaseDecodeHead(3, 16, num_classes=19, loss_decode=['CrossEntropyLoss'])
+
+ inputs = torch.randn(2, 19, 8, 8).float()
+ data_samples = [
+ SegDataSample(gt_sem_seg=PixelData(data=torch.ones(64, 64).long()))
+ for _ in range(2)
+ ]
+
+ head = BaseDecodeHead(
+ 3,
+ 16,
+ num_classes=19,
+ loss_decode=dict(
+ type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0))
+ if torch.cuda.is_available():
+ head, inputs = to_cuda(head, inputs)
+ loss = head.loss_by_feat(
+ seg_logits=inputs, batch_data_samples=data_samples)
+ assert 'loss_ce' in loss
+
+ # test multi-loss, loss_decode is list of dict
+ inputs = torch.randn(2, 19, 8, 8).float()
+ data_samples = [
+ SegDataSample(gt_sem_seg=PixelData(data=torch.ones(64, 64).long()))
+ for _ in range(2)
+ ]
+ head = BaseDecodeHead(
+ 3,
+ 16,
+ num_classes=19,
+ loss_decode=[
+ dict(type='CrossEntropyLoss', loss_name='loss_1'),
+ dict(type='CrossEntropyLoss', loss_name='loss_2')
+ ])
+ if torch.cuda.is_available():
+ head, inputs = to_cuda(head, inputs)
+
+ loss = head.loss_by_feat(
+ seg_logits=inputs, batch_data_samples=data_samples)
+ assert 'loss_1' in loss
+ assert 'loss_2' in loss
+
+ # 'loss_decode' must be a dict or sequence of dict
+ with pytest.raises(TypeError):
+ BaseDecodeHead(3, 16, num_classes=19, loss_decode=['CrossEntropyLoss'])
+ with pytest.raises(TypeError):
+ BaseDecodeHead(3, 16, num_classes=19, loss_decode=0)
+
+ # test multi-loss, loss_decode is list of dict
+ inputs = torch.randn(2, 19, 8, 8).float()
+ data_samples = [
+ SegDataSample(gt_sem_seg=PixelData(data=torch.ones(64, 64).long()))
+ for _ in range(2)
+ ]
+ head = BaseDecodeHead(
+ 3,
+ 16,
+ num_classes=19,
+ loss_decode=(dict(type='CrossEntropyLoss', loss_name='loss_1'),
+ dict(type='CrossEntropyLoss', loss_name='loss_2'),
+ dict(type='CrossEntropyLoss', loss_name='loss_3')))
+ if torch.cuda.is_available():
+ head, inputs = to_cuda(head, inputs)
+ loss = head.loss_by_feat(
+ seg_logits=inputs, batch_data_samples=data_samples)
+ assert 'loss_1' in loss
+ assert 'loss_2' in loss
+ assert 'loss_3' in loss
+
+ # test multi-loss, loss_decode is list of dict, names of them are identical
+ inputs = torch.randn(2, 19, 8, 8).float()
+ data_samples = [
+ SegDataSample(gt_sem_seg=PixelData(data=torch.ones(64, 64).long()))
+ for _ in range(2)
+ ]
+ head = BaseDecodeHead(
+ 3,
+ 16,
+ num_classes=19,
+ loss_decode=(dict(type='CrossEntropyLoss', loss_name='loss_ce'),
+ dict(type='CrossEntropyLoss', loss_name='loss_ce'),
+ dict(type='CrossEntropyLoss', loss_name='loss_ce')))
+ if torch.cuda.is_available():
+ head, inputs = to_cuda(head, inputs)
+ loss_3 = head.loss_by_feat(
+ seg_logits=inputs, batch_data_samples=data_samples)
+
+ head = BaseDecodeHead(
+ 3,
+ 16,
+ num_classes=19,
+ loss_decode=(dict(type='CrossEntropyLoss', loss_name='loss_ce')))
+ if torch.cuda.is_available():
+ head, inputs = to_cuda(head, inputs)
+ loss = head.loss_by_feat(
+ seg_logits=inputs, batch_data_samples=data_samples)
+ assert 'loss_ce' in loss
+ assert 'loss_ce' in loss_3
+ assert loss_3['loss_ce'] == 3 * loss['loss_ce']

tests/test_models/test_segmentors/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) OpenMMLab. All rights reserved.

tests/test_models/test_segmentors/test_cascade_encoder_decoder.py

@@ -0,0 +1,57 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from mmengine import ConfigDict
+
+from mmseg.models import build_segmentor
+from .utils import _segmentor_forward_train_test
+
+
+def test_cascade_encoder_decoder():
+
+ # test 1 decode head, w.o. aux head
+ cfg = ConfigDict(
+ type='CascadeEncoderDecoder',
+ num_stages=2,
+ backbone=dict(type='ExampleBackbone'),
+ decode_head=[
+ dict(type='ExampleDecodeHead'),
+ dict(type='ExampleCascadeDecodeHead')
+ ])
+ cfg.test_cfg = ConfigDict(mode='whole')
+ segmentor = build_segmentor(cfg)
+ _segmentor_forward_train_test(segmentor)
+
+ # test slide mode
+ cfg.test_cfg = ConfigDict(mode='slide', crop_size=(3, 3), stride=(2, 2))
+ segmentor = build_segmentor(cfg)
+ _segmentor_forward_train_test(segmentor)
+
+ # test 1 decode head, 1 aux head
+ cfg = ConfigDict(
+ type='CascadeEncoderDecoder',
+ num_stages=2,
+ backbone=dict(type='ExampleBackbone'),
+ decode_head=[
+ dict(type='ExampleDecodeHead'),
+ dict(type='ExampleCascadeDecodeHead')
+ ],
+ auxiliary_head=dict(type='ExampleDecodeHead'))
+ cfg.test_cfg = ConfigDict(mode='whole')
+ segmentor = build_segmentor(cfg)
+ _segmentor_forward_train_test(segmentor)
+
+ # test 1 decode head, 2 aux head
+ cfg = ConfigDict(
+ type='CascadeEncoderDecoder',
+ num_stages=2,
+ backbone=dict(type='ExampleBackbone'),
+ decode_head=[
+ dict(type='ExampleDecodeHead'),
+ dict(type='ExampleCascadeDecodeHead')
+ ],
+ auxiliary_head=[
+ dict(type='ExampleDecodeHead'),
+ dict(type='ExampleDecodeHead')
+ ])
+ cfg.test_cfg = ConfigDict(mode='whole')
+ segmentor = build_segmentor(cfg)
+ _segmentor_forward_train_test(segmentor)