본문 바로가기

개발/딥러닝

[논문 구현] U-Net: Convolutional Networks for BiomedicalImage Segmentation

U-net은 정말 다양한 곳에 쓰입니다. 저는 U-net을 노이즈 제거 작업에 사용할 것입니다.

 

먼저 GPU를 설정해 줍니다.

 

전에는 MNIST 데이터를 사용했습니다. 이번에는 fashion_mnist 데이터를 사용하겠습니다.

 

 

노이즈 제거 작업을 한다고 했었죠. 이미지에 노이즈를 입혀줍니다.

 

노이즈가 잘 적용되었는지 확인해 봅시다.

 

 

원본 이미지랑 비교했을 때 조금 더 얼룩이 들어가 있는 것을 확인할 수 있습니다.

 

이제 모델 구현으로 가봅시다. 이미지의 크기가 작으니 깊이를 2단계만 내려가도록 하겠습니다.

먼저 CNN을 2번 적용하고 MaxPool을 적용합니다.

 

똑같이 CNN 2번에 MaxPool을 적용합니다. 다른 구현에는 함수를 사용해 중복을 줄였지만 깊이가 깊지 않으니 굳이 함수로 빼지는 않겠습니다.

 

이제 축소 부분이 끝났습니다. CNN을 2번 적용하고 이제부터 확장 단계로 넘어갑니다.

 

확장은 Conv2DTranspose를 이용해서 진행합니다. 픽셀들을 strides만큼 띄워준 뒤 CNN을 적용한다고 생각하면 됩니다. 이렇게 하면 크기가 맞지 않으니 Cropping2D를 이용해 일부 잘라 줍니다. Concatenate를 하고 CNN을 2번 적용해 줍니다.

strides = (2,2)

 

한번 더 반복해 줍니다.

 

이제 CNN을 2번 적용해서 최종 출력을 만들어 줍니다. 저희는 노이즈를 제거하는 문제를 풀기 때문에 출력 형태는 28*28*1이여야 합니다.

 

0에서 1사이 값을 예측하는 것이니 mse 로스를 사용합니다. noise 데이터에서 원본 데이터를 뽑아내도록 합니다..

 

 

val loss가 상당히 작게 잡힌 것을 확인할 수 있습니다. 그리고 이미지 역시 노이즈가 어느 정도 사라진 것을 확인할 수 있습니다.