|
- import os
- 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)
-
-
- 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
-
|