File size: 5,661 Bytes
d2410ba |
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 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
# Please cite "4D Spatio-Temporal ConvNets: Minkowski Convolutional Neural
# Networks", CVPR'19 (https://arxiv.org/abs/1904.08755) if you use any part
# of the code.
import torch
import numpy as np
from collections.abc import Sequence
def fnv_hash_vec(arr):
'''
FNV64-1A
'''
assert arr.ndim == 2
# Floor first for negative coordinates
arr = arr.copy()
arr = arr.astype(np.uint64, copy=False)
hashed_arr = np.uint64(14695981039346656037) * \
np.ones(arr.shape[0], dtype=np.uint64)
for j in range(arr.shape[1]):
hashed_arr *= np.uint64(1099511628211)
hashed_arr = np.bitwise_xor(hashed_arr, arr[:, j])
return hashed_arr
def ravel_hash_vec(arr):
'''
Ravel the coordinates after subtracting the min coordinates.
'''
assert arr.ndim == 2
arr = arr.copy()
arr -= arr.min(0)
arr = arr.astype(np.uint64, copy=False)
arr_max = arr.max(0).astype(np.uint64) + 1
keys = np.zeros(arr.shape[0], dtype=np.uint64)
# Fortran style indexing
for j in range(arr.shape[1] - 1):
keys += arr[:, j]
keys *= arr_max[j + 1]
keys += arr[:, -1]
return keys
def sparse_quantize(coords,
feats=None,
labels=None,
ignore_label=255,
set_ignore_label_when_collision=False,
return_index=False,
hash_type='fnv',
quantization_size=1):
r'''Given coordinates, and features (optionally labels), the function
generates quantized (voxelized) coordinates.
Args:
coords (:attr:`numpy.ndarray` or :attr:`torch.Tensor`): a matrix of size
:math:`N \times D` where :math:`N` is the number of points in the
:math:`D` dimensional space.
feats (:attr:`numpy.ndarray` or :attr:`torch.Tensor`, optional): a matrix of size
:math:`N \times D_F` where :math:`N` is the number of points and
:math:`D_F` is the dimension of the features.
labels (:attr:`numpy.ndarray`, optional): labels associated to eah coordinates.
ignore_label (:attr:`int`, optional): the int value of the IGNORE LABEL.
set_ignore_label_when_collision (:attr:`bool`, optional): use the `ignore_label`
when at least two points fall into the same cell.
return_index (:attr:`bool`, optional): True if you want the indices of the
quantized coordinates. False by default.
hash_type (:attr:`str`, optional): Hash function used for quantization. Either
`ravel` or `fnv`. `ravel` by default.
quantization_size (:attr:`float`, :attr:`list`, or
:attr:`numpy.ndarray`, optional): the length of the each side of the
hyperrectangle of of the grid cell.
.. note::
Please check `examples/indoor.py` for the usage.
'''
use_label = labels is not None
use_feat = feats is not None
if not use_label and not use_feat:
return_index = True
assert hash_type in [
'ravel', 'fnv'
], "Invalid hash_type. Either ravel, or fnv allowed. You put hash_type=" + hash_type
assert coords.ndim == 2, \
"The coordinates must be a 2D matrix. The shape of the input is " + str(coords.shape)
if use_feat:
assert feats.ndim == 2
assert coords.shape[0] == feats.shape[0]
if use_label:
assert coords.shape[0] == len(labels)
# Quantize the coordinates
dimension = coords.shape[1]
if isinstance(quantization_size, (Sequence, np.ndarray, torch.Tensor)):
assert len(
quantization_size
) == dimension, "Quantization size and coordinates size mismatch."
quantization_size = [i for i in quantization_size]
elif np.isscalar(quantization_size): # Assume that it is a scalar
quantization_size = [quantization_size for i in range(dimension)]
else:
raise ValueError('Not supported type for quantization_size.')
discrete_coords = np.floor(coords / np.array(quantization_size))
# Hash function type
if hash_type == 'ravel':
key = ravel_hash_vec(discrete_coords)
else:
key = fnv_hash_vec(discrete_coords)
if use_label:
_, inds, counts = np.unique(key, return_index=True, return_counts=True)
filtered_labels = labels[inds]
if set_ignore_label_when_collision:
filtered_labels[counts > 1] = ignore_label
if return_index:
return inds, filtered_labels
else:
return discrete_coords[inds], feats[inds], filtered_labels
else:
_, inds, inds_reverse = np.unique(key, return_index=True, return_inverse=True)
# NOTE:
if use_feat:
voxel_feats = np.zeros((len(np.unique(key)), feats.shape[1]), dtype=feats.dtype)
for i in range(len(np.unique(key))):
# voxel_feats[i] = np.mean(feats[inds_reverse == i], axis=0)
# voxel_feats[i] = np.median(feats[inds_reverse == i], axis=0)
voxel_center = np.mean(coords[inds_reverse == i], axis=0)
distances = np.linalg.norm(coords[inds_reverse == i] - voxel_center, axis=1)
central_point_idx = np.argmin(distances)
voxel_feats[i] = feats[inds_reverse == i][central_point_idx]
if return_index:
return inds, inds_reverse, voxel_feats
##############
if return_index:
return inds, inds_reverse
else:
if use_feat:
return discrete_coords[inds], feats[inds]
else:
return discrete_coords[inds] |