diff --git a/pcdet/utils/common_utils.py b/pcdet/utils/common_utils.py new file mode 100644 index 0000000..af70728 --- /dev/null +++ b/pcdet/utils/common_utils.py @@ -0,0 +1,295 @@ +import logging +import os +import pickle +import random +import shutil +import subprocess +import SharedArray + +import numpy as np +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + + +def check_numpy_to_torch(x): + if isinstance(x, np.ndarray): + return torch.from_numpy(x).float(), True + return x, False + + +def limit_period(val, offset=0.5, period=np.pi): + val, is_numpy = check_numpy_to_torch(val) + ans = val - torch.floor(val / period + offset) * period + return ans.numpy() if is_numpy else ans + + +def drop_info_with_name(info, name): + ret_info = {} + keep_indices = [i for i, x in enumerate(info['name']) if x != name] + for key in info.keys(): + ret_info[key] = info[key][keep_indices] + return ret_info + + +def rotate_points_along_z(points, angle): + """ + Args: + points: (B, N, 3 + C) + angle: (B), angle along z-axis, angle increases x ==> y + Returns: + + """ + points, is_numpy = check_numpy_to_torch(points) + angle, _ = check_numpy_to_torch(angle) + + cosa = torch.cos(angle) + sina = torch.sin(angle) + zeros = angle.new_zeros(points.shape[0]) + ones = angle.new_ones(points.shape[0]) + rot_matrix = torch.stack(( + cosa, sina, zeros, + -sina, cosa, zeros, + zeros, zeros, ones + ), dim=1).view(-1, 3, 3).float() + points_rot = torch.matmul(points[:, :, 0:3], rot_matrix) + points_rot = torch.cat((points_rot, points[:, :, 3:]), dim=-1) + return points_rot.numpy() if is_numpy else points_rot + + +def angle2matrix(angle): + """ + Args: + angle: angle along z-axis, angle increases x ==> y + Returns: + rot_matrix: (3x3 Tensor) rotation matrix + """ + + cosa = torch.cos(angle) + sina = torch.sin(angle) + rot_matrix = torch.tensor([ + [cosa, -sina, 0], + [sina, cosa, 0], + [ 0, 0, 1] + ]) + return rot_matrix + + +def mask_points_by_range(points, limit_range): + mask = (points[:, 0] >= limit_range[0]) & (points[:, 0] <= limit_range[3]) \ + & (points[:, 1] >= limit_range[1]) & (points[:, 1] <= limit_range[4]) + return mask + + +def get_voxel_centers(voxel_coords, downsample_times, voxel_size, point_cloud_range): + """ + Args: + voxel_coords: (N, 3) + downsample_times: + voxel_size: + point_cloud_range: + + Returns: + + """ + assert voxel_coords.shape[1] == 3 + voxel_centers = voxel_coords[:, [2, 1, 0]].float() # (xyz) + voxel_size = torch.tensor(voxel_size, device=voxel_centers.device).float() * downsample_times + pc_range = torch.tensor(point_cloud_range[0:3], device=voxel_centers.device).float() + voxel_centers = (voxel_centers + 0.5) * voxel_size + pc_range + return voxel_centers + + +def create_logger(log_file=None, rank=0, log_level=logging.INFO): + logger = logging.getLogger(__name__) + logger.setLevel(log_level if rank == 0 else 'ERROR') + formatter = logging.Formatter('%(asctime)s %(levelname)5s %(message)s') + console = logging.StreamHandler() + console.setLevel(log_level if rank == 0 else 'ERROR') + console.setFormatter(formatter) + logger.addHandler(console) + if log_file is not None: + file_handler = logging.FileHandler(filename=log_file) + file_handler.setLevel(log_level if rank == 0 else 'ERROR') + file_handler.setFormatter(formatter) + logger.addHandler(file_handler) + logger.propagate = False + return logger + + +def set_random_seed(seed): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + +def worker_init_fn(worker_id, seed=666): + if seed is not None: + random.seed(seed + worker_id) + np.random.seed(seed + worker_id) + torch.manual_seed(seed + worker_id) + torch.cuda.manual_seed(seed + worker_id) + torch.cuda.manual_seed_all(seed + worker_id) + + +def get_pad_params(desired_size, cur_size): + """ + Get padding parameters for np.pad function + Args: + desired_size: int, Desired padded output size + cur_size: int, Current size. Should always be less than or equal to cur_size + Returns: + pad_params: tuple(int), Number of values padded to the edges (before, after) + """ + assert desired_size >= cur_size + + # Calculate amount to pad + diff = desired_size - cur_size + pad_params = (0, diff) + + return pad_params + + +def keep_arrays_by_name(gt_names, used_classes): + inds = [i for i, x in enumerate(gt_names) if x in used_classes] + inds = np.array(inds, dtype=np.int64) + return inds + + +def init_dist_slurm(tcp_port, local_rank, backend='nccl'): + """ + modified from https://github.com/open-mmlab/mmdetection + Args: + tcp_port: + backend: + + Returns: + + """ + proc_id = int(os.environ['SLURM_PROCID']) + ntasks = int(os.environ['SLURM_NTASKS']) + node_list = os.environ['SLURM_NODELIST'] + num_gpus = torch.cuda.device_count() + torch.cuda.set_device(proc_id % num_gpus) + addr = subprocess.getoutput('scontrol show hostname {} | head -n1'.format(node_list)) + os.environ['MASTER_PORT'] = str(tcp_port) + os.environ['MASTER_ADDR'] = addr + os.environ['WORLD_SIZE'] = str(ntasks) + os.environ['RANK'] = str(proc_id) + dist.init_process_group(backend=backend) + + total_gpus = dist.get_world_size() + rank = dist.get_rank() + return total_gpus, rank + + +def init_dist_pytorch(tcp_port, local_rank, backend='nccl'): + if mp.get_start_method(allow_none=True) is None: + mp.set_start_method('spawn') + # os.environ['MASTER_PORT'] = str(tcp_port) + # os.environ['MASTER_ADDR'] = 'localhost' + num_gpus = torch.cuda.device_count() + torch.cuda.set_device(local_rank % num_gpus) + + dist.init_process_group( + backend=backend, + # init_method='tcp://127.0.0.1:%d' % tcp_port, + # rank=local_rank, + # world_size=num_gpus + ) + rank = dist.get_rank() + return num_gpus, rank + + +def get_dist_info(return_gpu_per_machine=False): + if torch.__version__ < '1.0': + initialized = dist._initialized + else: + if dist.is_available(): + initialized = dist.is_initialized() + else: + initialized = False + if initialized: + rank = dist.get_rank() + world_size = dist.get_world_size() + else: + rank = 0 + world_size = 1 + + if return_gpu_per_machine: + gpu_per_machine = torch.cuda.device_count() + return rank, world_size, gpu_per_machine + + return rank, world_size + + +def merge_results_dist(result_part, size, tmpdir): + rank, world_size = get_dist_info() + os.makedirs(tmpdir, exist_ok=True) + + dist.barrier() + pickle.dump(result_part, open(os.path.join(tmpdir, 'result_part_{}.pkl'.format(rank)), 'wb')) + dist.barrier() + + if rank != 0: + return None + + part_list = [] + for i in range(world_size): + part_file = os.path.join(tmpdir, 'result_part_{}.pkl'.format(i)) + part_list.append(pickle.load(open(part_file, 'rb'))) + + ordered_results = [] + for res in zip(*part_list): + ordered_results.extend(list(res)) + ordered_results = ordered_results[:size] + shutil.rmtree(tmpdir) + return ordered_results + + +def scatter_point_inds(indices, point_inds, shape): + ret = -1 * torch.ones(*shape, dtype=point_inds.dtype, device=point_inds.device) + ndim = indices.shape[-1] + flattened_indices = indices.view(-1, ndim) + slices = [flattened_indices[:, i] for i in range(ndim)] + ret[slices] = point_inds + return ret + + +def generate_voxel2pinds(sparse_tensor): + device = sparse_tensor.indices.device + batch_size = sparse_tensor.batch_size + spatial_shape = sparse_tensor.spatial_shape + indices = sparse_tensor.indices.long() + point_indices = torch.arange(indices.shape[0], device=device, dtype=torch.int32) + output_shape = [batch_size] + list(spatial_shape) + v2pinds_tensor = scatter_point_inds(indices, point_indices, output_shape) + return v2pinds_tensor + + +def sa_create(name, var): + x = SharedArray.create(name, var.shape, dtype=var.dtype) + x[...] = var[...] + x.flags.writeable = False + return x + + +class AverageMeter(object): + """Computes and stores the average and current value""" + def __init__(self): + self.reset() + + def reset(self): + self.val = 0 + self.avg = 0 + self.sum = 0 + self.count = 0 + + def update(self, val, n=1): + self.val = val + self.sum += val * n + self.count += n + self.avg = self.sum / self.count