Randomly remove part of the point cloud. Similar to PyTorch RandomErasing but operating on 3D point clouds. Erases fronto-parallel cuboid. Instead of erasing we set coords of removed points to (0, 0, 0) to retain the same number of points
| 254 | |
| 255 | |
| 256 | class RemoveRandomBlock: |
| 257 | """ |
| 258 | Randomly remove part of the point cloud. Similar to PyTorch RandomErasing but operating on 3D point clouds. |
| 259 | Erases fronto-parallel cuboid. |
| 260 | Instead of erasing we set coords of removed points to (0, 0, 0) to retain the same number of points |
| 261 | """ |
| 262 | def __init__(self, p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3)): |
| 263 | self.p = p |
| 264 | self.scale = scale |
| 265 | self.ratio = ratio |
| 266 | |
| 267 | def get_params(self, coords): |
| 268 | # Find point cloud 3D bounding box |
| 269 | flattened_coords = coords.view(-1, 3) |
| 270 | min_coords, _ = torch.min(flattened_coords, dim=0) |
| 271 | max_coords, _ = torch.max(flattened_coords, dim=0) |
| 272 | span = max_coords - min_coords |
| 273 | area = span[0] * span[1] |
| 274 | erase_area = random.uniform(self.scale[0], self.scale[1]) * area |
| 275 | aspect_ratio = random.uniform(self.ratio[0], self.ratio[1]) |
| 276 | |
| 277 | h = math.sqrt(erase_area * aspect_ratio) |
| 278 | w = math.sqrt(erase_area / aspect_ratio) |
| 279 | |
| 280 | x = min_coords[0] + random.uniform(0, 1) * (span[0] - w) |
| 281 | y = min_coords[1] + random.uniform(0, 1) * (span[1] - h) |
| 282 | |
| 283 | return x, y, w, h |
| 284 | |
| 285 | def __call__(self, coords): |
| 286 | if random.random() < self.p: |
| 287 | x, y, w, h = self.get_params(coords) # Fronto-parallel cuboid to remove |
| 288 | mask = (x < coords[..., 0]) & (coords[..., 0] < x+w) & (y < coords[..., 1]) & (coords[..., 1] < y+h) |
| 289 | coords[mask] = torch.zeros_like(coords[mask]) |
| 290 | return coords |