SAU-Net: A Unified Network for Cell Counting in 2D and 3D Microscopy Images

Abstract

Image-based cell counting is a fundamental yet challenging task with wide applications in biological research. In this paper, we propose a novel unified deep network framework designed to solve this problem for various cell types in both 2D and 3D images. Specifically, we first propose SAU-Net for cell counting by extending the segmentation network U-Net with a Self-Attention module. Second, we design an extension of Batch Normalization (BN) to facilitate the training process for small datasets. In addition, a new 3D benchmark dataset based on the existing mouse blastocyst (MBC) dataset is developed and released to the community. Our SAU-Net achieves state-of-the-art results on four benchmark 2D datasets - synthetic fluorescence microscopy (VGG) dataset, Modified Bone Marrow (MBM) dataset, human subcutaneous adipose tissue (ADI) dataset, and Dublin Cell Counting (DCC) dataset, and the new 3D dataset, MBC. The BN extension is validated using extensive experiments on the 2D datasets, since GPU memory constraints preclude use of 3D datasets. The source code is available at https://github.com/mzlr/sau-net.

Fig. 1:

A sample image from Modified Bone Marrow (MBM) dataset [11] with dot annotations shown as red cross overlays (left image) and the corresponding dot annotations used in training (right image).

Fig. 2:

Overview of our regression based cell counting method with a 2D image containing two cells. The counting problem is addressed in two steps: 1) learn a mapping F from the input image I to the density map D and 2) integrate D to predict the final count. We apply convolution with a Gaussian kernel on the binary labels L to obtain D for the learning process.

Fig. 3:

The overall structure of SAU-Net (left) with details of each block (right). SAU-Net seamlessly unifies 2D and 3D image-based cell counting. The number after conv, maxpooling, deconv indicates the filter size, e.g., 3 for 3 × 3 or 3 × 3 × 3 depending on the dimension of the input images. The output layer is implemented as conv_1. All the convolutions and deconvolutions use ReLU nonlinearity except the linear embeddings, f, g and h in the Self-Attention module. ⊗ denotes matrix multiplication.

Fig. 4:

Sample images from the five datasets used in the experiment. The dot annotations are represented as the red cross overlays. The projections are generated using Volume Viewer in ImageJ.

Fig. 5:

The evaluation comparison between the proposed BN extension and standard BN with three common moving average momentums, m, on VGG, MBM, ADI, and DCC. All models were initialized with zero-mean and unit-variance at Step 0 (not shown) and are trained on SAU-Net under the same settings.

Fig. 6:

The L1 distance of the statistics between standard BN and our extension under three common moving average momentums, m.

Fig. 7:

Sample predicted density map on the test set for VGG dataset. Groundtruth Cell Count: 100, Predicted: 100.9

Fig. 8:

Sample predicted density map on the test set for MBM dataset. Groundtruth Cell Count: 134, Predicted: 135.8

Fig. 9:

Sample predicted density map on the test set for ADI dataset. Groundtruth Cell Count: 149, Predicted: 142.1

Fig. 10:

Sample predicted density map on the test set for DCC dataset. Groundtruth Cell Count: 55, Predicted: 58.9

Fig. 11:

Sample predicted density map on the test set for MBC dataset. Groundtruth Cell Count: 59, Predicted: 60.4

Fig. 12:

The impact of batch size B on our proposed BN extension with various datasets. B is set to 15 for MBM and 75 for the rest three 2D datasets as in the previous section.