Spaces:
Running on Zero
Running on Zero
File size: 2,068 Bytes
5e5d775 | 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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | import random
from PIL import Image
import cv2
import numpy as np
import torch
from torchvision import transforms
from torchvision.transforms import functional as F
from torchvision.transforms import InterpolationMode
from loguru import logger
class ScaleShiftDepthNormalizer():
"""
Use near and far plane to linearly normalize depth,
i.e. d' = d * s + t,
where near plane is mapped to `norm_min`, and far plane is mapped to `norm_max`
Near and far planes are determined by taking quantile values.
"""
is_absolute = False
far_plane_at_max = True
def __init__(
self, norm_min=-1.0, norm_max=1.0, min_max_quantile=0.02, clip=True
) -> None:
self.norm_min = norm_min
self.norm_max = norm_max
self.norm_range = self.norm_max - self.norm_min
self.min_quantile = min_max_quantile
self.max_quantile = 1.0 - self.min_quantile
self.clip = clip
def __call__(self, depth_linear, valid_mask=None, clip=None):
clip = clip if clip is not None else self.clip
if valid_mask is None:
valid_mask = torch.ones_like(depth_linear).bool()
valid_mask = valid_mask & (depth_linear > 0)
# Take quantiles as min and max
_min, _max = torch.quantile(
depth_linear[valid_mask],
torch.tensor([self.min_quantile, self.max_quantile]),
)
# scale and shift
depth_norm_linear = (depth_linear - _min) / (
_max - _min
) * self.norm_range + self.norm_min
if clip:
depth_norm_linear = torch.clip(
depth_norm_linear, self.norm_min, self.norm_max
)
return depth_norm_linear
def scale_back(self, depth_norm):
# scale to [0, 1]
depth_linear = (depth_norm - self.norm_min) / self.norm_range
return depth_linear
def denormalize(self, depth_norm, **kwargs):
logger.warning(f"{self.__class__} is not revertible without GT")
return self.scale_back(depth_norm=depth_norm)
|