""" 
@Date: 2021/08/12
@description:
"""
import torch
import torch.nn as nn


class LEDLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.loss = nn.L1Loss()

    def forward(self, gt, dt):
        camera_height = 1.6
        gt_depth = gt['depth'] * camera_height

        dt_ceil_depth = dt['ceil_depth'] * camera_height * gt['ratio']
        dt_floor_depth = dt['depth'] * camera_height

        ceil_loss = self.loss(gt_depth, dt_ceil_depth)
        floor_loss = self.loss(gt_depth, dt_floor_depth)

        loss = floor_loss + ceil_loss

        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(),
        'ceil_depth': gt['depth'] / gt['ratio']
    }
    # dummy_dt['depth'][..., :20] *= 3  # some different

    led_loss = LEDLoss()
    loss = led_loss(gt, dummy_dt)
    print(loss)