File size: 4,292 Bytes
3ef1661 |
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 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
import torch
import torch.nn as nn
EPSILON = 1e-6
"""
# @Zhengqi Li version.
def GradientLoss(self, log_prediction_d, mask, log_gt):
log_d_diff = log_prediction_d - log_gt
v_gradient = torch.abs(log_d_diff[:, :-2, :] - log_d_diff[:, 2:, :])
v_mask = torch.mul(mask[:, :-2, :], mask[:, 2:, :])
v_gradient = torch.mul(v_gradient, v_mask)
h_gradient = torch.abs(log_d_diff[:, :, :-2] - log_d_diff[:, :, 2:])
h_mask = torch.mul(mask[:, :, :-2], mask[:, :, 2:])
h_gradient = torch.mul(h_gradient, h_mask)
N = torch.sum(h_mask) + torch.sum(v_mask) + EPSILON
gradient_loss = torch.sum(h_gradient) + torch.sum(v_gradient)
gradient_loss = gradient_loss / N
return gradient_loss
"""
def gradient_log_loss(log_prediction_d, log_gt, mask):
log_d_diff = log_prediction_d - log_gt
v_gradient = torch.abs(log_d_diff[:, :, :-2, :] - log_d_diff[:, :, 2:, :])
v_mask = torch.mul(mask[:, :, :-2, :], mask[:, :, 2:, :])
v_gradient = torch.mul(v_gradient, v_mask)
h_gradient = torch.abs(log_d_diff[:, :, :, :-2] - log_d_diff[:, :, :, 2:])
h_mask = torch.mul(mask[:, :, :, :-2], mask[:, :, :, 2:])
h_gradient = torch.mul(h_gradient, h_mask)
N = torch.sum(h_mask) + torch.sum(v_mask) + EPSILON
gradient_loss = torch.sum(h_gradient) + torch.sum(v_gradient)
gradient_loss = gradient_loss / N
return gradient_loss
class GradientLoss_Li(nn.Module):
def __init__(self, scale_num=1, loss_weight=1, data_type = ['lidar', 'stereo'], **kwargs):
super(GradientLoss_Li, self).__init__()
self.__scales = scale_num
self.loss_weight = loss_weight
self.data_type = data_type
self.eps = 1e-6
def forward(self, prediction, target, mask, **kwargs):
total = 0
target_trans = target + (~mask) * 100
pred_log = torch.log(prediction)
gt_log = torch.log(target_trans)
for scale in range(self.__scales):
step = pow(2, scale)
total += gradient_log_loss(pred_log[:, ::step, ::step], gt_log[:, ::step, ::step], mask[:, ::step, ::step])
loss = total / self.__scales
if torch.isnan(loss).item() | torch.isinf(loss).item():
raise RuntimeError(f'VNL error, {loss}')
return loss * self.loss_weight
######################################################
# Multi-scale gradient matching loss, @Ke Xian implementation.
#####################################################
def gradient_loss(prediction, target, mask):
M = torch.sum(mask, (1, 2))
diff = prediction - target
diff = torch.mul(mask, diff)
grad_x = torch.abs(diff[:, :, 1:] - diff[:, :, :-1])
mask_x = torch.mul(mask[:, :, 1:], mask[:, :, :-1])
grad_x = torch.mul(mask_x, grad_x)
grad_y = torch.abs(diff[:, 1:, :] - diff[:, :-1, :])
mask_y = torch.mul(mask[:, 1:, :], mask[:, :-1, :])
grad_y = torch.mul(mask_y, grad_y)
image_loss = torch.sum(grad_x, (1, 2)) + torch.sum(grad_y, (1, 2))
valid = M.nonzero()
if image_loss[valid].numel() > 0:
image_loss[valid] = image_loss[valid] / M[valid]
loss = torch.mean(image_loss)
else:
loss = 0 * torch.sum(prediction)
return loss
class GradientLoss(nn.Module):
def __init__(self, scale_num=4, loss_weight=1, **kwargs):
super(GradientLoss, self).__init__()
self.__scales = scale_num
self.loss_weight = loss_weight
def forward(self, prediction, target, mask, **kwargs):
total = 0
for scale in range(self.__scales):
step = pow(2, scale)
total += gradient_loss(prediction[:, ::step, ::step], target[:, ::step, ::step], mask[:, ::step, ::step])
return total * self.loss_weight
if __name__ == '__main__':
import numpy as np
gradient = GradientLoss_Li(4)
pred_depth = np.random.random([2, 1, 480, 640])
gt_depth = np.ones_like(pred_depth) * (-1) #np.random.random([2, 1, 480, 640]) - 0.5 #
#gt_depth = np.abs(gt_depth)
intrinsic = [[100, 100, 200, 200], [200, 200, 300, 300]]
pred = torch.from_numpy(pred_depth).cuda()
gt = torch.from_numpy(gt_depth).cuda()
mask = gt > 0
loss = gradient(gt, gt, mask)
print(loss) |