2.1 VQ-VAE

Variational Autoencoders [5] can compress an image dataset into a latent multivariate gaussian distribution. For further image compression, quantization methods have been applied successfully [6, 7]. Similarly, Van den Oord et al. proposed the Vector Quantised-Variational AutoEncoder (VQ-VAE) [8] that also encodes input data to a vector of discrete latent variables. However, the VQ-VAE employs the discrete latent variables as indexes to a memory table to recover a set of embeddings. The embeddings are decoded to produce an image. Posterior work [9] proposed a hierarchical VQ-VAE with several latent codes to improve image quality.

2.2 Cell instance segmentation

Segmentation is considered the first critical step for biomedical microscopy image analysis. Although standard deep learning approaches have been applied to the task [10, 11], some proposals explore alternative data representations that might be more suitable for microscopy. For example, StarDist [12, 13] segments cell nuclei by describing them as star-convex polygons. StarDist first predicts for each pixel the distance to the cell nucleus boundary along a set of predefined equidistant angles, and afterward performs non-maximum suppression (NMS) to discard duplicated polygons. Multistar [14] extends Stardist by identifying overlapping objects, while SplineDist [15] modifies StarDist by representing objects as planar spline curves. Another example is Cellpose [16, 17] learns during training to predict the gradients of a diffusion process starting at the cell centre. Later during inference, Cellpose backtracks the predicted gradient to see which pixels converge to the same cell centre.

Finally, in amodal blastomere instance segmentation, a VQ-VAE has been used after a typical detection pipeline to generate mask representations from the image features [18]. In contrast, our method employs the VQ-VAE to generate individual objects representations and later performs the segmentation.

2.3 Image decomposition

Early work [19] already proposed image decomposition as an alternative to object detection pipelines. Other approaches tackle unsupervised learning of object representations through image reconstruction [20–23]. While MONet [20] and IODINE [22] apply a network recurrently, Slot-attention [21] and Capsule networks [23] apply an algorithm iteratively. Although Kaizen is similar since it iteratively applies a VQ-VAE, it is trained in a supervised manner. Unlike those models wich incorporate an explicit attention mechanism, Kaizen relies on a VQ-VAE to generate only individual objects. More recently Composer [24] decomposes an image into several decoupled representations including object instances, and then trains a diffusion model conditioned to these representations, to improve control on image generation.

3.1 VQ-VAE

As a generative model for Kaizen a VQ-VAE [8] was chosen, favouring inference speed over quality of the samples. The VQ-VAE was trained with small patches containing a few cells from the hundreds in a microscopy image. The purpose of the training was for the VQ-VAE to produce a single cell image as output when given an image containing multiple cells. Thus, for a given input training patch the corresponding training output was generated by multiplying the input patch by the cell’s mask touching the central pixel. As shown in Fig 2, after training the VQ-VAE encodes a single central cell from small patches containing multiple cells while disregarding other cells and noise present on the input. To avoid encoding non-existent cells, twenty percent of the training patches did not have any cell touching the central pixel; in these cases the VQ-VAE was trained to output an empty image.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 2. Samples of the VQ-VAE encoding individual cells for the U2OS dataset.

The first row corresponds to training image patches employed as input to the VQ-VAE. In the second row, under each input, the corresponding ground truth is shown. Finally, the third row shows the corresponding VQ-VAE ouput after training.

https://doi.org/10.1371/journal.pone.0313549.g002

3.2 Core algorithm

Using the VQ-VAE that encodes individual cells described above as predictor and taking inspiration from predictive coding in the brain, the core algorithm of Kaizen was implemented. Algorithm 1 calculates an error image as the difference between the input and a internal image prediction. At each iteration, several points are proposed where the error image reaches its maximum, and distant between them. Then at each proposed point in the image an object prediction is made. Finally, only predictions that diminish the loss when added to the internal prediction, are kept.

3.3 Predicting on the error

Notably, keeping an internal image allows estimation of an error image, which is derived by subtracting the internal prediction from the input image. Predicting on the error image filters out already predicted objects. Thus, predicting on the error image is easier, given its reduced complexity compared to the original image, minimizing potential confusion for the predictor. However, erroneous predictions in the initial stages could propagate into subsequent ones, compounding any inaccuracies committed by the predictor.

Algorithm 1. Core algorithm

To enhance object detection while minimizing error compounding, Kaizen first predicts on the input image with the core algorithm until no more object predictions are found. With this first set of predictions an error image is generated. Then, with the error image as input, we apply Kaizen again generating new predictions and a second error image. We repeat this process until no new objects are found in a given iteration, aiming to detect all objects even in the most densely populated images.

3.4 Datasets

U2OS dataset: Fluorescent microscopy images from U2OS cell lines. Image set BBBC039 version 1, available from the Broad Bioimage Benchmark Collection [25]. The ground truth consists of individual nucleus instances. Although the ground truth might contain some errors, we did not alter the ground truth. We randomly selected 100 images for the training set, 50 for the validation set, and 50 for the test set.

Neuroblastoma dataset: Fluorescent images of cultured neuroblastoma cells, available from the Cell Image Library [26]. The ground truth consists of manually annotated cell boundaries. Although the ground truth might contain some errors, we did not alter the ground truth. However, since the ground truth size was half the weight and height of the images, we shrink the images with cubic interpolation to match the ground truth size. We randomly selected 71 images for the training set, 12 for the validation set, and 17 for the test set.

3.5 Evaluation metric

A predicted object is considered a true positive (TP) if his intersection over union IoU with a ground truth object is above certain threshold . For the same threshold predicted objects without ground truth are considered false positives (FP), and finally ground truth not predicted is considered false negative (FN). The average precision (AP) for one image is given by:

(1)

The average precision reported is the average over all images in the test set.

3.6 Implementation details

The VQ-VAE encoder consists first of 5 strided convolutional layers with padding 1. The first three layers have stride 2 and window size 4 4. Fourth layer has stride 1 and window size . Fifth layer has stride 1 and window size . The first two layers have 64 hidden units while the rest of the layers have 128 hidden units. The convolutional layers are followed by two residual blocks (implemented as ReLU, conv, ReLU, conv)

The VQ-VAE decoder first has 1 convolutional layer of stride 1, window size , padding 1 and 128 hidden units. Then it has two residual blocks. Next, three transposed convolutions follow with stride 2, window size , and the first one with padding 1. The final layer is a transposed convolution with stride 1, window size and 64 hidden units.

The VQ-VAE codebook contains 128 embeddings of dimension 2, and we use exponential moving averages [8] to update the dictionary with = 0.25. We use the ADAM optimiser [27] with learning rate 1e–4, L1 loss, and train for 500,000 steps with batch-size 32. L1 loss was also employed for the Kaizen algorithm.

To predict at the border of the image, the input image was padded by half the input size of the VQ-VAE, eight pixels of mirror-padding followed by zero-padding (20 pixels in total for the U2OS dataset and 60 for the Neuroblastoma dataset).

To select several distant points in parallel, we convolve the error image with a kernel of ones with stride one. The highest value in the resulting image is the first point. The region surrounding this point ( pixels) is set to zero, and the process is repeated to select subsequent points. This process aims to minimize the occurrence of predictions on background noise and cell boundaries.

For the U2OS dataset the number of simultaneous points of prediction was set to ten (N in Algorithm 1). To avoid a hypothetical infinite loop, the repeat loop in Algorithm 1 was set to a maximum of 30 iterations. For the neuroblastoma dataset the number of simultaneous points of prediction was set to one (N in Algorithm 1) and the repeat loop in Algorithm 1 was set to a maximum of 100 iterations.

As pre-processing the images of both datasets were normalized. No data augmentation was applied. As post-processing to avoid empty predictions in the U2OS dataset, we eliminate all the masks with less than 20 pixels. For the same reason, in the neuroblastoma dataset, all masks with less than 20 blue pixels or less than one green pixel are eliminated.