| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class DepthLoss(nn.Module): |
| def __init__(self, type='l1'): |
| super(DepthLoss, self).__init__() |
| self.type = type |
|
|
|
|
| def forward(self, depth_pred, depth_gt, mask=None): |
| if (depth_gt < 0).sum() > 0: |
| |
| return torch.tensor(0.0).to(depth_pred.device) |
| if mask is not None: |
| mask_d = (depth_gt > 0).float() |
|
|
| mask = mask * mask_d |
|
|
| mask_sum = mask.sum() + 1e-5 |
| depth_error = (depth_pred - depth_gt) * mask |
| depth_loss = F.l1_loss(depth_error, torch.zeros_like(depth_error).to(depth_error.device), |
| reduction='sum') / mask_sum |
| else: |
| depth_error = depth_pred - depth_gt |
| depth_loss = F.l1_loss(depth_error, torch.zeros_like(depth_error).to(depth_error.device), |
| reduction='mean') |
| return depth_loss |
|
|
| def forward(self, depth_pred, depth_gt, mask=None): |
| if mask is not None: |
| mask_d = (depth_gt > 0).float() |
|
|
| mask = mask * mask_d |
|
|
| mask_sum = mask.sum() + 1e-5 |
| depth_error = (depth_pred - depth_gt) * mask |
| depth_loss = F.l1_loss(depth_error, torch.zeros_like(depth_error).to(depth_error.device), |
| reduction='sum') / mask_sum |
| else: |
| depth_error = depth_pred - depth_gt |
| depth_loss = F.l1_loss(depth_error, torch.zeros_like(depth_error).to(depth_error.device), |
| reduction='mean') |
| return depth_loss |
|
|
| class DepthSmoothLoss(nn.Module): |
| def __init__(self): |
| super(DepthSmoothLoss, self).__init__() |
|
|
| def forward(self, disp, img, mask): |
| """ |
| Computes the smoothness loss for a disparity image |
| The color image is used for edge-aware smoothness |
| :param disp: [B, 1, H, W] |
| :param img: [B, 1, H, W] |
| :param mask: [B, 1, H, W] |
| :return: |
| """ |
| grad_disp_x = torch.abs(disp[:, :, :, :-1] - disp[:, :, :, 1:]) |
| grad_disp_y = torch.abs(disp[:, :, :-1, :] - disp[:, :, 1:, :]) |
|
|
| grad_img_x = torch.mean(torch.abs(img[:, :, :, :-1] - img[:, :, :, 1:]), 1, keepdim=True) |
| grad_img_y = torch.mean(torch.abs(img[:, :, :-1, :] - img[:, :, 1:, :]), 1, keepdim=True) |
|
|
| grad_disp_x *= torch.exp(-grad_img_x) |
| grad_disp_y *= torch.exp(-grad_img_y) |
|
|
| grad_disp = (grad_disp_x * mask[:, :, :, :-1]).mean() + (grad_disp_y * mask[:, :, :-1, :]).mean() |
|
|
| return grad_disp |
|
|
