pointcept/datasets/utils.py

"""
Utils for Datasets

Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com), Yujia Zhang (yujia.zhang.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
import torch.nn.functional as F

from torch_scatter import scatter_min
from pointcept.models.utils import offset2batch


def collate_fn(batch):
    """
    collate function for point cloud which support dict and list,
    'coord' is necessary to determine 'offset'
    """
    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], list):
        batch = [torch.tensor(data) for data in batch]
        return torch.cat(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):
        if "img_num" in batch[0].keys():
            max_img_num = max([d["img_num"] for d in batch])
        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,
                    )
                )
                if "correspondence" not in key
                else collate_fn(
                    [
                        F.pad(
                            d[key].permute(0, 2, 1),
                            (0, max_img_num - d[key].shape[1]),
                            value=-1,
                        ).permute(0, 2, 1)
                        for d in batch
                    ]
                )
            )
            for key in batch[0]
        }
        return batch
    else:
        return default_collate(batch)


def pairwise_concatenate(tensors_list, dim=0):
    new_list = []
    n = len(tensors_list)
    for i in range(0, n - 1, 2):
        tensor1 = tensors_list[i]
        tensor2 = tensors_list[i + 1]
        concatenated_tensor = torch.cat([tensor1, tensor2], dim=dim)
        new_list.append(concatenated_tensor)
    if n % 2 != 0:
        new_list.append(tensors_list[-1])
    return new_list


def regroup_batch(batch, N, original_offsets, data_keys):
    num_segments = len(original_offsets)
    grouped_segments = {key: [[] for _ in range(N)] for key in data_keys}
    start_idx = 0
    for i in range(num_segments):
        end_idx = original_offsets[i]
        group_idx = i % N
        for key in data_keys:
            segment = batch[key][start_idx:end_idx]
            grouped_segments[key][group_idx].append(segment)
        start_idx = end_idx

    segment_lengths = original_offsets[1:] - original_offsets[:-1]
    segment_lengths = torch.cat([original_offsets[:1], segment_lengths])
    new_lengths_order = []
    split_lengths_order = []
    for i in range(N):
        segment_lengths_i = segment_lengths[i::N]
        n = segment_lengths_i.shape[0]
        num_pairs = n // 2
        paired_part = segment_lengths_i[: num_pairs * 2]
        reshaped_part = paired_part.view(num_pairs, 2, *segment_lengths_i.shape[1:])
        sums = torch.sum(reshaped_part, dim=1)
        first_parts = reshaped_part[:, 0]
        if n % 2 != 0:
            last_element = segment_lengths_i[-1]
            last_element = last_element.unsqueeze(0)
            final_result = torch.cat([sums, last_element], dim=0)
        else:
            final_result = sums
        new_lengths_order.append(final_result)
        split_lengths_order.append(first_parts)
    new_lengths_order = torch.stack(new_lengths_order, dim=1)
    new_lengths_order = new_lengths_order.flatten()
    new_offsets = torch.cumsum(new_lengths_order, dim=0)
    split_lengths_order = torch.stack(split_lengths_order, dim=1).flatten()

    new_batch = {}
    new_batch_imgs = []
    new_batch_img_num = 0
    img_num_offset = torch.cat(
        [torch.tensor([0]), torch.cumsum(batch["img_num"], dim=0)]
    )
    for key in data_keys:
        final_segments_in_order = []
        for i in range(N):
            grouped_segments[key][i] = pairwise_concatenate(
                grouped_segments[key][i], dim=0
            )
        for i in range(len(grouped_segments[key][0])):
            for j in range(N):
                final_segments_in_order.append(grouped_segments[key][j][i])
        new_batch[key] = torch.vstack(final_segments_in_order)
        if "correspondence" in key:
            current_start = 0
            N0, v, n_dim = new_batch[key].shape
            v2 = v * 2
            batch_correspondence_mix = -torch.ones(
                (N0, v2, n_dim),
                dtype=new_batch[key].dtype,
                device=new_batch[key].device,
            )
            for k, end in enumerate(new_offsets):
                len_part1 = split_lengths_order[k]
                split_point = current_start + len_part1
                if split_point > current_start:
                    mask1 = torch.any(
                        new_batch[key][current_start:split_point]
                        != torch.tensor([-1, -1]),
                        dim=2,
                    )
                    valid_index1 = torch.where(mask1)
                    if len(valid_index1[1]) == 0:
                        count1 = 0
                    else:
                        count1 = max(valid_index1[1])
                    batch_correspondence_mix[current_start:split_point, 0:count1] = (
                        new_batch[key][current_start:split_point, 0:count1]
                    )
                    if k % N == 0 and N == 2:
                        new_batch_imgs.append(
                            batch["images"][
                                img_num_offset[k // N * 2] : img_num_offset[k // N * 2]
                                + count1
                            ]
                        )
                if end > split_point:
                    mask2 = torch.any(
                        new_batch[key][split_point:end] != torch.tensor([-1, -1]), dim=2
                    )
                    valid_index2 = torch.where(mask2)
                    if len(valid_index2[1]) == 0:
                        count2 = 0
                    else:
                        count2 = max(valid_index2[1])
                    batch_correspondence_mix[
                        split_point:end, count1 : count1 + count2
                    ] = new_batch[key][split_point:end, 0:count2]
                    if k % N == 0 and N == 2:
                        new_batch_imgs.append(
                            batch["images"][
                                img_num_offset[k // N * 2 + 1] : img_num_offset[
                                    k // N * 2 + 1
                                ]
                                + count2
                            ]
                        )
                current_start = end

            if N == 2:
                new_batch_img_num = torch.tensor([i.shape[0] for i in new_batch_imgs])
                new_batch_imgs = torch.vstack(new_batch_imgs)
            else:
                new_batch_img_num = None
                new_batch_imgs = None
            new_batch[key] = batch_correspondence_mix
    return new_batch, new_offsets, new_batch_imgs, new_batch_img_num


def point_collate_fn(batch, mix_prob=0):
    assert isinstance(
        batch[0], Mapping
    )  # currently, only support input_dict, rather than input_list
    batch = collate_fn(batch)
    if random.random() < mix_prob:
        valid_keys = [
            "coord",
            "grid_coord",
            "origin_coord",
            "color",
            "normal",
            "feat",
            "correspondence",
        ]
        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]
        offset_assets = [asset for asset in batch.keys() if "offset" in asset]
        for offset_asset in offset_assets:
            batch[offset_asset] = torch.cat(
                [batch[offset_asset][1:-1:2], batch[offset_asset][-1].unsqueeze(0)],
                dim=0,
            )

        # Recompute grid_coord after mixing, because each scene's grid_coord was
        # independently shifted before mixing and is no longer consistent with
        # the merged coord. Only done when grid_size is available (e.g. LitePT
        # configs); other configs are unaffected.
        if "grid_coord" in batch and "grid_size" in batch:
            batch_idx = offset2batch(batch["offset"])
            scaled_coord = batch["coord"] / batch["grid_size"][0]
            grid_coord = torch.floor(scaled_coord).to(torch.int64)
            min_coord, _ = scatter_min(grid_coord, batch_idx, dim=0)
            batch["grid_coord"] = grid_coord - min_coord[batch_idx]
        offset_assets = [asset for asset in batch.keys() if "_offset" in asset]
        for offset_asset in offset_assets:
            offset_prefix = offset_asset.split("_")[0]
            valid_keys_with_prefix = [
                offset_prefix + "_" + valid_key for valid_key in valid_keys
            ]
            valid_keys_with_prefix = [
                valid_key_with_prefix
                for valid_key_with_prefix in valid_keys_with_prefix
                if valid_key_with_prefix in batch.keys()
            ]
            if "global" in offset_asset:
                N = 2
            elif "local" in offset_asset:
                N = 4
            updated_batch, new_offset, imgs, img_num = regroup_batch(
                batch, N, batch[offset_asset], valid_keys_with_prefix
            )
            batch[offset_asset] = new_offset
            batch.update(updated_batch)
            if "global" in offset_asset:
                batch["images"] = imgs
                batch["img_num"] = img_num

        if "img_num" in batch.keys():
            n = batch["img_num"].shape[0]
            num_pairs = n // 2
            len_pairs = num_pairs * 2
            pairs_tensor = batch["img_num"][:len_pairs]

            if num_pairs == 0:
                pass
            else:
                summed_pairs = pairs_tensor.view(-1, 2).sum(dim=1)
                if n % 2 != 0:
                    last_element = batch["img_num"][-1:]
                    result = torch.cat((summed_pairs, last_element))
                else:
                    result = summed_pairs
                batch["img_num"] = result
    return batch


def gaussian_kernel(dist2: np.array, a: float = 1, c: float = 5):
    return a * np.exp(-dist2 / (2 * c**2))