3D UNet으로 비장 Segmentation 시작하기

지난 글에서는 MONAI 전처리 파이프라인을 이용해
3D CT 데이터를 불러오고, 방향과 spacing을 맞춘 뒤,
RandCropByPosNegLabeld로 학습용 patch를 만드는 과정까지 정리했다.

마지막으로 확인했던 patch shape는 다음과 같았다.

patch image shape: torch.Size([4, 1, 96, 96, 96])
patch label shape: torch.Size([4, 1, 96, 96, 96])

이 shape는 다음 의미를 가진다.

• 4: 한 volume에서 뽑힌 patch 개수
• 1: 입력 채널 수
• 96 x 96 x 96: 각 patch의 3차원 크기

이제 이 patch를 실제 3D UNet에 넣어,
모델이 어떤 형태의 출력을 만드는지 확인해보자.

---

1. 3D UNet 모델 정의

이번 실습에서는 MONAI에서 제공하는 UNet을 사용했다.

model = UNet(
    spatial_dims = 3,                     
    in_channels = 1,                      
    out_channels = 2,                     
    channels = (16, 32, 64, 128, 256),    
    strides = (2, 2, 2, 2),               
    num_res_units = 2,                    
    norm = Norm.BATCH                     
).to(device)

핵심은 “3차원 입력을 받아 voxel 단위로 2개 클래스 점수를 예측하는 U-Net” 이라는 것이다.

1.2. 각 파라미터가 의미하는 것

• spatial_dims=3

  • 이 모델이 3차원 데이터를 처리한다는 뜻이다.
  • 즉, 2D U-Net이 아니라 3D U-Net이라는 의미다.

• in_channels=1

  • 입력 채널 수가 1이라는 뜻이다.
  • 이번 데이터는 RGB 컬러 이미지가 아니라 CT grayscale 영상이므로 채널은 1개만 있으면 된다.

• out_channels=2

  • 모델이 최종적으로 2개 채널을 출력한다는 뜻이다.
  • 왜 2개일까?
    • 이번 문제는 비장을 찾는 segmentation이므로 클래스가 2개다.
      • 0: 배경
      • 1: 비장
    • 즉, 모델은 각 voxel마다
      • “이 위치가 배경일 점수”와
      • “이 위치가 비장일 점수”를 각각 출력한다.
    • 중요한 점은, 모델이 처음부터 0과 1을 바로 뱉는 게 아니라
      각 클래스에 대한 점수(logit) 를 출력한다는 것이다.
      즉, 각 voxel마다 이런 식의 출력이 생긴다고 이해하면 된다.
      [배경 점수, 비장 점수]
      그리고 나중에 후처리 과정에서 이 둘 중 더 큰 점수를 가진 클래스를 선택해 최종 mask를 만든다.

• channels=(16, 32, 64, 128, 256)

  • 이 값은 U-Net 각 단계에서 사용할 feature map 수를 정한다.
  • U-Net은 입력을 점점 압축하면서 더 복잡한 특징을 뽑아내고,
    다시 복원하면서 위치 정보를 살려 segmentation 결과를 만든다.
  • 여기서 channels=(16, 32, 64, 128, 256)라는 것은
    깊이 들어갈수록 feature map 수를 늘리겠다는 뜻이다.
  • 즉, 모델의 흐름은 대략 이런 느낌이다.
    • 처음엔 단순한 특징 몇 개만 본다
    • 더 깊은 단계로 갈수록 더 많은 특징을 본다
    • 다시 올라오면서 위치 정보를 복원한다
  • 비유하면,
    • 처음에는 “밝기나 경계 같은 단순 정보”를 보고,
    • 깊어질수록 “장기처럼 보이는 패턴”을 더 복잡하게 이해한다고 볼 수 있다.

• strides=(2, 2, 2, 2)

  • 이 값은 각 downsampling 단계에서 공간 크기를 얼마나 줄일지를 의미한다.
  • 여기서는 각 단계마다 2를 사용했으므로, 공간 크기를 절반씩 줄여가며 특징을 압축한다.
  • 예를 들어 입력이 96 x 96 x 96이라면 대략적으로:
    • 96 → 48
    • 48 → 24
    • 24 → 12
    • 12 → 6
      이런 식으로 줄어들 수 있다.
  • 왜 이렇게 줄일까?
    • 전체 구조를 점점 압축하면서 작은 지역 정보뿐 아니라 더 넓은 문맥 정보까지 보기 위해서다.
    • 즉, 모델이 단순히 “여기 픽셀이 밝네” 수준이 아니라
      “이 주변 구조 전체를 보니 비장처럼 보인다”까지 판단할 수 있게 된다.

• num_res_units=2

  • 이 값은 각 단계 안에 Residual Unit을 몇 개 사용할지를 의미한다.
  • Residual 구조는 입력 정보를 완전히 끊지 않고 뒤쪽 레이어로 일부 전달해주는 방식이다.
  • 쉽게 말하면, 모델이 깊어지더라도 기존 정보를 너무 잃지 않도록 도와주는 장치다.
  • 비유하면,
    • 중간 정리본을 매 단계마다 버리는 것이 아니라
    • 핵심 메모를 옆에 붙여가며 계속 참고하는 것에 가깝다.
      → 그래서 학습이 더 안정적이고, 깊은 네트워크에서도 정보 손실을 줄이는 데 도움이 된다.

• norm=Norm.BATCH

  • 이 값은 모델 내부에서 사용할 정규화 방식을 의미한다.
  • 여기서는 Batch Normalization을 사용한다.
    • Batch Normalization은 각 레이어의 출력 분포를 어느 정도 안정화해서
      학습이 더 잘 진행되도록 도와준다.
    • 아주 단순하게 말하면, 각 단계의 값이 너무 들쑥날쑥하지 않도록 정리해주는 역할이다.

---

2.모델 입력과 출력 확인하기

모델을 정의한 뒤에는 실제 patch를 넣어서 입력 shape와 출력 shape를 확인했다.

sample = next(iter(patch_loader))
x = sample["image"].to(device)

with torch.no_grad():
    y = model(x)

print("입력 shape: ", x.shape)
print("출력 shape: ", y.shape)

이 코드는 학습을 하는 것이 아니라,
모델에 데이터가 정상적으로 들어가고 나오는지 확인하는 테스트에 가깝다.

torch.no_grad()를 쓴 이유는
지금은 gradient 계산이 필요 없는 단순 forward 확인 단계이기 때문이다.
결과는 아래와 같다.

입력 shape:  torch.Size([4, 1, 96, 96, 96])    # (patch 개수, 입력 채널, 높이, 너비, 깊이)
출력 shape:  torch.Size([4, 2, 96, 96, 96])    # (patch 개수, 클래스 수, 높이, 너비, 깊이)

• RandCropByPosNegLabeld에서 num_samples=4로 설정했기 때문에
  입력 첫번째 차원이 4로 나타나고,
• 출력 shape을 보면 출력 채널이 2로 바뀐 것을 확인할 수 있다.

---

3. 학습을 위한 핵심 도구 준비하기

이제 모델이 3D patch를 입력받아 (배경 점수, 비장 점수) 형태의 출력을 만든다는 것까지 확인했다.
그 다음 단계는 아주 자연스럽다.

• 모델이 얼마나 틀렸는지 계산할 기준을 정하고,
• 그 틀린 정도를 바탕으로 어떻게 수정할지 정하고,
• 검증 단계에서는 성능을 어떤 점수로 평가할지 정해야 한다.

이번 실습에서는 이를 위해 MONAI에서 제공하는 몇 가지 핵심 도구를 사용했다.

from monai.losses import DiceLoss
from monai.metrics import DiceMetric
from monai.inferers import sliding_window_inference
from monai.utils import set_determinism

각 도구의 역할은 다음과 같다.

• DiceLoss : 모델이 얼마나 틀렸는지 계산하는 기준
• DiceMetric : validation에서 모델 성능을 숫자로 평가하는 기준
• sliding_window_inference : 큰 3D 볼륨을 작은 창으로 나누어 추론하는 도구
• set_determinism : 실험 재현성을 높이기 위해 랜덤성을 가능한 범위에서 고정하는 도구

3.1. 손실함수: DiceLoss

Segmentation에서는 단순히 픽셀 개수를 맞추는 것보다,
예측 영역과 정답 영역이 얼마나 잘 겹치는지가 훨씬 중요하다.

특히 의료영상에서는 배경이 대부분이고, 실제로 찾고 싶은 장기나 병변은 상대적으로 작을 수 있다.
이런 상황에서 일반적인 분류 손실만 쓰면 배경을 많이 맞히는 쪽으로 학습이 치우칠 수 있다.

그래서 이번 실습에서는 Dice 기반 손실함수를 사용했다.

loss_function = DiceLoss(to_onehot_y=True, softmax=True)

이 코드를 해석하면 다음과 같다.

• DiceLoss를 손실함수로 사용한다.
• to_onehot_y=True : 정답 라벨을 one-hot 형식으로 변환한다.
• softmax=True : 모델 출력에 softmax를 적용해 클래스별 확률처럼 해석 가능한 형태로 바꾼다.

여기서 one-hot은 예를 들어 클래스가 2개일 때,
각 voxel의 정답을 [1, 0] 또는 [0, 1]처럼 표현하는 방식이다.

즉, 현재 문제에서는:

• 배경이면 [1, 0]
• 비장이면 [0, 1]
  같은 형태로 바뀐다고 이해하면 된다.

softmax를 함께 적용하는 이유는 모델의 출력이 아직은 단순 점수(logit)이기 때문이다.
DiceLoss가 이를 클래스별 확률처럼 비교할 수 있게 정리한 뒤 손실을 계산하게 된다.

🧐 여기서 든 의문:

손실함수 안에 softmax를 넣는게 연산 안정화 때문인데, 그럼 logit으로 먼저 손실을 계산해서 gradient를 얻고 softmax는 나중에 처리하는 순서가 맞지 않나?

결론: CrossEntropyLoss 쪽에서는 맞고(수식에 로그가 들어가기 때문), MONAI DiceLoss는 성격이 다르다.

MONAI 문서에 따르면 DiceLoss의 입력 채널은 logits 또는 probabilities 모두 가능하고, logits를 넣는다면 softmax=True나 sigmoid=True를 켜라고 되어 있다.
즉, 선택 옵션을 준 것.

결과적으로 DiceLoss에 대해 요약하자면:

DiceLoss는 예측 mask와 정답 mask가 얼마나 잘 겹치는지를 기준으로 오차를 계산하는 손실함수다.

3.2. Optimizer: Adam

손실함수를 정했다면, 이제 모델이 틀렸을 때 어떤 방식으로 가중치를 수정할지도 정해야 한다.

이번 실습에서는 가장 많이 쓰이는 optimizer 중 하나인 Adam을 사용했다.

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

Adam이 어떤식으로 가중치를 조정하는지 다시 정리하자면:

Adam은 “방향 기억 + 흔들림 보정”을 같이 써서, 파라미터마다 업데이트 크기를 다르게 조절하는 optimizer다.

1) 최근 gradient의 평균적인 방향.
이건 “지금까지 보니 어느 방향으로 계속 내려가는 경향이 있었는가?”를 보는 역할이다.
즉, 순간적인 흔들림보다 전체적인 이동 방향을 더 믿게 만든다.
Adam 논문에서는 이를 1차 모멘트 추정치(first moment)라고 설명한다.

2) gradient가 얼마나 크게, 불안정하게 흔들리는가?
이건 “이 파라미터는 기울기가 자주 크고 요동치는가, 아니면 비교적 안정적인가?”를 보는 역할이다.
Adam은 gradient의 제곱값을 누적해 이런 흔들림의 크기를 추정하고,
흔들림이 큰 파라미터는 업데이트 폭을 줄이고, 흔들림이 작고 안정적인 파라미터는 상대적으로 더 움직일 수 있게 만든다.
이게 2차 모멘트 추정치(second moment)다.

그래서 Adam의 업데이트는 직관적으로 이렇게 이해하면 된다.

• 최근에 꾸준히 내려가던 방향은 살리고,
• 요동이 큰 축은 조심해서 조금만 움직이고,
• 비교적 안정적인 축은 더 자신 있게 움직인다.

수식으로 보면 Adam은 대략 이런 형태로 업데이트한다.

• $g_t$ : 현재 gradient
• $m_t$ : gradient의 이동평균, 즉 방향 기억
• $v_t$ : gradient 제곱의 이동평균, 즉 흔들림 크기 기억
• $\hat{m}_t, \hat{v}_t$ : 초반 편향을 보정한 값
• $\alpha$ : learning rate
• $\epsilon$ : 0으로 나누는 상황을 막는 작은 값

여기서 특히 중요한 건 bias correction이다.
처음 몇 step에서는 이동평균이 아직 충분히 쌓이지 않아서 $𝑚_𝑡$, $𝑣_𝑡$가 작게 잡히는 경향이 있다.
Adam은 이 초기 편향을 보정한 $\hat{m}_t, \hat{v}_t$를 사용해서 초반 업데이트가 지나치게 왜곡되지 않도록 한다.

비유로 설명하면 더 쉽다.

경사길을 내려간다고 해보자.

• SGD는 지금 발밑 경사만 보고 한 걸음 움직이는 방식에 가깝다.
• Adam은
  • 지금까지 어느 방향으로 계속 내리막이었는지 기억하고,
  • 어느 방향은 길이 울퉁불퉁해서 자꾸 흔들렸는지도 기억한 뒤,
  • 더 안정적인 방향으로 조절해서 걷는 방식에 가깝다.

그래서 Adam은 noisy한 gradient, sparse한 gradient, 스케일이 서로 다른 파라미터들 이 섞여 있는 상황에서 실무적으로 많이 쓰인다.

3.3.validation 성능 확인용 점수: DiceMetric

dice_metric = DiceMetric(include_background=False, reduction="mean")

이 설정의 의미는 다음과 같다.

• include_background=False : 배경 클래스는 평가에서 제외
• reduction="mean" : 여러 샘플의 Dice 점수를 평균으로 계산

왜 배경을 제외할까?

의료영상 segmentation에서는 배경이 워낙 많기 때문에, 배경까지 포함하면 점수가 지나치게 높게 나올 수 있다.
우리가 진짜 알고 싶은 것은 비장 영역을 얼마나 잘 맞췄는가이기 때문에
배경은 제외하는 것이 더 적절하다.

즉, 이 DiceMetric은 비장 클래스만 기준으로 성능을 평균 평가하는 점수라고 보면 된다.

예측값과 정답값을 비교 가능한 형태로 바꾸기

모델 출력은 바로 비교하기에는 아직 정리되지 않은 상태다.
앞에서 본 것처럼 모델 출력은 (배경 점수, 비장 점수) 형태의 연속값이다.

그래서 validation 단계에서는
예측값과 정답값을 Dice 계산이 쉬운 형태로 바꾸는 과정이 필요하다.

from monai.transforms import AsDiscrete

post_pred = AsDiscrete(argmax=True, to_onehot=2)
post_label = AsDiscrete(to_onehot=2)

• post_pred: 이건 모델 예측값을 최종 비교용 mask로 바꾸는 도구다.

  • argmax=True: 각 voxel마다 가장 큰 점수를 가진 클래스를 선택
  • to_onehot=2: 선택된 클래스를 다시 2채널 one-hot 형태로 변환
  • 즉, 모델이 출력한 점수맵을
    “이 voxel은 배경인가 비장인가?”라는 딱 떨어지는 class mask 형태로 바꾸는 과정이다.

• post_label: 이건 정답 라벨도 같은 기준으로 맞추기 위해 2채널 one-hot 형태로 바꾸는 도구다.

결국 validation에서는

• 예측값도 정리하고
• 정답값도 정리해서

같은 형식으로 Dice를 계산하게 된다.

---

4. 학습 루프 짜기

4.1. 한 epoch 동안 모델을 학습하는 함수

이제 실제 학습 루프를 보자.

def train_one_epoch(model, loader, loss_function, optimizer, device):
    model.train()
    epoch_loss = 0.0
    step = 0

    for batch_data in tqdm(loader, desc="Training"):
        step += 1

        inputs = batch_data["image"].to(device)
        labels = batch_data["label"].to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()

        epoch_loss += loss.item()

    epoch_loss /= max(step, 1)
    return epoch_loss

• model.train()

  • 모델을 학습 모드로 바꾼다.
  • Dropout이나 BatchNorm 같은 레이어가 학습용 방식으로 동작하게 된다.

• 배치마다 입력과 정답을 가져온다

inputs = batch_data["image"].to(device)
labels = batch_data["label"].to(device)

• 여기서 device는 보통 GPU이므로,

• 계산을 위해 데이터를 GPU로 올리는 단계다.

• 이전 gradient를 초기화한다(optimizer.zero_grad())

• 모델 예측을 만든다(outputs = model(inputs))

• loss를 계산한다(loss = loss_function(outputs, labels))

  • 위에서 정의한 손실함수(loss_function = DiceLoss(...))가 여기서 사용됨.

• 역전파를 수행한다(loss.backward())

  • 여기서 손실값을 기준으로 각 가중치를 얼마나 수정해야 하는지 gradient가 계산된다.

• 실제로 가중치를 업데이트한다(optimizer.step())

• epoch 평균 loss를 계산한다

  • 각 step의 loss를 누적한 뒤(epoch_loss += loss.item()),
  • 마지막에 평균을 내서(epoch_loss /= max(step, 1)) 반환한다.

4.2. 한 epoch 동안 validation을 수행하는 함수

validation을 수행하는 함수를 정의해주기 전에 선행되어야할 작업이 있다.

바로 inferer 함수를 만드는 것이다.

roi_size = (160, 160, 160)
sw_batch_size = 1

def inferer(input, model):
    return sliding_window_inference(
        inputs=input,
        roi_size=roi_size,
        sw_batch_size=sw_batch_size,
        predictor=model,
    )

🧐 왜 validation에서는 sliding window inference를 쓸까?

학습할 때는 patch 단위로 입력을 넣었지만,
validation에서는 보통 큰 3D 볼륨 전체를 대상으로 예측하고 싶다.

문제는 3D 볼륨 전체를 한 번에 모델에 넣으면 메모리가 너무 많이 든다는 점이다.

그래서 MONAI에서는 sliding_window_inference를 자주 사용한다.

위에서 정의한 이 'inferer 함수'는 큰 영상을 한 번에 넣는 대신,
작은 3D 창(window)으로 잘라서 예측한 후 다시 합쳐준다.

쉽게 말하면:

• 큰 CT 전체를 한 번에 보지 않고
• 작은 조각으로 나눠서 본 뒤
• 마지막에 이어 붙이는 방식
  이다.

이 코드를 한 줄씩 해석하면 이렇다.

• def inferer(input, model):: sliding window 추론을 매번 길게 쓰지 않기 위해, 별도 함수로 감싼 것.

  • validation 코드에서는 그냥 inferer(val_inputs, model)처럼 간단히 호출할 수 있게 만든 래퍼(wrapper) 함수라고 보면 된다.

• inputs=input: 추론할 실제 입력 볼륨.

  • 입력은 channel-first, batch dimension 포함 형태를 가정한다.
    즉 3D라면 보통 (B, C, H, W, D) 형태다.

• roi_size = (160, 160, 160): 한 번에 잘라서 볼 3D 조각의 크기.

  • validation할 때 전체 CT를 통째로 넣는 대신,
    160×160×160 크기의 작은 큐브 단위로 나눠 보겠다는 뜻.

• sw_batch_size = 1: 잘라낸 window를 한 번에 몇 개씩 모델에 넣을지를 정하는 값.

  • 여기서는 1이므로 window를 하나씩 차례대로 넣는다.
  • 3D 데이터는 메모리 사용량이 크기 때문에, 안정적으로 추론하려고 이렇게 작게 잡는 경우가 많다.

• predictor=model: 각 window 조각을 실제로 예측할 주체가 바로 model이라는 뜻.

  • 즉, 잘라낸 patch를 받아 segmentation 결과를 내는 함수가 현재 3D UNet 모델이라는 의미다.

그럼 이제 본격적으로 한 epoch 동안 validation을 수행하는 함수를 만들어보자.

def validate_one_epoch(model, loader, device, inferer):
    model.eval()
    dice_metric.reset()

    with torch.no_grad():
        for val_data in tqdm(loader, desc="Validation"):
            val_inputs = val_data["image"].to(device)
            val_labels = val_data["label"].to(device)

            val_outputs = inferer(val_inputs, model)

            val_outputs_list = [post_pred(i) for i in decollate_batch(val_outputs)]
            val_labels_list = [post_label(i) for i in decollate_batch(val_labels)]

            dice_metric(y_pred=val_outputs_list, y=val_labels_list)

    metric = dice_metric.aggregate().item()
    return metric

여기서 중요한 포인트는 학습 함수와 validation 함수의 차이다.

• model.eval(): 모델을 평가 모드로 전환한다.

• dice_metric.reset(): DiceMetric은 내부적으로 결과를 누적할 수 있기 때문에,
  새 validation을 시작하기 전에 반드시 초기화해야 한다.

• torch.no_grad(): 검증 단계에서는 가중치를 업데이트하지 않으므로 gradient 계산이 필요 없다.

• val_outputs = inferer(val_inputs, model):

  • validation에서는 전체 volume을 대상으로 sliding window inference를 수행한다.

• decollate_batch

이 부분이 처음에는 꽤 헷갈린다.

val_outputs_list = [post_pred(i) for i in decollate_batch(val_outputs)]
val_labels_list = [post_label(i) for i in decollate_batch(val_labels)]

val_outputs와 val_labels는 batch 형태로 묶여 있다.
그런데 DiceMetric은 샘플 단위로 비교하는 것이 더 자연스럽기 때문에,
decollate_batch를 써서 batch를 다시 개별 샘플 리스트로 나눈다.

즉,

• batch 형태의 텐서를 샘플별 텐서 목록으로 쪼개고

• 각 샘플에 post_pred(모델 예측값을 최종 비교용 mask로 바꾸는 도구),
  post_label(정답 라벨도 같은 기준으로 맞추기 위해 2채널 one-hot 형태로 바꾸는 도구)을 적용한 뒤

• DiceMetric에 넣는(dice_metric(y_pred=val_outputs_list, y=val_labels_list)) 구조.

• dice_metric.aggregate().item():

  • validation 동안 누적된 Dice 결과를 최종 평균 점수로 합치고,
    그 값을 일반 숫자(float)로 꺼낸다.

4.3. 전체 학습 루프 구성하기

이제 학습과 검증 함수를 만들었으니, epoch 단위 반복만 붙이면 전체 학습 루프가 완성된다.

best_metric = -1
best_metric_epoch = -1
best_model_path = "/kaggle/working/best_metric_model.pth"

epoch_loss_values = []
metric_values = []

max_epochs = 5

for epoch in range(max_epochs):
    print("\n" + "=" * 60)
    print(f"Epoch {epoch + 1}/{max_epochs}")
    print("=" * 60)

    train_loss = train_one_epoch(model, train_loader, loss_function, optimizer, device)
    epoch_loss_values.append(train_loss)

    val_dice = validate_one_epoch(model, val_loader, device, inferer)
    metric_values.append(val_dice)

    print(f"Epoch {epoch + 1}: loss={train_loss:.4f}, val_dice={val_dice:.4f}")

    if val_dice > best_metric:
        best_metric = val_dice
        best_metric_epoch = epoch + 1
        torch.save(model.state_dict(), best_model_path)
        print("최고 성능 모델 저장 완료!")

이 흐름은 아래처럼 요약할 수 있다.

1) 한 epoch 에 대해 학습
2) 한 epoch 에 대한 검증
3) loss와 Dice를 기록
4) 지금까지 중 최고 성능이면 모델 저장
5) 이 과정을 반복(정해진 epoch 수 만큼)

여기서 torch.save(model.state_dict(), best_model_path)는
모델 전체를 저장하는 것이 아니라 학습된 가중치 값(state_dict)을 저장하는 방식이다.

즉, 가장 좋은 validation Dice를 기록한 시점의 가중치만 따로 저장해두는 것이다.

---

5. 학습 곡선 확인하기

학습이 끝난 뒤에는 loss와 validation Dice를 그래프로 그려 훈련이 제대로 진행되었는지 확인했다.

plt.figure(figsize=(12, 5))

plt.subplot(1, 2, 1)
plt.title("Epoch Average Loss")
plt.plot(range(1, len(epoch_loss_values) + 1), epoch_loss_values, marker="o")
plt.xlabel("Epoch")
plt.ylabel("Loss")

plt.subplot(1, 2, 2)
plt.title("Validation Dice")
plt.plot(range(1, len(metric_values) + 1), metric_values, marker="o")
plt.xlabel("Epoch")
plt.ylabel("Dice")

plt.tight_layout()
plt.show()

이 그래프는 학습 상태를 해석하는 중요한 단서가 된다.

예를 들어:

• loss가 줄어드는지
• validation Dice가 올라가는지
• 어느 시점에서 가장 성능이 좋았는지
  같은 흐름을 눈으로 확인할 수 있다.

---

6. 최고 성능 모델을 다시 불러와 예측해보기

학습이 끝났다고 바로 마지막 epoch 모델을 쓰는 건 아니다.
우리는 validation Dice 기준으로 가장 좋았던 모델을 저장해뒀기 때문에,
그 모델을 다시 불러와야 한다.

if not os.path.exists(best_model_path):
    raise FileNotFoundError(f"저장된 모델 파일이 없습니다: {best_model_path}")

model.load_state_dict(torch.load(best_model_path, map_location=device))
model.eval()

이제부터는 “가장 잘 학습된 시점의 모델”로 validation 예측을 확인하게 된다.

---

7. 예측 결과를 시각화하기

이제 validation 예시 하나를 꺼내 실제 예측 mask를 확인한다.

with torch.no_grad():
    val_example = next(iter(val_loader))

    val_inputs = val_example["image"].to(device)
    val_labels = val_example["label"].to(device)

    val_outputs = inferer(val_inputs, model)
    pred_mask = torch.argmax(val_outputs, dim=1)

    image_np = val_inputs[0, 0].cpu().numpy()
    label_np = val_labels[0, 0].cpu().numpy()
    pred_np = pred_mask[0].cpu().numpy()

여기서 torch.argmax(val_outputs, dim=1)은
각 voxel마다 두 클래스 중 더 큰 점수를 가진 클래스를 선택하는 과정이다.
즉, 점수맵을 실제 segmentation mask로 바꾸는 단계다.

정답 라벨이 가장 많이 포함된 slice를 골라 시각화하기

slice_sums = [label_np[:, :, i].sum() for i in range(label_np.shape[2])]
best_slice = int(np.argmax(slice_sums))

이 의미는 단순하다.

• 각 z-slice마다 비장 픽셀 수를 세고, 가장 많이 포함된 slice를 고른다

즉, 비장이 가장 잘 보이는 단면을 골라 시각화에 사용한 것이다.

이 방식은 굉장히 직관적이다.
왜냐하면 아무 slice나 고르면 비장이 거의 안 보일 수도 있기 때문이다.

원본 / 정답 / 예측 / Overlay 비교

이제 그 slice를 기준으로 원본 이미지, 정답 라벨, 모델 예측을 나란히 본다.

plt.subplot(1, 3, 1)
plt.title("Image")
plt.imshow(image_np[:, :, best_slice], cmap="gray")

plt.subplot(1, 3, 2)
plt.title("Label")
plt.imshow(label_np[:, :, best_slice], cmap="gray")

plt.subplot(1, 3, 3)
plt.title("Prediction")
plt.imshow(pred_np[:, :, best_slice], cmap="gray")

# overlay
plt.imshow(image_np[:, :, best_slice], cmap="gray")
plt.imshow(pred_np[:, :, best_slice], cmap="Reds", alpha=0.35)

이렇게 하면 예측 위치가 장기와 얼마나 잘 맞는지 가장 직관적으로 확인할 수 있다.

---

8. validation 샘플 전체에서 best case / worst case 찾기

이제는 더 자세히 모델의 성능을 분석하기 위해 validation 전체를 돌면서 각 샘플의 Dice를 계산하고,
가장 잘 맞춘 경우와 가장 못 맞춘 경우를 비교하고 차이를 분석하여 원인을 생각해
좀 더 높은 성능을 얻기 위한 인사이트를 얻어야 한다.

이를 위해 먼저 binary Dice 함수를 만들었다.

def binary_dice(pred_np, label_np, target_class=1):
    pred_bin = (pred_np == target_class)
    label_bin = (label_np == target_class)

    intersection = np.sum(pred_bin & label_bin)
    pred_sum = np.sum(pred_bin)
    label_sum = np.sum(label_bin)

    if pred_sum + label_sum == 0:
        return 1.0
    return 2.0 * intersection / (pred_sum + label_sum)

이 함수는 클래스 1, 즉 '비장 영역만 기준'으로 예측과 정답이 얼마나 잘 겹치는지를 직접 계산한다.

이렇게 각 validation 샘플의 Dice를 구한 뒤 정렬하면:

• 가장 낮은 Dice 샘플
• 가장 높은 Dice 샘플
  을 찾을 수 있다.

실 사용

case_results = []

model.eval()

with torch.no_grad():
    for idx, val_data in enumerate(val_loader):
        val_inputs = val_data["image"].to(device)
        val_labels = val_data["label"].to(device)

        val_outputs = inferer(val_inputs, model)
        pred_mask = torch.argmax(val_outputs, dim=1)

        image_np = val_inputs[0, 0].cpu().numpy()
        # batch 축과 channel 축을 제거하고 cpu로 옮긴 뒤 numpy 배열로 변경
        # 원본 CT를 시각화하기 쉬운 배열 형태로 바꾸는 코드
        
        label_np = val_labels[0, 0].cpu().numpy()
        pred_np = pred_mask[0].cpu().numpy()

        dice = binary_dice(pred_np, label_np, target_class=1)

        case_results.append({
            "idx": idx,
            "dice": dice,
            "image": image_np,
            "label": label_np,
            "pred": pred_np,
        })

case_results = sorted(case_results, key=lambda x: x["dice"])

worst_case = case_results[0]
best_case = case_results[-1]

print("가장 낮은 Dice 샘플 index:", worst_case["idx"], "| Dice:", worst_case["dice"])
print("가장 높은 Dice 샘플 index:", best_case["idx"], "| Dice:", best_case["dice"])

시각화

def get_best_slice_from_label(label_np):
    slice_sums = [label_np[:, :, i].sum() for i in range(label_np.shape[2])]
    return int(np.argmax(slice_sums))
def plot_best_vs_worst_case(worst_case, best_case):
    worst_slice = get_best_slice_from_label(worst_case["label"])
    best_slice = get_best_slice_from_label(best_case["label"])
    fig, axes = plt.subplots(2, 4, figsize=(20, 10))
    # 상단: worst
    axes[0, 0].imshow(worst_case["image"][:, :, worst_slice], cmap="gray")
    axes[0, 0].set_title(f"Worst - Original\nDice={worst_case['dice']:.4f}")
    axes[0, 0].axis("off")
    axes[0, 1].imshow(worst_case["label"][:, :, worst_slice], cmap="gray")
    axes[0, 1].set_title("Worst - Label")
    axes[0, 1].axis("off")
    axes[0, 2].imshow(worst_case["pred"][:, :, worst_slice], cmap="gray")
    axes[0, 2].set_title("Worst - Prediction")
    axes[0, 2].axis("off")
    axes[0, 3].imshow(worst_case["image"][:, :, worst_slice], cmap="gray")
    axes[0, 3].imshow(worst_case["pred"][:, :, worst_slice], cmap="Reds", alpha=0.35)
    axes[0, 3].set_title("Worst - Overlay")
    axes[0, 3].axis("off")
    # 하단: best
    axes[1, 0].imshow(best_case["image"][:, :, best_slice], cmap="gray")
    axes[1, 0].set_title(f"Best - Original\nDice={best_case['dice']:.4f}")
    axes[1, 0].axis("off")
    axes[1, 1].imshow(best_case["label"][:, :, best_slice], cmap="gray")
    axes[1, 1].set_title("Best - Label")
    axes[1, 1].axis("off")
    axes[1, 2].imshow(best_case["pred"][:, :, best_slice], cmap="gray")
    axes[1, 2].set_title("Best - Prediction")
    axes[1, 2].axis("off")
    axes[1, 3].imshow(best_case["image"][:, :, best_slice], cmap="gray")
    axes[1, 3].imshow(best_case["pred"][:, :, best_slice], cmap="Reds", alpha=0.35)
    axes[1, 3].set_title("Best - Overlay")
    axes[1, 3].axis("off")
    plt.tight_layout()
    plt.show()
plot_best_vs_worst_case(worst_case, best_case)

---

9. patch 단위 성능 분석

이번 실습에서는 volume 전체뿐 아니라 patch 단위로도 Dice를 확인했다.

왜냐하면 학습 자체는 patch 단위로 진행되기 때문에,
모델이 전체 볼륨에서는 놓치는 부분이 있어도 개별 patch에서는 꽤 잘 맞추고 있을 수 있기 때문이다.

즉,

• 전체 volume 성능이 낮은 이유가
• 모델이 patch 수준에서도 못 맞추는 건지
• 아니면 전체 조합 과정에서 문제가 생기는 건지
  를 간접적으로 파악할 수 있다.

코드

def normalize_patch_batch(image_tensor, label_tensor):
    if image_tensor.ndim == 6 and image_tensor.shape[0] == 1:
        image_tensor = image_tensor.squeeze(0)
    if label_tensor.ndim == 6 and label_tensor.shape[0] == 1:
        label_tensor = label_tensor.squeeze(0)
    return image_tensor, label_tensor

patch_batch_results = []

model.eval()

with torch.no_grad():
    for batch_idx, batch_data in enumerate(patch_loader):
        inputs = batch_data["image"].to(device)
        labels = batch_data["label"].to(device)

        inputs, labels = normalize_patch_batch(inputs, labels)

        outputs = model(inputs)
        pred_mask = torch.argmax(outputs, dim=1)

        for patch_idx in range(inputs.shape[0]):
            image_np = inputs[patch_idx, 0].cpu().numpy()
            label_np = labels[patch_idx, 0].cpu().numpy()
            pred_np = pred_mask[patch_idx].cpu().numpy()

            dice = binary_dice(pred_np, label_np, target_class=1)

            patch_batch_results.append({
                "batch_idx": batch_idx,
                "patch_idx": patch_idx,
                "dice": dice,
                "image": image_np,
                "label": label_np,
                "pred": pred_np,
            })

patch_batch_results = sorted(patch_batch_results, key=lambda x: x["dice"])

print("가장 낮은 patch Dice:", patch_batch_results[0]["dice"])
print("가장 높은 patch Dice:", patch_batch_results[-1]["dice"])
print("worst patch batch_idx / patch_idx:", patch_batch_results[0]["batch_idx"], patch_batch_results[0]["patch_idx"])
print("best patch batch_idx / patch_idx:", patch_batch_results[-1]["batch_idx"], patch_batch_results[-1]["patch_idx"])

추가로 고려해보면 좋을 성능 지표

1) Tversky index

Tversky index는 Dice를 조금 더 유연하게 만든 지표다.

• Dice는 false positive와 false negative를 비교적 대칭적으로 다루는 편이다.
• 반면 Tversky는 두 오류에 서로 다른 패널티를 줄 수 있다.
• “조금 더 많이 잡더라도 놓치면 안 되는 문제”에서는 false negative 쪽 패널티를 더 크게 둘 수 있다.
• 반대로 “괜한 오탐이 더 위험한 문제”라면 false positive 쪽 패널티를 더 크게 둘 수 있다.

수식은 아래와 같다.

여기서 의미는 다음과 같다.

• TP: 정답도 비장이고 예측도 비장인 voxel 수
• FP: 비장이 아닌데 비장이라고 예측한 수
• FN: 비장인데 놓친 수

또한,

• $\beta$를 더 크게 잡을수록 놓친 경우(FN)에 더 민감해져서 recall 쪽을 더 중시하게 된다.
• $\alpha = \beta = 0.5$이면 Dice/F1과 같은 계열로 연결된다.

2) F-beta score

F-beta score는 precision과 recall의 균형점을 직접 조절하는 지표다.

• beta > 1이면 recall을 더 중요하게 본다.
• beta < 1이면 precision을 더 중요하게 본다.
• 예를 들어 beta = 2인 F2 score는 recall 비중을 더 키운 경우다.
• beta = 0.5인 F0.5 score는 precision 비중을 더 키운 경우다.

수식은 아래와 같다.

해석은 다음과 같다.

• beta가 커질수록 FN 쪽 영향이 더 커지므로, 결과적으로 놓치는 것에 더 민감한 점수가 된다.
• 반대로 beta가 작아지면 불필요한 양성 예측(FP)을 줄이는 쪽, 즉 precision 중심 해석에 더 가까워진다.

---

10. 가장 큰 connected component만 남기는 후처리

마지막으로, 예측 결과를 조금 더 깔끔하게 만들기 위한 후처리도 시도했다.

여기서 사용한 아이디어는 단순하다.

• 비장은 보통 하나의 큰 연결된 덩어리로 나타날 가능성이 높다.
• 그렇다면 모델이 작은 잡음처럼 띄엄띄엄 예측한 부분은 지우고,
• 가장 큰 덩어리만 남기면 결과가 좋아질 수 있지 않을까?

이 아이디어를 코드로 구현한 것이 'keep_largest_component 함수'다.

from scipy import ndimage

def keep_largest_component(binary_mask):
    labeled, num_features = ndimage.label(binary_mask)

    if num_features == 0:
        return binary_mask

    sizes = ndimage.sum(binary_mask, labeled, range(1, num_features + 1))
    largest_label = np.argmax(sizes) + 1
    largest_component = (labeled == largest_label).astype(np.uint8)
    return largest_component

이 함수의 흐름은 다음과 같다.

1) 예측 mask 안의 연결된 덩어리들을 찾는다.
2) 각 덩어리 크기를 계산한다.
3) 가장 큰 덩어리 하나만 남긴다.
4) 나머지는 모두 배경으로 처리한다.

즉, 모델이 비장 외의 곳을 비장으로 잘못 예측했다면 그런 잡음을 줄이는 후처리라고 볼 수 있다.

후처리 전/후 Dice 비교

worst_pred_bin = (worst_case["pred"] == 1).astype(np.uint8)
# worst_case["pred"]는 가장 못 맞춘 샘플의 예측 결과.
# 그 중에서 클래스 1. 즉, 비장으로 예측한 부분만 1로 만든다.

worst_pred_post = keep_largest_component(worst_pred_bin)
# 방금 만든 0 또는 1을 마스크에 적용.

dice_before = binary_dice(worst_case["pred"], worst_case["label"], target_class=1)
dice_after = binary_dice(worst_pred_post, worst_case["label"], target_class=1)

print("후처리 전 Dice: ", dice_before)
print("후처리 후 Dice: ", dice_after)

시각화

best_slice = get_best_slice_from_label(worst_case["label"])

plt.figure(figsize=(16, 8))

# 원본
plt.subplot(2, 3, 1)
plt.title("Original")
plt.imshow(worst_case["image"][:, :, best_slice], cmap="gray")
plt.axis("off")

# 정답
plt.subplot(2, 3, 2)
plt.title("Label")
plt.imshow(worst_case["label"][:, :, best_slice], cmap="gray")
plt.axis("off")

# 후처리 전 예측 결과
plt.subplot(2, 3, 3)
plt.title(f"Prediction Before\nDice={dice_before:.4f}")
plt.imshow(worst_case["pred"][:, :, best_slice], cmap="gray")
plt.axis("off")

# 후처리 전 오버레이
plt.subplot(2, 3, 4)
plt.title("Overlay Before")
plt.imshow(worst_case["image"][:, :, best_slice], cmap="gray")
plt.imshow(worst_case["pred"][:, :, best_slice], cmap="Reds", alpha=0.35)
plt.axis("off")

# 후처리 후 예측 결과
plt.subplot(2, 3, 5)
plt.title(f"Prediction After\nDice={dice_after:.4f}")
plt.imshow(worst_pred_post[:, :, best_slice], cmap="gray")
plt.axis("off")

# 후처리 후 오버레이
plt.subplot(2, 3, 6)
plt.title("Overlay After")
plt.imshow(worst_case["image"][:, :, best_slice], cmap="gray")
plt.imshow(worst_pred_post[:, :, best_slice], cmap="Reds", alpha=0.35)
plt.axis("off")

plt.tight_layout()
plt.show()

---

11. 요약 ⭐️

---

12. 고찰 🧐

이번 실습을 통해 3D 영상 segmentation 방법에 대해 연습해볼 수 있었고,
후처리를 통해서도 성능을 얼마나 향상시킬 수 있는지 체감할 수 있었다.

그럼에도 안정적인 분할 성능을 얻지는 못했어서 아래와 같이 생각을 해보았다.

• 아무리 MONAI가 분할 학습으로도 어느정도 공간적인 정보를 학습한다지만
  patch로 나눠 학습하면서 충분히 공간적인 특징(주변 장기와의 위치적 연결성)이 학습이 덜 된 것은 아닐까?
  이를 개선하는 방향으로 학습 방법을 조절해보는 방법을 찾아보자.

• 위와 연결되는 맥락으로 우리는 정확히 비장을 segmentation 해야하고,
  사람마다 조금씩 차이는 있을지라도 큰 틀에서의 위치적 특징은 존재한다.
  그러면 비장의 대략적 위치를 먼저 좁힌 뒤 segmentation하는 coarse-to-fine 접근을 취하면 어떨까?

모델의 성능 어떻게 판단할지도 그렇고 생각해야할 점이 많이 남은 실습이었다.