Rohit
Rohit PhD student at University of Pennsylvania

TopNet: Topology Preserving Metric Learning

TopNet: Topology Preserving Metric Learning

These are some of my notes for the MICCAI 2020 paper TopNet: Topology Preserving Metric Learning.

Motivation

  • Reconstructing Portal Vein and Hepatic Vein trees from abdominal CT scans is a prerequisite for liver surgery.
  • Locally, both types of veins look similar, so dense segmentation based on patch wise DL wouldn’t work, and the trees need to be traced to their source.
  • Proposed Solution: Learn to predict centerlines, and a metric that connects these centerlines. Use source vessel detection to assign labels to vascular trees.

Introduction

Normal segmentation makes local mistakes which are amplifies the error across the image.

/blog/assets/images/topnet/Untitled.png

Contributions:

  1. Multi-task network for tree reconstruction which detects centerlines and connectivity between centerlines.
  2. Topology metric which learns both intra and interclass topological distance between vascular voxel pairs.
  3. Verified better accuracy

Method

There are 3 parts in their multi-task network.

/blog/assets/images/topnet/Untitled%201.png

1. Vessel extraction decoder

Separate vessels from background. Formulated as a semantic segmentation problem. Use dice loss to train this task.

/blog/assets/images/topnet/Untitled%202.png

where \(v\) is the ground truth binary vessel volume patch, \(v'\) is the predicted volume patch.

2. Centerness decoder

The centerness score can be formulated as a regression problem (it will be unstable to train as a binary segmentation problem).

Let \(m_i \in M\) be the binary volume representing center voxels.

For a given voxel \(j\), define distance transform:

\[S_j = [ \text{distance to closest } m_i = 1 ]\]

The values are considered only for points inside vessel (a lot of points outside vessel will skew the output otherwise).

Use smooth L1 loss:

/blog/assets/images/topnet/Untitled%203.png

The denominator ensures that the outputs are not skewed towards voxels away from center (there are \(\propto S^2\) voxels at a distance \(S\) from centerline).

Perform NMS (non-maximal suppression) over a window size of \(5\times 5 \times 5\).

3. Topology distance Decoder

This network outputs a \(8d\) feature vector for each voxel in CT volume.

The idea here is that the distance between the features \(x_i, x_j\) at voxels \(i, j\) given by \(\| x_i - x_j\|\) corresponds to topological distance in the vessel space if they are in the same tree.

If they are in different trees, the distance should be large.

Given labels \(l_i, l_j\) for voxels \(i, j\), the loss function is given by:

/blog/assets/images/topnet/Untitled%204.png

For a given input image patch, the loss is:

/blog/assets/images/topnet/Untitled%205.png

where \(N_i\) is a neighborhood of voxel \(i\).

Neighborhood is defined inside a radius of 15 voxels ⇒ topological distance is bounded.

They set \(\alpha = \frac{1}{15}\) to approximately normalize topological distance \(\alpha D_{ij}\) from 0 to 1. Also set \(K = 3\). To balance loss terms, set \(\gamma = \frac{1}{3}\).

Tree reconstruction (during test-time post-training)

  1. Use vessel prediction to mask center-line distance.
  2. Use NMS on center-line distance \(S_i\) to find center voxels.
  3. For voxel \(i\), create edges for all nodes \(j\) where \(\|x_i - x_j\| \le 2\).
  4. Use Dijkstra’s multi-source shortest path tree algorithm.

Experiments

Datasets

IRCAD public dataset, internal dataset.

Methods

  1. Comparison with single task 3D UNet (dice loss).
  2. For multi-task, replace topology distance with 2-channel probability map (dice loss over all center-voxels).
  3. Multi-task cosine metric learning (\(S_{ij}\) = cosine distance)

    /blog/assets/images/topnet/Untitled%206.png

    For tree reconstruction, check where \(S_{ij} <= 0.5\) and take edge weights as:

    \[w_{ij} = E_{ij} (1 - S_{ij})\]

where \(E_{ij}\) is the euclidean distance between voxels \(i, j\).

/blog/assets/images/topnet/Untitled%207.png

  • Without metric learning, there is abrupt change in vessel labels along vascular tree.
  • Probably reason is that cosine based or single-task try to classify based on frequently observed vessel patterns.
  • But topology metric explicitly constrains to tracing vessels (so that the topology distance is correct)

/blog/assets/images/topnet/Untitled%208.png

comments powered by Disqus