Python/PyTorch 공부

[PyTorch] skimage모듈 mark_boundaries 함수를 사용하여 segmentation 경계 표시하기

AI 꿈나무 2021. 6. 11. 00:45
반응형

 안녕하세요! 이번 포스팅에서는 skimage 모듈의 mark_boundaries 함수를 사용하여 senmentation 경계 표시하는 법을 살펴보겠습니다 ㅎㅎ !

 

 mark_boundaries 함수를 사용하면, 이미지와 마스크 파일을 하나로 결합하여, 이미지 내에서 마스크 경계 부분만 표시합니다. 아래 사진이 예시입니다 ㅎㅎ

 

 

 우선 필요한 라이브러리를 import 합니다.

import os
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from scipy import ndimage as ndi
from skimage.segmentation import mark_boundaries

 

 데이터 셋을 다운로드 받습니다.

 위 홈페이지에서 training_set.zip와 test_test.zip 파일을 다운로드 받습니다. 

 해당 데이터셋은 태아의 초음파 이미지와 태아 머리의 mask 이미지가 담겨있습니다.

 압축을 풉니다.

 

 데이터 경로를 지정하고 개수를 확인합니다.

# train data 경로
path2train = '/content/segnet/MyDrive/data/fetus/training_set'

# 개수 확인
imgsList = [pp for pp in os.listdir(path2train) if 'Annotation' not in pp]
anntsList = [pp for pp in os.listdir(path2train) if 'Annotation' in pp]

print('number of images:', len(imgsList))
print('number of annotations:', len(anntsList))

 

 랜덤 이미지를 선택합니다.

# select random images
np.random.seed(2019)
rndImgs = np.random.choice(imgsList, 4)
rndImgs

 

 이미지 시각화 함수를 정의합니다.

# show image function
def show_img_mask(img, mask):
    img_mask = mark_boundaries(np.array(img), np.array(mask), outline_color=(0,1,0), color=(0,1,0))
    plt.imshow(img_mask)

 

 랜덤으로 선택된 이미지를 시각화 합니다.

# display the images and masks
for fn in rndImgs:
    path2img = os.path.join(path2train, fn)
    path2annt = path2img.replace('.png', '_Annotation.png')
    img = Image.open(path2img)
    annt_edges = Image.open(path2annt)
    mask = ndi.binary_fill_holes(annt_edges)

    plt.figure(figsize=(10,10))
    plt.subplot(1,3,1)
    plt.imshow(img, cmap='gray')

    plt.subplot(1, 3, 2)
    plt.imshow(mask, cmap='gray')

    plt.subplot(1, 3, 3)
    show_img_mask(img,mask)

 

 

 감사합니다.

반응형