"""
摘自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
语义分割计算label图的边缘
最新推荐文章于 2023-06-25 19:40:18 发布