Spaces:
Sleeping
Sleeping
File size: 1,642 Bytes
98a77e0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
# Copyright (c) 2020-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
import torch
#----------------------------------------------------------------------------
# HDR image losses
#----------------------------------------------------------------------------
def _tonemap_srgb(f):
return torch.where(f > 0.0031308, torch.pow(torch.clamp(f, min=0.0031308), 1.0/2.4)*1.055 - 0.055, 12.92*f)
def _SMAPE(img, target, eps=0.01):
nom = torch.abs(img - target)
denom = torch.abs(img) + torch.abs(target) + 0.01
return torch.mean(nom / denom)
def _RELMSE(img, target, eps=0.1):
nom = (img - target) * (img - target)
denom = img * img + target * target + 0.1
return torch.mean(nom / denom)
def image_loss_fn(img, target, loss, tonemapper):
if tonemapper == 'log_srgb':
img = _tonemap_srgb(torch.log(torch.clamp(img, min=0, max=65535) + 1))
target = _tonemap_srgb(torch.log(torch.clamp(target, min=0, max=65535) + 1))
if loss == 'mse':
return torch.nn.functional.mse_loss(img, target)
elif loss == 'smape':
return _SMAPE(img, target)
elif loss == 'relmse':
return _RELMSE(img, target)
else:
return torch.nn.functional.l1_loss(img, target)
|