pytorch_edge_weighting_segmentation

import torch
import torch.nn as nn
import torch.nn.functional as F

import matplotlib.pyplot as plt


output = F.log_softmax(torch.randn(1, 3, 24, 24), 1)
target = torch.zeros(1, 24, 24, dtype=torch.long)
target[0, 4:12, 4:12] = 1
target[0, 14:20, 14:20] = 2

# Edge calculation
bin_target = torch.where(target > 0, torch.tensor(1), torch.tensor(0))
plt.imshow(bin_target[0])

o = F.avg_pool2d(bin_target.float(), kernel_size=3, padding=1, stride=1)
plt.imshow(o[0])

edge_idx = (o.ge(0.01) * o.le(0.99)).float()
plt.imshow(edge_idx[0])

weights = torch.ones_like(edge_idx, dtype=torch.float)
weights_sum0 = weights.sum()
weights = weights + edge_idx * 2.
weights_sum1 = weights.sum()
weights = weights / weights_sum1 * weights_sum0 # Rescale weigths
plt.imshow(weights[0])

# Calculate loss
criterion = nn.NLLLoss(reduce=False)
loss = criterion(output, target)
loss = loss * weights
loss = loss.sum() / weights.sum()