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

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],距离为 2 \sqrt{2} 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,进行数据增强

©️2020 CSDN 皮肤主题: 游动-白 设计师:上身试试 返回首页