PseudoEdgeNet: Nuclei Segmentation only with Point Annotations
이 글을 2019년 8월에 루닛 블로그에 올렸던 포스트입니다. (https://blog.lunit.io/2019/08/20/pseudoedgenet-nuclei-segmentation-only-with-point-annotations/)
이번 글에서는 올해 MICCAI2019 학회에서 발표될 예정인 논문(PseudoEdgeNet: Nuclei Segmentation only with Point Annotations)을 소개하고자 합니다. 세포핵 분할(nuclei segmentation)은 whole slide image (WSI) 병리 조직 분석에서 중요한 역할을 합니다. 하지만 WSI는 일반적으로 giga-pixel의 영상이기 때문에 의사가 일일히 모든 세포핵을 찾아 표시하기는 어렵습니다. 딥러닝 모델들은 대체로 좋은 성능을 보이지만, 학습을 위한 데이터가 많이 필요하다는 단점이 있기 때문입니다. 이번 논문에서는 딥러닝 모델들이 가지는 capacity의 차이를 이용하여, 세포핵의 경계를 전부 표시한 mask가 아닌 point label만을 써서 세포핵 분할을 학습시킵니다.
Introduction
Semantic segmentation 문제는 영상에서 각 픽셀이 어떤 물체의 category인지를 구분하는 문제입니다. 하지만 semantic segmentation의 data annotation은 물체의 테두리를 일일히 따야 하기 때문에 시간이 많이 걸립니다. 때문에 최근에는 다양한 weakly supervised learning 방법이 제안되고 있습니다. Weakly supervised learning을 통한 영상 분할의 대표적인 최근 논문으로는 Tang et al. [1] 이 있습니다. 이 논문에서는 PASCAL-Scribble dataset에 regularization loss를 적용해 기존 PASCAL VOC 데이터로 학습시킨 모델 성능의 97.6%를 달성하는 모델을 학습시킵니다. Mask 대신 scribble만으로 높은 성능을 냈다는 면에서 고무적인 결과입니다. 그러나 scribble annotation은 길이가 짧을수록 성능이 떨어지는 것을 실험 결과에서 볼 수 있습니다.
기존 세포 영상 데이터는 point annotation을 이용한 접근이 많습니다. 세포핵의 경계를 전부 그리는 것보다는 각 세포핵에 점을 찍는 게 annotation이 더 편하기 때문입니다. 그래서 세포를 찾는 다양한 접근과 모델은 기존에 다양하게 제안되었습니다. 하지만 이를 이용해 segmentation을 하는 모델은 제안되지 않았습니다. 최근 point annotation을 이용한 Laradji et al. [2] 은 segmentation scheme을 이용해 blob을 찾아 물체의 수를 세지만, 이 논문의 방법으로 찾은 blob은 물체의 경계를 맞추지 못합니다.
물체의 blob만 찾아진다는 것은 모델이 물체가 있을 것이라 여겨지는 정도를 pixel마다 표시해주는 것에 지나지 않습니다. 하지만 언뜻 보면 이는 이상합니다. 물체의 경계선을 찾는 것이 물체의 texture를 찾는 것보다 어려울까요? 직관적으로는 그렇지 않아 보입니다. 그럼에도 모델은 texture를 비교적 잘 찾는 반면에, (label로 주어지지 않은) 경계선은 못 찾고 있습니다.
이 논문에서는 network-capacity-gap이란 개념을 고안하여 이러한 문제를 해결하고자 합니다.
PseudoEdgeNet
논문에서 제시된 모델입니다. Segmentation은 [2]를 baseline으로 두고 같은 방식으로 모델을 학습시킵니다. Point annotation 부분에 작은 원을 그려 positive label을 주고, negative label은 세포핵을 가리키는 point들로 만든 Voronoi diagram의 경계로 구성하였습니다. 나머지 pixel들은 무시하고 앞에서 정의한 positive label과 negative label에만 binary cross entropy로 loss를 구합니다.
이렇게 학습시키면 [2]의 결과와 같은 문제가 생깁니다. 세포핵의 경계선을 따르지 않은 위양성 (false positive)와 위음성(false negative) pixel이 더 많이 나타나기 때문입니다. 그래서 PseudoEdgeNet이라는 모델을 추가로 사용합니다. 이 모델은 segmentation network가 낸 결과의 Sobel 경계선을 예측하도록 학습합니다. 학습에는 함수는 pixel 단위의
loss를 사용했습니다. Sobel filter는 미분가능하므로 edge network의 학습은 segmentation network에도 영향을 줍니다.
여기서 edge network는 segmentation network보다 capacity가 작습니다. Edge network가 segmentation network의 결과를 예측하는 것이 더 어려우므로, segmentation network가 edge network의 결과에 맞도록 학습됩니다. Edge network가 작기 때문에, 잘못된 edge 결과가 noise로 나타나는데, 이를 해결하기 위해 edge network보다 강한 attention module을 추가로 사용합니다. 아래는 학습을 위해 사용한 loss function입니다. I, P는 입력 영상과 point annotation이고, f, g, h는 각각 segmentation network, edge network, attention module입니다.
Experiments
맨 왼쪽 열과 오른쪽 열은 각각 입력 영상과 사람의 annotation 결과입니다 (a, g). PseudoEdgeNet과 함께 학습시킨 모델은 baseline 모델에서는 잡지 못하는 세포의 모양을 더 잘 잡습니다 (e, f). 또, edge network 가 segmentation 결과의 sobel edge를 잘 찾도록 학습되는 것을 볼 수 있습니다 ©. Edge network의 기본 결과는 다소 많은 noise가 끼어 있습니다. 이는 edge network가 적은 receptive field를 보도록 학습되기 때문인데, attention module을 통해 이를 보완할 수 있습니다 (b, d).
Baseline [2]과 DenseCRF [1] 모델에 비해 제시된 PseudoEdgeNet은 point annotation 더 좋은 성능을 보입니다. 특이한 점은, 큰 edge network는 PseudoEdgeNet을 써서 성능향상을 거의 얻을 수 없다는 점입니다. 이는 network-capacity-gap에 의해 나타나는 regularization 효과가 사라졌기 때문입니다.
이러한 network-capacity-gap은 단순히 network의 layer 개수보단, network 자체의 특성에서 크게 기인하는 것으로 보입니다. 단순하고 작은 CNN 구조에서 layer를 2, 4, 6, 8개로 차이를 주었지만 성능에 유의미한 차이를 보이지 않습니다. 하지만, segmentation network와 같은 구조인 FPN 구조로 가니 성능이 크게 떨어지는 것을 볼 수 있습니다.
Conclusion
이 논문에서는 처음으로 point annotation을 사용하여 세포핵 분할을 시도합니다. 또, 기존의 weakly supervised learning 방식과 달리 network-capacity-gap을 이용하는 방법을 제안합니다. Point annotation은 더 쉽게 annotation을 할 수 있다는 장점이 있으나, 같은 수의 데이터에서는 아직 supervised learning보다 뛰어나지 못합니다. 차후 연구에서는 적은 수의 contour annotation과 많은 수의 point annotation을 함께 사용함으로써 더 좋은 성능을 내볼 수 있지 않을까 기대해봅니다.
References
- On Regularized Losses for Weakly-supervised CNN Segmentation
- Where are the Blobs: Counting by Localization with Point Supervision
