CS 공부/AI

[Segmentation] Baseline Code

imsmile2000 2023. 6. 5. 23:51

평가 Metric : Dice


  • A = Ground Truth
  • B = Prediction

def dice_coef(y_true,y_pred):  #(B,C,H,W)
	y_true_f=y_true.flatten(2)
    y_pred_f=y_pred.flatten(2)
    intersection=torch.sum(y_true_f * y_pred_f,-1)
    
    eps=0.0001
    return (2. * intersection + eps) / (torch.sum(y_true_f,-1) + torch.sum(y_pred_f,-1) +eps)
    # 각 채널(클래스) 별로 수행

 

 

Baseline


  • Dataloader
    • __init__
      • filenames: 각 이미지 경로가 담긴 list
      • labelnames: 각 라벨 경로가 담긴 list
      • transforms: Transform
      • is_train: 생성될 객체의 용도가 train인지 valid인지 (True면 train)
    • __getitem__
      • Dataloader에서 사용할 data 반환
      • image와 label data 불러옴
      • label 생성시 각 point를 mask로 변환 (cv.fillPoly 함수 이용)
      • train이면 transform된 image와 label data 반환
      • valid면 transform된 image와 원본 label data 반환
      • Image와 label 객체의 채널을 (512,512,3) -> (3,512,512)로 transpose
    • Train/valid split: Group K-fold
    • Group K-fold로 split된 데이터로 torch dataloader 정의

 

  • Model: Torchvision을 이용한 pretrained model 불러오기 및 output class 수정
from torchvision import models
model=models.segmentation.fcn_resnet50(pretrained=True)

model.classifier[4]=nn.Conv2d(512,len(CLASSES),kernel_size=1) # output class data class에 맞게 수정

 

  • Loss / Optimizer
# Loss function 정의
criterion=nn.BCEWithLogitsLoss()
# Optimizer 정의
optimizer=optim.Adam(params=model.parameters(),lr=LR, weight_decay=1e-6)

 

  • Train
    • 정해진 epoch만큼 model을 이요해 forward 수행
    • 앞서 정의한 loss와 optimizer로 backward 수행

 

  • Validation
    • 모델 size를 original size로 복원 (512 -> 2048)
    • Sigmoid 함수를 사용해 prediction 값을 0~1범위로 변환 후, threshold(=0.5) 초과는 1로, 나머지는 0으로 설정
    • 평가 Metric Dice coefficient 계산

 

  • Best model 불러오기
    • 학습된 model이 저장된 pt 파일을 불러와 inference에 활용

 

  • Test (Submission.csv)
    • Output에 대한 RLE 인코딩 사용
    • RLE(Run-Length-Encoding): 파일 size 줄이기 위해 각 pixel 값 대신 'start position, total length'로 output 저장

18 3 25 5 33 6 41 7 50 6으로 인코딩 됨