File size: 1,916 Bytes
57746f1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
"""
Utils for Datasets
Author: Xiaoyang Wu ([email protected])
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):
"""
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], 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]) for key in batch[0]}
for key in batch.keys():
if "offset" in key:
batch[key] = torch.cumsum(batch[key], dim=0)
return batch
else:
return default_collate(batch)
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 "offset" in batch.keys():
# Mix3d (https://arxiv.org/pdf/2110.02210.pdf)
if random.random() < mix_prob:
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):
return a * np.exp(-dist2 / (2 * c**2))
|