import csv import copy import torch import einops import numpy as np from torch import nn import torch.nn.functional as F def get_activation_fn(activation_type): if activation_type not in ["relu", "gelu", "glu"]: raise RuntimeError(f"activation function currently support relu/gelu, not {activation_type}") return getattr(F, activation_type) def get_mlp_head(input_size, hidden_size, output_size, dropout=0): return nn.Sequential(*[ nn.Linear(input_size, hidden_size), nn.ReLU(), nn.LayerNorm(hidden_size, eps=1e-12), nn.Dropout(dropout), nn.Linear(hidden_size, output_size) ]) def layer_repeat(module, N, share_layer=False): if share_layer: return nn.ModuleList([module] * N) else: return nn.ModuleList([copy.deepcopy(module) for _ in range(N - 1)] + [module]) def calc_pairwise_locs(obj_centers, obj_whls, eps=1e-10, pairwise_rel_type='center', spatial_dist_norm=True, spatial_dim=5): if pairwise_rel_type == 'mlp': obj_locs = torch.cat([obj_centers, obj_whls], 2) pairwise_locs = torch.cat( [einops.repeat(obj_locs, 'b l d -> b l x d', x=obj_locs.size(1)), einops.repeat(obj_locs, 'b l d -> b x l d', x=obj_locs.size(1))], dim=3 ) return pairwise_locs pairwise_locs = einops.repeat(obj_centers, 'b l d -> b l 1 d') \ - einops.repeat(obj_centers, 'b l d -> b 1 l d') pairwise_dists = torch.sqrt(torch.sum(pairwise_locs ** 2, 3) + eps) # (b, l, l) if spatial_dist_norm: max_dists = torch.max(pairwise_dists.view(pairwise_dists.size(0), -1), dim=1)[0] norm_pairwise_dists = pairwise_dists / einops.repeat(max_dists, 'b -> b 1 1') else: norm_pairwise_dists = pairwise_dists if spatial_dim == 1: return norm_pairwise_dists.unsqueeze(3) pairwise_dists_2d = torch.sqrt(torch.sum(pairwise_locs[..., :2] ** 2, 3) + eps) if pairwise_rel_type == 'center': pairwise_locs = torch.stack( [norm_pairwise_dists, pairwise_locs[..., 2] / pairwise_dists, pairwise_dists_2d / pairwise_dists, pairwise_locs[..., 1] / pairwise_dists_2d, pairwise_locs[..., 0] / pairwise_dists_2d], dim=3 ) elif pairwise_rel_type == 'vertical_bottom': bottom_centers = torch.clone(obj_centers) bottom_centers[:, :, 2] -= obj_whls[:, :, 2] bottom_pairwise_locs = einops.repeat(bottom_centers, 'b l d -> b l 1 d') \ - einops.repeat(bottom_centers, 'b l d -> b 1 l d') bottom_pairwise_dists = torch.sqrt(torch.sum(bottom_pairwise_locs ** 2, 3) + eps) # (b, l, l) bottom_pairwise_dists_2d = torch.sqrt(torch.sum(bottom_pairwise_locs[..., :2] ** 2, 3) + eps) pairwise_locs = torch.stack( [norm_pairwise_dists, bottom_pairwise_locs[..., 2] / bottom_pairwise_dists, bottom_pairwise_dists_2d / bottom_pairwise_dists, pairwise_locs[..., 1] / pairwise_dists_2d, pairwise_locs[..., 0] / pairwise_dists_2d], dim=3 ) if spatial_dim == 4: pairwise_locs = pairwise_locs[..., 1:] return pairwise_locs 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) # centering obj_pcds[:, :, :3].sub_(obj_pcds[:, :, :3].mean(1, keepdim=True)) # normalization 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] # True as masked. return padded_sequence, mask else: return padded_sequence