|
from torch import nn |
|
|
|
from mmseg.core import add_prefix |
|
from mmseg.ops import resize |
|
from .. import builder |
|
from ..builder import SEGMENTORS |
|
from .encoder_decoder import EncoderDecoder |
|
|
|
|
|
@SEGMENTORS.register_module() |
|
class CascadeEncoderDecoder(EncoderDecoder): |
|
"""Cascade Encoder Decoder segmentors. |
|
|
|
CascadeEncoderDecoder almost the same as EncoderDecoder, while decoders of |
|
CascadeEncoderDecoder are cascaded. The output of previous decoder_head |
|
will be the input of next decoder_head. |
|
""" |
|
|
|
def __init__(self, |
|
num_stages, |
|
backbone, |
|
decode_head, |
|
neck=None, |
|
auxiliary_head=None, |
|
train_cfg=None, |
|
test_cfg=None, |
|
pretrained=None): |
|
self.num_stages = num_stages |
|
super(CascadeEncoderDecoder, self).__init__( |
|
backbone=backbone, |
|
decode_head=decode_head, |
|
neck=neck, |
|
auxiliary_head=auxiliary_head, |
|
train_cfg=train_cfg, |
|
test_cfg=test_cfg, |
|
pretrained=pretrained) |
|
|
|
def _init_decode_head(self, decode_head): |
|
"""Initialize ``decode_head``""" |
|
assert isinstance(decode_head, list) |
|
assert len(decode_head) == self.num_stages |
|
self.decode_head = nn.ModuleList() |
|
for i in range(self.num_stages): |
|
self.decode_head.append(builder.build_head(decode_head[i])) |
|
self.align_corners = self.decode_head[-1].align_corners |
|
self.num_classes = self.decode_head[-1].num_classes |
|
|
|
def init_weights(self, pretrained=None): |
|
"""Initialize the weights in backbone and heads. |
|
|
|
Args: |
|
pretrained (str, optional): Path to pre-trained weights. |
|
Defaults to None. |
|
""" |
|
self.backbone.init_weights(pretrained=pretrained) |
|
for i in range(self.num_stages): |
|
self.decode_head[i].init_weights() |
|
if self.with_auxiliary_head: |
|
if isinstance(self.auxiliary_head, nn.ModuleList): |
|
for aux_head in self.auxiliary_head: |
|
aux_head.init_weights() |
|
else: |
|
self.auxiliary_head.init_weights() |
|
|
|
def encode_decode(self, img, img_metas): |
|
"""Encode images with backbone and decode into a semantic segmentation |
|
map of the same size as input.""" |
|
x = self.extract_feat(img) |
|
out = self.decode_head[0].forward_test(x, img_metas, self.test_cfg) |
|
for i in range(1, self.num_stages): |
|
out = self.decode_head[i].forward_test(x, out, img_metas, |
|
self.test_cfg) |
|
out = resize( |
|
input=out, |
|
size=img.shape[2:], |
|
mode='bilinear', |
|
align_corners=self.align_corners) |
|
return out |
|
|
|
def _decode_head_forward_train(self, x, img_metas, gt_semantic_seg): |
|
"""Run forward function and calculate loss for decode head in |
|
training.""" |
|
losses = dict() |
|
|
|
loss_decode = self.decode_head[0].forward_train( |
|
x, img_metas, gt_semantic_seg, self.train_cfg) |
|
|
|
losses.update(add_prefix(loss_decode, 'decode_0')) |
|
|
|
for i in range(1, self.num_stages): |
|
|
|
prev_outputs = self.decode_head[i - 1].forward_test( |
|
x, img_metas, self.test_cfg) |
|
loss_decode = self.decode_head[i].forward_train( |
|
x, prev_outputs, img_metas, gt_semantic_seg, self.train_cfg) |
|
losses.update(add_prefix(loss_decode, f'decode_{i}')) |
|
|
|
return losses |
|
|
