|
import csv |
|
import numpy as np |
|
import torch |
|
|
|
|
|
def convert_pc_to_box(obj_pc): |
|
xmin = np.min(obj_pc[:,0]) |
|
ymin = np.min(obj_pc[:,1]) |
|
zmin = np.min(obj_pc[:,2]) |
|
xmax = np.max(obj_pc[:,0]) |
|
ymax = np.max(obj_pc[:,1]) |
|
zmax = np.max(obj_pc[:,2]) |
|
center = [(xmin+xmax)/2, (ymin+ymax)/2, (zmin+zmax)/2] |
|
box_size = [xmax-xmin, ymax-ymin, zmax-zmin] |
|
return center, box_size |
|
|
|
class LabelConverter(object): |
|
def __init__(self, file_path): |
|
self.raw_name_to_id = {} |
|
self.nyu40id_to_id = {} |
|
self.nyu40_name_to_id = {} |
|
self.scannet_name_to_scannet_id = {'cabinet':0, 'bed':1, 'chair':2, 'sofa':3, 'table':4, |
|
'door':5, 'window':6,'bookshelf':7,'picture':8, 'counter':9, 'desk':10, 'curtain':11, |
|
'refrigerator':12, 'shower curtain':13, 'toilet':14, 'sink':15, 'bathtub':16, 'others':17} |
|
self.id_to_scannetid = {} |
|
self.scannet_raw_id_to_raw_name = {} |
|
self.raw_name_to_scannet_raw_id = {} |
|
|
|
with open(file_path, encoding='utf-8') as fd: |
|
rd = list(csv.reader(fd, delimiter="\t", quotechar='"')) |
|
for i in range(1, len(rd)): |
|
raw_id = i - 1 |
|
scannet_raw_id = int(rd[i][0]) |
|
raw_name = rd[i][1] |
|
nyu40_id = int(rd[i][4]) |
|
nyu40_name = rd[i][7] |
|
self.raw_name_to_id[raw_name] = raw_id |
|
self.scannet_raw_id_to_raw_name[scannet_raw_id] = raw_name |
|
self.raw_name_to_scannet_raw_id[raw_name] = scannet_raw_id |
|
self.nyu40id_to_id[nyu40_id] = raw_id |
|
self.nyu40_name_to_id[nyu40_name] = raw_id |
|
if nyu40_name not in self.scannet_name_to_scannet_id: |
|
self.id_to_scannetid[raw_id] = self.scannet_name_to_scannet_id['others'] |
|
else: |
|
self.id_to_scannetid[raw_id] = self.scannet_name_to_scannet_id[nyu40_name] |
|
|
|
def build_rotate_mat(split, rot_aug=True, rand_angle='axis'): |
|
if rand_angle == 'random': |
|
theta = np.random.rand() * np.pi * 2 |
|
else: |
|
ROTATE_ANGLES = [0, np.pi/2, np.pi, np.pi*3/2] |
|
theta_idx = np.random.randint(len(ROTATE_ANGLES)) |
|
theta = ROTATE_ANGLES[theta_idx] |
|
if (theta is not None) and (theta != 0) and (split == 'train') and rot_aug: |
|
rot_matrix = np.array([ |
|
[np.cos(theta), -np.sin(theta), 0], |
|
[np.sin(theta), np.cos(theta), 0], |
|
[0, 0, 1] |
|
], dtype=np.float32) |
|
else: |
|
rot_matrix = None |
|
return rot_matrix |
|
|
|
def obj_processing_post(obj_pcds, rot_aug=True): |
|
obj_pcds = torch.from_numpy(obj_pcds) |
|
rot_matrix = build_rotate_mat('val', rot_aug) |
|
if rot_matrix is not None: |
|
rot_matrix = torch.from_numpy(rot_matrix.transpose()) |
|
obj_pcds[:, :, :3] @= rot_matrix |
|
|
|
xyz = obj_pcds[:, :, :3] |
|
center = xyz.mean(1) |
|
xyz_min = xyz.min(1).values |
|
xyz_max = xyz.max(1).values |
|
box_center = (xyz_min + xyz_max) / 2 |
|
size = xyz_max - xyz_min |
|
obj_locs = torch.cat([center, size], dim=1) |
|
obj_boxes = torch.cat([box_center, size], dim=1) |
|
|
|
|
|
obj_pcds[:, :, :3].sub_(obj_pcds[:, :, :3].mean(1, keepdim=True)) |
|
|
|
|
|
max_dist = (obj_pcds[:, :, :3]**2).sum(2).sqrt().max(1).values |
|
max_dist.clamp_(min=1e-6) |
|
obj_pcds[:, :, :3].div_(max_dist[:, None, None]) |
|
|
|
return obj_pcds, obj_locs, obj_boxes, rot_matrix |
|
|
|
|
|
def pad_sequence(sequence_list, max_len=None, pad=0, return_mask=False): |
|
lens = [x.shape[0] for x in sequence_list] |
|
if max_len is None: |
|
max_len = max(lens) |
|
|
|
shape = list(sequence_list[0].shape) |
|
shape[0] = max_len |
|
shape = [len(sequence_list)] + shape |
|
dtype = sequence_list[0].dtype |
|
device = sequence_list[0].device |
|
padded_sequence = torch.ones(shape, dtype=dtype, device=device) * pad |
|
for i, tensor in enumerate(sequence_list): |
|
padded_sequence[i, :tensor.shape[0]] = tensor |
|
padded_sequence = padded_sequence.to(dtype) |
|
|
|
if return_mask: |
|
mask = torch.arange(max_len).to(device)[None, :] >= torch.LongTensor(lens).to(device)[:, None] |
|
return padded_sequence, mask |
|
else: |
|
return padded_sequence |
