def forward( self, xyz: torch.Tensor, features: Optional[torch.Tensor] ) -> Tuple[torch.Tensor, torch.Tensor]: r""" Parameters ---------- xyz : torch.Tensor (B, N, 3) tensor of the xyz coordinates of the features features : torch.Tensor (B, C, N) tensor of the descriptors of the the features Returns ------- new_xyz : torch.Tensor (B, npoint, 3) tensor of the new features' xyz new_features : torch.Tensor (B, \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors """ new_features_list = [] xyz_flipped = xyz.transpose(1, 2).contiguous() new_xyz = (pointnet2_utils.gather_operation( xyz_flipped, pointnet2_utils.furthest_point_sample( xyz, self.npoint)).transpose(1, 2).contiguous() if self.npoint is not None else None) if self.fuse == 'add': for i in range(len(self.groupers)): new_features = self.groupers[i]( xyz, new_xyz, features) # (B, C, npoint, nsample) new_features = self.mlps[i]( new_features) # (B, mlp[-1], npoint, nsample) new_features = F.max_pool2d(new_features, kernel_size=[ 1, new_features.size(3) ]) # (B, mlp[-1], npoint, 1) new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint) if i == 0: new_features_sum = new_features else: new_features_sum += new_features return new_xyz, new_features_sum for i in range(len(self.groupers)): new_features = self.groupers[i]( xyz, new_xyz, features) # (B, C, npoint, nsample) new_features = self.mlps[i]( new_features) # (B, mlp[-1], npoint, nsample) new_features = F.max_pool2d(new_features, kernel_size=[ 1, new_features.size(3) ]) # (B, mlp[-1], npoint, 1) new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint) new_features_list.append(new_features) return new_xyz, torch.cat(new_features_list, dim=1)
def pool(xyz, points, k, npoint): xyz_flipped = xyz.transpose(1, 2).contiguous() new_xyz = pointnet2_utils.gather_operation( xyz_flipped, pointnet2_utils.furthest_point_sample( xyz_flipped, npoint)).transpose(1, 2).contiguous() _, idx = knn_point(k, xyz, new_xyz) new_points = torch.max(pointnet2_utils.grouping_operation( points.permute(0, 2, 1).contiguous(), idx.int().permute(0, 2, 1).contiguous()).permute(0, 3, 2, 1), dim=2).values return new_xyz, new_points
def fps_subsample(pcd, n_points=2048): """ Args pcd: (b, 16384, 3) returns new_pcd: (b, n_points, 3) """ new_pcd = gather_operation( pcd.permute(0, 2, 1).contiguous(), furthest_point_sample(pcd, n_points)) new_pcd = new_pcd.permute(0, 2, 1).contiguous() return new_pcd
def forward(self, x, feature, num_pool): x = x.contiguous() xyz_flipped = x.transpose(1, 2).contiguous() x_sub = (pointnet2_utils.gather_operation( x, pointnet2_utils.furthest_point_sample(xyz_flipped, num_pool)).transpose( 1, 2).contiguous()) # sub_index = self.knn(x_sub, xyz_flipped).int() sub_index = pointnet2_utils.ball_query(0.2 * self.nlayer, self.k, xyz_flipped, x_sub) x = pointnet2_utils.grouping_operation(x, sub_index) x = torch.max(x, dim=-1)[0] # x = soft_pool2d(x, [1, x.shape[-1]]).squeeze(-1) feature = pointnet2_utils.grouping_operation(feature, sub_index) feature = self.mlp(feature) feature = torch.max(feature, dim=-1)[0] # feature = soft_pool2d(feature, [1, feature.shape[-1]]).squeeze(-1) return x, feature
def sample_and_group(xyz, points, npoint, nsample, radius, use_xyz=True): """ Args: xyz: Tensor, (B, 3, N) points: Tensor, (B, f, N) npoint: int nsample: int radius: float use_xyz: boolean Returns: new_xyz: Tensor, (B, 3, npoint) new_points: Tensor, (B, 3 | f+3 | f, npoint, nsample) idx_local: Tensor, (B, npoint, nsample) grouped_xyz: Tensor, (B, 3, npoint, nsample) """ xyz_flipped = xyz.permute(0, 2, 1).contiguous() # (B, N, 3) new_xyz = gather_operation(xyz, furthest_point_sample(xyz_flipped, npoint)) # (B, 3, npoint) idx = ball_query(radius, nsample, xyz_flipped, new_xyz.permute(0, 2, 1).contiguous()) # (B, npoint, nsample) grouped_xyz = grouping_operation(xyz, idx) # (B, 3, npoint, nsample) grouped_xyz -= new_xyz.unsqueeze(3).repeat(1, 1, 1, nsample) if points is not None: grouped_points = grouping_operation(points, idx) # (B, f, npoint, nsample) if use_xyz: new_points = torch.cat([grouped_xyz, grouped_points], 1) else: new_points = grouped_points else: new_points = grouped_xyz return new_xyz, new_points, idx, grouped_xyz
vis.add_geometry(pcd) ctr = vis.get_view_control() ctr.rotate(10.0, 0.0) vis.update_geometry() vis.poll_events() vis.update_renderer() time.sleep(1) vis.capture_screen_image('{}b.png'.format(i)) vis.remove_geometry(pcd) point = torch.from_numpy(pcd.points) x_sub = ( pointnet2_utils.gather_operation( xyz_flipped, pointnet2_utils.furthest_point_sample(x, num_pool) ).transpose(1, 2).contiguous() ) vis.update_geometry(pcd) vis.poll_events() vis.update_renderer() time.sleep(1) vis.capture_screen_image('{}b.png'.format(i)) vis.remove_geometry(pcd) # pcd = o3d.geometry.PointCloud() # pcd.points = o3d.utility.Vector3dVector(data) # o3d.visualization.draw_geometries([pcd], # zoom=1, # front=[1, 1, -1], # lookat=[0, 0, 0],
def forward(self, xyz, rgb, istrain=False): hs = [] #xyz_copy = xyz.clone() #rgb_copy = rgb.clone() batch_size, n_points, _ = xyz.shape part_length = n_points // self.npart last_point = -1 last_feature_dim = -1 #h = self.proj_in(rgb) h = rgb s2_count = 0 for i in range(self.num_layers): h_input = h.clone() xyz_input = xyz.clone() batch_size, n_points, feature_dim = h.shape ######## Build Graph ######### last_point = n_points ######### Dynamic Graph Conv ######### xyz = xyz.transpose(1, 2).contiguous() #print(h.shape) # batchsize x point_number x feature_dim h = h.transpose(1, 2).contiguous() for j in range(self.num_conv): index = self.num_conv * i + j ####### BN + ReLU ##### if self.pre_act == True: if self.norm == 'ln': h = h.transpose(1, 2).contiguous() h = self.bn[index](h) h = h.transpose(1, 2).contiguous() else: h = self.bn[index](h) h = F.leaky_relu(h, 0.2) ####### Graph Feature ########### if self.k == 1 and j == 0: h = h.unsqueeze(-1) else: if i == self.num_layers - 1: if self.cluster == 'xyz': h = get_graph_feature(xyz, h, k=self.k) elif self.cluster == 'xyzrgb' or self.cluster == 'allxyzrgb': h = get_graph_feature(torch.cat((xyz, h), 1), h, k=self.k) else: # Common Layers if self.cluster == 'allxyzrgb': h = get_graph_feature(torch.cat((xyz, h), 1), h, k=self.k) else: h = get_graph_feature(xyz, h, k=self.k) ####### Conv ########## if self.light == True and i > 0: #shuffle after the first layer h = channel_shuffle(h, 2) h = self.conv[index](h) else: h = self.conv[index](h) h = h.max(dim=-1, keepdim=False)[0] ####### BN + ReLU ##### if self.pre_act == False: if self.norm == 'ln': h = h.transpose(1, 2).contiguous() h = self.bn[index](h) h = h.transpose(1, 2).contiguous() else: h = self.bn[index](h) h = F.leaky_relu(h, 0.2) h = h.transpose(1, 2).contiguous() #print(h.shape) # batchsize x point_number x feature_dim batch_size, n_points, feature_dim = h.shape ######### Residual Before Downsampling############# if self.id_skip == 1: if istrain and self.drop_connect_rate > 0: h = drop_connect(h, p=self.drop_connect_rate, training=istrain) if feature_dim != last_feature_dim: h_input = self.conv_s2[s2_count](h_input) h = h_input + self.res_scale * h ######### PointNet++ MSG ######## if feature_dim != last_feature_dim: h = h.transpose(1, 2).contiguous() xyz, h = self.sa[s2_count](xyz_input, h) h = h.transpose(1, 2).contiguous() if self.id_skip == 2: h_input = pointnet2_utils.gather_operation( h_input.transpose(1, 2).contiguous(), pointnet2_utils.furthest_point_sample( xyz_input, h.shape[1])).transpose(1, 2).contiguous() else: xyz = xyz.transpose(1, 2).contiguous() ######### Residual After Downsampling (Paper) ############# if self.id_skip == 2: if istrain and self.drop_connect_rate > 0: h = drop_connect(h, p=self.drop_connect_rate, training=istrain) if feature_dim != last_feature_dim: h_input = self.conv_s2[s2_count](h_input) h = h_input + self.res_scale * h if feature_dim != last_feature_dim: s2_count += 1 last_feature_dim = feature_dim #print(xyz.shape, h.shape) if self.npart == 1: # Pooling h_max, _ = torch.max(h, 1) h_avg = torch.mean(h, 1) hs.append(h_max) hs.append(h_avg) h = torch.cat(hs, 1) h = self.embs[0](h) h = self.bn_embs[0](h) h = self.dropouts[0](h) h = self.proj_output(h) else: # Sort #batch_size, n_points, _ = h.shape #y_index = torch.argsort(xyz[:, :, 1],dim = 1).view(batch_size * n_points, -1) #h = h.view(batch_size * n_points, -1) #h = h[y_index, :].view(batch_size, n_points, -1) h = h.transpose(1, 2) # Part Pooling h = self.partpool(h) for i in range(self.npart): part_h = h[:, :, i] part_h = self.embs[i](part_h) part_h = self.bn_embs[i](part_h) part_h = self.dropouts[i](part_h) part_h = self.proj_outputs[i](part_h) hs.append(part_h) h = hs return h
def forward(self, xyz, rgb, istrain=False): hs = [] #xyz_copy = xyz.clone() #rgb_copy = rgb.clone() batch_size, n_points, _ = xyz.shape part_length = n_points//self.npart last_point = -1 #h = self.proj_in(rgb) h = rgb s2_count = 0 for i in range(self.num_layers): h_input = h.clone() xyz_input = xyz.clone() batch_size, n_points, _ = h.shape if self.k>1: if i == self.num_layers-1: if self.cluster == 'xyz': g = self.nng(xyz, istrain = istrain and self.graph_jitter) elif self.cluster == 'rgb': g = self.nng(h, istrain=istrain and self.graph_jitter) elif self.cluster == 'xyzrgb': g = self.nng( torch.cat((xyz,h), 2), istrain=istrain and self.graph_jitter) elif i==0 or n_points != last_point: g = self.nng(xyz, istrain=istrain and self.graph_jitter) last_point = n_points h = h.view(batch_size * n_points, -1) if self.k==1: h = self.conv[i](h) elif self.conv_type == 'GatedGCN': h = self.conv[i](g, h, g.edata['feat'], snorm_n = 1/g.number_of_nodes() , snorm_e = 1/g.number_of_edges()) else: h = self.conv[i](g, h) h = F.leaky_relu(h, 0.2) h = h.view(batch_size, n_points, -1) h = h.transpose(1, 2).contiguous() xyz, h = self.sa[i](xyz_input, h) h = h.transpose(1, 2).contiguous() #h = self.conv_s1[i](h) if self.id_skip and h.shape[1] <= self.init_points//4: # We could use identity mapping Here or add connect drop if istrain and self.drop_connect_rate>0: h = drop_connect(h, p=self.drop_connect_rate, training=istrain) if h.shape[1] == n_points: h = h_input + self.res_scale * h # Here I borrow the idea from Inception-ResNet-v2 elif h.shape[1] == n_points//2: h_input_s2 = pointnet2_utils.gather_operation( h_input.transpose(1, 2).contiguous(), pointnet2_utils.furthest_point_sample(xyz_input, h.shape[1] ) ).transpose(1, 2).contiguous() h = self.conv_s2[s2_count](h_input_s2) + self.res_scale * h s2_count +=1 if self.npart==1: # Pooling h_max, _ = torch.max(h, 1) h_avg = torch.mean(h, 1) hs.append(h_max) hs.append(h_avg) h = torch.cat(hs, 1) h = self.embs[0](h) h = self.bn_embs[0](h) h = self.dropouts[0](h) h = self.proj_output(h) else: # Sort batch_size, n_points, _ = h.shape y_index = torch.argsort(xyz[:, :, 1],dim = 1).view(batch_size * n_points, -1) h = h.view(batch_size * n_points, -1) h = h[y_index, :].view(batch_size, n_points, -1) h = h.transpose(1, 2) # Part Pooling h = self.partpool(h) for i in range(self.npart): part_h = h[:,:,i] part_h = self.embs[i](part_h) part_h = self.bn_embs[i](part_h) part_h = self.dropouts[i](part_h) part_h = self.proj_outputs[i](part_h) hs.append(part_h) h = hs return h