├── .gitignore ├── README.md ├── box_np_ops.py ├── docs └── methods.jpg ├── geometry.py └── part_aware_augmentation.py /.gitignore: -------------------------------------------------------------------------------- 1 | 2 | __pycache__/ 3 | 4 | .idea/ 5 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Part-Aware Data Augmentation for 3D Object Detection in Point Cloud 2 | 3 | This repository contains a reference implementation of our [Part-Aware Data Augmentation for 3D Object Detection in Point Cloud](https://ieeexplore.ieee.org/document/9635887) (IROS 2021). 4 | 5 |

6 | 7 |

8 | 9 | 10 | If you find this code useful in your research, please consider citing our work: 11 | ``` 12 | @inproceedings{choi2021part, 13 | title={Part-aware data augmentation for 3d object detection in point cloud}, 14 | author={Choi, Jaeseok and Song, Yeji and Kwak, Nojun}, 15 | booktitle={2021 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)}, 16 | pages={3391--3397}, 17 | year={2021}, 18 | organization={IEEE} 19 | } 20 | ``` 21 | ## Prerequisites 22 | Our code was tested on [second.pytorch](https://github.com/traveller59/second.pytorch) and [OpenPCDet](https://github.com/open-mmlab/OpenPCDet). 23 | This repository contains only part-aware data augmentation code. 24 | Refer to the link above for code such as data loader or detector. 25 | 26 | ## Usage 27 | ``` 28 | Args: 29 | ** only supports KITTI format ** 30 | points: lidar points (N, 4), 31 | gt_boxes: ground truth boxes (B, 7), 32 | gt_names: ground truth classes (B, 1), 33 | class_names: list of classes to augment (3), 34 | pa_aug_param: parameters for PA_AUG (string). 35 | 36 | Returns: 37 | points: augmented lidar points (N', 4), 38 | gt_boxes_mask: mask for gt_boxes (B) 39 | 40 | 41 | class_names = ['Car', 'Pedestrian', 'Cyclist'] 42 | pa_aug_param = "dropout_p02_swap_p02_mix_p02_sparse40_p01_noise10_p01" 43 | 44 | pa_aug = PartAwareAugmentation(points, gt_boxes, gt_names, class_names=class_names) 45 | points, gt_boxes_mask = pa_aug.augment(pa_aug_param=pa_aug_param) 46 | gt_boxes = gt_boxes[gt_boxes_mask] 47 | ``` 48 | 49 | ## Example 50 | Follow [this repo](https://github.com/sky77764/PA-AUG-MD3D) if you want to check the implementation on OpenPCDet. 51 | -------------------------------------------------------------------------------- /box_np_ops.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | 3 | import numba 4 | import numpy as np 5 | from spconv.utils import rbbox_iou 6 | 7 | from geometry import points_in_convex_polygon_3d_jit, points_count_convex_polygon_3d_jit 8 | 9 | 10 | def riou_cc(rbboxes, qrbboxes, standup_thresh=0.0): 11 | # less than 50ms when used in second one thread. 10x slower than gpu 12 | boxes_corners = center_to_corner_box2d(rbboxes[:, :2], rbboxes[:, 2:4], 13 | rbboxes[:, 4]) 14 | boxes_standup = corner_to_standup_nd(boxes_corners) 15 | qboxes_corners = center_to_corner_box2d(qrbboxes[:, :2], qrbboxes[:, 2:4], 16 | qrbboxes[:, 4]) 17 | qboxes_standup = corner_to_standup_nd(qboxes_corners) 18 | # if standup box not overlapped, rbbox not overlapped too. 19 | standup_iou = iou_jit(boxes_standup, qboxes_standup, eps=0.0) 20 | return rbbox_iou(boxes_corners, qboxes_corners, standup_iou, 21 | standup_thresh) 22 | 23 | def second_box_encode(boxes, 24 | anchors, 25 | encode_angle_to_vector=False, 26 | smooth_dim=False, 27 | cylindrical=False): 28 | """box encode for VoxelNet in lidar 29 | Args: 30 | boxes ([N, 7 + ?] Tensor): normal boxes: x, y, z, w, l, h, r, custom values 31 | anchors ([N, 7] Tensor): anchors 32 | """ 33 | # need to convert boxes to z-center format 34 | box_ndim = anchors.shape[-1] 35 | cas, cgs = [], [] 36 | if box_ndim > 7: 37 | xa, ya, za, wa, la, ha, ra, *cas = np.split(anchors, box_ndim, axis=1) 38 | xg, yg, zg, wg, lg, hg, rg, *cgs = np.split(boxes, box_ndim, axis=1) 39 | else: 40 | xa, ya, za, wa, la, ha, ra = np.split(anchors, box_ndim, axis=1) 41 | xg, yg, zg, wg, lg, hg, rg = np.split(boxes, box_ndim, axis=1) 42 | 43 | diagonal = np.sqrt(la**2 + wa**2) # 4.3 44 | xt = (xg - xa) / diagonal 45 | yt = (yg - ya) / diagonal 46 | zt = (zg - za) / ha # 1.6 47 | lt = np.log(lg / la) 48 | wt = np.log(wg / wa) 49 | ht = np.log(hg / ha) 50 | rt = rg - ra 51 | cts = [g - a for g, a in zip(cgs, cas)] 52 | if smooth_dim: 53 | lt = lg / la - 1 54 | wt = wg / wa - 1 55 | ht = hg / ha - 1 56 | else: 57 | lt = np.log(lg / la) 58 | wt = np.log(wg / wa) 59 | ht = np.log(hg / ha) 60 | if encode_angle_to_vector: 61 | rgx = np.cos(rg) 62 | rgy = np.sin(rg) 63 | rax = np.cos(ra) 64 | ray = np.sin(ra) 65 | rtx = rgx - rax 66 | rty = rgy - ray 67 | return np.concatenate([xt, yt, zt, wt, lt, ht, rtx, rty, *cts], axis=1) 68 | else: 69 | rt = rg - ra 70 | return np.concatenate([xt, yt, zt, wt, lt, ht, rt, *cts], axis=1) 71 | 72 | 73 | 74 | def second_box_decode(box_encodings, 75 | anchors, 76 | encode_angle_to_vector=False, 77 | smooth_dim=False): 78 | """box decode for VoxelNet in lidar 79 | Args: 80 | boxes ([N, 7] Tensor): normal boxes: x, y, z, w, l, h, r 81 | anchors ([N, 7] Tensor): anchors 82 | """ 83 | # need to convert box_encodings to z-bottom format 84 | box_ndim = anchors.shape[-1] 85 | cas, cts = [], [] 86 | if box_ndim > 7: 87 | xa, ya, za, wa, la, ha, ra, *cas = np.split(anchors, box_ndim, axis=-1) 88 | if encode_angle_to_vector: 89 | xt, yt, zt, wt, lt, ht, rtx, rty, *cts = np.split(box_encodings, box_ndim + 1, axis=-1) 90 | else: 91 | xt, yt, zt, wt, lt, ht, rt, *cts = np.split(box_encodings, box_ndim, axis=-1) 92 | else: 93 | xa, ya, za, wa, la, ha, ra = np.split(anchors, box_ndim, axis=-1) 94 | if encode_angle_to_vector: 95 | xt, yt, zt, wt, lt, ht, rtx, rty = np.split(box_encodings, box_ndim + 1, axis=-1) 96 | else: 97 | xt, yt, zt, wt, lt, ht, rt = np.split(box_encodings, box_ndim, axis=-1) 98 | 99 | diagonal = np.sqrt(la**2 + wa**2) 100 | xg = xt * diagonal + xa 101 | yg = yt * diagonal + ya 102 | zg = zt * ha + za 103 | if smooth_dim: 104 | lg = (lt + 1) * la 105 | wg = (wt + 1) * wa 106 | hg = (ht + 1) * ha 107 | else: 108 | lg = np.exp(lt) * la 109 | wg = np.exp(wt) * wa 110 | hg = np.exp(ht) * ha 111 | if encode_angle_to_vector: 112 | rax = np.cos(ra) 113 | ray = np.sin(ra) 114 | rgx = rtx + rax 115 | rgy = rty + ray 116 | rg = np.arctan2(rgy, rgx) 117 | else: 118 | rg = rt + ra 119 | cgs = [t + a for t, a in zip(cts, cas)] 120 | return np.concatenate([xg, yg, zg, wg, lg, hg, rg, *cgs], axis=-1) 121 | 122 | 123 | def bev_box_encode(boxes, 124 | anchors, 125 | encode_angle_to_vector=False, 126 | smooth_dim=False): 127 | """box encode for VoxelNet in lidar 128 | Args: 129 | boxes ([N, 7] Tensor): normal boxes: x, y, z, w, l, h, r 130 | anchors ([N, 7] Tensor): anchors 131 | encode_angle_to_vector: bool. increase aos performance, 132 | decrease other performance. 133 | """ 134 | # need to convert boxes to z-center format 135 | xa, ya, wa, la, ra = np.split(anchors, 5, axis=-1) 136 | xg, yg, wg, lg, rg = np.split(boxes, 5, axis=-1) 137 | diagonal = np.sqrt(la**2 + wa**2) # 4.3 138 | xt = (xg - xa) / diagonal 139 | yt = (yg - ya) / diagonal 140 | if smooth_dim: 141 | lt = lg / la - 1 142 | wt = wg / wa - 1 143 | else: 144 | lt = np.log(lg / la) 145 | wt = np.log(wg / wa) 146 | if encode_angle_to_vector: 147 | rgx = np.cos(rg) 148 | rgy = np.sin(rg) 149 | rax = np.cos(ra) 150 | ray = np.sin(ra) 151 | rtx = rgx - rax 152 | rty = rgy - ray 153 | return np.concatenate([xt, yt, wt, lt, rtx, rty], axis=-1) 154 | else: 155 | rt = rg - ra 156 | return np.concatenate([xt, yt, wt, lt, rt], axis=-1) 157 | 158 | 159 | def bev_box_decode(box_encodings, 160 | anchors, 161 | encode_angle_to_vector=False, 162 | smooth_dim=False): 163 | """box decode for VoxelNet in lidar 164 | Args: 165 | boxes ([N, 7] Tensor): normal boxes: x, y, z, w, l, h, r 166 | anchors ([N, 7] Tensor): anchors 167 | """ 168 | # need to convert box_encodings to z-bottom format 169 | xa, ya, wa, la, ra = np.split(anchors, 5, axis=-1) 170 | if encode_angle_to_vector: 171 | xt, yt, wt, lt, rtx, rty = np.split(box_encodings, 6, axis=-1) 172 | else: 173 | xt, yt, wt, lt, rt = np.split(box_encodings, 5, axis=-1) 174 | diagonal = np.sqrt(la**2 + wa**2) 175 | xg = xt * diagonal + xa 176 | yg = yt * diagonal + ya 177 | if smooth_dim: 178 | lg = (lt + 1) * la 179 | wg = (wt + 1) * wa 180 | else: 181 | lg = np.exp(lt) * la 182 | wg = np.exp(wt) * wa 183 | if encode_angle_to_vector: 184 | rax = np.cos(ra) 185 | ray = np.sin(ra) 186 | rgx = rtx + rax 187 | rgy = rty + ray 188 | rg = np.arctan2(rgy, rgx) 189 | else: 190 | rg = rt + ra 191 | return np.concatenate([xg, yg, wg, lg, rg], axis=-1) 192 | 193 | 194 | def corners_nd(dims, origin=0.5): 195 | """generate relative box corners based on length per dim and 196 | origin point. 197 | 198 | Args: 199 | dims (float array, shape=[N, ndim]): array of length per dim 200 | origin (list or array or float): origin point relate to smallest point. 201 | 202 | Returns: 203 | float array, shape=[N, 2 ** ndim, ndim]: returned corners. 204 | point layout example: (2d) x0y0, x0y1, x1y0, x1y1; 205 | (3d) x0y0z0, x0y0z1, x0y1z0, x0y1z1, x1y0z0, x1y0z1, x1y1z0, x1y1z1 206 | where x0 < x1, y0 < y1, z0 < z1 207 | """ 208 | ndim = int(dims.shape[1]) 209 | corners_norm = np.stack( 210 | np.unravel_index(np.arange(2**ndim), [2] * ndim), 211 | axis=1).astype(dims.dtype) 212 | # now corners_norm has format: (2d) x0y0, x0y1, x1y0, x1y1 213 | # (3d) x0y0z0, x0y0z1, x0y1z0, x0y1z1, x1y0z0, x1y0z1, x1y1z0, x1y1z1 214 | # so need to convert to a format which is convenient to do other computing. 215 | # for 2d boxes, format is clockwise start with minimum point 216 | # for 3d boxes, please draw lines by your hand. 217 | if ndim == 2: 218 | # generate clockwise box corners 219 | corners_norm = corners_norm[[0, 1, 3, 2]] 220 | elif ndim == 3: 221 | corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]] 222 | corners_norm = corners_norm - np.array(origin, dtype=dims.dtype) 223 | corners = dims.reshape([-1, 1, ndim]) * corners_norm.reshape( 224 | [1, 2**ndim, ndim]) 225 | return corners 226 | 227 | 228 | @numba.njit 229 | def corners_2d_jit(dims, origin=0.5): 230 | ndim = 2 231 | corners_norm = np.array([[0, 0], [0, 1], [1, 1], [1, 0]], dtype=dims.dtype) 232 | corners_norm = corners_norm - np.array(origin, dtype=dims.dtype) 233 | corners = dims.reshape((-1, 1, ndim)) * corners_norm.reshape( 234 | (1, 2**ndim, ndim)) 235 | return corners 236 | 237 | 238 | @numba.njit 239 | def corners_3d_jit(dims, origin=0.5): 240 | ndim = 3 241 | corners_norm = np.array([ 242 | 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 1, 1, 1, 0, 0, 1, 0, 1, 1, 1, 0, 1, 1, 1 243 | ], 244 | dtype=dims.dtype).reshape((8, 3)) 245 | corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]] 246 | corners_norm = corners_norm - np.array(origin, dtype=dims.dtype) 247 | corners = dims.reshape((-1, 1, ndim)) * corners_norm.reshape( 248 | (1, 2**ndim, ndim)) 249 | return corners 250 | 251 | 252 | @numba.njit 253 | def corner_to_standup_nd_jit(boxes_corner): 254 | num_boxes = boxes_corner.shape[0] 255 | ndim = boxes_corner.shape[-1] 256 | result = np.zeros((num_boxes, ndim * 2), dtype=boxes_corner.dtype) 257 | for i in range(num_boxes): 258 | for j in range(ndim): 259 | result[i, j] = np.min(boxes_corner[i, :, j]) 260 | for j in range(ndim): 261 | result[i, j + ndim] = np.max(boxes_corner[i, :, j]) 262 | return result 263 | 264 | 265 | def corner_to_standup_nd(boxes_corner): 266 | assert len(boxes_corner.shape) == 3 267 | standup_boxes = [] 268 | standup_boxes.append(np.min(boxes_corner, axis=1)) 269 | standup_boxes.append(np.max(boxes_corner, axis=1)) 270 | return np.concatenate(standup_boxes, -1) 271 | 272 | 273 | def rbbox2d_to_near_bbox(rbboxes): 274 | """convert rotated bbox to nearest 'standing' or 'lying' bbox. 275 | Args: 276 | rbboxes: [N, 5(x, y, xdim, ydim, rad)] rotated bboxes 277 | Returns: 278 | bboxes: [N, 4(xmin, ymin, xmax, ymax)] bboxes 279 | """ 280 | rots = rbboxes[..., -1] 281 | rots_0_pi_div_2 = np.abs(limit_period(rots, 0.5, np.pi)) 282 | cond = (rots_0_pi_div_2 > np.pi / 4)[..., np.newaxis] 283 | bboxes_center = np.where(cond, rbboxes[:, [0, 1, 3, 2]], rbboxes[:, :4]) 284 | bboxes = center_to_minmax_2d(bboxes_center[:, :2], bboxes_center[:, 2:]) 285 | return bboxes 286 | 287 | 288 | def rotation_3d_in_axis(points, angles, axis=0): 289 | if len(points.shape) == 2: 290 | points = np.expand_dims(points, axis=0) 291 | # points: [N, point_size, 3] 292 | rot_sin = np.sin(angles) 293 | rot_cos = np.cos(angles) 294 | ones = np.ones_like(rot_cos) 295 | zeros = np.zeros_like(rot_cos) 296 | if axis == 1: 297 | rot_mat_T = np.stack([[rot_cos, zeros, -rot_sin], [zeros, ones, zeros], 298 | [rot_sin, zeros, rot_cos]]) 299 | elif axis == 2 or axis == -1: 300 | rot_mat_T = np.stack([[rot_cos, -rot_sin, zeros], 301 | [rot_sin, rot_cos, zeros], [zeros, zeros, ones]]) 302 | elif axis == 0: 303 | rot_mat_T = np.stack([[zeros, rot_cos, -rot_sin], 304 | [zeros, rot_sin, rot_cos], [ones, zeros, zeros]]) 305 | else: 306 | raise ValueError("axis should in range") 307 | 308 | return np.einsum('aij,jka->aik', points, rot_mat_T) 309 | 310 | 311 | def rotation_points_single_angle(points, angle, axis=0): 312 | # points: [N, 3] 313 | rot_sin = np.sin(angle) 314 | rot_cos = np.cos(angle) 315 | if axis == 1: 316 | rot_mat_T = np.array( 317 | [[rot_cos, 0, -rot_sin], [0, 1, 0], [rot_sin, 0, rot_cos]], 318 | dtype=points.dtype) 319 | elif axis == 2 or axis == -1: 320 | rot_mat_T = np.array( 321 | [[rot_cos, -rot_sin, 0], [rot_sin, rot_cos, 0], [0, 0, 1]], 322 | dtype=points.dtype) 323 | elif axis == 0: 324 | rot_mat_T = np.array( 325 | [[1, 0, 0], [0, rot_cos, -rot_sin], [0, rot_sin, rot_cos]], 326 | dtype=points.dtype) 327 | else: 328 | raise ValueError("axis should in range") 329 | 330 | return points @ rot_mat_T 331 | 332 | 333 | def rotation_2d(points, angles): 334 | """rotation 2d points based on origin point clockwise when angle positive. 335 | 336 | Args: 337 | points (float array, shape=[N, point_size, 2]): points to be rotated. 338 | angles (float array, shape=[N]): rotation angle. 339 | 340 | Returns: 341 | float array: same shape as points 342 | """ 343 | rot_sin = np.sin(angles) 344 | rot_cos = np.cos(angles) 345 | rot_mat_T = np.stack([[rot_cos, -rot_sin], [rot_sin, rot_cos]]) 346 | return np.einsum('aij,jka->aik', points, rot_mat_T) 347 | 348 | 349 | def rotation_box(box_corners, angle): 350 | """rotation 2d points based on origin point clockwise when angle positive. 351 | 352 | Args: 353 | points (float array, shape=[N, point_size, 2]): points to be rotated. 354 | angle (float): rotation angle. 355 | 356 | Returns: 357 | float array: same shape as points 358 | """ 359 | rot_sin = np.sin(angle) 360 | rot_cos = np.cos(angle) 361 | rot_mat_T = np.array([[rot_cos, -rot_sin], [rot_sin, rot_cos]], 362 | dtype=box_corners.dtype) 363 | return box_corners @ rot_mat_T 364 | 365 | 366 | def center_to_corner_box3d(centers, 367 | dims, 368 | angles=None, 369 | origin=(0.5, 0.5, 0.5), 370 | axis=2): 371 | """convert kitti locations, dimensions and angles to corners 372 | 373 | Args: 374 | centers (float array, shape=[N, 3]): locations in kitti label file. 375 | dims (float array, shape=[N, 3]): dimensions in kitti label file. 376 | angles (float array, shape=[N]): rotation_y in kitti label file. 377 | origin (list or array or float): origin point relate to smallest point. 378 | use [0.5, 1.0, 0.5] in camera and [0.5, 0.5, 0] in lidar. 379 | axis (int): rotation axis. 1 for camera and 2 for lidar. 380 | Returns: 381 | [type]: [description] 382 | """ 383 | # 'length' in kitti format is in x axis. 384 | # yzx(hwl)(kitti label file)<->xyz(lhw)(camera)<->z(-x)(-y)(wlh)(lidar) 385 | # center in kitti format is [0.5, 1.0, 0.5] in xyz. 386 | corners = corners_nd(dims, origin=origin) 387 | # corners: [N, 8, 3] 388 | if angles is not None: 389 | corners = rotation_3d_in_axis(corners, angles, axis=axis) 390 | corners += centers.reshape([-1, 1, 3]) 391 | return corners 392 | 393 | 394 | def center_to_corner_box2d(centers, dims, angles=None, origin=0.5): 395 | """convert kitti locations, dimensions and angles to corners. 396 | format: center(xy), dims(xy), angles(clockwise when positive) 397 | 398 | Args: 399 | centers (float array, shape=[N, 2]): locations in kitti label file. 400 | dims (float array, shape=[N, 2]): dimensions in kitti label file. 401 | angles (float array, shape=[N]): rotation_y in kitti label file. 402 | 403 | Returns: 404 | [type]: [description] 405 | """ 406 | # 'length' in kitti format is in x axis. 407 | # xyz(hwl)(kitti label file)<->xyz(lhw)(camera)<->z(-x)(-y)(wlh)(lidar) 408 | # center in kitti format is [0.5, 1.0, 0.5] in xyz. 409 | corners = corners_nd(dims, origin=origin) 410 | # corners: [N, 4, 2] 411 | if angles is not None: 412 | corners = rotation_2d(corners, angles) 413 | corners += centers.reshape([-1, 1, 2]) 414 | return corners 415 | 416 | 417 | @numba.jit(nopython=True) 418 | def box2d_to_corner_jit(boxes): 419 | num_box = boxes.shape[0] 420 | corners_norm = np.zeros((4, 2), dtype=boxes.dtype) 421 | corners_norm[1, 1] = 1.0 422 | corners_norm[2] = 1.0 423 | corners_norm[3, 0] = 1.0 424 | corners_norm -= np.array([0.5, 0.5], dtype=boxes.dtype) 425 | corners = boxes.reshape(num_box, 1, 5)[:, :, 2:4] * corners_norm.reshape( 426 | 1, 4, 2) 427 | rot_mat_T = np.zeros((2, 2), dtype=boxes.dtype) 428 | box_corners = np.zeros((num_box, 4, 2), dtype=boxes.dtype) 429 | for i in range(num_box): 430 | rot_sin = np.sin(boxes[i, -1]) 431 | rot_cos = np.cos(boxes[i, -1]) 432 | rot_mat_T[0, 0] = rot_cos 433 | rot_mat_T[0, 1] = -rot_sin 434 | rot_mat_T[1, 0] = rot_sin 435 | rot_mat_T[1, 1] = rot_cos 436 | box_corners[i] = corners[i] @ rot_mat_T + boxes[i, :2] 437 | return box_corners 438 | 439 | 440 | def rbbox3d_to_corners(rbboxes, origin=[0.5, 0.5, 0.5], axis=2): 441 | return center_to_corner_box3d( 442 | rbboxes[..., :3], 443 | rbboxes[..., 3:6], 444 | rbboxes[..., 6], 445 | origin, 446 | axis=axis) 447 | 448 | 449 | def rbbox3d_to_bev_corners(rbboxes, origin=0.5): 450 | return center_to_corner_box2d(rbboxes[..., :2], rbboxes[..., 3:5], 451 | rbboxes[..., 6], origin) 452 | 453 | 454 | def minmax_to_corner_2d(minmax_box): 455 | ndim = minmax_box.shape[-1] // 2 456 | center = minmax_box[..., :ndim] 457 | dims = minmax_box[..., ndim:] - center 458 | return center_to_corner_box2d(center, dims, origin=0.0) 459 | 460 | 461 | def minmax_to_corner_2d_v2(minmax_box): 462 | # N, 4 -> N 4 2 463 | return minmax_box[..., [0, 1, 0, 3, 2, 3, 2, 1]].reshape(-1, 4, 2) 464 | 465 | 466 | def minmax_to_corner_3d(minmax_box): 467 | ndim = minmax_box.shape[-1] // 2 468 | center = minmax_box[..., :ndim] 469 | dims = minmax_box[..., ndim:] - center 470 | return center_to_corner_box3d(center, dims, origin=0.0) 471 | 472 | 473 | def minmax_to_center_2d(minmax_box): 474 | ndim = minmax_box.shape[-1] // 2 475 | center_min = minmax_box[..., :ndim] 476 | dims = minmax_box[..., ndim:] - center_min 477 | center = center_min + 0.5 * dims 478 | return np.concatenate([center, dims], axis=-1) 479 | 480 | 481 | def center_to_minmax_2d_0_5(centers, dims): 482 | return np.concatenate([centers - dims / 2, centers + dims / 2], axis=-1) 483 | 484 | 485 | def center_to_minmax_2d(centers, dims, origin=0.5): 486 | if origin == 0.5: 487 | return center_to_minmax_2d_0_5(centers, dims) 488 | corners = center_to_corner_box2d(centers, dims, origin=origin) 489 | return corners[:, [0, 2]].reshape([-1, 4]) 490 | 491 | 492 | def limit_period(val, offset=0.5, period=np.pi): 493 | return val - np.floor(val / period + offset) * period 494 | 495 | 496 | def projection_matrix_to_CRT_kitti(proj): 497 | # P = C @ [R|T] 498 | # C is upper triangular matrix, so we need to inverse CR and use QR 499 | # stable for all kitti camera projection matrix 500 | CR = proj[0:3, 0:3] 501 | CT = proj[0:3, 3] 502 | RinvCinv = np.linalg.inv(CR) 503 | Rinv, Cinv = np.linalg.qr(RinvCinv) 504 | C = np.linalg.inv(Cinv) 505 | R = np.linalg.inv(Rinv) 506 | T = Cinv @ CT 507 | return C, R, T 508 | 509 | 510 | def get_frustum(bbox_image, C, near_clip=0.001, far_clip=100): 511 | fku = C[0, 0] 512 | fkv = -C[1, 1] 513 | u0v0 = C[0:2, 2] 514 | z_points = np.array( 515 | [near_clip] * 4 + [far_clip] * 4, dtype=C.dtype)[:, np.newaxis] 516 | b = bbox_image 517 | box_corners = np.array( 518 | [[b[0], b[1]], [b[0], b[3]], [b[2], b[3]], [b[2], b[1]]], 519 | dtype=C.dtype) 520 | near_box_corners = (box_corners - u0v0) / np.array( 521 | [fku / near_clip, -fkv / near_clip], dtype=C.dtype) 522 | far_box_corners = (box_corners - u0v0) / np.array( 523 | [fku / far_clip, -fkv / far_clip], dtype=C.dtype) 524 | ret_xy = np.concatenate([near_box_corners, far_box_corners], 525 | axis=0) # [8, 2] 526 | ret_xyz = np.concatenate([ret_xy, z_points], axis=1) 527 | return ret_xyz 528 | 529 | 530 | def get_frustum_v2(bboxes, C, near_clip=0.001, far_clip=100): 531 | fku = C[0, 0] 532 | fkv = -C[1, 1] 533 | u0v0 = C[0:2, 2] 534 | num_box = bboxes.shape[0] 535 | z_points = np.array( 536 | [near_clip] * 4 + [far_clip] * 4, 537 | dtype=C.dtype)[np.newaxis, :, np.newaxis] 538 | z_points = np.tile(z_points, [num_box, 1, 1]) 539 | box_corners = minmax_to_corner_2d_v2(bboxes) 540 | near_box_corners = (box_corners - u0v0) / np.array( 541 | [fku / near_clip, -fkv / near_clip], dtype=C.dtype) 542 | far_box_corners = (box_corners - u0v0) / np.array( 543 | [fku / far_clip, -fkv / far_clip], dtype=C.dtype) 544 | ret_xy = np.concatenate([near_box_corners, far_box_corners], 545 | axis=1) # [8, 2] 546 | ret_xyz = np.concatenate([ret_xy, z_points], axis=-1) 547 | return ret_xyz 548 | 549 | 550 | def create_anchors_3d_stride(feature_size, 551 | sizes=[1.6, 3.9, 1.56], 552 | anchor_strides=[0.4, 0.4, 0.0], 553 | anchor_offsets=[0.2, -39.8, -1.78], 554 | rotations=[0, np.pi / 2], 555 | dtype=np.float32): 556 | """ 557 | Args: 558 | feature_size: list [D, H, W](zyx) 559 | sizes: [N, 3] list of list or array, size of anchors, xyz 560 | 561 | Returns: 562 | anchors: [*feature_size, num_sizes, num_rots, 7] tensor. 563 | """ 564 | # almost 2x faster than v1 565 | x_stride, y_stride, z_stride = anchor_strides 566 | x_offset, y_offset, z_offset = anchor_offsets 567 | z_centers = np.arange(feature_size[0], dtype=dtype) 568 | y_centers = np.arange(feature_size[1], dtype=dtype) 569 | x_centers = np.arange(feature_size[2], dtype=dtype) 570 | z_centers = z_centers * z_stride + z_offset 571 | y_centers = y_centers * y_stride + y_offset 572 | x_centers = x_centers * x_stride + x_offset 573 | sizes = np.reshape(np.array(sizes, dtype=dtype), [-1, 3]) 574 | rotations = np.array(rotations, dtype=dtype) 575 | rets = np.meshgrid( 576 | x_centers, y_centers, z_centers, rotations, indexing='ij') 577 | tile_shape = [1] * 5 578 | tile_shape[-2] = int(sizes.shape[0]) 579 | for i in range(len(rets)): 580 | rets[i] = np.tile(rets[i][..., np.newaxis, :], tile_shape) 581 | rets[i] = rets[i][..., np.newaxis] # for concat 582 | sizes = np.reshape(sizes, [1, 1, 1, -1, 1, 3]) 583 | tile_size_shape = list(rets[0].shape) 584 | tile_size_shape[3] = 1 585 | sizes = np.tile(sizes, tile_size_shape) 586 | rets.insert(3, sizes) 587 | ret = np.concatenate(rets, axis=-1) 588 | return np.transpose(ret, [2, 1, 0, 3, 4, 5]) 589 | 590 | 591 | def create_anchors_3d_range(feature_size, 592 | anchor_range, 593 | sizes=[1.6, 3.9, 1.56], 594 | rotations=[0, np.pi / 2], 595 | dtype=np.float32): 596 | """ 597 | Args: 598 | feature_size: list [D, H, W](zyx) 599 | sizes: [N, 3] list of list or array, size of anchors, xyz 600 | 601 | Returns: 602 | anchors: [*feature_size, num_sizes, num_rots, 7] tensor. 603 | """ 604 | anchor_range = np.array(anchor_range, dtype) 605 | z_centers = np.linspace( 606 | anchor_range[2], anchor_range[5], feature_size[0], dtype=dtype) 607 | y_centers = np.linspace( 608 | anchor_range[1], anchor_range[4], feature_size[1], dtype=dtype) 609 | x_centers = np.linspace( 610 | anchor_range[0], anchor_range[3], feature_size[2], dtype=dtype) 611 | sizes = np.reshape(np.array(sizes, dtype=dtype), [-1, 3]) 612 | rotations = np.array(rotations, dtype=dtype) 613 | rets = np.meshgrid( 614 | x_centers, y_centers, z_centers, rotations, indexing='ij') 615 | tile_shape = [1] * 5 616 | tile_shape[-2] = int(sizes.shape[0]) 617 | for i in range(len(rets)): 618 | rets[i] = np.tile(rets[i][..., np.newaxis, :], tile_shape) 619 | rets[i] = rets[i][..., np.newaxis] # for concat 620 | sizes = np.reshape(sizes, [1, 1, 1, -1, 1, 3]) 621 | tile_size_shape = list(rets[0].shape) 622 | tile_size_shape[3] = 1 623 | sizes = np.tile(sizes, tile_size_shape) 624 | rets.insert(3, sizes) 625 | ret = np.concatenate(rets, axis=-1) 626 | res = np.transpose(ret, [2, 1, 0, 3, 4, 5]) 627 | return res 628 | 629 | 630 | def project_to_image(points_3d, proj_mat): 631 | points_shape = list(points_3d.shape) 632 | points_shape[-1] = 1 633 | points_4 = np.concatenate([points_3d, np.zeros(points_shape)], axis=-1) 634 | point_2d = points_4 @ proj_mat.T 635 | point_2d_res = point_2d[..., :2] / point_2d[..., 2:3] 636 | return point_2d_res 637 | 638 | 639 | def camera_to_lidar(points, r_rect, velo2cam): 640 | points_shape = list(points.shape[0:-1]) 641 | if points.shape[-1] == 3: 642 | points = np.concatenate([points, np.ones(points_shape + [1])], axis=-1) 643 | lidar_points = points @ np.linalg.inv((r_rect @ velo2cam).T) 644 | return lidar_points[..., :3] 645 | 646 | 647 | def lidar_to_camera(points, r_rect, velo2cam): 648 | points_shape = list(points.shape[:-1]) 649 | if points.shape[-1] == 3: 650 | points = np.concatenate([points, np.ones(points_shape + [1])], axis=-1) 651 | camera_points = points @ (r_rect @ velo2cam).T 652 | return camera_points[..., :3] 653 | 654 | 655 | def box_camera_to_lidar(data, r_rect, velo2cam): 656 | xyz = data[:, 0:3] 657 | l, h, w = data[:, 3:4], data[:, 4:5], data[:, 5:6] 658 | r = data[:, 6:7] 659 | xyz_lidar = camera_to_lidar(xyz, r_rect, velo2cam) 660 | return np.concatenate([xyz_lidar, w, l, h, r], axis=1) 661 | 662 | 663 | def box_lidar_to_camera(data, r_rect, velo2cam): 664 | xyz_lidar = data[:, 0:3] 665 | w, l, h = data[:, 3:4], data[:, 4:5], data[:, 5:6] 666 | r = data[:, 6:7] 667 | xyz = lidar_to_camera(xyz_lidar, r_rect, velo2cam) 668 | return np.concatenate([xyz, l, h, w, r], axis=1) 669 | 670 | 671 | def remove_outside_points(points, rect, Trv2c, P2, image_shape): 672 | # 5x faster than remove_outside_points_v1(2ms vs 10ms) 673 | C, R, T = projection_matrix_to_CRT_kitti(P2) 674 | image_bbox = [0, 0, image_shape[1], image_shape[0]] 675 | frustum = get_frustum(image_bbox, C) 676 | frustum -= T 677 | frustum = np.linalg.inv(R) @ frustum.T 678 | frustum = camera_to_lidar(frustum.T, rect, Trv2c) 679 | frustum_surfaces = corner_to_surfaces_3d_jit(frustum[np.newaxis, ...]) 680 | indices = points_in_convex_polygon_3d_jit(points[:, :3], frustum_surfaces) 681 | points = points[indices.reshape([-1])] 682 | return points 683 | 684 | 685 | @numba.jit(nopython=True) 686 | def iou_jit(boxes, query_boxes, eps=1.0): 687 | """calculate box iou. note that jit version runs 2x faster than cython in 688 | my machine! 689 | Parameters 690 | ---------- 691 | boxes: (N, 4) ndarray of float 692 | query_boxes: (K, 4) ndarray of float 693 | Returns 694 | ------- 695 | overlaps: (N, K) ndarray of overlap between boxes and query_boxes 696 | """ 697 | N = boxes.shape[0] 698 | K = query_boxes.shape[0] 699 | overlaps = np.zeros((N, K), dtype=boxes.dtype) 700 | for k in range(K): 701 | box_area = ((query_boxes[k, 2] - query_boxes[k, 0] + eps) * 702 | (query_boxes[k, 3] - query_boxes[k, 1] + eps)) 703 | for n in range(N): 704 | iw = (min(boxes[n, 2], query_boxes[k, 2]) - max( 705 | boxes[n, 0], query_boxes[k, 0]) + eps) 706 | if iw > 0: 707 | ih = (min(boxes[n, 3], query_boxes[k, 3]) - max( 708 | boxes[n, 1], query_boxes[k, 1]) + eps) 709 | if ih > 0: 710 | ua = ( 711 | (boxes[n, 2] - boxes[n, 0] + eps) * 712 | (boxes[n, 3] - boxes[n, 1] + eps) + box_area - iw * ih) 713 | overlaps[n, k] = iw * ih / ua 714 | return overlaps 715 | 716 | 717 | def points_in_rbbox(points, rbbox, z_axis=2, origin=(0.5, 0.5, 0.5)): 718 | rbbox_corners = center_to_corner_box3d( 719 | rbbox[:, :3], rbbox[:, 3:6], rbbox[:, 6], origin=origin, axis=z_axis) 720 | surfaces = corner_to_surfaces_3d(rbbox_corners) 721 | indices = points_in_convex_polygon_3d_jit(points[:, :3], surfaces) 722 | return indices 723 | 724 | def points_in_corners(points, corners): 725 | surfaces = corner_to_surfaces_3d(corners) 726 | indices = points_in_convex_polygon_3d_jit(points[:, :3], surfaces) 727 | return indices 728 | 729 | def points_corners_in_rbbox(points, rbbox, z_axis=2, origin=(0.5, 0.5, 0.5)): 730 | rbbox_corners = center_to_corner_box3d( 731 | rbbox[:, :3], rbbox[:, 3:6], rbbox[:, 6], origin=origin, axis=z_axis) 732 | surfaces = corner_to_surfaces_3d(rbbox_corners) 733 | indices = points_in_convex_polygon_3d_jit(points[:, :3], surfaces) 734 | return indices, rbbox_corners 735 | 736 | def points_count_rbbox(points, rbbox, z_axis=2, origin=(0.5, 0.5, 0.5)): 737 | rbbox_corners = center_to_corner_box3d( 738 | rbbox[:, :3], rbbox[:, 3:6], rbbox[:, 6], origin=origin, axis=z_axis) 739 | surfaces = corner_to_surfaces_3d(rbbox_corners) 740 | return points_count_convex_polygon_3d_jit(points[:, :3], surfaces) 741 | 742 | 743 | def corner_to_surfaces_3d(corners): 744 | """convert 3d box corners from corner function above 745 | to surfaces that normal vectors all direct to internal. 746 | 747 | Args: 748 | corners (float array, [N, 8, 3]): 3d box corners. 749 | Returns: 750 | surfaces (float array, [N, 6, 4, 3]): 751 | """ 752 | # box_corners: [N, 8, 3], must from corner functions in this module 753 | surfaces = np.array([ 754 | [corners[:, 0], corners[:, 1], corners[:, 2], corners[:, 3]], 755 | [corners[:, 7], corners[:, 6], corners[:, 5], corners[:, 4]], 756 | [corners[:, 0], corners[:, 3], corners[:, 7], corners[:, 4]], 757 | [corners[:, 1], corners[:, 5], corners[:, 6], corners[:, 2]], 758 | [corners[:, 0], corners[:, 4], corners[:, 5], corners[:, 1]], 759 | [corners[:, 3], corners[:, 2], corners[:, 6], corners[:, 7]], 760 | ]).transpose([2, 0, 1, 3]) 761 | return surfaces 762 | 763 | 764 | @numba.jit(nopython=True) 765 | def corner_to_surfaces_3d_jit(corners): 766 | """convert 3d box corners from corner function above 767 | to surfaces that normal vectors all direct to internal. 768 | 769 | Args: 770 | corners (float array, [N, 8, 3]): 3d box corners. 771 | Returns: 772 | surfaces (float array, [N, 6, 4, 3]): 773 | """ 774 | # box_corners: [N, 8, 3], must from corner functions in this module 775 | num_boxes = corners.shape[0] 776 | surfaces = np.zeros((num_boxes, 6, 4, 3), dtype=corners.dtype) 777 | corner_idxes = np.array([ 778 | 0, 1, 2, 3, 7, 6, 5, 4, 0, 3, 7, 4, 1, 5, 6, 2, 0, 4, 5, 1, 3, 2, 6, 7 779 | ]).reshape(6, 4) 780 | for i in range(num_boxes): 781 | for j in range(6): 782 | for k in range(4): 783 | surfaces[i, j, k] = corners[i, corner_idxes[j, k]] 784 | return surfaces 785 | 786 | 787 | def assign_label_to_voxel(gt_boxes, coors, voxel_size, coors_range): 788 | """assign a 0/1 label to each voxel based on whether 789 | the center of voxel is in gt_box. LIDAR. 790 | """ 791 | voxel_size = np.array(voxel_size, dtype=gt_boxes.dtype) 792 | coors_range = np.array(coors_range, dtype=gt_boxes.dtype) 793 | shift = coors_range[:3] 794 | voxel_origins = coors[:, ::-1] * voxel_size + shift 795 | voxel_centers = voxel_origins + voxel_size * 0.5 796 | gt_box_corners = center_to_corner_box3d( 797 | gt_boxes[:, :3] - voxel_size * 0.5, 798 | gt_boxes[:, 3:6] + voxel_size, 799 | gt_boxes[:, 6], 800 | origin=[0.5, 0.5, 0.5], 801 | axis=2) 802 | gt_surfaces = corner_to_surfaces_3d(gt_box_corners) 803 | ret = points_in_convex_polygon_3d_jit(voxel_centers, gt_surfaces) 804 | return np.any(ret, axis=1).astype(np.int64) 805 | 806 | 807 | def assign_label_to_voxel_v3(gt_boxes, coors, voxel_size, coors_range): 808 | """assign a 0/1 label to each voxel based on whether 809 | the center of voxel is in gt_box. LIDAR. 810 | """ 811 | voxel_size = np.array(voxel_size, dtype=gt_boxes.dtype) 812 | coors_range = np.array(coors_range, dtype=gt_boxes.dtype) 813 | shift = coors_range[:3] 814 | voxel_origins = coors[:, ::-1] * voxel_size + shift 815 | voxel_maxes = voxel_origins + voxel_size 816 | voxel_minmax = np.concatenate([voxel_origins, voxel_maxes], axis=-1) 817 | voxel_corners = minmax_to_corner_3d(voxel_minmax) 818 | gt_box_corners = center_to_corner_box3d( 819 | gt_boxes[:, :3], 820 | gt_boxes[:, 3:6], 821 | gt_boxes[:, 6], 822 | origin=[0.5, 0.5, 0.5], 823 | axis=2) 824 | gt_surfaces = corner_to_surfaces_3d(gt_box_corners) 825 | voxel_corners_flat = voxel_corners.reshape([-1, 3]) 826 | ret = points_in_convex_polygon_3d_jit(voxel_corners_flat, gt_surfaces) 827 | ret = ret.reshape([-1, 8, ret.shape[-1]]) 828 | return ret.any(-1).any(-1).astype(np.int64) 829 | 830 | 831 | def image_box_region_area(img_cumsum, bbox): 832 | """check a 2d voxel is contained by a box. used to filter empty 833 | anchors. 834 | Summed-area table algorithm: 835 | ==> W 836 | ------------------ 837 | | | | 838 | |------A---------B 839 | | | | 840 | | | | 841 | |----- C---------D 842 | Iabcd = ID-IB-IC+IA 843 | Args: 844 | img_cumsum: [M, H, W](yx) cumsumed image. 845 | bbox: [N, 4](xyxy) bounding box, 846 | """ 847 | N = bbox.shape[0] 848 | M = img_cumsum.shape[0] 849 | ret = np.zeros([N, M], dtype=img_cumsum.dtype) 850 | ID = img_cumsum[:, bbox[:, 3], bbox[:, 2]] 851 | IA = img_cumsum[:, bbox[:, 1], bbox[:, 0]] 852 | IB = img_cumsum[:, bbox[:, 3], bbox[:, 0]] 853 | IC = img_cumsum[:, bbox[:, 1], bbox[:, 2]] 854 | ret = ID - IB - IC + IA 855 | return ret 856 | 857 | 858 | def get_minimum_bounding_box_bv(points, 859 | voxel_size, 860 | bound, 861 | downsample=8, 862 | margin=1.6): 863 | x_vsize = voxel_size[0] 864 | y_vsize = voxel_size[1] 865 | max_x = points[:, 0].max() 866 | max_y = points[:, 1].max() 867 | min_x = points[:, 0].min() 868 | min_y = points[:, 1].min() 869 | max_x = np.floor(max_x / 870 | (x_vsize * downsample) + 1) * (x_vsize * downsample) 871 | max_y = np.floor(max_y / 872 | (y_vsize * downsample) + 1) * (y_vsize * downsample) 873 | min_x = np.floor(min_x / (x_vsize * downsample)) * (x_vsize * downsample) 874 | min_y = np.floor(min_y / (y_vsize * downsample)) * (y_vsize * downsample) 875 | max_x = np.minimum(max_x + margin, bound[2]) 876 | max_y = np.minimum(max_y + margin, bound[3]) 877 | min_x = np.maximum(min_x - margin, bound[0]) 878 | min_y = np.maximum(min_y - margin, bound[1]) 879 | return np.array([min_x, min_y, max_x, max_y]) 880 | 881 | 882 | @numba.jit(nopython=True) 883 | def get_anchor_bv_in_feature_jit(anchors_bv, voxel_size, coors_range, 884 | grid_size): 885 | anchors_bv_coors = np.zeros(anchors_bv.shape, dtype=np.int32) 886 | anchor_coor = np.zeros(anchors_bv.shape[1:], dtype=np.int32) 887 | grid_size_x = grid_size[0] - 1 888 | grid_size_y = grid_size[1] - 1 889 | for i in range(anchors_bv.shape[0]): 890 | anchor_coor[0] = np.floor( 891 | (anchors_bv[i, 0] - coors_range[0]) / voxel_size[0]) 892 | anchor_coor[1] = np.floor( 893 | (anchors_bv[i, 1] - coors_range[1]) / voxel_size[1]) 894 | anchor_coor[2] = np.floor( 895 | (anchors_bv[i, 2] - coors_range[0]) / voxel_size[0]) 896 | anchor_coor[3] = np.floor( 897 | (anchors_bv[i, 3] - coors_range[1]) / voxel_size[1]) 898 | anchor_coor[0] = max(anchor_coor[0], 0) 899 | anchor_coor[1] = max(anchor_coor[1], 0) 900 | anchor_coor[2] = min(anchor_coor[2], grid_size_x) 901 | anchor_coor[3] = min(anchor_coor[3], grid_size_y) 902 | anchors_bv_coors[i] = anchor_coor 903 | return anchors_bv_coors 904 | 905 | 906 | def get_anchor_bv_in_feature(anchors_bv, voxel_size, coors_range, grid_size): 907 | vsize_bv = np.tile(voxel_size[:2], 2) 908 | anchors_bv[..., [1, 3]] -= coors_range[1] 909 | anchors_bv_coors = np.floor(anchors_bv / vsize_bv).astype(np.int64) 910 | anchors_bv_coors[..., [0, 2]] = np.clip( 911 | anchors_bv_coors[..., [0, 2]], a_max=grid_size[0] - 1, a_min=0) 912 | anchors_bv_coors[..., [1, 3]] = np.clip( 913 | anchors_bv_coors[..., [1, 3]], a_max=grid_size[1] - 1, a_min=0) 914 | anchors_bv_coors = anchors_bv_coors.reshape([-1, 4]) 915 | return anchors_bv_coors 916 | 917 | 918 | @numba.jit(nopython=True) 919 | def sparse_sum_for_anchors_mask(coors, shape): 920 | ret = np.zeros(shape, dtype=np.float32) 921 | for i in range(coors.shape[0]): 922 | ret[coors[i, 1], coors[i, 2]] += 1 923 | return ret 924 | 925 | 926 | @numba.jit(nopython=True) 927 | def fused_get_anchors_area(dense_map, anchors_bv, stride, offset, grid_size): 928 | anchor_coor = np.zeros(anchors_bv.shape[1:], dtype=np.int32) 929 | grid_size_x = grid_size[0] - 1 930 | grid_size_y = grid_size[1] - 1 931 | N = anchors_bv.shape[0] 932 | ret = np.zeros((N), dtype=dense_map.dtype) 933 | for i in range(N): 934 | anchor_coor[0] = np.floor((anchors_bv[i, 0] - offset[0]) / stride[0]) 935 | anchor_coor[1] = np.floor((anchors_bv[i, 1] - offset[1]) / stride[1]) 936 | anchor_coor[2] = np.floor((anchors_bv[i, 2] - offset[0]) / stride[0]) 937 | anchor_coor[3] = np.floor((anchors_bv[i, 3] - offset[1]) / stride[1]) 938 | anchor_coor[0] = max(anchor_coor[0], 0) 939 | anchor_coor[1] = max(anchor_coor[1], 0) 940 | anchor_coor[2] = min(anchor_coor[2], grid_size_x) 941 | anchor_coor[3] = min(anchor_coor[3], grid_size_y) 942 | ID = dense_map[anchor_coor[3], anchor_coor[2]] 943 | IA = dense_map[anchor_coor[1], anchor_coor[0]] 944 | IB = dense_map[anchor_coor[3], anchor_coor[0]] 945 | IC = dense_map[anchor_coor[1], anchor_coor[2]] 946 | ret[i] = ID - IB - IC + IA 947 | return ret 948 | 949 | 950 | @numba.jit(nopython=True) 951 | def distance_similarity(points, 952 | qpoints, 953 | dist_norm, 954 | with_rotation=False, 955 | rot_alpha=0.5): 956 | N = points.shape[0] 957 | K = qpoints.shape[0] 958 | dists = np.zeros((N, K), dtype=points.dtype) 959 | rot_alpha_1 = 1 - rot_alpha 960 | for k in range(K): 961 | for n in range(N): 962 | if np.abs(points[n, 0] - qpoints[k, 0]) <= dist_norm: 963 | if np.abs(points[n, 1] - qpoints[k, 1]) <= dist_norm: 964 | dist = np.sum((points[n, :2] - qpoints[k, :2])**2) 965 | dist_normed = min(dist / dist_norm, dist_norm) 966 | if with_rotation: 967 | dist_rot = np.abs( 968 | np.sin(points[n, -1] - qpoints[k, -1])) 969 | dists[ 970 | n, 971 | k] = 1 - rot_alpha_1 * dist_normed - rot_alpha * dist_rot 972 | else: 973 | dists[n, k] = 1 - dist_normed 974 | return dists 975 | 976 | 977 | def box3d_to_bbox(box3d, rect, Trv2c, P2): 978 | box3d_to_cam = box_lidar_to_camera(box3d, rect, Trv2c) 979 | box_corners = center_to_corner_box3d( 980 | box3d[:, :3], box3d[:, 3:6], box3d[:, 6], [0.5, 1.0, 0.5], axis=1) 981 | box_corners_in_image = project_to_image(box_corners, P2) 982 | # box_corners_in_image: [N, 8, 2] 983 | minxy = np.min(box_corners_in_image, axis=1) 984 | maxxy = np.max(box_corners_in_image, axis=1) 985 | bbox = np.concatenate([minxy, maxxy], axis=1) 986 | return bbox 987 | 988 | 989 | def change_box3d_center_(box3d, src, dst): 990 | dst = np.array(dst, dtype=box3d.dtype) 991 | src = np.array(src, dtype=box3d.dtype) 992 | box3d[..., :3] += box3d[..., 3:6] * (dst - src) 993 | 994 | 995 | -------------------------------------------------------------------------------- /docs/methods.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sky77764/pa-aug.pytorch/45b730c32673e2bde6b7d340042ef0e842a0cb5f/docs/methods.jpg -------------------------------------------------------------------------------- /geometry.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import numba 3 | 4 | 5 | @numba.njit 6 | def is_line_segment_intersection_jit(lines1, lines2): 7 | """check if line segments1 and line segments2 have cross point 8 | 9 | Args: 10 | lines1 (float, [N, 2, 2]): [description] 11 | lines2 (float, [M, 2, 2]): [description] 12 | 13 | Returns: 14 | [type]: [description] 15 | """ 16 | 17 | # Return true if line segments AB and CD intersect 18 | N = lines1.shape[0] 19 | M = lines2.shape[0] 20 | ret = np.zeros((N, M), dtype=np.bool_) 21 | for i in range(N): 22 | for j in range(M): 23 | A = lines1[i, 0] 24 | B = lines1[i, 1] 25 | C = lines2[j, 0] 26 | D = lines2[j, 1] 27 | acd = (D[1] - A[1]) * (C[0] - A[0]) > (C[1] - A[1]) * (D[0] - A[0]) 28 | bcd = (D[1] - B[1]) * (C[0] - B[0]) > (C[1] - B[1]) * (D[0] - B[0]) 29 | if acd != bcd: 30 | abc = (C[1] - A[1]) * (B[0] - A[0]) > (B[1] - A[1]) * (C[0] - A[0]) 31 | abd = (D[1] - A[1]) * (B[0] - A[0]) > (B[1] - A[1]) * (D[0] - A[0]) 32 | if abc != abd: 33 | ret[i, j] = True 34 | return ret 35 | 36 | @numba.njit 37 | def line_segment_intersection(line1, line2, intersection): 38 | A = line1[0] 39 | B = line1[1] 40 | C = line2[0] 41 | D = line2[1] 42 | BA0 = B[0] - A[0] 43 | BA1 = B[1] - A[1] 44 | DA0 = D[0] - A[0] 45 | CA0 = C[0] - A[0] 46 | DA1 = D[1] - A[1] 47 | CA1 = C[1] - A[1] 48 | acd = DA1 * CA0 > CA1 * DA0 49 | bcd = (D[1] - B[1]) * (C[0] - B[0]) > (C[1] - B[1]) * (D[0] - B[0]) 50 | if acd != bcd: 51 | abc = CA1 * BA0 > BA1 * CA0 52 | abd = DA1 * BA0 > BA1 * DA0 53 | if abc != abd: 54 | DC0 = D[0] - C[0] 55 | DC1 = D[1] - C[1] 56 | ABBA = A[0]*B[1] - B[0]*A[1] 57 | CDDC = C[0]*D[1] - D[0]*C[1] 58 | DH = BA1 * DC0 - BA0 * DC1 59 | intersection[0] = (ABBA * DC0 - BA0 * CDDC) / DH 60 | intersection[1] = (ABBA * DC1 - BA1 * CDDC) / DH 61 | return True 62 | return False 63 | 64 | 65 | 66 | def _ccw(A, B, C): 67 | return (C[..., 1] - A[..., 1]) * (B[..., 0] - A[..., 0]) > ( 68 | B[..., 1] - A[..., 1]) * (C[..., 0] - A[..., 0]) 69 | 70 | 71 | def is_line_segment_cross(lines1, lines2): 72 | # 10x slower than jit version with 1000-1000 random lines input. 73 | # lines1, [N, 2, 2] 74 | # lines2, [M, 2, 2] 75 | A = lines1[:, 0, :][:, np.newaxis, :] 76 | B = lines1[:, 1, :][:, np.newaxis, :] 77 | C = lines2[:, 0, :][np.newaxis, :, :] 78 | D = lines2[:, 1, :][np.newaxis, :, :] 79 | return np.logical_and( 80 | _ccw(A, C, D) != _ccw(B, C, D), 81 | _ccw(A, B, C) != _ccw(A, B, D)) 82 | 83 | 84 | @numba.jit(nopython=False) 85 | def surface_equ_3d_jit(polygon_surfaces): 86 | # return [a, b, c], d in ax+by+cz+d=0 87 | # polygon_surfaces: [num_polygon, num_surfaces, num_points_of_polygon, 3] 88 | surface_vec = polygon_surfaces[:, :, :2, :] - polygon_surfaces[:, :, 1:3, :] 89 | # normal_vec: [..., 3] 90 | normal_vec = np.cross(surface_vec[:, :, 0, :], surface_vec[:, :, 1, :]) 91 | # print(normal_vec.shape, points[..., 0, :].shape) 92 | # d = -np.inner(normal_vec, points[..., 0, :]) 93 | d = np.einsum('aij, aij->ai', normal_vec, polygon_surfaces[:, :, 0, :]) 94 | return normal_vec, -d 95 | 96 | 97 | @numba.jit(nopython=False) 98 | def points_in_convex_polygon_3d_jit_v1(points, 99 | polygon_surfaces, 100 | num_surfaces=None): 101 | """check points is in 3d convex polygons. 102 | Args: 103 | points: [num_points, 3] array. 104 | polygon_surfaces: [num_polygon, max_num_surfaces, 105 | max_num_points_of_surface, 3] 106 | array. all surfaces' normal vector must direct to internal. 107 | max_num_points_of_surface must at least 3. 108 | num_surfaces: [num_polygon] array. indicate how many surfaces 109 | a polygon contain 110 | Returns: 111 | [num_points, num_polygon] bool array. 112 | """ 113 | max_num_surfaces, max_num_points_of_surface = polygon_surfaces.shape[1:3] 114 | num_points = points.shape[0] 115 | num_polygons = polygon_surfaces.shape[0] 116 | if num_surfaces is None: 117 | num_surfaces = np.full((num_polygons,), 9999999, dtype=np.int64) 118 | normal_vec, d = surface_equ_3d_jit(polygon_surfaces[:, :, :3, :]) 119 | # normal_vec: [num_polygon, max_num_surfaces, 3] 120 | # d: [num_polygon, max_num_surfaces] 121 | ret = np.ones((num_points, num_polygons), dtype=np.bool_) 122 | sign = 0.0 123 | for i in range(num_points): 124 | for j in range(num_polygons): 125 | for k in range(max_num_surfaces): 126 | if k > num_surfaces[j]: 127 | break 128 | sign = points[i, 0] * normal_vec[j, k, 0] \ 129 | + points[i, 1] * normal_vec[j, k, 1] \ 130 | + points[i, 2] * normal_vec[j, k, 2] + d[j, k] 131 | if sign >= 0: 132 | ret[i, j] = False 133 | break 134 | return ret 135 | 136 | def surface_equ_3d(polygon_surfaces): 137 | # return [a, b, c], d in ax+by+cz+d=0 138 | # polygon_surfaces: [num_polygon, num_surfaces, num_points_of_polygon, 3] 139 | surface_vec = polygon_surfaces[:, :, :2, :] - polygon_surfaces[:, :, 1:3, :] 140 | # normal_vec: [..., 3] 141 | normal_vec = np.cross(surface_vec[:, :, 0, :], surface_vec[:, :, 1, :]) 142 | # print(normal_vec.shape, points[..., 0, :].shape) 143 | # d = -np.inner(normal_vec, points[..., 0, :]) 144 | d = np.einsum('aij, aij->ai', normal_vec, polygon_surfaces[:, :, 0, :]) 145 | return normal_vec, -d 146 | 147 | 148 | 149 | def points_in_convex_polygon_3d_jit(points, 150 | polygon_surfaces, 151 | num_surfaces=None): 152 | """check points is in 3d convex polygons. 153 | Args: 154 | points: [num_points, 3] array. 155 | polygon_surfaces: [num_polygon, max_num_surfaces, 156 | max_num_points_of_surface, 3] 157 | array. all surfaces' normal vector must direct to internal. 158 | max_num_points_of_surface must at least 3. 159 | num_surfaces: [num_polygon] array. indicate how many surfaces 160 | a polygon contain 161 | Returns: 162 | [num_points, num_polygon] bool array. 163 | """ 164 | max_num_surfaces, max_num_points_of_surface = polygon_surfaces.shape[1:3] 165 | num_points = points.shape[0] 166 | num_polygons = polygon_surfaces.shape[0] 167 | if num_surfaces is None: 168 | num_surfaces = np.full((num_polygons,), 9999999, dtype=np.int64) 169 | normal_vec, d = surface_equ_3d_jitv2(polygon_surfaces[:, :, :3, :]) 170 | # normal_vec: [num_polygon, max_num_surfaces, 3] 171 | # d: [num_polygon, max_num_surfaces] 172 | return _points_in_convex_polygon_3d_jit(points, polygon_surfaces, normal_vec, d, num_surfaces) 173 | 174 | 175 | def points_count_convex_polygon_3d_jit(points, 176 | polygon_surfaces, 177 | num_surfaces=None): 178 | """check points is in 3d convex polygons. 179 | Args: 180 | points: [num_points, 3] array. 181 | polygon_surfaces: [num_polygon, max_num_surfaces, 182 | max_num_points_of_surface, 3] 183 | array. all surfaces' normal vector must direct to internal. 184 | max_num_points_of_surface must at least 3. 185 | num_surfaces: [num_polygon] array. indicate how many surfaces 186 | a polygon contain 187 | Returns: 188 | [num_polygon] array. 189 | """ 190 | max_num_surfaces, max_num_points_of_surface = polygon_surfaces.shape[1:3] 191 | num_points = points.shape[0] 192 | num_polygons = polygon_surfaces.shape[0] 193 | if num_surfaces is None: 194 | num_surfaces = np.full((num_polygons,), 9999999, dtype=np.int64) 195 | normal_vec, d = surface_equ_3d_jitv2(polygon_surfaces[:, :, :3, :]) 196 | # normal_vec: [num_polygon, max_num_surfaces, 3] 197 | # d: [num_polygon, max_num_surfaces] 198 | return _points_count_convex_polygon_3d_jit(points, polygon_surfaces, normal_vec, d, num_surfaces) 199 | 200 | 201 | @numba.njit 202 | def _points_in_convex_polygon_3d_jit(points, 203 | polygon_surfaces, 204 | normal_vec, d, 205 | num_surfaces=None): 206 | """check points is in 3d convex polygons. 207 | Args: 208 | points: [num_points, 3] array. 209 | polygon_surfaces: [num_polygon, max_num_surfaces, 210 | max_num_points_of_surface, 3] 211 | array. all surfaces' normal vector must direct to internal. 212 | max_num_points_of_surface must at least 3. 213 | num_surfaces: [num_polygon] array. indicate how many surfaces 214 | a polygon contain 215 | Returns: 216 | [num_points, num_polygon] bool array. 217 | """ 218 | max_num_surfaces, max_num_points_of_surface = polygon_surfaces.shape[1:3] 219 | num_points = points.shape[0] 220 | num_polygons = polygon_surfaces.shape[0] 221 | ret = np.ones((num_points, num_polygons), dtype=np.bool_) 222 | sign = 0.0 223 | for i in range(num_points): 224 | for j in range(num_polygons): 225 | for k in range(max_num_surfaces): 226 | if k > num_surfaces[j]: 227 | break 228 | sign = points[i, 0] * normal_vec[j, k, 0] \ 229 | + points[i, 1] * normal_vec[j, k, 1] \ 230 | + points[i, 2] * normal_vec[j, k, 2] + d[j, k] 231 | if sign >= 0: 232 | ret[i, j] = False 233 | break 234 | return ret 235 | 236 | @numba.njit 237 | def _points_count_convex_polygon_3d_jit(points, 238 | polygon_surfaces, 239 | normal_vec, d, 240 | num_surfaces=None): 241 | """count points in 3d convex polygons. 242 | Args: 243 | points: [num_points, 3] array. 244 | polygon_surfaces: [num_polygon, max_num_surfaces, 245 | max_num_points_of_surface, 3] 246 | array. all surfaces' normal vector must direct to internal. 247 | max_num_points_of_surface must at least 3. 248 | num_surfaces: [num_polygon] array. indicate how many surfaces 249 | a polygon contain 250 | Returns: 251 | [num_polygon] array. 252 | """ 253 | max_num_surfaces, max_num_points_of_surface = polygon_surfaces.shape[1:3] 254 | num_points = points.shape[0] 255 | num_polygons = polygon_surfaces.shape[0] 256 | ret = np.full((num_polygons,), num_points, dtype=np.int64) 257 | sign = 0.0 258 | for i in range(num_points): 259 | for j in range(num_polygons): 260 | for k in range(max_num_surfaces): 261 | if k > num_surfaces[j]: 262 | break 263 | sign = points[i, 0] * normal_vec[j, k, 0] \ 264 | + points[i, 1] * normal_vec[j, k, 1] \ 265 | + points[i, 2] * normal_vec[j, k, 2] + d[j, k] 266 | if sign >= 0: 267 | ret[j] -= 1 268 | break 269 | return ret 270 | 271 | 272 | @numba.jit 273 | def points_in_convex_polygon_jit(points, polygon, clockwise=True): 274 | """check points is in 2d convex polygons. True when point in polygon 275 | Args: 276 | points: [num_points, 2] array. 277 | polygon: [num_polygon, num_points_of_polygon, 2] array. 278 | clockwise: bool. indicate polygon is clockwise. 279 | Returns: 280 | [num_points, num_polygon] bool array. 281 | """ 282 | # first convert polygon to directed lines 283 | num_points_of_polygon = polygon.shape[1] 284 | num_points = points.shape[0] 285 | num_polygons = polygon.shape[0] 286 | if clockwise: 287 | vec1 = polygon - polygon[:, [num_points_of_polygon - 1] + 288 | list(range(num_points_of_polygon - 1)), :] 289 | else: 290 | vec1 = polygon[:, [num_points_of_polygon - 1] + 291 | list(range(num_points_of_polygon - 1)), :] - polygon 292 | # vec1: [num_polygon, num_points_of_polygon, 2] 293 | ret = np.zeros((num_points, num_polygons), dtype=np.bool_) 294 | success = True 295 | cross = 0.0 296 | for i in range(num_points): 297 | for j in range(num_polygons): 298 | success = True 299 | for k in range(num_points_of_polygon): 300 | cross = vec1[j, k, 1] * (polygon[j, k, 0] - points[i, 0]) 301 | cross -= vec1[j, k, 0] * (polygon[j, k, 1] - points[i, 1]) 302 | if cross >= 0: 303 | success = False 304 | break 305 | ret[i, j] = success 306 | return ret 307 | 308 | def points_in_convex_polygon(points, polygon, clockwise=True): 309 | """check points is in convex polygons. may run 2x faster when write in 310 | cython(don't need to calculate all cross-product between edge and point) 311 | Args: 312 | points: [num_points, 2] array. 313 | polygon: [num_polygon, num_points_of_polygon, 2] array. 314 | clockwise: bool. indicate polygon is clockwise. 315 | Returns: 316 | [num_points, num_polygon] bool array. 317 | """ 318 | # first convert polygon to directed lines 319 | num_lines = polygon.shape[1] 320 | polygon_next = polygon[:, [num_lines - 1] + list(range(num_lines - 1)), :] 321 | if clockwise: 322 | vec1 = (polygon - polygon_next)[np.newaxis, ...] 323 | else: 324 | vec1 = (polygon_next - polygon)[np.newaxis, ...] 325 | vec2 = polygon[np.newaxis, ...] - points[:, np.newaxis, np.newaxis, :] 326 | # [num_points, num_polygon, num_points_of_polygon, 2] 327 | cross = np.cross(vec1, vec2) 328 | return np.all(cross > 0, axis=2) 329 | 330 | 331 | @numba.njit 332 | def surface_equ_3d_jitv2(surfaces): 333 | # polygon_surfaces: [num_polygon, num_surfaces, num_points_of_polygon, 3] 334 | num_polygon = surfaces.shape[0] 335 | max_num_surfaces = surfaces.shape[1] 336 | normal_vec = np.zeros((num_polygon, max_num_surfaces, 3), dtype=surfaces.dtype) 337 | d = np.zeros((num_polygon, max_num_surfaces), dtype=surfaces.dtype) 338 | sv0 = surfaces[0, 0, 0] - surfaces[0, 0, 1] 339 | sv1 = surfaces[0, 0, 0] - surfaces[0, 0, 1] 340 | for i in range(num_polygon): 341 | for j in range(max_num_surfaces): 342 | sv0[0] = surfaces[i, j, 0, 0] - surfaces[i, j, 1, 0] 343 | sv0[1] = surfaces[i, j, 0, 1] - surfaces[i, j, 1, 1] 344 | sv0[2] = surfaces[i, j, 0, 2] - surfaces[i, j, 1, 2] 345 | sv1[0] = surfaces[i, j, 1, 0] - surfaces[i, j, 2, 0] 346 | sv1[1] = surfaces[i, j, 1, 1] - surfaces[i, j, 2, 1] 347 | sv1[2] = surfaces[i, j, 1, 2] - surfaces[i, j, 2, 2] 348 | normal_vec[i, j, 0] = (sv0[1] * sv1[2] - sv0[2] * sv1[1]) 349 | normal_vec[i, j, 1] = (sv0[2] * sv1[0] - sv0[0] * sv1[2]) 350 | normal_vec[i, j, 2] = (sv0[0] * sv1[1] - sv0[1] * sv1[0]) 351 | 352 | d[i, j] = -surfaces[i, j, 0, 0] * normal_vec[i, j, 0] - \ 353 | surfaces[i, j, 0, 1] * normal_vec[i, j, 1] - \ 354 | surfaces[i, j, 0, 2] * normal_vec[i, j, 2] 355 | return normal_vec, d 356 | 357 | 358 | @numba.njit 359 | def points_in_convex_polygon_3d_jit_v2(points, 360 | surfaces, 361 | num_surfaces=None): 362 | """check points is in 3d convex polygons. 363 | Args: 364 | points: [num_points, 3] array. 365 | polygon_surfaces: [num_polygon, max_num_surfaces, 366 | max_num_points_of_surface, 3] 367 | array. all surfaces' normal vector must direct to internal. 368 | max_num_points_of_surface must at least 3. 369 | num_surfaces: [num_polygon] array. indicate how many surfaces 370 | a polygon contain 371 | Returns: 372 | [num_points, num_polygon] bool array. 373 | """ 374 | num_polygon = surfaces.shape[0] 375 | max_num_surfaces = surfaces.shape[1] 376 | num_points = points.shape[0] 377 | normal_vec = np.zeros((num_polygon, max_num_surfaces, 3), dtype=surfaces.dtype) 378 | d = np.zeros((num_polygon, max_num_surfaces), dtype=surfaces.dtype) 379 | sv0 = surfaces[0, 0, 0] - surfaces[0, 0, 1] 380 | sv1 = surfaces[0, 0, 0] - surfaces[0, 0, 1] 381 | ret = np.ones((num_points, num_polygon), dtype=np.bool_) 382 | for i in range(num_polygon): 383 | for j in range(max_num_surfaces): 384 | sv0[0] = surfaces[i, j, 0, 0] - surfaces[i, j, 1, 0] 385 | sv0[1] = surfaces[i, j, 0, 1] - surfaces[i, j, 1, 1] 386 | sv0[2] = surfaces[i, j, 0, 2] - surfaces[i, j, 1, 2] 387 | sv1[0] = surfaces[i, j, 1, 0] - surfaces[i, j, 2, 0] 388 | sv1[1] = surfaces[i, j, 1, 1] - surfaces[i, j, 2, 1] 389 | sv1[2] = surfaces[i, j, 1, 2] - surfaces[i, j, 2, 2] 390 | normal_vec[i, j, 0] = (sv0[1] * sv1[2] - sv0[2] * sv1[1]) 391 | normal_vec[i, j, 1] = (sv0[2] * sv1[0] - sv0[0] * sv1[2]) 392 | normal_vec[i, j, 2] = (sv0[0] * sv1[1] - sv0[1] * sv1[0]) 393 | 394 | d[i, j] = -surfaces[i, j, 0, 0] * normal_vec[i, j, 0] - \ 395 | surfaces[i, j, 0, 1] * normal_vec[i, j, 1] - \ 396 | surfaces[i, j, 0, 2] * normal_vec[i, j, 2] 397 | 398 | sign = 0.0 399 | for i in range(num_points): 400 | for j in range(num_polygon): 401 | for k in range(max_num_surfaces): 402 | sign = points[i, 0] * normal_vec[j, k, 0] \ 403 | + points[i, 1] * normal_vec[j, k, 1] \ 404 | + points[i, 2] * normal_vec[j, k, 2] + d[j, k] 405 | if sign >= 0: 406 | ret[i, j] = False 407 | break 408 | return ret 409 | 410 | -------------------------------------------------------------------------------- /part_aware_augmentation.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import box_np_ops 3 | 4 | def get_partition_corners(augmented_box_corners, partition_idx, gt_names): 5 | partition_corners = [] 6 | if gt_names == "Car": 7 | if partition_idx == 0: 8 | corner_keys = ['0', '01', '02', '03', '04', '05', '06', '07'] 9 | elif partition_idx == 1: 10 | corner_keys = ['01', '1', '12', '02', '05', '15', '16', '06'] 11 | elif partition_idx == 2: 12 | corner_keys = ['02', '12', '2', '23', '06', '16', '26', '27'] 13 | elif partition_idx == 3: 14 | corner_keys = ['03', '02', '23', '3', '07', '06', '27', '37'] 15 | elif partition_idx == 4: 16 | corner_keys = ['04', '05', '06', '07', '4', '45', '46', '47'] 17 | elif partition_idx == 5: 18 | corner_keys = ['05', '15', '16', '06', '45', '5', '56', '46'] 19 | elif partition_idx == 6: 20 | corner_keys = ['06', '16', '26', '27', '46', '56', '6', '67'] 21 | elif partition_idx == 7: 22 | corner_keys = ['07', '06', '27', '37', '47', '46', '67', '7'] 23 | else: 24 | print("ERROR: wrong partition_idx {}".format(partition_idx)) 25 | exit() 26 | elif gt_names == "Pedestrian": 27 | if partition_idx == 0: 28 | corner_keys = ['0', '01', '23', '3', '04', '05', '27', '37'] 29 | elif partition_idx == 1: 30 | corner_keys = ['01', '1', '2', '23', '05', '15', '26', '27'] 31 | elif partition_idx == 2: 32 | corner_keys = ['05', '15', '26', '27', '45', '5', '6', '67'] 33 | elif partition_idx == 3: 34 | corner_keys = ['04', '05', '27', '37', '4', '45', '67', '7'] 35 | else: 36 | print("ERROR: wrong partition_idx {}".format(partition_idx)) 37 | exit() 38 | elif gt_names == "Cyclist": 39 | if partition_idx == 0: 40 | corner_keys = ['0', '01', '02', '03', '4', '45', '46', '47'] 41 | elif partition_idx == 1: 42 | corner_keys = ['01', '1', '12', '02', '45', '5', '56', '46'] 43 | elif partition_idx == 2: 44 | corner_keys = ['02', '12', '2', '23', '46', '56', '6', '67'] 45 | elif partition_idx == 3: 46 | corner_keys = ['03', '02', '23', '3', '47', '46', '67', '7'] 47 | else: 48 | print("ERROR: wrong partition_idx {}".format(partition_idx)) 49 | exit() 50 | else: 51 | print("ERROR: wrong gt_names {}".format(gt_names)) 52 | 53 | for key in corner_keys: 54 | partition_corners.append(augmented_box_corners[key]) 55 | partition_corners = np.expand_dims(np.array(partition_corners), axis=0) 56 | return partition_corners 57 | 58 | 59 | def augment_box_corners(box_corners): 60 | augmented_box_corners = {} 61 | for i, corner in enumerate(box_corners): 62 | augmented_box_corners[str(i)] = corner 63 | augmented_box_corners["01"] = (box_corners[0] + box_corners[1]) / 2 64 | augmented_box_corners["02"] = (box_corners[0] + box_corners[2]) / 2 65 | augmented_box_corners["03"] = (box_corners[0] + box_corners[3]) / 2 66 | augmented_box_corners["04"] = (box_corners[0] + box_corners[4]) / 2 67 | augmented_box_corners["05"] = (box_corners[0] + box_corners[5]) / 2 68 | augmented_box_corners["06"] = (box_corners[0] + box_corners[6]) / 2 69 | augmented_box_corners["07"] = (box_corners[0] + box_corners[7]) / 2 70 | augmented_box_corners["12"] = (box_corners[1] + box_corners[2]) / 2 71 | augmented_box_corners["15"] = (box_corners[1] + box_corners[5]) / 2 72 | augmented_box_corners["16"] = (box_corners[1] + box_corners[6]) / 2 73 | augmented_box_corners["23"] = (box_corners[2] + box_corners[3]) / 2 74 | augmented_box_corners["26"] = (box_corners[2] + box_corners[6]) / 2 75 | augmented_box_corners["27"] = (box_corners[2] + box_corners[7]) / 2 76 | augmented_box_corners["37"] = (box_corners[3] + box_corners[7]) / 2 77 | augmented_box_corners["45"] = (box_corners[4] + box_corners[5]) / 2 78 | augmented_box_corners["46"] = (box_corners[4] + box_corners[6]) / 2 79 | augmented_box_corners["47"] = (box_corners[4] + box_corners[7]) / 2 80 | augmented_box_corners["56"] = (box_corners[5] + box_corners[6]) / 2 81 | augmented_box_corners["67"] = (box_corners[6] + box_corners[7]) / 2 82 | return augmented_box_corners 83 | 84 | 85 | def assign_box_points_partition(points, gt_boxes, gt_names, num_partition): 86 | assert len(points.shape) == 2, 'Wrong points.shape' 87 | box_points_mask, box_corners_3d = box_np_ops.points_corners_in_rbbox(points, gt_boxes) 88 | bg_mask = np.logical_not(np.any(box_points_mask, axis=1)) 89 | assert len(box_points_mask.shape) == 2, 'Wrong box_points_mask.shape' 90 | 91 | fg_seperated_points = [] 92 | partition_corners_list = [] 93 | for i in range(box_points_mask.shape[1]): 94 | separated_box_points = [] 95 | box_points = points[box_points_mask[:, i]] 96 | box_corners = box_corners_3d[i] 97 | augmented_box_corners = augment_box_corners(box_corners) 98 | partition_corners_np = np.zeros((num_partition[gt_names[i]], 8, 3)) 99 | for j in range(num_partition[gt_names[i]]): 100 | partition_corners = get_partition_corners(augmented_box_corners, j, gt_names[i]) # [1, 8, 3] 101 | partition_corners_np[j] = partition_corners 102 | partition_points_mask = np.squeeze(box_np_ops.points_in_corners(box_points, partition_corners), axis=1) 103 | assert len(box_points[partition_points_mask].shape) == 2, 'Wrong box_points[partition_points_mask].shape' 104 | separated_box_points.append(box_points[partition_points_mask]) 105 | partition_corners_list.append(partition_corners_np) 106 | fg_seperated_points.append(separated_box_points) 107 | 108 | bg_points = points[bg_mask] 109 | return fg_seperated_points, bg_points, partition_corners_list 110 | 111 | def get_random_boxes(gt_boxes, gt_names, num_partition, box_corners_3d): 112 | dims = gt_boxes[:, 3:6] 113 | rys = gt_boxes[:, 6] 114 | 115 | corners_pair = [[1, 2], [1, 5], [1, 0]] 116 | corners_pair_opposite = [[5, 6], [2, 6], [6, 7]] 117 | 118 | random_partitions = [] 119 | for i in range(dims.shape[0]): 120 | l = np.linalg.norm(box_corners_3d[i, corners_pair[0][0], :] - box_corners_3d[i, corners_pair[0][1], :]) 121 | w = np.linalg.norm(box_corners_3d[i, corners_pair[1][0], :] - box_corners_3d[i, corners_pair[1][1], :]) 122 | h = np.linalg.norm(box_corners_3d[i, corners_pair[2][0], :] - box_corners_3d[i, corners_pair[2][1], :]) 123 | 124 | random_partition = [] 125 | for j in range(num_partition[gt_names[i]]): 126 | t_list = [] 127 | min_diff = 0.4 128 | for k in range(3): 129 | t = np.random.uniform(low=0, high=1, size=(1, 2)) 130 | if gt_names[i] == "Car": 131 | while True: 132 | t = np.random.uniform(low=0, high=1, size=(1, 2)) 133 | if np.abs(t[0, 0] - t[0, 1]) >= min_diff: 134 | break 135 | t_list.append(t) 136 | 137 | t_list = np.concatenate(t_list, axis=0) 138 | partition_l = l * np.abs(t_list[0][0] - t_list[0][1]) 139 | partition_w = w * np.abs(t_list[1][0] - t_list[1][1]) 140 | partition_h = h * np.abs(t_list[2][0] - t_list[2][1]) 141 | 142 | t_mean = np.mean(t_list, axis=1) 143 | point_l0 = t_mean[0] * box_corners_3d[i, corners_pair[0][0], :] + (1 - t_mean[0]) * box_corners_3d[i, corners_pair[0][1], :] 144 | point_l1 = t_mean[0] * box_corners_3d[i, corners_pair_opposite[0][0], :] + (1 - t_mean[0]) * box_corners_3d[i, corners_pair_opposite[0][1], :] 145 | 146 | point_top = t_mean[1] * point_l0 + (1 - t_mean[1]) * point_l1 147 | point_h = t_mean[2] * box_corners_3d[i, corners_pair[2][0], :] + (1 - t_mean[2]) * box_corners_3d[i, corners_pair[2][1], :] 148 | 149 | random_partition.append([point_top[0], point_top[1], point_h[2], partition_w, partition_l, partition_h, rys[i]]) 150 | 151 | random_partitions.append(random_partition) 152 | 153 | return random_partitions 154 | 155 | def assign_random_partition(points, gt_boxes, gt_names, num_partition): 156 | assert len(points.shape) == 2, 'Wrong points.shape' 157 | box_points_mask, box_corners_3d = box_np_ops.points_corners_in_rbbox(points, gt_boxes) 158 | bg_mask = np.logical_not(np.any(box_points_mask, axis=1)) 159 | assert len(box_points_mask.shape) == 2, 'Wrong box_points_mask.shape' 160 | 161 | random_partitions = get_random_boxes(gt_boxes, gt_names, num_partition, box_corners_3d) 162 | 163 | fg_seperated_points = [] 164 | fg_not_assigned_points = [] 165 | partition_corners_list = [] 166 | for i in range(box_points_mask.shape[1]): 167 | separated_box_points = [] 168 | box_points = points[box_points_mask[:, i]] 169 | partition_points_mask = np.zeros((box_points.shape[0]), dtype=np.bool) 170 | random_partition = np.array(random_partitions[i]) 171 | partition_corners = box_np_ops.center_to_corner_box3d(random_partition[:, 0:3], random_partition[:, 3:6], random_partition[:, 6], origin=(0.5, 0.5, 0.5), axis=2) 172 | partition_corners_list.append(partition_corners) 173 | for j in range(num_partition[gt_names[i]]): 174 | not_assigned_points_mask = np.logical_not(partition_points_mask) 175 | cur_box_points = box_points[not_assigned_points_mask] 176 | cur_partition_mask = np.squeeze(box_np_ops.points_in_corners(cur_box_points, partition_corners[j:j+1]), axis=1) 177 | assert len(box_points[partition_points_mask].shape) == 2, 'Wrong box_points[partition_points_mask].shape' 178 | separated_box_points.append(cur_box_points[cur_partition_mask]) 179 | partition_points_mask[not_assigned_points_mask] = np.logical_or(partition_points_mask[not_assigned_points_mask], cur_partition_mask) 180 | fg_seperated_points.append(separated_box_points) 181 | fg_not_assigned_points.append(box_points[np.logical_not(partition_points_mask)]) 182 | 183 | 184 | bg_points = points[bg_mask] 185 | check_empty = True 186 | for p in fg_not_assigned_points: 187 | if p.shape[0] > 0: 188 | check_empty = False 189 | if not check_empty: 190 | fg_not_assigned_points = np.concatenate(fg_not_assigned_points, axis=0) 191 | bg_points = np.concatenate([bg_points, fg_not_assigned_points], axis=0) 192 | return fg_seperated_points, bg_points, partition_corners_list, random_partitions 193 | 194 | def calc_distances(p0, points): 195 | return ((p0 - points)**2).sum(axis=1) 196 | 197 | def farthest_point_sampling(pts, K): 198 | farthest_pts = np.zeros((K, 3)) 199 | farthest_pts_idx = np.zeros(K, dtype=np.int) 200 | farthest_pts_idx[0] = np.random.randint(len(pts)) 201 | farthest_pts[0] = pts[farthest_pts_idx[0]] 202 | 203 | distances = calc_distances(farthest_pts[0], pts) 204 | 205 | for i in range(1, K): 206 | farthest_pts_idx[i] = np.argmax(distances) 207 | farthest_pts[i] = pts[farthest_pts_idx[i]] 208 | distances = np.minimum(distances, calc_distances(farthest_pts[i], pts)) 209 | return farthest_pts_idx 210 | 211 | class PartAwareAugmentation(object): 212 | def __init__(self, points, gt_boxes, gt_names, class_names=None, random_partition=False): 213 | self.points = points 214 | self.gt_boxes = gt_boxes 215 | self.gt_names = gt_names 216 | self.num_gt_boxes = gt_boxes.shape[0] 217 | self.gt_boxes_mask = [True] * self.num_gt_boxes 218 | self.num_partition = {} 219 | self.num_partition['Car'] = 8 220 | self.num_partition['Pedestrian'] = 4 221 | self.num_partition['Cyclist'] = 4 222 | self.num_classes = len(class_names) 223 | self.random_partition = random_partition 224 | if random_partition: 225 | self.separated_box_points, self.bg_points, self.partition_corners, self.random_partition_boxes = \ 226 | assign_random_partition(self.points, gt_boxes, gt_names, self.num_partition) 227 | else: 228 | self.separated_box_points, self.bg_points, self.partition_corners = assign_box_points_partition(self.points, gt_boxes, gt_names, self.num_partition) 229 | self.aug_flag = np.zeros((self.num_gt_boxes, 8, 6), dtype=np.bool) 230 | 231 | def interpret_pa_aug_param(self, pa_aug_param): 232 | pa_aug_param_dict={} 233 | method_list = ['dropout', 'sparse', 'noise', 'swap', 'mix', 'jitter', 'random', 'distance'] 234 | for method in method_list: 235 | if method == 'distance': 236 | pa_aug_param_dict[method] = 100 237 | continue 238 | if method == 'random': 239 | pa_aug_param_dict[method] = False 240 | continue 241 | pa_aug_param_dict[method] = 0 242 | pa_aug_param_dict[method + '_p'] = 0 243 | 244 | if pa_aug_param is None: 245 | return pa_aug_param_dict 246 | 247 | param_list = pa_aug_param.split('_') 248 | for i, param in enumerate(param_list): 249 | if param.startswith('p'): 250 | continue 251 | for method in method_list: 252 | if param.startswith(method): 253 | if method == 'random': 254 | pa_aug_param_dict[method] = True 255 | method_list.remove(method) 256 | break 257 | number = param.replace(method, "") 258 | if len(number) == 0: 259 | if method == "jitter": 260 | pa_aug_param_dict[method] = 0.1 261 | else: 262 | pa_aug_param_dict[method] = 1 263 | else: 264 | if method == "jitter": 265 | pa_aug_param_dict[method] = float(number) / 10 ** (len(number) - 1) 266 | else: 267 | pa_aug_param_dict[method] = int(number) 268 | if method == 'distance': 269 | method_list.remove(method) 270 | break 271 | pa_aug_param_dict[method + '_p'] = 1.0 272 | if param_list[i + 1].startswith('p'): 273 | number = param_list[i + 1].replace('p', "") 274 | if len(number) == 2: 275 | pa_aug_param_dict[method + '_p'] = float(number) / 10.0 276 | elif len(number) == 3: 277 | pa_aug_param_dict[method + '_p'] = float(number) / 100.0 278 | method_list.remove(method) 279 | break 280 | 281 | return pa_aug_param_dict 282 | 283 | def remove_empty_gt_boxes(self): 284 | for i in range(self.num_gt_boxes): 285 | check_empty = True 286 | if self.random_partition: 287 | box_points_mask, _ = box_np_ops.points_corners_in_rbbox(self.points, self.gt_boxes[i:i+1]) 288 | if np.any(box_points_mask, axis=0)[0]: 289 | check_empty = False 290 | else: 291 | for j in range(self.num_partition[self.gt_names[i]]): 292 | if self.separated_box_points[i][j].shape[0] > 0: 293 | check_empty = False 294 | break 295 | if check_empty: 296 | self.gt_boxes_mask[i] = False 297 | 298 | self.gt_boxes = self.gt_boxes[self.gt_boxes_mask] 299 | self.gt_names = self.gt_names[self.gt_boxes_mask] 300 | self.num_gt_boxes = self.gt_boxes.shape[0] 301 | self.separated_box_points = [d for d, s in zip(self.separated_box_points, self.gt_boxes_mask) if s] 302 | self.aug_flag = self.aug_flag[self.gt_boxes_mask] 303 | 304 | def stack_fg_points(self): 305 | fg_points = np.zeros((0, 4)) 306 | for i in range(self.num_gt_boxes): 307 | for j in range(self.num_partition[self.gt_names[i]]): 308 | fg_points = np.vstack((fg_points, self.separated_box_points[i][j])) 309 | return fg_points 310 | 311 | def stack_fg_points_idx(self, idx=0): 312 | fg_points = np.zeros((0, 4)) 313 | for i in range(self.num_gt_boxes): 314 | if i == idx: 315 | for j in range(self.num_partition[self.gt_names[i]]): 316 | fg_points = np.vstack((fg_points, self.separated_box_points[idx][j])) 317 | return fg_points 318 | 319 | def stack_fg_points_mask(self): 320 | fg_points = np.zeros((0, 4)) 321 | for i in range(self.num_gt_boxes): 322 | if self.gt_boxes_mask[i]: 323 | for j in range(self.num_partition[self.gt_names[i]]): 324 | fg_points = np.vstack((fg_points, self.separated_box_points[i][j])) 325 | return fg_points 326 | 327 | ##################################### 328 | # remove points in random sub-parts # 329 | ##################################### 330 | def dropout_partitions(self, num_dropout_partition=4, distance_limit=100, p=1.0, gt_box_idx=None, robustness_test=False): 331 | if num_dropout_partition <= 0: 332 | return 333 | 334 | for i in range(self.num_gt_boxes): 335 | if gt_box_idx is not None and i != gt_box_idx: 336 | continue 337 | if self.gt_boxes[i][0] > distance_limit: 338 | continue 339 | if np.random.rand(1) > p: 340 | continue 341 | if robustness_test: 342 | max_num_points = 0 343 | for j, separated_box_points in enumerate(self.separated_box_points[i]): 344 | if separated_box_points.shape[0] >= max_num_points: 345 | max_num_points = separated_box_points.shape[0] 346 | dropout_partition_idxes = [j] 347 | print(max_num_points) 348 | 349 | else: 350 | dropout_partition_idxes = np.random.choice(range(self.num_partition[self.gt_names[i]]), num_dropout_partition, replace=False) 351 | for j in range(len(dropout_partition_idxes)): 352 | self.separated_box_points[i][dropout_partition_idxes[j]] = np.zeros((0, 4)) 353 | self.aug_flag[i, dropout_partition_idxes[j], 0] = True 354 | self.remove_empty_gt_boxes() 355 | 356 | 357 | ####################################### 358 | # swap points in non-empty partitions # 359 | ####################################### 360 | def swap_partitions(self, num_swap=1, distance_limit=100, p=1.0, gt_box_idx=None): 361 | if num_swap <= 0: 362 | return 363 | 364 | for i in range(self.num_gt_boxes): 365 | if gt_box_idx is not None and i != gt_box_idx: 366 | continue 367 | if self.gt_boxes[i][0] > distance_limit: 368 | continue 369 | if np.random.rand(1) > p: 370 | continue 371 | gt_idxes = list(range(self.num_gt_boxes)) 372 | 373 | if self.num_classes > 1: 374 | # remove idxes of different classes 375 | same_class_mask = self.gt_names == self.gt_names[i] 376 | same_class_mask[i] = False 377 | gt_idxes = [i for (i, m) in zip(gt_idxes, same_class_mask) if m] 378 | else: 379 | gt_idxes.remove(i) 380 | 381 | # find non-empty partition idx for both gt 382 | non_empty_partition_idx = -1 383 | while len(gt_idxes) > 0: 384 | target_gt_idx = np.random.choice(gt_idxes, 1, replace=False)[0] 385 | non_empty_partition_idxes = [idx for idx, points in enumerate(self.separated_box_points[i]) if 386 | points.shape[0] != 0] 387 | while len(non_empty_partition_idxes) > 0: 388 | non_empty_partition_idx_candidate = \ 389 | np.random.choice(non_empty_partition_idxes, 1, replace=False)[0] 390 | if self.separated_box_points[target_gt_idx][non_empty_partition_idx_candidate].shape[0] != 0: 391 | non_empty_partition_idx = non_empty_partition_idx_candidate 392 | break 393 | non_empty_partition_idxes.remove(non_empty_partition_idx_candidate) 394 | if non_empty_partition_idx != -1: 395 | break 396 | gt_idxes.remove(target_gt_idx) 397 | 398 | if non_empty_partition_idx == -1: 399 | continue 400 | 401 | # normalize points 402 | target_partition_points = np.copy(self.separated_box_points[target_gt_idx][non_empty_partition_idx]) 403 | target_gt_center, target_gt_dim, target_gt_angle = self.gt_boxes[target_gt_idx][:3], self.gt_boxes[target_gt_idx][3:6], \ 404 | self.gt_boxes[target_gt_idx][6:7] 405 | target_partition_points[:, :3] -= target_gt_center 406 | target_partition_points[:, :3] = box_np_ops.rotation_3d_in_axis(target_partition_points[:, :3], -target_gt_angle, axis=2) 407 | target_partition_points[:, :3] /= target_gt_dim 408 | 409 | # swap points 410 | cur_gt_center, cur_gt_dim, cur_gt_angle = self.gt_boxes[i][:3], self.gt_boxes[i][3:6], self.gt_boxes[i][6:7] 411 | target_partition_points[:, :3] *= cur_gt_dim 412 | target_partition_points[:, :3] = box_np_ops.rotation_3d_in_axis(target_partition_points[:, :3], 413 | cur_gt_angle, axis=2) 414 | target_partition_points[:, :3] += cur_gt_center 415 | self.separated_box_points[i][non_empty_partition_idx] = target_partition_points 416 | self.aug_flag[i, non_empty_partition_idx, 1] = True 417 | 418 | def swap_partitions_random(self, num_swap=1, distance_limit=100, p=1.0, gt_box_idx=None): 419 | if num_swap <= 0: 420 | return 421 | 422 | for i in range(self.num_gt_boxes): 423 | if gt_box_idx is not None and i != gt_box_idx: 424 | continue 425 | if self.gt_boxes[i][0] > distance_limit: 426 | continue 427 | if np.random.rand(1) > p: 428 | continue 429 | gt_idxes = list(range(self.num_gt_boxes)) 430 | 431 | if self.num_classes > 1: 432 | # remove idxes of different classes 433 | same_class_mask = self.gt_names == self.gt_names[i] 434 | same_class_mask[i] = False 435 | gt_idxes = [i for (i, m) in zip(gt_idxes, same_class_mask) if m] 436 | else: 437 | gt_idxes.remove(i) 438 | 439 | # find non-empty partition idx for both gt 440 | non_empty_partition_idx = -1 441 | 442 | if len(gt_idxes) > 0: 443 | target_gt_idx = np.random.choice(gt_idxes, 1, replace=False)[0] 444 | non_empty_partition_idxes = [idx for idx, points in enumerate(self.separated_box_points[i]) if points.shape[0] != 0] 445 | if len(non_empty_partition_idxes) == 0: 446 | continue 447 | non_empty_partition_idx = np.random.choice(non_empty_partition_idxes, 1, replace=False)[0] 448 | target_non_empty_partition_idxes = [idx for idx, points in enumerate(self.separated_box_points[target_gt_idx]) if points.shape[0] != 0] 449 | if len(target_non_empty_partition_idxes) == 0: 450 | continue 451 | target_non_empty_partition_idx = np.random.choice(target_non_empty_partition_idxes, 1, replace=False)[0] 452 | 453 | if non_empty_partition_idx == -1: 454 | continue 455 | 456 | # normalize points 457 | target_partition_points = np.copy(self.separated_box_points[target_gt_idx][target_non_empty_partition_idx]) 458 | target_partition_boxes = np.array(self.random_partition_boxes[target_gt_idx][target_non_empty_partition_idx]) 459 | target_gt_center, target_gt_dim, target_gt_angle = target_partition_boxes[:3], target_partition_boxes[3:6], \ 460 | target_partition_boxes[6:7] 461 | target_partition_points[:, :3] -= target_gt_center 462 | target_partition_points[:, :3] = box_np_ops.rotation_3d_in_axis(target_partition_points[:, :3], -target_gt_angle, axis=2) 463 | target_partition_points[:, :3] /= target_gt_dim 464 | 465 | # swap points 466 | cur_partition_boxes = np.array(self.random_partition_boxes[i][non_empty_partition_idx]) 467 | cur_gt_center, cur_gt_dim, cur_gt_angle = cur_partition_boxes[:3], cur_partition_boxes[3:6], cur_partition_boxes[6:7] 468 | target_partition_points[:, :3] *= cur_gt_dim 469 | target_partition_points[:, :3] = box_np_ops.rotation_3d_in_axis(target_partition_points[:, :3], 470 | cur_gt_angle, axis=2) 471 | target_partition_points[:, :3] += cur_gt_center 472 | self.separated_box_points[i][non_empty_partition_idx] = target_partition_points 473 | self.aug_flag[i, non_empty_partition_idx, 1] = True 474 | 475 | ####################################### 476 | # mix points in non-empty partitions # 477 | ####################################### 478 | def mix_partitions(self, num_mix=1, distance_limit=100, p=1.0, gt_box_idx=None): 479 | if num_mix <= 0: 480 | return 481 | 482 | for i in range(self.num_gt_boxes): 483 | if gt_box_idx is not None and i != gt_box_idx: 484 | continue 485 | if self.gt_boxes[i][0] > distance_limit: 486 | continue 487 | if np.random.rand(1) > p: 488 | continue 489 | gt_idxes = list(range(self.num_gt_boxes)) 490 | if self.num_classes > 1: 491 | # remove idxes of different classes 492 | same_class_mask = self.gt_names == self.gt_names[i] 493 | same_class_mask[i] = False 494 | gt_idxes = [i for (i, m) in zip(gt_idxes, same_class_mask) if m] 495 | else: 496 | gt_idxes.remove(i) 497 | 498 | # find non-empty partition idx for both gt 499 | non_empty_partition_idx = -1 500 | while len(gt_idxes) > 0: 501 | target_gt_idx = np.random.choice(gt_idxes, 1, replace=False)[0] 502 | non_empty_partition_idxes = [idx for idx, points in enumerate(self.separated_box_points[i]) if 503 | points.shape[0] != 0] 504 | while len(non_empty_partition_idxes) > 0: 505 | non_empty_partition_idx_candidate = \ 506 | np.random.choice(non_empty_partition_idxes, 1, replace=False)[0] 507 | if self.separated_box_points[target_gt_idx][non_empty_partition_idx_candidate].shape[0] != 0: 508 | non_empty_partition_idx = non_empty_partition_idx_candidate 509 | break 510 | non_empty_partition_idxes.remove(non_empty_partition_idx_candidate) 511 | if non_empty_partition_idx != -1: 512 | break 513 | gt_idxes.remove(target_gt_idx) 514 | 515 | if non_empty_partition_idx == -1: 516 | continue 517 | 518 | # normalize points 519 | target_partition_points = np.copy(self.separated_box_points[target_gt_idx][non_empty_partition_idx]) 520 | target_gt_center, target_gt_dim, target_gt_angle = self.gt_boxes[target_gt_idx][:3], self.gt_boxes[ 521 | target_gt_idx][ 522 | 3:6], \ 523 | self.gt_boxes[target_gt_idx][6:7] 524 | target_partition_points[:, :3] -= target_gt_center 525 | target_partition_points[:, :3] = box_np_ops.rotation_3d_in_axis(target_partition_points[:, :3], 526 | -target_gt_angle, axis=2) 527 | target_partition_points[:, :3] /= target_gt_dim 528 | 529 | # mix points 530 | cur_gt_center, cur_gt_dim, cur_gt_angle = self.gt_boxes[i][:3], self.gt_boxes[i][3:6], self.gt_boxes[i][ 531 | 6:7] 532 | target_partition_points[:, :3] *= cur_gt_dim 533 | target_partition_points[:, :3] = box_np_ops.rotation_3d_in_axis(target_partition_points[:, :3], 534 | cur_gt_angle, axis=2) 535 | target_partition_points[:, :3] += cur_gt_center 536 | self.separated_box_points[i][non_empty_partition_idx] = np.concatenate( 537 | (self.separated_box_points[i][non_empty_partition_idx], target_partition_points), axis=0) 538 | self.aug_flag[i, non_empty_partition_idx, 2] = True 539 | 540 | def mix_partitions_random(self, num_mix=1, distance_limit=100, p=1.0, gt_box_idx=None): 541 | if num_mix <= 0: 542 | return 543 | 544 | for i in range(self.num_gt_boxes): 545 | if gt_box_idx is not None and i != gt_box_idx: 546 | continue 547 | if self.gt_boxes[i][0] > distance_limit: 548 | continue 549 | if np.random.rand(1) > p: 550 | continue 551 | gt_idxes = list(range(self.num_gt_boxes)) 552 | if self.num_classes > 1: 553 | # remove idxes of different classes 554 | same_class_mask = self.gt_names == self.gt_names[i] 555 | same_class_mask[i] = False 556 | gt_idxes = [i for (i, m) in zip(gt_idxes, same_class_mask) if m] 557 | else: 558 | gt_idxes.remove(i) 559 | 560 | # find non-empty partition idx for both gt 561 | non_empty_partition_idx = -1 562 | 563 | if len(gt_idxes) > 0: 564 | target_gt_idx = np.random.choice(gt_idxes, 1, replace=False)[0] 565 | non_empty_partition_idxes = [idx for idx, points in enumerate(self.separated_box_points[i]) if 566 | points.shape[0] != 0] 567 | if len(non_empty_partition_idxes) == 0: 568 | continue 569 | non_empty_partition_idx = np.random.choice(non_empty_partition_idxes, 1, replace=False)[0] 570 | target_non_empty_partition_idxes = [idx for idx, points in 571 | enumerate(self.separated_box_points[target_gt_idx]) if 572 | points.shape[0] != 0] 573 | if len(target_non_empty_partition_idxes) == 0: 574 | continue 575 | target_non_empty_partition_idx = \ 576 | np.random.choice(target_non_empty_partition_idxes, 1, replace=False)[0] 577 | 578 | if non_empty_partition_idx == -1: 579 | continue 580 | 581 | # normalize points 582 | target_partition_points = np.copy(self.separated_box_points[target_gt_idx][target_non_empty_partition_idx]) 583 | target_partition_boxes = np.array( 584 | self.random_partition_boxes[target_gt_idx][target_non_empty_partition_idx]) 585 | target_gt_center, target_gt_dim, target_gt_angle = target_partition_boxes[:3], target_partition_boxes[3:6], \ 586 | target_partition_boxes[6:7] 587 | target_partition_points[:, :3] -= target_gt_center 588 | target_partition_points[:, :3] = box_np_ops.rotation_3d_in_axis(target_partition_points[:, :3], 589 | -target_gt_angle, axis=2) 590 | target_partition_points[:, :3] /= target_gt_dim 591 | 592 | # mix points 593 | cur_partition_boxes = np.array(self.random_partition_boxes[i][non_empty_partition_idx]) 594 | cur_gt_center, cur_gt_dim, cur_gt_angle = cur_partition_boxes[:3], cur_partition_boxes[3:6], cur_partition_boxes[6:7] 595 | target_partition_points[:, :3] *= cur_gt_dim 596 | target_partition_points[:, :3] = box_np_ops.rotation_3d_in_axis(target_partition_points[:, :3], 597 | cur_gt_angle, axis=2) 598 | target_partition_points[:, :3] += cur_gt_center 599 | self.separated_box_points[i][non_empty_partition_idx] = np.concatenate( 600 | (self.separated_box_points[i][non_empty_partition_idx], target_partition_points), axis=0) 601 | self.aug_flag[i, non_empty_partition_idx, 2] = True 602 | 603 | ############################# 604 | # dense partition to sparse # 605 | ############################# 606 | def make_points_sparse(self, num_points_limit=1000, distance_limit=100, p=1.0, FPS=False, gt_box_idx=None): 607 | if num_points_limit <= 0: 608 | return 609 | 610 | for i in range(self.num_gt_boxes): 611 | if gt_box_idx is not None and i != gt_box_idx: 612 | continue 613 | if self.gt_boxes[i][0] > distance_limit: 614 | continue 615 | for j in range(self.num_partition[self.gt_names[i]]): 616 | if self.separated_box_points[i][j].shape[0] > num_points_limit: 617 | if np.random.rand(1) > p: 618 | continue 619 | if FPS: 620 | sparse_points_idx = farthest_point_sampling(self.separated_box_points[i][j][:, :3], num_points_limit) 621 | else: 622 | sparse_points_idx = np.random.choice(range(self.separated_box_points[i][j].shape[0]), num_points_limit, replace=False) 623 | self.separated_box_points[i][j] = self.separated_box_points[i][j][sparse_points_idx] 624 | self.aug_flag[i, j, 3] = True 625 | 626 | ######################################## 627 | # translate points with gaussian noise # 628 | ######################################## 629 | def jittering(self, sigma=0.01, p=0.5, distance_limit=100, gt_box_idx=None): 630 | for i in range(self.num_gt_boxes): 631 | if gt_box_idx is not None and i != gt_box_idx: 632 | continue 633 | if self.gt_boxes[i][0] > distance_limit: 634 | continue 635 | for j in range(self.num_partition[self.gt_names[i]]): 636 | if self.separated_box_points[i][j].shape[0] <= 0: 637 | continue 638 | 639 | if np.random.rand(1) > p: 640 | continue 641 | 642 | translation_noise = np.random.normal(0, sigma, size=self.separated_box_points[i][j].shape) 643 | 644 | self.separated_box_points[i][j] += translation_noise 645 | self.aug_flag[i, j, 4] = True 646 | 647 | 648 | ################################################### 649 | # generate random points in a specified partition # 650 | ################################################### 651 | def generate_random_noise(self, num_points=30, distance_limit=100, p=0.5, gt_box_idx=None): 652 | if num_points <= 0: 653 | return 654 | 655 | for i in range(self.num_gt_boxes): 656 | if gt_box_idx is not None and i != gt_box_idx: 657 | continue 658 | if self.gt_boxes[i][0] > distance_limit: 659 | continue 660 | 661 | for j in range(self.num_partition[self.gt_names[i]]): 662 | if np.random.rand(1) > p: 663 | continue 664 | center = np.expand_dims(self.gt_boxes[i][:3], axis=0) 665 | dim = np.expand_dims(self.gt_boxes[i][3:6], axis=0) 666 | angle = self.gt_boxes[i][6:7] 667 | 668 | corners = np.squeeze(box_np_ops.corners_nd(dim, origin=(0.5, 0.5, 0.5)), axis=0) 669 | augmented_box_corners = augment_box_corners(corners) 670 | partition_corners = get_partition_corners(augmented_box_corners, j, self.gt_names[i]) 671 | 672 | x_min, x_max = partition_corners[0, :, 0].min(), partition_corners[0, :, 0].max() 673 | y_min, y_max = partition_corners[0, :, 1].min(), partition_corners[0, :, 1].max() 674 | z_min, z_max = partition_corners[0, :, 2].min(), partition_corners[0, :, 2].max() 675 | 676 | generated_points = np.zeros((1, num_points, 4)) 677 | generated_points[0, :, 0] = np.random.uniform(low=x_min, high=x_max, size=(num_points,)) 678 | generated_points[0, :, 1] = np.random.uniform(low=y_min, high=y_max, size=(num_points,)) 679 | generated_points[0, :, 2] = np.random.uniform(low=z_min, high=z_max, size=(num_points,)) 680 | generated_points[0, :, 3] = np.random.uniform(low=0.0, high=1.0, size=(num_points,)) 681 | 682 | generated_points[:, :, :3] = box_np_ops.rotation_3d_in_axis(generated_points[:, :, :3], angle, axis=2) 683 | generated_points[:, :, :3] += center.reshape([-1, 1, 3]) 684 | self.separated_box_points[i][j] = np.concatenate((self.separated_box_points[i][j], generated_points[0]), axis=0) 685 | self.aug_flag[i, j, 5] = True 686 | 687 | 688 | def generate_random_noise_random(self, num_points=30, distance_limit=100, p=0.5, gt_box_idx=None): 689 | if num_points <= 0: 690 | return 691 | 692 | for i in range(self.num_gt_boxes): 693 | if gt_box_idx is not None and i != gt_box_idx: 694 | continue 695 | if self.gt_boxes[i][0] > distance_limit: 696 | continue 697 | 698 | for j in range(self.num_partition[self.gt_names[i]]): 699 | if np.random.rand(1) > p: 700 | continue 701 | center = np.expand_dims(np.array(self.random_partition_boxes[i][j][:3]), axis=0) 702 | dim = np.expand_dims(np.array(self.random_partition_boxes[i][j][3:6]), axis=0) 703 | angle = np.array(self.random_partition_boxes[i][j][6:7]) 704 | 705 | corners = np.squeeze(box_np_ops.corners_nd(dim, origin=(0.5, 0.5, 0.5)), axis=0) 706 | 707 | x_min, x_max = corners[:, 0].min(), corners[:, 0].max() 708 | y_min, y_max = corners[:, 1].min(), corners[:, 1].max() 709 | z_min, z_max = corners[:, 2].min(), corners[:, 2].max() 710 | 711 | generated_points = np.zeros((1, num_points, 4)) 712 | generated_points[0, :, 0] = np.random.uniform(low=x_min, high=x_max, size=(num_points,)) 713 | generated_points[0, :, 1] = np.random.uniform(low=y_min, high=y_max, size=(num_points,)) 714 | generated_points[0, :, 2] = np.random.uniform(low=z_min, high=z_max, size=(num_points,)) 715 | generated_points[0, :, 3] = np.random.uniform(low=0.0, high=1.0, size=(num_points,)) 716 | 717 | generated_points[:, :, :3] = box_np_ops.rotation_3d_in_axis(generated_points[:, :, :3], angle, axis=2) 718 | generated_points[:, :, :3] += center.reshape([-1, 1, 3]) 719 | self.separated_box_points[i][j] = np.concatenate((self.separated_box_points[i][j], generated_points[0]), axis=0) 720 | self.aug_flag[i, j, 5] = True 721 | 722 | 723 | 724 | ############################### 725 | # apply methods independently # 726 | ############################### 727 | def augment(self, pa_aug_param): 728 | pa_aug_param = self.interpret_pa_aug_param(pa_aug_param) 729 | self.dropout_partitions(num_dropout_partition=pa_aug_param['dropout'], p=pa_aug_param['dropout_p'], distance_limit=pa_aug_param['distance']) 730 | if self.random_partition: 731 | self.swap_partitions_random(num_swap=pa_aug_param['swap'], p=pa_aug_param['swap_p'], distance_limit=pa_aug_param['distance']) 732 | self.mix_partitions_random(num_mix=pa_aug_param['mix'], p=pa_aug_param['mix_p'], distance_limit=pa_aug_param['distance']) 733 | else: 734 | self.swap_partitions(num_swap=pa_aug_param['swap'], p=pa_aug_param['swap_p'], distance_limit=pa_aug_param['distance']) 735 | self.mix_partitions(num_mix=pa_aug_param['mix'], p=pa_aug_param['mix_p'], distance_limit=pa_aug_param['distance']) 736 | self.make_points_sparse(num_points_limit=pa_aug_param['sparse'], p=pa_aug_param['sparse_p'], FPS=True, distance_limit=pa_aug_param['distance']) 737 | self.jittering(sigma=pa_aug_param['jitter'], p=pa_aug_param['jitter_p'], distance_limit=pa_aug_param['distance']) 738 | if self.random_partition: 739 | self.generate_random_noise_random(num_points=pa_aug_param['noise'], p=pa_aug_param['noise_p'], distance_limit=pa_aug_param['distance']) 740 | else: 741 | self.generate_random_noise(num_points=pa_aug_param['noise'], p=pa_aug_param['noise_p'], distance_limit=pa_aug_param['distance']) 742 | 743 | fg_points = self.stack_fg_points() 744 | self.points = np.vstack((fg_points, self.bg_points)) 745 | return self.points, self.gt_boxes_mask 746 | 747 | 748 | 749 | def generate_noise_robustness_test(self, noise_ratio=0.1, remove_original_points=False): 750 | num_points = self.points.shape[0] 751 | x_min, x_max = self.points[:, 0].min(), self.points[:, 0].max() 752 | y_min, y_max = self.points[:, 1].min(), self.points[:, 1].max() 753 | z_min, z_max = self.points[:, 2].min(), self.points[:, 2].max() 754 | r_min, r_max = self.points[:, 3].min(), self.points[:, 3].max() 755 | false_idx = np.random.choice(range(num_points), int(num_points * noise_ratio), replace=False) 756 | mask = np.ones(num_points, dtype=np.bool) 757 | mask[false_idx] = False 758 | self.points = self.points[mask] 759 | 760 | num_noise = int(num_points * noise_ratio) 761 | generated_points = np.zeros((num_noise, 4)) 762 | generated_points[:, 0] = np.random.uniform(low=x_min, high=x_max, size=(num_noise,)) 763 | generated_points[:, 1] = np.random.uniform(low=y_min, high=y_max, size=(num_noise,)) 764 | generated_points[:, 2] = np.random.uniform(low=z_min, high=z_max, size=(num_noise,)) 765 | generated_points[:, 3] = np.random.uniform(low=r_min, high=r_max, size=(num_noise,)) 766 | 767 | self.points = np.concatenate((self.points, generated_points), axis=0) 768 | 769 | def sparse_robustness_test(self, sparse_ratio=0.8): 770 | num_points = self.points.shape[0] 771 | num_points_limit = int(num_points * sparse_ratio) 772 | sparse_points_idx = farthest_point_sampling(self.points[:, :3], num_points_limit) 773 | self.points = self.points[sparse_points_idx] 774 | 775 | def jitter_robustness_test(self, sigma=0.01): 776 | num_points = self.points.shape[0] 777 | translation_noise = np.random.normal(0, sigma, size=[num_points, 3]) 778 | self.points[:, :3] += translation_noise 779 | 780 | 781 | 782 | def create_robusteness_test_data(self, test_name="KITTI-D"): 783 | if test_name == "KITTI-D": 784 | self.dropout_partitions(num_dropout_partition=1, p=1.0, robustness_test=True) 785 | fg_points = self.stack_fg_points() 786 | self.points = np.vstack((fg_points, self.bg_points)) 787 | elif test_name == "KITTI-N": 788 | self.generate_noise_robustness_test(noise_ratio=0.2, remove_original_points=True) 789 | elif test_name == "KITTI-S": 790 | self.sparse_robustness_test(sparse_ratio=0.3) 791 | elif test_name == "KITTI-J": 792 | self.jitter_robustness_test(sigma=0.1) 793 | else: 794 | print() 795 | 796 | return self.points, self.gt_boxes_mask, self.aug_flag, self.partition_corners 797 | 798 | --------------------------------------------------------------------------------