11institutetext:  Computer Vision Lab, IFI CAIDAS, University of Würzburg,
Germany 11email<EMAIL_ADDRESS>
https://github.com/mv-lab/AISP/ 22institutetext: Couger Inc., Tokyo, Japan

# Real-Time Under-Display Cameras
Image Restoration and HDR on Mobile Devices

Marcos V. Conde 11 Florin Vasluianu 11 Sabari Nathan 22 Radu Timofte 11

###### Abstract

The new trend of full-screen devices implies positioning the camera behind the
screen to bring a larger display-to-body ratio, enhance eye contact, and
provide a notch-free viewing experience on smartphones, TV or tablets. On the
other hand, the images captured by under-display cameras (UDCs) are degraded
by the screen in front of them. Deep learning methods for image restoration
can significantly reduce the degradation of captured images, providing
satisfying results for the human eyes. However, most proposed solutions are
unreliable or efficient enough to be used in real-time on mobile devices. In
this paper, we aim to solve this image restoration problem using efficient
deep learning methods capable of processing FHD images in real-time on
commercial smartphones while providing high-quality results. We propose a
lightweight model for blind UDC Image Restoration and HDR, and we also provide
a benchmark comparing the performance and runtime of different methods on
smartphones. Our models are competitive on UDC benchmarks while using $\times
4$ less operations than others. To the best of our knowledge, we are the first
work to approach and analyze this real-world single image restoration problem
from the efficiency and production point of view.

###### Keywords:

Computational Photography, HDR, Image Restoration, UDC, Low-level vision,
Mobile AI

##  1 Introduction

The consumer demand for smartphones with bezel-free, notch-less display has
sparked a surge of interest from the phone manufacturers in a newly-defined
imaging system, Under-Display Camera (UDC). Besides smartphones, UDC also
demonstrates its practical applicability in other scenarios, _i.e_., for
videoconferencing with UDC TVs, laptops, or tablets, enabling more natural
gaze focus as they place cameras at the center of the displays [20, 44]. The
UDC system consists of a camera module placed underneath and closely attached
to the semi-transparent Organic Light-Emitting Diode (OLED) display [44]. This
solution provides an advantage when it comes to the user experience analysis,
with the full-screen design providing a higher level of comfort, and better
display properties. The disadvantage of this solution is the fact that the
OLED display acts as an obstacle for the light interacting with the camera
sensor, inducing additional reflections, refractions and other connected
effects to the Image Signal Processing (ISP) [7] model characterizing the
camera. Even though the display is partially transparent, there are opaque
regions (i.e. the regions between the display pixels) that substantially
affect the incoming light through a multitude of diffractions that are
extremely difficult to model [20, 44].

In this work, we aim to address the aforementioned issues. We present a novel
solution for the UDC camera image restoration problem, that overcomes the
difficulties induced by the real-world degradation in the UDC images. Our
method provides high-quality results in real-time working conditions (under
one second for image inference). We experiment with one of the world’s first
production UDC device, ZTE Axon 20, which incorporates a UDC system into its
selfie camera [9, 10].

In summary, our contributions are as follows:

  1. 1.

We propose a new U-Net based architecture characterized by a lower number of
parameters, enhancing its performance level by using an attention branch.

  2. 2.

We optimize our model in terms of used parameters, the number of FLOPS and the
used operations, thus being able to achieve real-time performance on current
smartphone GPUs at Full-HD input image resolution.

  3. 3.

We propose a new type of analysis, from a production point of view. This
analysis looks at the behaviour of the different models when deployed on
smartphone hardware platforms targeting different market segments and various
price levels.

##  2 Related Work

In recent years, several works characterized and analyzed the diffraction
effects of UDC systems [28, 34, 9, 20, 18, 23]. Kwon _et al_. [17] modeled the
edge spread function of transparent OLED systems. Qin _et al_. [27] identified
a pixel structure design that can potentially reduce the light diffraction.
Additionally, several works [33, 32] proposed UDC designs for enhanced
interaction with flat displays. Zhou _et al_. [45] and their 2020 ECCV
challenge [44] were the first works that directly addressed this novel
restoration problem using deep learning. Baidu Research team [44] proposed the
Residual dense based on Shade-Correction for T-OLED UDC Image Restoration.
They used the matching coefficient of the input image to do a light intensity
correction, removing the input shade. In [45], the authors devised a Monitor
Camera Imaging System (MCIS) to capture paired images, and solve the UDC image
restoration problem as a blind deconvolution problem. Using the acquired data,
they proposed a fully-supervised UNet [29] like model for this task. More
recently, Feng _et al_. [9] proposed one of the world’s first production UDC
device for data collection, experiments and evaluations. They also proposed a
new model called DISCNet [9] for non-blind restoration, and provided a
benchmark for multiple blind and non-blind methods on their brand-new SYNTH
dataset [9]

However, the aforementioned UDC image restoration works [34, 9, 20, 17, 27]
suffer from several drawbacks. The first drawback is introduced by the usage
of an MCIS system to estimate the camera Point Spread Function (PSF) [15],
given the acquisition of non-realistic paired data to develop a model [26, 1].
In [22], the authors identified MCISs limitations in terms of High Dynamic
Range (HDR), which is crucial for a realistic diffraction artifacts removal.
The second drawback comes as the authors are trying to simulate a UDC system
by manually covering a camera with an OLED display. Even if this setup
provides them with quasi-realistic data, the properties of an actual UDC
system are much more complex than this approximation. Regular events like
tilting, regular rotations or camera motion will produce a PSF that varies
during recording, and thus, their models will suffer when asked to handle
different degradations (as they will rely on a rough approximation of the PSF
given by the simulated data). In [9], Feng _et al_. proposed the usage of one
of the first commercially available UDC devices for data collection, enabling
realistic data conducted experiments, that further enhanced the performance of
the evaluated UDC image restoration models.

Lately, the tendency has switched to focus the efforts in computing reliable
PSFs, to be then used in data synthesis. This enables solving the UDC image
restoration in the fully supervised framework, with various UNet-like models
being used as image-to-image translators. However, the real world domain PSF
is still able to provide useful information for PSF based models. For example,
Feng _et al_. [9] proposed DISCNet, a method which achieved state-of-the-art
results by leveraging the computed PSF to improve performance. They provide a
benchmark for multiple blind and non-blind methods, developed using their real
synthetic UDC data [9].

##### Non-blind Image Restoration.

The Non-Blind Image Restoration implies the usage of the real PSF
characterizing the system in the model design. Once the PSF is known, the
deconvolution approach can be used. Several works have follow this approach
[25, 19, 5, 36], where the authors imposed prior knowledge regarding the
acquisition system to limit the solution space size for the estimation of the
unknown noise model. Lately, different authors [31, 38, 40] have focused their
efforts in the design of neural networks aiming for non-blind image
restoration. For example, Zhang _et al_. [41] proposed SRMD, using a single
neural network to handle multiple degradations. Gu _et al_. [12] proposed a
method called SFTMD, using the Iterative Kernel Correction (IKC) to gradually
correct the existing degradations. Generative Adversarial Networks (GANs) were
used in various works [2, 39, 43], to produce a more realistic reconstruction
by tackling the remaining unknown degradations which are not described by the
PSF. The possibility of using the PSF kernel in the model design was explored
in works like [41]. To date, Feng _et al_. DISCNet [9] is the state-of-the-art
for non-blind UDC image restoration. This method defines a physics-based image
formation model to better understand the degradation. The authors measure the
real-world Point Spread Function (PSF) from smartphone prototypes, and design
a Dynamic Skip Connection Network (DISCNet) to restore the UDC images. This
approach shows strong results on both synthetic and real data samples.

##### Blind Image Restoration.

Solutions tackling the Blind Image Restoration problem do not rely on any
prior knowledge about the Point Spread Function (PSF) to describe the light
transport effects characteristic to a UDC setup. This makes impossible the
usage of the PSF for data synthesis, and the amount of available training data
becomes another limitation of this difficult problem. This type of solution
was proposed by Zhou _et al_. [45, 44] (_i.e_. DE-UNet). Most recently Koh _et
al_. [16] use a two-branch neural network tackling specifically the light
diffraction produced by the pixels grid, and the diffuse intensity caused by
the additional film layers used for the display system.

However, none of all the previous mentioned methods have analyzed the problem
from an efficiency point of view. We believe this is due to the fact that the
current proposed models can generate high-quality results but cannot be
integrated into modern smartphones due to their complexity (_i.e_. number of
FLOPs, memory requirements).

##  3 Proposed Method

We propose two solutions for tackling the UDC image restoration problem. Both
solutions are blind image restoration methods, and therefore, they do not rely
on the PSF information about the camera:

  1. 1.

DRM-UDCNet: A UNet-like model inspired in previous solutions [9, 45, 20, 44].
For this network, we propose a 2.2 Million parameters dual-branch model that
employs our own formulation of Residual Dense Attention Blocks (RDAB) [42].
This represents a reduction of 40% of the number of parameters with respect to
the models that achieved the current state-of-the-art results [9, 45].

  2. 2.

LUDCNet: A compact version of _DRM-UDCNet_ , tailored specifically to be
efficient when deployed on various smartphone devices. As this model is
designed to perform UDC image restoration in real-time conditions on mobile
devices, we optimize it considering SSIM performance and a lower number of
FLOPs. As we will prove in Section 5, this model achieves competitive
performance while processing Full-HD images in real-time on a wide range of
commercial smartphones.

Our methods combine ideas from deblurring [24] and HDR [22, 4] networks, and
attention methods [37, 21]. We believe these tasks are extremely correlated
with the UDC restoration problem and applications. Our method is illustrated
in Figure 1. The designed dual-branch model allows to adapt it depending on
the trade-off between resources and performance. For instance, we can remove
the attention branch and obtain competitive results with a lower inference
time. In exchange of slightly higher inference time, we can improve notably
our performance by using the proposed attention branch.

###  3.1 Method Description

As we show in Figure 1, _DRM-UDCNet_ consists of:

  1. 1.

The main image restoration branch: It is composed of three encoder blocks
($E_{1}$, $E_{2}$, $E_{3}$) and three decoder blocks ($D_{1}$, $D_{2}$,
$D_{3}$). Each encoder block consists of two Dense Residual Modules (DRM) [42,
30] and a downsampling layer. The decoder blocks $D_{1}$, $D_{2}$ have two DRM
followed by a bilinear upsampling layer. Note that the input image spatial
coordinates are mapped using the Coordinate Convolution Layer [21], and we use
skip connections between encoder-decoder blocks [30]. The output of the final
decoder block is mapped into a 3-channel residual using a convolution layer
and _tanh_ activation. We further activate this residual, multiplying it by an
attention map $\mathcal{A}$. This is a technique used in state-of-the-art
single image HDR [22]. We finally apply global residual learning as a standard
technique in image restoration.

  2. 2.

Our attention branch: Similar to HDR methods [22], it aims at generating a
3-channel attention map to control artifacts and hallucination in overexposed
areas. The attention map is generated after applying a CBAM block [37] to the
input image. The final layer is activated using a sigmoid function to produce
a per-channel attention map $\mathcal{A}$.

Figure 1: Architecture of the proposed network. The top of the figure shows
the overall _DRM-UDCNet_ model for blind UDC Image Restoration. The bottom of
the figure presents the Dense Residual Module [42]. Note that we apply pruning
and sequentially remove layers to reduce its size after training. We provide
more details in the supplementary material.

We further compact our challenge model _DRM-UDCNet_ into _LUDCNet_ , which
only uses two DRM blocks and a reduced number of filters. We remove Batch
Normalization (BN) layers, which consume the same amount of GPU memory as
convolutional layers, and they also increase computational complexity [42].
Furthermore, we replace classical CNN blocks by Inverted Linear Residual
Blocks [13]. _LUDCNet_ will process the images at half-resolution, with the
input image being downsampled by a factor of two, and the output image
upsampled to the original resolution.

##  4 Experimental Setup

###  4.1 Implementation details

The original RGB images from SYNTH [9] have resolution $800\times 800$, we
extracted non-overlapping patches of size $400\times 400$. We extracted
patches (instead of downsampling the image) to preserve the high-frequency
details. Moreover, images were transformed from the original domain to the
tone-mapped domain using the following function $f(x)=x/(x+0.25)$. Therefore
the pixel intensity is in the range $[0,1)$. To avoid artifacts, the output of
our models is clipped into range $[0,1-p]$ , where $p$ is $10^{-5}$. Our
models were implemented in Tensorflow 2 and trained using a single Tesla P100
GPU (16Gb). Some experiments were performed using a TPU v3-8. We used Adam
optimizer with default hyper-parameters, and an initial learning rate of 1e-3.
The learning rate is reduced by 50% during plateaus up to a minimum learning
rate of 1e-6. We use basic augmentations: horizontal flip, vertical flip, and
rotations. We set 4 as mini-batch size and trained using only the SYNTH
dataset [9] to convergence for a few days.

##### Loss function

The model is trained using a weighted sum of $\mathcal{L}1$ loss, SSIM loss,
and gradient loss. We define the loss functions as follows:

$\mathcal{L}_{1}=\frac{1}{N}\sum_{i=1}^{N}\left\|\mathcal{I}-\mathcal{\hat{I}}\right\|_{1}$
(1)

where $\mathcal{D}$ is the input degraded image,
$f(\mathcal{D})=\hat{\mathcal{I}}$ is the restored image using our model $f$,
and $\mathcal{I}$ is the ground truth image. The SSIM loss is defined as:

$\mathcal{L}_{SSIM}=\frac{1}{N}\sum_{i=1}^{N}(1-SSIM(\mathcal{I},\hat{\mathcal{I}}))$
(2)

where SSIM is the structural similarity index function defined in [35].

The gradient loss is the $\mathcal{L}_{1}$ distance between the gradients of
the restored image and the ground-truth, in the $xy$ axis:

$\mathcal{L}_{Grad}=\frac{1}{N}\sum_{i=1}^{N}\left\|\nabla_{x}\mathcal{I}-\nabla_{x}\hat{\mathcal{I}}\right\|_{1}+\left\|\nabla_{y}\mathcal{I}-\nabla_{y}\hat{\mathcal{I}}\right\|_{1}$
(3)

Therefore, the final loss function is:

$\mathcal{L}=0.1~{}\mathcal{L}_{SSIM}+\mathcal{L}_{1}+\mathcal{L}_{Grad}$ (4)

###  4.2 Experimental Results

####  4.2.1 Quantitative Results

We evaluate our models on two different benchmarks: (i) the SYNTH dataset [9],
also used in the “UDC MIPI 2022 Challenge” [10]; (ii) the T-OLED UDC 2020
Dataset [20, 44]. We will mainly focus on the SYNTH dataset [9], where we
compare against current state-of-the-art methods (excluding contemporary
methods, and methods better than DISCNet [9] in the 2022 Challenge). Note that
these results are not fully reproducible since multiple approaches do not
provide open-sourced code. Moreover, the SYNTH dataset [9] test set (and PSF)
is public, which makes difficult a fair comparison. As we show in Tables 1 and
2, we achieve competitive results in both benchmarks, while using models with
fewer parameters than the other competitive methods. As an example, we can
look at the comparison against DISCNet [9] base version, as we used the same
initial setup to build up our model. We can see how our method outperforms
this one by almost 2dBs. We extend this analysis in Section 5.

Table 1: Quantitative Results for Under-display Camera (UDC) Image Restoration
using the SYNTH dataset [9]. We show fidelity and perceptual metrics. (*)
indicates methods proposed at the 2022 UDC Challenge [10]. We consult some
numbers from [9].

Method | PSNR (dB) $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ | # Params | Uses PSF
---|---|---|---|---|---
Wiener Filter (WF) [25] | 27.30 | 0.83 | 0.330 | - | yes
SRMDNF [9] | 34.80 | 0.96 | 0.036 | 1.5 | yes
SFTMD [9] | 42.35 | 0.98 | 0.012 | 3.9 | yes
DISCNet (PSF) [9] | 42.77 | 0.98 | 0.012 | 3.8 | yes
RDUNet [42, 30] | 34.37 | 0.95 | 0.040 | 8.1 | no
DE-UNet [45] | 38.11 | 0.97 | 0.021 | 9.0 | no
DISCNet (w/o PSF) [9] | 38.55 | 0.97 | 0.030 | 2.0 | no
RushRushRush * | 39.52 | 0.98 | 0.021 | - | no
eye3 * | 36.69 | 0.97 | 0.032 | - | no
FMS Lab * | 35.77 | 0.97 | 0.045 | - | no
EDLC2004 * | 35.50 | 0.96 | 0.045 | - | no
SAU_LCFC * | 32.75 | 0.96 | 0.056 | - | no
DRM-UDCNet | 37.45 | 0.97 | 0.030 | 2.0 | no
DRM-UDCNet + Attn | 40.21 | 0.98 | 0.020 | 2.9 | no

Table 2: Quantitative Results and comparison of methods at the T-OLED UDC 2020
Challenge [20, 44]. TT denotes the training time, and IT the inference time
per image. Note that we only compare with simple methods trained in similar
conditions.

Team | PSNR (dB) $\uparrow$ | SSIM $\uparrow$ | TT(h) | IT (s/frame) | CPU/GPU
---|---|---|---|---|---
CILab IITM | $36.91$ | $0.9734$ | 96 | 1.72 | 1080 Ti
lyl | $36.72$ | $0.9776$ | 72 | 3.0 | -
DRM-UDCNet | $36.50$ | $0.9730$ | 14 | 1.0 | Tesla P100
Image Lab | $34.35$ | $0.9645$ | - | 1.6 | 1080 Ti
San Jose Earthquakes | $33.78$ | $0.9324$ | 18 | 180 | -
UNet [44] | $32.42$ | $0.9343$ | 24 | - | Titan X
DeP [44] | $28.50$ | $0.9117$ | 24 | - | Titan X

####  4.2.2 Qualitative Results

Figure 2 provides a visual comparison for the results of our model, compared
to other state-of-the-art methods. As you can see, our method is able to
provide a better reconstruction, with results characterized by better
properties in terms of textures, colors and the geometric properties of the
hallucinated objects. In Figure 6, we provide equivalent images for a
comparison of the model prediction to the ground truth images, given each
input sample. On the last column (Fig. 6), we provide scaled error maps
computed between the model prediction and the ground truth, observing that the
most of the reconstruction error is concentrated in areas characterized by
high light intensity, and therefore affected by over-exposure and glare.

---
Figure 2: Visual comparison on synthetic validation images. Our method
restores fine details, recovers information from HDR, and renders high
perceptual quality results without notable artifacts. Images selection from
DISCNet [9]. Zoom in for better view.

##  5 Efficiency Analysis

We evaluate the performance of different methods using three smartphone
devices. In Table 3, we provide all the details regarding the target devices
used in our benchmark. We choose two different smartphones that target the
_entry-level_ segment of the market (1,2), and we also consider a mobile phone
designed for the _flagship_ segment (3). To observe the evolution of the
entry-level smartphones in terms of AI deployment capabilities, we compare the
Samsung A50 to the OnePlus Nord 2 5G. As a reference metric to measure the
performance of each mobile device in the scenario of the deep learning models
deployment, we provide the AI Score [14]. This AI score is computed over a set
of experiments (_i.e_. category recognition, semantic segmentation or image
enhancement), where inference precision and hardware-software behaviour are
observed and quantified into a score metric depending on the importance of
each of the factors.

All the tests are described by the authors in [14], and are available for the
Android smartphones through the application _AI Benchmark_ [14], which also
provides support for deep learning models deployment on different hardware
options (CPU vs. GPU), and with various precision types (_i.e_. FP16, FP32).
Models are tested on CPU and GPU, and FP16 precision after default Tensorflow-
lite conversion and optimization, using as input a tensor image of different
resolutions. The FLOPs are calculated at each resolution. We show in Figure 4
the evaluation process using _AI Benchmark app_ [14].

DISCNet [9] cannot be directly converted to the required TFlite format due to
unsupported operations. For this reason, we use DISCNet “baseline (a)” [9] to
evaluate the model, this provides a fair lower-bound of the complete method’s
performance. Note that for an input of resolution $800\times 800$, the authors
report 364 GFLOPs [9], meanwhile in our re-implementation we obtain 442
GFLOPs111FLOPs measured using keras-flops library. However, for a fair
comparison, we compare with a different possible implementation with fewer
operations. SRMDNF [45] and SFTMD [45] cannot be directly converted to the
requested standard format [14], therefore, we use a canonical UNet model with
the same number of FLOPs and memory requirements to approximate their
performance.

As we show in Table 4, our proposed model _LUDCNet_ can process Full-HD images
in real-time in commercial mobile devices GPUs, being $\times 5$ faster than
DISCNet [9]. Our model can also process low-resolution images on CPU in real-
time (under 1 second per image), and its performance only decays $0.04$ with
respect to DISCNet [9] in terms of SSIM [35].

In Figure 5 we provide qualitative samples in challenging scenarios, our model
is able to improve notably the perceptual quality of the UDC degraded images
in real-time. Note that in the perception-distortion tradeoff [3], we focus on
perceptual metrics [11, 6], as we aim for pleasant results for users.

Table 3: Description of the selected commercial smartphone devices.

Phone Model | Launch | Chipset | CPU | GPU | RAM (GB) | AI Score $\uparrow$
---|---|---|---|---|---|---
(# 1) Samsung A50 | 03/2019 | Exynos 9610 | 8 cores | Mali-G71 GP3 | 4 | 45.4
(# 2) OnePlus Nord 2 5G | 07/2021 | MediaTek Dimensity 1200 | 8 cores | Mali-G77 MC9 | 8 | 194.3
(# 3) OnePlus 8 Pro | 04/2020 | Qualcomm Snapdragon 865 5G | 8 cores | Adreno 650 | 12 | 137.0

Table 4: Efficiency benchmark for UDC Image Restoration on different
commercial smartphones. We show the performance of different SOTA methods in
terms of runtime at different image resolutions, device architectures and
running scenarios (CPU, GPU). Runtimes are the average of at least 5
iterations. SRMDNF and SFTMD [45] failed the test (✗) due to exceed the memory
limit. Our method is the only one that can perform Full-HD real-time
processing, while we achieving high perceptual quality results. SSIM
perceptual metric results are reported for the SYNTH dataset [9].

|  |  |  | Runtime (s)
---|---|---|---|---
Method | FLOPs $\downarrow$ | Resolution | SSIM $\uparrow$ | Phone #1 | Phone #2 | Phone #3
name | (G) | (px) |  | CPU | GPU | CPU | GPU | CPU | GPU
SRMDNF [45] | 951 | 800$\times$800 | 0.965 | ✗ | ✗ | ✗ | ✗ | ✗ | ✗
SFTMD [45] | 2460 | 800$\times$800 | 0.986 | ✗ | ✗ | ✗ | ✗ | ✗ | ✗
DISCNet (a) [9] | 442 | 800$\times$800 | 0.974 | 52.4 | 4.0 | 13.5 | 0.63 | 15.6 | 1.4
DISCNet (a)* [9] | 300 | 800$\times$800 | 0.974 | 13.4 | 2.3 | 4.4 | 0.37 | 5.5 | 0.9
DISCNet (a) [9] | 256 | 256$\times$256 | 1.1 | 0.41 | 0.3 | 0.07 | 0.6 | 0.14
DISCNet (a) [9] | 1434 | 1920$\times$1080 | 1055 | 13.7 | 261.5 | 2.0 | 61.1 | 5.3
LUDCNet |  100 | 1920$\times$1080 | 0.930 | 9.1 | 1.6 | 3.1 | 0.26 | 3.8 | 0.5
LUDCNet | 30 | 800$\times$800 | 2.6 | 0.5 | 0.9 | 0.095 | 1.2 | 0.19
LUDCNet | 3 | 256$\times$256 | 0.3 | 0.05 | 0.1 | 0.023 | 0.1 | 0.02

Figure 3: Comparison of parameters, FLOPs and performance of state-of-the art methods for UDC Image Restoration. Our method achieves high perceptual quality results while being notable smaller and computationally “cheap”. |  |
---|---|---
Figure 4: AI Benchmark [14] model performance evaluation. |  |
---|---|---
|  |
|  |
|  |
Input | DRM-UDCNet | LUDCNet
Figure 5: Qualitative samples from the SYNTH test benchmark [9]. Our models can improve substantially the quality of the input UDC degraded images in indoor and outdoor, natural or artificial illumination, and day-night scenarios. As we explain in Section 5, our lightweight model _LUDCNet_ can process Full-HD images in real-time. |  |  |
---|---|---|---
|  |  |
Input | Prediction | Reference | Error Map
Figure 6: SYNTH Dataset [9] validation samples. Images visualized after
applying tone-mapping. The error maps are multiplied by $16$ for a better
visualization.

Perceptual Image Enhancement for Smartphone Applications
Besides the proposed model, we also developed a considerably more powerful
architecture called LPIENet [8], and a extended efficiency benchmark [8]. Our
new model can process 2K resolution images under 1 second in mid-level
commercial smartphones, and tackles: denoising, deblurring and UDC image
restoration.

##  6 Conclusions

First, we develop DRM-UDCNet, a U-Net based model to solve this problem and
achieve the most competitive performance on well-known benchmarks. Next, we
focus on compacting such model, aiming to be the first work to tackle real-
time UDC image restoration in the smartphone deployment scenario. As a result,
our model LUDCNet uses $\times 4$ less operations than other state-of-the-art
methods. Efficiency analysis shows that our compact model is $\times 5$ faster
than DISCNet. We set up a new benchmark to show our method’s potential, using
three commercial smartphones showcasing different hardware architectures and
capabilities. We show that our model does not depend on specific hardware,
achieving real-time performance on current mobile phone CPUs at lower
resolutions. To the best of our knowledge, we are the first work to approach
and analyze the UDC restoration problem from the efficiency and production
point of view.

Acknowledgments This work was partly supported by The Alexander von Humboldt
Foundation (AvH).

## References

  * [1] Asif, M.S., Ayremlou, A., Sankaranarayanan, A., Veeraraghavan, A., Baraniuk, R.G.: Flatcam: Thin, lensless cameras using coded aperture and computation. IEEE Transactions on Computational Imaging 3(3), 384–397 (2016)
  * [2] Bell-Kligler, S., Shocher, A., Irani, M.: Blind super-resolution kernel estimation using an internal-gan. In: Advances in Neural Information Processing Systems. pp. 284–293 (2019)
  * [3] Blau, Y., Michaeli, T.: The perception-distortion tradeoff. In: Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 6228–6237 (2018)
  * [4] Catley-Chandar, S., Tanay, T., Vandroux, L., Leonardis, A., Slabaugh, G., Pérez-Pellitero, E.: Flexhdr: Modelling alignment and exposure uncertainties for flexible hdr imaging. arXiv preprint arXiv:2201.02625 (2022)
  * [5] Cho, S., Wang, J., Lee, S.: Handling outliers in non-blind image deconvolution. In: 2011 International Conference on Computer Vision. pp. 495–502. IEEE (2011)
  * [6] Conde, M.V., Burchi, M., Timofte, R.: Conformer and blind noisy students for improved image quality assessment. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 940–950 (2022)
  * [7] Conde, M.V., McDonagh, S., Maggioni, M., Leonardis, A., Pérez-Pellitero, E.: Model-based image signal processors via learnable dictionaries. Proceedings of the AAAI Conference on Artificial Intelligence 36(1), 481–489 (Jun 2022). https://doi.org/10.1609/aaai.v36i1.19926, https://ojs.aaai.org/index.php/AAAI/article/view/19926
  * [8] Conde, M.V., Vasluianu, F., Vazquez-Corral, J., Timofte, R.: Perceptual image enhancement for smartphone real-time applications. In: Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (2023)
  * [9] Feng, R., Li, C., Chen, H., Li, S., Loy, C.C., Gu, J.: Removing diffraction image artifacts in under-display camera via dynamic skip connection network. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). pp. 662–671 (June 2021)
  * [10] Feng, R., Li, C., Zhou, S., Sun, W., Zhu, Q., Jiang, J., Yang, Q., Loy, C.C., Gu, J.: Mipi 2022 challenge on under-display camera image restoration: Methods and results. arXiv preprint arXiv:2209.07052 (2022)
  * [11] Gu, J., Cai, H., Dong, C., Ren, J.S., Timofte, R., Gong, Y., Lao, S., Shi, S., Wang, J., Yang, S., et al.: Ntire 2022 challenge on perceptual image quality assessment. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 951–967 (2022)
  * [12] Gu, J., Lu, H., Zuo, W., Dong, C.: Blind super-resolution with iterative kernel correction. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. pp. 1604–1613 (2019)
  * [13] Howard, A.G., Zhu, M., Chen, B., Kalenichenko, D., Wang, W., Weyand, T., Andreetto, M., Adam, H.: Mobilenets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint arXiv:1704.04861 (2017)
  * [14] Ignatov, A.D., Timofte, R., Chou, W., Wang, K., Wu, M., Hartley, T., Gool, L.V.: Ai benchmark: Running deep neural networks on android smartphones. In: ECCV Workshops (2018)
  * [15] Joshi, N., Szeliski, R., Kriegman, D.J.: Psf estimation using sharp edge prediction. In: 2008 IEEE Conference on Computer Vision and Pattern Recognition. pp. 1–8. IEEE (2008)
  * [16] Koh, J., Lee, J., Yoon, S.: Bnudc: A two-branched deep neural network for restoring images from under-display cameras. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). pp. 1950–1959 (June 2022)
  * [17] Kwon, H.J., Yang, C.M., Kim, M.C., Kim, C.W., Ahn, J.Y., Kim, P.R.: Modeling of luminance transition curve of transparent plastics on transparent oled displays. Electronic Imaging 2016(20), 1–4 (2016)
  * [18] Kwon, K., Kang, E., Lee, S., Lee, S.J., Lee, H.E., Yoo, B., Han, J.J.: Controllable image restoration for under-display camera in smartphones. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). pp. 2073–2082 (June 2021)
  * [19] Levin, A., Weiss, Y., Durand, F., Freeman, W.T.: Understanding and evaluating blind deconvolution algorithms. In: 2009 IEEE Conference on Computer Vision and Pattern Recognition. pp. 1964–1971. IEEE (2009)
  * [20] Lim, S., Zhou, Y., Emerton, N., Ghioni, L., Large, T., Bathiche, S.: Camera in display (2020 (accessed Nov 9, 2020)), https://www.microsoft.com/applied-sciences/projects/camera-in-display
  * [21] Liu, R., Lehman, J., Molino, P., Petroski Such, F., Frank, E., Sergeev, A., Yosinski, J.: An intriguing failing of convolutional neural networks and the coordconv solution. Advances in neural information processing systems 31 (2018)
  * [22] Liu, Y.L., Lai, W.S., Chen, Y.S., Kao, Y.L., Yang, M.H., Chuang, Y.Y., Huang, J.B.: Single-image hdr reconstruction by learning to reverse the camera pipeline. In: IEEE Conference on Computer Vision and Pattern Recognition (2020)
  * [23] Luo, J., Ren, W., Wang, T., Li, C., Cao, X.: Under-display camera image enhancement via cascaded curve estimation. IEEE Transactions on Image Processing 31, 4856–4868 (2022)
  * [24] Nah, S., Son, S., Lee, S., Timofte, R., Lee, K.M.: Ntire 2021 challenge on image deblurring. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 149–165 (2021)
  * [25] Orieux, F., Giovannelli, J.F., Rodet, T.: Bayesian estimation of regularization and point spread function parameters for wiener–hunt deconvolution. JOSA A 27(7), 1593–1607 (2010)
  * [26] Peng, Y., Sun, Q., Dun, X., Wetzstein, G., Heidrich, W., Heide, F.: Learned large field-of-view imaging with thin-plate optics. ACM Trans. Graph. 38(6), 219–1 (2019)
  * [27] Qin, Z., Tsai, Y.H., Yeh, Y.W., Huang, Y.P., Shieh, H.P.D.: See-through image blurring of transparent organic light-emitting diodes display: calculation method based on diffraction and analysis of pixel structures. Journal of Display Technology 12(11), 1242–1249 (2016)
  * [28] Qin, Z., Xie, J., Lin, F.C., Huang, Y.P., Shieh, H.P.D.: Evaluation of a transparent display’s pixel structure regarding subjective quality of diffracted see-through images. IEEE Photonics Journal 9(4), 1–14 (2017)
  * [29] Ronneberger, O., Fischer, P., Brox, T.: U-net: Convolutional networks for biomedical image segmentation. In: International Conference on Medical image computing and computer-assisted intervention. pp. 234–241. Springer (2015)
  * [30] Ronneberger, O., Fischer, P., Brox, T.: U-net: Convolutional networks for biomedical image segmentation. In: International Conference on Medical image computing and computer-assisted intervention. pp. 234–241. Springer (2015)
  * [31] Schuler, C.J., Christopher Burger, H., Harmeling, S., Scholkopf, B.: A machine learning approach for non-blind image deconvolution. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. pp. 1067–1074 (2013)
  * [32] Suh, S., Choi, C., Park, D., Kim, C.: 50.2: Adding depth-sensing capability to an oled display system based on coded aperture imaging. In: SID Symposium Digest of Technical Papers. vol. 44, pp. 697–700. Wiley Online Library (2013)
  * [33] Suh, S., Choi, C., Yi, K., Park, D., Kim, C.: P-135: An lcd display system with depth-sensing capability based on coded aperture imaging. In: SID Symposium Digest of Technical Papers. vol. 43, pp. 1574–1577. Wiley Online Library (2012)
  * [34] Tang, Q., Jiang, H., Mei, X., Hou, S., Liu, G., Li, Z.: 28-2: Study of the image blur through ffs lcd panel caused by diffraction for camera under panel. In: SID Symposium Digest of Technical Papers. vol. 51, pp. 406–409. Wiley Online Library (2020)
  * [35] Wang, Z., Bovik, A., Sheikh, H., Simoncelli, E.: Image quality assessment: from error visibility to structural similarity. IEEE Transactions on Image Processing 13(4), 600–612 (2004). https://doi.org/10.1109/TIP.2003.819861
  * [36] Whyte, O., Sivic, J., Zisserman, A.: Deblurring shaken and partially saturated images. International journal of computer vision 110(2), 185–201 (2014)
  * [37] Woo, S., Park, J., Lee, J.Y., Kweon, I.S.: Cbam: Convolutional block attention module. In: Proceedings of the European conference on computer vision (ECCV). pp. 3–19 (2018)
  * [38] Xu, L., Ren, J.S., Liu, C., Jia, J.: Deep convolutional neural network for image deconvolution. In: Advances in neural information processing systems. pp. 1790–1798 (2014)
  * [39] Yuan, Y., Liu, S., Zhang, J., Zhang, Y., Dong, C., Lin, L.: Unsupervised image super-resolution using cycle-in-cycle generative adversarial networks. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops. pp. 701–710 (2018)
  * [40] Zhang, J., Pan, J., Lai, W.S., Lau, R.W., Yang, M.H.: Learning fully convolutional networks for iterative non-blind deconvolution. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. pp. 3817–3825 (2017)
  * [41] Zhang, K., Zuo, W., Zhang, L.: Learning a single convolutional super-resolution network for multiple degradations. In: IEEE Conference on Computer Vision and Pattern Recognition. vol. 6 (2018)
  * [42] Zhang, Y., Tian, Y., Kong, Y., Zhong, B., Fu, Y.: Residual dense network for image super-resolution. In: Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 2472–2481 (2018)
  * [43] Zhou, R., Susstrunk, S.: Kernel modeling super-resolution on real low-resolution images. In: Proceedings of the IEEE International Conference on Computer Vision. pp. 2433–2443 (2019)
  * [44] Zhou, Y., Kwan, M., Tolentino, K., Emerton, N., Lim, S., Large, T., Fu, L., Pan, Z., Li, B., Yang, Q., et al.: Udc 2020 challenge on image restoration of under-display camera: Methods and results. In: European Conference on Computer Vision. pp. 337–351. Springer (2020)
  * [45] Zhou, Y., Ren, D., Emerton, N., Lim, S., Large, T.: Image restoration for under-display camera. In: 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). pp. 9175–9184 (2021). https://doi.org/10.1109/CVPR46437.2021.00906


# MiTU-Net: A fine-tuned U-Net with SegFormer backbone for segmenting pubic
symphysis-fetal head

Fangyijie Wang
School of Medicine
University College Dublin
<EMAIL_ADDRESS>
& Guénolé Silvestre
School of Computer Science
University College Dublin
<EMAIL_ADDRESS>
& Kathleen Curran
School of Medicine
University College Dublin
<EMAIL_ADDRESS>

###### Abstract

Ultrasound measurements have been examined as potential tools for predicting
the likelihood of successful vaginal delivery. The angle of progression (AoP)
is a measurable parameter that can be obtained during the initial stage of
labor. The AoP is defined as the angle between a straight line along the
longitudinal axis of the pubic symphysis (PS) and a line from the inferior
edge of the PS to the leading edge of the fetal head (FH). However, the
process of measuring AoP on ultrasound images is time consuming and prone to
errors. To address this challenge, we propose the Mix Transformer U-Net (MiTU-
Net) network, for automatic fetal head-pubic symphysis segmentation and AoP
measurement. The MiTU-Net model is based on an encoder-decoder framework,
utilizing a pre-trained efficient transformer to enhance feature
representation. Within the efficient transformer encoder, the model
significantly reduces the trainable parameters of the encoder-decoder model.
The effectiveness of the proposed method is demonstrated through experiments
conducted on a recent transperineal ultrasound dataset. Our model achieves
competitive performance, ranking 5th compared to existing approaches. The
MiTU-Net presents an efficient method for automatic segmentation and AoP
measurement, reducing errors and assisting sonographers in clinical practice.
Reproducibility: Framework implementation and models available on
https://github.com/13204942/MiTU-Net.

_Keywords_ Fetal Ultrasound $\cdot$ Fine-tuning $\cdot$ Transformer $\cdot$
Semantic Segmentation

##  1 Introduction

Assessing the progression of labor is vital to ensure proper advancement in
peripartum care. This assessment enables early detection of any deviations
from the expected course, aiming to minimize potential complications for both
the mother and fetus. The International Society of Ultrasound in Obstetrics
and Gynecology (ISUOG) recommends the implementation of ultrasound biometric
technology, particularly the angle of progression (AoP), when assessing the
fetal craniopubic symphysis structure Ghi et al. (2018), shown in Figure 1.
Thus, the measurement of AoP proves to be a superior method for assessing
labor progression. AOP is defined as the angle between a straight line
extending along the longitudinal axis of the pubic symphysis (PS) and a line
connecting the inferior edge of the PS to the leading edge of the fetal head
(FH) Lu et al. (2022a). The current measurement of AoP is manual assessments
conducted by experienced physicians. However, this process is time-consuming,
cumbersome, and prone to measurement errors. To address this problem, an
algorithm for the automatic measurement of AoP is highly demand in clinical
practice. This contribution improves the accuracy of labor progress
assessment. The accurate segmentation of the FH and PS regions from ultrasound
images is an essential prerequisite for the automated measurement of AoP.

Figure 1: The illustration of the process of measuring the angle of
progression (AoP) using transperineal ultrasound. (A) Schematic diagram of
calculating AoP. (B) An image showing the symphysis pubis and fetal head.

In recent years, deep learning (DL) technology experienced significant
advancements. In the field of medical imaging analysis, convolutional neural
networks (CNNs), a type of deep learning technology, have demonstrated
impressive performance in a wide range of medical imaging applications. Three
notable architectures are fully convolutional networks (FCNs) Long et al.
(2015), U-Net Ronneberger et al. (2015), and the three-dimensional V-Net
Milletari et al. (2016). These architectures have gained recognition as
classic models for analyzing medical images.

Recently, researchers have made attempts to combine CNN and Transformer
Vaswani et al. (2017b) architectures for the purpose of medical image
segmentation. Chen et al. (2021) propose to combine U-Net and Transformer to
obtain TransUnet for medical image segmentation. To analysis pregnancy
ultrasound images with automatic DL algorithms, Chen et al. (2024) propose a
model incorporating a dual attention module, a multi-scale feature screening
module and a direction guidance block for segmenting the fetal head and pubic
symphysis. This model offers significant advantages in extracting global
contextual information and refining segmentation results with direction
guidance.

This paper introduces an efficient segmentation network based on the encoder-
decoder framework, U-Net, for the automated segmentation of ultrasound images
and the measurement of AoP. We enhance the U-Net architecture by replacing its
encoder path with a pre-trained Mix Transformer Xie et al. (2021). The Mix
Transformer encoder includes attention mechanisms to enhance feature
representation. Furthermore, it reduces computational cost and memory usage by
selecting and processing only important parts of the input. Additionally, it
significantly reduces the number of trainable parameters from 31 million to 5
million in comparison to the traditional U-Net.

The rest of this paper is organized as follows. Section 2 presents the
detailed description of the proposed model. Section 3 focuses on the analysis
and evaluation of the experimental results.

##  2 Methodology

This section introduces an efficient and powerful segmentation fine-tuned
model with low computationally demanding modules, MiTU-Net, see Figure 2. The
model consists of a U-Net architecture with a pre-trained Mix Transformer
encoder for segmenting pubic symphysis-fetal head on a ultrasound dataset,
JNU-IFM Lu et al. (2022b). The U-Net architecture is followed the design from
Ronneberger Ronneberger et al. (2015) in 2015. The contracting path
("encoder") of the U-Net is replaced with a pre-trained Mix Transformer
encoder (MiT-B0) designed by Xie in SegFormer framework Xie et al. (2021) in
2021. The framework is developed by programming language Python, library
PyTorch, and library Segmentation Models Iakubovskii (2019).

###  2.1 Model Achitecture

Our U-Net architecture contains the contracting path ("encoder") and the
expansive path ("decoder"). They both have the same depth of four layers.
Between encoder and decoder, we use skip connections to concatenating features
from the Mix Transformer encoder to recover spatial information lost during
downsampling. We first transpose the original ultrasound image of size
$3\times 256\times 256$ to size of $256\times 256\times 3$. Then the Mix
Transformer encoder divides the ultrasound image into small patches of size
$4\times 4$. These patches are the input to the Mix Transformer encoder to
obtain multi-level features at ${\\{1/4,1/8,1/16,1/32\\}}$ of the original
image resolution. The encoder passes the multi-level features to the decoder
to predict the segmentation mask at a $\frac{H}{4}\times\frac{W}{4}\times
N_{cls}$ resolution, where $N_{cls}$ is the number of categories. The JNU-IFM
dataset contains three categories labeled as 0, 1, or 2, with 0 representing
the background, 1 representing the pubic symphysis, and 2 representing the
fetal head.

#### Feature Representation

The Mix Transformer encoder is a hierarchical type of Transformer encoder that
produces multi-level features. It provides both high-resolution coarse
features and low-resolution fine features. Specifically, the encoder creates a
hierarchical feature map consisting of a resolution of
$\frac{H}{2^{i+1}}\times\frac{W}{2^{i+1}}\times C_{i}$, where
$i\in\\{1,2,3,4\\}$, by merging multiple patches extracted from the input
ultrasound image.

#### Overlapped Patch Merging

We reproduce the Overlapped Patch Merging (OPM) process proposed by Xie Xie et
al. (2021) in the patch merging process for SegFormer. It requires $K,S,$ and
$P$, where $K$ is the patch size, $S$ is the stride between two adjacent
patches, and $P$ is the padding size. In the first OPM step, we set
$K=7,S=4,P=3$, and $K=3,S=2,P=1$ for the rest of OPM steps to perform patch
merging to produces features.

#### Self-Attention

We apply the multi-head self-attention process Vaswani et al. (2017a), each of
the heads $Q,K,V$ have the same dimensions $H\times W\times C$, where $H$ is
the height of input image, $W$ is the width of input image, and $C$ is the
number of channels. The self-attention is estimated as:

$\operatorname{Attention}(Q,K,V)=\operatorname{Softmax}\left(\frac{QK^{\mathrm{\top}}}{\sqrt{d_{head}}}\right)V$
(1)

The computational complexity of this process is $O\left(N^{2}\right)$,

#### Feed-forward Network

We use a feed-forward network (FFN) after the the self-attention module. This
FFN contains a $3\times 3$ depth-wise convolution that can provide positional
information for Transformers sufficiently and reduce computational cost Xie et
al. (2021). It is formulated as follows:

$\begin{split}\mathbf{\hat{x}}_{out}&=\operatorname{Linear}(C_{in},C_{out})(\mathbf{x}_{in})\\\
\mathbf{x}_{out}&=\operatorname{Linear}(C_{out},C_{in})(\operatorname{GELU}(\operatorname{DWConv}_{3\times
3}(\mathbf{\hat{x}}_{out})))\\\ \end{split}$ (2)

where $\mathbf{x}_{in}$ is the feature map from the self-attention module, and
Linear $\left(C_{\text{in }},C_{\text{out }}\right)(\cdot)$ is a linear layer
with $C_{in}$ and $C_{out}$ as input and output vector dimensions
respectively. $\operatorname{DWConv}_{3\times 3}$ refers to the $3\times 3$
depth-wise convolution.

###  2.2 Decoder

The decoder path is symmetric to the encoder path. In total it has four
convolutional layers. Every layer consists of an upsampling of the feature map
followed by a $2\times 2$ up-convolution that halves the number of feature
channels, a concatenation with the correspondingly cropped feature map from
the encoder path, and two $3\times 3$ convolutions, each followed by a ReLU.
The bottom of decoder path is a single $3\times 3$ ("segmentation head") up-
convolution followed by a Softmax that outputs asegmentation map with 3
channels. We formulate the decoder’s single layer as:

$\begin{split}\hat{X}_{out}&={\operatorname{Conv}_{3\times
3}}(C_{F},C_{out})(F)\\\
\hat{X}_{out}&={\operatorname{ReLU}}({\operatorname{BatchNorm2d}}(\hat{X}_{out}))\\\
X_{out}&={\operatorname{Conv}_{3\times 3}}(C_{out},C_{out})(\hat{X}_{out})\\\
X_{out}&={\operatorname{ReLU}}(\operatorname{BatchNorm2d}(X_{out}))\\\
\end{split}$ (3)

where $F$ refers to the feature maps from Mix Transformer encoder path
concatenated with the upsampling feature maps from decoder path, $C_{F}$
refers to the dimensions of the concatenated feature maps $F$, and $C_{out}$
is the output vector dimensions.

###  2.3 Data pre-processing and augmentation

The JNU-IFM dataset has 4000 ultrasound images in total. All of them are used
for training. We randomly select 30 ultrasound ultrasound images for
validation during training. The size of images is $3\times 256\times 256$. We
apply one hot encoding on ground truth masks for semantic segmentation,
accordingly the ground truth prediction has 3 channels, background prediction,
pubic symphysis prediction, and fetal head prediction.

Initially, we are inspired by the regularization technique of randomly masking
out square regions of input during training DeVries and Taylor (2017),
therefore, we randomly mask out $4\times 4$ square regions of the input
images. The number of masked regions is in range $[0,4]$. Additionally, we
pick a random angle from $[-25^{\circ},25^{\circ}]$ to rotate the input image.
Also, we randomly flip the input image horizontally with 50% probability, and
vertically with 30% probability. All of input pixels are normalized with mean
values [0.0, 0.0, 0.0], and standard deviation values [1.0, 1.0, 1.0]. The
validation images have no augmentation, but their input pixels are also
normalized with mean values [0.0, 0.0, 0.0], and standard deviation values
[1.0, 1.0, 1.0].

###  2.4 Fine-tuning

We implemented all of our experiments using PyTorch, then we trained the
entire U-Net with pre-trained MiT-B0 encoder with 20 epochs from scratch. The
training dataset and validation dataset both have a batch size of 10. The Adam
optimiser was used in training processes with a learning rate of $1e-4$. All
training processes were performed on an NVIDIA Tesla T4 graphics card.

#### Loss Function

During the training process, we construct the loss function for this
segmentation task with the Binary Cross Entropy loss function and Jaccard
Index. Given a predicted mask $\hat{Y}$ and a ground truth mask $Y$, the loss
function can be formulated as follows:

$\mathcal{L}=\operatorname{BCE}(Y,\hat{Y})+(1-\operatorname{J}(Y,\hat{Y}))$
(4)

where $\operatorname{BCE}$ refers to the Binary Cross Entropy loss function,
and $\operatorname{J}$ refers to the Jaccard Index (or Intersection over
Union).

#### Evaluation Metrics

The typical metrics applied to evaluate the performance of this U-Net
architecture are Pixel Accuracy (PA), Dice coefficient Dice (1945), and Mean
Intersection over Union (IoU) Jaccard (1912). Mean IoU is defined as the
average IoU over all classes $K$. These metrics can be computed using true
positive(TP), false positive (FP), true negative (TN) and false negative (FP)
of the segmentation results Tavakoli et al. (2021). TP represents the pixels
considered as being in the object and being really in the object, conversely,
TN is the pixels outside the object both in the segmentation and the ground
truth. FP represents the pixels considered by the segmentation in the object,
but which in reality are not part of it. Finally, FN represents the pixels of
the object that the segmentation has classified outside.

Figure 2: The illustration of the proposed MiTU-Net for automatic fetal head
(FH)-pubic symphysis (PS) segmentation.

##  3 Experiments Results

###  3.1 Evaluation Metrics

The typical metrics applied to evaluate the performance of segmentation models
are Dice similarity coefficient (DSC) Dice (1945), Hausdorff Distance (HD)
Birsan and Tiba (2006), Average Surface Distance (ASD) Yeghiazaryan and
Voiculescu (2018), and the AoP difference ($\Delta$AoP) between predicted and
manually measured AoP for prediction. The DSC is employed to assess the level
of similarity between a predicted segmentation mask and the ground truth
segmentation mask. The DSC values range from 0, representing no overlap, to 1,
indicating perfect overlap. HD measures the maximum spatial dissimilarity
between corresponding points in the predicted and ground truth masks. On the
other hand, the ASD quantifies the average spatial dissimilarity between the
predicted and ground truth masks, offering a different perspective on
segmentation accuracy. The final evaluation metrics can be defined by 5:

$\begin{split}S&=0.25\left(\frac{DSC_{FH}+DSC_{PS}+DSC_{ALL}}{3.0}\right)\\\
&+0.25\left[\frac{0.5\left(1-\frac{HD_{FH}}{100}+1-\frac{HD_{PS}}{100}+1-\frac{HD_{ALL}}{100}\right)}{3.0}+\frac{0.5\left(1-\frac{ASD_{FH}}{100}+1-\frac{ASD_{PS}}{100}+1-\frac{ASD_{ALL}}{100}\right)}{3.0}\right]\\\
&+0.5\left(1-\frac{\Delta AoP}{180}\right)\end{split}$ (5)

Besides, the visualization of the segmentation results obtained by MiTU-Net is
shown in Figure 3. The second row displays the ground truth masks, while the
third row presents the segmentation details obtained from MiTU-Net. The
results indicate a competitive performance of MiTU-Net in segmenting the pubic
symphysis and fetal head shapes.

Figure 3: Visualization of segmentation results on the JNU-IFM dataset. GT:
Ground Truth

###  3.2 Compare with SOTA Methods

Compared to other methods, our proposed method achieved the 5th place in the
FH-PS-AOP challenge, shown in Table 1. The MiTU-Net achieves a final score of
0.9283, demonstrating its effectiveness. Moreover, its average Dice score of
0.9247 for the fetal head and public symphysis demonstrates its
competitiveness, surpassing even the rank 1 method. However, it is noteworthy
that our method has big gap between top-ranking methods on HD for fetal head
and pubic symphysis. We believe the reason for this is that HD is not
considered as a loss during fine-tuning process.

Table 1: The final segmentation results of top 7 methods evaluated on the FH-PS-AOP dataset. Rank |  Author |  Final Score |  AOP |  HD FH |  HD PS |  HD All |  ASD FH |  ASD PS |  ASD All |  Dice FH |  Dice PS |  Dice All
---|---|---|---|---|---|---|---|---|---|---|---|---
1 |  Stock |  0.9418 |  6.5437 |  12.6313 |  7.6385 |  13.4477 |  3.8963 |  2.4086 |  3.4857 |  0.9303 |  0.8833 |  0.9236
2 |  Elbatel |  0.9416 |  7.9698 |  10.6985 |  7.5586 |  12.0595 |  3.3069 |  2.9945 |  2.9811 |  0.9403 |  0.8881 |  0.9346
3 |  Qiu |  0.9388 |  7.6472 |  12.4586 |  7.6612 |  13.6147 |  3.6155 |  2.2572 |  3.2383 |  0.9358 |  0.8696 |  0.9295
4 |  Chen |  0.9312 |  8.5581 |  14.0105 |  9.0512 |  15.3342 |  3.8692 |  2.6199 |  3.5173 |  0.9314 |  0.8596 |  0.9240
5 |  Ours |  0.9283 |  8.7188 |  14.0093 |  10.8286 |  15.8089 |  3.9837 |  2.9824 |  3.5785 |  0.9313 |  0.8580 |  0.9247
6 |  Sun |  0.9230 |  9.2760 |  15.7951 |  11.5361 |  17.5598 |  4.7231 |  3.1139 |  4.2646 |  0.9179 |  0.8370 |  0.9096
7 |  Cai |  0.8974 |  12.1990 |  20.0307 |  14.0684 |  21.8731 |  7.0988 |  4.2081 |  6.0576 |  0.8787 |  0.8036 |  0.8725

## Acknowledgments

This publication has emanated from research conducted with the financial
support of Science Foundation Ireland under Grant number 18/CRT/6183. For the
purpose of Open Access, the author has applied a CC BY public copyright
licence to any Author Accepted Manuscript version arising from this
submission.

## References

  * Birsan and Tiba [2006] T. Birsan and D. Tiba.  One hundred years since the introduction of the set distance by dimitrie pompeiu.  In _System Modeling and Optimization_ , pages 35–39. Springer US, 2006.
  * Chen et al. [2021] J. Chen, Y. Lu, Q. Yu, X. Luo, E. Adeli, Y. Wang, L. Lu, A. L. Yuille, and Y. Zhou.  Transunet: Transformers make strong encoders for medical image segmentation.  _arXiv_ , 2021.
  * Chen et al. [2024] Z. Chen, Z. Ou, Y. Lu, and J. Bai.  Direction-guided and multi-scale feature screening for fetal head–pubic symphysis segmentation and angle of progression calculation.  _Expert Systems with Applications_ , 245:123096, July 2024.
  * DeVries and Taylor [2017] T. DeVries and G. W. Taylor.  Improved regularization of convolutional neural networks with cutout, 2017.
  * Dice [1945] L. R. Dice.  Measures of the amount of ecologic association between species.  _Ecology_ , 26(3):297–302, July 1945.
  * Ghi et al. [2018] T. Ghi, T. Eggebø, C. Lees, K. Kalache, P. Rozenberg, A. Youssef, L. J. Salomon, and B. Tutschek.  ISUOG practice guidelines: intrapartum ultrasound.  _Ultrasound Obstet. Gynecol._ , 52(1):128–139, July 2018.
  * Iakubovskii [2019] P. Iakubovskii.  Segmentation models pytorch.  https://github.com/qubvel/segmentation_models.pytorch, 2019.
  * Jaccard [1912] P. Jaccard.  The distribution of the flora in the alpine zone.  _New Phytol._ , 11(2):37–50, Feb. 1912.
  * Long et al. [2015] J. Long, E. Shelhamer, and T. Darrell.  Fully convolutional networks for semantic segmentation.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , June 2015.
  * Lu et al. [2022a] Y. Lu, M. Zhou, D. Zhi, M. Zhou, X. Jiang, R. Qiu, Z. Ou, H. Wang, D. Qiu, M. Zhong, X. Lu, G. Chen, and J. Bai.  The JNU-IFM dataset for segmenting pubic symphysis-fetal head.  _Data Brief_ , 41(107904):107904, Apr. 2022a.
  * Lu et al. [2022b] Y. Lu, M. Zhou, D. Zhi, M. Zhou, X. Jiang, R. Qiu, Z. Ou, H. Wang, D. Qiu, M. Zhong, X. Lu, G. Chen, and J. Bai.  The JNU-IFM dataset for segmenting pubic symphysis-fetal head.  _Data in Brief_ , 41:107904, 2022b.  ISSN 2352-3409.  doi:https://doi.org/10.1016/j.dib.2022.107904.  URL https://www.sciencedirect.com/science/article/pii/S2352340922001160.
  * Milletari et al. [2016] F. Milletari, N. Navab, and S.-A. Ahmadi.  V-net: Fully convolutional neural networks for volumetric medical image segmentation.  In _2016 Fourth International Conference on 3D Vision (3DV)_ , pages 565–571, 2016.  doi:10.1109/3DV.2016.79.
  * Ronneberger et al. [2015] O. Ronneberger, P. Fischer, and T. Brox.  U-Net: Convolutional networks for biomedical image segmentation.  In _LNCS: Medical Image Computing and Computer-Assisted Intervention (MICCAI)_ , volume 9351, pages 234–241. Springer, 2015.  ISBN 978-3-319-24574-4.
  * Tavakoli et al. [2021] S. Tavakoli, A. Ghaffari, Z. M. Kouzehkanan, and R. Hosseini.  New segmentation and feature extraction algorithm for classification of white blood cells in peripheral smear images.  _Sci. Rep._ , 11(1):19428, Sept. 2021.
  * Vaswani et al. [2017a] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin.  Attention is all you need.  In I. Guyon, U. von Luxburg, S. Bengio, H. M. Wallach, R. Fergus, S. V. N. Vishwanathan, and R. Garnett, editors, _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA_ , pages 5998–6008, 2017a.  URL https://proceedings.neurips.cc/paper/2017/hash/3f5ee243547dee91fbd053c1c4a845aa-Abstract.html.
  * Vaswani et al. [2017b] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. u. Kaiser, and I. Polosukhin.  Attention is all you need.  In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017b.  URL https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.
  * Xie et al. [2021] E. Xie, W. Wang, Z. Yu, A. Anandkumar, J. M. Alvarez, and P. Luo.  Segformer: Simple and efficient design for semantic segmentation with transformers.  In _Neural Information Processing Systems (NeurIPS)_ , 2021.
  * Yeghiazaryan and Voiculescu [2018] V. Yeghiazaryan and I. Voiculescu.  Family of boundary overlap metrics for the evaluation of medical image segmentation.  _J Med Imaging (Bellingham)_ , 5(1):015006, Feb. 2018.