Spaces:
Runtime error
Runtime error
File size: 3,843 Bytes
3b96cb1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
# Copyright (c) OpenCD. All rights reserved.
import torch
import torch.nn as nn
from mmseg.models.decode_heads.decode_head import BaseDecodeHead
from mmseg.models.losses import accuracy
from mmseg.models.utils import resize
from opencd.registry import MODELS
@MODELS.register_module()
class IdentityHead(BaseDecodeHead):
"""Identity Head."""
def __init__(self, **kwargs):
super().__init__(channels=1, **kwargs)
delattr(self, 'conv_seg')
def init_weights(self):
pass
def _forward_feature(self, inputs):
"""
Args:
inputs (list[Tensor]): List of multi-level img features.
Returns:
feats (Tensor): A tensor of shape (batch_size, self.channels,
H, W) which is feature map for last layer of decoder head.
"""
x = self._transform_inputs(inputs)
return x
def forward(self, inputs):
"""Forward function."""
output = self._forward_feature(inputs)
return output
@MODELS.register_module()
class DSIdentityHead(BaseDecodeHead):
"""Deep Supervision Identity Head."""
def __init__(self, **kwargs):
super().__init__(channels=1, **kwargs)
delattr(self, 'conv_seg')
def init_weights(self):
pass
def _forward_feature(self, inputs):
"""
Args:
inputs (list[Tensor]): List of multi-level img features.
Returns:
feats (Tensor): A tensor of shape (batch_size, self.channels,
H, W) which is feature map for last layer of decoder head.
"""
x = self._transform_inputs(inputs)
return x
def forward(self, inputs):
"""Forward function."""
output = self._forward_feature(inputs)
return output
def loss_by_feat(self, seg_logits, batch_data_samples):
"""Compute segmentation loss.
Args:
seg_logits (Tensor): The output from decode head forward function.
batch_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
"""
seg_label = self._stack_batch_gt(batch_data_samples)
loss = dict()
seg_label_size = seg_label.shape[2:]
for seg_idx, single_seg_logit in enumerate(seg_logits):
single_seg_logit = resize(
input=single_seg_logit,
size=seg_label_size,
mode='bilinear',
align_corners=self.align_corners)
if self.sampler is not None:
seg_weight = self.sampler.sample(single_seg_logit, seg_label)
else:
seg_weight = None
seg_label = seg_label.squeeze(1)
if not isinstance(self.loss_decode, nn.ModuleList):
losses_decode = [self.loss_decode]
else:
losses_decode = self.loss_decode
for loss_decode in losses_decode:
loss_name = f'aux_{seg_idx}_' + loss_decode.loss_name
if loss_decode.loss_name not in loss:
loss[loss_name] = loss_decode(
single_seg_logit,
seg_label,
weight=seg_weight,
ignore_index=self.ignore_index)
else:
loss[loss_name] += loss_decode(
single_seg_logit,
seg_label,
weight=seg_weight,
ignore_index=self.ignore_index)
loss['acc_seg'] = accuracy(
single_seg_logit, seg_label, ignore_index=self.ignore_index)
return loss
|