Skip to content

Commit

Permalink
add compatibility for mmdetection3d dev-1.x
Browse files Browse the repository at this point in the history
  • Loading branch information
@TillBeemelmanns
TillBeemelmanns committed May 11, 2024
1 parent 2c0ba9d commit 3fb6c28
Showing 1 changed file with 83 additions and 51 deletions.
134 changes: 83 additions & 51 deletions converter/lidar_converter.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,33 @@
LIDAR_CORRUPTIONS = ['pointsreducing', 'beamsreducing', 'snow', 'fog', 'copy',
'spatialmisalignment', 'temporalmisalignment', 'motionblur']

def update_pathes_dev1x(infos):
"""
For .pkl files from mmdetection3d framework on branch dev-1.x
Update the pathes for the lidar files.
"""
for info in infos:
info['lidar_path'] = os.path.join('samples', 'LIDAR_TOP', info['lidar_points']['lidar_path'])
sweep_pathes = []
if 'lidar_sweeps' in info:
for sweep_info in info['lidar_sweeps']:
sweep_pathes.append({
"data_path": sweep_info['lidar_points']['lidar_path'][16:]
})
info['sweeps'] = sweep_pathes
return infos

def update_pathes(infos):
"""
For .pkl files from mmdetection3d framework older than 1.x.
Remove the first 16 characters from the data_path.
"""
for info in infos:
info['lidar_path'] = info['lidar_path'][16:]
for sweep in info['sweeps']:
sweep['data_path'] = sweep['data_path'][16:]
return infos


def parse_arguments():
parser = argparse.ArgumentParser(description='Generate corrupted nuScenes dataset for LiDAR')
Expand Down Expand Up @@ -45,11 +72,17 @@ def parse_arguments():

nusc_info = NuScenes(version='v1.0-trainval', dataroot=args.root_folder, verbose=True)

imageset = os.path.join(args.root_folder, 'deepinter_infos_val.pkl')
imageset = os.path.join(args.root_folder, 'nuscenes_infos_val.pkl')
with open(imageset, 'rb') as f:
infos = pickle.load(f)
all_files = infos['infos']
all_paths = copy.deepcopy(all_files)
infos = pickle.load(f)

if 'infos' in infos:
# mmdetection3d 1.4
all_infos = update_pathes(infos['infos'])
elif 'data_list' in infos:
all_infos = update_pathes_dev1x(infos['data_list'])
else:
exit("This mmdetection3d version is not supported.")

Path(args.dst_folder).mkdir(parents=True, exist_ok=True)
lidar_save_root = os.path.join(args.dst_folder , 'samples/LIDAR_TOP')
Expand All @@ -63,43 +96,42 @@ def parse_arguments():


def sweep_map(i: int) -> None:
info = all_paths[i]
lidar_path = info['lidar_path'][16:]
point = np.fromfile(os.path.join(args.root_folder, lidar_path), dtype=np.float32, count=-1).reshape([-1, 5])
sweep_info = info['sweeps']
info = all_infos[i]
lidar_path = info['lidar_path']
point = np.fromfile(os.path.join(args.root_folder, lidar_path), dtype=np.float32, count=-1).reshape([-1, 5])

if args.corruption == 'pointsreducing':
new_point = lidar.pointsreducing(point, args.severity)
for n in sweep_info:
sweep_path = n['data_path'][16:]
sweep_point = np.fromfile(os.path.join(args.root_folder, sweep_path), dtype=np.float32, count=-1).reshape([-1, 5])
for n in info['sweeps']:
sweep_path = n['data_path']
sweep_point = np.fromfile(os.path.join(args.root_folder, sweep_path), dtype=np.float32, count=-1).reshape([-1, 5])
new_sweep_point = lidar.pointsreducing(sweep_point, args.severity)
sweep_path = os.path.join(args.dst_folder, sweep_path)
new_sweep_point.astype(np.float32).tofile(sweep_path)

elif args.corruption == 'beamsreducing':
new_point = lidar.reduce_LiDAR_beamsV2(point, args.severity)
for n in sweep_info:
sweep_path = n['data_path'][16:]
sweep_point = np.fromfile(os.path.join(args.root_folder, sweep_path), dtype=np.float32, count=-1).reshape([-1, 5])
for n in info['sweeps']:
sweep_path = n['data_path']
sweep_point = np.fromfile(os.path.join(args.root_folder, sweep_path), dtype=np.float32, count=-1).reshape([-1, 5])
new_sweep_point = lidar.reduce_LiDAR_beamsV2(sweep_point, args.severity)
sweep_path = os.path.join(args.dst_folder, sweep_path)
new_sweep_point.astype(np.float32).tofile(sweep_path)

elif args.corruption == 'motionblur':
new_point = lidar.pts_motion(point, args.severity)
for n in sweep_info:
sweep_path = n['data_path'][16:]
sweep_point = np.fromfile(os.path.join(args.root_folder, sweep_path), dtype=np.float32, count=-1).reshape([-1, 5])
for n in info['sweeps']:
sweep_path = n['data_path']
sweep_point = np.fromfile(os.path.join(args.root_folder, sweep_path), dtype=np.float32, count=-1).reshape([-1, 5])
new_sweep_point = lidar.pts_motion(sweep_point, args.severity)
sweep_path = os.path.join(args.dst_folder, sweep_path)
new_sweep_point.astype(np.float32).tofile(sweep_path)

elif args.corruption == 'snow':
lidar_sd_token = nusc_info.get('sample', info['token'])['data']['LIDAR_TOP']
lidar_seg_label_path = nusc_info.get('lidarseg', lidar_sd_token)['filename']
lidarseg_labels_filename = os.path.join(args.root_folder, lidar_seg_label_path)
label = np.fromfile(lidarseg_labels_filename, dtype=np.uint8).reshape([-1, 1])
lidar_seg_label_path = os.path.join(args.root_folder, lidar_seg_label_path)
label = np.fromfile(lidar_seg_label_path, dtype=np.uint8).reshape([-1, 1])

_, new_point, _ = lidar.simulate_snow(
point,
Expand All @@ -110,9 +142,9 @@ def sweep_map(i: int) -> None:
shuffle=True
)

for n in sweep_info:
sweep_path = n['data_path'][16:]
sweep_point = np.fromfile(os.path.join(args.root_folder, sweep_path), dtype=np.float32, count=-1).reshape([-1, 5])
for n in info['sweeps']:
sweep_path = n['data_path']
sweep_point = np.fromfile(os.path.join(args.root_folder, sweep_path), dtype=np.float32, count=-1).reshape([-1, 5])
_, new_sweep_point = lidar.simulate_snow_sweep(
sweep_point,
args.severity,
Expand All @@ -125,58 +157,58 @@ def sweep_map(i: int) -> None:

elif args.corruption == 'fog':
new_point, _, _, _ = lidar.simulate_fog(args.severity, point, 10)
for n in sweep_info:
sweep_path = n['data_path'][16:]
sweep_point = np.fromfile(os.path.join(args.root_folder, sweep_path), dtype=np.float32, count=-1).reshape([-1, 5])
for n in info['sweeps']:
sweep_path = n['data_path']
sweep_point = np.fromfile(os.path.join(args.root_folder, sweep_path), dtype=np.float32, count=-1).reshape([-1, 5])
new_sweep_point, _, _, _ = lidar.simulate_fog(args.severity, sweep_point, 10)
sweep_path = os.path.join(args.dst_folder, sweep_path)
new_sweep_point.astype(np.float32).tofile(sweep_path)

elif args.corruption == 'spatialmisalignment':
new_point = lidar.transform_points(point, args.severity)
for n in sweep_info:
sweep_path = n['data_path'][16:]
sweep_point = np.fromfile(os.path.join(args.root_folder, sweep_path), dtype=np.float32, count=-1).reshape([-1, 5])
for n in info['sweeps']:
sweep_path = n['data_path']
sweep_point = np.fromfile(os.path.join(args.root_folder, sweep_path), dtype=np.float32, count=-1).reshape([-1, 5])
new_sweep_point = lidar.transform_points(sweep_point, args.severity)
sweep_path = os.path.join(args.dst_folder, sweep_path)
new_sweep_point.astype(np.float32).tofile(sweep_path)

elif args.corruption == 'temporalmisalignment':
sweep_list = []
if len(sweep_info) == 0:
if len(info['sweeps']) == 0:
pass
else:
for m in sweep_info:
sweep_m_path = m['data_path'][16:]
for n in info['sweeps']:
sweep_m_path = n['data_path']
sweep_list.append(sweep_m_path)

if i >= 2:
prev_info = all_paths[i-1]
prev_info = all_infos[i-1]
prev_sweep_list = []
prev_sweep_info=prev_info['sweeps']
prev_sweep_info = prev_info['sweeps']
if len(prev_sweep_info) == 0:
pass
else:
for m in prev_sweep_info:
prev_sweep_m_path = m['data_path'][16:]
for n in prev_sweep_info:
prev_sweep_m_path = n['data_path']
prev_sweep_list.append(prev_sweep_m_path)

s = [0.2, 0.4, 0.6][args.severity - 1]

if np.random.rand() < s and i >= 2:
prev_lidar_path = prev_info['lidar_path'][16:]
prev_point = np.fromfile(os.path.join(args.root_folder, prev_lidar_path), dtype=np.float32, count=-1).reshape([-1, 5])
prev_lidar_path = prev_info['lidar_path']
prev_point = np.fromfile(os.path.join(args.root_folder, prev_lidar_path), dtype=np.float32, count=-1).reshape([-1, 5])
new_point = prev_point.copy()

if len(prev_sweep_list) != len(sweep_list):
for n in sweep_list:
new_sweep_point = np.fromfile(os.path.join(args.root_folder, n), dtype=np.float32, count=-1).reshape([-1, 5])
new_sweep_point = np.fromfile(os.path.join(args.root_folder, n), dtype=np.float32, count=-1).reshape([-1, 5])
sweep_path = os.path.join(args.dst_folder, n)
new_sweep_point.astype(np.float32).tofile(sweep_path)
else:
prev_points = []
for n in prev_sweep_list:
new_sweep_point = np.fromfile(os.path.join(args.root_folder, n), dtype=np.float32, count=-1).reshape([-1, 5])
new_sweep_point = np.fromfile(os.path.join(args.root_folder, n), dtype=np.float32, count=-1).reshape([-1, 5])
prev_points.append(new_sweep_point)

for idx, n in enumerate(sweep_list):
Expand All @@ -186,15 +218,15 @@ def sweep_map(i: int) -> None:
else:
new_point = point.copy()
for n in sweep_list:
new_sweep_point = np.fromfile(os.path.join(args.root_folder, n), dtype=np.float32, count=-1).reshape([-1, 5])
new_sweep_point = np.fromfile(os.path.join(args.root_folder, n), dtype=np.float32, count=-1).reshape([-1, 5])
sweep_path = os.path.join(args.dst_folder, n)
new_sweep_point.astype(np.float32).tofile(sweep_path)

elif args.corruption == "copy":
new_point = point
for n in sweep_info:
sweep_path = n['data_path'][16:]
sweep_point = np.fromfile(os.path.join(args.root_folder, sweep_path), dtype=np.float32, count=-1).reshape([-1, 5])
for n in info['sweeps']:
sweep_path = n['data_path']
sweep_point = np.fromfile(os.path.join(args.root_folder, sweep_path), dtype=np.float32, count=-1).reshape([-1, 5])
sweep_path = os.path.join(args.dst_folder, sweep_path)
sweep_point.astype(np.float32).tofile(sweep_path)

Expand All @@ -206,9 +238,9 @@ def sweep_map(i: int) -> None:


def sample_map(i: int) -> None:
info = all_paths[i]
lidar_path = info['lidar_path'][16:]
point = np.fromfile(os.path.join(args.root_folder, lidar_path), dtype=np.float32, count=-1).reshape([-1, 5])
info = all_infos[i]
lidar_path = info['lidar_path']
point = np.fromfile(os.path.join(args.root_folder, lidar_path), dtype=np.float32, count=-1).reshape([-1, 5])

if args.corruption == 'pointsreducing':
new_point = lidar.pointsreducing(point, args.severity)
Expand All @@ -235,20 +267,20 @@ def sample_map(i: int) -> None:
)

elif args.corruption == 'fog':
new_point, _, _, _ = lidar.simulate_fog(args.severity, point, 10)
new_point, _, _, _ = lidar.simulate_fog(args.severity, point, 10)

elif args.corruption == 'spatialmisalignment':
new_point = lidar.transform_points(point, args.severity)

elif args.corruption == 'temporalmisalignment':
if i >= 2:
prev_info = all_paths[i-1]
prev_info = all_infos[i-1]

s = [0.2, 0.4, 0.6][args.severity - 1]

if np.random.rand() < s and i >= 2:
prev_lidar_path = prev_info['lidar_path'][16:]
prev_point = np.fromfile(os.path.join(args.root_folder, prev_lidar_path), dtype=np.float32, count=-1).reshape([-1, 5])
prev_lidar_path = prev_info['lidar_path']
prev_point = np.fromfile(os.path.join(args.root_folder, prev_lidar_path), dtype=np.float32, count=-1).reshape([-1, 5])
new_point = prev_point.copy()

else:
Expand All @@ -264,7 +296,7 @@ def sample_map(i: int) -> None:
new_point.astype(np.float32).tofile(lidar_save_path)


length = len(all_files)
length = len(all_infos)
if args.sweep:
with mp.Pool(args.n_cpus) as pool:
l = list(tqdm(pool.imap(sweep_map, range(length)), total=length))
Expand Down

0 comments on commit 3fb6c28

Please sign in to comment.