Package: cellseg_models_pytorch
Python library for 2D cell/nuclei instance segmentation models written withPyTorch.Introductioncellseg-models.pytorchis a library built uponPyTorchthat contains multi-task encoder-decoder architectures along with dedicated post-processing methods for segmenting cell/nuclei instances. As the name might suggest, this library is heavily inspired bysegmentation_models.pytorchlibrary for semantic segmentation.FeaturesHigh level API to define cell/nuclei instance segmentation models.6 cell/nuclei instance segmentation models and more to come.Open source datasets for training and benchmarking.Pre-trained backbones/encoders from thetimmlibrary.Pre-trained transformer backbones likeDinoV2andSAM.All the architectures can be augmented topanoptic segmentation.Flexibility to modify the components of the model architectures.Sliding window inference for large images.Multi-GPU inference.Popular training losses and benchmarking metrics.Benchmarking utilities both for model latency & segmentation performance.Regularization techniques to tackle batch effects/domain shifts such asStrong Augment,Spectral decoupling,Label smoothing.Ability to add transformers to the decoder layers.Example notebooks to train models withlightningoraccelerate.InstallationBasic installationpipinstallcellseg-models-pytorchTo install extra dependencies (training utilities and datamodules for open-source datasets) usepipinstallcellseg-models-pytorch[all]ModelsModelPaper[1] HoVer-Nethttps://www.sciencedirect.com/science/article/pii/S1361841519301045?via%3Dihub[2] Cellposehttps://www.nature.com/articles/s41592-020-01018-x[3] Omniposehttps://www.biorxiv.org/content/10.1101/2021.11.03.467199v2[4] Stardisthttps://arxiv.org/abs/1806.03535[5] CellVit-SAMhttps://arxiv.org/abs/2306.15350.03535[6] CPP-Nethttps://arxiv.org/abs/2102.0686703535DatasetsDatasetPaper[7, 8] Pannukehttps://arxiv.org/abs/2003.10778,https://link.springer.com/chapter/10.1007/978-3-030-23937-4_2[9] Lizardhttp://arxiv.org/abs/2108.11195Notebook examplesTraining Hover-Net with PannukeTraining Hover-Net with Pannuke. Here we train theHover-Netnuclei segmentation model with animagenetpretrainedresnet50backbone from thetimmlibrary. The Pannuke dataset (fold 1 & fold 2) are used for training data and the fold 3 is used as validation data. The model is trained by utilizinglightning.Training Stardist with PannukeTraining Stardist with Pannuke. Here we train theStardistmulti-class nuclei segmentation model with animagenetpretrainedefficientnetv2_sbackbone from thetimmlibrary. The Pannuke dataset (fold 1 & fold 2) are used for training data and the fold 3 is used as validation data. The model is trained by utilizinglightning.Training CellPose with PannukeTraining CellPose with Pannuke. Here we train theCellPosemulti-class nuclei segmentation model with animagenetpretrainedconvnext_smallbackbone from thetimmlibrary. The Pannuke dataset (fold 1 & fold 2) are used for training data and the fold 3 is used as validation data. The model is trained (with checkpointing) by utilizingaccelerateby hugginface.Training OmniPose with PannukeTraining OmniPose with Pannuke. Here we train the OmniPose multi-class nuclei segmentation model with animagenetpretrainedfocalnet_small_lrfbackbone from thetimmlibrary. The Pannuke dataset (fold 1 & fold 2) are used for training data and the fold 3 is used as validation data. The model is trained (with checkpointing) by utilizingaccelerateby hugginface.Training CPP-Net with PannukeTraining CPP-Net with Pannuke. Here we train the CPP-Net multi-class nuclei segmentation model with animagenetpretrainedefficientnetv2_sbackbone from thetimmlibrary. The Pannuke dataset (fold 1 & fold 2) are used for training data and the fold 3 is used as validation data. The model is trained by utilizinglightning.Finetuning CellPose with DINOv2 backboneFinetuning CellPose with DINOv2 backbone for Pannuke. Here we finetune the CellPose multi-class nuclei segmentation model with aLVD-142MpretrainedDINOv2backbone. The Pannuke dataset (fold 1 & fold 2) are used for training data and the fold 3 is used as validation data. The model is trained (with checkpointing) by utilizinglightning.Finetuning CellVit-SAM with PannukeFinetuning CellVit-SAM with Pannuke. Here we finetune the CellVit-SAM multi-class nuclei segmentation model with aSA-1BpretrainedSAM-image-encoder backbone (checkoutSAM). The encoder is transformer basedVitDet-model. The Pannuke dataset (fold 1 & fold 2) are used for training data and the fold 3 is used as validation data. The model is trained (with checkpointing) by utilizingaccelerateby hugginface.Benchmarking Cellpose Trained on PannukeBenchmarking Cellpose Trained on Pannuke. Here we run benchmarking forCellposethat was trained on Pannuke. Both the model performance and latency benchmarking are covered.Training CellPose with LizardTraining CellPose with Lizard. Here we train theCellposemodel with Lizard dataset that is composed of varying sized images. This example is old and might not be up to date.Code ExamplesDefine Cellpose for cell segmentation.importcellseg_models_pytorchascsmpimporttorchmodel=csmp.models.cellpose_base(type_classes=5)x=torch.rand([1,3,256,256])# NOTE: the outputs still need post-processing.y=model(x)# {"cellpose": [1, 2, 256, 256], "type": [1, 5, 256, 256]}Define Cellpose for cell and tissue area segmentation (Panoptic segmentation).importcellseg_models_pytorchascsmpimporttorchmodel=csmp.models.cellpose_plus(type_classes=5,sem_classes=3)x=torch.rand([1,3,256,256])# NOTE: the outputs still need post-processing.y=model(x)# {"cellpose": [1, 2, 256, 256], "type": [1, 5, 256, 256], "sem": [1, 3, 256, 256]}Define panoptic Cellpose model with more flexibility.importcellseg_models_pytorchascsmp# the model will include two decoder branches.decoders=("cellpose","sem")# and in total three segmentation heads emerging from the decoders.heads={"cellpose":{"cellpose":2,"type":5},"sem":{"sem":3}}model=csmp.CellPoseUnet(decoders=decoders,# cellpose and semantic decodersheads=heads,# three output headsdepth=5,# encoder depthout_channels=(256,128,64,32,16),# num out channels at each decoder stagelayer_depths=(4,4,4,4,4),# num of conv blocks at each decoder layerstyle_channels=256,# num of style vector channelsenc_name="resnet50",# timm encoderenc_pretrain=True,# imagenet pretrained encoderlong_skip="unetpp",# unet++ long skips ("unet", "unetpp", "unet3p")merge_policy="sum",# concatenate long skips ("cat", "sum")short_skip="residual",# residual short skips ("basic", "residual", "dense")normalization="bcn",# batch-channel-normalization.activation="gelu",# gelu activation.convolution="wsconv",# weight standardized conv.attention="se",# squeeze-and-excitation attention.pre_activate=False,# normalize and activation after convolution.)x=torch.rand([1,3,256,256])# NOTE: the outputs still need post-processing.y=model(x)# {"cellpose": [1, 2, 256, 256], "type": [1, 5, 256, 256], "sem": [1, 3, 256, 256]}Run HoVer-Net inference and post-processing with a sliding window approach.importcellseg_models_pytorchascsmp# define the modelmodel=csmp.models.hovernet_base(type_classes=5)# define the final activations for each model outputout_activations={"hovernet":"tanh","type":"softmax","inst":"softmax"}# define whether to weight down the predictions at the image boundaries# typically, models perform the poorest at the image boundaries and with# overlapping patches this causes issues which can be overcome by down-# weighting the prediction boundariesout_boundary_weights={"hovernet":True,"type":False,"inst":False}# define the infererinferer=csmp.inference.SlidingWindowInferer(model=model,input_folder="/path/to/images/",checkpoint_path="/path/to/model/weights/",out_activations=out_activations,out_boundary_weights=out_boundary_weights,instance_postproc="hovernet",# THE POST-PROCESSING METHODnormalization="percentile",# same normalization as in trainingpatch_size=(256,256),stride=128,padding=80,batch_size=8,)inferer.infer()inferer.out_masks# {"image1" :{"inst": [H, W], "type": [H, W]}, ..., "imageN" :{"inst": [H, W], "type": [H, W]}}Models APIGenerally, the model building API enables the effortless creation of hard-parameter sharing multi-task encoder-decoder CNN architectures. The general architectural schema is illustrated in the below image.Class APIThe class API enables the most flexibility in defining different model architectures. It borrows a lot fromsegmentation_models.pytorchmodels API.Model classes:csmp.CellPoseUnetcsmp.StarDistUnetcsmp.HoverNetcsmp.CellVitSAMAll of the models contain:model.encoder- pretrainedtimmbackbone for feature extraction.model.{decoder_name}_decoder- Models can have multiple decoders with unique names.model.{head_name}_seg_head- Model decoders can have multiple segmentation heads with unique names.model.forward(x)- forward pass.model.forward_features(x)- forward pass of the encoder and decoders. Returns enc and dec featuresDefining your own multi-task architectureFor example, to define a multi-task architecture that hasresnet50encoder, four decoders, and 5 output heads withCellPoseUnetarchitectural components, we could do this:importcellseg_models_pytorchascsmpimporttorchmodel=csmp.CellPoseUnet(decoders=("cellpose","dist","contour","sem"),heads={"cellpose":{"type":5,"cellpose":2},"dist":{"dist":1},"contour":{"contour":1},"sem":{"sem":4}},)x=torch.rand([1,3,256,256])model(x)# {#   "cellpose": [1, 2, 256, 256],#   "type": [1, 5, 256, 256],#   "dist": [1, 1, 256, 256],#   "contour": [1, 1, 256, 256],#   "sem": [1, 4, 256, 256]# }Function APIWith the function API, you can build models with low effort by calling the below listed functions. Under the hood, the function API simply calls the above classes with pre-defined decoder and head names. The training and post-processing tools of this library are built around these names, thus, it is recommended to use the function API, although, it is a bit more rigid than the class API. Basically, the function API only lacks the ability to define the output-tasks of the model, but allows for all the rest as the class API.Model functionsOutput namesTaskcsmp.models.cellpose_base"type","cellpose",instance segmentationcsmp.models.cellpose_plus"type","cellpose","sem",panoptic segmentationcsmp.models.omnipose_base"type","omnipose"instance segmentationcsmp.models.omnipose_plus"type","omnipose","sem",panoptic segmentationcsmp.models.hovernet_base"type","inst","hovernet"instance segmentationcsmp.models.hovernet_plus"type","inst","hovernet","sem"panoptic segmentationcsmp.models.hovernet_small"type","hovernet"instance segmentationcsmp.models.hovernet_small_plus"type","hovernet","sem"panoptic segmentationcsmp.models.stardist_base"stardist","dist"binary instance segmentationcsmp.models.stardist_base_multiclass"stardist","dist","type"instance segmentationcsmp.models.stardist_plus"stardist","dist","type","sem"panoptic segmentationcsmp.models.cppnet_base"stardist_refined","dist"binary instance segmentationcsmp.models.cppnet_base_multiclass"stardist_refined","dist","type"instance segmentationcsmp.models.cppnet_plus"stardist_refined","dist","type","sem"panoptic segmentationcsmp.models.cellvit_sam_base"type","inst","hovernet"instance segmentationcsmp.models.cellvit_sam_plus"type","inst","hovernet","sem"panoptic segmentationcsmp.models.cellvit_sam_small"type","hovernet"instance segmentationcsmp.models.cellvit_sam_small_plus"type","hovernet","sem"panoptic segmentationReferences[1] S. Graham, Q. D. Vu, S. E. A. Raza, A. Azam, Y-W. Tsang, J. T. Kwak and N. Rajpoot. "HoVer-Net: Simultaneous Segmentation and Classification of Nuclei in Multi-Tissue Histology Images." Medical Image Analysis, Sept. 2019.[2] Stringer, C.; Wang, T.; Michaelos, M. & Pachitariu, M. Cellpose: a generalist algorithm for cellular segmentation Nature Methods, 2021, 18, 100-106[3] Cutler, K. J., Stringer, C., Wiggins, P. A., & Mougous, J. D. (2022). Omnipose: a high-precision morphology-independent solution for bacterial cell segmentation. bioRxiv. doi:10.1101/2021.11.03.467199[4] Uwe Schmidt, Martin Weigert, Coleman Broaddus, & Gene Myers (2018). Cell Detection with Star-Convex Polygons. In Medical Image Computing and Computer Assisted Intervention - MICCAI 2018 - 21st International Conference, Granada, Spain, September 16-20, 2018, Proceedings, Part II (pp. 265–273).[5] Hörst, F., Rempe, M., Heine, L., Seibold, C., Keyl, J., Baldini, G., Ugurel, S., Siveke, J., Grünwald, B., Egger, J., & Kleesiek, J. (2023). CellViT: Vision Transformers for Precise Cell Segmentation and Classification (Version 1). arXiv.https://doi.org/10.48550/ARXIV.2306.15350.[6] Chen, S., Ding, C., Liu, M., Cheng, J., & Tao, D. (2023). CPP-Net: Context-Aware Polygon Proposal Network for Nucleus Segmentation. In IEEE Transactions on Image Processing (Vol. 32, pp. 980–994). Institute of Electrical and Electronics Engineers (IEEE).https://doi.org/10.1109/tip.2023.3237013[7] Gamper, J., Koohbanani, N., Benet, K., Khuram, A., & Rajpoot, N. (2019) PanNuke: an open pan-cancer histology dataset for nuclei instance segmentation and classification. In European Congress on Digital Pathology (pp. 11-19).[8] Gamper, J., Koohbanani, N., Graham, S., Jahanifar, M., Khurram, S., Azam, A.,Hewitt, K., & Rajpoot, N. (2020). PanNuke Dataset Extension, Insights and Baselines. arXiv preprint arXiv:2003.10778.[9] Graham, S., Jahanifar, M., Azam, A., Nimir, M., Tsang, Y.W., Dodd, K., Hero, E., Sahota, H., Tank, A., Benes, K., & others (2021). Lizard: A Large-Scale Dataset for Colonic Nuclear Instance Segmentation and Classification. In Proceedings of the IEEE/CVF International Conference on Computer Vision (pp. 684-693).Citation@misc{csmp2022,title={{cellseg_models.pytorch}: Cell/Nuclei Segmentation Models and Benchmark.},author={Oskari Lehtonen},howpublished={\url{https://github.com/okunator/cellseg_models.pytorch}},doi={10.5281/zenodo.7064617}year={2022}}LicenceThis project is distributed underMIT LicenseThe project contains code from the original cell segmentation and 3rd-party libraries that have permissive licenses:timm(Apache-2)HoVer-Net(MIT)Cellpose(BSD-3)Stardist(BSD-3)If you find this library useful in your project, it is your responsibility to ensure you comply with the conditions of any dependent licenses. Please create an issue if you think something is missing regarding the licenses.