Skip to content

GitLab

Unverified Commit 5a658d61 authored by ZLTJohn's avatar ZLTJohn Committed by GitHub
Browse files

[Feature]: Support waymo_info_test.pkl and ImageSets generating (#2180)

parent aecc7a17
Show whitespace changes
Inline Side-by-side
Showing with 130 additions and 90 deletions
tools/create_data.py
View file @ 5a658d61
......@@ -199,6 +199,10 @@ def waymo_data_prep(root_path,
test_mode=(split
in ['testing', 'testing_3d_camera_only_detection']))
converter.convert()
from tools.dataset_converters.waymo_converter import \
create_ImageSets_img_ids
create_ImageSets_img_ids(osp.join(out_dir, 'kitti_format'), splits)
# Generate waymo infos
out_dir = osp.join(out_dir, 'kitti_format')
kitti.create_waymo_info_file(
......@@ -206,9 +210,11 @@ def waymo_data_prep(root_path,
info_train_path = osp.join(out_dir, f'{info_prefix}_infos_train.pkl')
info_val_path = osp.join(out_dir, f'{info_prefix}_infos_val.pkl')
info_trainval_path = osp.join(out_dir, f'{info_prefix}_infos_trainval.pkl')
test_path = osp.join(out_dir, f'{info_prefix}_infos_test.pkl')
update_pkl_infos('waymo', out_dir=out_dir, pkl_path=info_train_path)
update_pkl_infos('waymo', out_dir=out_dir, pkl_path=info_val_path)
update_pkl_infos('waymo', out_dir=out_dir, pkl_path=info_trainval_path)
update_pkl_infos('waymo', out_dir=out_dir, pkl_path=test_path)
GTDatabaseCreater(
'WaymoDataset',
out_dir,
......
tools/dataset_converters/update_infos_to_v2.py
View file @ 5a658d61
......@@ -919,8 +919,9 @@ def update_waymo_infos(pkl_path, out_dir):
'point_cloud']['num_features']
temp_data_info['lidar_points']['timestamp'] = ori_info_dict[
'timestamp']
temp_data_info['lidar_points']['lidar_path'] = Path(
ori_info_dict['point_cloud']['velodyne_path']).name
velo_path = ori_info_dict['point_cloud'].get('velodyne_path')
if velo_path is not None:
temp_data_info['lidar_points']['lidar_path'] = Path(velo_path).name
# TODO discuss the usage of Tr_velo_to_cam in lidar
Trv2c = ori_info_dict['calib']['Tr_velo_to_cam'].astype(np.float32)
......@@ -953,10 +954,11 @@ def update_waymo_infos(pkl_path, out_dir):
temp_data_info['lidar_sweeps'].append(lidar_sweep)
temp_data_info['image_sweeps'].append(image_sweep)
anns = ori_info_dict['annos']
anns = ori_info_dict.get('annos')
ignore_class_name = set()
if anns is not None:
num_instances = len(anns['name'])
ignore_class_name = set()
instance_list = []
for instance_id in range(num_instances):
empty_instance = get_empty_instance()
......@@ -974,19 +976,20 @@ def update_waymo_infos(pkl_path, out_dir):
loc = anns['location'][instance_id]
dims = anns['dimensions'][instance_id]
rots = anns['rotation_y'][:, None][instance_id]
gt_bboxes_3d = np.concatenate([loc, dims,
rots]).astype(np.float32).tolist()
gt_bboxes_3d = np.concatenate([loc, dims, rots
]).astype(np.float32).tolist()
empty_instance['bbox_3d'] = gt_bboxes_3d
empty_instance['bbox_label_3d'] = copy.deepcopy(
empty_instance['bbox_label'])
empty_instance['bbox'] = anns['bbox'][instance_id].tolist()
empty_instance['truncated'] = int(
anns['truncated'][instance_id].tolist())
empty_instance['occluded'] = anns['occluded'][instance_id].tolist()
empty_instance['occluded'] = anns['occluded'][
instance_id].tolist()
empty_instance['alpha'] = anns['alpha'][instance_id].tolist()
empty_instance['index'] = anns['index'][instance_id].tolist()
empty_instance['group_id'] = anns['group_ids'][instance_id].tolist(
)
empty_instance['group_id'] = anns['group_ids'][
instance_id].tolist()
empty_instance['difficulty'] = anns['difficulty'][
instance_id].tolist()
empty_instance['num_lidar_pts'] = anns['num_points_in_gt'][
......@@ -998,9 +1001,10 @@ def update_waymo_infos(pkl_path, out_dir):
temp_data_info['instances'] = instance_list
# waymo provide the labels that sync with cam
anns = ori_info_dict['cam_sync_annos']
num_instances = len(anns['name'])
anns = ori_info_dict.get('cam_sync_annos')
ignore_class_name = set()
if anns is not None:
num_instances = len(anns['name'])
instance_list = []
for instance_id in range(num_instances):
empty_instance = get_empty_instance()
......@@ -1018,19 +1022,20 @@ def update_waymo_infos(pkl_path, out_dir):
loc = anns['location'][instance_id]
dims = anns['dimensions'][instance_id]
rots = anns['rotation_y'][:, None][instance_id]
gt_bboxes_3d = np.concatenate([loc, dims,
rots]).astype(np.float32).tolist()
gt_bboxes_3d = np.concatenate([loc, dims, rots
]).astype(np.float32).tolist()
empty_instance['bbox_3d'] = gt_bboxes_3d
empty_instance['bbox_label_3d'] = copy.deepcopy(
empty_instance['bbox_label'])
empty_instance['bbox'] = anns['bbox'][instance_id].tolist()
empty_instance['truncated'] = int(
anns['truncated'][instance_id].tolist())
empty_instance['occluded'] = anns['occluded'][instance_id].tolist()
empty_instance['occluded'] = anns['occluded'][
instance_id].tolist()
empty_instance['alpha'] = anns['alpha'][instance_id].tolist()
empty_instance['index'] = anns['index'][instance_id].tolist()
empty_instance['group_id'] = anns['group_ids'][instance_id].tolist(
)
empty_instance['group_id'] = anns['group_ids'][
instance_id].tolist()
empty_instance['camera_id'] = anns['camera_id'][
instance_id].tolist()
empty_instance = clear_instance_unused_keys(empty_instance)
......
tools/dataset_converters/waymo_converter.py
View file @ 5a658d61
......@@ -9,8 +9,9 @@ except ImportError:
raise ImportError('Please run "pip install waymo-open-dataset-tf-2-5-0" '
'>1.4.5 to install the official devkit first.')
import os
from glob import glob
from os.path import join
from os.path import exists, join
import mmengine
import numpy as np
......@@ -132,8 +133,6 @@ class Waymo2KITTI(object):
self.save_image(frame, file_idx, frame_idx)
self.save_calib(frame, file_idx, frame_idx)
if 'testing_3d_camera_only_detection' not in self.load_dir:
# the camera only split doesn't contain lidar points.
self.save_lidar(frame, file_idx, frame_idx)
self.save_pose(frame, file_idx, frame_idx)
self.save_timestamp(frame, file_idx, frame_idx)
......@@ -230,6 +229,9 @@ class Waymo2KITTI(object):
range_images, camera_projections, seg_labels, range_image_top_pose = \
parse_range_image_and_camera_projection(frame)
if range_image_top_pose is None:
# the camera only split doesn't contain lidar points.
return
# First return
points_0, cp_points_0, intensity_0, elongation_0, mask_indices_0 = \
self.convert_range_image_to_point_cloud(
......@@ -601,3 +603,30 @@ class Waymo2KITTI(object):
else:
raise ValueError(mat.shape)
return ret
def create_ImageSets_img_ids(root_dir, splits):
save_dir = join(root_dir, 'ImageSets/')
if not exists(save_dir):
os.mkdir(save_dir)
idx_all = [[] for i in splits]
for i, split in enumerate(splits):
path = join(root_dir, splits[i], 'calib')
if not exists(path):
RawNames = []
else:
RawNames = os.listdir(path)
for name in RawNames:
if name.endswith('.txt'):
idx = name.replace('.txt', '\n')
idx_all[int(idx[0])].append(idx)
idx_all[i].sort()
open(save_dir + 'train.txt', 'w').writelines(idx_all[0])
open(save_dir + 'val.txt', 'w').writelines(idx_all[1])
open(save_dir + 'trainval.txt', 'w').writelines(idx_all[0] + idx_all[1])
open(save_dir + 'test.txt', 'w').writelines(idx_all[2])
# open(save_dir+'test_cam_only.txt','w').writelines(idx_all[3])
print('created txt files indicating what to collect in ', splits)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment