torch-em

# Train a 2d U-Net for foreground and boundary segmentation of nuclei, using data from# https://github.com/mpicbg-csbd/stardist/releases/download/0.1.0/dsb2018.zipimporttorch_emfromtorch_em.modelimportUNet2dfromtorch_em.data.datasetsimportget_dsb_loadermodel=UNet2d(in_channels=1,out_channels=2)# Transform to convert from instance segmentation labels to foreground and boundary probabilties.label_transform=torch_em.transform.BoundaryTransform(add_binary_target=True,ndim=2)# Create the training and validation data loader.data_path="./dsb"# The training data will be downloaded and saved here.train_loader=get_dsb_loader(data_path,patch_shape=(1,256,256),batch_size=8,split="train",download=True,label_transform=label_transform,)val_loader=get_dsb_loader(data_path,patch_shape=(1,256,256),batch_size=8,split="test",label_transform=label_transform,)# The trainer handles the details of the training process.# It will save checkpoints in "checkpoints/dsb-boundary-model"# and the tensorboard logs in "logs/dsb-boundary-model".trainer=torch_em.default_segmentation_trainer(name="dsb-boundary-model",model=model,train_loader=train_loader,val_loader=val_loader,learning_rate=1e-4,)trainer.fit(iterations=5000)# Fit for 5000 iterations.# Export the trained model to the bioimage.io model format.fromglobimportglobimportimageiofromtorch_em.utilimportexport_bioimageio_model# Load one of the images to use as reference image.# Crop it to a shape that is guaranteed to fit the network.test_im=imageio.imread(glob(f"{data_path}/test/images/*.tif")[0])[:256, :256]# Export the model.export_bioimageio_model("./checkpoints/dsb-boundary-model","./bioimageio-model",test_im)