import pickle
import matplotlib.pyplot as plt
import torch
position_range = [-32.0, -32.0, -10.0, 32.0, 32.0, 10.0]
depth_num = 64
depth_start = 1
def position_embedding(features, img_features):
eps = 1e-5
img_features = img_features.unsqueeze(1)
B, N, C, tar_H, tar_W = img_features.shape
device = img_features.device
crop_top = 28
crop_left = 416
H = [16 for _ in range(3)]
W = [
64 * 1088 // (1088 * 2 + 1920),
64 * 1920 // (1088 * 2 + 1920),
64 * 1088 // (1088 * 2 + 1920)
]
# 左视图(16,17)
coords_h_l = torch.arange(H[0], device=device).float() * 1080 / H[0] + crop_top / H[0]
coords_w_l = torch.arange(W[0], device=device).float() * 1920 / W[0] + crop_left / W[0]
# 前视图(16,30)
coords_h_f = torch.arange(H[1], device=device).float() * 1080 / H[1] + crop_top / H[1]
coords_w_f = torch.arange(W[1], device=device).float() * 1920 / W[1]
# 右视图(16,17)
coords_h_r = torch.arange(H[2], device=device).float() * 1080 / H[2] + crop_top / H[2]
coords_w_r = torch.arange(W[2], device=device).float() * 1920 / W[2] + crop_left / W[2]
index = torch.arange(start=0, end=depth_num, step=1, device=img_features.device).float()
index_1 = index + 1
bin_size = (position_range[3] - depth_start) / (depth_num * (1 + depth_num))
coords_d = depth_start + bin_size * index * index_1
D = coords_d.shape[0]
coords = [1] * 3 # 0,1,2 -> front, left, right
coords[0] = torch.stack(torch.meshgrid([coords_w_l, coords_h_l, coords_d])).permute(1, 2, 3, 0) # W, H, D, 3
coords[1] = torch.stack(torch.meshgrid([coords_w_f, coords_h_f, coords_d])).permute(1, 2, 3, 0) # W, H, D, 3
coords[2] = torch.stack(torch.meshgrid([coords_w_r, coords_h_r, coords_d])).permute(1, 2, 3, 0) # W, H, D, 3
# coords = torch.cat((coords, torch.ones_like(coords[..., :1])), -1)
coords[0][..., :2] = coords[0][..., :2] * torch.max(coords[0][..., 2:3], torch.ones_like(coords[0][..., 2:3]) * eps)
coords[1][..., :2] = coords[1][..., :2] * torch.max(coords[1][..., 2:3], torch.ones_like(coords[1][..., 2:3]) * eps)
coords[2][..., :2] = coords[2][..., :2] * torch.max(coords[2][..., 2:3], torch.ones_like(coords[2][..., 2:3]) * eps)
# img_meta
# img2lidars = ?
pos_3d_embed = None
for i in range(3):
sensor2lidar_rotation = features["sensor2lidar_rotation"][i]
sensor2lidar_translation = features["sensor2lidar_translation"][i]
intrinsics = features["intrinsics"][i]
combine = torch.matmul(sensor2lidar_rotation, torch.inverse(intrinsics)).float() # (B, 1, 3, 3) ?
# print(combine.shape)
# coords_front,coords_fleft,coords_fright (W, H, D, 3)
# coords3d = torch.stack((coords_front, coords_fleft, coords_fright), dim=0) # (N, W, H, D, 3) -> (B, N, W, H, D, 3, 1)
# coords = coords.view(1, H, W, D, 1, 3).repeat(B, 1, 1, 1, 1, 1)
coords3d = coords[i].view(1, N, W[i], H[i], D, 3, 1).repeat(B, 1, 1, 1, 1, 1,
1) # (B, N, W, H, D, 3, 1) -> (B, N, W, H, D, 3, 3)
combine = combine.view(B, N, 1, 1, 1, 3, 3).repeat(1, 1, W[i], H[i], D, 1, 1)
coords3d = torch.matmul(combine, coords3d).squeeze(-1) # (B, N, W, H, D, 3)
sensor2lidar_translation = sensor2lidar_translation.view(B, N, 1, 1, 1, 3)
coords3d += sensor2lidar_translation
coords3d[..., 0:1] = (coords3d[..., 0:1] - position_range[0]) / (
position_range[3] - position_range[0])
coords3d[..., 1:2] = (coords3d[..., 1:2] - position_range[1]) / (
position_range[4] - position_range[1])
coords3d[..., 2:3] = (coords3d[..., 2:3] - position_range[2]) / (
position_range[5] - position_range[2])
# coords_mask = (coords3d > 1.0) | (coords3d < 0.0)
# coords_mask = coords_mask.flatten(-2).sum(-1) > (D * 0.5)
# coords_mask = coords_mask.permute(0, 1, 3, 2)
# for j in range(1000000):
# print(coords3d.shape)
# (2, 1, 17, 16, 64, 3) -> (B, N, W, H, D, 3)
# (2, 1, 30, 16, 64, 3)
# -> (2, 1, 17+30+17, 16, 64, 3)
# coords3d = coords3d.permute(0, 1, 4, 5, 3, 2).contiguous().view(B * N, -1, H[i], W[i])
if pos_3d_embed is None:
pos_3d_embed = coords3d
else:
pos_3d_embed = torch.cat((pos_3d_embed, coords3d), dim=2)
# for i in range(100000):
# print(img_features.shape)
pos_3d_embed = pos_3d_embed.permute(0, 1, 4, 5, 3, 2).contiguous().view(B * N, -1, tar_H, tar_W)
return pos_3d_embed
if __name__ == '__main__':
H, W = 16, 64
logs = pickle.load(open('/mnt/g/navsim/navsim_logs/tiny/2021.05.12.22.28.35_veh-35_00620_01164.pkl', 'rb'))
log = logs[0]
features = {
'sensor2lidar_rotation': [
torch.from_numpy(log['cams']['CAM_L0']['sensor2lidar_rotation']),
torch.from_numpy(log['cams']['CAM_F0']['sensor2lidar_rotation']),
torch.from_numpy(log['cams']['CAM_R0']['sensor2lidar_rotation']),
],
'sensor2lidar_translation': [
torch.from_numpy(log['cams']['CAM_L0']['sensor2lidar_translation']),
torch.from_numpy(log['cams']['CAM_F0']['sensor2lidar_translation']),
torch.from_numpy(log['cams']['CAM_R0']['sensor2lidar_translation']),
],
'intrinsics': [
torch.from_numpy(log['cams']['CAM_L0']['cam_intrinsic']),
torch.from_numpy(log['cams']['CAM_F0']['cam_intrinsic']),
torch.from_numpy(log['cams']['CAM_R0']['cam_intrinsic']),
]
}
img_features = torch.randn((1, 3, H, W))
coords_3d = position_embedding(
features, img_features
)
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
for i in range(H):
for j in range(W):
frustum_points = coords_3d.permute(2, 3, 1, 0).reshape(H, W, depth_num, 3)
pixel_points = frustum_points[i, j]
x_points = pixel_points[:, 0]
y_points = pixel_points[:, 1]
z_points = pixel_points[:, 2]
ax.scatter(x_points, y_points, z_points)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.view_init(elev=90, azim=0)
ax.set_zlabel('Z')
plt.show()