Spaces:
Build error
Build error
""" | |
@Date: 2021/08/12 | |
@description: For HorizonNet, using latitudes to calculate loss. | |
""" | |
import torch | |
import torch.nn as nn | |
from utils.conversion import depth2xyz, xyz2lonlat | |
class BoundaryLoss(nn.Module): | |
def __init__(self): | |
super().__init__() | |
self.loss = nn.L1Loss() | |
def forward(self, gt, dt): | |
gt_floor_xyz = depth2xyz(gt['depth']) | |
gt_ceil_xyz = gt_floor_xyz.clone() | |
gt_ceil_xyz[..., 1] = -gt['ratio'] | |
gt_floor_boundary = xyz2lonlat(gt_floor_xyz)[..., -1:] | |
gt_ceil_boundary = xyz2lonlat(gt_ceil_xyz)[..., -1:] | |
gt_boundary = torch.cat([gt_floor_boundary, gt_ceil_boundary], dim=-1).permute(0, 2, 1) | |
dt_boundary = dt['boundary'] | |
loss = self.loss(gt_boundary, dt_boundary) | |
return loss | |
if __name__ == '__main__': | |
import numpy as np | |
from dataset.mp3d_dataset import MP3DDataset | |
mp3d_dataset = MP3DDataset(root_dir='../src/dataset/mp3d', mode='train') | |
gt = mp3d_dataset.__getitem__(0) | |
gt['depth'] = torch.from_numpy(gt['depth'][np.newaxis]) # batch size is 1 | |
gt['ratio'] = torch.from_numpy(gt['ratio'][np.newaxis]) # batch size is 1 | |
dummy_dt = { | |
'depth': gt['depth'].clone(), | |
'boundary': torch.cat([ | |
xyz2lonlat(depth2xyz(gt['depth']))[..., -1:], | |
xyz2lonlat(depth2xyz(gt['depth'], plan_y=-gt['ratio']))[..., -1:] | |
], dim=-1).permute(0, 2, 1) | |
} | |
# dummy_dt['boundary'][:, :, :20] /= 1.2 # some different | |
boundary_loss = BoundaryLoss() | |
loss = boundary_loss(gt, dummy_dt) | |
print(loss) | |