def transform_points_to_voxels(self,
data_dict=None,
config=None,
voxel_generator=None):
if data_dict is None:
try:
from spconv.utils import VoxelGeneratorV2 as VoxelGenerator
except:
from spconv.utils import VoxelGenerator
voxel_generator = VoxelGenerator(
voxel_size=config.VOXEL_SIZE,
point_cloud_range=self.point_cloud_range,
max_num_points=config.MAX_POINTS_PER_VOXEL,
max_voxels=config.MAX_NUMBER_OF_VOXELS[self.mode])
grid_size = (self.point_cloud_range[3:6] -
self.point_cloud_range[0:3]) / np.array(
config.VOXEL_SIZE)
self.grid_size = np.round(grid_size).astype(np.int64)
self.voxel_size = config.VOXEL_SIZE
return partial(self.transform_points_to_voxels,
voxel_generator=voxel_generator)
points = data_dict['points']
voxel_output = voxel_generator.generate(points)
if isinstance(voxel_output, dict):
voxels, coordinates, num_points = \
voxel_output['voxels'], voxel_output['coordinates'], voxel_output['num_points_per_voxel']
else:
voxels, coordinates, num_points = voxel_output
if not data_dict['use_lead_xyz']:
voxels = voxels[..., 3:] # remove xyz in voxels(N, 3)
data_dict['voxels'] = voxels
data_dict['voxel_coords'] = coordinates
data_dict['voxel_num_points'] = num_points
voxel_features, voxel_num_points = data_dict['voxels'], data_dict[
'voxel_num_points']
points_mean = np.sum(voxel_features[:, :, :], axis=1, keepdims=False)
num_points = num_points.astype(voxels.dtype).reshape((-1, 1))
normalizer = np.clip(num_points, a_min=1.0, a_max=6.0)
points_mean = points_mean / normalizer
data_dict['voxel_features'] = points_mean
return data_dict
def getVoxel(self,points):
max_voxels=100000
voxel_generator = VoxelGeneratorV2(
voxel_size=list(self.args_vg.voxel_size),
point_cloud_range=list(self.args_vg.point_cloud_range),
max_num_points=self.args_vg.max_number_of_points_per_voxel,
max_voxels=max_voxels,
full_mean=self.args_vg.full_empty_part_with_mean,
block_filtering=self.args_vg.block_filtering,
block_factor=self.args_vg.block_factor,
block_size=self.args_vg.block_size,
height_threshold=self.args_vg.height_threshold)
res = voxel_generator.generate(
points, max_voxels)
return res
class DepthContrastDataset(Dataset):
"""Base Self Supervised Learning Dataset Class."""
def __init__(self, cfg):
self.split = "train" ### Default is training
self.label_objs = []
self.data_paths = []
self.label_paths = []
self.cfg = cfg
self.batchsize_per_replica = cfg["BATCHSIZE_PER_REPLICA"]
self.label_sources = [] #cfg["LABEL_SOURCES"]
self.dataset_names = cfg["DATASET_NAMES"]
self.label_type = cfg["LABEL_TYPE"]
self.AUGMENT_COORDS_TO_FEATS = False #optional
self._labels_init = False
self._get_data_files("train")
self.data_objs = np.load(
self.data_paths[0]) ### Only load the first one for now
#### Add the voxelizer here
if ("Lidar" in cfg) and cfg["VOX"]:
self.VOXEL_SIZE = [0.1, 0.1, 0.2]
self.point_cloud_range = POINT_RANGE #np.array([ 0. , -75. , -3. , 75.0, 75. , 3. ], dtype=np.float32)
self.MAX_POINTS_PER_VOXEL = 5
self.MAX_NUMBER_OF_VOXELS = 16000
self.voxel_generator = VoxelGenerator(
voxel_size=self.VOXEL_SIZE,
point_cloud_range=self.point_cloud_range,
max_num_points=self.MAX_POINTS_PER_VOXEL,
max_voxels=self.MAX_NUMBER_OF_VOXELS)
grid_size = (self.point_cloud_range[3:6] -
self.point_cloud_range[0:3]) / np.array(
self.VOXEL_SIZE)
self.grid_size = np.round(grid_size).astype(np.int64)
self.voxel_size = self.VOXEL_SIZE
elif cfg["VOX"]:
augment_data = (self.split == "TRAIN")
#### Vox parameters here
self.VOXEL_SIZE = 0.05 #0.02 # 5cm
self.CLIP_BOUND = None #(-1000, -1000, -1000, 1000, 1000, 1000)
self.data_aug_color_trans_ratio = 0.1
self.data_aug_color_jitter_std = 0.05
self.ELASTIC_DISTORT_PARAMS = ((0.2, 0.4), (0.8, 1.6))
if augment_data:
self.prevoxel_transform_train = []
self.prevoxel_transform_train.append(
transforms.ElasticDistortion(self.ELASTIC_DISTORT_PARAMS))
self.prevoxel_transform = transforms.Compose(
self.prevoxel_transform_train)
self.input_transforms = []
self.input_transforms += [
transforms.RandomDropout(0.2),
transforms.RandomHorizontalFlip('z', False),
#transforms.ChromaticAutoContrast(),
transforms.ChromaticTranslation(
self.data_aug_color_trans_ratio),
transforms.ChromaticJitter(self.data_aug_color_jitter_std),
# t.HueSaturationTranslation(config.data_aug_hue_max, config.data_aug_saturation_max),
]
self.input_transforms = transforms.Compose(
self.input_transforms)
# Coordinate Augmentation Arguments: Unlike feature augmentation, coordinate
# augmentation has to be done before voxelization
self.SCALE_AUGMENTATION_BOUND = (0.9, 1.1)
self.ROTATION_AUGMENTATION_BOUND = ((-np.pi / 64, np.pi / 64),
(-np.pi / 64,
np.pi / 64), (-np.pi, np.pi))
self.TRANSLATION_AUGMENTATION_RATIO_BOUND = ((-0.2, 0.2),
(-0.2, 0.2), (0, 0))
self.voxelizer = Voxelizer(
voxel_size=self.VOXEL_SIZE,
clip_bound=self.CLIP_BOUND,
use_augmentation=augment_data,
scale_augmentation_bound=self.SCALE_AUGMENTATION_BOUND,
rotation_augmentation_bound=self.ROTATION_AUGMENTATION_BOUND,
translation_augmentation_ratio_bound=self.
TRANSLATION_AUGMENTATION_RATIO_BOUND,
ignore_label=True)
def _get_data_files(self, split):
local_rank = int(os.environ.get("LOCAL_RANK", 0))
self.data_paths = self.cfg["DATA_PATHS"]
self.label_paths = []
logging.info(f"Rank: {local_rank} Data files:\n{self.data_paths}")
logging.info(f"Rank: {local_rank} Label files:\n{self.label_paths}")
def _augment_coords_to_feats(self, coords, feats, labels=None):
# Center x,y
coords_center = coords.mean(0, keepdims=True)
coords_center[0, 2] = 0
norm_coords = coords - coords_center
feats = np.concatenate((feats, norm_coords), 1)
return coords, feats, labels
def toVox(self, coords, feats, labels):
if "Lidar" in self.cfg:
voxel_output = self.voxel_generator.generate(coords)
if isinstance(voxel_output, dict):
voxels, coordinates, num_points = \
voxel_output['voxels'], voxel_output['coordinates'], voxel_output['num_points_per_voxel']
else:
voxels, coordinates, num_points = voxel_output
data_dict = {}
data_dict['voxels'] = voxels
data_dict['voxel_coords'] = coordinates
data_dict['voxel_num_points'] = num_points
return data_dict
else:
precoords = np.copy(coords)
prefeats = np.copy(feats)
if (self.split == "TRAIN") and (self.prevoxel_transform
is not None):
coords, feats, labels = self.prevoxel_transform(
coords, feats, labels)
coords, feats, labels, transformation = self.voxelizer.voxelize(
coords, feats, labels)
if (self.split == "TRAIN") and (self.input_transforms is not None):
try:
coords, feats, labels = self.input_transforms(
coords, feats, labels)
except:
print("error with: ", coords.shape)
coords = np.zeros((100, 3), dtype=np.int32)
feats = np.zeros((100, 3), dtype=np.float64)
labels = np.zeros((100, ), dtype=np.int32)
if (self.split == "TRAIN") and (self.AUGMENT_COORDS_TO_FEATS):
coords, feats, labels = self._augment_coords_to_feats(
coords, feats, labels)
return (coords, feats, labels)
def load_data(self, idx):
is_success = True
point_path = self.data_objs[idx]
try:
if "Lidar" in self.cfg:
#point = np.load(point_path)
point = np.fromfile(str(point_path), dtype=np.float32).reshape(
-1, 4) #np.load(point_path)
if point.shape[1] != 4:
temp = np.zeros((point.shape[0], 4))
temp[:, :3] = point
point = np.copy(temp)
upper_idx = np.sum(
(point[:, 0:3] <= POINT_RANGE[3:6]).astype(np.int32),
1) == 3
lower_idx = np.sum(
(point[:, 0:3] >= POINT_RANGE[0:3]).astype(np.int32),
1) == 3
new_pointidx = (upper_idx) & (lower_idx)
point = point[new_pointidx, :]
else:
point = np.load(point_path)
### Add height
floor_height = np.percentile(point[:, 2], 0.99)
height = point[:, 2] - floor_height
point = np.concatenate([point, np.expand_dims(height, 1)], 1)
except Exception as e:
logging.warn(
f"Couldn't load: {self.point_dataset[idx]}. Exception: \n{e}")
point = np.zeros([50000, 7])
is_success = False
return point, is_success
def __getitem__(self, idx):
cfg = self.cfg
# TODO: this doesn't yet handle the case where the length of datasets
# could be different.
if cfg["DATA_TYPE"] == "point_vox":
item = {
"data": [],
"data_valid": [],
"data_moco": [],
"vox": [],
"vox_moco": []
}
data, valid = self.load_data(idx)
item["data"].append(data)
item["data_moco"].append(np.copy(data))
item["vox"].append(np.copy(data))
item["vox_moco"].append(np.copy(data))
item["data_valid"].append(1 if valid else -1)
else:
item = {
"data": [],
"data_moco": [],
"data_valid": [],
"data_idx": []
}
data, valid = self.load_data(idx)
item["data"].append(data)
item["data_moco"].append(np.copy(data))
item["data_valid"].append(1 if valid else -1)
### Make copies for moco setting
item["label"] = []
item["label"].append(idx)
### Apply the transformation here
if (cfg["DATA_TYPE"] == "point_vox"):
tempitem = {"data": item["data"]}
tempdata = get_transform3d(tempitem, cfg["POINT_TRANSFORMS"])
item["data"] = tempdata["data"]
tempitem = {"data": item["data_moco"]}
tempdata = get_transform3d(tempitem, cfg["POINT_TRANSFORMS"])
item["data_moco"] = tempdata["data"]
tempitem = {"data": item["vox"]}
tempdata = get_transform3d(tempitem,
cfg["POINT_TRANSFORMS"],
vox=True)
coords = tempdata["data"][0][:, :3]
feats = tempdata["data"][
0][:, 3:6] * 255.0 #np.ones(coords.shape)*255.0
labels = np.zeros(coords.shape[0]).astype(np.int32)
item["vox"] = [self.toVox(coords, feats, labels)]
tempitem = {"data": item["vox_moco"]}
tempdata = get_transform3d(tempitem,
cfg["POINT_TRANSFORMS"],
vox=True)
coords = tempdata["data"][0][:, :3]
feats = tempdata["data"][
0][:, 3:6] * 255.0 #np.ones(coords.shape)*255.0
labels = np.zeros(coords.shape[0]).astype(np.int32)
item["vox_moco"] = [self.toVox(coords, feats, labels)]
else:
tempitem = {"data": item["data"]}
tempdata = get_transform3d(tempitem,
cfg["POINT_TRANSFORMS"],
vox=cfg["VOX"])
if cfg["VOX"]:
coords = tempdata["data"][0][:, :3]
feats = tempdata["data"][0][:, 3:6] * 255.0
labels = np.zeros(coords.shape[0]).astype(np.int32)
item["data"] = [self.toVox(coords, feats, labels)]
else:
item["data"] = tempdata["data"]
tempitem = {"data": item["data_moco"]}
tempdata = get_transform3d(tempitem,
cfg["POINT_TRANSFORMS"],
vox=cfg["VOX"])
if cfg["VOX"]:
coords = tempdata["data"][0][:, :3]
feats = tempdata["data"][
0][:, 3:6] * 255.0 #np.ones(coords.shape)*255.0
labels = np.zeros(coords.shape[0]).astype(np.int32)
item["data_moco"] = [self.toVox(coords, feats, labels)]
else:
item["data_moco"] = tempdata["data"]
return item
def __len__(self):
return len(self.data_objs)
def get_available_splits(self, dataset_config):
return [
key for key in dataset_config if key.lower() in ["train", "test"]
]
def num_samples(self, source_idx=0):
return len(self.data_objs)
def get_batchsize_per_replica(self):
# this searches for batchsize_per_replica in self and then in self.dataset
return getattr(self, "batchsize_per_replica", 1)
def get_global_batchsize(self):
if torch.distributed.is_available(
) and torch.distributed.is_initialized():
world_size = torch.distributed.get_world_size()
else:
world_size = 1
return self.get_batchsize_per_replica() * world_size
def transform_points_to_voxels(self,
data_dict=None,
config=None,
voxel_generator=None,
voxel_generator_2=None):
if data_dict is None:
try:
from spconv.utils import VoxelGeneratorV2 as VoxelGenerator
except:
from spconv.utils import VoxelGenerator
voxel_generator = VoxelGenerator(
voxel_size=config.VOXEL_SIZE,
point_cloud_range=self.point_cloud_range,
max_num_points=config.MAX_POINTS_PER_VOXEL,
max_voxels=config.MAX_NUMBER_OF_VOXELS[self.mode])
#voxel_generator_2 = VoxelGenerator(
# voxel_size=config.VOXEL_SIZE,
# point_cloud_range=self.point_cloud_range,
# max_num_points=config.MAX_POINTS_PER_VOXEL,
# max_voxels=60000
#)
grid_size = (self.point_cloud_range[3:6] -
self.point_cloud_range[0:3]) / np.array(
config.VOXEL_SIZE)
self.grid_size = np.round(grid_size).astype(np.int64)
self.voxel_size = config.VOXEL_SIZE
return partial(self.transform_points_to_voxels,
voxel_generator=voxel_generator,
voxel_generator_2=voxel_generator_2)
points = data_dict['points']
points = data_dict['points_sp']
indices = data_dict['indices']
#print(points.shape)
"""
add code for visibility
"""
ori_points = points[:, [0, 1, 2, 4]]
#print(points[:,-1])
voxel_size = self.voxel_size
pc_range = self.point_cloud_range
#print(pc_range)
#print(self.voxel_size)
origins = data_dict['origins']
num_points = points.shape[0]
num_original = num_points
time_stamps = np.array([0], dtype=np.float32)
time_stamps = points[
indices[:-1],
-1] # counting on the fact we do not miss points from any intermediate time_stamps
time_stamps = (time_stamps[:-1] + time_stamps[1:]) / 2
time_stamps = [-1000.0] + time_stamps.tolist() + [1000.0
] # add boundaries
time_stamps = np.array(time_stamps)
num_original = indices[-1]
#print(time_stamps)
#print(points.shape)
#sys.exit()
if num_points > num_original:
#print("this is test sample")
original_points, sampled_points = ori_points[:
num_original, :], ori_points[
num_original:, :]
visibility, original_mask, sampled_mask = mapping.compute_logodds_and_masks(
original_points, sampled_points, origins, time_stamps,
pc_range, min(voxel_size))
points = np.concatenate(
(original_points[original_mask], sampled_points[sampled_mask]))
else:
visibility = mapping.compute_logodds(ori_points, origins,
time_stamps, pc_range, 0.2)
np.set_printoptions(threshold=sys.maxsize)
visi_map = np.zeros([512, 512, 3])
visibility = np.int64(visibility)
visibility = np.reshape(visibility, (40, 512, 512))[0:40, :, :]
visibility = np.transpose(visibility, (2, 1, 0))
#print(visibility)
#sys.exit()
mask_occ = (visibility >= 1).nonzero()
#print(mask_occ)
mask_free = (visibility == 0).nonzero()
mask_unknown = (visibility == -1).nonzero()
visi_map[np.int64(mask_free[0]),
np.int64(mask_free[1]), :] = np.array([255, 0, 0]) / 255
visi_map[np.int64(mask_occ[0]),
np.int64(mask_occ[1]), :] = np.array([0, 255, 0]) / 255
visi_map[mask_unknown[0],
mask_unknown[1], :] = np.array([0, 0, 255]) / 255
#print(.shape)
#visibility = np.pad(visibility, ((0,2),(0,0)), 'edge')
data_dict['vis'] = visibility
#print(data_dict.keys())
dense_points = data_dict['dense_point']
points = data_dict['points']
#print(points.shape)
voxel_output = voxel_generator.generate(points)
#voxel_dense = voxel_generator.generate(dense_points)
#print(voxel_dense[...,-1])
#print(voxel_output['voxels'].shape)
#sys.exit()
if isinstance(voxel_output, dict):
voxels, coordinates, num_points = \
voxel_output['voxels'], voxel_output['coordinates'], voxel_output['num_points_per_voxel']
else:
voxels, coordinates, num_points = voxel_output
if not data_dict['use_lead_xyz']:
voxels = voxels[..., 3:] # remove xyz in voxels(N, 3)
#data_dict['dense_pillar'] = voxel_dense['voxels']
#data_dict['dense_pillar_coords'] = voxel_dense['coordinates']
data_dict['voxels'] = voxels
data_dict['voxel_coords'] = coordinates
data_dict['voxel_num_points'] = num_points
return data_dict
from spconv.utils import VoxelGeneratorV2
kitti_voxel_generator = VoxelGeneratorV2(
voxel_size=[0.05, 0.05, 0.1],
point_cloud_range=[0.0, -40.0, -3.0, 70.4, 40.0, 1.0],
max_num_points=5,
max_voxels=20000,
full_mean=False,
block_filtering=False,
block_factor=0,
block_size=0,
height_threshold=0.0)
