import numpy as np import torch from scipy.spatial.transform import Rotation from camviz import Camera as CameraCV import numpy as np from anydata.utils.data import cat_channel_ones, unnorm_pixel_grid from torch_scatter import scatter_min # from pytorch3d.renderer.fisheyecameras import FishEyeCameras from anydata.utils.geometry import invert_extrinsics from anydata.geometry.camera_utils import multiply_extrinsics def project_pointcloud(uv_all, z_all, hw): """Creates a depth map from a pointcloud (requires torch-scatter)""" i = 0 # Loop over point-cloud batch-wise depths_tgt = [] # Flatten coordinates and depth values uv, z = uv_all[i].reshape(-1, 2), z_all[i].reshape(-1, 1) # Remove out-of-bounds coordinates and points behind the camera idx = (uv[:, 0] > 0) & (uv[:, 0] < 300) & \ (uv[:, 1] > 0) & (uv[:, 1] < 300) & (z[:, 0] > 0.0) # Unormalize and stack coordinates for scatter operation # uv = (unnorm_pixel_grid(uv[idx], hw)).floor().long() uv = uv[idx].floor().long() uv = uv[:, 0] + uv[:, 1] * hw[1] # Min scatter operation (only keep the closest depth) depth_tgt = 1e10 * torch.ones((hw[0] * hw[1], 1), device=z.device, dtype=z.dtype) depth_tgt, argmin = scatter_min(src=z[idx], index=uv.unsqueeze(1), dim=0, out=depth_tgt) depth_tgt[depth_tgt == 1e10] = 0. num_valid = (depth_tgt > 0).sum() # Substitute invalid values with zero invalid = argmin == argmin.max() argmin[invalid] = 0 # Reshape outputs depth_tgt = depth_tgt.reshape(1, 1, hw[0], hw[1]) depths_tgt.append(depth_tgt) # Return stacked depth maps return torch.cat(depths_tgt, 0) class FThetaCameraModel: def __init__(self, cx, cy, coefficients): """ cx, cy: Principal point coefficients: [k1, k2, k3, k4] for r = k1*theta + k2*theta^2 ... """ self.cx = cx self.cy = cy self.k = coefficients def project(self, pcl): x, y, z = pcl[:, 0], pcl[:, 1], pcl[:, 2] r = np.sqrt(x ** 2 + y ** 2) theta = np.arctan2(r, z) theta_d = self.k[0] * theta \ + self.k[1] * theta ** 2 \ + self.k[2] * theta ** 3 \ + self.k[3] * theta ** 4 # # Avoid division by zero # if r == 0: # return self.cx, self.cy u = self.cx + (x / r) * theta_d v = self.cy + (y / r) * theta_d return u, v, z def backproject(self, u, v): """Pixel coordinate to 3D Ray""" dx = u - self.cx dy = v - self.cy r_d = np.sqrt(dx**2 + dy**2) # This requires solving the polynomial for theta_d, # often done numerically or by using a precomputed look-up table. # ... (Inverse polynomial mapping) ... return unit_ray def pose_to_matrix(trans, quat): """Convert 7D pose [x, y, z, qx, qy, qz, qw] to 4x4 transformation matrix""" T = np.eye(4) T[:3, :3] = Rotation.from_quat(quat).as_matrix() T[:3, 3] = trans return T from transforms3d.euler import euler2quat, quat2euler from transforms3d.quaternions import quat2mat, mat2quat, axangle2quat, quat2axangle import cv2 def pose_array_quat_2_matrix(pose:np.ndarray): '''transform pose array of quaternion to transformation matrix Param: pose: 7d vector, with t(3d) + q(4d) ---------- Return: mat: 4x4 matrix, with R,T,0,1 form ''' if pose.shape != (7,): raise ValueError mat = quat2mat([pose[3], pose[4], pose[5], pose[6]]) return np.array([[mat[0][0], mat[0][1], mat[0][2], pose[0]], [mat[1][0], mat[1][1], mat[1][2], pose[1]], [mat[2][0], mat[2][1], mat[2][2], pose[2]], [0,0,0,1]]) clip_id = "/data/cv_downloaded/PhysicalAI/seq01/2edf278f-d5e3-4b83-b5df-923a04335725" data = np.load(f"{clip_id}.npz", allow_pickle=True)["data"][()] intrinsics = data['camera_intrinsics'] extrinsics = data['sensor_extrinsics'] time_series = data['time_series'][0] # print(time_series.keys()) images = time_series['camera_images'] pointcloud = time_series['point_cloud'] # print(time_series['ego_state']['position']) # print(time_series['ego_state']['rotation']) # print(time_series['ego_state'].keys()) # import sys # sys.exit() val = data['sensor_extrinsics']['lidar_top_360fov'] quat = np.array([val['qx'], val['qy'], val['qz'], val['qw']]) trans = np.array([val['x'], val['y'], val['z']]) lidar_pose = pose_to_matrix(trans, quat) # lidar_pose = invert_extrinsics(lidar_pose) lidar_pose = torch.tensor(lidar_pose).unsqueeze(0).float() from anydata.utils.write import write_image from anydata.utils.viz import viz_depth cvcams = [] rgbs, cams, depths = [], [], [] backs = [] pointcloud = torch.tensor(pointcloud).unsqueeze(0).float().permute(0, 2, 1) pointcloud = multiply_extrinsics(pointcloud, lidar_pose) for key in images.keys(): rgb = images[key] val = intrinsics[key] cx, cy = val['cx'], val['cy'] k0f = val['fw_poly_0'] k1f = val['fw_poly_1'] k2f = val['fw_poly_2'] k3f = val['fw_poly_3'] k4f = val['fw_poly_4'] k0b = val['bw_poly_0'] k1b = val['bw_poly_1'] k2b = val['bw_poly_2'] k3b = val['bw_poly_3'] k4b = val['bw_poly_4'] K = torch.tensor([ cx, cy, k0f, k1f, k2f, k3f, k4f, k0b, k1b, k2b, k3b, k4b, ]).unsqueeze(0) from camviz import Draw draw = Draw((1600, 600)) draw.add3Dworld('wld', (0.6, 0.0, 1.0, 1.0), # pose=(-8.16675, -14.55655, 17.26451, 0.44646, 0.72559, -0.46922, 0.23241), ) draw.add2Dimage('rgb', (0.0, 0.0, 0.6, 0.5), res=(1920,1080)) draw.add2Dimage('dep', (0.0, 0.5, 0.6, 1.0), res=(1920,1080)) val = extrinsics[key] quat = np.array([val['qx'], val['qy'], val['qz'], val['qw']]) trans = np.array([val['x'], val['y'], val['z']]) pose = pose_to_matrix(trans, quat) pose = torch.tensor(pose).unsqueeze(0).float() rgb = images[key] / 255 rgb = torch.tensor(rgb).permute(2,0,1).unsqueeze(0) print(K) from anydata.geometry.camera import Camera cam = Camera(K=K, Tcw=pose, hw=rgb, geometry='ftheta') depth = cam.project_pointcloud(pointcloud) back = cam.reconstruct_depth_map(depth, to_world=True) backs.append(back) viz = viz_depth(depth[0], filter_zeros=True) viz = torch.tensor(viz).permute(2, 0, 1).unsqueeze(0) rgbs.append(rgb) depths.append(viz) cams.append(cam) cvcams.append(CameraCV.from_anydata(cam)) draw.addBufferf('pcl', pointcloud[0].permute(1, 0)) for i, rgb in enumerate(rgbs): draw.addTexture(f'rgb{i}', rgb[0]) for i, depth in enumerate(depths): draw.addTexture(f'dep{i}', depth[0]) for i, back in enumerate(backs): draw.addBufferf(f'bck{i}', back[0]) t = 0 while draw.input(): if draw.UP: t = (t + 1) % len(images) draw.halt(100) if draw.DOWN: t = (t - 1) % len(images) draw.halt(100) draw.clear() draw['wld'].size(3).color('whi').points('pcl') for i, cam in enumerate(cvcams): draw['wld'].color('red').size(3).points(f'bck{i}') draw['wld'].object(cam, tex=f'rgb{i}') draw['rgb'].image(f'rgb{t}') draw['dep'].image(f'dep{t}') draw.update(30)