from monai.networks.nets import DynUNet

# 5 levels; the first downsample skips z (see below).
kernels = [[3, 3, 3]] * 5
strides = [[1, 1, 1], [1, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]]

model = DynUNet(
    spatial_dims=3,
    in_channels=3,          # CT + line + endpoint blobs
    out_channels=1,         # one probability per voxel
    kernel_size=kernels,
    strides=strides,
    upsample_kernel_size=strides[1:],
    filters=features,
    res_block=True,
)