File size: 1,405 Bytes
938e515 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
import torch
from torch.nn import functional as F
def generate_edge_tensor(label, edge_width=3):
label = label.type(torch.cuda.FloatTensor)
if len(label.shape) == 2:
label = label.unsqueeze(0)
n, h, w = label.shape
edge = torch.zeros(label.shape, dtype=torch.float).cuda()
# right
edge_right = edge[:, 1:h, :]
edge_right[(label[:, 1:h, :] != label[:, :h - 1, :]) & (label[:, 1:h, :] != 255)
& (label[:, :h - 1, :] != 255)] = 1
# up
edge_up = edge[:, :, :w - 1]
edge_up[(label[:, :, :w - 1] != label[:, :, 1:w])
& (label[:, :, :w - 1] != 255)
& (label[:, :, 1:w] != 255)] = 1
# upright
edge_upright = edge[:, :h - 1, :w - 1]
edge_upright[(label[:, :h - 1, :w - 1] != label[:, 1:h, 1:w])
& (label[:, :h - 1, :w - 1] != 255)
& (label[:, 1:h, 1:w] != 255)] = 1
# bottomright
edge_bottomright = edge[:, :h - 1, 1:w]
edge_bottomright[(label[:, :h - 1, 1:w] != label[:, 1:h, :w - 1])
& (label[:, :h - 1, 1:w] != 255)
& (label[:, 1:h, :w - 1] != 255)] = 1
kernel = torch.ones((1, 1, edge_width, edge_width), dtype=torch.float).cuda()
with torch.no_grad():
edge = edge.unsqueeze(1)
edge = F.conv2d(edge, kernel, stride=1, padding=1)
edge[edge!=0] = 1
edge = edge.squeeze()
return edge
|