+
+ +
+

nn-UNet

+
+

Overview

+

nn-UNet is a UNet based library designed to segment medical images, refer to +github and the following citation:

+
    +

  • Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2021). nnU-Net: a self-configuring +method for deep learning-based biomedical image segmentation. Nature methods, 18(2), 203-211.

    • +
+

nn-UNet is easiest to use with their command line interface with three commands nnUNetv2_plan_and_preprocess, +nnUNetv2_train and nnUNetv2_predict.

+

For cvpl_tools, cvpl_tools/nnunet/cli.py provides two +wrapper command line interface commands train and predict that simplify the three commands into +two and hides unused parameters for SPIMquant workflow.

+

cvpl_tools/nnunet needs torch library and pip install nnunetv2. GPU is automatically used when +nnUNetv2_train and nnUNetv2_predict are called directly or indirectly through train and +predict and when you have a GPU available on the computer.

+

For those unfamiliar, nn-UNet has the following quirks:

+
    +

  • Residual encoder is available for nnunetv2 but we prefer without it since it costs more to train

    • +

  • Due to limited training data, 2d instead of 3d_fullres mode is used in cvpl_tools

    • +

  • It trains on images pairs of input size (C, Y, X) and output size (Y, X) where C is number of color channels +(1 in our case), and Y, X are spatial coordinates; specifically, N pairs of images will be provided as training +set and a 80%-20% split will be done for train-validation split which is automatically done by nnUNet. It should +be noted in our case we draw Z images from a single scan volume (C, Z, Y, X), so a random split will have +training set distribution correlated with validation set generated by nnUNet, but such thing is hard to avoid

    • +

  • The algorithm is not scale-invariant, meaning during prediction, if we zoom the input image by a factor of 2x or +0.5x we get much worse output results. For best results, use the same input/output image sizes as the training +phase. In our mousebrain lightsheet dataset, we downsample the original >200GB dataset by a factor of (4, 8, 8) +before running the nnUNet for training or prediction.

    • +

  • The algorithm supports the following epochs, useful for small-scale training in our case: +link +if you input number of epochs not listed in this page to the predict command, an error will occur

    • +

  • nn-UNet supports 5-fold ensemble, which is to run nnUNetv2_train command 5 times each on a different +80%-20% split to obtain 5 models to ensemble the prediction. This does not require rerun nnUNetv2_plan_and_preprocess +and is supported by the --fold argument of cvpl_tools’ train command so +you don’t need to run it 5 times. If you finish training all folds, you may use the --fold argument of +cvpl_tools’ predict command to specify all for better accuracy after ensemble or +0 to specify using the first fold trained for comparison.

    • +

  • Running the nn-UNet’s command nnUNetv2_train or cvpl_tools’ train generates one +nnUNet_results folder, which contains a model (of size a few hundred MBs) and a folder of results +including a loss/DICE graph and a log file containing training losses per epoch and per class. The +same model file is used later for prediction.

    • +
+
+
+

Negative Masking for Mouse-brain Lightsheet

+

In this section, we focus primarily on the usage of nn-UNet within cvpl_tools. This part of the +library is designed with handling mouse-brain lightsheet scans in mind. These scans are large (>200GB) +volumes of scans in the format of 4d arrays of data type np.uint16 which is of shape (C, Z, Y, X). An +example is in the google storage bucket +“gcs://khanlab-lightsheet/data/mouse_appmaptapoe/bids/sub-F4A1Te3/micr/sub-F4A1Te3_sample-brain_acq-blaze4x_SPIM.ome.zarr” +with an image shape of (3, 1610, 9653, 9634).

+

The objective of our algorithm is to quantify the locations and sizes of beta-amyloid plaques in a volume +of lightsheet scan like the above, which appear as small-sized round-shaped bright spots in the image +volume, and can be detected using a simple thresholding method.

+

Problem comes, however, since the scanned mouse brain edges areas are as bright as the plaques, they +will be marked as false positives. These edges are relatively easier to detect by a UNet algorithm, which +results in the following segmentation workflow we use:

+
    +

  1. For N mousebrain scans M1, …, MN we have at hand, apply bias correction to smooth out within image brightness +difference caused by imaging artifacts

    2. +

  2. Then select one of N scans, say M1

    3. +
+
    +

  1. Downsample M1 and use a GUI to paint a binary mask, which contains 1 on regions of edges and 0 on plaques and +elsewhere

    2. +

  2. Split the M1 volume and its binary mask annotation vertically to Z slices, and train an nnUNet model on these slices

    3. +

  3. Above produces a model that can predict negative masks on any mousebrain scans of the same format; for the rest N-1 +mouse brains, they are down-sampled and we use this model to predict on them to obtain their corresponding negative +masks

    4. +

  4. These masks are used to remove edge areas of the image before we apply thresholding to find plaque objects. +Algorithmically, we compute M’ where M'[z, y, x] = M[z, y, x] * (1 - NEG_MASK[z, y, x]) for each +voxel location (z, y, x); then, we apply threshold on M’ and take connected component of value of 1 as individual +plaque objects; their centroid locations and sizes (in number of voxels) are summarized in a numpy table and +reported

    5. +
+

In this next part, we discuss the annotation part 2, training part 3 and prediction part 4.

+
+
+

TODO

+
+
+ + +
+