SAFRON: Stitching Across the Frontier Network for Generating Colorectal Cancer Histology Images
Paper Review
이번에 리뷰하고자 하는 논문은 SAFRON: Stitching Across the Frontier Network for Generating Colorectal Cancer Histology Images 입니다. 해당 논문은 GAN을 통해서 Histology image에 대해서 데이터를 생성하고자 하는 Medical ai의 논문 중 하나입니다.
Abstract
Computational pathology 분야에서 Histology image의 generation은 상당히 중요한 영역으로 여겨지고 있습니다. 이 때, Histology image의 경우 사이즈가 매우 크고 고해상도 이미지 중 하나입니다. 하지만 기존까지의 generation 연구의 경우, 해상도가 높아지면 높아질 수록 generation 성능이 한계점에 도달한다는 문제점들이 존재하였습니다. 더불어 해상도가 높아 이로 인한 당시의 리소스 제약 또한 존재하였습니다.
이를 해결하고자 본 논문은 SAFRON Framework를 제안하고 있습니다. 해당 프레임워크는 Tissue component masks을 조건으로 하여 사실적이고 큰 고해상도 조직 이미지에 대해서 generation 성능을 끌어올리고자 하였습니다. 더불어 Image 전체를 한번에 모델에 넣는 것이 아닌 Patch 단위로 나눠 훈련한 다음, 이를 Stitching 기법을 통해서 생성하고자 합니다. 위 두가지의 기술적 mechanism으로 인해서 형탵학적 특징을 보존하고 Boundary artifact를 최소화 할 수 있습니다.
Introduction
의료 이미지 분석 분야에서의 WSI의 자동화된 분석이 상당한 주목을 받고 있습니다. 그러나 앞서서 이야기 했듯이 WSI의 특징은 크기가 크며 고해상도의 이미지로 구성되어 있습니다. 더불어 의료 이미지 분석을 위해 모델 학습 시, 대용량의 고품질 어노테이션 된 데이터셋이 필요한 상황입니다. 하지만 medical domain의 특성 상 대용량의 고품질 어노테이션 된 데이터셋을 구하기 어려운 환경입니다. 그렇기에 이를 생성형 모델을 통한 generation 연구를 통해서 해결하고자 하며 이러한 연구들로 개인 정보 문제를 해결하고 데이터 증강이 가능한 점이 존재하는 이점이 존재하였습니다. 그러나 WSI의 고해상도 이미지에서는 생성형 모델의 성능이 좋지 않은 문제점들이 존재합니다.
Related work
해당 논문의 특징은 본 논문 이전에 이루어진 연구들에 대해서 많이 언급을 하고 있습니다. 먼저 THeCOT 연구부터 확인해보도록 하겠습니다.
1. THeCOT(Tumor Heterogeneity in Colorectal Tissue) model
해당 연구에서는 사용자가 다양한 파라미터를 직접 조정하는 생성형 모델 연구인 Rule based model입니다. 해당 연구에서는 다음과 같은 파라미터들을 직접 조정할 수 있습니다.
- 암 등급(Cancer Grade)
- 세포 밀도(Cellularity)
- 세포 중첩 비율(Cell overlap ratio)
- 이미지 해상도(Image resolution)
- 렌즈 배율(Objective level)
하지만 사용자가 직접 다양한 파라미터들을 조정하기에 사용자에 따른 모델 성능이 영향을 받으며 더불어 rule based이기 때문에 이미지가 고해상도일수록 생성된 이미지의 Realism이 부족한 단점이 존재하였습니다.
2. GAN, cGAN(Conditional GAN) + Unsupervised learning
이전에 이루어진 Rule-based model의 한계점인 Realistic Image 성능에 집중한 연구입니다. 특히 Conditional GAN의 경우, 임이의 노이즈에서 이미지를 생성하는 GAN과 달리 어떤 종류의 이미지가 생성될 지 제어하는 Condition을 함께 활용한 방식입니다. 이로 인해서 주어진 조건에 맞는 이미지 생성이 가능하게 됩니다. 특히, GAN 계열의 모델은 Unsupervised Learning으로 학습을 진행하는 경우 성능이 가낭 높았음을 확인할 수 있습니다. 해당 연구의 경우 Relistic image 생성에는 성공하였으나 여전히 해상도가 높을수록 생성된 이미지가 실제 이미지와 구분이 쉽다는 한계점이 여전히 잔재하였습니다.
3. ProGAN, SR-GAN
ProGAN의 경우, 처음에는 저해상도 이미지를 생성한 다음 훈련 전반에 걸쳐 레이어를 점진적으로 확장하며 고해상도 이미지 생성을 수행합니다.
SR-GAN의 경우, 저해상도 이미지로부터 최대 4배까지 확대하여 고해상도 이미지를 생성하고자 합니다. 특히, 해당 방식에서는 skip-connection, PixelShuffle로 업샘플링을 수행합니다.
위 두 가지 연구 모두 높은 메모리 연산량을 요구하는 한계점이 존재하였습니다.
4. Multi-scale GAN
이미지 전체의 해상도를 계층적으로 확장하면서 학습을 진행합니다. 이는 앞서서 나온 ProGAN과는 반대로 동작하고 있습니다. 즉, 낮은 해상도부터 높은 해상도까지 생성을하게 되며 위 사진과 같이 각 스케일에 해당하는 Discriminator가 존재하게 됩니다. 하지만 Multi-scale GAN 또한 앞선 연구들과 동일하게 높은 메모리 연산량을 요구하게 됩니다.
Methodology
SAFRON framework의 경우, 크게 Generator 모델이 U-Net 구조를 가지고 있습니다. 각 mechanism의 동작 단계에 대해서 보겠습니다.
- 학습을 위해서 728 x 728 마스크 이미지를 입력하게 됩니다.(Tissue Component Mask)
- 입력 마스크에 대해서 zero-padding을 수행하게 됩니다. 즉, 해당 단계를 통해서 Patch 테두리의 경계 아티팩트 손실에 대해서 최소화할 수 있습니다.
- Overlap하여 여러개의 Patch를 생성합니다. 이 때, Overlap하지 않고 Patch를 나누게 된다면 Patch 사이의 이음새 부분에서 손실이 발생하기에 overlap으로 수행하게 됩니다.
- U-Net 구조의 Generator를 동작시킵니다. 해당 U-Net의 구조 특징은 인코더와 매칭되는 디코더가 존재하여 이들의 Skip connection을 통해서 해상도 손실 최소화 및 global 및 local한 feature에 대해서 학습에서 유리함을 가져갈 수 있습니다.
- 이렇게 작은 Patch들이 생성되어지게 됩니다.
- 이후, 만들어진 작은 Patch에 대해서 Stitching 기법을 수행하게 됩니다. 공간 좌표에 따라 Overlap하여 연겨결하게 되며 중첩된 부분읜 경우, 공간 평균을 수행하게 됩니다.
- 이렇게 연결한 Patch들을 Tile로 부르게 되며 해당 Tile level에서 동작하는 Discriminator를 수행하게 됩니다. 이로 인해서 입력된 이미지와 Mask로 Real/Fake를 판별하며 Tile level Discriminator를 동작함으로 Gloval Coherence를 부여할 수 있습니다.
이번에는 해당 framework에서 활용하고 있는 Loss function에 대해서 알아보도록 하겠습니다. 해당 framework에서 활용하고 있는 Loss function의 경우, 크게 3가지가 존재하며 GAN에서 활용하는 loss function과 유사하게 동작하고 있습니다.
해당 loss function의 경우 실제 이미지 Y와 생성된 이미지 사이의 L1 거리를 측정하게 됩니다.
두 번째 loss function의 경우, 기존 GAN의 손실함수와 동일한 구조입니다.
아미작 loss function의 경우 이미지의 재구성 정확도와 realism 정도를 균형있게 최적화하기 위해서 활용되며 저희가 GAN을 통해서 배운 MinMax를 동일하게 수행하고 있음을 확인할 수 있습니다.
Experiments
set up
해당 실험에서 활용한 데이터셋은 CRAG, DigestPath의 데이터셋을 활용하고 있으며 이미지 사이즈는 각각 동일하게 728 x 728을 활용하고 있습니다. 특히, labelling의 비율의 경우 양성과 악성이 975 : 1025로 라벨링 균형을 이루고 있으며 사용한 파라미터들은 다음과 같습니다.
Training params
- Epochs – 100
- Optimizer – Adam
- Learning Rate – 0.0001
- Momentum – 0.5
- BatchNorm 수행
1. 고해상도 4096 x 4096 이미지 비교
해당 실험에서는 Multi-scale GAN과 비교하였으며 해당 모델의 경우 고해상도에서 조직의 형태학적 세부 정보 손실이 발생했음을 원본 이미지와 비교했을 때 확인할 수 있습니다. 그러나 해당 논문에서 제안하고 있는 SAFRON의 경우, 실제 이미지와 비교했을 때 훨씬 더 현실적으로 보이기도 하며 조직의 형태학적 세부 정보까지 높은 수준으로 생성하고 있음을 확인할 수 있었습니다.
2. FID 성능 지표 비교
양성 및 악성 조직에 대해서 FID 평가 및 해상도에 따른 FID 평가를 수행하고 있습니다. 먼저, 양성 및 악성 조직에 대한 FID 평가에서 본 논문에서 제안하고 있는 기법이 가장 낮은 FID 점수를 기록했음을 확인할 수 있습니다.
더불어 해상도에 따른 FID 평가를 보시게 되면 ProGAN과 CycleGAN의 경우, 2048, 4096 해상도에서 가용 메모리 초과로 인한 실행 불가와 Multi-Scale GAN의 경우, 해상도 증가에 따라 FID 점수가 크게 증가하고 있음을 볼 수 있습니다. 그러나 본 논문에서 제안하고 있는 기법은 해상도가 증가함에 불구하고 가장 낮은 점수를 기록한 것을 볼 수 있습니다.
FID의 경우, 실제 원본 데이터의 분포와 생성해낸 이미지의 분포 간의 거리를 구한 성능 지표로 점수가 낮을수록 원래의 데이터 분포와 유사함을 의미하고 있습니다.
왼쪽 사진이 GT이며 오른쪽 사진이 본 논문에서 제안하는 모델이 생성한 이미지와 해당 이미지의 THeCoT 모델을 적용하여 Gland mask를 구한 것입니다. 이처럼 FID가 가장 낮은 점수를 기록했기에 실제 이미지와 비교했을 때, 매우 유사함을 확인할 수 있습니다.
3. Discriminator의 level에 따른 결과 비교
Discriminator를 동작 시키는 level이 Patch와 Tile가 있습니다.
Patch level Discriminator로 동작 시키면 각 패치는 단순히 이어 붙였을 때, 뚜렷한 이음새 및 아티팩트가 검출되었습니다. 그러나 Tile 단위로 동작시킨다면 패치 간의 이음새 및 아티팩트 검출이 되지 않았으며 생성한 이미지의 일관성을 확보할 수 있음을 위 사진에서 확인할 수 있습니다.
4. Augmentation에 따른 ResNet 성능 비교
SAFRON을 통해서 augmentation을 수행하고 Resnet을 통해서 성능 비교를 수행했을 때, augmentation을 수행하지 않은 baseline의 경우 93% 성능을 기록햇으나 augmentation을 수행하게 된 경우에 97%로 성능 향상이 이루어짐을 확인할 수 있었습니다.
5. 해상도에 따른 컴퓨터 연산량 및 RAM 용량
SAFON이 다른 모델들 보다 적은 연산량과 적은 RAM 요구량을 보였으며 이미지 생성 소요 시간 또한 고해상도에서도 매우 빠른 속도로 수행됨을 확인 할 수 있었습니다.
Conclusion
해당 논문의 기여는 다음과 같습니다.
-
Patch 단위 이미지 생성 - 메모리와 연산량을 낮게 유지
-
Stitching Mechanism - 겹치는 영역의 픽셀들을 공간적으로 평균화 수행으로 이음새, 경계 아티팩트 없이 Seamless Image 생성 가능
-
Tile Level Discriminator - Patch 단위 생성된 이미지를 Stiching을 통한 Tile 단위로 Discriminator 동작하여 Patch 경계의 이음새를 최소화하고 Local, Global 형태학적 Coherence 유지 가능
-
적은 메모리와 낮은 연산량으로 고해상도에서도 높은 수준의 이미지 생성
