matting之trimap生成_膨胀腐蚀

在抠图技术中三分图（trimap)经常被用到，通常使用的方法是膨胀腐蚀(一般在去除噪声的时候先腐蚀再膨胀)。

1.

import os
import numpy as np
import cv2

def random_dilate(alpha, low=1, high=5, mode='constant'):
    """Dilation. erode"""
    iterations = np.random.randint(1, 20)
    erode_ksize = np.random.randint(low=low, high=high)
    dilate_ksize = np.random.randint(low=low, high=high)
    erode_kernel = cv2.getStructuringElement(
        cv2.MORPH_ELLIPSE, (erode_ksize, erode_ksize)) # 椭圆 (尺寸)
    dilate_kernel = cv2.getStructuringElement(
        cv2.MORPH_ELLIPSE, (dilate_ksize, dilate_ksize))
    alpha_eroded = cv2.erode(alpha, erode_kernel, iterations=iterations)
    alpha_dilated = cv2.dilate(alpha, dilate_kernel, iterations=iterations)
    if mode == 'constant':
        alpha_noise = 128 * np.ones_like(alpha)
        alpha_noise[alpha_eroded >= 255] = 255 ###250
        alpha_noise[alpha_dilated <= 0] = 0
    else:
        value = np.random.randint(low=100, high=255)
        alpha_noise = value * ((alpha_dilated - alpha_eroded) / 255.)
        alpha_noise += alpha_eroded
    return alpha_noise

2.

def erode_dilate(msk, struc="ELLIPSE", size=(10, 10)):
    if struc == "RECT":
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, size)
    elif struc == "CORSS":
        kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, size)
    else:
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, size)

    #msk = msk.astype(np.float32)
    #print(msk)
    msk = msk / 255
    #print(msk.shape)

    #msk = msk.astype(np.uint8)

    # val in 0 or 255
    iterations = 3
    dilated = cv2.dilate(msk, kernel, iterations=iterations) * 255
    eroded = cv2.erode(msk, kernel, iterations=iterations) * 255 # 腐蚀
    # 一般在去噪声时先用腐蚀再用膨胀。

    # cnt1 = len(np.where(msk >= 0)[0])
    # cnt2 = len(np.where(msk == 0)[0])
    # cnt3 = len(np.where(msk == 1)[0])
    # #print("all:{} bg:{} fg:{}".format(cnt1, cnt2, cnt3))
    # assert(cnt1 == cnt2 + cnt3) # 边缘模糊所以存在两者之间的像素
    #
    # cnt1 = len(np.where(dilated >= 0)[0])
    # cnt2 = len(np.where(dilated == 0)[0])
    # cnt3 = len(np.where(dilated == 255)[0])
    # #print("all:{} bg:{} fg:{}".format(cnt1, cnt2, cnt3))
    # assert(cnt1 == cnt2 + cnt3)
    #
    # cnt1 = len(np.where(eroded >= 0)[0])
    # cnt2 = len(np.where(eroded == 0)[0])
    # cnt3 = len(np.where(eroded == 255)[0])
    # #print("all:{} bg:{} fg:{}".format(cnt1, cnt2, cnt3))
    # assert(cnt1 == cnt2 + cnt3)

    res = dilated.copy()
    #res[((dilated == 255) & (msk == 0))] = 128
    res[((dilated == 255) & (eroded == 0))] = 128 # alpha

    return res

一些细节上的处理，导致结果有些不同，根据数据的实际情况，灵活运用，酌情修改

---

更新

from scipy.ndimage import morphology    
def getTrimap(self, alpha):
	fg = np.array(np.equal(alpha, 255).astype(np.float32))
	unknown = np.array(np.not_equal(alpha, 0).astype(np.float32))  # unknown = alpha > 0
	unknown = unknown - fg
	unknown = morphology.distance_transform_edt(unknown == 0) <= np.random.randint(1, 20)
	trimap = fg
	trimap[unknown] = 0.5
	# print(trimap[:, :, :1].shape)
	return trimap[:, :, :1]

---

二次更新：

from scipy.ndimage import grey_dilation, grey_erosion
low, high = 2, 8  ### 根据具体情况再定义
d_size = np.random.randint(low=low, high=high)
e_size = np.random.randint(low=low, high=high)
matte = label / 255.  # $numpy array of your matte(with values between[0, 1])$

trimap = (matte >= 0.9).astype('float32')
not_bg = (matte > 0).astype('float32')

trimap[np.where(
    (grey_dilation(not_bg, size=(d_size, d_size)) - grey_erosion(trimap, size=(e_size, e_size))) != 0)] = 0.5

Reference:

Python - OpenCV 之图像形态学(膨胀与腐蚀)