| import numpy as np |
| from nltk.tree import Tree |
| import random |
|
|
|
|
| class Compose(object): |
| def __init__(self, transforms): |
| self.transforms = transforms |
|
|
| def __call__(self, coords_and_feats): |
| for t in self.transforms: |
| coords_and_feats = t(coords_and_feats) |
| return coords_and_feats |
|
|
| def __str__(self) -> str: |
| sep = "\n\t" |
| msg = "Compose of [" |
| for t in self.transforms: |
| msg += f"{sep}{t}," |
| msg += "]\n" |
| return msg |
|
|
| def __repr__(self) -> str: |
| sep = "\n\t" |
| msg = "Compose of [" |
| for t in self.transforms: |
| msg += f"{sep}{t}" |
| msg += "]\n" |
| return msg |
|
|
|
|
| class RandomDropSubTrees(object): |
| def __init__(self, probs=[0.1, 0.2, 0.3, 0.4, 0.5], max_cnt=5): |
| self.probs = probs |
| self.max_cnt = max_cnt |
| self.cnt = 0 |
|
|
| def remove_subtrees(self, root, level_idx): |
| reduced_root = Tree(root.label(), []) |
| if len(root) == 0: |
| return reduced_root |
| p = np.random.uniform(0, 1, len(root)) |
| |
| if level_idx >= len(self.probs): |
| level_idx = len(self.probs) - 1 |
| for child_idx in range(len(root)): |
| if self.cnt > self.max_cnt: |
| reduced_root.append(root[child_idx]) |
| continue |
| else: |
| if p[child_idx] > self.probs[level_idx]: |
| reduced_root.append( |
| self.remove_subtrees(root[child_idx], level_idx + 1) |
| ) |
| else: |
| self.cnt += 1 |
| return reduced_root |
|
|
| def __call__(self, tree): |
| self.cnt = 0 |
| return self.remove_subtrees(tree, level_idx=0) |
|
|
| def __str__(self) -> str: |
| return f"RandomDropSubTrees(probs={self.probs}, max_cnt={self.max_cnt})" |
|
|
| def __repr__(self) -> str: |
| return f"RandomDropSubTrees(probs={self.probs}, max_cnt={self.max_cnt})" |
|
|
|
|
| class RandomSkipParentNode(object): |
| def __init__(self, probs=[0.1], max_cnt=5): |
| self.probs = probs |
| self.max_cnt = max_cnt |
| self.cnt = 0 |
|
|
| def move_grandson_to_son(self, root, level_idx): |
| if len(root) == 0 or len(root) == 1: |
| return root |
| p = np.random.uniform(0, 1, len(root)) |
| |
| if level_idx >= len(self.probs): |
| level_idx = len(self.probs) - 1 |
| for child_idx in range(len(root)): |
| if self.cnt >= self.max_cnt: |
| break |
| if p[child_idx] < self.probs[level_idx]: |
| if len(root[child_idx]) == 0 or len(root[child_idx]) == 1: |
| continue |
| else: |
| idx = random.randint(0, len(root[child_idx]) - 1) |
| root[child_idx] = root[child_idx][idx] |
| self.cnt += 1 |
| else: |
| root[child_idx] = self.move_grandson_to_son( |
| root[child_idx], level_idx + 1 |
| ) |
| return root |
|
|
| def __call__(self, tree): |
| self.cnt = 0 |
| return self.move_grandson_to_son(tree, level_idx=0) |
|
|
| def __str__(self) -> str: |
| return f"RandomSkipParentNode(probs={self.probs}, max_cnt={self.max_cnt})" |
|
|
| def __repr__(self) -> str: |
| return f"RandomSkipParentNode(probs={self.probs}, max_cnt={self.max_cnt})" |
|
|
|
|
| class RandomSwapSiblingSubTrees(object): |
| def __init__(self, probs=[0.1], max_cnt=5): |
| self.probs = probs |
| self.max_cnt = max_cnt |
| self.cnt = 0 |
|
|
| def swap_sibling_subtrees(self, root, level_idx): |
| if len(root) < 2: |
| return root |
| p = np.random.uniform(0, 1, len(root)) |
| |
| if level_idx >= len(self.probs): |
| level_idx = len(self.probs) - 1 |
| for child_idx in range(len(root)): |
| if self.cnt >= self.max_cnt: |
| break |
| if p[child_idx] < self.probs[level_idx]: |
| if len(root[child_idx]) < 2: |
| continue |
| else: |
| my_subtree_idx = random.randint(0, len(root[child_idx]) - 1) |
| sibling_idx = random.randint(0, len(root) - 1) |
| if len(root[sibling_idx]) == 0: |
| continue |
| sibling_subtree_idx = random.randint(0, len(root[sibling_idx]) - 1) |
| my_subtree = root[child_idx][my_subtree_idx].copy() |
| sibling_subtree = root[sibling_idx][sibling_subtree_idx].copy() |
| root[child_idx][my_subtree_idx] = sibling_subtree |
| root[sibling_idx][sibling_subtree_idx] = my_subtree |
| self.cnt += 1 |
| else: |
| root[child_idx] = self.swap_sibling_subtrees( |
| root[child_idx], level_idx + 1 |
| ) |
| return root |
|
|
| def __call__(self, tree): |
| self.cnt = 0 |
| return self.swap_sibling_subtrees(tree, level_idx=0) |
|
|
| def __str__(self) -> str: |
| return f"RandomSwapSiblingSubTrees(probs={self.probs}, max_cnt={self.max_cnt})" |
|
|
| def __repr__(self) -> str: |
| return f"RandomSwapSiblingSubTrees(probs={self.probs}, max_cnt={self.max_cnt})" |
|
|
|
|
| class RandomRotateAligned(object): |
| def __init__(self, p=0.5, axis=2): |
| self.prob = p |
| self.axis = axis |
|
|
| def __call__(self, coords_and_feats): |
| coord = coords_and_feats[:, :3] |
| if np.random.rand() < self.prob: |
| angle = np.random.uniform() * 2 * np.pi |
| cos, sin = np.cos(angle), np.sin(angle) |
| R_x = np.array([[1, 0, 0], [0, cos, -sin], [0, sin, cos]]) |
| R_y = np.array([[cos, 0, sin], [0, 1, 0], [-sin, 0, cos]]) |
| R_z = np.array([[cos, -sin, 0], [sin, cos, 0], [0, 0, 1]]) |
| R = [R_x, R_y, R_z][self.axis] |
| coord = np.dot(coord, R) |
| coords_and_feats[:, :3] = coord |
| return coords_and_feats |
|
|
| def __str__(self) -> str: |
| return f"RandomRotateAligned(p={self.prob},axis={self.axis})" |
|
|
| def __repr__(self) -> str: |
| return f"RandomRotateAligned(p={self.prob},axis={self.axis})" |
|
|
|
|
| class RandomRotate(object): |
| def __init__(self, sigma=0.03, clip=0.09, p=0.5): |
| self.sigma = sigma |
| self.clip = clip |
| self.prob = p |
|
|
| def __call__(self, coords_and_feats): |
| coord = coords_and_feats[:, :3] |
| if np.random.rand() < self.prob: |
| angle_x = np.random.uniform() * 2 * np.pi |
| angle_y = np.random.uniform() * 2 * np.pi |
| angle_z = np.random.uniform() * 2 * np.pi |
| cos_x, sin_x = np.cos(angle_x), np.sin(angle_x) |
| cos_y, sin_y = np.cos(angle_y), np.sin(angle_y) |
| cos_z, sin_z = np.cos(angle_z), np.sin(angle_z) |
| R_x = np.array([[1, 0, 0], [0, cos_x, -sin_x], [0, sin_x, cos_x]]) |
| R_y = np.array([[cos_y, 0, sin_y], [0, 1, 0], [-sin_y, 0, cos_y]]) |
| R_z = np.array([[cos_z, -sin_z, 0], [sin_z, cos_z, 0], [0, 0, 1]]) |
| R = np.dot(R_z, np.dot(R_y, R_x)) |
| coord = np.dot(coord, R) |
| coords_and_feats[:, :3] = coord |
| return coords_and_feats |
|
|
| def __str__(self) -> str: |
| return f"RandomRotate(p={self.prob})" |
|
|
| def __repr__(self) -> str: |
| return f"RandomRotate(p={self.prob})" |
|
|
|
|
| class RandomMaskFeats(object): |
| def __init__(self, p=0.2): |
| self.prob = p |
|
|
| def __call__(self, coords_and_feats): |
| if len(coords_and_feats[0]) > 5: |
| feats = coords_and_feats[:, 5:] |
| feats[ |
| :, |
| np.random.choice( |
| np.arange(len(feats[0])), int(len(feats[0]) * self.prob) |
| ), |
| ] = 0 |
| coords_and_feats[:, 5:] = feats |
| return coords_and_feats |
|
|
| def __str__(self) -> str: |
| return f"RandomMaskFeats(p={self.prob})" |
|
|
| def __repr__(self) -> str: |
| return f"RandomMaskFeats(p={self.prob})" |
|
|
|
|
| class RandomElasticate(object): |
| def __init__(self, p=0.2, scales=[0.8, 1.2]): |
| self.prob = p |
| self.scales = scales |
|
|
| def __call__(self, coords_and_feats): |
| if len(coords_and_feats[0]) > 5: |
| if np.random.rand() < self.prob: |
| branches = coords_and_feats[:, 5:] |
| scales = np.random.uniform( |
| self.scales[0], self.scales[1], branches.shape |
| ) |
| branches *= scales |
| coords_and_feats[:, 5:] = branches |
| return coords_and_feats |
|
|
| def __str__(self) -> str: |
| return f"RandomElasticate(p={self.prob}, scales={self.scales})" |
|
|
| def __repr__(self) -> str: |
| return f"RandomElasticate(p={self.prob}, scales={self.scales})" |
|
|
|
|
| class RandomScaleCoords(object): |
| def __init__(self, scale=[0.8, 1.2], p=0.5): |
| self.scale = scale |
| self.prob = p |
|
|
| def __call__(self, coords_and_feats): |
| if np.random.rand() < self.prob: |
| scale = np.random.uniform(self.scale[0], self.scale[1]) |
| coords_and_feats[:, :4] *= scale |
| if len(coords_and_feats[0]) > 5: |
| coords_and_feats[:, 5:] *= scale |
| return coords_and_feats |
|
|
| def __str__(self) -> str: |
| return f"RandomScaleCoords(p={self.prob}, scale={self.scale})" |
|
|
| def __repr__(self) -> str: |
| return f"RandomScaleCoords(p={self.prob}, scale={self.scale})" |
|
|
|
|
| class RandomScaleCoordsTranslation(object): |
| def __init__(self, scale=[0.5, 2], p=0.5): |
| self.scale = scale |
| self.prob = p |
|
|
| def __call__(self, coords_and_feats): |
| if np.random.rand() < self.prob: |
| scale = np.random.uniform(self.scale[0], self.scale[1]) |
| coord1 = coords_and_feats[:, :4] |
| coord1 *= scale |
| coords_and_feats[:, :4] = coord1 |
| if len(coords_and_feats[0]) > 5: |
| coord2 = coords_and_feats[:, 5:] |
| coord2 *= scale |
| coords_and_feats[:, 5:] = coord2 |
| return coords_and_feats |
|
|
| def __str__(self) -> str: |
| return f"RandomScaleCoordsTranslation(p={self.prob}, scale={self.scale})" |
|
|
| def __repr__(self) -> str: |
| return f"RandomScaleCoordsTranslation(p={self.prob}, scale={self.scale})" |
|
|
|
|
| class RandomScaleFeats(object): |
| def __init__(self, scale=[0.5, 2], p=0.5): |
| self.scale = scale |
| self.prob = p |
|
|
| def __call__(self, coords_and_feats): |
| feats = coords_and_feats[:, 4:] |
| if np.random.rand() < self.prob: |
| scale = np.random.uniform(self.scale[0], self.scale[1]) |
| feats *= scale |
| coords_and_feats[:, 4:] = feats |
| return coords_and_feats |
|
|
| def __str__(self) -> str: |
| return f"RandomScaleFeats(p={self.prob}, scale={self.scale})" |
|
|
| def __repr__(self) -> str: |
| return f"RandomScaleFeats(p={self.prob}, scale={self.scale})" |
|
|
|
|
| class RandomShift(object): |
| def __init__(self, shift=[5, 5, 5], p=0.5): |
| self.shift = shift |
| self.prob = p |
|
|
| def __call__(self, coords_and_feats): |
| coord = coords_and_feats[:, :3] |
| if np.random.rand() < self.prob: |
| shift_x = np.random.uniform(-self.shift[0], self.shift[0]) |
| shift_y = np.random.uniform(-self.shift[1], self.shift[1]) |
| shift_z = np.random.uniform(-self.shift[2], self.shift[2]) |
| coord += [shift_x, shift_y, shift_z] |
| coords_and_feats[:, :3] = coord |
| return coords_and_feats |
|
|
| def __str__(self) -> str: |
| return f"RandomShift(p={self.prob}, shift={self.shift})" |
|
|
| def __repr__(self) -> str: |
| return f"RandomShift(p={self.prob}, shift={self.shift})" |
|
|
|
|
| class RandomFlip(object): |
| def __init__(self, p=0.5): |
| self.prob = p |
|
|
| def __call__(self, coords_and_feats): |
| coord = coords_and_feats[:, :3] |
| if np.random.rand() < self.prob: |
| if np.random.rand() < 0.5: |
| coord[:, 0] = -coord[:, 0] |
| if np.random.rand() < 0.5: |
| coord[:, 1] = -coord[:, 1] |
| coords_and_feats[:, :3] = coord |
| return coords_and_feats |
|
|
| def __str__(self) -> str: |
| return f"RandomFlip(p={self.prob})" |
|
|
| def __repr__(self) -> str: |
| return f"RandomFlip(p={self.prob})" |
|
|
|
|
| class RandomJitter(object): |
| def __init__(self, sigma=1, clip=5, p=0.5): |
| self.sigma = sigma |
| self.clip = clip |
| self.prob = p |
|
|
| def __call__(self, coords_and_feats): |
| coord = coords_and_feats[:, :3] |
| assert self.clip > 0 |
| if np.random.rand() < self.prob: |
| jitter = np.clip( |
| self.sigma * np.random.randn(coord.shape[0], 3), -self.clip, self.clip |
| ) |
| coord += jitter |
| coords_and_feats[:, :3] = coord |
| return coords_and_feats |
|
|
| def __str__(self) -> str: |
| return f"RandomJitter(p={self.prob}, sigma={self.sigma}, clip={self.clip})" |
|
|
| def __repr__(self) -> str: |
| return f"RandomJitter(p={self.prob}, sigma={self.sigma}, clip={self.clip})" |
|
|
|
|
| class RandomJitterLength(object): |
| def __init__(self, sigma=0.1, clip=1, p=0.5): |
| self.sigma = sigma |
| self.clip = clip |
| self.prob = p |
|
|
| def __call__(self, coords_and_feats): |
| feats1 = coords_and_feats[:, 3:4] |
| assert self.clip > 0 |
| if np.random.rand() < self.prob: |
| jitter1 = np.clip( |
| self.sigma * np.random.randn(*feats1.shape), -self.clip, self.clip |
| ) |
| feats1 += jitter1 |
| if len(coords_and_feats[0]) > 5: |
| feats2 = coords_and_feats[:, 5:] |
| jitter2 = np.clip( |
| self.sigma * np.random.randn(*feats2.shape), -self.clip, self.clip |
| ) |
| feats2 += jitter2 |
| coords_and_feats[:, 5:] = feats2 |
| coords_and_feats[:, 3:4] = feats1 |
| return coords_and_feats |
|
|
| def __str__(self) -> str: |
| return ( |
| f"RandomJitterLength(p={self.prob}, sigma={self.sigma}, clip={self.clip})" |
| ) |
|
|
| def __repr__(self) -> str: |
| return ( |
| f"RandomJitterLength(p={self.prob}, sigma={self.sigma}, clip={self.clip})" |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| transform = RandomScaleCoordsTranslation(p=1) |
|
|
| coords_feats = np.random.rand(1024, 3) |
| |
| |
| |
| |
| print(coords_feats[-1]) |
| transformed = transform(coords_feats) |
| print(transformed[-1]) |
| print(coords_feats[-1]) |
|
|
