def __init__(
self,
encoder: EncoderModule,
unet_channels: Union[int, List[int]],
num_classes: int = 1,
dropout=0.25,
full_size_mask=True,
activation=ACT_RELU,
upsample_block: Union[
Type[nn.Upsample],
Type[ResidualDeconvolutionUpsample2d]] = nn.UpsamplingNearest2d,
need_supervision_masks=False,
last_upsample_block=None,
):
super().__init__()
self.encoder = encoder
abn_block = partial(ABN, activation=activation)
self.decoder = UNetDecoder(
feature_maps=encoder.channels,
decoder_features=unet_channels,
unet_block=partial(UnetBlock, abn_block=abn_block),
upsample_block=upsample_block,
)
if last_upsample_block is not None:
self.last_upsample_block = last_upsample_block(unet_channels[0])
self.mask = nn.Sequential(
OrderedDict([
("drop", nn.Dropout2d(dropout)),
("conv",
conv1x1(self.last_upsample_block.out_channels,
num_classes)),
]))
else:
self.last_upsample_block = None
self.mask = nn.Sequential(
OrderedDict([("drop", nn.Dropout2d(dropout)),
("conv", conv1x1(unet_channels[0],
num_classes))]))
if need_supervision_masks:
self.supervision = nn.ModuleList([
conv1x1(channels, num_classes)
for channels in self.decoder.channels
])
self.supervision_names = [
output_mask_name_for_stride(stride)
for stride in self.encoder.strides
]
else:
self.supervision = None
self.supervision_names = None
self.full_size_mask = full_size_mask
def __init__(
self,
encoder: EncoderModule,
num_classes: int,
dropout=0.25,
fpn_channels=256,
abn_block=ABN,
activation=ACT_RELU,
full_size_mask=True,
):
super().__init__()
self.encoder = encoder
abn_block = partial(abn_block, activation=activation)
self.decoder = FPNCatDecoder(
encoder.channels,
context_block=partial(FPNContextBlock, abn_block=abn_block),
bottleneck_block=partial(FPNBottleneckBlock, abn_block=abn_block),
fpn_channels=fpn_channels,
)
self.fuse = FPNFuse()
self.mask = nn.Sequential(
OrderedDict([("drop", nn.Dropout2d(dropout)), ("conv", conv1x1(sum(self.decoder.channels), num_classes))])
)
self.full_size_mask = full_size_mask
def __init__(
self,
encoder: EncoderModule,
unet_channels: Union[int, List[int]],
num_classes: int = 1,
dropout=0.25,
full_size_mask=True,
activation=ACT_RELU,
):
super().__init__()
self.encoder = encoder
abn_block = partial(ABN, activation=activation)
self.decoder = UNetDecoder(
feature_maps=encoder.channels,
decoder_features=unet_channels,
unet_block=partial(UnetBlock, abn_block=abn_block),
upsample_block=nn.UpsamplingNearest2d,
)
self.mask = nn.Sequential(
OrderedDict([("drop", nn.Dropout2d(dropout)), ("conv", conv1x1(unet_channels[0], num_classes))])
)
self.full_size_mask = full_size_mask
def __init__(
self, encoder: EncoderModule, num_classes: int, dropout=0.25, full_size_mask=True, fpn_channels=256,
):
super().__init__()
self.encoder = encoder
self.decoder = FPNSumDecoder(feature_maps=encoder.output_filters, fpn_channels=fpn_channels,)
self.mask = nn.Sequential(
OrderedDict([("drop", nn.Dropout2d(dropout)), ("conv", conv1x1(fpn_channels, num_classes))])
)
self.full_size_mask = full_size_mask
def __init__(self,
encoder: E.EncoderModule,
features=256,
num_classes: int = 1,
dropout=0.25,
full_size_mask=True):
super().__init__()
self.encoder = encoder
self.decoder = CANDecoder(encoder.channels, out_channels=features)
self.mask = nn.Sequential(
OrderedDict([("drop", nn.Dropout2d(dropout)),
("conv", conv1x1(features, num_classes))]))
self.full_size_mask = full_size_mask
def __init__(
self,
encoder: EncoderModule,
unet_channels: Union[int, List[int]],
num_classes: int = 1,
dropout=0.25,
full_size_mask=True,
activation=Swish,
):
super().__init__()
self.encoder = encoder
self.decoder = EfficientUNetDecoder(
feature_maps=encoder.channels, decoder_features=unet_channels, activation=activation
)
self.mask = nn.Sequential(
OrderedDict([("drop", nn.Dropout2d(dropout)), ("conv", conv1x1(self.decoder.channels[0], num_classes))])
)
self.full_size_mask = full_size_mask