LION/utils/data_helper.py

36 lines
1.6 KiB
Python
Raw Normal View History

2023-01-23 05:14:49 +00:00
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
def normalize_point_clouds(pcs, mode='shape_bbox'): # can be a property func
'''
copied from https://github.com/luost26/diffusion-point-cloud/blob/0bfd688379e78ac75fa75e6a2c5029e362496169/test_gen.py#L16
Args:
pcs: list of [N,3] or tensor in shape: B,N,3
'''
# logger.debug('Normalization mode: %s' % mode)
assert(type(pcs) == list or len(pcs.shape) ==
3), f'expect pcs to be list, get: {type(pcs)} or 3d tensor; '
output_list = []
for i in range(len(pcs)): # , desc='Normalize'):
pc = pcs[i]
pc = pc.detach().clone()
assert(mode == 'shape_bbox')
assert(len(pc.shape) == 2 and pc.shape[-1] in [3, 4,
6, 9]), f'expect get (N,3 or 6), get {pc.shape}'
pc_max, _ = pc.max(dim=0, keepdim=True) # (1, 3)
pc_min, _ = pc.min(dim=0, keepdim=True) # (1, 3)
pc_min = pc_min[:, :3]
pc_max = pc_max[:, :3]
shift = ((pc_min + pc_max) / 2).view(1, 3)
scale = (pc_max - pc_min).max().reshape(1, 1) / 2
pc[:, :3] = (pc[:, :3] - shift) / scale
# pcs[i] = pc
output_list.append(pc)
return output_list