| 19 | |
| 20 | |
| 21 | class Augment_Point_Data(): |
| 22 | def __init__(self, mode=0, is_train=False): |
| 23 | if is_train: |
| 24 | if mode == 0: |
| 25 | transform = [transforms.ToTensor()] |
| 26 | elif mode == 1: |
| 27 | transform = [transforms.ToTensor(), |
| 28 | JitterPoints(sigma=0.001, clip=0.002), |
| 29 | RemoveRandomPoints(r=(0.0, 0.1)), |
| 30 | RandomTranslation(max_delta=0.01), |
| 31 | RemoveRandomBlock(p=0.4), |
| 32 | RandomRotation(axis=np.array([0,0,1]), |
| 33 | max_theta=15, |
| 34 | max_theta2=None)] |
| 35 | else: |
| 36 | raise NotImplementedError(f'Uncognized data augmentation mode') |
| 37 | else: |
| 38 | transform = [transforms.ToTensor()] |
| 39 | self.transform = transforms.Compose(transform) |
| 40 | |
| 41 | def __call__(self, input): |
| 42 | output = self.transform(input) |
| 43 | return output |
| 44 | |
| 45 | |
| 46 | class Augment_RGB_Data(): |