Spaces:
Running on Zero
Running on Zero
File size: 3,403 Bytes
0122a25 | 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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """Normalize Transform."""
from __future__ import annotations
import torch
from mapdet3d.common.typing import NDArrayF32
from ..const import CommonKeys as K
from .base import Transform
@Transform(K.images, K.images)
class NormalizeImages:
"""Normalize a list of image tensor with given mean and std.
Image tensor is of shape [N, H, W, C] and range (0, 255).
"""
def __init__(
self,
mean: tuple[float, float, float] = (123.675, 116.28, 103.53),
std: tuple[float, float, float] = (58.395, 57.12, 57.375),
epsilon: float = 1e-08,
) -> None:
"""Creates an instance of NormalizeImage.
Args:
mean (Tuple[float, float, float], optional): Mean value. Defaults
to (123.675, 116.28, 103.53).
std (Tuple[float, float, float], optional): Standard deviation
value. Defaults to (58.395, 57.12, 57.375).
epsilon (float, optional): Epsilon for numerical stability of
division. Defaults to 1e-08.
"""
self.mean = mean
self.std = std
self.epsilon = epsilon
def __call__(self, images: list[NDArrayF32]) -> list[NDArrayF32]:
"""Normalize image tensor."""
for i, image in enumerate(images):
img = torch.from_numpy(image).permute(0, 3, 1, 2)
pixel_mean = torch.tensor(self.mean).view(-1, 1, 1)
pixel_std = torch.tensor(self.std).view(-1, 1, 1)
img = (img - pixel_mean) / (pixel_std + self.epsilon)
images[i] = img.permute(0, 2, 3, 1).numpy()
return images
@Transform(K.depth_maps, [K.depth_maps, "depth_shift_scale"])
class NormalizeDepthMaps:
"""Normalize a list of depth map tensor with given mean and std."""
def __init__(self, trunc_value: float = 0.1, eps: float = 1e-2) -> None:
"""Init."""
self.trunc_value = trunc_value
self.eps = eps
def __call__(self, depth_maps: list[NDArrayF32]):
"""Standardize depth maps."""
normalized_depth_maps = []
depth_shift_scale = []
for depth in depth_maps:
depth_img = torch.from_numpy(depth)
# Set invalid depth to nan
depth_img[depth_img <= 0.0] = torch.nan
sorted_img = torch.sort(torch.flatten(depth_img))[0]
# Remove nan, nan at the end of sort
num_nan = sorted_img.isnan().sum()
if num_nan > 0:
sorted_img = sorted_img[:-num_nan]
# Remove outliers
trunc_img = sorted_img[
int(self.trunc_value * len(sorted_img)) : int(
(1 - self.trunc_value) * len(sorted_img)
)
]
if len(trunc_img) <= 1:
# guard against no valid Jasper.
trunc_mean = 0.0
trunc_std = 1.0
else:
trunc_mean = trunc_img.mean().item()
trunc_std = torch.sqrt(trunc_img.var() + self.eps).item()
# Replace nan by mean
depth_img = torch.nan_to_num(depth_img, nan=trunc_mean)
# Standardize
depth_img = (depth_img - trunc_mean) / trunc_std
normalized_depth_maps.append(depth_img.numpy())
depth_shift_scale.append([trunc_mean, trunc_std])
return normalized_depth_maps, depth_shift_scale
|