Add File

pcdet/utils/transform_utils.py

@@ -0,0 +1,91 @@
+import math
+import torch
+
+try:
+    from kornia.geometry.conversions import (
+        convert_points_to_homogeneous,
+        convert_points_from_homogeneous,
+    )
+except:
+    pass 
+    # print('Warning: kornia is not installed. This package is only required by CaDDN')
+
+
+def project_to_image(project, points):
+    """
+    Project points to image
+    Args:
+        project [torch.tensor(..., 3, 4)]: Projection matrix
+        points [torch.Tensor(..., 3)]: 3D points
+    Returns:
+        points_img [torch.Tensor(..., 2)]: Points in image
+        points_depth [torch.Tensor(...)]: Depth of each point
+    """
+    # Reshape tensors to expected shape
+    points = convert_points_to_homogeneous(points)
+    points = points.unsqueeze(dim=-1)
+    project = project.unsqueeze(dim=1)
+
+    # Transform points to image and get depths
+    points_t = project @ points
+    points_t = points_t.squeeze(dim=-1)
+    points_img = convert_points_from_homogeneous(points_t)
+    points_depth = points_t[..., -1] - project[..., 2, 3]
+
+    return points_img, points_depth
+
+
+def normalize_coords(coords, shape):
+    """
+    Normalize coordinates of a grid between [-1, 1]
+    Args:
+        coords: (..., 3), Coordinates in grid
+        shape: (3), Grid shape
+    Returns:
+        norm_coords: (.., 3), Normalized coordinates in grid
+    """
+    min_n = -1
+    max_n = 1
+    shape = torch.flip(shape, dims=[0])  # Reverse ordering of shape
+
+    # Subtract 1 since pixel indexing from [0, shape - 1]
+    norm_coords = coords / (shape - 1) * (max_n - min_n) + min_n
+    return norm_coords
+
+
+def bin_depths(depth_map, mode, depth_min, depth_max, num_bins, target=False):
+    """
+    Converts depth map into bin indices
+    Args:
+        depth_map: (H, W), Depth Map
+        mode: string, Discretiziation mode (See https://arxiv.org/pdf/2005.13423.pdf for more details)
+            UD: Uniform discretiziation
+            LID: Linear increasing discretiziation
+            SID: Spacing increasing discretiziation
+        depth_min: float, Minimum depth value
+        depth_max: float, Maximum depth value
+        num_bins: int, Number of depth bins
+        target: bool, Whether the depth bins indices will be used for a target tensor in loss comparison
+    Returns:
+        indices: (H, W), Depth bin indices
+    """
+    if mode == "UD":
+        bin_size = (depth_max - depth_min) / num_bins
+        indices = ((depth_map - depth_min) / bin_size)
+    elif mode == "LID":
+        bin_size = 2 * (depth_max - depth_min) / (num_bins * (1 + num_bins))
+        indices = -0.5 + 0.5 * torch.sqrt(1 + 8 * (depth_map - depth_min) / bin_size)
+    elif mode == "SID":
+        indices = num_bins * (torch.log(1 + depth_map) - math.log(1 + depth_min)) / \
+            (math.log(1 + depth_max) - math.log(1 + depth_min))
+    else:
+        raise NotImplementedError
+
+    if target:
+        # Remove indicies outside of bounds
+        mask = (indices < 0) | (indices > num_bins) | (~torch.isfinite(indices))
+        indices[mask] = num_bins
+
+        # Convert to integer
+        indices = indices.type(torch.int64)
+    return indices