279 lines
9.2 KiB
Python
279 lines
9.2 KiB
Python
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
"""
|
|
@Author : Peike Li
|
|
@Contact : peike.li@yahoo.com
|
|
@File : lovasz_softmax.py
|
|
@Time : 8/30/19 7:12 PM
|
|
@Desc : Lovasz-Softmax and Jaccard hinge loss in PyTorch
|
|
Maxim Berman 2018 ESAT-PSI KU Leuven (MIT License)
|
|
@License : This source code is licensed under the license found in the
|
|
LICENSE file in the root directory of this source tree.
|
|
"""
|
|
|
|
from __future__ import print_function, division
|
|
|
|
import torch
|
|
from torch.autograd import Variable
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
from torch import nn
|
|
|
|
try:
|
|
from itertools import ifilterfalse
|
|
except ImportError: # py3k
|
|
from itertools import filterfalse as ifilterfalse
|
|
|
|
|
|
def lovasz_grad(gt_sorted):
|
|
"""
|
|
Computes gradient of the Lovasz extension w.r.t sorted errors
|
|
See Alg. 1 in paper
|
|
"""
|
|
p = len(gt_sorted)
|
|
gts = gt_sorted.sum()
|
|
intersection = gts - gt_sorted.float().cumsum(0)
|
|
union = gts + (1 - gt_sorted).float().cumsum(0)
|
|
jaccard = 1. - intersection / union
|
|
if p > 1: # cover 1-pixel case
|
|
jaccard[1:p] = jaccard[1:p] - jaccard[0:-1]
|
|
return jaccard
|
|
|
|
|
|
def iou_binary(preds, labels, EMPTY=1., ignore=None, per_image=True):
|
|
"""
|
|
IoU for foreground class
|
|
binary: 1 foreground, 0 background
|
|
"""
|
|
if not per_image:
|
|
preds, labels = (preds,), (labels,)
|
|
ious = []
|
|
for pred, label in zip(preds, labels):
|
|
intersection = ((label == 1) & (pred == 1)).sum()
|
|
union = ((label == 1) | ((pred == 1) & (label != ignore))).sum()
|
|
if not union:
|
|
iou = EMPTY
|
|
else:
|
|
iou = float(intersection) / float(union)
|
|
ious.append(iou)
|
|
iou = mean(ious) # mean accross images if per_image
|
|
return 100 * iou
|
|
|
|
|
|
def iou(preds, labels, C, EMPTY=1., ignore=None, per_image=False):
|
|
"""
|
|
Array of IoU for each (non ignored) class
|
|
"""
|
|
if not per_image:
|
|
preds, labels = (preds,), (labels,)
|
|
ious = []
|
|
for pred, label in zip(preds, labels):
|
|
iou = []
|
|
for i in range(C):
|
|
if i != ignore: # The ignored label is sometimes among predicted classes (ENet - CityScapes)
|
|
intersection = ((label == i) & (pred == i)).sum()
|
|
union = ((label == i) | ((pred == i) & (label != ignore))).sum()
|
|
if not union:
|
|
iou.append(EMPTY)
|
|
else:
|
|
iou.append(float(intersection) / float(union))
|
|
ious.append(iou)
|
|
ious = [mean(iou) for iou in zip(*ious)] # mean accross images if per_image
|
|
return 100 * np.array(ious)
|
|
|
|
|
|
# --------------------------- BINARY LOSSES ---------------------------
|
|
|
|
|
|
def lovasz_hinge(logits, labels, per_image=True, ignore=None):
|
|
"""
|
|
Binary Lovasz hinge loss
|
|
logits: [B, H, W] Variable, logits at each pixel (between -\infty and +\infty)
|
|
labels: [B, H, W] Tensor, binary ground truth masks (0 or 1)
|
|
per_image: compute the loss per image instead of per batch
|
|
ignore: void class id
|
|
"""
|
|
if per_image:
|
|
loss = mean(lovasz_hinge_flat(*flatten_binary_scores(log.unsqueeze(0), lab.unsqueeze(0), ignore))
|
|
for log, lab in zip(logits, labels))
|
|
else:
|
|
loss = lovasz_hinge_flat(*flatten_binary_scores(logits, labels, ignore))
|
|
return loss
|
|
|
|
|
|
def lovasz_hinge_flat(logits, labels):
|
|
"""
|
|
Binary Lovasz hinge loss
|
|
logits: [P] Variable, logits at each prediction (between -\infty and +\infty)
|
|
labels: [P] Tensor, binary ground truth labels (0 or 1)
|
|
ignore: label to ignore
|
|
"""
|
|
if len(labels) == 0:
|
|
# only void pixels, the gradients should be 0
|
|
return logits.sum() * 0.
|
|
signs = 2. * labels.float() - 1.
|
|
errors = (1. - logits * Variable(signs))
|
|
errors_sorted, perm = torch.sort(errors, dim=0, descending=True)
|
|
perm = perm.data
|
|
gt_sorted = labels[perm]
|
|
grad = lovasz_grad(gt_sorted)
|
|
loss = torch.dot(F.relu(errors_sorted), Variable(grad))
|
|
return loss
|
|
|
|
|
|
def flatten_binary_scores(scores, labels, ignore=None):
|
|
"""
|
|
Flattens predictions in the batch (binary case)
|
|
Remove labels equal to 'ignore'
|
|
"""
|
|
scores = scores.view(-1)
|
|
labels = labels.view(-1)
|
|
if ignore is None:
|
|
return scores, labels
|
|
valid = (labels != ignore)
|
|
vscores = scores[valid]
|
|
vlabels = labels[valid]
|
|
return vscores, vlabels
|
|
|
|
|
|
class StableBCELoss(torch.nn.modules.Module):
|
|
def __init__(self):
|
|
super(StableBCELoss, self).__init__()
|
|
|
|
def forward(self, input, target):
|
|
neg_abs = - input.abs()
|
|
loss = input.clamp(min=0) - input * target + (1 + neg_abs.exp()).log()
|
|
return loss.mean()
|
|
|
|
|
|
def binary_xloss(logits, labels, ignore=None):
|
|
"""
|
|
Binary Cross entropy loss
|
|
logits: [B, H, W] Variable, logits at each pixel (between -\infty and +\infty)
|
|
labels: [B, H, W] Tensor, binary ground truth masks (0 or 1)
|
|
ignore: void class id
|
|
"""
|
|
logits, labels = flatten_binary_scores(logits, labels, ignore)
|
|
loss = StableBCELoss()(logits, Variable(labels.float()))
|
|
return loss
|
|
|
|
|
|
# --------------------------- MULTICLASS LOSSES ---------------------------
|
|
|
|
|
|
def lovasz_softmax(probas, labels, classes='present', per_image=False, ignore=255, weighted=None):
|
|
"""
|
|
Multi-class Lovasz-Softmax loss
|
|
probas: [B, C, H, W] Variable, class probabilities at each prediction (between 0 and 1).
|
|
Interpreted as binary (sigmoid) output with outputs of size [B, H, W].
|
|
labels: [B, H, W] Tensor, ground truth labels (between 0 and C - 1)
|
|
classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average.
|
|
per_image: compute the loss per image instead of per batch
|
|
ignore: void class labels
|
|
"""
|
|
if per_image:
|
|
loss = mean(lovasz_softmax_flat(*flatten_probas(prob.unsqueeze(0), lab.unsqueeze(0), ignore), classes=classes, weighted=weighted)
|
|
for prob, lab in zip(probas, labels))
|
|
else:
|
|
loss = lovasz_softmax_flat(*flatten_probas(probas, labels, ignore), classes=classes, weighted=weighted )
|
|
return loss
|
|
|
|
|
|
def lovasz_softmax_flat(probas, labels, classes='present', weighted=None):
|
|
"""
|
|
Multi-class Lovasz-Softmax loss
|
|
probas: [P, C] Variable, class probabilities at each prediction (between 0 and 1)
|
|
labels: [P] Tensor, ground truth labels (between 0 and C - 1)
|
|
classes: 'all' for all, 'present' for classes present in labels, or a list of classes to average.
|
|
"""
|
|
if probas.numel() == 0:
|
|
# only void pixels, the gradients should be 0
|
|
return probas * 0.
|
|
C = probas.size(1)
|
|
losses = []
|
|
class_to_sum = list(range(C)) if classes in ['all', 'present'] else classes
|
|
for c in class_to_sum:
|
|
fg = (labels == c).float() # foreground for class c
|
|
if (classes is 'present' and fg.sum() == 0):
|
|
continue
|
|
if C == 1:
|
|
if len(classes) > 1:
|
|
raise ValueError('Sigmoid output possible only with 1 class')
|
|
class_pred = probas[:, 0]
|
|
else:
|
|
class_pred = probas[:, c]
|
|
errors = (Variable(fg) - class_pred).abs()
|
|
errors_sorted, perm = torch.sort(errors, 0, descending=True)
|
|
perm = perm.data
|
|
fg_sorted = fg[perm]
|
|
if weighted is not None:
|
|
losses.append(weighted[c]*torch.dot(errors_sorted, Variable(lovasz_grad(fg_sorted))))
|
|
else:
|
|
losses.append(torch.dot(errors_sorted, Variable(lovasz_grad(fg_sorted))))
|
|
return mean(losses)
|
|
|
|
|
|
def flatten_probas(probas, labels, ignore=None):
|
|
"""
|
|
Flattens predictions in the batch
|
|
"""
|
|
if probas.dim() == 3:
|
|
# assumes output of a sigmoid layer
|
|
B, H, W = probas.size()
|
|
probas = probas.view(B, 1, H, W)
|
|
B, C, H, W = probas.size()
|
|
probas = probas.permute(0, 2, 3, 1).contiguous().view(-1, C) # B * H * W, C = P, C
|
|
labels = labels.view(-1)
|
|
if ignore is None:
|
|
return probas, labels
|
|
valid = (labels != ignore)
|
|
vprobas = probas[valid.nonzero().squeeze()]
|
|
vlabels = labels[valid]
|
|
return vprobas, vlabels
|
|
|
|
|
|
def xloss(logits, labels, ignore=None):
|
|
"""
|
|
Cross entropy loss
|
|
"""
|
|
return F.cross_entropy(logits, Variable(labels), ignore_index=255)
|
|
|
|
|
|
# --------------------------- HELPER FUNCTIONS ---------------------------
|
|
def isnan(x):
|
|
return x != x
|
|
|
|
|
|
def mean(l, ignore_nan=False, empty=0):
|
|
"""
|
|
nanmean compatible with generators.
|
|
"""
|
|
l = iter(l)
|
|
if ignore_nan:
|
|
l = ifilterfalse(isnan, l)
|
|
try:
|
|
n = 1
|
|
acc = next(l)
|
|
except StopIteration:
|
|
if empty == 'raise':
|
|
raise ValueError('Empty mean')
|
|
return empty
|
|
for n, v in enumerate(l, 2):
|
|
acc += v
|
|
if n == 1:
|
|
return acc
|
|
return acc / n
|
|
|
|
# --------------------------- Class ---------------------------
|
|
class LovaszSoftmax(nn.Module):
|
|
def __init__(self, per_image=False, ignore_index=255, weighted=None):
|
|
super(LovaszSoftmax, self).__init__()
|
|
self.lovasz_softmax = lovasz_softmax
|
|
self.per_image = per_image
|
|
self.ignore_index=ignore_index
|
|
self.weighted = weighted
|
|
|
|
def forward(self, pred, label):
|
|
pred = F.softmax(pred, dim=1)
|
|
return self.lovasz_softmax(pred, label, per_image=self.per_image, ignore=self.ignore_index, weighted=self.weighted) |