Source code for pimm.datasets.utils

"""Collation helpers for ragged point-cloud batches.

Dataset items are flat dictionaries with per-point arrays and one-sample
``offset`` tensors. These helpers concatenate ragged fields and convert offsets
to cumulative batch offsets expected by Pointcept/pimm model heads.

Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com)
Please cite our work if the code is helpful to you.
"""

import random
from collections.abc import Mapping, Sequence
import numpy as np
import torch
from torch.utils.data.dataloader import default_collate




def collate_fn(batch, mix_prob=0):
    """Recursively collate tensors, mappings, and list-style point samples.

    Tensor leaves are concatenated instead of stacked. Keys containing
    ``offset`` are treated as per-sample lengths and converted into cumulative
    offsets after concatenation.
    """
    if not isinstance(batch, Sequence):
        raise TypeError(f"{batch.dtype} is not supported.")

    if isinstance(batch[0], torch.Tensor):
        return torch.cat(list(batch))
    elif isinstance(batch[0], str):
        # str is also a kind of Sequence, judgement should before Sequence
        return list(batch)
    elif isinstance(batch[0], Sequence):
        for data in batch:
            data.append(torch.tensor([data[0].shape[0]]))
        batch = [collate_fn(samples) for samples in zip(*batch)]
        batch[-1] = torch.cumsum(batch[-1], dim=0).int()
        return batch
    elif isinstance(batch[0], Mapping):
        batch = {
            key: (
                collate_fn([d[key] for d in batch])
                if "offset" not in key
                # offset -> bincount -> concat bincount-> concat offset
                else torch.cumsum(
                    collate_fn([d[key].diff(prepend=torch.tensor([0])) for d in batch]),
                    dim=0,
                )
            )
            for key in batch[0]
            if not key.startswith("_")  # skip non-tensor metadata
        }
        return batch
    else:
        return default_collate(batch)



def point_collate_fn(batch, mix_prob=0):
    """Collate point-cloud dictionaries and optionally apply pairwise mixup."""
    assert isinstance(
        batch[0], Mapping
    )  # currently, only support input_dict, rather than input_list
    batch = collate_fn(batch)
    if random.random() < mix_prob:
        if "instance" in batch.keys():
            offset = batch["offset"]
            start = 0
            num_instance = 0
            for i in range(len(offset)):
                if i % 2 == 0:
                    num_instance = max(batch["instance"][start : offset[i]])
                if i % 2 != 0:
                    mask = batch["instance"][start : offset[i]] != -1
                    batch["instance"][start : offset[i]] += num_instance * mask
                start = offset[i]
        if "offset" in batch.keys():
            batch["offset"] = torch.cat(
                [batch["offset"][1:-1:2], batch["offset"][-1].unsqueeze(0)], dim=0
            )
    return batch


def gaussian_kernel(dist2: np.array, a: float = 1, c: float = 5):
    """Evaluate a scalar Gaussian kernel for squared distances."""
    return a * np.exp(-dist2 / (2 * c**2))


def inseg_collate_fn(batch, mix_prob=0):
    """Collate instance-segmentation query lists into one flattened batch.

    Each incoming item is a list of dictionaries, one per query/candidate.
    """
    assert isinstance(
        batch[0], list
    )  # currently, only support input_dict, rather than input_list
    assert isinstance(
        batch[0][0], Mapping
    )  # currently, only support input_dict, rather than input_list
    # Flatten the batch - each item is already a list of dictionaries

    flattened_batch = []
    for item_list in batch:
        flattened_batch.extend(item_list)

    # Use the original collate function on the flattened batch
    return collate_fn(flattened_batch)