| import torch |
| import torch.nn as nn |
|
|
| class SSILoss(nn.Module): |
| """ |
| Scale shift invariant MAE loss. |
| loss = MAE((d-median(d)/s - (d'-median(d'))/s'), s = mean(d- median(d)) |
| """ |
| def __init__(self, loss_weight=1, data_type=['sfm', 'stereo', 'lidar'], **kwargs): |
| super(SSILoss, self).__init__() |
| self.loss_weight = loss_weight |
| self.data_type = data_type |
| self.eps = 1e-6 |
| |
| def ssi_mae(self, target, prediction, mask): |
| valid_pixes = torch.sum(mask) + self.eps |
|
|
| gt_median = torch.median(target) if target.numel() else 0 |
| gt_s = torch.abs(target - gt_median).sum() / valid_pixes |
| gt_trans = (target - gt_median) / (gt_s + self.eps) |
|
|
| pred_median = torch.median(prediction) if prediction.numel() else 0 |
| pred_s = torch.abs(prediction - pred_median).sum() / valid_pixes |
| pred_trans = (prediction - pred_median) / (pred_s + self.eps) |
| |
| ssi_mae_sum = torch.sum(torch.abs(gt_trans - pred_trans)) |
| return ssi_mae_sum, valid_pixes |
|
|
| def forward(self, prediction, target, mask=None, **kwargs): |
| """ |
| Calculate loss. |
| """ |
| B, C, H, W = prediction.shape |
| loss = 0 |
| valid_pix = 0 |
| for i in range(B): |
| mask_i = mask[i, ...] |
| gt_depth_i = target[i, ...][mask_i] |
| pred_depth_i = prediction[i, ...][mask_i] |
| ssi_sum, valid_pix_i = self.ssi_mae(pred_depth_i, gt_depth_i, mask_i) |
| loss += ssi_sum |
| valid_pix += valid_pix_i |
| loss /= (valid_pix + self.eps) |
| return loss * self.loss_weight |
| |
| if __name__ == '__main__': |
| torch.manual_seed(1) |
| torch.cuda.manual_seed_all(1) |
|
|
| ssil = SSILoss() |
| pred = torch.rand((2, 1, 256, 256)).cuda() |
| gt = torch.rand((2, 1, 256, 256)).cuda() |
| gt[:, :, 100:256, 0:100] = -1 |
| mask = gt > 0 |
| out = ssil(pred, gt, mask) |
| print(out) |
|
|
