def resnet34_cls(input_channels=6,
num_classes=5,
dropout=0.0,
pretrained=True):
encoder = E.Resnet34Encoder(pretrained=pretrained, layers=[1, 2, 3, 4])
return DamageTypeClassificationModel(encoder,
damage_type_classes=len(DAMAGE_TYPES),
dropout=dropout)
def fpn128_resnet34(num_classes=1, num_channels=3):
assert num_channels == 3
encoder = E.Resnet34Encoder()
decoder = D.FPNDecoder(features=encoder.output_filters,
prediction_block=DoubleConvRelu,
bottleneck=FPNBottleneckBlockBN,
fpn_features=128)
return SegmentationModel(encoder, decoder, num_classes)
def resnet34_unet32(input_channels=3,
num_classes=1,
dropout=0.0,
pretrained=True):
encoder = E.Resnet34Encoder(pretrained=pretrained, layers=[0, 1, 2, 3, 4])
if input_channels != 3:
encoder.change_input_channels(input_channels)
return UnetSegmentationModel(encoder,
num_classes=num_classes,
unet_channels=[32, 64, 128, 256],
dropout=dropout)
def resnet34_fpncatv2_256(num_classes=5,
dropout=0.0,
pretrained=True,
classifiers=True):
encoder = E.Resnet34Encoder(pretrained=pretrained)
return FPNCatSegmentationModelV2(
encoder,
num_classes=num_classes,
disaster_type_classes=len(DISASTER_TYPES) if classifiers else None,
damage_type_classes=len(DAMAGE_TYPES) if classifiers else None,
fpn_channels=256,
dropout=dropout,
abn_block=partial(ABN, activation=ACT_RELU),
)
def resnet34_unet64_s4(input_channels=3,
num_classes=1,
dropout=0.2,
pretrained=True):
encoder = E.Resnet34Encoder(pretrained=pretrained, layers=[1, 2, 3, 4])
if input_channels != 3:
encoder.change_input_channels(input_channels)
return UnetSegmentationModel(
encoder,
num_classes=num_classes,
unet_channels=[64, 128, 256],
activation=ACT_SWISH,
dropout=dropout,
)
def resnet34_unet_v2(input_channels=6,
num_classes=5,
dropout=0.0,
pretrained=True,
classifiers=True):
encoder = E.Resnet34Encoder(pretrained=pretrained, layers=[0, 1, 2, 3, 4])
return UnetV2SegmentationModel(
encoder,
num_classes=num_classes,
disaster_type_classes=len(DISASTER_TYPES) if classifiers else None,
damage_type_classes=len(DAMAGE_TYPES) if classifiers else None,
unet_channels=[64, 128, 256, 256],
dropout=dropout,
abn_block=partial(ABN, activation=ACT_RELU),
)
def __init__(self, num_classes, in_channels=3):
super().__init__()
model = encoders.Resnet34Encoder()
self.conv1_ = model.layer0[0]
# if in_channels == 3 else nn.Conv2d(in_channels,
# 64, 7, 2, 3)
self.bn1 = model.layer0[1]
self.act = nn.ReLU(inplace=True)
self.pool = nn.MaxPool2d(2, 2)
self.enc1 = model.layer1 # 64
self.enc2 = model.layer2 # 128
self.enc3 = model.layer3 # 256
self.enc4 = model.layer4 # 512
self.up = nn.Upsample(scale_factor=2,
mode='bilinear',
align_corners=True)
self.dec3 = DecoderBlockGated(512, 256, 256)
self.conv3 = Conv2dBn(512 + 256, 256, 3, 1, 1) # 256
self.dec2 = DecoderBlockGated(256, 128, 128)
self.conv2 = Conv2dBn(256 + 128, 128, 3, 1, 1) # 128
self.dec1 = DecoderBlockGated(128, 64, 64)
self.conv1 = Conv2dBn(128 + 64, 64, 3, 1, 1) # 64
self.full_conv = DecoderBlockGated(64, 64, 32)
self.convbn3 = Conv2dBn(32 + 64, 16, 3, 1, 1)
self.out = nn.Sequential(
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
nn.Conv2d(16, num_classes, 2, stride=2, bias=False))
for (name, m) in self.named_modules():
if 'enc' not in name and 'conv1_' not in name:
if isinstance(m, nn.Conv2d) or isinstance(
m, nn.ConvTranspose2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2.0 / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
def resnet34_fpn(num_classes=1, fpn_features=128):
encoder = E.Resnet34Encoder()
return FPNSegmentationModel(encoder, num_classes, fpn_features)
def resnet34_fpncat128(num_classes=5, dropout=0.0, pretrained=True):
encoder = E.Resnet34Encoder(pretrained=pretrained)
return FPNCatSegmentationModel(encoder, num_classes=num_classes, fpn_channels=128, dropout=dropout)
def resnet34_deeplab128(num_classes=1, dropout=0.0):
encoder = E.Resnet34Encoder()
return DeeplabV3SegmentationModel(encoder,
num_classes=num_classes,
high_level_bottleneck=128,
dropout=dropout)
