Skip to content

Commit 462af4e

Browse files
Kin-ZhangKin-Zhang
committed
feat(data): add av2 color info into pointcloud.
check av2 data issue here for more info: argoverse/av2-api#166
1 parent 161b555 commit 462af4e

2 files changed

Lines changed: 103 additions & 12 deletions

File tree

‎scripts/py/data/extract_argoverse2.py‎

Lines changed: 88 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9,23 +9,101 @@
99
# This script is licensed under the terms of the MIT license.
1010
# For a copy, see <https://opensource.org/licenses/MIT>.
1111

12-
import av2.utils.io as io_utils
13-
from av2.utils.typing import NDArrayFloat
12+
1413
import numpy as np
1514
from scipy.spatial.transform import Rotation as R
1615
from pathlib import Path
17-
16+
import pandas as pd
1817
from tqdm import tqdm
1918
import sys, os
2019
BASE_DIR = os.path.abspath(os.path.join(os.path.dirname( __file__ ), '..' ))
2120
sys.path.append(BASE_DIR)
2221

22+
import av2.utils.io as io_utils
23+
from av2.utils.typing import NDArrayFloat, NDArrayByte
24+
from av2.datasets.sensor.constants import RingCameras, StereoCameras
25+
from av2.utils.synchronization_database import SynchronizationDB, get_timestamps_from_sensor_folder
26+
from av2.geometry.se3 import SE3
27+
import av2.geometry.geometry as geometry_utils
28+
from av2.geometry.camera.pinhole_camera import PinholeCamera
29+
2330
from utils.pcdpy3 import save_pcd
2431
from utils import bc
2532

26-
DATA_FOLDER = "/home/kin/bags/av2/07YOTznatmYypvQYpzviEcU3yGPsyaGg__Spring_2020" # The one we used in paper
33+
DATA_FOLDER = "/home/kin/bags/av2/Dynamic_Map/07YOTznatmYypvQYpzviEcU3yGPsyaGg__Spring_2020" # The one we used in paper
2734
SAVE_FOLDER = "/home/kin/data/av2/pcd"
2835

36+
def convert_pose_dataframe_to_SE3(pose_df: pd.DataFrame) -> SE3:
37+
"""Convert a dataframe with parameterization of a single pose, to an SE(3) object.
38+
39+
Args:
40+
pose_df: parameterization of a single pose.
41+
42+
Returns:
43+
SE(3) object representing the egovehicle's pose in the city frame.
44+
"""
45+
qw, qx, qy, qz = pose_df[["qw", "qx", "qy", "qz"]].to_numpy().squeeze()
46+
tx_m, ty_m, tz_m = pose_df[["tx_m", "ty_m", "tz_m"]].to_numpy().squeeze()
47+
city_q_ego: NDArrayFloat = np.array([qw, qx, qy, qz])
48+
city_t_ego: NDArrayFloat = np.array([tx_m, ty_m, tz_m])
49+
city_R_ego = geometry_utils.quat_to_mat(quat_wxyz=city_q_ego)
50+
city_SE3_ego = SE3(rotation=city_R_ego, translation=city_t_ego)
51+
return city_SE3_ego
52+
53+
def get_city_SE3_ego(data_folder: str, timestamp_ns: int) -> SE3:
54+
log_poses_df = io_utils.read_feather(Path(data_folder/ "city_SE3_egovehicle.feather"))
55+
pose_df = log_poses_df.loc[log_poses_df["timestamp_ns"] == timestamp_ns]
56+
57+
if len(pose_df) == 0:
58+
raise RuntimeError("Pose was not available for the requested timestamp.")
59+
60+
city_SE3_ego = convert_pose_dataframe_to_SE3(pose_df)
61+
return city_SE3_ego
62+
63+
def get_colored_info(data_folder, lidar_timestamp_ns, sweep_lidar):
64+
lidar_timestamp_ns = int(lidar_timestamp_ns)
65+
log_id = data_folder.split('/')[-1]
66+
sweep_rgb: NDArrayByte = np.zeros((sweep_lidar.shape[0], 3), dtype=np.uint8)
67+
_sdb = SynchronizationDB("./")
68+
69+
# we will fix SynchronizationDB here:
70+
_sdb.per_log_cam_timestamps_index[log_id] = {}
71+
for cam_name in list(RingCameras) + list(StereoCameras):
72+
sensor_folder_wildcard = f"{data_folder}/sensors/cameras/{cam_name}/*.jpg"
73+
cam_timestamps = get_timestamps_from_sensor_folder(sensor_folder_wildcard)
74+
_sdb.per_log_cam_timestamps_index[log_id][cam_name] = cam_timestamps
75+
sensor_folder_wildcard = f"{data_folder}/sensors/lidar/*.feather"
76+
lidar_timestamps = get_timestamps_from_sensor_folder(sensor_folder_wildcard)
77+
_sdb.per_log_lidar_timestamps_index[log_id] = lidar_timestamps
78+
79+
for cam_enum in list(RingCameras):
80+
cam_name = cam_enum.value
81+
cam_timestamp_ns = _sdb.get_closest_cam_channel_timestamp(
82+
lidar_timestamp_ns, cam_name, log_id
83+
)
84+
if cam_timestamp_ns is None:
85+
continue
86+
img_fpath = Path(f"{DATA_FOLDER}/sensors/cameras/{cam_name}/{cam_timestamp_ns}.jpg")
87+
if not os.path.exists(img_fpath):
88+
continue
89+
cam_timestamp_ns = int(img_fpath.stem)
90+
91+
pinhole_camera = PinholeCamera.from_feather(Path(data_folder), cam_name=cam_name)
92+
city_SE3_ego_cam_t = get_city_SE3_ego(data_folder=Path(data_folder), timestamp_ns=cam_timestamp_ns)
93+
city_SE3_ego_lidar_t = get_city_SE3_ego(data_folder=Path(data_folder), timestamp_ns=lidar_timestamp_ns)
94+
95+
uv, points_cam, is_valid = pinhole_camera.project_ego_to_img_motion_compensated(
96+
points_lidar_time=data_xyz,
97+
city_SE3_ego_cam_t=city_SE3_ego_cam_t,
98+
city_SE3_ego_lidar_t=city_SE3_ego_lidar_t,
99+
)
100+
uv_valid = np.round(uv[is_valid]).astype(np.int64)
101+
u = uv_valid[:, 0]
102+
v = uv_valid[:, 1]
103+
img = io_utils.read_img(img_fpath, channel_order="RGB")
104+
sweep_rgb[is_valid] = img[v, u]
105+
return sweep_rgb
106+
29107
if __name__ == "__main__":
30108
cetner_scense_xyz = np.array([0,0,0])
31109
save_folder = DATA_FOLDER.split('/')[-1][:4]
@@ -57,21 +135,22 @@
57135

58136
# extract point cloud data with intensity
59137
sweep_df = io_utils.read_feather(abs_path_pts_file)
60-
data_xyzi: NDArrayFloat = sweep_df[["x", "y", "z", "intensity"]].to_numpy().astype(np.float64)
138+
data_xyz: NDArrayFloat = sweep_df[["x", "y", "z"]].to_numpy().astype(np.float64)
139+
data_color: NDArrayByte = get_colored_info(DATA_FOLDER, frame_index, data_xyz)
61140

62141
# detail view you can know the data frame now is based on ego
63142
ego2up_lidar.rotation = np.eye(3)
64-
data = ego2up_lidar.inverse().transform_point_cloud(data_xyzi[:, :3])
143+
data = ego2up_lidar.inverse().transform_point_cloud(data_xyz[:, :3])
65144

66145
# NOTE: We transform the point cloud to the world frame based on pose
67146
data = pose2origin.transform_point_cloud(data[:,:3])
68147

69-
data = np.hstack((data, data_xyzi[:, 3].reshape(-1, 1)))
148+
# data = np.hstack((data, data_xyzi[:, 3].reshape(-1, 1)))
70149

71150
# need qw, qx, qy, qz to save pcd in VIEWPOINT
72151
qxyzw = R.from_matrix(pose2origin.rotation).as_quat() # quat
73152
pose_array = [pose2origin.translation[0], pose2origin.translation[1], pose2origin.translation[2], \
74153
qxyzw[3], qxyzw[0], qxyzw[1], qxyzw[2]]
75-
save_pcd(f"{SAVE_FOLDER}/{i:06d}.pcd", data, pose_array)
76-
154+
save_pcd(f"{SAVE_FOLDER}/{i:06d}.pcd", data, pose_array, rgb=data_color)
155+
# break
77156
print(f"{bc.OKGREEN}Done{bc.ENDC} Check: {SAVE_FOLDER}..")

‎scripts/py/utils/pcdpy3.py‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ def load_pcd(path):
2424
data.xyzi2np()
2525
return data
2626

27-
def save_pcd(path, data: np.array, pos=np.array([0,0,0,1,0,0,0])):
27+
def save_pcd(path, data: np.array, pos=np.array([0,0,0,1,0,0,0]), rgb: np.array = None):
2828
# need be data: [N,3] or [N,4] [x,y,z,intensity]
2929
# pos: [x,y,z, qw,qx,qy,qz]
3030
# assert with print
@@ -55,13 +55,25 @@ def save_pcd(path, data: np.array, pos=np.array([0,0,0,1,0,0,0])):
5555
pc_data['y'] = data_float32[:, 1]
5656
pc_data['z'] = data_float32[:, 2]
5757
pc_data['intensity'] = data_float32[:, 3]
58-
58+
59+
elif(data.shape[1] == 3 and rgb is not None and data.shape[0] == rgb.shape[0]):
60+
md['fields'].append('rgb')
61+
md['size'].append(4)
62+
md['type'].append('F')
63+
md['count'].append(1)
64+
pc_data = np.zeros(data_float32.shape[0], dtype=[('x', np.float32), ('y', np.float32), ('z', np.float32), ('rgb', np.uint32)])
65+
pc_data['x'] = data_float32[:, 0]
66+
pc_data['y'] = data_float32[:, 1]
67+
pc_data['z'] = data_float32[:, 2]
68+
rgb = rgb.astype(np.uint32)
69+
pc_data['rgb'] = (rgb[:, 0] << 16) | (rgb[:, 1] << 8) | rgb[:, 2]
70+
5971
elif(data.shape[1] == 3):
6072
pc_data = np.zeros(data_float32.shape[0], dtype=[('x', np.float32), ('y', np.float32), ('z', np.float32)])
6173
pc_data['x'] = data_float32[:, 0]
6274
pc_data['y'] = data_float32[:, 1]
6375
pc_data['z'] = data_float32[:, 2]
64-
76+
6577
if pc_data is not None:
6678
pc = PointCloud(md, pc_data)
6779
point_cloud_to_fileobj(pc, fileobj)

0 commit comments

Comments
 (0)