11institutetext: Department of Cybernetics,
Faculty of Electrical Engineering,
Czech Technical University in Prague

# Learning to segment from object sizes

Denis Baručić  Jan Kybic

###### Abstract

Deep learning has proved particularly useful for semantic segmentation, a
fundamental image analysis task. However, the standard deep learning methods
need many training images with ground-truth pixel-wise annotations, which are
usually laborious to obtain and, in some cases (e.g., medical images), require
domain expertise. Therefore, instead of pixel-wise annotations, we focus on
image annotations that are significantly easier to acquire but still
informative, namely the size of foreground objects. We define the object size
as the maximum Chebyshev distance between a foreground and the nearest
background pixel. We propose an algorithm for training a deep segmentation
network from a dataset of a few pixel-wise annotated images and many images
with known object sizes. The algorithm minimizes a discrete (non-
differentiable) loss function defined over the object sizes by sampling the
gradient and then using the standard back-propagation algorithm. Experiments
show that the new approach improves the segmentation performance.

###### keywords:

semantic segmentation, weakly-supervised learning, deep learning, distance
transform

00footnotetext: Copyright ©2022 for this paper by its authors. Use permitted
under Creative Commons License Attribution 4.0 International (CC BY 4.0).

##  1 Introduction

Semantic segmentation is the process of associating a class label to each
pixel of an image. With the advent of deep learning, deep networks have
achieved incredible performance on many image processing tasks, including
semantic segmentation. Deep learning for semantic segmentation has many
benefits; for example, it is flexible w.r.t. the model architecture and scales
particularly well [5, 6]. On the contrary, the standard deep learning demands
many ground-truth (GT) pixel-wise annotations to prevent overfitting. Since a
human expert annotator must usually provide the GT annotations, acquiring a
good-quality training dataset can be difficult. To combat this issue, we focus
on learning from GT image annotations that are easier to produce but still
informative enough, namely the sizes of foreground objects. In practice, our
approach assumes a training dataset that consists of relatively few pixel-wise
annotated images and many images with known object sizes. We present a work-
in-progress solution.

###  1.1 Proposed approach

Suppose a standard convolutional network for image segmentation (e.g., a U-Net
[10]). Given an input image, we feed it to the network and collect the output
prediction. The prediction is then thresholded to obtain a binary mask, which
is processed by a distance transform, assigning to each foreground pixel the
shortest distance to the background. Finally, the object size is defined as
double the maximum of the computed distances.

Due to the thresholding, the cost function is not differentiable and it is
therefore not possible to use the standard gradient descent for learning. We
overcome this obstacle by adding random noise to the output of our network.
The predicted binary masks then become stochastic and the gradient can be
sampled. A detailed description of our method is given later in Sec. 2 and 3.

###  1.2 Related work

Cano-Espinosa et al. [1] considered a similar learning problem. They proposed
a network architecture that performs a biomarker (fat contents) regression and
image segmentation after being trained directly on images annotated by
biomarker values only. Similarly to ours, their method derives the biomarker
value from the predicted segmentation deterministically. The difference is
that their biomarker, equivalent to the foreground area, can be obtained by a
simple summation. Furthermore, the method assumes that the foreground objects
can be roughly segmented using thresholding. Pérez-Pelegrí et al. [7] took a
similar approach. Although their method does not involve thresholding to
produce approximate segmentation, it was tailored explicitly for learning from
images annotated by the foreground volume (as their images are 3D).

Karam et al. [4] implemented a differentiable distance transform via a
combination of the convolution operations. The method is fast but exhibits
numerical instabilities for bigger images. Resolving the numerical
instabilities, Pham et al. [8] later proposed a cascaded procedure with
locally restricted convolutional distance transforms. Nonetheless, both
methods substitute the minimum function with the log-sum-exp operation, which
leads to inaccurate results.

The way our method deals with a non-differentiable cost function is borrowed
from stochastic binary networks [9]. In a stochastic binary network, one needs
to deal with zero gradient after each layer of the network. However, methods
such as ARM [13] or PSA [11] are unnecessarily complex. Instead, we employ a
single sample estimation, which has been discussed in [2].

##  2 Model

The proposed model consists of (1) a segmentation network, $f_{\bm{\theta}}$,
parametrized by $\bm{\theta}$, and (2) a deterministic algorithm to derive the
object size based on distance transform, denoted as $g$.

Given an input image $\bm{x}=(x_{1},\ldots,x_{V})$, the network produces a
pixel-wise segmentation

$\bm{a}=f_{\bm{\theta}}(\bm{x}),$ (1)

such that $a_{i}\in{\mathbb{R}},\,1\leq i\leq V$, where $V$ is the number of
pixels. The method does not make any assumptions about the network’s technical
details, except that it can be trained using the standard back-propagation
algorithm and gradient descent. In our experiments, we always employed a U-Net
[10] with a residual network encoder [3] and a mirroring decoder.

To obtain a binary mask $\bm{\hat{y}}\in\\{\pm 1\\}^{V}$, the network response
$\bm{a}$ is thresholded,

$\hat{y}_{i}=\operatorname{sign}a_{i}.$ (2)

###  2.1 Object size

We use a distance transform of the binary mask to define the object size (see
Fig. 1). Distance transform assigns to each pixel the shortest distance to the
background, i.e.,

$d_{i}=\min_{j,\hat{y}_{j}=-1}\delta(i,j),\quad i=1,\ldots,V,$ (3)

where $\delta(i,j)$ is the Chebyshev $\ell_{\infty}$ distance. After that, we
take double the maximum distance to define the object size,

$\hat{s}=2\,\max_{i}\,d_{i}.$ (4)

The composition of the distance transform and the maximum aggregation is the
object size, denoted as $g\colon\\{\pm 1\\}^{V}\to\mathbb{R}$,

$g(\hat{\bm{y}})=2\,\max_{i}\min_{j,\hat{y}_{j}=-1}\delta(i,j).$ (5)

$i$$d_{i}$$\tilde{s}$ Figure 1: Illustrative example of an object and its
derived size. The object is outlined by the thick boundary line. The point $i$
denotes the foreground pixel whose shortest distance to the background,
$d_{i}$, is the highest among the pixels. The derived object size
$\hat{s}=2d_{i}$.

####  2.1.1 Implementation details

There is an efficient, two-pass algorithm that computes the distance transform
in $\Theta(V)$ time. Furthermore, when evaluating a batch of images, it is
possible to compute the distance transform on all images in parallel.

We have implemented a CPU version111https://github.com/barucden/chdt of this
algorithm that works with PyTorch tensors and is faster than, e.g., the SciPy
implementation.

image $\bm{x}$Segmentationnetwork $f_{\bm{\theta}}$noise
$Z$DistancetransformMaxLoss $l$$s$$l(s,g(Y))$Size derivation $g$$Y$$g(Y)$
Figure 2: An overview of the proposed probabilistic model.

##  3 Learning

Suppose a training dataset $\mathcal{D}=\mathcal{D}_{f}\cup\mathcal{D}_{w}$
consists of fully- and weakly-annotated subsets $\mathcal{D}_{f}$ and
$\mathcal{D}_{w}$. The fully-annotated subset $\mathcal{D}_{f}$ contains pairs
$(\bm{x},\bm{y})$, where $\bm{x}$ is an input image and $\bm{y}$ the
corresponding GT pixel-wise segmentation, while $\mathcal{D}_{w}$ comprises of
pairs $(\bm{x},s)$, where $s$ is the size of the object present in the image
$\bm{x}$. We focus on situations when
$\lvert\mathcal{D}_{f}\rvert\ll\lvert\mathcal{D}_{w}\rvert$.

###  3.1 Supervised pre-training

Our method starts by optimizing a pixel-wise loss w.r.t. the network
parameters $\bm{\theta}$ on the small subset $\mathcal{D}_{f}$, as in the
standard supervised learning. For a particular training pair
$(\bm{x},\bm{y})\in\mathcal{D}_{f}$ and the corresponding prediction
$\bm{a}\in\mathbb{R}^{V}$, the loss function reads

$\sum_{i=1}^{V}\left(a_{i}(1-y_{i})+\log(1+\exp(-a_{i}))\right),$ (6)

which is sometimes referred to as the binary cross-entropy with logits loss.
The optimization continues until convergence.

Using proper data augmentation to extend the training dataset, the network
tends to recognize useful features and produces decent predictions after this
initial stage (see Sec. 4.2).

###  3.2 Weakly-supervised training

Consider a training pair $(\bm{x},s)\in\mathcal{D}_{w}$. As described in Sec.
2, one can obtain a prediction of the object size, $\hat{s}=g(\hat{\bm{y}})$,
from the thresholded network response $\hat{\bm{y}}$. We penalize the
prediction error by the square loss

$l(s,\hat{s})=(s-\hat{s})^{2}.$ (7)

We propose to follow an approach similar to those used in binary neural
networks [11] and subtract random noise $Z$ from the real predictions $a_{i}$
before thresholding. Consequently, the binary segmentation becomes a
collection $\bm{Y}=(Y_{1},\ldots,Y_{V})$ of $V$ independent Bernoulli
variables,

$Y_{i}=\operatorname{sign}(a_{i}-Z),$ (8)

with

$\Pr(Y_{i}=+1\mid\bm{x};\bm{\theta})=\Pr(Z\leq a_{i})=F_{Z}(a_{i}),$ (9)

where $F_{Z}$ is the cumulative distribution function (CDF) of the noise $Z$
(see Fig. 2).

Then, instead of minimizing the loss $l$ (7), we minimize the expected loss
$\mathcal{L}=\operatorname{\mathbb{E}}_{\bm{Y}}[l(s,g(\bm{Y}))]$,

$\mathcal{L}=\sum_{\bm{y}\in\\{\pm
1\\}^{V}}\Pr(\bm{Y}=\bm{y}\mid\bm{x};\bm{\theta})l(s,g(\bm{y})).$ (10)

Contrary to (7), the expected loss (10) is differentiable, assuming a smooth
$F_{Z}$.

####  3.2.1 Noise distribution

Following [11], we sample the noise $Z$ from the logistic distribution with
mean $\mu=0$ and scale $s=1$. Hence, the CDF of $Z$ is a smooth, sigmoid
function,

$F_{Z}(a)=\frac{1}{1+\exp(-a)}.$ (11)

####  3.2.2 Exact gradient

To compute the gradient $\nabla_{\bm{\theta}}\mathcal{L}$, we need to evaluate
the derivative

$\frac{\partial\operatorname{\mathbb{E}}_{\bm{Y}}[l(s,g(\bm{Y}))]}{\partial
F_{Z}(a_{i})}$ (12)

for each pixel $i=1,\ldots,V$. The gradient can be then computed automatically
by the back-propagation algorithm. However, an exact computation of (12) leads
to

$\sum_{\bm{y}\in\\{\pm
1\\}^{V}}\frac{\Pr(\bm{Y}=\bm{y}\mid\bm{x};\bm{\theta})}{\Pr(Y_{i}=y_{i}\mid\bm{x};\bm{\theta})}l(s,g(\bm{y}))y_{i},$
(13)

which involves summing $2^{V}$ terms and is thus tractable only for very small
images. Instead, we resort to a single sample estimator.

####  3.2.3 Single sample estimator

The single sample estimator is based on Lemma 3.1, which is, in fact, a
specific form of [11, Lemma B.1].

######  Lemma 3.1.

Let $\bm{Y}=(Y_{1},\ldots,Y_{V})$ be a collection of $V$ independent $\\{\pm
1\\}$-valued Bernoulli variables with probabilities $\Pr(Y_{i}=+1)=p_{i}$. Let
$h$ be a function $h\colon\\{\pm 1\\}^{V}\to\mathbb{R}$. Let
$\bm{y}=(y_{1},\ldots,y_{V})$ denote a random sample of $\bm{Y}$ and
$\bm{y}_{\downarrow i}=(y_{1},\ldots,y_{i-1},-y_{i},y_{i+1},\ldots,y_{V})$.
Then

$y_{i}\left(h(\bm{y})-h(\bm{y}_{\downarrow i})\right)$ (14)

is an unbiased estimate of $\frac{\partial}{\partial
p_{i}}\operatorname{\mathbb{E}}_{\bm{y}\sim\bm{Y}}[h(\bm{y})]$.

######  Proof 3.2.

We take the derivative of the expectation,

$\frac{\partial}{\partial
p_{i}}\operatorname{\mathbb{E}}_{\bm{y}\sim\bm{Y}}[h(\bm{y})]=\sum_{\bm{y}}\frac{\Pr(\bm{y})}{\Pr(y_{i})}h(\bm{y})y_{i},$
(15)

and write out the sum over $y_{i}$,

$\sum_{\bm{y}_{\neg i}}\sum_{y_{i}}\Pr(\bm{y}_{\neg
i})h(\bm{y})y_{i}=\sum_{\bm{y}_{\neg i}}\Pr(\bm{y}_{\neg
i})\sum_{y_{i}}h(\bm{y})y_{i}$ (16)

where $\bm{y}_{\neg i}$ denotes vector $\bm{y}$ with the $i$-th component
omitted. Notice that the inner sum simplifies and no longer depends on
$y_{i}$,

$\sum_{\bm{y}_{\neg i}}\Pr(\bm{y}_{\neg
i})(h(\bm{y}_{i=+1})-h(\bm{y}_{i=-1})),$ (17)

where $\bm{y}_{i=z}$ is the vector $\bm{y}$ with the $i$-th component set to
$z$. Then, we multiply the inner subtraction by the constant factor
$1=p_{i}+(1-p_{i})=\sum_{y_{i}}\Pr(y_{i})$,

$\sum_{\bm{y}_{\neg i}}\Pr(\bm{y}_{\neg
i})\sum_{y_{i}}\Pr(y_{i})(h(\bm{y}_{i=+1})-h(\bm{y}_{i=-1})),$ (18)

ultimately leading to the following expression for (15):

$\sum_{\bm{y}}\Pr(\bm{y})(h(\bm{y}_{i=+1})-h(\bm{y}_{i=-1})),$ (19)

which can be written as

$\sum_{\bm{y}}\Pr(\bm{y})y_{i}\left[h(\bm{y})-h(\bm{y}_{\downarrow
i})\right].$ (20)

Thus, (14) is a single sample unbiased estimate of (15).

$F_{Z}$

$n=1$

$n=8$

$n=64$

$n=512$

Figure 3: Examples of derivatives (12) computed according to (21) for
different number of samples $n$, given the output of $F_{Z}$, for a small,
$6\times 6$ image. The red frame outlines the object.

According to Lemma 3.1, an unbiased estimate of the derivative (12) is

$\frac{\partial\operatorname{\mathbb{E}}_{Y}[l(s,g(Y))]}{\partial
F_{Z}(a_{i})}\approx y_{i}\left[l(s,g(\bm{y}))-l(s,g(\bm{y}_{\downarrow
i}))\right],$ (21)

where $\bm{y}$ is a random sample of Bernoulli variables with probabilities
(9) (see a few examples of sampled derivatives in Fig. 3).

##  4 Experiments

The proposed method was implemented in the PyTorch Lightning
framework222https://github.com/Lightning-AI/lightning using a ResNet
implementation from the Segmentation Models PyTorch
library333https://github.com/qubvel/segmentation_models.pytorch. The presented
experiments were perfomed on a server equipped with Intel Xeon Silver 4214R
(2.40GHz) and NVIDIA GeForce RTX 2080 Ti.

Figure 4: Example of a hippocampus image [12] with the object outlined in red.

The data for our experiments was based on a dataset of 3D MRI images of the
hippocampus [12]. The dataset consists of 394 volumes provided with GT
segmentation of classes hippocampus head, hippocampus body, and background. We
decomposed the volumes into individual 2D slices of size $48\times 32$ pixels
and kept only those with at least 1% foreground, obtaining a total of 6093
images. Next, we merged the hippocampus classes to get a binary segmentation
problem (see Fig. 4). Afterward, we derived the object sizes from the GT
pixel-wise annotations to use in training. Finally, we randomly split the data
into training, validation, and testing subsets containing 70%, 10%, and 20% of
the images.

Given a GT segmentation $\bm{y}$ and a predicted segmentation $\hat{\bm{y}}$,
we evaluate two metrics, the squared size prediction error $E$ and the
intersection-over-union $IoU$,

$\displaystyle E(\bm{y},\hat{\bm{y}})$
$\displaystyle=l(g(\bm{y}),g(\hat{\bm{y}})),$ (22) $\displaystyle
IoU(\bm{y},\hat{\bm{y}})$
$\displaystyle=\frac{\sum_{i=1}^{V}1+y_{i}+\hat{y}_{i}+y_{i}\hat{y}_{i}}{\sum_{i=1}^{V}3+y_{i}+\hat{y}_{i}-y_{i}\hat{y}_{i}}.$
(23)

In the case of standard supervised method, vertical and horizontal flipping
was randomly applied to augment the training dataset. The proposed method did
not apply any augmentation.

###  4.1 Number of derivative samples

Figure 5: Average epoch duration for the proposed method with different number
of gradient samples. The duration of the standard method is given as a
reference.

Figure 6: Development of the squared size prediction error $E$ and the
intersection-over-union $IoU$ on the validation images over the course of
learning for different numbers of derivative samples $n$.

A toy example (see Fig. 3) indicated that taking more samples of the
derivatives (21) might lead to better results than taking just one. This
experiment investigates how the number of derivative samples $n$ impacts
learning speed and prediction quality.

We considered four different numbers of samples $n$, $n\in\\{1,2,4,8\\}$. For
each $n$, the other parameters (such as the batch size or the learning rate)
were the same, and the learning began with the same segmentation network
$f_{\bm{\theta}}$ that was pre-trained in the standard way on $85$ pixel-wise
annotated images from the training subset. The proposed method always ran
until the squared error $E$ on the validation data stopped improving.

To assess the learning speed, we measured the duration of one learning epoch.
For $n=1$, an epoch took $\approx 10\times$ longer than the standard
supervised learning. Generally, the duration grew roughly exponentially with
$n$ (see Fig. 5).

Higher values of $n$ did not lead to a lower $E$ or a faster convergence speed
(see Fig. 6). In fact, $n=1$ and $n=2$ achieved the lowest $E$, but not by a
large margin. Given the speed benefits, we use $n=1$ always. Interestingly,
even though $E$ kept decreasing over the course of learning for all $n$, $IoU$
improved only slightly and started declining after $\approx 20$ epochs. This
observation suggests that the squared error of the object size is not a
sufficient objective for learning the segmentation.

###  4.2 Pre-training impact

This experiment tests the essential question: given a segmentation model
trained on a few pixel-level annotated images, can we improve its testing
performance by further learning from size annotations?

We trained different segmentation networks until convergence on randomly
selected training subsets of size $m$. Then, we fine-tuned these networks on
the whole training dataset using the proposed method. We measured the test
performance in terms of $IoU$.

The proposed method led to a $\approx 5\%$ increase of $IoU$ for small $m<100$
(see Fig. 7), improving the segmentation quality. For higher $m$, the effect
was negligible, which complements the observation from the previous experiment
that improving the size estimate does not necessarily improve the segmentation
quality.

Figure 7: $IoU$ on the test data for different sizes $m$ of the pre-training
dataset. The plot shows results achieved by a network after pre-training and
after subsequent fine-tuning by the proposed method.

##  5 Discussion

The method is promising but there is definitely potential for improvement in
both speed and prediction performance.

The proposed method samples the derivatives according to (21) for each pixel
$i$. Flipping the prediction, $y_{i}\mapsto-y_{i}$, changes the derived size
only for some $i$; particularly those within and on the border of the
predicted object. Therefore, given a sample $\bm{y}$,
$l(s,g(\bm{y}))=l(s,g(\bm{y}_{\downarrow i}))$ for many pixels $i$, and the
sampled derivatives (21) are sparse. The method might sample only those
derivatives that are potentially non-zero and set the rest to zero directly,
which would save much computational time.

We have seen in the experiments that lower size prediction error does not
strictly imply better segmentation. We need to closely investigate in what
cases the size prediction loss is insufficient and adjust the objective. The
adjustment might involve adding an L1 regularization (as in [1]) or drawing
inspiration from unsupervised methods (e.g., demand for the segmentation to
respect edges in images, etc.).

The proposed approach entails some principled limitations. For example, it
allows only a single object in an image. We also expect the method to be ill-
suited for complex object shapes, but we have not performed any experiments in
that regard yet.

##  6 Conclusion

We proposed a weakly-supervised method for training a segmentation network
from a few pixel-wise annotated images and many images annotated by the object
size. The key ingredients is a method for evaluating the object size from a
probabilistic segmentation and a method for optimizing a deep network using a
non-differentiable objective.

The achieved results seem promising. We believe the improvements suggested in
the discussion will improve performance, rendering the method valuable for
training segmentation models for biomedical images.

### Acknowledgments

The authors acknowledge the support of the OP VVV funded project
“CZ.02.1.01/0.0/0.0/16_019/0000765 Research Center for Informatics”, the Czech
Science Foundation project 20-08452S, and the Grant Agency of the Czech
Technical University in Prague, grant No. SGS20/170/OHK3/3T/13.

## References

  * [1] C. Cano-Espinosa et al.  Biomarker localization from deep learning regression networks.  IEEE Transactions on Medical Imaging, 39(6):2121–2132, 2020.
  * [2] Y. Cong, M. Zhao, K. Bai, and L. Carin.  GO gradient for expectation-based objectives.  In 7th International Conference on Learning Representations, 2019\.
  * [3] K. He, X. Zhang, S. Ren, and J. Sun.  Deep residual learning for image recognition.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 770–778, 2016.
  * [4] C. Karam, K. Sugimoto, and K. Hirakawa.  Fast convolutional distance transform.  IEEE Signal Processing Letters, 26(6):853–857, 2019.
  * [5] X. Liu et al.  A review of deep-learning-based medical image segmentation methods.  Sustainability, 13(3):1224, 2021.
  * [6] S. Minaee et al.  Image segmentation using deep learning: A survey.  IEEE transactions on pattern analysis and machine intelligence, 2021\.
  * [7] M. Pérez-Pelegrí et al.  Automatic left ventricle volume calculation with explainability through a deep learning weak-supervision methodology.  Computer Methods and Programs in Biomedicine, 208:106275, 2021.
  * [8] D. D. Pham, G. Dovletov, and J. Pauli.  A differentiable convolutional distance transform layer for improved image segmentation.  In DAGM German Conference on Pattern Recognition, pages 432–444. Springer, 2020.

  * [9] T. Raiko, M. Berglund, G. Alain, and L. Dinh.  Techniques for learning binary stochastic feedforward neural networks.  In 3rd International Conference on Learning Representations, 2015\.
  * [10] O. Ronneberger, P. Fischer, and T. Brox.  U-net: Convolutional networks for biomedical image segmentation.  In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 234–241. Springer, 2015.
  * [11] A. Shekhovtsov, V. Yanush, and B. Flach.  Path sample-analytic gradient estimators for stochastic binary networks.  Advances in Neural Information Processing Systems, 33:12884–12894, 2020.
  * [12] A. L. Simpson et al.  A large annotated medical image dataset for the development and evaluation of segmentation algorithms, 2019, arXiv:1902.09063.
  * [13] M. Yin and M. Zhou.  ARM: augment-REINFORCE-merge gradient for stochastic binary networks.  In 7th International Conference on Learning Representations, 2019\.