Masked Registration and Autoencoding of CT Images for Predictive Tibia Reconstruction
Extracted source text · arXiv 2512.09525
% \documentclass[conference]{IEEEtran}
% \IEEEoverridecommandlockouts
\documentclass[runningheads]{llncs}
\usepackage[T1]{fontenc}
% The preceding line is only needed to identify funding in the first footnote. If that is unneeded, please comment it out.
% \usepackage{mathptmx} % assumes new font selection scheme installed
\usepackage{times} % assumes new font selection scheme installed
\usepackage{amsmath} % assumes amsmath package installed
\usepackage{amssymb} % assumes amsmath package installed
\usepackage{amstext}
\usepackage{graphicx}
\usepackage{multicol}
\usepackage[bookmarks=true]{hyperref}
\usepackage[usenames]{color}
\usepackage[x11names, table]{xcolor}
\usepackage{xspace}
\usepackage{balance}
\usepackage[lofdepth,lotdepth]{subfig}
\usepackage{pifont}
% \usepackage{natbib}
\usepackage{derivative}
\usepackage{algorithm}
\usepackage{algpseudocode}
\usepackage[multi-part-units=single]{siunitx}
\usepackage{enumitem}
%%% TABLE %%
\usepackage{multirow}
\usepackage{multicol}
\usepackage{makecell}
\usepackage{tabularx}
\usepackage{array}
%%% GRAPHICS %%%
\usepackage{graphicx}
\usepackage{subcaption}
\usepackage{pgfplots}
\pgfplotsset{compat=1.18}
% \usepackage{tikz}
% \usetikzlibrary{math,3d}
% \usetikzlibrary{matrix,calc}
% \usetikzlibrary{shapes,arrows,chains}
% \usetikzlibrary{positioning}
% \usetikzlibrary{intersections}
\usepackage{pgfplots}
\usepgfplotslibrary{fillbetween}
\usepackage{keyval}
\usepackage{ifthen}
% \usepackage{tex/networks}
\usepackage[dvipsnames]{xcolor}
\usepackage[export]{adjustbox}
\usepackage{tikzscale}
\usepackage{svg}
\usepackage{svg-extract}
% \usepackage{wrapfig}
%%% Commands %%%
% \usepackage{import}
% \subimport{layers/}{init}
\newcommand{\tx}{\text}
\newcommand{\cmark}{\ding{51}}%
\newcommand{\xmark}{\ding{55}}%
\newcommand{\RRR}{\mathbb{R}}
\newcommand{\todo}[1]{\textbf{\color{red}#1}}
\newcommand{\T}{\top}
\newcommand{\marc}[1]{{\textcolor{blue}{marc: #1}}}
\newcommand{\hy}[1]{{\textcolor{cyan}{hongyou: #1}}}
\newcommand{\TT}{\mathcal{T}}
\newcommand{\NN}{\mathcal{N}}
\renewcommand{\AA}{\mathcal{A}}
\renewcommand{\t}{\theta}
\newcommand{\Exp}[2][]{\mathbb{E}_{#1}\!\left\{#2\right\}}
%%============= marc's space tricks
%\renewcommand{\baselinestretch}{0.99}
%\renewcommand{\arraystretch}{1.2}
%\renewcommand{\textfloatsep}{3ex}
%\renewcommand{\floatpagefraction}{.6}
%\renewcommand{\dblfloatpagefraction}{.6}
%\setlength{\mathindent}{2.5em}
%\setlength{\jot}{0pt} %zwischen den math zeilen
%\setlength{\abovedisplayskip}{-10pt}
%\setlength{\belowdisplayskip}{-10pt}
%\setlength{\mathsurround}{-10pt}
%\renewcommand{\floatsep}{-1ex}
%\renewcommand{\topfraction}{1}
%\renewcommand{\bottomfraction}{1}
%\renewcommand{\textfraction}{0}
%\columnsep 5ex
%\parindent 3ex
%\parskip 1ex
%%=============
\usepackage{mathtools}
\DeclarePairedDelimiter{\norm}{\Vert}{\Vert}
\DeclareMathOperator*{\argmin}{argmin}
\usepackage{comment}
\newif\ifdisclosure
\disclosuretrue
\ifdisclosure
\includecomment{disclosure}
\excludecomment{anonymous}
\else
\includecomment{anonymous}
\excludecomment{disclosure}
\fi
\def\BibTeX{{\rm B\kern-.05em{\sc i\kern-.025em b}\kern-.08em
T\kern-.1667em\lower.7ex\hbox{E}\kern-.125emX}}
\begin{document}
\title{Masked Registration and Autoencoding of CT Images for Predictive Tibia Reconstruction}
\titlerunning{Predictive Tibia Reconstruction}
%%%%%%%%%%%%%%%% For disclosure %%%%%%%%%%%%%%%
\begin{disclosure}
\author{
Hongyou Zhou\index{Zhou, Hongyou}\inst{1} \and
Cederic Aßmann\index{Aßmann, Cederic}\inst{1} \and
Alaa Bejaoui\index{Bejaoui, Alaa}\inst{2} \and
Heiko Tzschätzsch\index{Tzschätzsch, Heiko}\inst{2} \and
Mark Heyland\index{Heyland, Mark}\inst{4} \and
Julian Zierke\index{Zierke, Julian}\inst{4} \and
Niklas Tuttle\index{Tuttle, Niklas}\inst{3} \and
Sebastian Hölzl\index{Hölzl, Sebastian}\inst{3} \and
Timo Auer\index{Auer, Timo}\inst{3} \and
David A. Back\index{Back, Alexander David}\inst{3} \and
Marc Toussaint\index{Toussaint, Marc}\inst{1}
}
\authorrunning{Zhou and Toussaint et al.}
\institute{
Technical University Berlin, Learning and Intelligent Systems, Berlin, Germany
\and
Charité – Universitätsmedizin Berlin, corporate member of Freie Universität Berlin and Humboldt-Universität zu Berlin, Institute of Medical Informatics, Berlin, Germany
\and
Charité – Universitätsmedizin Berlin, corporate member of Freie Universität Berlin and Humboldt-Universität zu Berlin, Center for Musculoskeletal Surgery, Berlin, Germany
\and
Julius Wolff Institute, Berlin Institute of Health at Charité – Universitätsmedizin Berlin, Germany}
\end{disclosure}
%%%%%%%%%%%%% Anonymized Authors %%%%%%%%%%%%%%%%%%%%%
\begin{anonymous}
\author{Anonymized Authors}
\authorrunning{Anonymized Author et al.}
\institute{Anonymized Affiliations \\
\email{email@anonymized.com}}
\end{anonymous}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\maketitle
% \input{tex/tikz}
% \begin{IEEEkeywords}
% component, formatting, style, styling, insert
% \end{IEEEkeywords}
\input{sections/00_abstract}
% \input{sections/01_introduction}
\input{sections/02_related_work}
\input{sections/03_methodology}
\input{sections/04_experiments}
\input{sections/05_conclusion}
%===============================================================================
% \newpage
% \bibliographystyle{IEEEtranN}
% \footnotesize
\bibliographystyle{splncs04}
\bibliography{citations}
%===============================================================================
% \input{sections/99_appendix}
\end{document}
\begin{abstract}
%% Bone repair surgery planning requires reconstructing fragmented bones to approximate a healthy structure, often formulated as completing masked information.
%% Existing methods, such as U-Net-based architectures, Fast Fourier Convolutions, and Masked Autoencoders, show promise but fail to align data to a unified reference framework. Similarly, traditional registration methods like ICP and deep learning approaches such as VoxelMorph and Vision Transformers struggle with consistent alignment across datasets.
%% This study bridges data registration and 3D analysis through a novel pipeline that includes data preprocessing, normalization, and autoencoder training to complete masked information.
%% Key contributions are: (i) a registration model for non-deformable objects, enabling dataset alignment to a unified framework; (ii) evaluation of autoencoder architectures on normalized data; and (iii) extending registration techniques to masked data. This work enhances the precision of bone state prediction and 3D data analysis.
% Surgical planning for complex tibial fractures is challenging due to the difficulty in visualizing the desired outcome of fragmented CT scans. We propose a novel system using 3D generative models to predict the healthy tibial morphology from fracture-masked CT data, helping in preoperative planning. Our approach first uses a modified spatial transformer network (STN) to register the raw CT to a standardized coordinate system using a prototype tibia. Subsequently, various autoencoder (AE) architectures, including Masked Autoencoders (MAE), are trained to model healthy tibia variations. Both the STN and AE models are designed to be robust to masked input, enabling them to process fractured CT and generate a prediction of the corresponding healthy bone in standard coordinates. This system offers a new method for creating patient-specific 3D models of healthy tibial morphology from fractured data, which potentially improves surgical planning accuracy and efficiency. Our contributions include: i) a 3D-adapted STN for global spatial registration, ii) a comparative analysis of AEs for bone CT modeling, and iii) the extension of both to handle masked inputs for predictive generation of healthy bone structures.
Surgical planning for complex tibial fractures can be challenging for surgeons, as the 3D structure of the later desirable bone alignment may be difficult to imagine.
To assist in such planning, we address the challenge of predicting a patient-specific reconstruction target from a CT of the fractured tibia.
Our approach combines neural registration and autoencoder models.
Specifically, we first train a modified spatial transformer network (STN) to register a raw CT to a standardized coordinate system of a jointly trained tibia prototype.
Subsequently, various autoencoder (AE) architectures are trained to model healthy tibial variations.
Both the STN and AE models are further designed to be robust to masked input, allowing us to apply them to fractured CTs and decode to a prediction of the patient-specific healthy bone in standard coordinates.
Our contributions include: i) a 3D-adapted STN for global spatial registration, ii) a comparative analysis of AEs for bone CT modeling, and iii) the extension of both to handle masked inputs for predictive generation of healthy bone structures. Project page: \href{https://github.com/HongyouZhou/repair}{https://github.com/HongyouZhou/repair}
\keywords{Rigid Objects Registration \and Bone Object Reconstruction \and Bone Surgical Planning}
\end{abstract}
\section{Introduction}
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{imgs/overview/overview.png}
\caption{\textbf{Overview of the proposed methods}. (See details in Sec.~\ref{sec:methods}.) Top left: Training data $D$ of non-aligned CT images. Step 1 trains an STN and prototype jointly for consistent registration. Step 2 trains alternative AEs to provide latent representations of individual tibia. Step 3 retrains STN and AEs to become masked-robust and enable predicting a tibia reconstruction.}
\label{fig:overview}
\end{figure}
Fractures are a relevant burden for all societies around the world, with a notable increase in numbers internationally within the last decades~\cite{wu2021global}.
Surgical interventions are often indicated, especially when larger bones and joints are involved~\cite{jones2020principles}.
In the joint area in particular, exact restoration of the initial anatomical situation and maintenance of alignment and length by surgical fixation are essential to prevent early osteoarthritis~\cite{schenker2014pathogenesis}.
For this reason, planning is of decisive importance with regard to adequate fragment repositioning, possible bony supplementation, and final osteosynthesis based on the prediction of the ideal result~\cite{scolozzi2012maxillofacial,knopf2001adaptive}.
However, planning such operations based solely on computer tomography (CT) imaging poses significant challenges. The difficulty in reconstructing the desired bone alignment, especially in multi-fragmentary situations, often lead to post-traumatic arthrosis, especially if the joint surface is not reconstructed accurately~\cite{keil2023intraoperative}.
Predicting the entire healthy structure from partial input (e.g.\ fracture-masked) can be considered an instance of (3D) image completion~\cite{liu2018imageinpaintingirregularholes}.
Previous approaches~\cite{liu2018imageinpaintingirregularholes} in computer vision include using U-Net-based~\cite{ronneberger2015u} architectures to complete the partially masked images.
More recently, a framework called Resolution-Robust Large Mask Inpainting (LaMa)~\cite{suvorov2022resolution} has been introduced with fast Fourier convolutions (FFC)~\cite{chi2020fast} for image completion.
%% This approach leverages FFC by splitting channels into local and global branches, applying traditional convolutions in the local branch, and using FFT for frequency transformations in the global branch.
%% The outputs from the two branches are then fused. The primary advantage of Fast Fourier Convolutions lies in its full-image receptive field from the early layers, facilitating a better understanding of global context.
In general, Masked Autoencoders (MAE)~\cite{kingma2013auto,he2022masked} are a
strong approach to learn latent representations that are robust to
partially masked information.
The encoder operates only on the visible patches and generates latent codes, which are then
processed by the decoder along with mask tokens to reconstruct the
target image.
In our approach, we study various AE architectures for 3D image tensors, trained to directly reconstruct masked CT images.
However, naively applying autoencoders (AEs) on raw tibia CTs performs poorly, as the latter are commonly not aligned to standardized coordinates.
Structure registration -- that is, identifying the spatial transform of some structure to standard coordinates -- has extensively been studied in medical imaging~\cite{miao2016cnn,miao2018dilated} and many other areas~\cite{peng2019pvnet,deng2020self}.
To avoid issues with local optima of traditional methods like Iterative Closest Point (ICP)~\cite{besl1992method}, modern approaches train neural nets to directly predict the spatial transform from the image input: For 2D image data, Spatial Transformer Network (STN) architectures~\cite{lee2019image,jaderberg2016spatialtransformernetworks} show impressive results for deformable object registration.
VoxelMorph~\cite{balakrishnan2019voxelmorph} is a U-Net-based model to align a moving against a fixed 3D volumetric data from CT scans,
%which also employs STN to compute the partial derivatives of each voxel under affine transformations, thereby creating a registration field for the deformable objects at the voxel level.
and C2FViT~\cite{mok2022affine} applied attention mechanisms, using Vision Transformers (ViT)~\cite{dosovitskiy2020image}.
%for multi-resolution registration of deformable objects.
%% For mesh data, Iterative Closest Point (ICP)~\cite{besl1992method} remains a widely popular registration method even today.
Further deep learning-based registration methods include~\cite{choy2020deep,zhou2016fast,kim2022diffusemorph}.
However, the above studies only address registration against a known target, whereas our focus lies on inferring a consistent prototype from an entire dataset as detailed below.
In our work, we propose a modified STN that is trained to predict a registration of a raw tibia CT to shared coordinates, such that they align best with a co-optimized \emph{prototype} tibia CT.
We then consider various AE to model the distribution of tibia varieties, essentially modeling the variance around the prototype.
Both, the STN-registration as well as the AE are then trained to be mask-invariant, to be applied to a fracture-masked tibia and to decode it to the predicted reconstruction.
%We are seeking for the assistance of 3D generative models for such planning.
%More specifically, we aim for a system that takes raw CTs of fractured Tibia as input, predicts an affine registration to standard coordinates, and generates a healthy variant that provides a reference for reconstruction.
Our novel contributions include: i) the modified STN to predict the global spatial registration of a 3D image tensor,
ii) a comparison of alternative AE architectures to model bone CTs,
and iii) an extension of both of the above to masked CT images, enabling the predictive decoding to a healthy bone in standard coordinates.
% \begin{figure}[t]
% \centering
% \begin{tikzpicture}
% \node (u) [block]{U};
% \node (wrapper) [wrapper, right=3.75cm of u.south,anchor=south] {};
% \node (v) [block, right=3.25cm of wrapper.south, anchor=south]{V};
% \node (loc_wrapper) [loc_wrapper, above=-2.5cm of wrapper, xshift=-0.925cm] {};
% \node (loc_wrapper_title) [above=-0.05cm of loc_wrapper, color=red] {\footnotesize LocNet};
% \node (grid_wrapper) [grid_wrapper, right=-0.175cm of loc_wrapper] {};
% \node (grid_wrapper_title) [above=-0.05cm of grid_wrapper] {\footnotesize grid generator};
% \node (loc_conv1) [conv, right=0.25cm of loc_wrapper.west, anchor=west] {\footnotesize\rotatebox{90}{\tiny Conv}};
% \node (loc_pool1) [pool, right=0.cm of loc_conv1.east, anchor=west] {\footnotesize\rotatebox{90}{\tiny Pool}};
% \node (loc_conv2) [conv, right=0.1cm of loc_pool1.east, anchor=west] {\footnotesize\rotatebox{90}{\tiny Conv}};
% \node (loc_pool2) [pool, right=0.cm of loc_conv2.east, anchor=west] {\footnotesize\rotatebox{90}{\tiny Pool}};
% \node (loc_fc) [conv, right=0.1cm of loc_pool2.east, anchor=west] {\footnotesize\rotatebox{90}{\tiny FC}};
% \node (theta) [right=0.2cm of grid_wrapper.west, anchor=west] {$\theta$};
% \draw [arrow, color=lightblue, line width=0.3mm,] ([xshift=0.1cm,yshift=0.cm]loc_fc.east) -- ([xshift=0.1cm,yshift=0.cm]theta.west);
% \node (tau) [right=0.3cm of theta.east, anchor=west] {$\mathcal{T}_{\theta}(G)$};
% \draw [arrow, color=lightblue, line width=0.3mm,] ([xshift=0cm,yshift=0cm]theta.east) -- ([xshift=0.1cm,yshift=0.cm]tau.west);
% \node (cross) [draw,circle,cross,minimum width=0.2cm] at (u.east-|tau.south) {};
% \node (sample) [below=0cm of cross.south, anchor=north] {\footnotesize Sample};
% \draw [arrow, color=lightblue, line width=0.3mm,] ([xshift=0cm,yshift=0cm]tau.south) -- ([xshift=0cm,yshift=0.cm]cross.north);
% \draw [arrow, color=lightblue, line width=0.3mm,] ([xshift=0.15cm]u.east |- cross.west) -- (cross.west);
% \draw [arrow, color=lightblue, line width=0.3mm,] ([xshift=0.15cm]u.east |- cross.west) -- ++(0.5,0) |- (loc_conv1.west);
% \draw [arrow, color=lightblue, line width=0.3mm,] (cross.east) -- (cross.east -| v.west);
% \node (p) [param, above=0.75cm of v.north, anchor=south]{P};
% \node (l) [above=0.25cm of v.north, anchor=south, xshift=1cm] {$\ell$};
% \draw [arrow, color=lightblue, line width=0.3mm,] ([xshift=0.15cm]p.east) -| (l.north);
% \draw [arrow, color=lightblue, line width=0.3mm,] ([xshift=0.15cm]v.east) -| (l.south);
% \end{tikzpicture}
% \caption{Proposed neural network architecture for jointly prototype optimization and registration.
% The CT input labeled as U is processed through a LocNet to extract features, which are then passed through a fully connected (FC) layer to obtain the parameters of the affine transformation.
% These parameters are subsequently fed into a grid sampler, resulting in the affine-transformed tensor.}
% \label{fig:stn}
% \end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Methods}
\label{sec:methods}
In this section, we now present our method for end-to-end reconstruction.
Figure~\ref{fig:overview} outlines our method, which builds on training data $\mathcal{D}=\{(x_i,m_i)\}_{i=1}^N$ of $N$ CT scans of healthy tibiae (Fig.~\ref{fig:overview}, top left).
Each CT scan $x_i \in \mathbb{R}^{W \times H \times S}$ is a single-channel tensor.
%\footnote{In our evaluation data CT scans have pixel dimension $0.3\times0.3\times 1.0$ millimeters, and $N\approx 250$.}
We further assume to have a binary mask $m_i\in \{0,1\}^{W \times H \times D}$ indicating the tibia bone voxels while masking out neighboring bones like the femur, fibula, patella.
%\footnote{E.g.\ provided by 3D Total Segmentator.}.
At query time, we are given a fractured tibia CT (Fig.~\ref{fig:overview}, bottom left) with fractures coarsely masked out and aim to predict its reconstruction and affine transform to standard coordinates. To this end, we train models in the following steps:
\begin{enumerate}[label=\textsc{Step \arabic*:}, wide]
\item We jointly train a tibia prototype tensor $\bar x$ and a Spatial Transformer Network \cite{jaderberg2016spatialtransformernetworks} that predicts an affine transformation to match the input tibia CTs best to the prototype. Output of this step is the tibia prototype $\bar x$ and a consistently aligned data set $\hat{\mathcal{D}}=\{(\hat x_i,\hat m_i)\}_{i=1}^N$.
\item We train autoencoder models on the aligned dataset to capture the variance of healthy tibia around the prototype in their latent code (Fig.~\ref{fig:overview}, top right).
\item We then train both, the STN and autoencoders to become robust to masking, so that they can also be applied on a fracture-masked tibia CT, mapping it first to standard coordinates, then a latent code, and then decoding to its reconstruction (Fig.~\ref{fig:overview}, bottom).
\end{enumerate}
\subsection{Joint Prototype and Registration Net Optimization} \label{sec:proto}
An essential part of a \emph{Spatial Transformer Net (STN)}~\cite{jaderberg2016spatialtransformernetworks} is a transformation $\hat x = \AA(x, \theta)$ which performs a \emph{differentiable} affine transformation by grid re-sampling an image or 3D tensor $x$ for given transform parameters $\t$.\footnote{If $x$ was not discretized and $\TT_\t$ is the affine transformation, we had $\hat x(p) = x(\TT_\t^{-1}(p))$ for $p\in\RRR^3$.
An AGS $\hat x = \AA(x, \theta)$ approximates this operation for a discretized tensor $x$ using tri-linear interpolation and providing differentiation w.r.t.\ both $x$ and $\t$.}
We call $\AA$ \emph{Affine Grid Sampler (AGS)}.
This differentiability can be leveraged to directly optimize for individual affine parameters $\t_i$ for each data CT $x_i$, or a \emph{Localization Network (LocNet)}~\cite{jaderberg2016spatialtransformernetworks} can be trained to predict $\t$ as a function of $x$. We evaluate both approaches.
Using a LocNet can be less prone to local optima and will allow us to directly generalize to potentially masked query CTs. However, we consider directly using AGS for registration as an evaluation baseline, which we explain first for simplicity:
\paragraph{Direct AGS Registration:} We solve
\begin{align}\label{eq_reg_p}
\min_{\bar x, \theta_{1:N}}~ \textstyle\sum_{i=1}^{N} \ell(\bar x, \mathcal{A}(x_i \cdot m_i, \theta_i)) ~,
\end{align}
where $\bar x$ becomes the tibia prototype tensor, $\t_i$ are individual transforms for each data CT $x_i$, $x\cdot m$ is the elem-wise multiplication, and $\ell$ is a loss function comparing two CT tensors (we use D-SSIM).
\paragraph{STN Registration:} We solve
\begin{align}\label{eq_reg_p_locnet}
\min_{\bar x, \phi}~ \textstyle\sum_{i=1}^{N} \ell(\bar x, \mathcal{A}(x_i \cdot m_i, \tx{LocNet}_{\phi} (x_i \cdot m_i))) ~,
\end{align}
and use a modified $\tx{LocNet}_{\phi}$, which is based on 3D convolutions, see Fig.~\ref{fig:overview} that maps the masked data CT to affine transform parameters; $\t$ are the network parameters of $\tx{LocNet}_{\phi}$.
\subsection{Autoencoders and Healthy Tibia Latent Codes}
A \emph{variational autoencoder} (VAE) trains an encoder $q_\phi(z|x)$ and decoder $p_\theta(x'|z)$ to minimize
\begin{align}\label{eq:vac}
\min_{\phi,\theta}~ \Exp[x\sim \mathcal{D}, z \sim q_\phi(z|\hat x)]{ \ell(x', \hat x) - \beta D_\text{KL}\left(q_\phi(z| \hat x) \,\|\, p(z) \right) }~,
\end{align}
where $\ell(x',\hat x)$ is the reconstruction error between the aligned data $\hat x$ and decoded $x'$, and $p(z) = \NN(0, \mathrm{I})$ is a latent code prior. Training an autoencoder implies training a latent code that represents the variance of $\hat x$, such that we can encode $\hat x\mapsto z$ and decode $z\mapsto x'$.
Note that this VAE is trained on unmasked, aligned CTs, while a mask-robust version is later retrained with masked loss \ref{eq:vae_masked} to support fracture completion.
In step 2 of our pipeline, we train autoencoders on our aligned and masked data CTs $\{\hat x_i \circ \hat m_i\}_{i=1}^N$.
We compare three alternative approaches in this work:
\paragraph{VAE:} We train a standard variational autoencoder minimizing \eqref{eq:vac}, where the latent code maintains spatial consistency, i.e. $z \in \RRR^{E_{\tx{VAE}} \times \frac{WHS}{L^3}}$, where $E_{\tx{VAE}}$ is the predefined dimension of the latent space and $L$ is the total level of downsampling. The specific network used in this work is a U-Net-based architecture~\cite{ronneberger2015u,he2016deep}
% that contains down/up-sampling phases.
%
%The downsampling is implemented at each level via convolutional layers, while the upsampling is achieved through linear interpolation. Mirroring the U-Net architecture, we also incorporate skip connections.
% Each layer is constructed using a ResNet block~\cite{he2016deep}, implying that during the upsampling process, features from the corresponding level in the downsampling path are concatenated and subsequently processed by the ResNet block at the receiving upsampling level.
%See appendix \ref{app:VAE} for our specific implementation.
\paragraph{Vector Quantized VAE:}
A Vector Quantized VAE (VQ-VAE) provides an alternative latent structure to capture data variance, where the latent variable $z$ is categorical, i.e., discrete.
%
% The loss function of a VQ-VAE differs slightly from a standard VAE to account for the discrete latent code; we refer to \cite{van2017neural} for details.
We define a latent embedding space $z \in \RRR^{E_{\tx{VQVAE}} \times \frac{WHD}{L^3}}$ where $E_{\tx{VQVAE}}$ is the dimensionality of each latent embedding vector $z_i$, and a learnable codebook $C \in \RRR^{K \times E_{\tx{VQVAE}}}$ which has $K$ vectors.
The model takes a aligned input CT image $\hat x$, through an encoder, resulting in the output $z_e(\hat x)$. The discrete latent variables $z_i$ are then calculated by a nearest neighbor lookup using the vectors in $C$.
\paragraph{PCA:} Making the encoder and decoder deterministic and linear with limited rank $p$, and choosing the loss function as square error and $\lambda=0$, we retrieve PCA with $p$ principal components. We include this approach as a classical baseline to model the tibia CT variance around the prototype.
% See appendix \ref{app:PCA} for our specific implementation.
%See appendix \ref{app:VQVAE} for our specific implementation.
%
% \marc{needs complete shortening -- just refer to a VQ-VAE paper. The only thing that needs to get across is the structure of the latent code.} The posterior categorical distribution $q(z|x)$ probabilities are then defined as one-hot as follows:
% \begin{equation}
% \label{eq:vqvae_post}
% q(z=k|x)=
% \begin{cases}
% 1 & \tx{for} \quad k=\argmin_{j} \norm{z_e(x) - e_j}\\
% 0 & otherwise
% \end{cases}
% \end{equation}
% where $z_e(x)$ is the output of the encoder network.
% The representation $z_e(x)$ is processed through the discretization bottleneck and subsequently mapped to the nearest element of the embedding $e$, as described in Equations~\ref{eq:vqvae_post} and~\ref{eq:vqvae_emb}.
% \begin{equation}
% \label{eq:vqvae_emb}
% z_q(x) = \argmin_{e_i} \norm{z_e(x) - e_i} ~.
% \end{equation}
%
% \normalsize
%
We will evaluate these three autoencoder approaches, with the PCA being our baseline, for their performance in reconstructing healthy tibia CTs.
\subsection{Extension to Fracture-Masked CTs via Data Augmentation} \label{sec:data_aug}
At application time, we have a query CT $y$ of a fractured tibia together with a mask $m$ that (coarsely) masks out the fractured parts. Our system aims to predict the reconstruction as well as affine parameters $\t$ to transform to standard coordinates.
Applying the previously trained LocNet directly on the facture-masked CT performs mediocre, as it has never been trained on heavily masked CTs (see Table~\ref{tab:qual_comp_val_set} for masking examples). Therefore, in step 3 we train a second, \emph{mask-robust} LocNet using data augmentation. To this end, we start with our aligned training data $\hat{\mathcal{D}} = \{(\hat x_i, \hat m_i)\}_{i=1}^N$ which is the output of step 1. We now consider these $\hat x_i$ as ground truth data and generate a large data set of randomly masked and transformed instances of these $\hat x_i$ to train the second LocNet to reversely predict these transformations.
More specifically, we apply random masks $m \sim \mathrm{Bernoulli}(m)$ sampled from a distribution over blocks of voxels with random size in $[10,20]$, and sample random affine transforms $\t\sim p(\t)$ uniformly with Euler angles $\le 15^\circ$, translations $\le 2.5cm$ and scalar scaling $[0.9,1.1]$. A data CT $\hat x$ is therefore transformed to $\tilde x = \TT_{\t}^{-1}(\hat x \cdot m)$. We train the second LocNet to predict this transformation by minimizing
\begin{align}\label{eq_frac_reg}
\min_{\phi}~ \Exp[\hat x \sim \hat D,p(m), p(\t)]{ \ell(\hat x \cdot m, \mathcal{A}(\tilde x, \text{LocNet}_{\phi} (\tilde x))) + \alpha \norm{\t - \text{LocNet}_\phi(\tilde x)}^2 } ~,
\end{align}
where the second term of the loss function is a standard square error regression loss to the true affine transform. During the experiments, we trained all our models from scratch.
Finally, we also train our autoencoders alternatives from scratch to become mask-robust, using the masked reconstruction loss as follow:
\begin{align}\label{eq:vae_masked}
\min_{\phi,\theta}~ \Exp[x\sim D, z \sim q_\phi(z|\hat x)]{ \ell(x' \cdot \neg m,\hat x \cdot \neg m) - \beta D_\text{KL}\left(q_\phi(z| \hat x \cdot \neg m) \,\|\, p_\theta(z) \right) }~,
\end{align}
where the first term differs from~\ref{eq:vac} with only the non-masked area been used for calculating the loss.
% \tiny
% \marc{I don't understand the remaining text}
% Although the STN can help us identify the location of tibial fractures, it cannot generate the full tibia conditioned on the tibial fracture. In the following section, we will introduce three methods for this purpose. The input to the next phase is assumed to be affine normalized. Thus, we will disable the affine augmentations in the next section.
% Furthermore, the decision to separate the registration of healthy bones and bone fragments into two distinct steps is based on a comparison of the two registration strategies.
% As demonstrated in the subsequent experimental results, the neural network-based registration model is more robust than directly optimizing the affine transformation.
% This is because the registration loss function contains numerous local minima and introducing additional parameters allows the optimizer to explore gradients in additional dimensions, making it easier to reach the global optimum.
% However, the required number of parameters does not scale linearly with the complexity of the problem.
% In other words, combining the two registration tasks in a single step would require a significantly larger neural network to achieve global optimization.
% In the next section, we will further explore this issue.
% \normalsize
% \section{Methodology}
% In this section, we introduce our methods for registration of the CT image, prediction of the healthy status and finding the position of the bone fractures.
% \subsection{Joint Prototype optimization and data registration}\label{sec:jpodr}
% Bring the 3D data into the same pose and normalizing its size will benefit all the rest of the pipeline. In this step, we employ an unsupervised approach to normalize all the training data while obtaining the tibia prototype.
% \subsubsection{Local registration using STN}
% Let $m \in \mathbb{R}^{1 \times W \times H \times D}$ be the moving volume with single channel and fixed $W$, $H$ and $D$ and $f \in \mathbb{R}^{1 \times W \times H \times D}$ be the moving volume which can be transformed. And $t \in \mathbb{R}^{1 \times W \times H \times D}$ be the masking of the tibia segmentation. $t_m$ and $t_f$ are tibia segmentations for moving and fixed volume, respectively.
% The objective could be defined as
% \begin{equation}\label{eq_reg}
% \argmin_{\mathcal{\theta}} \ell(f * t_f,\mathcal{A}(m * t_m, \theta))
% \end{equation}
% where $\ell$ is the object defined later, and $\mathcal{A}$ is applying affine transformation along with its parameters $\theta$ on a volume.
% By using STN, we can get the partial derivative w.r.t. the affine parameters $\theta$, thus getting the local optimal.
% Additionally, since obtaining a CT scan requires the patient to maintain a consistent position, the resulting images are generally well-aligned.
% Consequently, we hypothesize that the local optimum corresponds to the global optimum, thereby eliminating the need for a global search.
% \subsubsection{Initialization of the Prototype}
% A tibia prototype $P \in \mathbb{R}^{1 \times W \times H \times D}$ can be acquired by solving
% \begin{equation}\label{eq_reg_p}
% \argmin_{P, \theta_i}{\sum_{i}^{N}{\ell(P, \mathcal{A}(x_i * t_{x_i}, \theta_i)})}
% \end{equation}
% where $x_i \in \mathbb{R}^{1 \times W \times H \times D}$ is a single tibia volume from a dataset which has $N$ tibia volumes, and $t_{x_i}$ is its tibia segmentations.
% Unlike Equation ~\eqref{eq_reg}, here we jointly optimize $P$ and the affine transformation parameters $\theta$ across all datasets, thus $P$ would contain features from tibia across the whole datasets.
% \subsection{Prediction of Healthy Variant}\label{sec:pred}
% To estimate a pre-fracture condition based on the CT scan of fractures, we decided to use an encoder-decoder model.
% We preprocessed the training data by randomly masking parts of it, thereby enabling the encoder-decoder model to develop the capability of completing missing information.
% We defined a function for getting the affine transformation parameters using ~\eqref{eq_reg} as
% \begin{equation}
% \tx{reg}(f, m) = \argmin_{\mathcal{\theta}} \ell(f,\mathcal{A}(m, \theta))
% \end{equation}
% Also we defined the fractured data as $x \in \mathcal{R}^{1 \times W \times H \times D}$ with
% \subsubsection{Linear based Prediction}
% By using singular value decomposition on our tibia datasets, we get
% \begin{equation}
% \Tilde{X} = U \Sigma W
% \end{equation}
% where $\Tilde{X} \in \mathbb{R}^{N \times 1 \times W \times H \times D}$ are the original tibia volume normalized by affine transformation with parameters gets from ~\eqref{eq_reg_p}, namely, $\Tilde{x} = \mathcal{A}(x, \tx{reg}(P, x))$, and $U \in \mathbb{R}^{N \times N}$, $W \in \mathbb{R}^{1 \times W \times H \times D}$, $\Sigma \in \mathbb{R}^{N \times N \times W \times H \times D}$, are the left/right signal vectors and singular values respectively.
% The completing of missing information is done by
% \begin{equation}
% T = (\Tilde{x} * f - \mu * f) W \qquad x_r = \Tilde{T} \Tilde{W}
% \end{equation}
% where $\Tilde{x}$ is the affine-transformed fracture data, and $f$ is the masking for the largest fracture in $\Tilde{x}$, $\mu = \frac{1}{N}\sum_{i}^{N}{x}$.
% The fracture bone $\Tilde{x}$ is encoded with $W$, with the resulted $T$, we can get the reconstructed bone by projecting $T$ using selected basis $\Tilde{W}$
% \subsubsection{Variational Autoencoder}
% For the nonlinear approach, we use a VAE to encode and decode the tibia. The loss function employed is MSE.
% To enhance the robustness of the Autoencoder against incomplete information, we use random masking to augment the training data.
% This approach is crucial because the Autoencoder will be validated on fractured bones.
% Specifically, we take a portion of the fractured bone as the input to the Autoencoder and use the output as a prediction of the healthy state of the fractured bone.
% \subsection{Reconstruction}
% The reconstruction of bone fragments, under the assumption of no path planning, can be viewed as a straightforward optimization problem.
% However, as described earlier, this study addresses a low-overlap registration challenge.
% In scenarios with low overlap, the number of outliers tends to be significantly higher compared to standard registration tasks.
% \subsubsection{Adaptive Structural Similarity Index Measure}
% When calculating the registration score, the similarity of the background can cause the valuation to become excessively high, thereby introducing challenges in gradient computation during the optimization process.
% To better adapt to specialized registration algorithms, we need to adjust the loss function accordingly.
% Building upon the original SSIM framework, we introduce an additional layer of masking.
% This modification allows us to control the loss function returned by SSIM at the voxel level. The adaptive SSIM
% \begin{equation}
% \tx{ADAPT\_SSIM}(f, m) = f \cdot \tx{SSIM}(f, m)
% \end{equation}
% where $m$, $f$ are the moving volume and fixed volume respectively, SSIM is the original SSIM in which the luminance difference being removed. Multiplying by $f$ will emphasize the area where voxel has high density value, implicitly means putting more focus on the bone structure.
% \subsubsection{Low-overlap registration}
% We defined the problem as a registration problem for bone fracture which has a CT image $V \in \mathcal{R}^{1 \times W \times H \times D}$ and the fracture segmentations $S_{i \in I} \in {0, 1} ^ {1 \times W \times H \times D}$.
% The steps for this approach are as follows:
% \begin{itemize}
% \item Finding the affine transformation $\theta_{l}$ for the largest fragment $S_{\Tilde{i}}$ w.r.t. the prototype $P$ which generated in \ref{sec:jpodr}
% \item Apply affine transformation $\theta_{l}$ for the rest of fragments $S_{I \setminus \Tilde{i}}$
% \item Use the En/decoder from \ref{sec:pred} to get the healthy variant on the affine transformed largest fragment, namely, $V_p = reg(P, V * S_{\Tilde{i}})$.
% \item for each other fragments, using alg. \ref{algo:ransac_reg} to reason the reconstruction position and orientation.
% \end{itemize}
% \begin{algorithm}\label{algo:ransac_reg}
% \caption{RANSAC based Registration on supervoxel}
% \begin{algorithmic}[1]
% \State \textbf{Input:} Moving Volume $V_m$
% \State \textbf{Input:} Fixed Volume $V_f$
% \State \textbf{Input:} Inlier Thereshold $t$
% \State \textbf{Input:} Number of Inlier Thereshold $d$
% \State \textbf{Output:} Affine transformation parameters $\Theta$
% \State Initialize mask $M$ with the same size of $V_f$
% \State $SV_m \gets \tx{oversegmentation}(V_m)$
% \State $\tx{bestFit} \gets \tx{None}$
% \For{Iteration}
% \State $SV_{rand} \gets \tx{randomPick}(SV_m)$ \Comment{Exploration}
% \State $SvMFit \gets \tx{optStn}(SV_{rand}, V_f)$
% \State $\tx{inlier} \gets SvMFit > t$
% \If{$num(\tx{inlier}) > d$}
% \State $\tx{betterSV} \gets SV_{rand} + inlier$ \Comment{Exploitation}
% \State $\tx{SvMFit} \gets \tx{optStn}(betterSV, V_f)$
% \If{loss is the current best}
% \State $\tx{bestFit} \gets \tx{SvMFit}$
% \EndIf
% \EndIf
% \EndFor
% \State \Return $bestFit$
% \end{algorithmic}
% \end{algorithm}
% In the problem we are addressing, the complex structure of the bones means that fragments may contain various biological materials, including cancellous bone.
% During the formation of bone fragments, cancellous bone can detach from the main skeletal structure and form separate micro-fragments.
% These micro-fragments are often overlooked or discarded during surgery. We aim for the algorithm to exhibit a certain degree of robustness towards such fragments.
% The RANSAC algorithm employs a random approach to identify potential outliers, which can be understood as an exploratory phase, before switching to a refinement phase to identify the optimal alignment.
% Optimization algorithms are often hindered by local optima, and this case is no exception.
% Given that the parameters of the optimization algorithm are limited to the nine parameters associated with affine transformations, the likelihood of escaping local optima is relatively low.
% However, this does not prevent us from using the algorithm to determine the optimal positioning of fragments. Since bones are typically surrounded by muscle and skin tissues, their positional changes are significantly constrained even if the skeletal structure is damaged.
% Therefore, a local search for the alignment of bone fragments is sufficient to find the global optimal solution.
% \subsubsection{Registration with Supervoxel}\label{algo:watershed}
% Building upon the superpixel methodology, we extend the application to three dimensions, resulting in supervoxels.
% Before applying RANSAC registration, we first compute supervoxels based on watershed features. ~\eqref{algo:watershed}
% \begin{algorithm}
% \caption{Watershed Oversegmentation using Voxels}
% \begin{algorithmic}[1]
% \State \textbf{Input:} 3D Volume $V$
% \State \textbf{Output:} Segmented 3D Volume $S$
% \State Initialize marker volume $M$ with local minima of $V$
% \State Initialize label volume $L$ with unique labels for each marker in $M$
% \State Initialize priority queue $Q$ with all voxels from markers, prioritized by intensity in $V$
% \While{$Q$ is not empty}
% \State Pop voxel $v$ with the lowest intensity from $Q$
% \State Let $N(v)$ be the set of neighbors of $v$
% \For{each neighbor $n \in N(v)$}
% \If{$L(n) = 0$} \Comment{If $n$ is unlabeled}
% \State $L(n) \gets L(v)$
% \State Insert $n$ into $Q$ with priority $V(n)$
% \EndIf
% \EndFor
% \EndWhile
% \State \Return $L$ as the segmented 3D volume $S$
% \end{algorithmic}
% \end{algorithm}
% \begin{algorithm}
% \caption{Fracture Registration}\label{algo:ransac_reg}
% \begin{algorithmic}[1]
% \State \textbf{Input:} Moving Volume (fracture) $V_m$
% \State \textbf{Input:} Segmentation of Moving Volume $S_m$
% \State \textbf{Input:} Fixed Volume $V_f$
% \State \textbf{Output:} Affine transformation parameters $\Theta$ for each fracture
% \State $res \gets emptySet()$
% \For{$iteration < numFractures$}
% \State $SV\_m\_fit, \theta \gets ransac_reg(V_m, V_f)$
% \State $fractureSeg \gets getBestFitMask(SvMFit, S_m)$
% \State $V_m \gets V_m * (1 - fractureSeg)$
% \State add $(fractureSeg, \theta)$ to res
% \EndFor
% \State \Return res
% \end{algorithmic}
% \end{algorithm}
% During the registration process, we integrate supervoxels with SSIM, normalizing the SSIM values within each supervoxel.
% Finally, the normalized SSIM is utilized for registration.
\section{Experiments}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\newcommand{\maxf}[1]{{\cellcolor[gray]{0.8}} #1}
\begin{table}[t]
\renewrobustcmd{\bfseries}{\fontseries{b}\selectfont}
\centering
\caption{Reconstruction performances (mean with $95\%$ confidence interval) for the investigated registration and autoencoding methods on a left-out validation dataset.
%(for training set, see Appendix \ref{app:results}).
Autoencoding method $P$ means predicting the constant prototype (directly relates to losses \eqref{eq_reg_p} and \eqref{eq_reg_p_locnet}), respectively; while the performances for other autoencoding methods relate to \eqref{eq:vac}. The bottom block evaluates mask-robust methods on randomly masked samples of the validation dataset.}
% The best-performing method is denoted in bold in each group, while the optimal method under both masked and unmasked conditions is highlighted with a gray background.}
\label{tab:quat_comp_val_set}
\sisetup{
separate-uncertainty=true,
round-mode=uncertainty,
group-separator=\text{},
round-precision=1,
multi-part-units=single,
bracket-numbers = false,
print-zero-integer = true,
detect-all=true,
% detect-weight=true,
detect-inline-weight=math
% table-alignment-mode=format,
% table-number-alignment = center,
% scientific-notation=true,
% table-number-alignment=left
}
\include{tables/val_val_set}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{table}[t]
\caption{Qualitative comparison of reconstructed tibia for investigated registration and autoencoding methods. 3D models were the isosurfaces generated using a Marching Cubes methods.}
%(see details in appendix \ref{app:march}).}
\label{tab:qual_comp_val_set}
\centering
\begin{tabular}{|c|cc|ccccc|}
\hline
\multicolumn{3}{|c|}{Registration Method} &
\multicolumn{5}{c|}{Autoencoder Method}
\\
\hline
&
Input&
Prototype &
$\tx{PCA}_{32}$ &
$\tx{PCA}_{16}$ &
VAE &
VQ-VAE &
Ground Truth \\
\hline
&&&&&&&\\[-2.5ex]
\makecell{w/o\\Registration} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_no_reg_gt.png} &
N/A &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_no_reg_pca_32.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_no_reg_pca_16.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_no_reg_vae.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/vqvae/woReg_out.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_no_reg_gt.png} % val_set batch_id: 4 id: 2 global_id: 27
\\
\makecell{Direct AGS~\ref{eq_reg_p}\\Registration} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_no_reg_gt.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_no_locnet_p.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_no_locnet_pca_32.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_no_locnet_pca_16.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_no_locnet_vae.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/vqvae/BaseReg_out.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_no_locnet_gt.png} % val_set batch_id: 4 id: 2 global_id:27
\\
\makecell{STN~\ref{eq_reg_p_locnet}\\Registration} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_no_reg_gt.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_locnet_p.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_locnet_pca_32.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_locnet_pca_16.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_locnet_vae.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/vqvae/RegNet_out.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_locnet_gt.png}
\\
\hline
\hline
\makecell{Mask-robust\\STN~\ref{eq_frac_reg}\\Registration} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_locnet_m_input.png} &
N/A &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_locnet_m_pca_32.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_locnet_m_pca_16.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_locnet_m_vae.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/vqvae/64_cnn_out.png} &
\includegraphics[width=0.10\linewidth,valign=c]{imgs/tibia_locnet_m_gt.png}
\\[1.5ex]
\hline
\end{tabular}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Data \& Training Details}
\label{sec:ex_pred_healthy}
\begin{anonymous}
Our dataset consists of 249 lower-body CT scans\footnote{The study was approved by the institutional ethics board as part of the research project.}.
\end{anonymous}
\begin{disclosure}
Our dataset consists of 249 full-body CT scans provided by the Department of Musculoskeletal Surgery at Charité\footnote{The study was approved by the institutional ethics board as part of the REPAIR research project. \url{https://repair.charite.de/}}.
\end{disclosure}
Using TotalSegmentator~\cite{wasserthal2023totalsegmentator} we get the data masks $\{m_i\}$ which crop out the tibia and its bounding box.
To make our training more robust, we re-sample these tibia CTs to a homogeneous voxel size of 1mm and finally use padding or cropping to make them isotropic, resulting in a tensor of the size of $x_i \in \mathbb{R}^{128 \times 128 \times 256}$ together with binary tibia voxel masks $m_i$ of the same dimension.
To increase the amount of available training data, we performed a
mirroring of the right-side skeleton along the median plane ($xy$ plane). Including both legs, we
have in total 498 samples. We used 442 samples as the training set,
and the remaining 56 samples as the validation set.
For the \emph{Direct Registration}~\eqref{eq_reg_p} we use the AdamW~\cite{loshchilov2017decoupled} optimizer with a learning rate of $1e^{-5}$, a batch size of 16, and let the optimizer run 50 steps for each batch and 100 epochs through the whole dataset. For the \emph{Registration Network}~\eqref{eq_reg_p_locnet} we use the AdamW optimizer with a learning rate of $1e^{-5}$, batch size of 16, and train the network for 1000 epochs.
For training the VAE we use the AdamW optimizer with a learning rate of $1e^{-5}$, the KL weight $\beta = 1e^{-2}$, and we train the network for 500 epochs for each data set. For VQ-VAE training we use the AdamW optimizer with a learning rate of $1e^{-5}$ and a quantizer weight $\beta$ following a step scheduler from $\beta_{\text{start}} = 0.5$ to $\beta_{\text{max}} = 1.0$ over the first 50 epochs, we train the network for 500 epochs for each data set. We use \emph{MSE} loss for all training task of AEs.
We train the mask-robust models using the same parameters but freeze the prototype parameters.
For the training data, we use data that have been aligned using STN registration.
\subsection{Reconstruction Performance}
Table~\ref{tab:quat_comp_val_set} reports on quantitative reconstruction performances for the investigated registration and autoencoding methods on the validation data set. The \emph{w/o Registration} method naively applies autoencoders on the non-aligned raw data.
The \emph{Direct} and \emph{STN Registration} methods follow Sec.~\ref{sec:proto}.
The \emph{Mask-robust STN Registration} tests the mask-robust LocNet we trained for randomly masked and transformed data (Sec.~\ref{sec:data_aug}), but applied to the non-masked evaluation set as the other methods. The autoencoding method \emph{P} denotes the baseline of constantly predicting the prototype only (which is not available w/o registration).
As \textbf{metrics} we use the D-SSIM~\cite{wang2003multiscale} with varying window sizes $3, 7, 11, 15$,%
\footnote{The window size of D-SSIM affects the scale of comparison, which smaller windows emphasizing finer-grained similarity, whereas larger windows focussing more on overall similarity. Evaluating multiple window sizes allows us to assess the reconstruction performance in terms of both detail and structure.}
PSNR \cite{wang2009mean}, root mean squared error (RMSE) and root mean surface distance (RMSD). The latter denotes the distance between surfaces constructed using the marching cube algorithm~\cite{lorensen1998marching} with value window from $0.0$ to $1.0$ in millimeters.
Table~\ref{tab:quat_comp_val_set} shows results on the validation.
% Appendix \ref{app:results} presents the analogous table for the training data set.
Table~\ref{tab:qual_comp_val_set} more qualitatively illustrates
reconstructions for a sample CT for the different methods. The 3D
surface models were created using a multi-threshold marching cube
method.
% (see details in Appendix \ref{app:march})
\paragraph{Discussion:}
For the \emph{Direct AGS Registration}, the optimization process (\ref{eq_reg_p})
often converges to local optima, resulting in a prototype that fails
to converge to a unified skeletal model.
In contrast, the trained \emph{STN} achieves a better solution, as indicated by the more uniform 3D reconstruction of the prototype.
This superiority is further supported by the quantitative analysis on
both the training and validation datasets, where the registration
results obtained using the STN (\ref{eq_reg_p_locnet}) consistently outperform the
direct parameter optimization method (\ref{eq_reg_p}).
%% \subsubsection{Prediction of the Healthy} \label{sec:ex_pred_healthy}
%% In this section, we use VAE and VQ-VAE to generate the healthy variant of the augmented data that we mentioned in~\ref{sec:data_aug} without applying random affine transformation, i.e. $r={[1,0,0,0]}^T$, $t={[0,0,0]}^T$ and $s={[1,1,1]}^T$.
%% Both VAE and VQ-VAE share the same UNet configurations.
%% The only difference is the bottleneck of the network, where VAE is using a reparameterization trick, and VQ-VAE is using the vector quantizer.
%% The details of the U-Net configuration used for the VAE and the VQ-VAE can be found in the Appendix.
When comparing autoencoder methods, the PCA results clearly indicate that registration significantly enhances reconstruction quality.
In constrast, the VAE and VQ-VAE methods are more robust to inaccurate registration.
The quantitative results show that the combination of using the trained STN registration with a (VQ-)VAE encoder consistently outperforms the alternatives.
The lower blocks of tables \ref{tab:quat_comp_val_set} and \ref{tab:qual_comp_val_set} show results when input CTs were randomly masked -- only the trained mask-robust STN registration is applicable.
Our results indicate that reconstruction results are still of high quality despite input masking.
Overall we find:
\begin{itemize}
\item Without masking: VQ-VAE achieves better quantitative performance on unmasked data, as the use of discrete latent embeddings allows it to memorize high-frequency structural patterns. These hard assignments enable the decoder to reconstruct sharper edges and localized anatomical details more accurately compared to the continuous latent codes of a VAE.
\item With masking: In contrast, the continuous latent space of a VAE facilitates better generalization to unseen or partially missing inputs. During inpainting, the smoothness of the latent manifold allows the decoder to interpolate plausible structures, leading to more anatomically coherent results and fewer artifacts, especially under large fracture masks.
\item Artifacts in VQ-VAE with masking arise due to limited codebook flexibility, making it struggle with unseen or missing patterns. We hypothesize that increasing the codebook size and applying EMA updates can reduce hallucinations by improving coverage of anatomical variations. Additionally, encouraging balanced codeword usage through entropy regularization may lead to more reliable reconstructions under masked inputs.
\end{itemize}
Finally, Fig.~\ref{fig:qualitative_comp} presents results on applying our reconstruction pipeline on three exemplary fracture-masked CTs.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{figure}[t]
\centering
\begin{tabularx}{\linewidth}{c@{}c @{\hspace{2\tabcolsep}} c@{}c @{\hspace{2\tabcolsep}} c@{}c}
\includegraphics[width=0.15\linewidth,valign=c]{imgs/frac_input_6.png} &
\includegraphics[width=0.15\linewidth,valign=c]{imgs/frac_recon_6.png} &
\includegraphics[width=0.15\linewidth,valign=c]{imgs/frac_input_2.png} &
\includegraphics[width=0.15\linewidth,valign=c]{imgs/frac_recon_2.png} &
\includegraphics[width=0.15\linewidth,valign=c]{imgs/frac_input_7.png} &
\includegraphics[width=0.15\linewidth,valign=c]{imgs/frac_recon_7.png} \\
Input & Reconstruction & Input & Reconstruction & Input & Reconstruction \\
\end{tabularx}
\caption{Application of our mask-robust STN+VAE pipeline on three exemplary CTs where fractures have been masked out.}
\label{fig:qualitative_comp}
\end{figure}
% \begin{figure}[t]
% \centering
% \begin{tabularx}{\linewidth}{ccc}
% (a) Reference $\hat x_i$ & (b) Augmented Data $\tilde x$ & (c) Predicted Transform \\
% \includegraphics[width=0.3\linewidth,valign=c]{imgs/tibia_locnet_p.png} &
% \includegraphics[width=0.3\linewidth,valign=c]{imgs/masked_reg_input.png} &
% \includegraphics[width=0.3\linewidth,valign=c]{imgs/masked_reg_pred.png} \\
% \end{tabularx}
% \caption{Qualitative result of the Registration for Masked Data. (a) Registered reference bone (Prototype); (b) Test bone after data augmentation; (c) Re-localized test bone.}
% \label{fig:mask_qual_comp}
% \end{figure}
% The previous step assumes that the input is correctly repositioned.
% This is done by training a slightly modified registration model described in Sec.~\ref{sec:data_aug}.
% We augment the data with the configurations described in Sec.~\ref{sec:data_aug}.
% Unlike the model for generating the prototype, the model here only has a LocNet to input the affine parameters.
% We then freeze the prototype parameters in this section, and the rest of the configurations are as we used in sec~\ref{sec:ex_pred_healthy}.
% We use the AdamW optimizer with a learning rate of $1e^{-5}$ for 500 epochs.
% The last section in Table~\ref{tab:quat_comp_val_set} presents the results of training the encoder-decoder after performing data augmentation on data that were previously registered using a registration network.
% Its performance on the validation dataset is slightly worse compared to the unmasked case, which is due to information loss and registration-induced errors.
% Finally, we can apply the entire pipeline to real-world data, with the results shown in the lower part of Fig.~\ref{fig:overview}.
% \begin{figure}[htbp]
% \centering
% \includegraphics[width=\linewidth]{imgs/plots/reg_masked_train_loss.png}
% \caption{The training loss of the registration model for fractured Tibiae. The figure shows the occurrence of significant pseudo-convergence phenomena during training, which increases the difficulty of model optimization.}
% \label{plt:reg_train_loss}
% \end{figure}
% \subsubsection{Test on real tibia fracture}
% In this section, a fracture from the real fracture dataset (ID: 011, SegName: Tibia\_L) is used as the input volume for the VAE and VQ-VAE. The data augmentations used for the VQ-VAE are listed in Tab.~\ref{tab:data_aug_vqvae_random}. A distinction is made between the random masking strategy, upper and lower region masking, and a combination of both. The upper and lower region masking experiments include a comparison of results with max. 30\% and 50\% of the upper and lower regions of the bone fully masked, shown in Fig.~\ref{fig:realfractureGeneration}.
% \begin{figure}[t]
% \centering
% \begin{tabularx}{\linewidth}{cc}
% % \includegraphics[width=0.32\linewidth,valign=c]{imgs/real_fracture/real_uplow03_in.png} &
% % \includegraphics[width=0.32\linewidth,valign=c]{imgs/real_fracture/real_uplow03_out.png} &
% % \includegraphics[width=0.32\linewidth,valign=c]{imgs/real_fracture/real_randuplow03_out.png} \\
% % (a) Input Fracture & (b) up/low max. $0.3$ & (c) random + up/low max. $0.3$ \\
% % \includegraphics[width=0.32\linewidth,valign=c]{imgs/real_fracture/real_wMask_out.png} &
% % \includegraphics[width=0.32\linewidth,valign=c]{imgs/real_fracture/real_uplow05_out.png} &
% % \includegraphics[width=0.32\linewidth,valign=c]{imgs/real_fracture/real_randuplow05_out.png} \\
% % (d) random Masking & (e) up/low max. $0.5$ & (f) random + up/low max. $0.5$
% \includegraphics[width=0.32\linewidth,valign=c]{imgs/real_fracture/real_uplow03_in.png} &
% \includegraphics[width=0.32\linewidth,valign=c]{imgs/real_fracture/real_randuplow05_out.png} \\
% (a) Input Fracture & (b) Reconstruction
% \end{tabularx}
% \caption{Reconstruction on real bone fragments. Since there is no corresponding ground truth, we only present qualitative results.}
% \label{fig:realfractureGeneration}
% \end{figure}
% \subsubsection{Fracture Registration}
% After obtaining the predictive model for healthy bones, we proceed to the third step of the experiment: optimizing the positions of the bone fragments to determine their original locations.
% Here, the data we use differs from the previous steps.
% Due to the requirements of clinical doctors, the dimensions of the CT scans of the damaged bones are not uniform.
% Most of the time, we can only obtain the upper part of the tibia, while the lower part is often omitted from the CT scan for efficiency reasons.
% Thus, the first issue we need to address is how to perform the initial registration of the fragmented bones.
% The bone prototype generated in the first step can serve as a reference for the affine transformation of the bone model.
% Therefore, before using the model trained in the second step to predict the healthy state, we need to register the fragmented bones to the bone prototype.
% Here, we use a registration method combined with RANSAC, primarily to counteract the information noise caused by bone fragmentation.
% We use RANSAC combined with STN for the registration of the moving volume to the fixed volume.
% Additionally, we pad the moving volume with zeros to match the size of the fixed volume.
% We then use our registration method, combined with adaptive SSIM, to align the fragmented bones.
% Due to the structural characteristics of bones, fragmented bones often consist of a larger main body along with many smaller fragments.
% In the preliminary registration, we aim to align the main part of the fragments to the desired position.
% As shown in the figure, after registration, the main part of the damaged bone has been regularized through affine transformation.
% Additionally, the bone fragments will also be subjected to the same affine transformation parameters as the main part.
% After normalizing the fragmented bones, we retain the main part of the bone and mask the other parts.
% This serves as the input for the healthy state prediction model. We evaluate the accuracy of the predictions using the healthy bone model from the patient's other side.
% Due to the limited amount of data, we have only 20 sets of data to validate the accuracy of the algorithm in our experiments.
% For this reason, we predominantly use optimization algorithms as the main method.
% % TODO
% In using the mirrored leg as validation data, the obtained SSIM, PSNR, and MSE are 0.001, 0.002, and 0.003, respectively.
% The actual results predicted by the model are shown in the figure.
% \subsubsection{Planning for Reconstruction}
% We also use RANSAC combined with the optimized mode of STN to find the reasonable positions for the bone fragments.
% Using the previously mentioned algorithms, we can register the bone fragments one by one.
% Additionally, we need to address the issue of volume inconsistency caused by segmentation.
% We can utilize adaptive SSIM to reduce the registration weight for already occupied positions, thereby allowing the algorithm to consider the overlap issue without hindrance.
\resizebox{\linewidth}{!}{%
\begin{tabular}{|c|c|
S[scientific-notation=fixed,fixed-exponent=-2, table-format=1.3(1.3)]
S[scientific-notation=fixed,fixed-exponent=-2, table-format=1.3(1.3)]
S[scientific-notation=fixed,fixed-exponent=-2, table-format=1.3(1.3)]
S[scientific-notation=fixed,fixed-exponent=-2, table-format=1.3(1.3)]
|
S[table-format=2.2(1.2)]
S[scientific-notation=fixed,fixed-exponent=-2, table-format=1.3(1.3)]
S[table-format=1.3(1.3)]
|}
\hline
Registration & Autoencoding & \multicolumn{4}{c|}{D-SSIM $\times 10^{-2}$} & && \\
Method &
Method &
{$\tx{D-SSIM}_{15} \downarrow$} &
{$\tx{D-SSIM}_{11}\downarrow$} &
{$\tx{D-SSIM}_{7}\downarrow$} &
{$\tx{D-SSIM}_3\downarrow$} &
{$\tx{PSNR}\uparrow$} &
{$\tx{RMSE}\downarrow \times 10^{-2}$} &
{$\tx{RMSD}\downarrow$}\\
\hline
{\multirow{5}{*}{\makecell{w/o \\ Registration}}}
& P & N/A & N/A & N/A & N/A & N/A & N/A & N/A\\
& $\tx{PCA}_{32}$ &
0.01664875613318549(0.0007912259296683309) &
0.016647324517921166(0.0007898485529793968) &
0.016689325924272896(0.0007566921572770006) &
0.017105474516197476(0.0005964196563127651) &
31.991376311690715(0.35830097039510245) &
0.025455816855861076(0.003398967530621464) &
1.5949316473803754(0.13023528405505114)\\
& $\tx{PCA}_{16}$ &
0.016985931882151844(0.0008365385186977965) &
0.01698143228336617(0.0008353968375617851) &
0.016947547594706215(0.0008061797522638038) &
0.016957825532665963(0.0006645116966271725) &
31.463969972398544(0.3563413692818919) &
0.02704199617383656(0.0037492360968660807) &
1.5884868517511581(0.13537030195568564)\\
& VAE &
0.007190997557093699(0.00034211831242498555) &
0.007211656723585393(0.00033997894152110886) &
0.007289892735166683(0.00034175782447836816) &
0.00772983513565527(0.00036338943975606434) &
33.065775553385414(0.36990255277776757) &
0.022648220881819725(0.004717212861609128) &
0.7399625664813341(0.03247699874245162)\\
& VQ-VAE &
\bfseries 0.0014544372679665685(0.00009338633558732737) &
\bfseries 0.001456961061598526(0.00009348113906498797) &
\bfseries 0.0015430379038055737(0.00009545097049256508) &
\bfseries 0.0026010833163228296(0.00011470780671469184) &
\bfseries 43.76724667019314(0.6494134946130806) &
\bfseries 0.006848246169586976(0.001110340568190161) &
\bfseries 0.7441745011632422(0.021380441143050044)\\
\hline
{\multirow{5}{*}{\makecell{Direct AGS~\ref{eq_reg_p}\\Registration}}}
& P &
0.01135346572846174(0.0004735046701158597) &
0.01131404777643857(0.0004713639387237148) &
0.01052123977354279(0.00043263599315175393) &
0.008160593884962574(0.0003079250545932524) &
28.29331814801251(0.2743336983180342) &
0.03900840340389146(0.00719745167912021) &
2.0235141930611795(0.8703873268857633)\\
& $\tx{PCA}_{32}$ &
0.01633695706173225(0.0007880602611303825) &
0.016336032085948513(0.0007866982120127982) &
0.01638820546644705(0.0007533548980669755) &
0.01683376674298886(0.0005918307370763068) &
32.34716401276764(0.3698492412887344) &
0.02445481193286401(0.0031539549231793837) &
1.56617499336258(0.12809033482952706)\\
& $\tx{PCA}_{16}$ &
0.016760126860053452(0.0008277762120175264) &
0.01675623600129728(0.0008266375204466819) &
0.016732495692041185(0.0007975822128124037) &
0.01676245640825342(0.0006574767190777763) &
31.784438945628978(0.3666125278204015) &
0.026081271952501043(0.003493921688341166) &
1.5687894794425024(0.1336431934810827)\\
& VAE &
0.011574677502115566(0.00064369278643893) &
0.011579266956282986(0.0006438202452334441) &
0.011658345038692156(0.0006403313171045792) &
0.012489944903386964(0.0006199870477620393) &
40.15658442179362(0.327861765645936) &
0.010012523064182864(0.0006411456004977963) &
0.568795047150313(0.008133955507814296)\\
& VQ-VAE &
\bfseries 0.001140505670466357(0.00012333135322701802) &
\bfseries 0.0011414256813522014(0.00012349329011092548) &
\bfseries 0.0011791245002920432(0.00012660437818623563) &
\bfseries 0.0016894926084205506(0.00015145878857139635) &
\bfseries 45.76874160766601(0.25035728384401146) &
\bfseries 0.005202674048228397(0.00016221666344963274) &
\bfseries 0.38406632309453703(0.0074613715065239864)\\
\hline
{\multirow{5}{*}{\makecell{STN~\ref{eq_reg_p_locnet}\\Registration}}}
& P &
0.0036250502043576145(0.000496090911094132) &
0.0036179037350747316(0.0004939899420505251) &
0.0034819560004743155(0.0004491611256700097) &
0.003033307975985937(0.00029935578978056637) &
35.66410629837601(0.5067043353210079) &
0.0175191990479275(0.0022159121679690944) &
0.7719681750030623(0.30453800823402055) \\
& $\tx{PCA}_{32}$ &
0.0026021682553821128(0.00047365022053762093) &
0.0025992735668464947(0.00047247123175165893) &
0.002580253614319696(0.00044125209863467993) &
0.0026331775718265098(0.0003019113333110822) &
39.49588235219319(0.5596755041440498) &
0.010959273963062847(0.0014636276014536346) &
0.6938937416008498(0.11294655714607955)\\
& $\tx{PCA}_{16}$ &
0.0027392484523631896(0.00047062812825493435) &
0.002735732330216301(0.0004694258352554797) &
0.0027026815546883475(0.0004379703645521335) &
0.0026873030044414366(0.0003005437716524351) &
38.87070980778448(0.5340410003347925) &
0.01173761423194298(0.0015473381390749565) &
0.6938267032995306(0.11121547523164899)\\
& VAE &
0.001181295004155902(0.00017046009269404807) &
0.00118420807768901(0.0001702841030314989) &
0.00120908794634872(0.00016498369753885784) &
0.001579695243802335(0.00013174335897814777) &
42.26615227593315(0.3640199523731436) &
0.007890183271633254(0.007890183271633254) &
0.5252590217266043(0.027200714897046886)\\
& VQ-VAE &
\bfseries 0.0006209678378783994(0.00011220383381063144) &
\bfseries 0.0006221207408493178(0.00011227553813101833) &
\bfseries 0.0006628980433257917(0.00011436952786564886) &
\bfseries 0.0010917667273638977(0.00012666321471132275) &
\bfseries 51.16865242852105(0.22523856559723482) &
\bfseries 0.0027897320946471556(0.00005859414559308857) &
\bfseries 0.39472172035524766(0.00516740208326045)\\
\hline
\hline
{\multirow{5}{*}{\makecell{Mask-robust~\ref{eq_frac_reg}\\STN\\Registration}}}
& P & N/A & N/A & N/A & N/A & N/A & N/A & N/A\\
& $\tx{PCA}_{32}$ &
0.004249622424443559(0.00019954539715890987) &
0.0042396220895979135(0.00019889973516680614) &
0.004056279857953387(0.00018465944814718825) &
0.00353793634308709(0.0001288944517001471) &
36.25393705014828(0.4765798399722228) &
0.015713219599867306(0.0017993836352386435) &
0.6710157162299232(0.03207359413791263)\\
& $\tx{PCA}_{16}$ &
0.004404680596457585(0.0002205984671801287) &
0.004394674190768489(0.00021993522566233464) &
0.00420772808569449(0.00020508026525064215) &
0.0036624493422331614(0.0001480871166792448) &
36.242796933209426(0.466674787487145) &
0.01572106887275973(0.0017465669539955295) &
0.6724579791638271(0.03126404095235117)\\
& VAE &
\bfseries 0.0014298437163233757(0.00008090300722387914) &
\bfseries 0.0014294855645857751(0.00008077349112719437) &
\bfseries 0.0014381750952452421(0.00007713142205422591) &
\bfseries 0.0015914186369627714(0.000059420488192200475) &
43.25596046447754(0.34627400462768554) &
0.007025767117738724(0.00043580540846641395) &
\bfseries 0.5227560698986053(0.010329485039599929) \\
& VQ-VAE &
0.003018034958384103(0.0003290086925340603) &
0.003019261277384228(0.00032908133071834475) &
0.0030634503345936537(0.0003288607830890938) &
0.0035465702983654207(0.0003329653247226302) &
\bfseries 44.03749889797635(0.7235716493297008) &
\bfseries 0.006809761385536856(0.0004895386424482849) &
0.6035331036411872(0.043848543521306634)\\
\hline
\end{tabular}
}
\section{Conclusion}
Our work demonstrates that robust predictive reconstruction of a healthy tibia from fractured CT images is achievable by combining deep models for spatial registration and autoencoders.
Specifically, by introducing a modified STN for consistent registration, evaluating alternative autoencoder models, and introducing a novel pipeline to train mask-robust models we achieve high-quality reconstructions also for fracture-masked CTs.
Our experimental results indicate that while VQ-VAE delivers sharper details in unmasked scenarios, the continuous latent space of the VAE offers enhanced performance when handling masked inputs, effectively reducing reconstruction artifacts.
Our combination of STN with VAEs not only improves the fidelity of patient-specific skeletal models but also sets the stage for integrating such techniques into surgical planning workflows.
%Future work will focus on refining registration robustness, exploring more flexible latent representations, and seamlessly embedding these reconstructions into clinical decision-making systems to ultimately guide and improve bone repair surgeries.