mirror of
https://github.com/open-mmlab/mmsegmentation.git
synced 2025-06-03 22:03:48 +08:00
857f854b61
Thanks for your contribution and we appreciate it a lot. The following instructions would make your pull request more healthy and more easily get feedback. If you do not understand some items, don't worry, just make the pull request and seek help from maintainers. ## Motivation https://github.com/open-mmlab/mmsegmentation/issues/3181 https://github.com/open-mmlab/mmsegmentation/issues/2965 https://github.com/open-mmlab/mmsegmentation/issues/2644 https://github.com/open-mmlab/mmsegmentation/issues/1645 https://github.com/open-mmlab/mmsegmentation/issues/1444 https://github.com/open-mmlab/mmsegmentation/issues/1370 https://github.com/open-mmlab/mmsegmentation/issues/125 ## Modification Remove the assertion at data_preprocessor ## BC-breaking (Optional) Does the modification introduce changes that break the backward-compatibility of the downstream repos? If so, please describe how it breaks the compatibility and how the downstream projects should modify their code to keep compatibility with this PR. ## Use cases (Optional) If this PR introduces a new feature, it is better to list some use cases here, and update the documentation. ## Checklist 1. Pre-commit or other linting tools are used to fix the potential lint issues. 2. The modification is covered by complete unit tests. If not, please add more unit test to ensure the correctness. 3. If the modification has potential influence on downstream projects, this PR should be tested with downstream projects, like MMDet or MMDet3D. 4. The documentation has been modified accordingly, like docstring or example tutorials.
365 lines
14 KiB
Python
365 lines
14 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import logging
|
|
from typing import List, Optional
|
|
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from mmengine.logging import print_log
|
|
from torch import Tensor
|
|
|
|
from mmseg.registry import MODELS
|
|
from mmseg.utils import (ConfigType, OptConfigType, OptMultiConfig,
|
|
OptSampleList, SampleList, add_prefix)
|
|
from .base import BaseSegmentor
|
|
|
|
|
|
@MODELS.register_module()
|
|
class EncoderDecoder(BaseSegmentor):
|
|
"""Encoder Decoder segmentors.
|
|
|
|
EncoderDecoder typically consists of backbone, decode_head, auxiliary_head.
|
|
Note that auxiliary_head is only used for deep supervision during training,
|
|
which could be dumped during inference.
|
|
|
|
1. The ``loss`` method is used to calculate the loss of model,
|
|
which includes two steps: (1) Extracts features to obtain the feature maps
|
|
(2) Call the decode head loss function to forward decode head model and
|
|
calculate losses.
|
|
|
|
.. code:: text
|
|
|
|
loss(): extract_feat() -> _decode_head_forward_train() -> _auxiliary_head_forward_train (optional)
|
|
_decode_head_forward_train(): decode_head.loss()
|
|
_auxiliary_head_forward_train(): auxiliary_head.loss (optional)
|
|
|
|
2. The ``predict`` method is used to predict segmentation results,
|
|
which includes two steps: (1) Run inference function to obtain the list of
|
|
seg_logits (2) Call post-processing function to obtain list of
|
|
``SegDataSample`` including ``pred_sem_seg`` and ``seg_logits``.
|
|
|
|
.. code:: text
|
|
|
|
predict(): inference() -> postprocess_result()
|
|
infercen(): whole_inference()/slide_inference()
|
|
whole_inference()/slide_inference(): encoder_decoder()
|
|
encoder_decoder(): extract_feat() -> decode_head.predict()
|
|
|
|
3. The ``_forward`` method is used to output the tensor by running the model,
|
|
which includes two steps: (1) Extracts features to obtain the feature maps
|
|
(2)Call the decode head forward function to forward decode head model.
|
|
|
|
.. code:: text
|
|
|
|
_forward(): extract_feat() -> _decode_head.forward()
|
|
|
|
Args:
|
|
|
|
backbone (ConfigType): The config for the backnone of segmentor.
|
|
decode_head (ConfigType): The config for the decode head of segmentor.
|
|
neck (OptConfigType): The config for the neck of segmentor.
|
|
Defaults to None.
|
|
auxiliary_head (OptConfigType): The config for the auxiliary head of
|
|
segmentor. Defaults to None.
|
|
train_cfg (OptConfigType): The config for training. Defaults to None.
|
|
test_cfg (OptConfigType): The config for testing. Defaults to None.
|
|
data_preprocessor (dict, optional): The pre-process config of
|
|
:class:`BaseDataPreprocessor`.
|
|
pretrained (str, optional): The path for pretrained model.
|
|
Defaults to None.
|
|
init_cfg (dict, optional): The weight initialized config for
|
|
:class:`BaseModule`.
|
|
""" # noqa: E501
|
|
|
|
def __init__(self,
|
|
backbone: ConfigType,
|
|
decode_head: ConfigType,
|
|
neck: OptConfigType = None,
|
|
auxiliary_head: OptConfigType = None,
|
|
train_cfg: OptConfigType = None,
|
|
test_cfg: OptConfigType = None,
|
|
data_preprocessor: OptConfigType = None,
|
|
pretrained: Optional[str] = None,
|
|
init_cfg: OptMultiConfig = None):
|
|
super().__init__(
|
|
data_preprocessor=data_preprocessor, init_cfg=init_cfg)
|
|
if pretrained is not None:
|
|
assert backbone.get('pretrained') is None, \
|
|
'both backbone and segmentor set pretrained weight'
|
|
backbone.pretrained = pretrained
|
|
self.backbone = MODELS.build(backbone)
|
|
if neck is not None:
|
|
self.neck = MODELS.build(neck)
|
|
self._init_decode_head(decode_head)
|
|
self._init_auxiliary_head(auxiliary_head)
|
|
|
|
self.train_cfg = train_cfg
|
|
self.test_cfg = test_cfg
|
|
|
|
assert self.with_decode_head
|
|
|
|
def _init_decode_head(self, decode_head: ConfigType) -> None:
|
|
"""Initialize ``decode_head``"""
|
|
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``"""
|
|
if auxiliary_head is not None:
|
|
if isinstance(auxiliary_head, list):
|
|
self.auxiliary_head = nn.ModuleList()
|
|
for head_cfg in auxiliary_head:
|
|
self.auxiliary_head.append(MODELS.build(head_cfg))
|
|
else:
|
|
self.auxiliary_head = MODELS.build(auxiliary_head)
|
|
|
|
def extract_feat(self, inputs: Tensor) -> List[Tensor]:
|
|
"""Extract features from images."""
|
|
x = self.backbone(inputs)
|
|
if self.with_neck:
|
|
x = self.neck(x)
|
|
return x
|
|
|
|
def encode_decode(self, inputs: Tensor,
|
|
batch_img_metas: List[dict]) -> Tensor:
|
|
"""Encode images with backbone and decode into a semantic segmentation
|
|
map of the same size as input."""
|
|
x = self.extract_feat(inputs)
|
|
seg_logits = self.decode_head.predict(x, batch_img_metas,
|
|
self.test_cfg)
|
|
|
|
return seg_logits
|
|
|
|
def _decode_head_forward_train(self, inputs: List[Tensor],
|
|
data_samples: SampleList) -> dict:
|
|
"""Run forward function and calculate loss for decode head in
|
|
training."""
|
|
losses = dict()
|
|
loss_decode = self.decode_head.loss(inputs, data_samples,
|
|
self.train_cfg)
|
|
|
|
losses.update(add_prefix(loss_decode, 'decode'))
|
|
return losses
|
|
|
|
def _auxiliary_head_forward_train(self, inputs: List[Tensor],
|
|
data_samples: SampleList) -> dict:
|
|
"""Run forward function and calculate loss for auxiliary head in
|
|
training."""
|
|
losses = dict()
|
|
if isinstance(self.auxiliary_head, nn.ModuleList):
|
|
for idx, aux_head in enumerate(self.auxiliary_head):
|
|
loss_aux = aux_head.loss(inputs, data_samples, self.train_cfg)
|
|
losses.update(add_prefix(loss_aux, f'aux_{idx}'))
|
|
else:
|
|
loss_aux = self.auxiliary_head.loss(inputs, data_samples,
|
|
self.train_cfg)
|
|
losses.update(add_prefix(loss_aux, 'aux'))
|
|
|
|
return losses
|
|
|
|
def loss(self, inputs: Tensor, data_samples: SampleList) -> dict:
|
|
"""Calculate losses from a batch of inputs and data samples.
|
|
|
|
Args:
|
|
inputs (Tensor): Input images.
|
|
data_samples (list[:obj:`SegDataSample`]): The seg data samples.
|
|
It usually includes information such as `metainfo` and
|
|
`gt_sem_seg`.
|
|
|
|
Returns:
|
|
dict[str, Tensor]: a dictionary of loss components
|
|
"""
|
|
|
|
x = self.extract_feat(inputs)
|
|
|
|
losses = dict()
|
|
|
|
loss_decode = self._decode_head_forward_train(x, data_samples)
|
|
losses.update(loss_decode)
|
|
|
|
if self.with_auxiliary_head:
|
|
loss_aux = self._auxiliary_head_forward_train(x, data_samples)
|
|
losses.update(loss_aux)
|
|
|
|
return losses
|
|
|
|
def predict(self,
|
|
inputs: Tensor,
|
|
data_samples: OptSampleList = None) -> SampleList:
|
|
"""Predict results from a batch of inputs and data samples with post-
|
|
processing.
|
|
|
|
Args:
|
|
inputs (Tensor): Inputs with shape (N, C, H, W).
|
|
data_samples (List[:obj:`SegDataSample`], optional): The seg data
|
|
samples. It usually includes information such as `metainfo`
|
|
and `gt_sem_seg`.
|
|
|
|
Returns:
|
|
list[:obj:`SegDataSample`]: Segmentation results of the
|
|
input images. Each SegDataSample usually contain:
|
|
|
|
- ``pred_sem_seg``(PixelData): Prediction of semantic segmentation.
|
|
- ``seg_logits``(PixelData): Predicted logits of semantic
|
|
segmentation before normalization.
|
|
"""
|
|
if data_samples is not None:
|
|
batch_img_metas = [
|
|
data_sample.metainfo for data_sample in data_samples
|
|
]
|
|
else:
|
|
batch_img_metas = [
|
|
dict(
|
|
ori_shape=inputs.shape[2:],
|
|
img_shape=inputs.shape[2:],
|
|
pad_shape=inputs.shape[2:],
|
|
padding_size=[0, 0, 0, 0])
|
|
] * inputs.shape[0]
|
|
|
|
seg_logits = self.inference(inputs, batch_img_metas)
|
|
|
|
return self.postprocess_result(seg_logits, data_samples)
|
|
|
|
def _forward(self,
|
|
inputs: Tensor,
|
|
data_samples: OptSampleList = None) -> Tensor:
|
|
"""Network forward process.
|
|
|
|
Args:
|
|
inputs (Tensor): Inputs with shape (N, C, H, W).
|
|
data_samples (List[:obj:`SegDataSample`]): The seg
|
|
data samples. It usually includes information such
|
|
as `metainfo` and `gt_sem_seg`.
|
|
|
|
Returns:
|
|
Tensor: Forward output of model without any post-processes.
|
|
"""
|
|
x = self.extract_feat(inputs)
|
|
return self.decode_head.forward(x)
|
|
|
|
def slide_inference(self, inputs: Tensor,
|
|
batch_img_metas: List[dict]) -> Tensor:
|
|
"""Inference by sliding-window with overlap.
|
|
|
|
If h_crop > h_img or w_crop > w_img, the small patch will be used to
|
|
decode without padding.
|
|
|
|
Args:
|
|
inputs (tensor): the tensor should have a shape NxCxHxW,
|
|
which contains all images in the batch.
|
|
batch_img_metas (List[dict]): List of image metainfo where each may
|
|
also contain: 'img_shape', 'scale_factor', 'flip', 'img_path',
|
|
'ori_shape', and 'pad_shape'.
|
|
For details on the values of these keys see
|
|
`mmseg/datasets/pipelines/formatting.py:PackSegInputs`.
|
|
|
|
Returns:
|
|
Tensor: The segmentation results, seg_logits from model of each
|
|
input image.
|
|
"""
|
|
|
|
h_stride, w_stride = self.test_cfg.stride
|
|
h_crop, w_crop = self.test_cfg.crop_size
|
|
batch_size, _, h_img, w_img = inputs.size()
|
|
out_channels = self.out_channels
|
|
h_grids = max(h_img - h_crop + h_stride - 1, 0) // h_stride + 1
|
|
w_grids = max(w_img - w_crop + w_stride - 1, 0) // w_stride + 1
|
|
preds = inputs.new_zeros((batch_size, out_channels, h_img, w_img))
|
|
count_mat = inputs.new_zeros((batch_size, 1, h_img, w_img))
|
|
for h_idx in range(h_grids):
|
|
for w_idx in range(w_grids):
|
|
y1 = h_idx * h_stride
|
|
x1 = w_idx * w_stride
|
|
y2 = min(y1 + h_crop, h_img)
|
|
x2 = min(x1 + w_crop, w_img)
|
|
y1 = max(y2 - h_crop, 0)
|
|
x1 = max(x2 - w_crop, 0)
|
|
crop_img = inputs[:, :, y1:y2, x1:x2]
|
|
# change the image shape to patch shape
|
|
batch_img_metas[0]['img_shape'] = crop_img.shape[2:]
|
|
# the output of encode_decode is seg logits tensor map
|
|
# with shape [N, C, H, W]
|
|
crop_seg_logit = self.encode_decode(crop_img, batch_img_metas)
|
|
preds += F.pad(crop_seg_logit,
|
|
(int(x1), int(preds.shape[3] - x2), int(y1),
|
|
int(preds.shape[2] - y2)))
|
|
|
|
count_mat[:, :, y1:y2, x1:x2] += 1
|
|
assert (count_mat == 0).sum() == 0
|
|
seg_logits = preds / count_mat
|
|
|
|
return seg_logits
|
|
|
|
def whole_inference(self, inputs: Tensor,
|
|
batch_img_metas: List[dict]) -> Tensor:
|
|
"""Inference with full image.
|
|
|
|
Args:
|
|
inputs (Tensor): The tensor should have a shape NxCxHxW, which
|
|
contains all images in the batch.
|
|
batch_img_metas (List[dict]): List of image metainfo where each may
|
|
also contain: 'img_shape', 'scale_factor', 'flip', 'img_path',
|
|
'ori_shape', and 'pad_shape'.
|
|
For details on the values of these keys see
|
|
`mmseg/datasets/pipelines/formatting.py:PackSegInputs`.
|
|
|
|
Returns:
|
|
Tensor: The segmentation results, seg_logits from model of each
|
|
input image.
|
|
"""
|
|
|
|
seg_logits = self.encode_decode(inputs, batch_img_metas)
|
|
|
|
return seg_logits
|
|
|
|
def inference(self, inputs: Tensor, batch_img_metas: List[dict]) -> Tensor:
|
|
"""Inference with slide/whole style.
|
|
|
|
Args:
|
|
inputs (Tensor): The input image of shape (N, 3, H, W).
|
|
batch_img_metas (List[dict]): List of image metainfo where each may
|
|
also contain: 'img_shape', 'scale_factor', 'flip', 'img_path',
|
|
'ori_shape', 'pad_shape', and 'padding_size'.
|
|
For details on the values of these keys see
|
|
`mmseg/datasets/pipelines/formatting.py:PackSegInputs`.
|
|
|
|
Returns:
|
|
Tensor: The segmentation results, seg_logits from model of each
|
|
input image.
|
|
"""
|
|
assert self.test_cfg.get('mode', 'whole') in ['slide', 'whole'], \
|
|
f'Only "slide" or "whole" test mode are supported, but got ' \
|
|
f'{self.test_cfg["mode"]}.'
|
|
ori_shape = batch_img_metas[0]['ori_shape']
|
|
if not all(_['ori_shape'] == ori_shape for _ in batch_img_metas):
|
|
print_log(
|
|
'Image shapes are different in the batch.',
|
|
logger='current',
|
|
level=logging.WARN)
|
|
if self.test_cfg.mode == 'slide':
|
|
seg_logit = self.slide_inference(inputs, batch_img_metas)
|
|
else:
|
|
seg_logit = self.whole_inference(inputs, batch_img_metas)
|
|
|
|
return seg_logit
|
|
|
|
def aug_test(self, inputs, batch_img_metas, rescale=True):
|
|
"""Test with augmentations.
|
|
|
|
Only rescale=True is supported.
|
|
"""
|
|
# aug_test rescale all imgs back to ori_shape for now
|
|
assert rescale
|
|
# to save memory, we get augmented seg logit inplace
|
|
seg_logit = self.inference(inputs[0], batch_img_metas[0], rescale)
|
|
for i in range(1, len(inputs)):
|
|
cur_seg_logit = self.inference(inputs[i], batch_img_metas[i],
|
|
rescale)
|
|
seg_logit += cur_seg_logit
|
|
seg_logit /= len(inputs)
|
|
seg_pred = seg_logit.argmax(dim=1)
|
|
# unravel batch dim
|
|
seg_pred = list(seg_pred)
|
|
return seg_pred
|