| 193 | |
| 194 | |
| 195 | class JitterPoints: |
| 196 | def __init__(self, sigma=0.01, clip=None, p=1.): |
| 197 | assert 0 < p <= 1. |
| 198 | assert sigma > 0. |
| 199 | |
| 200 | self.sigma = sigma |
| 201 | self.clip = clip |
| 202 | self.p = p |
| 203 | |
| 204 | def __call__(self, e): |
| 205 | """ Randomly jitter points. jittering is per point. |
| 206 | Input: |
| 207 | BxNx3 array, original batch of point clouds |
| 208 | Return: |
| 209 | BxNx3 array, jittered batch of point clouds |
| 210 | """ |
| 211 | |
| 212 | sample_shape = (e.shape[0],) |
| 213 | if self.p < 1.: |
| 214 | # Create a mask for points to jitter |
| 215 | m = torch.distributions.categorical.Categorical(probs=torch.tensor([1 - self.p, self.p])) |
| 216 | mask = m.sample(sample_shape=sample_shape) |
| 217 | else: |
| 218 | mask = torch.ones(sample_shape, dtype=torch.int64 ) |
| 219 | |
| 220 | mask = mask == 1 |
| 221 | jitter = self.sigma * torch.randn_like(e[mask]) |
| 222 | |
| 223 | if self.clip is not None: |
| 224 | jitter = torch.clamp(jitter, min=-self.clip, max=self.clip) |
| 225 | |
| 226 | e[mask] = e[mask] + jitter |
| 227 | return e |
| 228 | |
| 229 | |
| 230 | class RemoveRandomPoints: |