An U-Net - based regression model
incorporating a parametric description of theprostate
Artur Kos
Department of Telecommunications
AGH University of Science and Technology
kosar@agh.edu.plAbstract. Accurate prostate segmentation in the CT data remains achallenging task due to organ variation, typically unclear boarder between neighbouring structures and low signal-to-noise ratio of the CTdata. We propose an U-Net based algorithm, which combines its encoder
part with the parametric shape model of prostate. The presented algorithm acts as a regression model, has a linear activation function ofthe output layer and could be directly optimized by minimization of theMSE distance to ground-truth parameters. Segmentation accuracy of theproposed method was measured using Dice Similarity Coefficient (DSC)
and was equal 0.73 with STD equal 0.18.
Keywords: prostate segmentation · CT · U-NET· CNN · semantic segmentation.

1

Introduction

A prostate cancer remains one of the most common cancer among the men
population [1]. Often practised treatment strategy is the radiotherapy whichneeds a several of CT volumes with precise manual delineation of the radiatedstructures. Since such task is time consuming, an automatic (semiautomatic)
image segmentation techniques is desired. It is, however, not a trivial task due
to organ variation, typically unclear boarder between neighbouring structuresand low signal-to-noise ratio in the CT data.
To overcome these problems, a several segmentation algorithms (operating inCT or other modalities) has been developed. Part of them can be described as aconventional [2], while other are based on convolutional neural networks (CNNs).
The latter group have gained a significant attention in the field of medical image
segmentation, including CT images of the pelvis area [3, 6].

2Methods and Data

Available image database was consisted of 42 real-patient CT volumes without
contrast enhancement, each having 512×512 pixels in axial plane with the inslice resolution ranged from 0.9678 to 1.225 mm. Number of slices in one volume

2

Artur Kos

Fig. 1. Segmentation stages of the proposed method. Left image: original CT slice;
central image: ground truth - manually delineated by the doctor (white contour), parametric model calculated from ground truth (red dots); right image: parametric model
calculated from ground truth (red dots), parametric model inferred by CNN from theoriginal CT slice (green dots).

was varied from 63 to 116 across the set (84 at average), number of delineatedslices varied from 6 to 18 (9.1 at average), while their slice thickness was always
equal 5 mm. Each volume came from a different patient and was delineatedmanually by a medical doctor (see Fig. 1).
Along with other researchers [4, 5], in our ongoing work we state, that U-Net
architecture is suitable for prostate segmentation. It means, that latent space inthat model is capable enough to describe the shape of the prostate from a single
CT slice. This observation is explored in the proposed model.
The original U-Net architecture is the so-called fully-convolutional networkand consists of the encoder and the decoder. Since the information is locatedmostly in the latent space (the layer between the encoder and the decoder) theproposed model ends on this layer.
In order to recover the shape of the prostate slice, a simple parametric model
was proposed. It consists of 18 parameters. The first two parameters (a1 anda2 ) correspond to the ”x” and ”y” coordinate of the prostate center and theremaining 16 corresponds to the uniformly spaced radius which spans the shapeof the prostate contour (see Fig. 2). This model is justified by the observationof the typical shape of the prostate slice. It is a continuous round shape witha lack or a single small cavity. It could be well described by this parametricmodel. Another justification of that prostate model is the possibility to minimize
parameters by the MSE algorithm.
Finally, the proposed model is shown in the Fig. 3. It consists of the encoder
part of the original U-Net. The layer corresponding to the latent space is followedby three fully-connected layers. The output layer corresponds to the parametricmodel of the prostate shape. In this example, the size of the output layer is set
to 18.
In contrast to the original U-Net which was the so-called semantic segmentation model, the proposed one acts as a regression model and thus, the activation
function of the output layer is a linear one. Since all parameters of the prostateAn U-Net - based regression model3Fig. 2. Parametric model description. The first two parameters correspond to the ”x”
and ”y” coordinate of the prostate center and the remaining n-2 (in tested implementation 16) corresponds to the uniformly spaced radius which spans the shape of theprostate contour.

model are located on the x-y plane, the model could be optimized by the MSE
distance to ground-truth parameters.

input image
1632512128

MaxPooling + Dropout

128

64

1842512
62

256
82

256
162

128
182

128

20264
882

902

922

186 x 186

184 x 184

188 x 188

378 x 378

376 x 376

380 x 38064

402

32422

16442

1output parameters

Fig. 3. Network architecture. Data flow is from left to right. Blue boxes correspond tothe convolutional layer. The number of filters (channels) is denoted above each box,
while the size of the layer - on the left-hand side. The max-pooling operation is stacked
with dropout regularization. The Green color denotes the fully connected layer. Size ofthe output layer corresponds to the number of parameters of the prostate parameters.

The network presented in Fig. 3 was trained and validated only on slices
from CT scans on which the prostate intersection was located and delineatedby the professional clinics. The training and the validation set consists of 275
and 83 slices respectively. The training procedure was the Adam algorithm withthe standard parameters. The loss function was the standard MSE distance tothe ground truth. To prevent model overfitting, the dropout layers were added
as well as L2 regularization to fully-connected layers. Additionally, the data
augmentation procedure was applied to the input data. Input images (DICOM4Artur Kosslices) were cropped from 512x512 to 380x380 pixels. Additionally, this window
was randomly shifted in the ”x” and ”y” direction, such as the ground truth
region was still located into the output window. Additionally, random Gaussian
noise was added to the image. The distribution of Gaussian noise was set tomatch the distribution of pixels intensity of the prostate. Both augmentations
were applied simultaneously to all slices from all batches in all epochs. The
proposed computational model is relatively simple and well understood, andthus, the standard training and preprocessing procedures could be applied. It has
the potential to be used with other, more sophisticated parametric descriptionof the prostate contour.

3Results and Discussion

Table 1. Comparison of results obtained from the proposed method and the U-Net
methodDiceU-Net
0.720+-0.17
Proposed method 0.73+-0.18

The obtained results presented in Table 1 are not superb, however, in ouropinion they prove that it is possible to extract the correct shape and locationof the prostate slice from the latent space of the U-Net architecture and obtaincomparable results with the regular U-net. What is more, this approach could
provide the so-called ”model explainability” required for application in a real
medical case.

References
1. Bray, Freddie, et al. Global cancer statistics 2018: GLOBOCAN estimates of incidence and mortality worldwide for 36 cancers in 185 countries. CA: a cancer journalfor clinicians 68.6 (2018): 394-424
2. Soumya Ghose, Arnau Oliver, et al. A survey of prostate segmentation methodologies in ultrasound, magnetic resonance and computed tomography images, Computer Methods and Programs in Biomedicine, Volume 108, Issue 1, 2012
3. Yamashita, Rikiya, et al. Convolutional neural networks: an overview and application in radiology. Insights into imaging 9.4 (2018): 611-629.
4. O. Ronneberger, P. Fischer, T. Brox, U-net: Convolutional networks for biomedical
image segmentation, in: International Conference on Medical image computing andcomputer-assisted intervention, Springer, 2015
5. Balagopal A, Kazemifar S, Nguyen D, Lin MH, Hannan R, Owrangi A, Jiang S.
Fully automated organ segmentation in male pelvic CT images. Phys Med Biol.
2018 Dec 14
6. Sultana, Sharmin et al.CNN-based hierarchical coarse-to-fine segmentation of pelvic
CT images for prostate cancer radiotherapy. Proceedings of SPIE–the InternationalSociety for Optical Engineering vol. 11315 (2020): 113151I.