def get_model(large_model):
n_out = 2
if large_model:
print("Using large model")
model = AnisotropicUNet(
scale_factors=[
[1, 2, 2],
[1, 2, 2],
[2, 2, 2],
[2, 2, 2],
[2, 2, 2]
],
in_channels=1,
out_channels=n_out,
initial_features=128,
gain=2,
final_activation="Sigmoid"
)
else:
print("Using vanilla model")
model = AnisotropicUNet(
scale_factors=[
[1, 2, 2],
[1, 2, 2],
[2, 2, 2],
[2, 2, 2]
],
in_channels=1,
out_channels=n_out,
initial_features=64,
gain=2,
final_activation="Sigmoid"
)
return model
def train_shallow2deep(args):
name = f"cremi3d-v{args.version}"
# check if we need to train the rfs for preparation
rf_folder = os.path.join("checkpoints", name, "rfs")
have_rfs = len(glob(os.path.join(rf_folder, "*.pkl"))) == args.n_rfs
if not have_rfs:
prepare_shallow2deep_cremi(args, rf_folder)
assert os.path.exists(rf_folder)
model = AnisotropicUNet(in_channels=1,
out_channels=1,
final_activation="Sigmoid",
scale_factors=[[1, 2, 2], [1, 2, 2], [2, 2, 2],
[2, 2, 2]])
train_loader = get_cremi_loader(args, "train", rf_folder)
val_loader = get_cremi_loader(args, "val", rf_folder)
dice_loss = torch_em.loss.DiceLoss()
trainer = torch_em.default_segmentation_trainer(name,
model,
train_loader,
val_loader,
loss=dice_loss,
metric=dice_loss,
learning_rate=1.0e-4,
device=args.device,
log_image_interval=50)
trainer.fit(args.n_iterations)
def get_model():
model = AnisotropicUNet(scale_factors=4 * [[2, 2, 2]],
in_channels=1,
out_channels=2,
initial_features=32,
gain=2,
final_activation="Sigmoid")
return model
def test_anisotropic_unet(self):
from torch_em.model import AnisotropicUNet
scale_factors = [[1, 2, 2], [1, 2, 2], [2, 2, 2]]
for anisotropic_kernel in (False, True):
net = AnisotropicUNet(1,
1,
scale_factors,
initial_features=4,
anisotropic_kernel=anisotropic_kernel)
self._test_net(net, (1, 1, 8, 32, 32))
def get_model(use_diagonal_offsets):
n_out = len(get_offsets(use_diagonal_offsets))
def my_sampler(scale_factor, inc, outc):
return Upsampler3d(scale_factor, inc, outc, mode='bilinear')
model = AnisotropicUNet(
in_channels=1,
out_channels=n_out,
final_activation='Sigmoid',
pooler_impl=nn.AvgPool2d,
sampler_impl=my_sampler
)
return model
def get_model():
n_out = len(OFFSETS)
model = AnisotropicUNet(
scale_factors=[
[1, 2, 2],
[1, 2, 2],
[2, 2, 2],
[2, 2, 2],
[2, 2, 2]
],
in_channels=1,
out_channels=n_out,
initial_features=32,
gain=2,
final_activation="Sigmoid"
)
return model
def get_model():
# we have 1 channel per affinty offsets and another channel
# for foreground background segmentation
n_out = len(OFFSETS) + 1
# TODO adapt scale factors for data at native resolution
scale_factors = 4 * [[2, 2, 2]]
model = AnisotropicUNet(
in_channels=1,
out_channels=n_out,
scale_factors=scale_factors,
initial_features=64,
gain=2,
final_activation='Sigmoid',
anisotropic_kernel=
False # set anisotropic kernels if scale factors are anisotropic
)
return model