python--使用distance_transform_edt生成trimap(三值图)

distance_transform_edt介绍

distance_transform_edt是scipy库中的函数，调用方式为：

from scipy import ndimage
ndimage.morphology.distance_transform_edt()

作用是距离转换，计算图像中非零点到最近背景点（即0）的距离。举例如下：
下面展示一些 内联代码片。

import numpy as np
from scipy import ndimage

mat = np.array(([1, 0, 0, 0, 1],
                [1, 1, 1, 0, 1],
                [0, 1, 1, 1, 0],
                [0, 0, 0, 1, 1],
                [1, 1, 1, 1, 1]))
               
dis = ndimage.morphology.distance_transform_edt(mat)
print(dis)

结果：

[[1.         0.         0.         0.         1.        ]
 [1.         1.         1.         0.         1.        ]
 [0.         1.         1.         1.         0.        ]
 [0.         0.         0.         1.         1.        ]
 [1.         1.         1.         1.41421356 2.        ]]

从结果上来看，对于mat[4][3]，最近的0值点在mat[3][2]，距离为$\sqrt{2}$，对于mat[4][4]，最近的0值点在mat[2][4]，距离为2。

trimap生成

在抠图任务中既image matting任务中，常需要trimap作为输入来训练/生成抠图结果。一般trimap通过alpha图的腐蚀膨胀来生成，本文介绍的函数同样可以用来生成trimap。
废话不多说，直接放代码：

def generate_trimap(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 * 255
    trimap[unknown] = 128
    return trimap.astype(np.uint8)

其中
morphology.distance_transform_edt(unknown == 0) <= np.random.randint(1, 20)等同于腐蚀膨胀的效果；
通过“<=”后的数值控制生成的unknown区域的大小；
通过np.random.randint(1, 20)生成不同unknown区域大小的trimaps，进行数据增强。