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 \overrightarrow{d_p}$

   $p$ : gt boundary pixel/voxel

   $\overrightarrow{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}<\overrightarrow{d_p},\overrightarrow{e_p}(\theta )>||$

   where

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

Level Set

1. Level set evolution

   $\frac{\partial \varphi }{\partial t}=g_k(\kappa +c)|\mathrm{\nabla }\varphi |+\mathrm{\nabla }{g}_k·\mathrm{\nabla }\varphi$

   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)