ashawkey / dimr

[ECCV 2022] Disentangled Instance Mesh Reconstruction
Apache License 2.0
28 stars 4 forks source link

Confusion in RfsNet.forward() #4

Closed saladcat closed 2 years ago

saladcat commented 2 years ago

There is such code:


            # single-scale proposal gen (pointgroup)
            if self.training:

                ### BFS clustering on shifted coords
                idx_shift, start_len_shift = pointgroup_ops.ballquery_batch_p(coords_ + pt_offsets_, batch_idxs_, batch_offsets_, self.cluster_radius, self.cluster_shift_meanActive)
                proposals_idx_shift, proposals_offset_shift = pointgroup_ops.bfs_cluster(semantic_preds_, idx_shift.cpu(), start_len_shift.cpu(), self.cluster_npoint_thre)
                proposals_idx_shift[:, 1] = object_idxs[proposals_idx_shift[:, 1].long()].int() # remap: sumNPoint --> N

                # proposals_idx_shift: (sumNPoint, 2), int, dim 0 for cluster_id, dim 1 for corresponding point idxs in N
                # proposals_offset_shift: (nProposal + 1), int, start/end index for each proposal, e.g., [0, c1, c1+c2, ..., c1+...+c_nprop = sumNPoint], same information as cluster_id, just in the convinience of cuda operators.

                ### BFS clustering on original coords
                idx, start_len = pointgroup_ops.ballquery_batch_p(coords_, batch_idxs_, batch_offsets_, self.cluster_radius, self.cluster_meanActive)
                proposals_idx, proposals_offset = pointgroup_ops.bfs_cluster(semantic_preds_, idx.cpu(), start_len.cpu(), self.cluster_npoint_thre)
                proposals_idx[:, 1] = object_idxs[proposals_idx[:, 1].long()].int()

                # proposals_idx: (sumNPoint, 2), int, dim 0 for cluster_id, dim 1 for corresponding point idxs in N
                # proposals_offset: (nProposal + 1), int

                ### merge two type of clusters
                proposals_idx_shift[:, 0] += (proposals_offset.size(0) - 1)
                proposals_offset_shift += proposals_offset[-1]

                proposals_idx = torch.cat((proposals_idx, proposals_idx_shift), dim=0)#.long().cuda()
                proposals_offset = torch.cat((proposals_offset, proposals_offset_shift[1:]), dim=0)#.cuda()
                # why [1:]: offset is (0, c1, c2), offset_shift is (0, d1, d2) + c2, output is (0, c1, c2, c2+d1, c2+d2)

            # multi-scale proposal gen (naive, maskgroup)
            else:

                idx_shift, start_len_shift = pointgroup_ops.ballquery_batch_p(coords_ + pt_offsets_, batch_idxs_, batch_offsets_, 0.01, self.cluster_shift_meanActive)
                proposals_idx_shift_0, proposals_offset_shift_0 = pointgroup_ops.bfs_cluster(semantic_preds_, idx_shift.cpu(), start_len_shift.cpu(), self.cluster_npoint_thre)
                proposals_idx_shift_0[:, 1] = object_idxs[proposals_idx_shift_0[:, 1].long()].int()

                idx_shift, start_len_shift = pointgroup_ops.ballquery_batch_p(coords_ + pt_offsets_, batch_idxs_, batch_offsets_, 0.03, self.cluster_shift_meanActive)
                proposals_idx_shift_1, proposals_offset_shift_1 = pointgroup_ops.bfs_cluster(semantic_preds_, idx_shift.cpu(), start_len_shift.cpu(), self.cluster_npoint_thre)
                proposals_idx_shift_1[:, 1] = object_idxs[proposals_idx_shift_1[:, 1].long()].int()

                idx_shift, start_len_shift = pointgroup_ops.ballquery_batch_p(coords_ + pt_offsets_, batch_idxs_, batch_offsets_, 0.05, self.cluster_shift_meanActive)
                proposals_idx_shift_2, proposals_offset_shift_2 = pointgroup_ops.bfs_cluster(semantic_preds_, idx_shift.cpu(), start_len_shift.cpu(), self.cluster_npoint_thre)
                proposals_idx_shift_2[:, 1] = object_idxs[proposals_idx_shift_2[:, 1].long()].int()

                idx, start_len = pointgroup_ops.ballquery_batch_p(coords_, batch_idxs_, batch_offsets_, 0.03, self.cluster_meanActive)
                proposals_idx_0, proposals_offset_0 = pointgroup_ops.bfs_cluster(semantic_preds_, idx.cpu(), start_len.cpu(), self.cluster_npoint_thre)
                proposals_idx_0[:, 1] = object_idxs[proposals_idx_0[:, 1].long()].int()

                _offset = proposals_offset_0.size(0) - 1
                proposals_idx_shift_0[:, 0] += _offset
                proposals_offset_shift_0 += proposals_offset_0[-1]

                _offset += proposals_offset_shift_0.size(0) - 1
                proposals_idx_shift_1[:, 0] += _offset
                proposals_offset_shift_1 += proposals_offset_shift_0[-1]

                _offset += proposals_offset_shift_1.size(0) - 1
                proposals_idx_shift_2[:, 0] += _offset
                proposals_offset_shift_2 += proposals_offset_shift_1[-1]

                proposals_idx = torch.cat((proposals_idx_0, proposals_idx_shift_0, proposals_idx_shift_1, proposals_idx_shift_2), dim=0)
                proposals_offset = torch.cat((proposals_offset_0, proposals_offset_shift_0[1:], proposals_offset_shift_1[1:], proposals_offset_shift_2[1:]))

            #### proposals voxelization again
            input_feats, inp_map, proposal_angle, point_center, point_scale, proposal_semantics = self.clusters_voxelization(proposals_idx, proposals_offset, output_feats, semantic_scores_CAD, pt_angles, coords, self.score_fullscale, self.score_scale, self.mode)

I would like to know how to find the original batch_idx corresponding to each Proposal BBOX?