语义分割计算label图的边缘

"""
摘自GSCNN
Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
"""

import os

import cv2
import numpy as np
from PIL import Image
from scipy.ndimage.morphology import distance_transform_edt

def mask_to_onehot(mask, num_classes):
    """
    Converts a segmentation mask (H,W) to (K,H,W) where the last dim is a one
    hot encoding vector

    """
    _mask = [mask == (i + 1) for i in range(num_classes)]
    return np.array(_mask).astype(np.uint8)

def onehot_to_mask(mask):
    """
    Converts a mask (K,H,W) to (H,W)
    """
    _mask = np.argmax(mask, axis=0)
    _mask[_mask != 0] += 1
    return _mask

def onehot_to_multiclass_edges(mask, radius, num_classes):
    """
    Converts a segmentation mask (K,H,W) to an edgemap (K,H,W)

    """
    if radius < 0:
        return mask
    
    # We need to pad the borders for boundary conditions
    mask_pad = np.pad(mask, ((0, 0), (1, 1), (1, 1)), mode='constant', constant_values=0)
    
    channels = []
    for i in range(num_classes):
        dist = distance_transform_edt(mask_pad[i, :])+distance_transform_edt(1.0-mask_pad[i, :])
        dist = dist[1:-1, 1:-1]
        dist[dist > radius] = 0
        dist = (dist > 0).astype(np.uint8)
        channels.append(dist)
        
    return np.array(channels)


''' 多类别的label图中计算一个像素点宽的边缘'''
def onehot_to_binary_edges(mask, radius, num_classes):
    """
    Converts a segmentation mask (K,H,W) to a binary edgemap (H,W)

    """
    
    if radius < 0:
        return mask
    
    # We need to pad the borders for boundary conditions
    mask_pad = np.pad(mask, ((0, 0), (1, 1), (1, 1)), mode='constant', constant_values=0)
    
    edgemap = np.zeros(mask.shape[1:])

    for i in range(num_classes):
        dist = distance_transform_edt(mask_pad[i, :])+distance_transform_edt(1.0-mask_pad[i, :])
        dist = dist[1:-1, 1:-1]
        dist[dist > radius] = 0
        edgemap += dist
    edgemap = np.expand_dims(edgemap, axis=0)    
    edgemap = (edgemap > 0).astype(np.uint8)
    return edgemap

"""自己制作的简单版"""
# 用于单类别的标签图片
def edge_det(image_path=None,image_array=None,is_array=True,return_array=True):
    # image_path表示label图片的路径
    # image_array是label图片所生成的数组
    # is_array表示我们选用image_path还是image_array
    #return_array表示返回数组还是图片

    if is_array:
        img=image_array.astype(np.float)
    else:
        img=Image.open(image_path)
        img = np.array(img).astype(np.float)

    img=(img>=1)*1
    _edgemap = distance_transform_edt(img)  # distance_transform_edt：该函数用于求前景像素点距背景像素点的距离
    edgemap = (_edgemap == 1) * 1

    if return_array:
        return edgemap
    else:
        edgemap = edgemap.astype(np.uint8)
        edgemap = Image.fromarray(edgemap)
        return edgemap

#用canny算法检测出边缘
def edge_det_canny(image_path=None,image_array=None,is_array=True,return_array=True):
    if is_array:
        img = image_array.astype(np.uint8)
    else:
        img = Image.open(image_path)
        img = np.array(img)
    edgemap = cv2.Canny(img, 2, 10)
    edgemap = (edgemap>0)*1
    if return_array:
        return edgemap
    else:
        edgemap = edgemap.astype(np.uint8)*255
        edgemap = Image.fromarray(edgemap)
        return edgemap