[ISBI 2012 segmentation] U-Net 모델로 구현해보기-1

연구실에 처음 들어와서 공부할 때 구현해본, U-Net model 이다. 당시에는 CV 를 한창 공부해나가던 시절이기에 이런
간단한 모델을 구현하는 것 조차 상당히 애먹었던게 기억난다... 옛날에 노션에 정리한 것에 내용을 더 추가하였다.

 

논문에 대한 내용이 궁금하다면, 아래 포스팅을 참고하면 된다.

https://cvad.tistory.com/10

[논문리뷰] U-Net : Convolutional Networks for Biomedical Image Segmentation

---

0. Dataset download

먼저, 모델을 학습시키는 데 필요한 데이터셋을 다운 받아야한다. 구글링하면 많이 나오지만, 나의 경우 아래의 Github 링크에서 다운받았다.

https://github.com/alexklibisz/isbi-2012

GitHub - alexklibisz/isbi-2012: Image Segmentation Techniques on the ISBI 2012 dataset: http://brainiac2.mit.edu/isbi_challenge/

 

위의 링크로부터 파일을 받으면,  tif 형태로 저장된 파일들이 존재하는데 이 파일들을 이용하여 학습 이미지와 마스크 이미지를 생성해야한다. 참고로 해당 파일들은 아래와 같이 생겼다.

 

tif 파일들

 

주어진 데이터셋의 정보를 조금 더 확인해보자. 아래의 코드를 실행하면 된다.

import os
from PIL import Image

data_path = '../../Data/ISBI/'
source_img = 'train-volume.tif'
source_lbl = 'train-labels.tif'
source_test_img = 'test-volume.tif'

make_path_list = ['train/image', 'train/label', 'test']

for i in make_path_list:
    path = data_path + i
    # 만약 경로에 파일 존재 시, 파일 새로 생성
    if os.path.exists(path):
        print(f'{path} is already existed')
    else:
       print(f'{path} is created.')
         
       os.makedirs(path)

print()

tif_img = Image.open(os.path.join(data_path, source_img))
tif_lbl = Image.open(os.path.join(data_path, source_lbl))
tif_test_img = Image.open(os.path.join(data_path, source_test_img))


print(f'[Training data] # of image : {tif_img.n_frames} || # of label : {tif_lbl.n_frames}')
print(f'[Training data] (Hight, Width, "channel") || Image : {tif_img.size} || Label : {tif_lbl.size}')
print()
print(f'[Test data] # of image : {tif_test_img.n_frames} || # of label : {tif_test_img.n_frames}')
print(f'[Test data] (Hight, Width, "channel") || Image : {tif_test_img.size}')

======================================================================================
../../Data/ISBI/train/image is already existed
../../Data/ISBI/train/label is already existed
../../Data/ISBI/test is already existed

[Training data] # of image : 30 || # of label : 30
[Training data] (Hight, Width, "channel") || Image : (512, 512) || Label : (512, 512)

[Test data] # of image : 30 || # of label : 30
[Test data] (Hight, Width, "channel") || Image : (512, 512)

 

다행히 사진들의 크기 512x512로 일정하고 프레임도 마스크 이미지와 입 이미지가 동일하다.

 

(CARLA 데이터들 전처리 할때를 생각해보면, 이게 얼마나 다행인지 모른다....)

 

그럼 이제 프레임마다 이미지와 마스크를 한 장씩 저장해보자. 아래의 코드를 실행하면 된다.
나의 경우 train dataset에서 validation dataset을 따로 지정하여 저장하도록 코드를 짰다.

(저의 경우 따로 저장해놓는 것이 경험상 편하여 그런 것이니 torch의 random_split으로 쓰셔도 됩니다)

 

import numpy as np

# frame에서 랜덤하게 idx를 추출하기 위한 index 값 생성
random_idx = np.arange(tif_img.n_frames)
np.random.shuffle(random_idx)

# Validation data 저장할 path 생성
val_paths = ['valid/image', 'valid/label']

for path in val_paths:
    if os.path.exists(os.path.join(data_path, path)):
        pass
    else:
        os.makedirs(os.path.join(data_path, path))
        print(f'{path} is created.')


# train과 valid의 비율이 8:2가 되도록 설정
crit = int(0.8*tif_img.n_frames)
cnt = 1

for i in random_idx:
    tif_img.seek(i)
    tif_lbl.seek(i)
    tif_test_img.seek(i)

    img = tif_img.copy()
    lbl = tif_lbl.copy()
    test_img = tif_test_img.copy()
    
    if cnt <= crit:
        img_save_path = data_path + '/train/image'
        lbl_save_path = data_path + '/train/label'
    else:
        img_save_path = data_path + '/valid/image'
        lbl_save_path = data_path + '/valid/label'

    img.save(img_save_path+f'/image_{cnt-1}.jpg')    
    lbl.save(lbl_save_path+f'/image_{cnt-1}.png')
    test_img.save(data_path+f'/test/image_{cnt-1}.jpg')

    cnt+=1

 

이제 학습을 위한 데이터셋 준비는 다 끝났다.

학습에 들어가기 앞서서 논문에서 언급한 것처럼 마스크 데이터에 class imbalance가 실제로 큰 지 알아보고싶었다.

 

결론적으로 background로 정의된 픽셀(0값이 픽셀)보다 세포로 정의된 픽셀(255)가 거의 4배가량 많았다.
확인해보고 싶다면 아래의 코드를 실행시켜보자

import cv2
import re

lbl_path = '../../Data/ISBI/train/label/'

def sorter(text):
    num = re.findall(r'd\+', text)
    return int(num[0]) if num else 0

lbl_list = sorted(os.listdir(lbl_path), key=sorter)

total_class = {}

for lbl_name in lbl_list:
    lbl_file_path = os.path.join(lbl_path, lbl_name)
    lbl = cv2.imread(lbl_file_path, cv2.IMREAD_GRAYSCALE)

    # mask 데이터에 존재하는 class 확인
    uniq = np.unique(lbl)

    # mask에 존재하는 class의 픽셀 갯수 카운팅
    for cls in uniq:
        num = np.count_nonzero(lbl == cls)
        # 만약 처음보는 class의 경우 딕셔너리에 추가한 뒤 갯수 기록
        if not num in total_class:
            total_class[cls] = num
        # 이미 있는 class의 경우 값을 업데이트
        else:
            total_class[cls] += num

total_class

===============================================================================================
{0: 54700, 255: 207444}

 

물론 이상태에서 UNet을 학습시켜도 괜찮지만, 조금 더 논문의 흐름을 따라가보고 싶었다.

---

1. Patch division

 U-Net 논문에서는 4개의 augmentation( shift, rotate, gray value, random elasitc deformation)과 overlap-tile
strategy를 사용했다고 한다. 나머지는 이미 pytorch에서 구현되어있는데 문제는 overlap-tile이었다.

우선 기본 이미지를 무조건 잘라서 input tile 형태로 만들어줘야할 뿐만 아니라, 이 과정을 dataset 정의 이전에 처리해야 하지 않을까라는 생각이 들었다.

 

코드 작성 전에 한 가지 의아한 점은 U-Net architecture에서 input tile size는 572x572 이었고, 출력된 값의 크기, 즉 patch size는 388x388이다. 우리가 사용하는 image의 크기를 고려해본다면 patch size가 뭔가 이상했다. 
왜냐하면, 512x512를 겹치지 않게 patch 분할을 한다고 하면 저 크기값은 말이 안된다

 

가능한 경우의 수를 생각해보면

   (1) 한 이미지에 388x388로 random crop을 하여 single patch로 진행하였다.

         이 경우 모델의 input이 말이 되긴하지만, 굳이 한 이미지로 single patch를 만들거면 왜 원본 이미지에 바로
         mirroring을 하지 않고 굳이 patch로 잘라서 했을까 싶다.

   (2) ISBI가 아닌 다른 dataset 이거나, resizing을 통해 원본 이미지를 늘린 후 진행하였다.
        아마도 다른 dataset에서 사용한 크기일 수도 있겠다는 생각이 들었다. 

 

아직 이것에 대한 이유를 찾진 못했지만, 만약 찾을 경우 추가적으로 포스팅 하겠다.

 

다시 본론으로 돌아가서, overlap-tile을 적용할 코드는 아래와 같다.

 

나는 임의로 4개의 patch로 원본이미지를 분할하였고, 미러링의 경우 94개의 pixel 값을 기준으로 적용하였다.

( 제가 임의로 설정한 조건으로, 논문에서 제안하는 수치가 아닙니다!)

아래의 그림을 참고하면 더 정확하게 이해가 될 것이다.

 

 

patch 분할에 대한 코드는 아래와 같다.

def img2patch(image, patch_size):
    # RGB 이미지인지 마스크 이미지인지 확인
    if len(image.shape) == 2:  # Mask image
        num_patches_axis = image.shape[0] // patch_size
    elif len(image.shape) == 3:  # RGB image
        num_patches_axis = image.shape[0] // patch_size
    else:
        raise ValueError("Invalid image shape. Image should be either 2D (grayscale) or 3D (RGB).")

    # Prepare an array to hold the patches
    patches = []
    
    for i in range(num_patches_axis):
        for j in range(num_patches_axis):
            # 좌 상단부터 우측으로 이동하며 patch를 분할
            start_i = i * patch_size
            end_i = start_i + patch_size
            start_j = j * patch_size
            end_j = start_j + patch_size
            
            # patch 픽셀들 추출
            if len(image.shape) == 2:  # Grayscale image
                patch = image[start_i:end_i, start_j:end_j]
            else:  # RGB image
                patch = image[start_i:end_i, start_j:end_j, :]
            
            # Add the patch to the list
            patches.append(patch)
    
    # numpy 형태로 변환하여 저장
    patches_array = np.stack(patches)
    
    return patches_array

 

이 코드들을 사용하여 기존의 데이터셋에 적용해주면 된다.

train_img_path = '../../Data/ISBI/train/image'
train_lbl_path = '../../Data/ISBI/train/label'

valid_img_path = '../../Data/ISBI/valid/image/'
valid_lbl_path = '../../Data/ISBI/valid/label/'


for i in range(2):
    s=0
    idx = 0
    if i == 0:
        img_path = train_img_path
        lbl_path = train_lbl_path
    else:
        img_path = valid_img_path
        lbl_path = valid_lbl_path

    new_img_save_path = img_path+'/aug'
    new_lbl_save_path = lbl_path+'/aug'
    paths = [new_img_save_path, new_lbl_save_path]

    for path in paths:
        if os.path.exists(path):
            shutil.rmtree(path)
        
    img_list = sorted(os.listdir(img_path), key=sorter)
    lbl_list = sorted(os.listdir(lbl_path), key=sorter)

    if s == 0:
        for path in paths:
            if not os.path.exists(path):
                os.makedirs(path)
        s+=1

    for j in range(len(img_list)):
        img = cv2.imread(os.path.join(img_path, img_list[j]), cv2.IMREAD_COLOR)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        lbl = cv2.imread(os.path.join(lbl_path, lbl_list[j]), cv2.IMREAD_GRAYSCALE)

        img_patches = img2patch(img, 256)
        lbl_patches = img2patch(lbl, 256)

        for k in range(len(img_patches)):
            input_tile = img_patches[k]
            new_mask = lbl_patches[k]

            cv2.imwrite(new_img_save_path + f'/image_{idx}.jpg', input_tile)
            cv2.imwrite(new_lbl_save_path + f'/image_{idx}.png', new_mask)
            idx += 1

 

이렇게하면 4개로 분할된 patch들로 dataset을 나눌 수 있다.

이제 다음 포스팅부터는 이 데이터들을 가지고 전처리 및 dataset을 정의하는 내용을 다루겠다.