hustvl / ViTMatte

[Information Fusion (Vol.103, Mar. '24)] Boosting Image Matting with Pretrained Plain Vision Transformers
MIT License
363 stars 37 forks source link

how to get trimap? #4

Closed ZZfive closed 1 year ago

ZZfive commented 1 year ago

Model's inputs are images and trimap. The trimaps are obtained by expanding or corroding alpha,but when predicting,the alpha is the goal, how do i get image trimap without alpha? Is there an end-to-end method for generating trimap based on images?

JingfengYao commented 1 year ago

Trimap is manually generated hint map for natural image matting. As for other ways, maybe you could refer to our tech report, Matte Anything, for more detail. Hope it could help.

xiezhongzhao commented 1 year ago
def generate_trimap(mask, erode_kernel_size=10, dilate_kernel_size=10):
    erode_kernel = np.ones((erode_kernel_size, erode_kernel_size), np.uint8)
    dilate_kernel = np.ones((dilate_kernel_size, dilate_kernel_size), np.uint8)
    eroded = cv2.erode(mask, erode_kernel, iterations=5)
    dilated = cv2.dilate(mask, dilate_kernel, iterations=5)
    trimap = np.zeros_like(mask)
    trimap[dilated==255] = 128
    trimap[eroded==255] = 255
    return trimap