Python과 OpenCV – 29 : GrabCut 알고리즘을 이용한 전경(Foreground) 추출

이 글의 원문은 https://github.com/abidrahmank/OpenCV2-Python-Tutorials/blob/master/source/py_tutorials/py_imgproc/py_grabcut/py_grabcut.rst 입니다.

GrabCut 알고리즘은 이미지에서 배경이 아닌 전경에 해당하는 이미지를 추출해 내는 방법입니다. 이미지에서 한번에 전경을 추출해 내는 것이 아닌 사용자와의 상호 작용을 통해 단계적으로 전경을 추출합니다. 이 상호작용은 크게 2가지 단계로 진행되는데, 첫번째는 이미지에서 전경이 포함되는 영역을 사각형으로 대략적으로 지정합니다. 단, 이때 지정한 사각형 영역 안에는 전경이 모두 포함되어 있어야 합니다. 그리고 두번째는 첫번째에서 얻어진 전경 이미지의 내용중 포함되어진 배경 부분은 어디인지, 누락된 전경 부분은 어디인지를 마킹하면 이를 이용해 다시 전경 이미지가 새롭게 추출됩니다.

위의 GrabCut 알고리즘에 대한 OpenCV의 구현은 cv2.grabCut 입니다. 예제를 통해 이 함수를 살펴보겠습니다. 첫번째 예제는 앞서 언급한 GrabCut의 상호작용 중 전경이 포함된 사각 영역을 지정하는 것만을 이용해 전경을 추출하는 경우입니다.

import numpy as np
import cv2
from matplotlib import pyplot as plt

img = cv2.imread('./data/messi5.jpg')
mask = np.zeros(img.shape[:2],np.uint8)

bgdModel = np.zeros((1,65),np.float64)
fgdModel = np.zeros((1,65),np.float64)

rect = (50,50,450,290)
cv2.grabCut(img,mask,rect,bgdModel,fgdModel,5,cv2.GC_INIT_WITH_RECT)

mask2 = np.where((mask==2)|(mask==0),0,1).astype('uint8')
img = img*mask2[:,:,np.newaxis]

plt.imshow(img),plt.colorbar(),plt.show()

결과는 아래와 같습니다.

코드 중 전경이 포함된 사각영역은 11번 코드에서 rect 변수에 지정하고 있습니다. 이 rect 변수를 이용해 12번 코드에서 전경을 추출하는 cv2.grabCut 함수가 실행됩니다. 이 함수의 인자를 살펴보면 첫번째는 입력 이미지, 두번째는 6번 코드에서 생성한 모든 요소가 0인 마스크 이미지, 세번째는 사각형역, 네번째와 다섯번째는 이 알고리즘의 수행 과정중 활용할 메모리, 여섯번째는 해당 알고리즘의 반복 수행 횟수, 일곱번재는 사각형 영역 지정을 통한 GrabCut일 경우 cv2.GC_INIT_WITH_RECT를 지정합니다. 이 인자 중 6번째인 mask에는 전경에 해당하는 화소 위치에 값이 설정되는데, 0 또는 2 값은 배경이고 1 또는 3인 전경이라는 의미입니다. 이를 이용해 원본 이미지에 마스킹 처리를 해서 전경만을 표시하는 것이 14-17번 코드입니다.

다음은 사각형 영역의 지정 이후에 추출된 전경에 대해 다시 누락된 전경과 포함된 배경 부분을 마킹하여 보다 완전한 전경 이미지를 추출하는 전체적인 GrabCut 알고리즘에 대한 예제입니다.

import numpy as np
import cv2
from matplotlib import pyplot as plt

img = cv2.imread('./data/messi5.jpg')
mask = np.zeros(img.shape[:2],np.uint8)

bgdModel = np.zeros((1,65),np.float64)
fgdModel = np.zeros((1,65),np.float64)

# Step 1
rect = (50,50,450,290)
cv2.grabCut(img,mask,rect,bgdModel,fgdModel,1,cv2.GC_INIT_WITH_RECT)

# Step 2
newmask = cv2.imread('./data/newmask2.png',0)
mask[newmask == 0] = 0
mask[newmask == 255] = 1
cv2.grabCut(img,mask,None,bgdModel,fgdModel,1,cv2.GC_INIT_WITH_MASK)

mask2 = np.where((mask==2)|(mask==0),0,1).astype('uint8')
img = img*mask2[:,:,np.newaxis]
plt.imshow(img),plt.colorbar(),plt.show()

결과는 다음과 같습니다.

코드 중 Step1으로 처리된 부분은 앞서 설명한 사각영역 지정을 통한 전경의 추출이고 Step2가 마킹을 통한 전경 이미지 추출입니다. 마킹 방식은 newmask2.png에 마킹 정보를 나타내고 있는데 이 이미지는 아래와 같습니다.

즉, 하얀색 픽셀값(255)은 전경을, 검정색 픽셀값(0)은 배경임을 나타내는 마킹 정보를 포함하고 있는 이미지입니다.