| 229 | return loss |
| 230 | |
| 231 | def unified_step_focal_loss(prob_volume, depth_values, interval, depth_gt, mask, weight, gamma, alpha): |
| 232 | depth_gt_volume = depth_gt.unsqueeze(1).expand_as(depth_values) # (b, d, h, w) |
| 233 | |
| 234 | gt_index_volume = (depth_values-depth_gt_volume).abs()<=interval |
| 235 | |
| 236 | gt_unity_index_volume = torch.zeros_like(prob_volume, requires_grad=False) |
| 237 | gt_unity_index_volume[gt_index_volume] = 1.0 - (depth_gt_volume[gt_index_volume] - depth_values[gt_index_volume]).abs() / interval |
| 238 | |
| 239 | gt_unity, _ = torch.max(gt_unity_index_volume, dim=1, keepdim=True) |
| 240 | gt_unity = torch.where(gt_unity > 0.0, gt_unity, torch.ones_like(gt_unity)) # (b, 1, h, w) |
| 241 | pos_weight = (sigmoid((gt_unity - prob_volume).abs() / gt_unity, base=5) - 0.5) * 4 + 1 # [1, 3] |
| 242 | neg_weight = (sigmoid(prob_volume / gt_unity, base=5) - 0.5) * 2 # [0, 1] |
| 243 | focal_weight = (gt_unity_index_volume > 0.0).float() + alpha * (gt_unity_index_volume <= 0.0).float() |
| 244 | |
| 245 | mask = mask.unsqueeze(1).expand_as(depth_values).float() |
| 246 | # offset=prob_volume-1 |
| 247 | # torch.where |
| 248 | prob_volume=prob_volume/(prob_volume.max()) |
| 249 | loss = (F.binary_cross_entropy(prob_volume, gt_unity_index_volume, reduction="none") * focal_weight * mask).sum() / mask.sum() |
| 250 | loss = loss * weight |
| 251 | return loss |
| 252 | def unified_focal_loss(prob_volume, depth_values, interval, depth_gt, mask, weight, gamma, alpha): |
| 253 | depth_gt_volume = depth_gt.unsqueeze(1).expand_as(depth_values) # (b, d, h, w) |
| 254 | |