def execute(self, x, cls_label): batch_size, _, N = x.size() x = self.relu(self.bn1(self.conv1(x))) # B, D, N x = self.relu(self.bn2(self.conv2(x))) x1 = self.sa1(x) x2 = self.sa2(x1) x3 = self.sa3(x2) x4 = self.sa4(x3) x = concat((x1, x2, x3, x4), dim=1) x = self.conv_fuse(x) x_max = jt.max(x, 2) x_avg = jt.mean(x, 2) x_max_feature = x_max.view(batch_size, -1).unsqueeze(-1).repeat(1, 1, N) x_avg_feature = x_avg.view(batch_size, -1).unsqueeze(-1).repeat(1, 1, N) cls_label_one_hot = cls_label.view(batch_size, 16, 1) cls_label_feature = self.label_conv(cls_label_one_hot).repeat(1, 1, N) x_global_feature = concat( (x_max_feature, x_avg_feature, cls_label_feature), 1) # 1024 + 64 x = concat((x, x_global_feature), 1) # 1024 * 3 + 64 x = self.relu(self.bns1(self.convs1(x))) x = self.dp1(x) x = self.relu(self.bns2(self.convs2(x))) x = self.convs3(x) return x
def execute(self, point_cloud, label): B, D, N = point_cloud.size() trans = self.stn(point_cloud) point_cloud = point_cloud.transpose(0, 2, 1) point_cloud = nn.bmm(point_cloud, trans) point_cloud = point_cloud.transpose(0, 2, 1) out1 = self.relu(self.bn1(self.conv1(point_cloud))) out2 = self.relu(self.bn2(self.conv2(out1))) out3 = self.relu(self.bn3(self.conv3(out2))) trans_feat = self.fstn(out3) x = out3.transpose(0, 2, 1) net_transformed = nn.bmm(x, trans_feat) net_transformed = net_transformed.transpose(0, 2, 1) out4 = self.relu(self.bn4(self.conv4(net_transformed))) out5 = self.bn5(self.conv5(out4)) out_max = jt.argmax(out5, 2, keepdims=True)[1] out_max = out_max.view(-1, 2048) out_max = concat((out_max, label), 1) expand = out_max.view(-1, 2048 + 16, 1).repeat(1, 1, N) concat_feature = concat([expand, out1, out2, out3, out4, out5], 1) net = self.relu(self.bns1(self.convs1(concat_feature))) net = self.relu(self.bns2(self.convs2(net))) net = self.relu(self.bns3(self.convs3(net))) net = self.convs4(net) return net
def execute(self, x): imsize = x.shape _, _, _, x = self.backbone(x) b, c, h, w = x.shape x_k = self.conv_0(x) x1 = self.layer5a(x_k).reshape(b, c, -1) x2 = self.layer5b(x_k).reshape(b, c, -1) x3 = self.layer5c(x_k).reshape(b, c, -1) x4 = self.layer5d(x_k).reshape(b, c, -1) x_k = concat([x1, x2, x3, x4], 2).transpose(0, 2, 1) # b 110 c x_q = self.conv_1(x) x_q = x_q.reshape(b, c, -1) x_attention = nn.bmm(x_k, x_q) # b 110 N x_v = self.conv_2(x) x1 = self.layer5a(x_v).reshape(b, c, -1) x2 = self.layer5b(x_v).reshape(b, c, -1) x3 = self.layer5c(x_v).reshape(b, c, -1) x4 = self.layer5d(x_v).reshape(b, c, -1) x_v = concat([x1, x2, x3, x4], 2) # b c 110 x = nn.bmm(x_v, x_attention).reshape(b, c, h, w) x = self.final_conv(x) x = nn.resize(x, size=(imsize[2], imsize[3]), mode='bilinear', align_corners=True) return x
def sample_and_group(num_points, num_neighbors, points, features=None): """ sample `num-points` from points via furthest point sampling algorithm and consider them sampled points as centroids to construct KNN graphs on points where k=`num-neighbors` """ B, _, C = points.shape S = num_points fps_idx = farthest_point_sample(points, num_points) # _, fps_idx = sampler(points) # [B, npoint] sampled_points = index_points(points, fps_idx) point_knn_idx = knn_points(num_neighbors, points, sampled_points) grouped_points = index_points( points, point_knn_idx) # [B, num-point, num-neighbos, C] if features is None: reshaped_points = sampled_points.view(B, S, 1, C) grouped_points_relative = grouped_points - reshaped_points return sampled_points, concat([ grouped_points_relative, reshaped_points.repeat(1, 1, num_neighbors, 1) ], dim=-1) sampled_features = index_points(features, fps_idx) grouped_features = index_points(features, point_knn_idx) reshaped_features = sampled_features.view(B, S, 1, -1) grouped_feature_relative = grouped_features - reshaped_features return sampled_points, concat([ grouped_feature_relative, reshaped_features.repeat(1, 1, num_neighbors, 1) ], dim=-1)
def execute(self, x): #print (x.shape) # x = x.transpose(0, 2, 1) # b, n, c -> b, c, n batch_size = x.shape[0] x = get_graph_feature(x, knn=self.knn, k=self.k) x = self.conv1(x) x1 = x.max(dim=-1, keepdims=False) x = get_graph_feature(x1, knn=self.knn, k=self.k) x = self.conv2(x) x2 = x.max(dim=-1, keepdims=False) x = get_graph_feature(x2, knn=self.knn, k=self.k) x = self.conv3(x) x3 = x.max(dim=-1, keepdims=False) x = get_graph_feature(x3, knn=self.knn, k=self.k) x = self.conv4(x) x4 = x.max(dim=-1, keepdims=False) x = concat((x1, x2, x3, x4), dim=1) x = self.conv5(x) # ############ this line x1 = x.max(dim=2).reshape(batch_size, -1) x2 = x.mean(dim=2).reshape(batch_size, -1) x = concat((x1, x2), 1) x = nn.leaky_relu(self.bn6(self.linear1(x)), scale=0.2) x = self.dp1(x) x = nn.leaky_relu(self.bn7(self.linear2(x)), scale=0.2) x = self.dp2(x) x = self.linear3(x) return x
def execute(self, pcd, prev_s): """ Args: pcd: b, n, 3 prev_s: b, c, n """ b, n, _ = pcd.shape pcd_bcn = pcd.transpose(0, 2, 1) l0_xyz = pcd l0_points = pcd_bcn if self.if_noise: noise_points = init.gauss([b, 3, n], 'float', mean=0.0, std=self.noise_stdv) l0_points = jittor.concat([l0_points, noise_points], 1) l1_xyz, l1_points = self.sa_module_1(l0_xyz, l0_points) # b, 512, 128 (bnc) l2_xyz, l2_points = self.sa_module_2(l1_xyz, l1_points) l3_xyz, l3_points = self.sa_module_3(l2_xyz, l2_points) l2_points = self.fp_module_3(l2_xyz, l3_xyz, l2_points, l3_points) l2_points, prev_s['l2'] = self.unit_3(l2_points, prev_s['l2']) l1_points = self.fp_module_2(l1_xyz, l2_xyz, l1_points, l2_points) l1_points, prev_s['l1'] = self.unit_2(l1_points, prev_s['l1']) l0_points = self.fp_module_1(l0_xyz, l1_xyz, concat([pcd_bcn, pcd_bcn], dim=1), l1_points) l0_points, prev_s['l0'] = self.unit_1(l0_points, prev_s['l0']) # (B, 128, 2048) noise = init.gauss([b, 32, n], 'float', mean=0.0, std=1.0) feat = concat([l0_points, noise], dim=1) delta_xyz = self.tanh(self.mlp_conv(feat)) * 1.0 / 10 ** (self.step - 1) point_cloud = (pcd_bcn + delta_xyz).transpose(0, 2, 1) return point_cloud, delta_xyz
def execute(self, x): residual = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) spx = out outs = [] for i in range(self.nums): if i==0 or self.stype=='stage': sp = spx[:, i*self.width: (i+1)*self.width] else: sp = sp + spx[:, i*self.width: (i+1)*self.width] sp = self.convs[i](sp) sp = self.relu(self.bns[i](sp)) outs.append(sp) if self.stype=='normal' or self.stride==1: outs.append(spx[:, self.nums*self.width: (self.nums+1)*self.width]) elif self.stype=='stage': outs.append(self.pool(spx[:, self.nums*self.width: (self.nums+1)*self.width])) out = concat(outs, 1) out = self.conv3(out) out = self.bn3(out) if self.downsample is not None: residual = self.downsample(x) out += residual out = self.relu(out) return out
def execute(self, x): xyz = x.permute(0, 2, 1) batch_size, _, _ = x.size() x = self.relu(self.bn1(self.conv1(x))) # B, D, N x = self.relu(self.bn2(self.conv2(x))) # B, D, N x = x.permute(0, 2, 1) new_xyz, new_feature = sample_and_group(npoint=512, nsample=32, xyz=xyz, points=x) feature_0 = self.gather_local_0(new_feature) feature = feature_0.permute(0, 2, 1) new_xyz, new_feature = sample_and_group(npoint=256, nsample=32, xyz=new_xyz, points=feature) feature_1 = self.gather_local_1(new_feature) x = self.pt_last(feature_1) x = concat([x, feature_1], dim=1) x = self.conv_fuse(x) x = jt.max(x, 2) x = x.view(batch_size, -1) x = self.relu(self.bn6(self.linear1(x))) x = self.dp1(x) x = self.relu(self.bn7(self.linear2(x))) x = self.dp2(x) x = self.linear3(x) return x
def sample_and_group(npoint, nsample, xyz, points): B, N, C = xyz.shape S = npoint # xyz = xyz.contiguous() sampler = FurthestPointSampler(npoint) _, fps_idx = sampler(xyz) # [B, npoint] # print ('fps size=', fps_idx.size()) # fps_idx = sampler(xyz).long() # [B, npoint] new_xyz = index_points(xyz, fps_idx) new_points = index_points(points, fps_idx) # new_xyz = xyz[:] # new_points = points[:] idx = knn_point(nsample, xyz, new_xyz) #idx = query_ball_point(radius, nsample, xyz, new_xyz) grouped_xyz = index_points(xyz, idx) # [B, npoint, nsample, C] grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C) grouped_points = index_points(points, idx) grouped_points_norm = grouped_points - new_points.view(B, S, 1, -1) new_points = concat([ grouped_points_norm, new_points.view(B, S, 1, -1).repeat(1, 1, nsample, 1) ], dim=-1) return new_xyz, new_points
def execute(self, cur_x, prev_s): """ Args: cur_x: Tensor, (B, in_channel, N) prev_s: Tensor, (B, in_channel, N) Returns: h: Tensor, (B, in_channel, N) h: Tensor, (B, in_channel, N) """ if self.step == 1: return cur_x, cur_x z = self.conv_z(concat([cur_x, prev_s], dim=1)) r = self.conv_r(concat([cur_x, prev_s], dim=1)) h_hat = self.conv_h(concat([cur_x, r * prev_s], dim=1)) h = (1 - z) * cur_x + z * h_hat return h, h
def execute(self, x, l): batch_size = x.size(0) num_points = x.size(2) # print (x.size()) # x0 = get_graph_feature(x, k=self.k) # (batch_size, 3, num_points) -> (batch_size, 3*2, num_points, k) x = get_graph_feature(x, knn=self.knn, k=self.k) x = self.conv1(x) # (batch_size, 3*2, num_points, k) -> (batch_size, 64, num_points, k) x = self.conv2(x) # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points, k) x1 = x.max(dim=-1, keepdims=False) # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points) x = get_graph_feature(x1, knn=self.knn, k=self.k) x = self.conv3(x) # (batch_size, 64*2, num_points, k) -> (batch_size, 64, num_points, k) x = self.conv4(x) # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points, k) x2 = x.max(dim=-1, keepdims=False) # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points) x = get_graph_feature(x2, knn=self.knn, k=self.k) x = self.conv5(x) # (batch_size, 64*2, num_points, k) -> (batch_size, 64, num_points, k) x3 = x.max(dim=-1, keepdims=False) # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points) x = concat((x1, x2, x3), dim=1) # (batch_size, 64*3, num_points) x = self.conv6(x) # (batch_size, 64*3, num_points) -> (batch_size, emb_dims, num_points) x = x.max(dim=-1, keepdims=True) # (batch_size, emb_dims, num_points) -> (batch_size, emb_dims, 1) l = l.view(batch_size, -1, 1) # (batch_size, num_categoties, 1) l = self.conv7(l) # (batch_size, num_categoties, 1) -> (batch_size, 64, 1) x = concat((x, l), dim=1) # (batch_size, 1088, 1) x = x.repeat(1, 1, num_points) # (batch_size, 1088, num_points) x = concat((x, x1, x2, x3), dim=1) # (batch_size, 1088+64*3, num_points) x = self.conv8(x) # (batch_size, 1088+64*3, num_points) -> (batch_size, 256, num_points) x = self.dp1(x) x = self.conv9(x) # (batch_size, 256, num_points) -> (batch_size, 256, num_points) x = self.dp2(x) x = self.conv10(x) # (batch_size, 256, num_points) -> (batch_size, 128, num_points) x = self.conv11(x) # (batch_size, 256, num_points) -> (batch_size, seg_num_all, num_points) return x
def execute(self, x): x1 = self.aspp1(x) x2 = self.aspp2(x) x3 = self.aspp3(x) x4 = self.aspp4(x) x5 = self.global_avg_pool(x) x5 = x5.broadcast((1,1,x4.shape[2],x4.shape[3])) x = concat((x1, x2, x3, x4, x5), dim=1) x = self.conv1(x) x = self.bn1(x) x = self.relu(x) return x
def execute(self, x, low_level_feat): low_level_feat = self.conv1(low_level_feat) low_level_feat = self.bn1(low_level_feat) low_level_feat = self.relu(low_level_feat) #print (low_level_feat.shape) x = nn.resize(x, size=(low_level_feat.shape[2], low_level_feat.shape[3]), mode='bilinear') x = concat((x, low_level_feat), dim=1) x = self.last_conv(x) return x
def execute(self, x): size = x.shape _, _, _, x = self.backbone(x) x = self.final_conv( concat([ x, self.layer5a(x), self.layer5b(x), self.layer5c(x), self.layer5d(x), ], 1)) return nn.resize(x, size=(size[2], size[3]), mode='bilinear')
def execute(self, xyz1, xyz2, points1, points2): """ Input: xyz1: input points position data, [B, C, N] xyz2: sampled input points position data, [B, C, S] points1: input points data, [B, N, D] points2: input points data, [B, S, D] Return: new_points: upsampled points data, [B, N, D'] """ # xyz1 = xyz1.permute(0, 2, 1) # xyz2 = xyz2.permute(0, 2, 1) # points2 = points2.permute(0, 2, 1) B, N, C = xyz1.shape _, S, _ = xyz2.shape points2 = points2.transpose(0, 2, 1) # b, n, c if S == 1: interpolated_points = points2.repeat(1, N, 1) else: # dists = square_distance(xyz1, xyz2) # idx, dists = jt.argsort(dists, dim=-1) # dists, idx = dists[:, :, :3], idx[:, :, :3] # [B, N, 3] _, idx = three_nn(xyz1, xyz2) dists = ((xyz1.unsqueeze(-2) - xyz2['i0', idx]) ** 2).sum(-1) # dists = jt.ones([B, N, 3]) # idx = jt.index([B, N, 3], 2) dists = jt.clamp(dists, min_v=1e-10) dist_recip = 1.0 / dists norm = jt.sum(dist_recip, dim=2, keepdims=True) weight = dist_recip / norm interpolated_points = jt.sum(index_points(points2, idx) * weight.view(B, N, 3, 1), dim=2) interpolated_points = interpolated_points.transpose(0, 2, 1) # b, c, n if points1 is not None: # points1 = points1.permute(0, 2, 1) new_points = concat([interpolated_points, points1], dim=1) else: new_points = interpolated_points # new_points = new_points.permute(0, 2, 1) # b, c, n # l = len(self.mlp_convs) # for i, conv in self.mlp_convs.layers.items(): # conv = self.mlp_convs[i] # bn = self.mlp_bns[i] # new_points = self.relu(bn(conv(new_points))) new_points = self.mlp(new_points) return new_points # .permute(0, 2, 1) # b, n, c
def execute(self, features, points=None): # batch-size, num-points, num-dims -> batch-size, num-dims, num-points batch_size = features.size(0) if points is not None: # add position embedding # points: batch-size, num-points, num-dims -> batch-size, num-dims, num-points points = self.pos_conv(points) # end # features: batch-size, num-points, num-dims -> batch-size, num-dims, num-points features = features.permute(0, 2, 1) x = self.conv(features) # B, D, N x1 = self.oa1(x, points) x2 = self.oa2(x1, points) x3 = self.oa3(x2, points) x4 = self.oa4(x3, points) return self.fuse_conv(concat((features, x1, x2, x3, x4), dim=1))
def execute(self, context, feature): batch_size, c, h, w = feature.shape origin_feature = feature feature = feature.reshape(batch_size, c, -1).transpose(0, 2, 1) # b, h*w, c context = context.reshape(batch_size, context.shape[1], -1) # b, n_cls, h*w attention = self.softmax(context) ocr_context = nn.bmm(attention, feature).transpose(0, 2, 1) # b, c, n_cls relation = nn.bmm(feature, ocr_context).transpose(0, 2, 1) # b, n_cls, h*w attention = self.softmax(relation) #b , n_cls, h*w result = nn.bmm(ocr_context, attention).reshape(batch_size, c, h, w) result = self.conv_1x1(result) result = concat([result, origin_feature], dim=1) result = self.last_conv(result) return result
def execute(self, x_l, x_h): # (N, P, C_out) """ Given a point cloud, and its corresponding features, return a new set of randomly-sampled representative points with features projected from the point cloud. :param x: (pts, fts) where - pts: Regional point cloud such that fts[:,p_idx,:] is the feature associated with pts[:,p_idx,:]. - fts: Regional features such that pts[:,p_idx,:] is the feature associated with fts[:,p_idx,:]. :return: Randomly subsampled points and their features. """ pts_l, fts_l = x_l pts_h, fts_h = x_h rep_pts_fts = self.pointcnn((pts_h, pts_l, fts_l)) concat_feature = concat ((rep_pts_fts, fts_h), dim=2) rep_pts_fts = self.conv_fuse(concat_feature) return pts_h, rep_pts_fts
def execute(self, xyz1, xyz2, points1, points2): """ Input: xyz1: input points position data, [B, C, N] xyz2: sampled input points position data, [B, C, S] points1: input points data, [B, D, N] points2: input points data, [B, D, S] Return: new_points: upsampled points data, [B, D', N] """ # xyz1 = xyz1.permute(0, 2, 1) # xyz2 = xyz2.permute(0, 2, 1) # points2 = points2.permute(0, 2, 1) B, N, C = xyz1.shape _, S, _ = xyz2.shape if S == 1: interpolated_points = points2.repeat(1, N, 1) else: dists = square_distance(xyz1, xyz2) idx, dists = jt.argsort(dists, dim=-1) dists, idx = dists[:, :, :3], idx[:, :, :3] # [B, N, 3] dist_recip = 1.0 / (dists + 1e-8) norm = jt.sum(dist_recip, dim=2, keepdims=True) weight = dist_recip / norm interpolated_points = jt.sum(index_points(points2, idx) * weight.view(B, N, 3, 1), dim=2) if points1 is not None: # points1 = points1.permute(0, 2, 1) new_points = concat([points1, interpolated_points], dim=-1) else: new_points = interpolated_points new_points = new_points.permute(0, 2, 1) # l = len(self.mlp_convs) for i, conv in self.mlp_convs.layers.items(): # conv = self.mlp_convs[i] bn = self.mlp_bns[i] new_points = self.relu(bn(conv(new_points))) return new_points.permute(0, 2, 1)
def sample_and_group_with_density_scale(npoint, nsample, xyz, points, density_scale=None): """ Input: npoint: nsample: xyz: input points position data, [B, N, C] points: input points data, [B, N, D] Return: new_xyz: sampled points position data, [B, 1, C] new_points: sampled points data, [B, 1, N, C+D] """ B, N, C = xyz.shape S = npoint fps_idx = farthest_point_sample(xyz, npoint) # [B, npoint, C] # jt.sync_all(True) # print ('11111111111111111') new_xyz = index_points(xyz, fps_idx) # jt.sync_all(True) # print ('2222222222222222222') idx = knn_points(nsample, xyz, new_xyz) # jt.sync_all(True) # print ('333333333333333333') grouped_xyz = index_points(xyz, idx) # [B, npoint, nsample, C] grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C) if points is not None: grouped_points = index_points(points, idx) new_points = concat([grouped_xyz_norm, grouped_points], dim=-1) # [B, npoint, nsample, C+D] else: new_points = grouped_xyz_norm # jt.sync_all(True) # print ('44444444444444444444444') if density_scale is None: return new_xyz, new_points, grouped_xyz_norm, idx else: grouped_density = index_points(density_scale, idx) return new_xyz, new_points, grouped_xyz_norm, idx, grouped_density
def execute(self, x): # # b, 3, npoint, nsample # conv2d 3 -> 128 channels 1, 1 # b * npoint, c, nsample # permute reshape batch_size, _, N = x.size() x = self.relu(self.bn1(self.conv1(x))) # B, D, N x = self.relu(self.bn2(self.conv2(x))) x1 = self.sa1(x) x2 = self.sa2(x1) x3 = self.sa3(x2) x4 = self.sa4(x3) x = concat((x1, x2, x3, x4), dim=1) return x
def execute(self, x, xyz): # # b, 3, npoint, nsample # conv2d 3 -> 128 channels 1, 1 # b * npoint, c, nsample # permute reshape batch_size, _, N = x.size() # add position embedding xyz = xyz.permute(0, 2, 1) xyz = self.pos_xyz(xyz) # end x = self.relu(self.bn1(self.conv1(x))) # B, D, N x1 = self.sa1(x, xyz) x2 = self.sa2(x1, xyz) x3 = self.sa3(x2, xyz) x4 = self.sa4(x3, xyz) x = concat((x1, x2, x3, x4), dim=1) return x
def execute(self, xyz, feature, cls_label): # for module in self.pointnet_modules: # xyz, feature = module(xyz, feature) B, N, _ = xyz.shape l1_xyz, l1_feature = self.pointnet_modules[0](xyz, feature) l2_xyz, l2_feature = self.pointnet_modules[1](l1_xyz, l1_feature) l3_xyz, l3_feature = self.pointnet_modules[2](l2_xyz, l2_feature) # print ('before interpolate shape') # print (l2_xyz.shape, l2_feature.shape, l3_xyz.shape, l3_feature.shape) l2_feature = self.fp3(l2_xyz, l3_xyz, l2_feature, l3_feature) l1_feature = self.fp2(l1_xyz, l2_xyz, l1_feature, l2_feature) cls_label_one_hot = cls_label.view(B, 16, 1).repeat(1, 1, N).permute(0, 2, 1) # print ('before concat size ') # print (cls_label_one_hot.size(),xyz.size(),feature.size()) feature = self.fp1(xyz, l1_xyz, concat([cls_label_one_hot, xyz, feature], 2), l1_feature) feature = feature.permute(0, 2, 1) # print (feature.shape) return self.fc_layer(feature)
def get_graph_feature(x, knn=None, k=None, idx=None): batch_size = x.shape[0] num_points = x.shape[2] x = x.reshape(batch_size, -1, num_points) if idx is None: idx = knn(x,x) # (batch_size, num_points, k) idx = idx.permute(0, 2, 1) idx_base = jt.array(np.arange(0, batch_size)).reshape(-1, 1, 1)*num_points idx = idx + idx_base idx = idx.reshape(-1) _, num_dims, _ = x.shape x = x.transpose(0, 2, 1) # (batch_size, num_points, num_dims) -> (batch_size*num_points, num_dims) # batch_size * num_points * k + range(0, batch_size*num_points) feature = x.reshape(batch_size*num_points, -1)[idx, :] feature = feature.reshape(batch_size, num_points, k, num_dims) x = x.reshape(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1) feature = concat((feature-x, x), dim=3).transpose(0, 3, 1, 2) return feature
def sample_and_group_all(xyz, points, density_scale=None): """ Input: xyz: input points position data, [B, N, C] points: input points data, [B, N, D] Return: new_xyz: sampled points position data, [B, 1, C] new_points: sampled points data, [B, 1, N, C+D] """ #device = xyz.device B, N, C = xyz.shape #new_xyz = torch.zeros(B, 1, C).to(device) new_xyz = xyz.mean(dim=1, keepdims=True) grouped_xyz = xyz.reshape(B, 1, N, C) - new_xyz.reshape(B, 1, 1, C) if points is not None: new_points = concat([grouped_xyz, points.view(B, 1, N, -1)], dim=3) else: new_points = grouped_xyz if density_scale is None: return new_xyz, new_points, grouped_xyz else: grouped_density = density_scale.reshape(B, 1, N, 1) return new_xyz, new_points, grouped_xyz, grouped_density
def execute(self, x): """ inputs : x : input feature maps( B X C X H X W) returns : out : attention value + input feature attention: B X (HxW) X (HxW) """ m_batchsize, C, height, width = x.size() proj_query = self.query_conv(x).reshape(m_batchsize, -1, width * height).transpose( 0, 2, 1) proj_key = self.key_conv(x).reshape(m_batchsize, -1, width * height) energy = nn.bmm(proj_query, proj_key) attention = self.softmax(energy) proj_value = self.value_conv(x).reshape(m_batchsize, -1, width * height) out = nn.bmm(proj_value, attention.transpose(0, 2, 1)) out = out.reshape(m_batchsize, C, height, width) out = self.scale_conv(out) out = concat([out, x], 1) return out
def execute(self, x): batch_size, _, N = x.size() # print (x.size()) x = self.relu(self.bn1(self.conv1(x))) # B, D, N x = self.relu(self.bn2(self.conv2(x))) x1 = self.sa1(x) x2 = self.sa2(x1) x3 = self.sa3(x2) x4 = self.sa4(x3) x = concat((x1, x2, x3, x4), dim=1) x = self.conv_fuse(x) # x = F.adaptive_max_pool1d(x, 1).view(batch_size, -1) x = jt.max(x, 2) x = x.view(batch_size, -1) x = self.relu(self.bn6(self.linear1(x))) x = self.dp1(x) x = self.relu(self.bn7(self.linear2(x))) x = self.dp2(x) x = self.linear3(x) return x
def execute(self, x): # (N, K, C_out) """ Applies XConv to the input data. :param x: (rep_pt, pts, fts) where - rep_pt: Representative point. - pts: Regional point cloud such that fts[:,p_idx,:] is the feature associated with pts[:,p_idx,:]. - fts: Regional features such that pts[:,p_idx,:] is the feature associated with fts[:,p_idx,:]. :return: Features aggregated into point rep_pt. """ rep_pt, pts, fts = x # b, n, c // b ,n k, c // b, n, k, d if fts is not None: assert(rep_pt.size()[0] == pts.size()[0] == fts.size()[0]) # Check N is equal. assert(rep_pt.size()[1] == pts.size()[1] == fts.size()[1]) # Check P is equal. assert(pts.size()[2] == fts.size()[2] == self.K) # Check K is equal. assert(fts.size()[3] == self.cin) # Check C_in is equal. else: assert(rep_pt.size()[0] == pts.size()[0]) # Check N is equal. assert(rep_pt.size()[1] == pts.size()[1]) # Check P is equal. assert(pts.size()[2] == self.K) # Check K is equal. assert(rep_pt.size()[2] == pts.size()[3] == self.dims) # Check dims is equal. N = pts.size()[0] P = rep_pt.size()[1] # (N, P, K, dims) p_center = jt.unsqueeze(rep_pt, dim = 2) # (N, P, 1, dims) # print (p_center.size()) # # Move pts to local coordinate system of rep_pt. pts_local = pts - p_center.repeat(1, 1, self.K, 1) # (N, P, K, dims) # pts_local = self.pts_layernorm(pts - p_center) # Individually lift each point into C_mid space. # print (pts_local.size(), 'before size') pts_local = pts_local.permute(0, 3, 1, 2) # N, dim, P, K fts_lifted0 = self.dense1(pts_local) # ? # print (.size(), 'after size') fts_lifted = self.dense2(fts_lifted0) # N, C_mid, P, K fts = fts.permute(0, 3, 1, 2) if fts is None: fts_cat = fts_lifted else: fts_cat = concat((fts_lifted, fts), 1) # (N, C_mid + C_in, P, K) # Learn the (N, K, K) X-transformation matrix. X_shape = (N, P, self.K, self.K) # X = self.x_trans(pts_local) # N, K*K, 1, P x = self.x_trans_0(pts_local) x = self.x_trans_1(x) X = self.x_trans_2(x) # print ('X size ', X.size()) X = X.permute(0, 2, 3, 1) # n p 1 k X = X.view(X_shape) # N, P, K, K # print (fts_cat.shape) fts_cat = fts_cat.permute(0, 2, 3, 1) fts_X = jt.matmul(X, fts_cat) # # print ('fts X size =', fts_X.shape) fts_p = self.end_conv(fts_X).squeeze(dim = 2) # print ('xxxxxxxxxxx') # print ('result size') # print (fts_X.size(), fts_p.size()) return fts_p