-
A Probabilistic U-Net for Segmentation of Ambiguous Images
Simon A. A. Kohl1*,2, Bernardino Romera-Paredes1, Clemens
Meyer1, Jeffrey De Fauw1, Joseph R. Ledsam1, Klaus H.
Maier-Hein2,
S. M. Ali Eslami1, Danilo Jimenez Rezende1, Olaf
Ronneberger1
1DeepMind2German Cancer Research Center
*work done during an internship at DeepMind
-
A Probabilistic U-Net for Segmentation of Ambiguous Images
Simon A. A. Kohl1*,2, Bernardino Romera-Paredes1, Clemens
Meyer1, Jeffrey De Fauw1, Joseph R. Ledsam1, Klaus H.
Maier-Hein2,
S. M. Ali Eslami1, Danilo Jimenez Rezende1, Olaf
Ronneberger1
1DeepMind2German Cancer Research Center
Poster #127Medical Imaging Workshop Talk: Sat, Dec 8, 9:45
am
*work done during an internship at DeepMind
-
3
Images are often Ambiguous
-
4
Images are often AmbiguousPotential Cancer
Expert Graders
-
5
Images are often AmbiguousPotential Cancer
Expert Graders Segmentations from our model (U-Net + conditional
VAE)
-
6
U-Net
Image
Deterministic U-NetInference
-
7
Probabilistic U-Net
U-Net
Latent SpacePrior Net
Image
𝛍,𝞂prior
Sampling
-
8
Probabilistic U-Net
U-Net
Latent SpacePrior Net
Image
𝛍,𝞂prior
Sampling
11
Sample
*
1
-
9
Probabilistic U-Net
U-Net
Latent SpacePrior Net
Image
𝛍,𝞂prior
Sampling
11
*
1
2
2
Sample
-
10
Probabilistic U-Net
U-Net
Latent SpacePrior Net
Image
𝛍,𝞂prior 13
2
*
1
Sampling
2
3
Sample
-
11
Probabilistic U-Net
U-Net
Image
Sample 𝐳
Sample Groundtruth
Cross-Entropy
Prior Net
𝛍,𝞂prior
Training
-
12
Probabilistic U-Net
U-Net
Image
Sample 𝐳
Sample Groundtruth
Cross-Entropy
Prior Net
𝛍,𝞂prior
Training
Position in Latent Space for this GT example?
-
13
Probabilistic U-Net
U-Net
Image
Posterior Net
Sample 𝐳
Sample Groundtruth
KL
Cross-Entropy
Prior Net Latent Space
𝛍,𝞂prior
𝛍,𝞂post
Training
-
14
Image
Latent Space AnalysisProbabilistic U-Net
-
15
Image
Latent Space AnalysisProbabilistic U-Net
-
16
Graders
Image
Latent Space AnalysisProbabilistic U-Net
1
0
3
2
-
17
Lung Abnormalities Segmentation: Quantitative ResultsEn
ergy
dis
tanc
e (lo
wer
is b
ette
r)
-
18
Lung Abnormalities Segmentation: Quantitative ResultsEn
ergy
dis
tanc
e (lo
wer
is b
ette
r)
1
-
19
Lung Abnormalities Segmentation: Quantitative ResultsEn
ergy
dis
tanc
e (lo
wer
is b
ette
r)
1 4
-
20
Lung Abnormalities Segmentation: Quantitative ResultsEn
ergy
dis
tanc
e (lo
wer
is b
ette
r)
1 4 8
-
21
Lung Abnormalities Segmentation: Quantitative ResultsEn
ergy
dis
tanc
e (lo
wer
is b
ette
r)
1 4 8 16
-
22
Lung Abnormalities Segmentation: Quantitative ResultsEn
ergy
dis
tanc
e (lo
wer
is b
ette
r)
-
23
Lung Abnormalities Segmentation: Quantitative ResultsEn
ergy
dis
tanc
e (lo
wer
is b
ette
r)
-
24
Lung Abnormalities Segmentation: Quantitative ResultsEn
ergy
dis
tanc
e (lo
wer
is b
ette
r)
-
25
Cityscapes segmentation: Qualitative Results
sidewalk
person
car veget.
Input Image Ground-truth Grader Styles
Samples (Probabilistic U-Net)
person 2
car 2 veget. 2
sidewalk 2 47 %
41 %
35 %
29 %
stochastic flips:
road road 2 24 %
-
26
Cityscapes segmentation: Qualitative Results
sidewalk
person
car veget.
Input Image Ground-truth Grader Styles
Samples (Probabilistic U-Net)
person 2
car 2 veget. 2
sidewalk 2 47 %
41 %
35 %
29 %
stochastic flips:
road road 2 24 %
-
27
Cityscapes segmentation: Quantitative Results
-
28
Conclusions
● Learn conditional probability over segmentation maps
● Each sample is a valid & consistent segmentation
● The likelihoods are well calibrated
● Works on large-scale, real-world data
● Can also be trained with a uni-modal GT
● Can be used to asses annotations under the model
code: github.com/SimonKohl/probabilistic_unet
https://github.com/SimonKohl/probabilistic_unet
-
A Probabilistic U-Net for Segmentation of Ambiguous Images
Simon A. A. Kohl1*,2, Bernardino Romera-Paredes1, Clemens
Meyer1, Jeffrey De Fauw1, Joseph R. Ledsam1, Klaus H.
Maier-Hein2,
S. M. Ali Eslami1, Danilo Jimenez Rezende1, Olaf
Ronneberger1
1DeepMind2German Cancer Research Center
Poster #127Medical Imaging Workshop Talk: Sat, Dec 8, 9:45
am
*work done during an internship at DeepMind
-
30
Probabilistic Segmentation: Clinical Use-Cases
● Best-fit could be picked by clinician and adjusted if
necessary.
● Hypotheses could be propagated into next diagnostic pipeline
steps.
● Hypotheses could inform actions to resolve ambiguities.
-
31
Evaluation Metric for Quantitative ComparisonWe use the Energy
Distance1 statistic (aka MMD):
where d(x,y) = 1 - IoU(x,y) and
Pout Pgt
1 Székely, G.J., Rizzo, M.L.: Energy statistics: A class of
statistics based on distances. Journal of statistical planning and
inference 143(8) (2013) 1249–1272
-
32
Baselines
1
2
m
U-Net Ensemble
1
2
m
M-HeadsDropout U-Net
1,2,3,...
U-Net
Normal Prior
Sample 𝐳1,𝐳2,𝐳3,...
13
2
Image2Image VAE