A Probabilistic U-Net for Segmentation of Ambiguous Images Simon A. - PowerPoint PPT Presentation

A Probabilistic U-Net for Segmentation of Ambiguous Images Simon A. A. Kohl 1*,2 , Bernardino Romera-Paredes 1 , Clemens Meyer 1 , Jeffrey De Fauw 1 , Joseph R. Ledsam 1 , Klaus H. Maier-Hein 2 , S. M. Ali Eslami 1 , Danilo Jimenez Rezende 1 , Olaf Ronneberger 1 1 DeepMind 2 German Cancer Research Center *work done during an internship at DeepMind

A Probabilistic U-Net for Segmentation of Ambiguous Images Simon A. A. Kohl 1*,2 , Bernardino Romera-Paredes 1 , Clemens Meyer 1 , Jeffrey De Fauw 1 , Joseph R. Ledsam 1 , Klaus H. Maier-Hein 2 , S. M. Ali Eslami 1 , Danilo Jimenez Rezende 1 , Olaf Ronneberger 1 1 DeepMind 2 German Cancer Research Center *work done during an internship at DeepMind Poster #127 Medical Imaging Workshop Talk: Sat, Dec 8, 9:45 am

Images are often Ambiguous 3

Images are often Ambiguous Potential Expert Graders Cancer 4

Images are often Ambiguous Potential Expert Graders Segmentations from our model (U-Net + conditional VAE) Cancer 5

Deterministic U-Net Inference Image U-Net 6

Probabilistic U-Net Sampling 𝛎 , 𝞃 prior Prior Net Latent Space Image U-Net 7

Probabilistic U-Net Sampling 𝛎 , 𝞃 prior 1 Sample 1 * Prior Net Latent Space Image U-Net 8

Probabilistic U-Net Sampling 𝛎 , 𝞃 prior 2 1 Sample 1 * Prior Net Latent Space 2 Image U-Net 9

Probabilistic U-Net Sampling 𝛎 , 𝞃 prior 2 1 Sample 3 1 * Prior Net Latent Space 2 Image 3 U-Net 10

Probabilistic U-Net Training 𝛎 , 𝞃 prior Prior Net Sample 𝐴 Cross- Entropy Image Sample Groundtruth U-Net 11

Probabilistic U-Net Training Position in Latent Space for this GT example? 𝛎 , 𝞃 prior Prior Net Sample 𝐴 Cross- Entropy Image Sample Groundtruth U-Net 12

Probabilistic U-Net Training Posterior Net 𝛎 , 𝞃 post 𝛎 , 𝞃 prior KL Prior Net Latent Space Sample 𝐴 Cross- Entropy Image Sample Groundtruth U-Net 13

Latent Space Analysis Probabilistic U-Net Image 14

Latent Space Analysis Probabilistic U-Net Image 15

Latent Space Analysis Probabilistic U-Net Graders 0 Image 1 2 3 16

Lung Abnormalities Segmentation: Quantitative Results Energy distance ( lower is better ) 17

Lung Abnormalities Segmentation: Quantitative Results Energy distance ( lower is better ) 1 18

Lung Abnormalities Segmentation: Quantitative Results Energy distance ( lower is better ) 1 4 19

Lung Abnormalities Segmentation: Quantitative Results Energy distance ( lower is better ) 1 4 8 20

Lung Abnormalities Segmentation: Quantitative Results Energy distance ( lower is better ) 1 4 8 16 21

Lung Abnormalities Segmentation: Quantitative Results Energy distance ( lower is better ) 22

Lung Abnormalities Segmentation: Quantitative Results Energy distance ( lower is better ) 23

Lung Abnormalities Segmentation: Quantitative Results Energy distance ( lower is better ) 24

Cityscapes segmentation: Qualitative Results Input Image Ground-truth Grader Samples (Probabilistic U-Net) Styles stochastic flips: sidewalk sidewalk 2 47 % person 41 % person 2 35 % car car 2 veget. 29 % veget. 2 road road 2 24 % 25

Cityscapes segmentation: Qualitative Results Input Image Ground-truth Grader Samples (Probabilistic U-Net) Styles stochastic flips: sidewalk sidewalk 2 47 % person 41 % person 2 35 % car car 2 veget. 29 % veget. 2 road road 2 24 % 26

Cityscapes segmentation: Quantitative Results 27

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 28

A Probabilistic U-Net for Segmentation of Ambiguous Images Simon A. A. Kohl 1*,2 , Bernardino Romera-Paredes 1 , Clemens Meyer 1 , Jeffrey De Fauw 1 , Joseph R. Ledsam 1 , Klaus H. Maier-Hein 2 , S. M. Ali Eslami 1 , Danilo Jimenez Rezende 1 , Olaf Ronneberger 1 1 DeepMind 2 German Cancer Research Center *work done during an internship at DeepMind Poster #127 Medical Imaging Workshop Talk: Sat, Dec 8, 9:45 am

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. 30

Evaluation Metric for Quantitative Comparison We use the Energy Distance 1 statistic (aka MMD): P gt P out where d(x,y) = 1 - IoU(x,y) and 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 31

Baselines Normal Prior 1 1 2 1 3 2 2 Sample 𝐴 1 , 𝐴 2 , 𝐴 3 ,... 1,2,3,... m m U-Net Dropout U-Net M-Heads Image2Image VAE U-Net Ensemble 32