Semantic Segmentation with Noisy Boundary Annotations

Implemented boundary detection based on "Devil is in the Edges: Learning Semantic Boundaries from Noisy Annotations" (see link), generalized it to 3-dimension cases.

Implementation

• Basic 2D/3D CASENet network with weighted multilabel BCE loss

• 2D/3D Geodesic active contour inference

• Iterative update between network training and level-set refinement

• 3D UNET (obsolete code)

• NMS loss and direction losstr

Configuration example

1. Setup configuration for 2D CASENet with 2D level set: training (link), testing (link)
2. Setup configuration for 3D CASENet with 3D level set: trianing (link)
3. obsolete configuration code: UNet3D traning(link), testing (link)

Usage

Clone this repo

git clone http://gitlab.bj.sensetime.com/shenrui/edgeDL.git
cd edgeDL

Install dependencies

Require Python 3.6+ and Pytorch 1.0+. Please install dependencies by

conda env create -f environment.yml

Preprocessing

Resample the data into same resolution. This code requires Free Surfer mri_convert (see link, Free Surfer installation guide).

./utils/resample.sh

Generate file lists for traning, validation and testing sets.

python data2txt.py

Traning

Setup configuration file and run

python train_casenet.py --config PATH_TO_CONFIG_FILE

Testing

Setup configuration file and run

python predict_casenet.py --config PATH_TO_CONFIG_FILE

Loss function

1. Weighted multilabel BCE loss

   $\mathcal{L}_{BCE}(\theta) = - \sum_k\sum_m{\beta y_k^m\log f_k(m|x,\theta) + (1-\beta) (1-y_k^m)\log(1 - f_k(m|x,\theta))}$

   where

   $\beta$ : non-edge pixels/voxels ratio, $\beta = \frac{|Y^-|}{|Y|}$

   $k$ : class

   $m$ : pixel/voxel

2. NMS loss (edge thinning layers, to be implemented)

   $\mathcal{L}_{NSM}(\theta) = -\sum_k\sum_p \log h_k(p|x,\theta)$

   where

   $h_k(p|x,\theta) = \frac{\exp(f_k(p|x,\theta)/\tau)}{\sum_{t=-L}^L \exp(f_k(p_t|x,\theta)/\tau)}$ for normalization

   $x(p_t) = x(p) + t · \cos \vec{d_p} $ , $y(p_t) = y(p) + t · \sin \vec{d_p}$

   $p$ : gt boundary pixel/voxel

   $\vec{d_p}$ : normal direction at $p$ computed from gt boundary map

   $t \in {-L, -L+1, ... L}$

   Notes for implementation

   • Normal direction: use a fixed convolutoonal layer to estimate second derivatives, and then use trigonometry function to compute normal direction from the gt boundary map
   • code reference: edgesNMS(link)

3. Direction Loss (to be implemented)

   $\mathcal{L}_{Dir}(\theta) = \sum_k\sum_p ||\cos ^{-1} <\vec{d_p}, \vec{e_p}(\theta)>||$

   where

   $\vec{e_p}(\theta)$ : normal direction at p computed from prediction map

Level Set

1. Level set evolution

   $\frac{\partial \phi}{\partial t} = g_k(\kappa + c)|\nabla\phi| + \nabla g_k · \nabla \phi$

   solved by morphological approach (see link)

2. Energy (edge) map for level set alignment

   $g_k = \frac{1}{\sqrt{1+\alpha f_k}}+\frac{\lambda}{\sqrt{1+\alpha \sigma(y_k)}}$

   where

   $f_k$ : probability map predicted by neural network

   $\sigma(y_k)$ : (previous) ground truth annotation smoothed by gaussian filter with $\sigma$

Github reference:

1. STEAL (link)
2. edges (link)
3. Morphsnakes (link)