Spaces:
Sleeping
Sleeping
File size: 4,787 Bytes
b5cbaa6 8e6512c b5cbaa6 8e6512c b5cbaa6 8e6512c b5cbaa6 8e6512c b5cbaa6 8e6512c b5cbaa6 8e6512c b5cbaa6 8e6512c b5cbaa6 |
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 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 |
import os
import sys
import mock
import pytest
import torch
# mock detection module
sys.modules["torchvision._C"] = mock.Mock()
import segmentation_models_pytorch as smp
IS_TRAVIS = os.environ.get("TRAVIS", False)
def get_encoders():
travis_exclude_encoders = [
"senet154",
"resnext101_32x16d",
"resnext101_32x32d",
"resnext101_32x48d",
]
encoders = smp.encoders.get_encoder_names()
if IS_TRAVIS:
encoders = [e for e in encoders if e not in travis_exclude_encoders]
return encoders
ENCODERS = get_encoders()
DEFAULT_ENCODER = "resnet18"
def get_sample(model_class):
if model_class in [
smp.Unet,
smp.Linknet,
smp.FPN,
smp.PSPNet,
smp.UnetPlusPlus,
smp.MAnet,
]:
sample = torch.ones([1, 3, 64, 64])
elif model_class == smp.PAN:
sample = torch.ones([2, 3, 256, 256])
elif model_class == smp.DeepLabV3:
sample = torch.ones([2, 3, 128, 128])
else:
raise ValueError("Not supported model class {}".format(model_class))
return sample
def _test_forward(model, sample, test_shape=False):
with torch.no_grad():
out = model(sample)
if test_shape:
assert out.shape[2:] == sample.shape[2:]
def _test_forward_backward(model, sample, test_shape=False):
out = model(sample)
out.mean().backward()
if test_shape:
assert out.shape[2:] == sample.shape[2:]
@pytest.mark.parametrize("encoder_name", ENCODERS)
@pytest.mark.parametrize("encoder_depth", [3, 5])
@pytest.mark.parametrize(
"model_class", [smp.FPN, smp.PSPNet, smp.Linknet, smp.Unet, smp.UnetPlusPlus]
)
def test_forward(model_class, encoder_name, encoder_depth, **kwargs):
if (
model_class is smp.Unet
or model_class is smp.UnetPlusPlus
or model_class is smp.MAnet
):
kwargs["decoder_channels"] = (16, 16, 16, 16, 16)[-encoder_depth:]
model = model_class(
encoder_name, encoder_depth=encoder_depth, encoder_weights=None, **kwargs
)
sample = get_sample(model_class)
model.eval()
if encoder_depth == 5 and model_class != smp.PSPNet:
test_shape = True
else:
test_shape = False
_test_forward(model, sample, test_shape)
@pytest.mark.parametrize(
"model_class",
[
smp.PAN,
smp.FPN,
smp.PSPNet,
smp.Linknet,
smp.Unet,
smp.UnetPlusPlus,
smp.MAnet,
smp.DeepLabV3,
],
)
def test_forward_backward(model_class):
sample = get_sample(model_class)
model = model_class(DEFAULT_ENCODER, encoder_weights=None)
_test_forward_backward(model, sample)
@pytest.mark.parametrize(
"model_class",
[smp.PAN, smp.FPN, smp.PSPNet, smp.Linknet, smp.Unet, smp.UnetPlusPlus, smp.MAnet],
)
def test_aux_output(model_class):
model = model_class(
DEFAULT_ENCODER, encoder_weights=None, aux_params=dict(classes=2)
)
sample = get_sample(model_class)
label_size = (sample.shape[0], 2)
mask, label = model(sample)
assert label.size() == label_size
@pytest.mark.parametrize("upsampling", [2, 4, 8])
@pytest.mark.parametrize("model_class", [smp.FPN, smp.PSPNet])
def test_upsample(model_class, upsampling):
default_upsampling = 4 if model_class is smp.FPN else 8
model = model_class(DEFAULT_ENCODER, encoder_weights=None, upsampling=upsampling)
sample = get_sample(model_class)
mask = model(sample)
assert mask.size()[-1] / 64 == upsampling / default_upsampling
@pytest.mark.parametrize("model_class", [smp.FPN])
@pytest.mark.parametrize("encoder_name", ENCODERS)
@pytest.mark.parametrize("in_channels", [1, 2, 4])
def test_in_channels(model_class, encoder_name, in_channels):
sample = torch.ones([1, in_channels, 64, 64])
model = model_class(DEFAULT_ENCODER, encoder_weights=None, in_channels=in_channels)
model.eval()
with torch.no_grad():
model(sample)
assert model.encoder._in_channels == in_channels
@pytest.mark.parametrize("encoder_name", ENCODERS)
def test_dilation(encoder_name):
if (
encoder_name in ["inceptionresnetv2", "xception", "inceptionv4"]
or encoder_name.startswith("vgg")
or encoder_name.startswith("densenet")
or encoder_name.startswith("timm-res")
):
return
encoder = smp.encoders.get_encoder(encoder_name)
encoder.make_dilated(
stage_list=[5],
dilation_list=[2],
)
encoder.eval()
with torch.no_grad():
sample = torch.ones([1, 3, 64, 64])
output = encoder(sample)
shapes = [out.shape[-1] for out in output]
assert shapes == [64, 32, 16, 8, 4, 4] # last downsampling replaced with dilation
if __name__ == "__main__":
pytest.main([__file__])
|