같은 sequence의 MR 이미지라도 촬영 장비에 따라 contrast가 제각각이기 때문에 규격화된 하나의 contrast로 통일시키는 harmonization이 필요합니다. 이 논문은 이미지마다 다른 contrast를 style로 생각하고 규격화된 style로 변화시키는 방법을 제안합니다. 그리고 같은 환자가 대량의 다른 기계에서 동일한 포즈롤 촬영을 할 수 없기 때문에, unsupervised learning으로 연구를 진행했습니다.
https://www.biorxiv.org/content/10.1101/2021.03.17.435892v1.abstract
Introduction
MR image들은 MRI 장비를 통해 촬영된 영상이기 때문에 장비를 만드는 회사, 장비버전, 촬영기사의 세팅 등등에 따라 같은 sequence 영상일지라도 contrast가 제각각입니다.
위 그림처럼 같은 T1-weighted 영상일지라도 어떤 이미지에서는 밝고, 어떤 이미지에서는 어두울수가 있습니다. 그러므로 이런 제각각인 영상들을 하나로 통일시켜주는 harmonization 기술 연구가 활발히 이뤄지고 있습니다.
Architecture of Style-Encoding GAN
일반 영상분야에서 활발히 연구가 진행되고 있는 style transfer의 개념으로 이미지에 새로운 style을 입혀주는 방법을 이 논문도 활용했습니다.
위 그림처럼 저자가 제안한 모델에는 Generator, Discriminator, Style Encoder, Mapping Network 총 4개의 네트워크가 사용됩니다.
1. Mapping Network
그림에서 초록색인 mapping network는 간단한 MLP layer의 형태로 변화될 style을 만들어주는 네트워크입니다. latent vector z를 받아 layer를 통과해 바뀔 style vector s를 생성합니다. s = M(z)
2. Generator
Generator는 영상에 새로운 style을 입혀주는 형태로 학습이 진행됩니다.
Generator는 encoder - decoder 형태로 encoder에서는 이미지 사이즈를 줄여가며 high feature를 추출합니다. 그리고 decoder에서 style vector s를 AdaIn layer를 통해 입혀주면서 새로운 style의 output 영상을 만들어냅니다. y = G(x,s)
3. Discriminator
Discriminator는 generator를 통해 생성된 영상이 실제 target style의 영상과 비슷한지 여부를 판별해주는 네트워크입니다. 일반적인 GAN에서의 discriminator와 동일합니다.
4. Style Encoder
마지막으로 style encoder는 들어온 영상의 style을 추출하도록 학습하는 네트워크입니다. 즉, style encoder를 통해 생성되는 vector는 mapping network에서 생성하는 style vector와 동일한 형태입니다. s = E(x)
Network Training
총 4개의 loss function을 통해 각각의 네트워크가 각자의 역할을 잘 수행하도록 유도합니다.
1. Adversarial Loss
GAN을 base하는 모델이기 때문에 자연스럽게 계산되는 loss입니다. Generator를 통해 생성되는 target contrast의 이미지와 실제 이미지를 통해 generator, discriminator에서 계산됩니다.
2. Cycle-Consistency Loss
CycleGAN에서 사용한 loss로서, 원본영상에서 변환시킨 영상을 가지고 다시 원본영상처럼 바꾸면 실제 원본영상과 같아야한다는 개념입니다.
위 그림에서처럼 source 이미지를 가지고 generator를 통해 traget 이미지를 생성합니다. 그 뒤, 다시 원본영상의 contrast로 바꾸기 위해 generator를 한번더 통과시켜줍니다. 이 때, generator에 바꾸고 싶은 style vector를 함께 넣어줘야하는데, 이 style vector는 원본영상을 style encoder를 통해 추출한 것을 사용합니다. 즉, 원본 영상의 style을 추출한 뒤, 그 style을 입혀주면 원본 영상의 style처럼 바뀌어야만 한다는 개념입니다.
3. Style Reconstruction Loss
Style Reconstruction loss는 입혀준 style과 생성된 영상에서 추출한 style은 같아야한다는 개념입니다.
위 그림에서처럼 mapping network를 통해 생성된 style vector s를 generator에 넣어 target 이미지를 생성합니다. 그리고 생성된 이미지를 style encoder에 넣어주면 생성된 이미지의 style vector가 추출됩니다. 여기서 입혀줄때 넣어준 style가 생성된 이미지에서 추출된 style은 같아야만 합니다.
4. Style Diversification Loss
Style diversification loss는, generator가 좀 더 다양한 영상을 만들수 있도록 설계한 loss입니다.
랜덤한 latent vector z1, z2를 통해 생성된 style vector s1, s2를 가지고 genrator에 넣어 만들어진 G(x,s1)과 G(x,s2)의 이미지는 달라야만 한다는 개념입니다. 이 loss가 최대가 되도록 구성하여, generator가 획일된 이미지를 만드는 것이 아닌 주어진 style에 맞게 다양한 영상을 만들도록 학습하게 됩니다.
전체 loss로는 아래와 같이 adversarial, cycle-consistent, style reconstruction loss는 최소가되게하고, style diversification loss는 최대가 되도록 네트워크는 학습하게 됩니다.
Result
위 그림을 보면, 5가지의 source image들이 9개의 reference image의 contrast대로 변환이 잘되는 것을 확인할 수 있습니다. 이를 통해 하나의 규격화된 reference image만 있으면 모든 영상이 규격화된 contrast로 변환될 수 있음을 알 수 있습니다.
댓글