The Coherent Point Drift for Clustered Point Sets

TL;DR

This paper introduces an extension to the Coherent Point Drift algorithm that incorporates prior clustering information to improve non-rigid point set registration accuracy, particularly in medical imaging applications.

Contribution

It extends the probabilistic framework of CPD to include clustering priors via a Gaussian mixture model, enhancing registration accuracy without significant performance loss.

Findings

Improved registration accuracy with clustering priors.

Effective for medical data, especially heart model personalization.

Comparable or better than existing methods.

Abstract

The problem of non-rigid point set registration is a key problem for many computer vision tasks. In many cases the nature of the data or capabilities of the point detection algorithms can give us some prior information on point sets distribution. In non-rigid case this information is able to drastically improve registration results by limiting number of possible solutions. In this paper we explore use of prior information about point sets clustering, such information can be obtained with preliminary segmentation. We extend existing probabilistic framework for fitting two level Gaussian mixture model and derive closed form solution for maximization step of the EM algorithm. This enables us to improve method accuracy with almost no performance loss. We evaluate our approach and compare the Cluster Coherent Point Drift with other existing non-rigid point set registration methods and show it's advantages for digital medicine tasks, especially for heart template model personalization using patient's medical data.

Figures

Tables

Table 1. Table 1: Comparison of proposed method with ECPD and CPD.

Metric	ccpd	cpd	ecpd
Hausdorff	8.88	9.29	9.07
Cluster Hausdorff	8.80	23.09	19.13

Equations

$$p(x)=\sum _{m=1}^{M+1}P(m)p(x|m)$$

$$p(x)=\frac{\omega }{N}+(1-\omega )\sum _{m=1}^{M+1}p(x|m)$$

$$E(\Theta ,{\sigma }^2)=-\sum _{n=1}^{N}log(p(x_n))$$

$$Q=-\sum _{n=1}^N\sum _{m=1}^{M+1}{P}^{old}(m|x_n)log(P^{new}(m)p^{new}(x_n|m))$$

$$Q(\Theta ,{\sigma }^2)=\frac{1}{2{\sigma }^2}\sum _{n=1}^N\sum _{m=1}^M{P}^{old}(m|x_n)\Vert x_n-T(y_m,\Theta ){\Vert }^2+\frac{N_{p}D}{2}log{\sigma }^2$$

$$p_{mn}=\frac{ex{p}^{-\frac{\Vert x_n-T(y_m,\Theta ){\Vert }^2}{2{\sigma }^2}}}{{\sum }_{k=1}^{M}ex{p}^{-\frac{\Vert x_n-T(y_k,\Theta ){\Vert }^2}{2{\sigma }^2}}+c}$$

$$T(Y)=Y+v(Y)$$

$$f(v,{\sigma }^2)=E(v,{\sigma }^2)+\frac{\lambda }{2}\varphi (v)$$

$$\varphi (v)=\Vert v{\Vert }_{H_m}^2=\Vert Pv{\Vert }^2$$

$$v(z)=\sum _{m=1}^M{w}_{m}G(z,y_m)$$

$$(G+\lambda {\sigma }^{2}d(P1{)}^{-1})W=d(P1{)}^{-1}PX-Y$$

$$p_{mn}=\frac{ex{p}^{-\frac{\Vert x_n-T(y_m,\Theta ){\Vert }^2}{2{\sigma }^2}}}{{\sum }_{k=1}^{M}ex{p}^{-\frac{\Vert x_n-T(y_k,\Theta ){\Vert }^2}{2{\sigma }^2}}+\frac{{(2\pi {\sigma }^2)}^{\frac{D}{2}}\omega }{1-\omega }\frac{M}{N}}$$

$$P_c(N_c)=\prod _{(j,k)\in N_C}{p}_c(x_j,y_k)$$

$$\tilde{p}(x)=P_c(N_c)p(x)$$

$$\tilde{E}(\Theta ,{\sigma }^2)=E(\Theta ,{\sigma }^2)+\sum _{(j,k)\in N_c}log(p_c(x_j,y_k))$$

$$\tilde{Q}=Q+\frac{1}{2{\alpha }^2}\sum _{n=1}^N\sum _{m=1}^M{\tilde{P}}_{mn}\Vert x_n-T(y_m,\Theta ){\Vert }^2$$

$${\tilde{p}}_{mn}=\{\begin{array}{ll}1,& \text{mbox}if(m,n)\in N_c\\ 0,& \text{mbox}otherwise\end{array}$$

$$\tilde{Q}(\Theta ,{\sigma }^2)=$$

$$\frac{N_{p}D}{2}log{\sigma }^2+\frac{\lambda }{2}\varphi (v)$$

$$(d(P1)G+\frac{{\sigma }^2}{{\alpha }^2}d(\tilde{P}1)G+\lambda {\sigma }^{2}I)W=PX-d(P1)Y+\frac{{\sigma }^2}{{\alpha }^2}(\tilde{P}X-d(\tilde{P}1)Y)$$

$$\bar{p}(x)=\sum _{c=1}^{C+1}P(c)\sum _{m=1}^{M+1}p(x|m,c)P(m|c)$$

$$\bar{E}=-\sum _{n=1}^N\bar{p}(x)$$

$$\bar{Q}=-\sum _{c=1}^{C+1}P(c)\sum _{n=1}^N\sum _{m=1}^{M+1}{P}^{old}(m|x_n,c)log(p^{new}(m|x_n)p^{new}(m|c))$$

$$\bar{Q}(\Theta ,{\sigma }^2)=\frac{1}{{\sigma }^2}\sum _{c=1}^{C+1}P(c)\sum _{n=1}^N\sum _{m=1}^M{P}^{old}(m|x_n,c)\Vert x_n-T(y_m,\Theta ){\Vert }^2+\frac{{\bar{N}}_{p}D}{2}log({\sigma }^2)$$

$$P(m|x_n,c)=\frac{P(c|x_n,m)p(x|m)}{{\sum }_{k=1}^{M}P(c|x_n,m=k)p(x|m=k)}$$

$$\bar{Q}(W)=\frac{1}{{\sigma }^2}\sum _{c=1}^{C}P(c)\sum _{n=1}^N\sum _{m=1}^M{P}^{old}(m|x_n,c)\Vert x_n-(y_m+G(m,)W){\Vert }^2+\frac{\lambda }{2}tr(W^{T}GW)$$

$$\frac{\partial \bar{Q}(W)}{\partial W}=\frac{C}{{\sigma }^2}\sum _{c=1}^{C}P(c)(d(\bar{P}(c)1)(Y+GW)-\bar{P}(c)X)+\lambda GW=0$$

$$(\sum _{c=1}^{C}P(c)d(\bar{P}(c)1)W+\lambda {\sigma }^{2}I)W=\sum _{c=1}^{C}P(c)(\bar{P}(c)X-d(\bar{P}(c)1)Y)$$

$${\sigma }^2=\frac{1}{{\bar{N}}_{p}D}(tr(X^T(\sum _{c=1}^{C}P(c)d(\bar{P}(c)1))X)-2tr(((\sum _{c=1}^{C}P(c)\bar{P}(c))X{)}^{T}T)+tr(T^{T}d(\sum _{c=1}^{C}P(c)\bar{P}(c)1)T))$$

$$\beta >0,\lambda >0,0\le \omega \le 1,{\sigma }^2=\frac{1}{DMN}\sum _{n=1}^N\sum _{m=1}^M\Vert x_n-y_m{\Vert }^2.$$

$$G(y_i,y_j)=e^{-\frac{1}{2{\beta }^2}\Vert y_i-y)j{\Vert }^2},P(c|x_n,m)=\frac{1}{NM}\sum _{n=1}^{N}I(x_n\in c)\sum _{m=1}^{M}I(y_m\in c);$$

Full text

1 Introduction

Points set registration task often arises in many computer vision problems. Statistical shape reconstruction, medical image registration, template personalization and various number of other tasks incorporate point set registration. This task consists of computing point correspondences and spatial transformation between two given point sets. These sets could be artificially generated, obtained with scanner or extracted from an image using various point detection approaches. In practice point set registration is primarily used for obtaining transformation between pre-defined template and points extracted from some data. In this paper we consider a special case of point set registration - where information on point sets clustering is already known. As an example, points extracted from CT scan of the person’s heart can have such property. Heart ventricles and atriums are considered as clusters in this case.

The task of non-rigid point set registration is ill-posed by itself. So additional constrains are used to achieve a solution. In general, the type of the transformation between template set and extracted set is unknown, so the solution of the problem results in finding optimal function with respect to problem’s constrains in non-rigid transformations class. These constrains are motivated by assumptions on point sets distribution and target transformation type. Desired solution should describe original transformation the best, be robust to outliers, noise and missing data, be able to process high dimensional data and have reasonable computational complexity. The Coherent Point Drift [8] stands out of other methods[12] and is able to provide accurate results in reasonable time. Still, it fails to perform in some complex cases. In this paper we extend it’s capabilities by taking the idea of including prior information described by V. Golyanik et al. [5] and expand it further to include information on point sets clustering. This information often comes naturally with properties of the point extraction algorithm or nature of the data. For example, many biomedical tasks include segmentation as pre-processing tool but no information on clusters is actually used. In this paper we present CPD modification for clustered point sets that extends existing approach to adopt prior information. In the experiments section we evaluate proposed method using 20 CT samples from MMWHS-2017 [16, 15, 14] challenge and compare results with CPD[8] and ECPD [5]. For the results evaluation we are using Hausdorff distance [13] and it’s variant for clustered point sets. This shows advantages of the proposed method over existing ones for the given task.

2 Related Work

The Coherent Point drift algorithm as a working idea was formulated at 2007 by Myronenko et al. [9] for the first time. The idea behind the method is to formulate point set registration as a maximum likelihood estimation problem. The transformation (model parameters) aligning point set Y with point set X that maximizes likelihood with subject to constrains is chosen as the result. Only non-rigid registration case is considered in that work. The authors expand this approach further at [8] for both rigid and non-rigid cases. Techniques reducing algorithm’s complexity and significantly improving performance were proposed. Authors demonstrate method’s advantages over existing approaches like LM-ICP [4] and TPS-RPM [3] in presence of noise and outliers. Review of the method’s modern applications is given in [6].

Golyanik et al. extended existing approach to take prior knowledge into account [5]. In their work they adopted coarse-to-fine strategy for processing large point set in reasonable time. Authors considered a special case of the non-rigid registration problem, when a sparse set of correspondence priors is known in advance. A modification for Coherent Point Drift was proposed that allows to embed correspondence priors in a closed-form. ECPD employs correspondence preserving subsampling counterbalancing the polynomial complexity by splitting the problem into two subproblems of smaller size and reducing the number of operations by a linear factor.

In [10], two novel elements were proposed: (1) generalization of the CPD (gCPD) that used a procedure of group-wise registration (alignment) for statistical shape reconstruction and (2) more complex point sets, in which each point that belonged to the surface was associated with a normal; moreover, instead of the Gaussian mixture model (GMM), the hybrid mixture model (HMM) was employed that combined the Student t-distribution and the von Mises–Fisher distribution for a point and a normal, respectively. Experiments used clinical data on 27 brain ventricles (Neuro) and 15 hearts (Cardiac). The HMM was shown to provide a significant increase in accuracy compared to the methods (including the gCPD) that used conventional point sets. As a metric, the mean surface distance (MSD) was employed.

2.1 Coherent Point Drift

The Coherent point drift algorithm was introduced in [9] and improved later in [8]. This method is designed to align two point sets $X$ and $Y$ where one point set represents GMM centroids and another represents data. $X$ points set representing the data is considered to be fixed, $Y$ is moving. In order to make $Y$ points move coherently regularization is used.

Authors are using the following notation we will also stick to.

• •

$D$ - dimension of the point sets;

• •

$N,M$ - number of points in X and Y respectively;

• •

$X_{NxD=(x_{1},...,x_{N})^{T}}$ - the first point set representing the data;

• •

$Y_{MxD=(y_{1},...,y_{M})^{T}}$ - the second point set representing the centroids;

• •

$T(Y,\Theta)$ - transformation $T$ applied to $Y$, where $\Theta$ is the transformation parameters;

• •

$I$ - identity matrix;

• •

$1$ - column vector of all ones;

• •

$d(a)$ - diagonal matrix formed from the vector a.

The $Y$ set is GMM centroids with probability density written as

[TABLE]

where $p(x|m)=\frac{1}{{(2\pi\sigma^{2})}^{\frac{D}{2}}}exp^{-\frac{\lVert x-y_{m}\rVert^{2}}{2\sigma^{2}}}$. To account presence of the noise and outliers $P(x|M+1)=1/N$ was added to mixture model. Equal isotropic covariances $\sigma^{2}$ and equal membership probabilities $P(m)=1/M$ are used. Resulting mixture model density can be written as

[TABLE]

where $\omega$ is the fraction of the noise in the model.

The model parameters $\Theta$ are estimated by maximizing likelihood function. This task is equivalent to minimization of the negative log-likelihood $E(\Theta,\sigma^{2})$ that can be written as

[TABLE]

For parameters estimation EM algorithm [1] is used. It consists of two steps, first, posterior probabilities $p(m|x_{n})$ are calculated using old parameters, then new parameters values are estimated during second step. The upper bound of the negative log-likelihood can be written as

[TABLE]

rewriting it we can get the following:

[TABLE]

where $N_{p}=\sum_{n=1}^{N}{\sum_{m=1}^{M+1}{P^{old}(m|x_{n})}}\leq N$ and $P^{old}$ - posterior probabilities of mixture components calculated with old parameters values. $P^{old}$ can be written as

[TABLE]

where $c=\frac{{(2\pi\sigma^{2})}^{\frac{D}{2}}\omega}{1-\omega}\frac{M}{N}$. Minimizing the function $Q$, we necessarily decrease the negative log-likelihood function $E$, unless it is already at a local minimum. In the paper [8] both rigid and non-rigid cases are considered. From biomedical applications perspective non-rigid case is more valuable, so we will focus on non-rigid CPD version. In order to represent non-rigid transformation the displacement field $v$ is introduced.

[TABLE]

Adding regularization term to negative log-likelihood we obtain

[TABLE]

where $E$ is negative log-likelihood and $\phi$ is regularization term. Authors show that regularization term can be rewritten as

[TABLE]

where operator $P$ extracts the high frequency content [8, 2, 11]. Authors show that displacement function $v(z)$ can be written as

[TABLE]

where $w_{m}=\frac{1}{\lambda\sigma^{2}}\sum_{n=1}^{N}{P^{old}(m|x_{n})(x_{n}-(y_{m}+v(y_{m}))}$ and $G(z,y)=e^{-\frac{1}{2\beta^{2}}\lVert z-y\rVert^{2}}$. Matrix $W=(w_{1},...,w_{m})^{T}$ can be evaluated by solving following equation

[TABLE]

where $g_{i,j}=G(y_{i},y_{j})$ and

[TABLE]

2.2 Extended Coherent Point Drift

Extended Coherent Point Drift [5] integrates prior knowledge into registration algorithm. It requires sparse set of points correspondences between point sets to be defined.

These correspondences are defined as pairs $(x_{j},y_{k})$ where $(j,k)\in N_{c}\subseteq N^{2}$. Correspondence priors are modeled as a product of particular independent density functions

[TABLE]

where $p_{c}(x_{j},y_{k})=\frac{1}{{2\pi\alpha^{2}}^{\frac{D}{2}}}exp^{-\frac{\lVert x_{j}-T(y_{k},\Theta)\rVert^{2}}{2\alpha^{2}}}$ and $\alpha>0$ is priors’ degree of reliability. Correspondence priors are incorporated into GMM

[TABLE]

Modified negative log-likelihood will be written as

[TABLE]

Upper bound $Q$ will be defined as

[TABLE]

where $\tilde{P}_{mn}$ defined as

[TABLE]

$\tilde{P}_{mn}$ matrix can be pre-calculated. For non-rigid case upper-bound function $\tilde{Q}$ will be written as

[TABLE]

The following linear system is solved during maximization step of the EM algorithm

[TABLE]

3 Proposed method

In this paper we consider a special case for point set registration, where point sets are clustered and correspondences between clusters of the different sets are known. These points can be extracted form heart CT scans so clusters will represent different heart cavities. In general, such information can be extracted for various number of biomedical tasks using different segmentation techniques like neural nets [7]. In the paper we will be using following notation:

• •

$c\in N$ - number of the cluster, $C$ clusters in total;

• •

$P(c|x_{n},m)$ - probability of $x_{n}$ and $y_{m}$ both to be inside cluster $c$;

• •

$P(c|x_{n})$ - probability of $x_{n}$ to be inside cluster $c$;

• •

$P(c|m)$ - probability of $y_{m}$ to be inside cluster $c$;

• •

$P(c)$ - probability of the cluster $c$;

Let us define additional cluster $C+1$ so that $P(C+1|x_{n},m)=1-\sum_{c=1}^{C}{P(c|x_{n},m)}$ with $P(C+1)=0$. This cluster includes all pairs $(x_{n},y_{m})$ that don’t share any cluster. We will rewrite probability density function of the mixture model as

[TABLE]

The negative log-likelihood is defined as

[TABLE]

And it’s upper bound can be written as

[TABLE]

[TABLE]

where $\bar{N}_{p}=\sum_{c=1}^{C+1}{P(c)\sum_{n=1}^{N}{\sum_{m=1}^{M}{P^{old}(m|x_{n},c)}}}$. The posterior probabilities $P(m|x_{n},c)$ can be calculated as follows

[TABLE]

After adding regularization and removing the term that corresponds to $P(C+1)$ we can write upper bound as

[TABLE]

After rewriting equations in matrix form and setting derivative to zero we will get

[TABLE]

So during maximization step we iteratively update matrix $W$ using following equation:

[TABLE]

where $\bar{P}(c)=P^{old}(m|x_{n},c)$. Setting derivatives to zero we can get $\sigma^{2}$ update rule

[TABLE]

where $T=Y+GW$

Strictly speaking $P(c|x_{n},m)$ is unknown, but it has to be initialized. From it’s definition - probability of given $x_{n}$ and $y_{m}$ to be in the same cluster $c$; we can estimate it with product $P(c|x_{n},m)=\frac{1}{NM}\sum_{n=1}^{N}{I(x_{n}\in c)}\sum_{m=1}^{M}{I(y_{m}\in c)}$, where $I(x_{n}\in c)$ equals to 1 if $x_{n}$ belongs to cluster $c$ and equals to 0 otherwise.

As a result, Cluster Coherent Point drift can be described as

• •

inputs: point sets $X$ and $Y$, probabilities $P(c|x_{n})$ and $P(c|y_{m})$;

• •

outputs: aligned point set T(Y);

• •

algorithm: CCPD;

• •

Initialize:

[TABLE]

• •

Repeat until convergence;

• –

E-step: compute

[TABLE]

• –

M-step: solve

[TABLE]

• –

update $T,\bar{N}_{p},\sigma^{2}$

• •

Aligned set is $T=Y+GW$

The key difference between ECPD and proposed method is the assumption made about prior information.

4 Evaluation

We implemented CPD, ECPD and CCPD in C++ and tested them on Xeon E5-2660 with K40. Efficient GPU version was implemented with CUDA C++. All experiments were conducted using GPU version. We tested methods listed above with MMWHS-2017 dataset [16, 15, 14]. This dataset consists of multiple CT and MRI labeled heart scans (figure 1). Among other labels, heart ventricles and atriums are landmarked. For the experiment we will be using only these four labels. Experiments were conducted on 20 training CT volumes with corresponding ground-truth segmentation with goal to align STL template with blood cavities surfaces.

After aligning template’s ventricles and atriums, whole heart structure can be personalized using resulting transformation. Unfortunately we cannot measure quality of the alignment for neighborhood structures. An example of template personalization shown at figure 2.

The results of template registration are shown at figure 3. CPD doesn’t incorporate any prior information and ECPD uses modified $\tilde{P}$ matrix where

[TABLE]

. Proposed methods shows the best result among others.

For quantitative evaluation two metrics were used: Hausdorff distance and it’s cluster variant. Let us write Hausdorff distance as follows:

[TABLE]

where $h(A,B)=\max_{a\in A}{\min_{b\in B}{\lVert a-b\rVert}}$, and it’s cluster variant as a mean of $H(A(c),B(c))$:

[TABLE]

where $A(c)\subseteq A$ so that $a\in A(c)$ corresponds to cluster $c$.

Table 1 shows experimental results. For ECPD column the best result is present. First metric shows quality of alignment with no known information during evaluation stage. The second metric computes average between ground truth clusters, that is more relevant in our case. For both cases proposed method shows better results.

5 Conclusion & Discussion

The problem of non-rigid point set registration is a key problem for many computer vision tasks. In the paper we discussed point set registration problem for clustered point clouds. As an example, points extracted from CT scan of the person’s heart can have such property. Heart ventricles and atriums are considered as clusters in this case.

We presented modified version of Coherent Pint Drift for special type of point sets - sets that can be clustered. We defined loss function and showed the way to optimize it using EM algorithm. In the experiments section we compared three methods: proposed one, Coherent Point Drift and Extended Coherent Point Drift that is able to embed prior information on sparse correspondences. We evaluated them in the same environment using MMWHS-2017 data. Proposed method shows more then two times better accuracy than its competitor.

Presented approach is robust to outliers and noise presence and able to produce precise enough transformation even for heart computational models personalization.