Function logits_mask

modules/functional/sampling.py:51–84

Use logits to sample points :param coords: coords of points, FloatTensor[B, 3, N] :param logits: binary classification logits, FloatTensor[B, 2, N] :param num_points_per_object: M, #points per object after masking, int :return: selected_coords: FloatTensor[B, 3, M]

(coords, logits, num_points_per_object)

def logits_mask(coords, logits, num_points_per_object):
    """
    Use logits to sample points
    :param coords: coords of points, FloatTensor[B, 3, N]
    :param logits: binary classification logits, FloatTensor[B, 2, N]
    :param num_points_per_object: M, #points per object after masking, int
    :return:
        selected_coords: FloatTensor[B, 3, M]
        masked_coords_mean: mean coords of selected points, FloatTensor[B, 3]
        mask: mask to select points, BoolTensor[B, N]
    """
    batch_size, _, num_points = coords.shape
    mask = torch.lt(logits[:, 0, :], logits[:, 1, :])   # [B, N]
    num_candidates = torch.sum(mask, dim=-1, keepdim=True)  # [B, 1]
    masked_coords = coords * mask.view(batch_size, 1, num_points)  # [B, C, N]
    masked_coords_mean = torch.sum(masked_coords, dim=-1) / torch.max(num_candidates,
                                                                      torch.ones_like(num_candidates)).float()  # [B, C]
    selected_indices = torch.zeros((batch_size, num_points_per_object), device=coords.device, dtype=torch.int32)
    for i in range(batch_size):
        current_mask = mask[i]  # [N]
        current_candidates = current_mask.nonzero().view(-1)
        current_num_candidates = current_candidates.numel()
        if current_num_candidates >= num_points_per_object:
            choices = np.random.choice(current_num_candidates, num_points_per_object, replace=False)
            selected_indices[i] = current_candidates[choices]
        elif current_num_candidates > 0:
            choices = np.concatenate([
                np.arange(current_num_candidates).repeat(num_points_per_object // current_num_candidates),
                np.random.choice(current_num_candidates, num_points_per_object % current_num_candidates, replace=False)
            ])
            np.random.shuffle(choices)
            selected_indices[i] = current_candidates[choices]
    selected_coords = gather(masked_coords - masked_coords_mean.view(batch_size, -1, 1), selected_indices)
    return selected_coords, masked_coords_mean, mask