Biomedical Image Segmentation by Retina-like Sequential Attention Mechanism Using Only A Few Training Images
Shohei Hayashi Bisser Raytchev Toru Tamaki Kazufumi Kaneda




Biomedical Image Segmentation by Retina-like Sequential Attention




Department of Information Engineering, Hiroshima University, Japan












Biomedical Image Segmentation by Retina-like Sequential Attention Mechanism
Using Only A Few Training Images
    Shohei Hayashi
Bisser Raytchevcorresponding author, email: bisser@hiroshima-u.ac.jp
Toru Tamaki

Kazufumi Kaneda
    December 30, 2023
====================================================================================================================




In this paper we propose a novel deep learning-based algorithm for biomedical image segmentation
which uses a sequential attention mechanism able to shift the focus of attention across the image in a
selective way, allowing subareas which are more difficult to classify to be processed at increased resolution. 
The spatial distribution of class information in each subarea is learned using a retina-like representation where resolution decreases with distance from the center of attention.  
The final segmentation is achieved by averaging class predictions over overlapping subareas, utilizing the power of ensemble learning to increase segmentation accuracy.
Experimental results for  semantic segmentation task
for which only a few training images are available show that a CNN using the proposed method outperforms both a
patch-based classification CNN and a fully convolutional-based method.







§ INTRODUCTION

Recently deep learning methods <cit.>, which automatically extract hierarchical features 
capturing complex nonlinear relationships in the data, have managed to  
successfully replace most task-specific hand-crafted features. This has resulted in a significant
improvement in performance on a variety of biomedical image analysis tasks, like 
object detection, recognition and segmentation
(see e.g. <cit.> for a recent survey of the field and representative methods used in different applications),
and currently Convolutional Neural Network (CNN) based methods define the state-of-the-art in this area.

In this paper we concentrate on biomedical image segmentation. For segmentation, where each pixel needs 
to be classified into its corresponding class, initially patch-wise training/classification was used <cit.>. 
In patch-based methods, local patches of pre-determined size are extracted 
from the images, typically using a CNN as pixel-wise classifier. During training, the patch is used as an input to the network and it is assigned as a label the class of the pixel at the center of the patch (available from ground-truth
data provided by a human expert). During the test phase, a patch is fed into the trained net and the output layer of the
net provides the probabilities for each class. More recently, Fully Convolutional Networks (FCN) <cit.>, which 
replace the fully connected layers with convolutional ones, have
replaced the patch-wise approach by providing a more efficient way to train CNNs end-to-end, pixels-to-pixels, and methods stemming from this approach presently define the state-of-the-art in biomedical image segmentation, exemplified by conv-deconv-based methods like U-Net <cit.>.

  Although fully convolutional methods have shown state-of-the-art performance on many segmentation tasks, they
typically need to be trained on large datasets to achieve good accuracy. In many biomedical image segmentation tasks, however, only
a few training images are available – either data simply being not available, or providing pixel-level ground truth 
by experts being too costly to obtain. Here we are motivated by a similar problem (section 3), 
where less than 50 images are available for training.
On the other hand, patch-wise methods need only local patches, a huge number of which can be extracted 
even from a small number of training images. They however suffer from the following problem. While fully convolutional methods
learn a map from pixel areas (multiple input image values) to pixel areas (multiple classes of all the pixels in the area),
patch-wise methods learn a map from pixel areas (input image values) to a single pixel (class of the pixel in the center  of the patch). 
As illustrated in Fig. <ref> [i], 
this wastes the rich information about the topology 
of the class structure inside the patch for many of the samples which contain more than a 
single class, and these would typically be the most interesting/difficult samples <cit.>. 
Instead of trying to represent in a suitable way and learn the complex class topology, it just substitutes it by
a single class (the class of the pixel in the center of the patch). 




Regarding fully convolutional methods, they treat all locations in the images 
in the same way, which is in contrast with how human visual perception works.
It is known that humans employ attentional mechanisms to focus selectively on
subareas of interest and construct a global scene-representation by combining the information from different local subareas <cit.>.  



Based on the above observations, we propose a new method, which takes a middle ground between fully convolutional and patch-wise learning and combines the benefits of both 
of these strategies.
As shown in Fig. <ref>, as in the patch-wise approach we consider subareas
of the whole image at a time, which provides us with sufficient number of training
samples, even if only a few ground-truth labeled images are available. 
However, as illustrated in Fig. <ref> [ii], the class information is organized
as in the retina <cit.>: the spatial resolution is highest in the central area (corresponding to the fovea), and it
diminishes as we go to the periphery of the subarea.
We propose a sequential attention mechanism which shifts the focus of attention
in such a way that areas of the image which are difficult to classify (i.e. the
classification uncertainty is higher) are considered in much more detail than
areas which are easy to classify. Since the focus of attention moves the subarea 
under investigation much slower over difficult areas (i.e. with much smaller step), this results in many overlapping
subareas in these regions. The final segmentation is achieved by averaging the
class predictions over the overlapping subareas. In this way, the power of ensemble
learning  <cit.> is utilized by incorporating information from all 
overlapping subareas in the neighborhood (each of which provides slightly different 
views of the scene) to further improve accuracy. 










  
This is the basic
idea of the method proposed in the paper, and details how to implement it in a CNN will be given in the next section. 
Experimental results are reported in section 3, indicating that a significant improvement 
in segmentation accuracy can be achieved by the proposed method compared to both patch-based and fully convolutional-based methods. 



§ METHOD

  
  
  We represent a subarea S extracted from an input image (centered at
  the current focus of attention) as a tensor of size d × d × c,
  where d is the size of the subarea in pixels and c stands for the color channels if color images are used (as typically done, and c = 3 for RGB images).
  As shown in Fig. <ref> [ii], we can represent the class information
  corresponding to this subarea as grids of different resolution, where each
  cell in a grid contains a 
sample histogram h^(i) calculated so that the k-th element of  h^(i) 
gives the number of pixels from class k observed in the i-th cell of the grid.

If each histogram h^(i) is normalized to sum to 1 by dividing each bin by the total number of 
pixels covered by the cell, the corresponding vector  p^(i) can now be used to represent the 
probability mass function (pmf) for cell i: the k-th element of  p^(i), i.e. 
p_k^(i), 
can be interpreted as the probability of observing class k in the area covered by
the i-th cell of the grid. 
  
Next, we show how retina-like grids of different resolution levels can be created.
Let's start with a grid of size 4 × 4, as shown in the top row of Fig. <ref> [ii].
This we will call Resolution Level 1 and denote it as r = 1. At resolution level 1 all
the cells in the grid have the same resolution, i.e. the pmfs corresponding to each
cell are calculated from areas of the same size. For example, if the size of the local subarea
under consideration is d = 128 pixels (i.e. image patch of size 128 × 128 pixels), each cell in the grid at
resolution level r = 1 would correspond to an image area of size 32 × 32 pixels from which
a probability mass function p^(i) would be calculated, as explained above.
Next, we can create a grid at Resolution Level 2 (r = 2) by dividing in half 
the four cells in the center of the grid, so that they form an inner 4 × 4 grid, whose
resolution is double. We can continue this process of dividing the innermost 4 cells into 2
to obtain still higher resolution levels. It is easy to see that the number of cells N in a grid
obtained at resolution level r is N = 16 + 12(r-1). 
Of course, it is not necessary the initial grid at r = 1 to be of 
size 4 × 4, but choosing this number makes the process of creating different resolution
levels especially simple, since in this case the innermost cells are always 4 (2 × 2).

In our method, we train a CNN to learn the map between a local subarea image given as input to the network, and the corresponding pmfs p^(i) used as target values.
We use the cross-entropy between the pmfs of the targets p^(i) and the
corresponding output unit activations y^(i) as loss function L:
  
    L = - ∑_n∑_i∑_k p^(i)_k,nlogy^(i)_k,n (S_n ; w⃗),

  where n indexes the training subarea image patches (S_n being the n-th training subarea image
patch), i indexing the cells in the corresponding resolution grid, and k the classes.
Here, w⃗ represents the weights of the network, to be found by minimizing the loss function. To form probabilities, the network output units corresponding to each
cell are passed through the soft-max activation function.
  
  Finally, we describe the sequential attention mechanism we utilize, whose purpose is to move
the focus of attention across the image in such a way that those parts which are difficult
to classify (i.e. classification uncertainty is high) are observed at the highest possible
resolution, and the retina-like grid of pmfs moves with smaller steps across such areas.
To evaluate the classification uncertainty of the grid over the present subarea S,
we use the following function,  
  
    H(S) = - 1/N∑_i ∈ S∑_k p^(i)_klogp^(i)_k,

  which represents the average entropy obtained from the 
posterior pmfs p^(i) for each cell (indexed by i) inside the grid, and
k indexes the classes. Using H(S) as a measure of classification uncertainty, the position
of the next location where to move the focus of attention (horizontal shift in pixels) is given by
  
    f(H(S)) = d exp{-(H(S))^2/2σ^2}.

The whole process is illustrated in Fig. <ref>.
We start at the upper left corner of the input image, with a subarea of size d × d pixels
(the yellow patch (a) in the figure). The classification uncertainty for that subarea is calculated
using Eq. <ref>, and the step size in pixels to move in the horizontal direction
is calculated by Eq. <ref>. As illustrated in the graph in the center of Fig. <ref>, since in this case the classification uncertainty is 0
(all pixels in the subarea belong to the same class), the focus of attention moves d pixels
to the right, i.e. in this extreme case there is no overlap between the current and next subareas.
For the subarea (b) shown in green, the situation is very different. In this case the 
classification uncertainty is very high, and the focus of attention would move only slightly
to the right, allowing the image around this area to be assessed at the highest resolution.
This would result in very high level of overlap between neighboring subareas, as shown in the
heat map on the right (where intensity is proportional to the level of overlap). This process is repeated until the right corner of the image is reached.
Then the focus of attention is moved 10 pixels in the vertical direction to scan
the next row and everything is repeated until
the whole image is processed.






 

  While the above attention mechanism moves the focus of attention across the image, 
the posterior class pmfs from the grids corresponding to each subarea
are stored in a probability map of the same size as the image, i.e. to each pixel in
the image is allocated a pmf equal to the pmf of the cell from the grid positioned
above that pixel. In areas in the image where several subareas overlap, the probability map 
is computed by averaging for each pixel the pmfs of all cells which partially overlap over that pixel. 
Finally, the class of the pixel is
obtained by taking the class with highest probability from the final probability map, as
shown in the upper right corner of Fig. <ref> for the final segmented image.
  
      

  

  
  


  

§ EXPERIMENTS

  In this section we evaluate the proposed method in comparison with a standard patch-wise 
  classification-based CNN <cit.> and the fully convolutional-based U-Net <cit.> on the dataset described below.
  Additionally we implemented a U-Net version, called UNet-patch, which applies U-Net to local patches rather than to a whole image. The original U-Net method which takes as input the whole image we will call UNet-image. 

  Dataset:  Our dataset consists of 59 images showing colonies of undifferentiated and
differentiated iPS cells obtained through phase-contrast microscopy. 
Induced pluripotent stem (iPS) cells <cit.>, for whose discovery S. Yamanaka received the 
Nobel prize in Physiology and Medicine in 2012, contain great promise for regenerative medicine.
Still, in order to fulfill their promise a
steady supply of iPS cells obtained through harvesting of individual cell colonies
is needed and automating the detection of abnormalities arising during the cultivation process is crucial.
Thus, our task is to segment the input images into three categories: Good (undifferentiated), Bad (differentiated) and Background (BGD, the culture medium).
  Several representative images together with ground-truth provided by experts can be seen in 
  Fig. <ref>.
  All images in this dataset are of size 1600 × 1200 pixels. Several images contained a few locations where even the experts were not sure what the corresponding class was.
  Such ambiguous regions are shown in pink and these areas were not used during training and not evaluated during test. 

Network Architecture and Hyperparameters: We used a network architecture based on the VGG-16 CNN net, apart from the following differences.
  There are 13 convolutional layers in VGG-16, while we used 10 here. Also, in VGG-16 there are 3 fully-convolutional layers of which the first two consist of 4096 units,
  while those had 1024 units in our implementation.
  The learning rate was set to 0.0001 and for the optimization procedure we used ADAM.
  Batch size was 16, training for 20 epochs (U-Net-patch for 15 epochs and U-Net-image for 200 epochs). For the implementation of the CNNs we used TensorFlow and Keras.
  
  Four different sizes for the local subareas were tried: d = 96, 128, 192, 256 and resolution level was changed between r = 1 to r = 5. The width of the Gaussian in Eq. <ref> was empirically set to σ = 0.4 for all experimental results.
Evaluation procedure and criteria: The quality of the obtained segmentation results for each method were evaluated by 5-fold cross-validation using the following criteria: 
  Jaccard index (the most challenging one), Dice coefficient, True Positive Rate (TPR), True Negative Rate (TNR) and Accuracy. For each score the average and standard deviation are reported. 

Results: The results obtained on the iPS cell colonies dataset for each of the methods
are given in Table <ref>,
where, patch-center stands for the patch-wise classification method, and
results for resolution levels from r = 1 to r = 5 are given for the proposed method.
As can be seen from the results, the proposed method outperforms both patch-wise classification
and the U-Net-based methods.
  Fig. <ref> gives some examples of segmentation on images from the iPS
dataset, showing that very good accuracy of segmentation can be achieved by the proposed method. 
The heat maps given in the last column demonstrate that the proposed attentional mechanism
is able to focus the high-resolution parts of the retina-like grid on the boundaries between 
the classes which seem to be most difficult to classify, resulting in increased accuracy
of segmentation.













  

§ CONCLUSION

  In this paper we have shown that the combined power of (1) a sequential attention mechanism controlling the shift of the focus of attention, (2) local retina-like representation of the spatial 
  distribution of class information and (3) ensemble learning can lead to increased segmentation accuracy in biomedical segmentation tasks. We expect that the proposed method can be especially useful for datasets for which only a few training images are available.


























splncs04



