[DL] MNIST 데이터셋으로 U-Net 구현하기: 이미지 노이즈 제거 실험

MNIST 데이터셋으로 U-Net 구현하기: 이미지 노이즈 제거 실험

🔍 TL;DR: MNIST 데이터셋을 사용하여 U-Net 모델을 구현해 이미지 노이즈 제거 실험을 진행했습니다. 28×28 크기에 맞게 모델 구조를 조정하고, 랜덤 노이즈와 고정 위치 노이즈 실험을 비교한 결과, 랜덤 노이즈로 학습한 모델이 일반화 성능이 더 뛰어났습니다.

소개

U-Net은 의료 이미지 세그멘테이션을 위해 개발되었지만, 이미지 복원, 노이즈 제거 등 다양한 분야에서도 뛰어난 성능을 보이는 모델입니다. 이번 실험에서는 MNIST 데이터셋을 사용하여 이미지 노이즈 제거를 위한 U-Net 모델을 구현해보았습니다.

1. 데이터셋 준비

MNIST 데이터셋은 0부터 9까지의 손글씨 숫자 이미지와 해당 레이블로 구성되어 있습니다. 하지만 U-Net은 이미지-이미지 변환 모델이므로, 데이터셋 구성을 다음과 같이 변경했습니다:

변경 전	변경 후	설명
x_train(입력 이미지)	노이즈가 추가된 이미지	노이즈 제거 학습을 위한 입력 데이터
y_train(레이블)	원본 x_train 이미지	노이즈 없는 깨끗한 이미지(타겟)
x_test(테스트 이미지)	노이즈가 추가된 테스트 이미지	테스트용 입력 데이터
y_test(테스트 레이블)	원본 x_test 이미지	테스트용 타겟 데이터

즉, 원래 MNIST의 이미지 데이터(x_train, x_test)를 깨끗한 타겟 이미지로 사용하고, 이 이미지에 인위적으로 노이즈를 추가한 데이터를 입력 이미지로 사용했습니다.

2. U-Net 모델 구조 설계

기존 U-Net은 고해상도 의료 이미지를 위해 설계되어 인코더(다운샘플링)와 디코더(업샘플링) 단계가 각각 5단계로 구성되어 있습니다. 하지만 MNIST 이미지는 28×28 크기로 작기 때문에, 모델 구조를 다음과 같이 조정했습니다:

 

원래 U-Net 다운샘플링 경로:

572×572 → 284×284 → 142×142 → 71×71 → 35×35 → 28×28

 

MNIST용 조정된 U-Net 경로:

28×28 → 14×14 → 7×7 → 14×14 → 28×28

 

구현된 U-Net 모델 코드

from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Conv2D, BatchNormalization, MaxPooling2D, Conv2DTranspose, concatenate

def unet(input_size=(28,28,1)):
    inputs = Input(input_size)
    depth = 14

    # 인코더 부분(다운샘플링)
    # 첫 번째 블록: 28×28×1 → 28×28×14
    conv1 = Conv2D(depth*1, (3, 3), activation='relu', padding='same')(inputs)
    conv1 = BatchNormalization()(conv1)
    conv1 = Conv2D(depth*1, (3, 3), activation='relu', padding='same')(conv1)
    conv1 = BatchNormalization()(conv1)
    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)  # 28×28×14 → 14×14×14

    # 두 번째 블록: 14×14×14 → 14×14×28
    conv2 = Conv2D(depth*2, (3, 3), activation='relu', padding='same')(pool1)
    conv2 = BatchNormalization()(conv2)
    conv2 = Conv2D(depth*2, (3, 3), activation='relu', padding='same')(conv2)
    conv2 = BatchNormalization()(conv2)
    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)  # 14×14×28 → 7×7×28

    # 병목 부분: 7×7×28 → 7×7×56
    conv3 = Conv2D(depth*4, (3, 3), activation='relu', padding='same')(pool2)
    conv3 = BatchNormalization()(conv3)
    conv3 = Conv2D(depth*4, (3, 3), activation='relu', padding='same')(conv3)
    conv3 = BatchNormalization()(conv3)

    # 디코더 부분(업샘플링)
    # 첫 번째 업샘플링 블록: 7×7×56 → 14×14×28
    up1 = concatenate([Conv2DTranspose(depth*2, (2, 2), strides=(2, 2), padding='same')(conv3), conv2], axis=3)
    dconv1 = Conv2D(depth*2, (3, 3), activation='relu', padding='same')(up1)
    dconv1 = BatchNormalization()(dconv1)
    dconv1 = Conv2D(depth*2, (3, 3), activation='relu', padding='same')(dconv1)
    dconv1 = BatchNormalization()(dconv1)

    # 두 번째 업샘플링 블록: 14×14×28 → 28×28×14
    up2 = concatenate([Conv2DTranspose(depth*1, (2, 2), strides=(2, 2), padding='same')(dconv1), conv1], axis=3)
    dconv2 = Conv2D(depth*1, (3, 3), activation='relu', padding='same')(up2)
    dconv2 = BatchNormalization()(dconv2)
    dconv2 = Conv2D(depth*1, (3, 3), activation='relu', padding='same')(dconv2)
    dconv2 = BatchNormalization()(dconv2)

    # 출력 레이어: 28×28×14 → 28×28×1
    outputs = Conv2D(1, (1, 1), activation='sigmoid')(dconv2)
    
    return Model(inputs=[inputs], outputs=[outputs])



# 모델 컴파일
model = unet()
model.compile(loss="binary_crossentropy", optimizer="adam", metrics=['accuracy'])
model.summary()

 

주요 특징:

• 깊이(depth)를 14로 설정하여 첫 레이어에서 14개 필터로 시작
• 다운샘플링 및 업샘플링 단계를 각각 2단계로 축소
• 스킵 커넥션(Skip Connection)을 통해 인코더와 디코더 간 특성 맵 연결 유지

3. 첫 번째 실험: 고정 위치 노이즈

첫 번째 실험에서는 모든 이미지의 동일한 위치에 노이즈를 추가했습니다.

실험 결과

학습 후 모델을 평가한 결과, 동일 위치의 노이즈는 잘 제거하지만 다른 위치의 노이즈에 대해서는 제거 성능이 떨어지는 문제가 있었습니다. 정확도도 예상보다 낮았습니다.

이는 모델이 특정 위치의 노이즈 패턴만 학습하고, 일반적인 노이즈 제거 능력을 기르지 못했기 때문입니다. 이런 접근법은 처음부터 의도한 것이 아니었습니다.

4. 두 번째 실험: 랜덤 노이즈

첫 번째 실험의 결과를 바탕으로, 두 번째 실험에서는 모든 학습 이미지에 랜덤 위치에 노이즈를 적용했습니다. 노이즈의 크기는 첫 번째 실험보다 작게 설정했지만, 이미지 전체에 랜덤하게 분포시켰습니다.

실험 결과

랜덤 노이즈로 학습한 모델은 다양한 위치의 노이즈를 더 효과적으로 제거할 수 있었습니다. 학습 결과가 첫 번째 실험보다 훨씬 만족스러웠으며, 모델의 일반화 성능이 향상되었습니다.

5. 최종 학습 설정 및 결과

최종적으로 선택한 학습 설정은 다음과 같습니다:

• 에포크(Epochs): 10
• 배치 크기(Batch Size): 3
• 학습 데이터: 1,000개 이미지 (학습:검증 = 7:3)
• 노이즈 유형: 랜덤 위치 노이즈

이 설정으로 학습한 결과, 모델은 다양한 패턴의 노이즈에 대해 효과적인 제거 성능을 보여주었습니다.

결론 및 교훈

1. 데이터 준비의 중요성: U-Net과 같은 이미지-이미지 변환 모델을 학습시키기 위해서는 데이터셋을 적절히 재구성해야 합니다.
   
   
2. 모델 구조 조정: 원래의 U-Net은 고해상도 이미지를 위해 설계되었으므로, MNIST와 같은 작은 이미지에 맞게 구조를 조정해야 합니다.
   
   
3. 데이터 다양성: 고정 위치 노이즈보다 랜덤 노이즈로 학습한 모델이 더 뛰어난 일반화 성능을 보여주었습니다. 이는 모델이 다양한 패턴을 학습하는 것의 중요성을 보여줍니다.
   
   
4. 하이퍼파라미터 튜닝: 적은 에포크와 작은 배치 크기로도 충분히 좋은 결과를 얻을 수 있었습니다. 항상 큰 모델과 많은 에포크가 필요한 것은 아닙니다.
   
   

U-Net은 세그멘테이션뿐만 아니라 이미지 복원, 노이즈 제거와 같은 작업에서도 뛰어난 성능을 보여줍니다. 이 실험을 통해 간단한 MNIST 데이터셋으로도 U-Net의 강력한 특성을 확인할 수 있었습니다.

전체 코드와 더 자세한 구현 내용은 GitHub 저장소에서 확인하실 수 있습니다.

---

이 글이 U-Net과 이미지 노이즈 제거에 관심 있는 분들께 도움이 되길 바랍니다. 질문이나 의견은 언제든지 댓글로 남겨주세요!