File size: 4,445 Bytes
94da716 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
# Copyright (c) Facebook, Inc. and its affiliates.
""" Utility functions for processing point clouds.
Author: Charles R. Qi and Or Litany
"""
import os
import sys
import torch
# Point cloud IO
import numpy as np
from plyfile import PlyData, PlyElement
# Mesh IO
import trimesh
MEAN_COLOR_RGB = np.array([109.8, 97.2, 83.8])
# ----------------------------------------
# Point Cloud Sampling
# ----------------------------------------
def random_sampling(pc, num_sample, replace=None, return_choices=False):
"""Input is NxC, output is num_samplexC"""
if replace is None:
replace = pc.shape[0] < num_sample
choices = np.random.choice(pc.shape[0], num_sample, replace=replace)
if return_choices:
return pc[choices], choices
else:
return pc[choices]
def check_aspect(crop_range, aspect_min):
xy_aspect = np.min(crop_range[:2]) / np.max(crop_range[:2])
xz_aspect = np.min(crop_range[[0, 2]]) / np.max(crop_range[[0, 2]])
yz_aspect = np.min(crop_range[1:]) / np.max(crop_range[1:])
return (
(xy_aspect >= aspect_min)
or (xz_aspect >= aspect_min)
or (yz_aspect >= aspect_min)
)
class RandomCuboid(object):
"""
RandomCuboid augmentation from DepthContrast [https://arxiv.org/abs/2101.02691]
We slightly modify this operation to account for object detection.
This augmentation randomly crops a cuboid from the input and
ensures that the cropped cuboid contains at least one bounding box
"""
def __init__(
self,
min_points,
aspect=0.8,
min_crop=0.5,
max_crop=1.0,
box_filter_policy="center",
):
self.aspect = aspect
self.min_crop = min_crop
self.max_crop = max_crop
self.min_points = min_points
self.box_filter_policy = box_filter_policy
def __call__(self, point_cloud, target_boxes, per_point_labels=None):
range_xyz = np.max(point_cloud[:, 0:3], axis=0) - np.min(
point_cloud[:, 0:3], axis=0
)
for _ in range(100):
crop_range = self.min_crop + np.random.rand(3) * (
self.max_crop - self.min_crop
)
if not check_aspect(crop_range, self.aspect):
continue
sample_center = point_cloud[np.random.choice(len(point_cloud)), 0:3]
new_range = range_xyz * crop_range / 2.0
max_xyz = sample_center + new_range
min_xyz = sample_center - new_range
upper_idx = (
np.sum((point_cloud[:, 0:3] <= max_xyz).astype(np.int32), 1) == 3
)
lower_idx = (
np.sum((point_cloud[:, 0:3] >= min_xyz).astype(np.int32), 1) == 3
)
new_pointidx = (upper_idx) & (lower_idx)
if np.sum(new_pointidx) < self.min_points:
continue
new_point_cloud = point_cloud[new_pointidx, :]
# filtering policy is the only modification from DepthContrast
if self.box_filter_policy == "center":
# remove boxes whose center does not lie within the new_point_cloud
new_boxes = target_boxes
if (
target_boxes.sum() > 0
): # ground truth contains no bounding boxes. Common in SUNRGBD.
box_centers = target_boxes[:, 0:3]
new_pc_min_max = np.min(new_point_cloud[:, 0:3], axis=0), np.max(
new_point_cloud[:, 0:3], axis=0
)
keep_boxes = np.logical_and(
np.all(box_centers >= new_pc_min_max[0], axis=1),
np.all(box_centers <= new_pc_min_max[1], axis=1),
)
if keep_boxes.sum() == 0:
# current data augmentation removes all boxes in the pointcloud. fail!
continue
new_boxes = target_boxes[keep_boxes]
if per_point_labels is not None:
new_per_point_labels = [x[new_pointidx] for x in per_point_labels]
else:
new_per_point_labels = None
# if we are here, all conditions are met. return boxes
return new_point_cloud, new_boxes, new_per_point_labels
# fallback
return point_cloud, target_boxes, per_point_labels
|