ZiyuG/SAM2Point

1# Please cite "4D Spatio-Temporal ConvNets: Minkowski Convolutional Neural2# Networks", CVPR'19 (https://arxiv.org/abs/1904.08755) if you use any part3# of the code.4import torch5import numpy as np6from collections.abc import Sequence7 8 9def fnv_hash_vec(arr):10    '''11    FNV64-1A12    '''13    assert arr.ndim == 214    # Floor first for negative coordinates15    arr = arr.copy()16    arr = arr.astype(np.uint64, copy=False)17    hashed_arr = np.uint64(14695981039346656037) * \18                 np.ones(arr.shape[0], dtype=np.uint64)19    for j in range(arr.shape[1]):20        hashed_arr *= np.uint64(1099511628211)21        hashed_arr = np.bitwise_xor(hashed_arr, arr[:, j])22    return hashed_arr23 24 25def ravel_hash_vec(arr):26    '''27    Ravel the coordinates after subtracting the min coordinates.28    '''29    assert arr.ndim == 230    arr = arr.copy()31    arr -= arr.min(0)32    arr = arr.astype(np.uint64, copy=False)33    arr_max = arr.max(0).astype(np.uint64) + 134 35    keys = np.zeros(arr.shape[0], dtype=np.uint64)36    # Fortran style indexing37    for j in range(arr.shape[1] - 1):38        keys += arr[:, j]39        keys *= arr_max[j + 1]40    keys += arr[:, -1]41    return keys42 43 44def sparse_quantize(coords,45                    feats=None,46                    labels=None,47                    ignore_label=255,48                    set_ignore_label_when_collision=False,49                    return_index=False,50                    hash_type='fnv',51                    quantization_size=1):52    r'''Given coordinates, and features (optionally labels), the function53    generates quantized (voxelized) coordinates.54 55    Args:56        coords (:attr:`numpy.ndarray` or :attr:`torch.Tensor`): a matrix of size57        :math:`N \times D` where :math:`N` is the number of points in the58        :math:`D` dimensional space.59 60        feats (:attr:`numpy.ndarray` or :attr:`torch.Tensor`, optional): a matrix of size61        :math:`N \times D_F` where :math:`N` is the number of points and62        :math:`D_F` is the dimension of the features.63 64        labels (:attr:`numpy.ndarray`, optional): labels associated to eah coordinates.65 66        ignore_label (:attr:`int`, optional): the int value of the IGNORE LABEL.67 68        set_ignore_label_when_collision (:attr:`bool`, optional): use the `ignore_label`69        when at least two points fall into the same cell.70 71        return_index (:attr:`bool`, optional): True if you want the indices of the72        quantized coordinates. False by default.73 74        hash_type (:attr:`str`, optional): Hash function used for quantization. Either75        `ravel` or `fnv`. `ravel` by default.76 77        quantization_size (:attr:`float`, :attr:`list`, or78        :attr:`numpy.ndarray`, optional): the length of the each side of the79        hyperrectangle of of the grid cell.80 81    .. note::82        Please check `examples/indoor.py` for the usage.83 84    '''85    use_label = labels is not None86    use_feat = feats is not None87    if not use_label and not use_feat:88        return_index = True89 90    assert hash_type in [91        'ravel', 'fnv'92    ], "Invalid hash_type. Either ravel, or fnv allowed. You put hash_type=" + hash_type93    assert coords.ndim == 2, \94        "The coordinates must be a 2D matrix. The shape of the input is " + str(coords.shape)95    if use_feat:96        assert feats.ndim == 297        assert coords.shape[0] == feats.shape[0]98    if use_label:99        assert coords.shape[0] == len(labels)100 101    # Quantize the coordinates102    dimension = coords.shape[1]103    if isinstance(quantization_size, (Sequence, np.ndarray, torch.Tensor)):104        assert len(105            quantization_size106        ) == dimension, "Quantization size and coordinates size mismatch."107        quantization_size = [i for i in quantization_size]108    elif np.isscalar(quantization_size):  # Assume that it is a scalar109        quantization_size = [quantization_size for i in range(dimension)]110    else:111        raise ValueError('Not supported type for quantization_size.')112    discrete_coords = np.floor(coords / np.array(quantization_size))113 114    # Hash function type115    if hash_type == 'ravel':116        key = ravel_hash_vec(discrete_coords)117    else:118        key = fnv_hash_vec(discrete_coords)119 120    if use_label:121        _, inds, counts = np.unique(key, return_index=True, return_counts=True)122        filtered_labels = labels[inds]123        if set_ignore_label_when_collision:124            filtered_labels[counts > 1] = ignore_label125        if return_index:126            return inds, filtered_labels127        else:128            return discrete_coords[inds], feats[inds], filtered_labels129    else:130        _, inds, inds_reverse = np.unique(key, return_index=True, return_inverse=True)131        # NOTE:132        if use_feat:133            voxel_feats = np.zeros((len(np.unique(key)), feats.shape[1]), dtype=feats.dtype)134            for i in range(len(np.unique(key))):135            #     voxel_feats[i] = np.mean(feats[inds_reverse == i], axis=0)136            #     voxel_feats[i] = np.median(feats[inds_reverse == i], axis=0)137                voxel_center = np.mean(coords[inds_reverse == i], axis=0)138                distances = np.linalg.norm(coords[inds_reverse == i] - voxel_center, axis=1)139                central_point_idx = np.argmin(distances)140                voxel_feats[i] = feats[inds_reverse == i][central_point_idx]141            if return_index:142                return inds, inds_reverse, voxel_feats143        ##############144        if return_index:145            return inds, inds_reverse146        else:147            if use_feat:148                return discrete_coords[inds], feats[inds]149            else:150                return discrete_coords[inds]