Class InverseCDFRaySampling

mapping/src/functions/voxel_helpers.py:257–353

class InverseCDFRaySampling(Function):
    @staticmethod
    def forward(
            ctx,
            pts_idx,
            min_depth,
            max_depth,
            probs,
            steps,
            fixed_step_size=-1,
            deterministic=False,
    ):
        G, N, P = int(np.ceil(pts_idx.size(0) / 2048)), pts_idx.size(0), pts_idx.size(1)
        # if G == 0:
        #     G = 200
        H = int(np.ceil(N / G)) * G

        if H > N:
            pts_idx = torch.cat([pts_idx, pts_idx[:1].expand(H - N, P)], 0)
            min_depth = torch.cat(
                [min_depth, min_depth[:1].expand(H - N, P)], 0)
            max_depth = torch.cat(
                [max_depth, max_depth[:1].expand(H - N, P)], 0)
            probs = torch.cat([probs, probs[:1].expand(H - N, P)], 0)
            steps = torch.cat([steps, steps[:1].expand(H - N)], 0)

        # print(G, P, np.ceil(N / G), N, H, pts_idx.shape, min_depth.device)
        pts_idx = pts_idx.reshape(G, -1, P)
        min_depth = min_depth.reshape(G, -1, P)
        max_depth = max_depth.reshape(G, -1, P)
        probs = probs.reshape(G, -1, P)
        steps = steps.reshape(G, -1)

        # pre-generate noise
        max_steps = steps.ceil().long().max() + P
        noise = min_depth.new_zeros(*min_depth.size()[:-1], max_steps)
        if deterministic:
            noise += 0.5
        else:
            noise = noise.uniform_().clamp(min=0.001, max=0.999)  # in case

        # call cuda function
        # chunk_size = 4 * G  # to avoid oom?
        results = [
            _ext.inverse_cdf_sampling(
                pts_idx.contiguous(),
                min_depth.float().contiguous(),
                max_depth.float().contiguous(),
                noise.float().contiguous(),
                probs.float().contiguous(),
                steps.float().contiguous(),
                fixed_step_size,
            )
        ]
        #
        # # call cuda function
        # chunk_size = 4 * G  # to avoid oom?
        # results = [