diff --git a/Dockerfile b/Dockerfile index 72374ae02f70cae2864bf8287b0b85969d64d238..73d8f2b9d75732d508203dbbb4ef07e80d52be44 100644 --- a/Dockerfile +++ b/Dockerfile @@ -15,9 +15,9 @@ ENV PATH /app/miniconda/bin:$PATH SHELL ["conda", "run","--no-capture-output", "-p","/app/env", "/bin/bash", "-c"] -COPY --chown=1000:1000 ./web_server.py /app/web_server.py -COPY --chown=1000:1000 ./docker/web_server_config/scene-0383-medium-00.yaml /app/docker/web_server_config/scene-0383-medium-00.yaml -COPY --chown=1000:1000 ./download_pre_datas.py /app/download_pre_datas.py +COPY --chown=1000:1000 ./code /app/code +COPY --chown=1000:1000 ./web_server.py /app/code/web_server.py +COPY --chown=1000:1000 ./download_pre_datas.py /app/code/download_pre_datas.py ENV TCNN_CUDA_ARCHITECTURES 75 @@ -32,6 +32,6 @@ RUN ./.pixi/envs/default/bin/python3 -m pip install psutil RUN ./.pixi/envs/default/bin/python3 -m pip install moviepy -RUN ./.pixi/envs/default/bin/python /app/download_pre_datas.py +RUN ./.pixi/envs/default/bin/python /app/code/download_pre_datas.py -CMD ["./.pixi/envs/default/bin/python", "web_server.py"] +CMD ["./.pixi/envs/default/bin/python", "/app/code/web_server.py"] diff --git a/code/gaussian_renderer/__init__.py b/code/gaussian_renderer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..90e86ae48eaaa341f393d11319cfe26b0c630860 --- /dev/null +++ b/code/gaussian_renderer/__init__.py @@ -0,0 +1,231 @@ +import torch +from scene.gaussian_model import GaussianModel +from scene.ground_model import GroundModel +from gsplat.rendering import rasterization +import roma +from scene.cameras import Camera +from torch import Tensor + +def euler2matrix(yaw): + return torch.tensor([ + [torch.cos(-yaw), 0, torch.sin(-yaw)], + [0, 1, 0], + [-torch.sin(-yaw), 0, torch.cos(-yaw)] + ]).cuda() + +def cat_bgfg(bg, fg, only_xyz=False): + if only_xyz: + if bg.ground_model is None: + bg_feats = [bg.get_xyz] + else: + bg_feats = [bg.get_full_xyz] + else: + if bg.ground_model is None: + bg_feats = [bg.get_xyz, bg.get_opacity, bg.get_scaling, bg.get_rotation, bg.get_features, bg.get_3D_features] + else: + bg_feats = [bg.get_full_xyz, bg.get_full_opacity, bg.get_full_scaling, bg.get_full_rotation, bg.get_full_features, bg.get_full_3D_features] + + + if len(fg) == 0: + return bg_feats + + output = [] + for fg_feat, bg_feat in zip(fg, bg_feats): + if fg_feat is None: + output.append(bg_feat) + else: + if bg_feat.shape[1] != fg_feat.shape[1]: + fg_feat = fg_feat[:, :bg_feat.shape[1], :] + output.append(torch.cat((bg_feat, fg_feat), dim=0)) + + return output + +def concatenate_all(all_fg): + output = [] + for feat in list(zip(*all_fg)): + output.append(torch.cat(feat, dim=0)) + return output + +def proj_uv(xyz, cam): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + intr = torch.as_tensor(cam.K[:3, :3]).float().to(device) # (3, 3) + w2c = torch.linalg.inv(cam.c2w)[:3, :] # (3, 4) + + c_xyz = (w2c[:3, :3] @ xyz.T).T + w2c[:3, 3] + i_xyz = (intr @ c_xyz.mT).mT # (N, 3) + uv = i_xyz[:, [1,0]] / i_xyz[:, -1:].clip(1e-3) # (N, 2) + return uv + + +def unicycle_b2w(timestamp, model): + pred = model(timestamp) + if pred is None: + return None + pred_a, pred_b, pred_v, pitchroll, pred_yaw, pred_h = pred + rt = torch.eye(4).float().cuda() + rt[:3,:3] = roma.euler_to_rotmat('xzy', [-pitchroll[0]+torch.pi/2, -pitchroll[1]+torch.pi/2, -pred_yaw+torch.pi/2]) + rt[1, 3], rt[0, 3], rt[2, 3] = pred_h, pred_a, pred_b + return rt + + +def render(viewpoint:Camera, prev_viewpoint:Camera, pc:GaussianModel, dynamic_gaussians:dict, + unicycles:dict, bg_color:Tensor, render_optical=False, planning=[]): + """ + Render the scene. + + Background tensor (bg_color) must be on GPU! + """ + timestamp = viewpoint.timestamp + + all_fg = [None, None, None, None, None, None] + prev_all_fg = [None] + + if unicycles is None or len(unicycles) == 0: + track_dict = viewpoint.dynamics + if prev_viewpoint is not None: + prev_track_dict = prev_viewpoint.dynamics + else: + track_dict, prev_track_dict = {}, {} + for track_id, B2W in viewpoint.dynamics.items(): + if track_id in unicycles: + B2W = unicycle_b2w(timestamp, unicycles[track_id]['model']) + track_dict[track_id] = B2W + if prev_viewpoint is not None: + prev_B2W = unicycle_b2w(prev_viewpoint.timestamp, unicycles[track_id]['model']) + prev_track_dict[track_id] = prev_B2W + if len(planning) > 0: + for plan_id, B2W in planning[0].items(): + track_dict[plan_id] = B2W + if prev_viewpoint is not None: + for plan_id, B2W in planning[1].items(): + prev_track_dict[plan_id] = B2W + + all_fg, prev_all_fg = [], [] + for track_id, B2W in track_dict.items(): + w_dxyz = (B2W[:3, :3] @ dynamic_gaussians[track_id].get_xyz.T).T + B2W[:3, 3] + + drot = roma.quat_wxyz_to_xyzw(dynamic_gaussians[track_id].get_rotation) + drot = roma.unitquat_to_rotmat(drot) + w_drot = roma.quat_xyzw_to_wxyz(roma.rotmat_to_unitquat(B2W[:3, :3] @ drot)) + fg = [w_dxyz, + dynamic_gaussians[track_id].get_opacity, + dynamic_gaussians[track_id].get_scaling, + w_drot, + # dynamic_gaussians[track_id].get_rotation, + dynamic_gaussians[track_id].get_features, + dynamic_gaussians[track_id].get_3D_features] + + all_fg.append(fg) + + if render_optical and prev_viewpoint is not None: + if track_id in prev_track_dict: + prev_B2W = prev_track_dict[track_id] + prev_w_dxyz = torch.mm(prev_B2W[:3, :3], dynamic_gaussians[track_id].get_xyz.T).T + prev_B2W[:3, 3] + prev_all_fg.append([prev_w_dxyz]) + else: + prev_all_fg.append([w_dxyz]) + + all_fg = concatenate_all(all_fg) + xyz, opacities, scales, rotations, shs, feats3D = cat_bgfg(pc, all_fg) + + if render_optical and prev_viewpoint is not None: + prev_all_fg = concatenate_all(prev_all_fg) + prev_xyz = cat_bgfg(pc, prev_all_fg, only_xyz=True)[0] + uv = proj_uv(xyz, viewpoint) + prev_uv = proj_uv(prev_xyz, prev_viewpoint) + delta_uv = prev_uv - uv + delta_uv = torch.cat([delta_uv, torch.ones_like(delta_uv[:, :1], device=delta_uv.device)], dim=-1) + else: + delta_uv = torch.zeros_like(xyz) + + if pc.affine: + cam_xyz, cam_dir = viewpoint.c2w[:3, 3].cuda(), viewpoint.c2w[:3, 2].cuda() + o_enc = pc.pos_enc(cam_xyz[None, :] / 60) + d_enc = pc.dir_enc(cam_dir[None, :]) + appearance = pc.appearance_model(torch.cat([o_enc, d_enc], dim=1)) * 1e-1 + affine_weight, affine_bias = appearance[:, :9].view(3, 3), appearance[:, -3:] + affine_weight = affine_weight + torch.eye(3, device=appearance.device) + + if render_optical: + render_mode = 'RGB+ED+S+F' + else: + render_mode = 'RGB+ED+S' + + renders, render_alphas, info = rasterization( + means=xyz, + quats=rotations, + scales=scales, + opacities=opacities[:, 0], + colors=shs, + viewmats=torch.linalg.inv(viewpoint.c2w)[None, ...], # [C, 4, 4] + Ks=viewpoint.K[None, :3, :3], # [C, 3, 3] + width=viewpoint.width, + height=viewpoint.height, + smts=feats3D[None, ...], + flows= delta_uv[None, ...], + render_mode=render_mode, + sh_degree=pc.active_sh_degree, + near_plane=0.01, + far_plane=500, + packed=False, + backgrounds=bg_color[None, :], + ) + + renders = renders[0] + rendered_image = renders[..., :3].permute(2,0,1) + depth = renders[..., 3][None, ...] + smt = renders[..., 4:(4+feats3D.shape[-1])].permute(2,0,1) + + if pc.affine: + colors = rendered_image.view(3, -1).permute(1, 0) # (H*W, 3) + refined_image = (colors @ affine_weight + affine_bias).clip(0, 1).permute(1, 0).view(*rendered_image.shape) + else: + refined_image = rendered_image + + return {"render": refined_image, + "feats": smt, + "depth": depth, + "opticalflow": renders[..., -2:].permute(2,0,1) if render_optical else None, + "alphas": render_alphas, + "viewspace_points": info["means2d"], + "info": info, + } + + +def render_ground(viewpoint:Camera, pc:GroundModel, bg_color:Tensor): + xyz, opacities, scales = pc.get_xyz, pc.get_opacity, pc.get_scaling + rotations, shs, feats3D = pc.get_rotation, pc.get_features, pc.get_3D_features + + K = viewpoint.K[None, :3, :3] + renders, render_alphas, info = rasterization( + means=xyz, + quats=rotations, + scales=scales, + opacities=opacities[:, 0], + colors=shs, + viewmats=torch.linalg.inv(viewpoint.c2w)[None, ...], # [C, 4, 4] + Ks=K, # [C, 3, 3] + width=viewpoint.width, + height=viewpoint.height, + smts=feats3D[None, ...], + render_mode='RGB+ED+S', + sh_degree=pc.active_sh_degree, + near_plane=0.01, + far_plane=500, + packed=False, + backgrounds=bg_color[None, :], + ) + + renders = renders[0] + rendered_image = renders[..., :3].permute(2,0,1) + depth = renders[..., 3][None, ...] + smt = renders[..., 4:(4+feats3D.shape[-1])].permute(2,0,1) + + return {"render": rendered_image, + "feats": smt, + "depth": depth, + "opticalflow": None, + "alphas": render_alphas, + "viewspace_points": info["means2d"], + "info": info, + } \ No newline at end of file diff --git a/code/gaussian_renderer/__pycache__/__init__.cpython-311.pyc b/code/gaussian_renderer/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..140929fcb8c6252cb6c1680afa189e352b191334 Binary files /dev/null and b/code/gaussian_renderer/__pycache__/__init__.cpython-311.pyc differ diff --git a/code/scene/__init__.py b/code/scene/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a265fec07f3e26fc97c963e793381dcad3d6ea2f --- /dev/null +++ b/code/scene/__init__.py @@ -0,0 +1,111 @@ +import os +import random +import json +from utils.system_utils import searchForMaxIteration +from scene.dataset_readers import sceneLoadTypeCallbacks +from scene.gaussian_model import GaussianModel +from scene.obj_model import ObjModel +from scene.cameras import cameraList_from_camInfos +import torch +import open3d as o3d +import numpy as np +import shutil + + +def load_cameras(args, data_type, ignore_dynamic=False): + train_cameras = {} + test_cameras = {} + if os.path.exists(os.path.join(args.source_path, "meta_data.json")): + print("Found meta_data.json file, assuming HUGSIM format data set!") + scene_info = sceneLoadTypeCallbacks['HUGSIM'](args.source_path, data_type, ignore_dynamic) + else: + assert False, "Could not recognize scene type! "+args.source_path + + print("Loading Training Cameras") + train_cameras = cameraList_from_camInfos(scene_info.train_cameras, args) + print("Loading Test Cameras") + test_cameras = cameraList_from_camInfos(scene_info.test_cameras, args) + return train_cameras, test_cameras, scene_info + +class Scene: + + def __init__(self, args, gaussians:GaussianModel, load_iteration=None, shuffle=True, + data_type='kitti360', ignore_dynamic=False, planning=None): + """b + :param path: Path to colmap scene main folder. + """ + self.model_path = args.model_path + self.loaded_iter = None + self.gaussians = gaussians + self.data_type = data_type + + if load_iteration: + if load_iteration == -1: + self.loaded_iter = searchForMaxIteration(os.path.join(self.model_path, "ckpts")) + else: + self.loaded_iter = load_iteration + print("Loading trained model at iteration {}".format(self.loaded_iter)) + + self.train_cameras, self.test_cameras, scene_info = load_cameras(args, data_type, ignore_dynamic) + + self.dynamic_verts = scene_info.verts + self.dynamic_gaussians = {} + for track_id in scene_info.verts: + self.dynamic_gaussians[track_id] = ObjModel(args.model.sh_degree, feat_mutable=False) + if planning is not None: + for plan_id in planning.keys(): + self.dynamic_gaussians[plan_id] = ObjModel(args.model.sh_degree, feat_mutable=False) + + if not self.loaded_iter: + shutil.copyfile(scene_info.ply_path, os.path.join(self.model_path, "input.ply")) + shutil.copyfile(os.path.join(args.source_path, 'meta_data.json'), os.path.join(self.model_path, 'meta_data.json')) + shutil.copyfile(os.path.join(args.source_path, 'ground_param.pkl'), os.path.join(self.model_path, 'ground_param.pkl')) + + if shuffle: + random.shuffle(scene_info.train_cameras) + random.shuffle(scene_info.test_cameras) + + self.cameras_extent = scene_info.nerf_normalization["radius"] + + if self.loaded_iter: + (model_params, first_iter) = torch.load(os.path.join(self.model_path, "ckpts", f"chkpnt{self.loaded_iter}.pth")) + gaussians.restore(model_params, None) + for iid, dynamic_gaussian in self.dynamic_gaussians.items(): + if planning is None or iid not in planning: + (model_params, first_iter) = torch.load(os.path.join(self.model_path, "ckpts", f"dynamic_{iid}_chkpnt{self.loaded_iter}.pth")) + dynamic_gaussian.restore(model_params, None) + else: + (model_params, first_iter) = torch.load(planning[iid]) + model_params = list(model_params) + model_params.append(None) + dynamic_gaussian.restore(model_params, None) + # for iid, unicycle_pkg in self.unicycles.items(): + # model_params = torch.load(os.path.join(self.model_path, "ckpts", f"unicycle_{iid}_chkpnt{self.loaded_iter}.pth")) + # unicycle_pkg['model'].restore(model_params) + + else: + self.gaussians.create_from_pcd(scene_info.point_cloud, self.cameras_extent) + for track_id in self.dynamic_gaussians.keys(): + vertices = scene_info.verts[track_id] + + # init from template + l, h, w = vertices[:, 0].max() - vertices[:, 0].min(), vertices[:, 1].max() - vertices[:, 1].min(), vertices[:, 2].max() - vertices[:, 2].min() + pcd = o3d.io.read_point_cloud(f"utils/vehicle_template/benz_{data_type}.ply") + points = np.array(pcd.points) * np.array([l, h, w]) + pcd.points = o3d.utility.Vector3dVector(points) + pcd.colors = o3d.utility.Vector3dVector(np.ones_like(points) * 0.5) + + self.dynamic_gaussians[track_id].create_from_pcd(pcd, self.cameras_extent) + + def save(self, iteration): + # self.gaussians.save_ply(os.path.join(point_cloud_path, "point_cloud.ply")) + point_cloud_vis_path = os.path.join(self.model_path, "point_cloud_vis/iteration_{}".format(iteration)) + self.gaussians.save_vis_ply(os.path.join(point_cloud_vis_path, "point.ply")) + for iid, dynamic_gaussian in self.dynamic_gaussians.items(): + dynamic_gaussian.save_vis_ply(os.path.join(point_cloud_vis_path, f"dynamic_{iid}.ply")) + + def getTrainCameras(self): + return self.train_cameras + + def getTestCameras(self): + return self.test_cameras \ No newline at end of file diff --git a/code/scene/__pycache__/__init__.cpython-311.pyc b/code/scene/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9906ef04b0ed104d42c01493a291924e675febb7 Binary files /dev/null and b/code/scene/__pycache__/__init__.cpython-311.pyc differ diff --git a/code/scene/__pycache__/cameras.cpython-311.pyc b/code/scene/__pycache__/cameras.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..16d39e936ffcfd2b9282a14c0653716c4e29b404 Binary files /dev/null and b/code/scene/__pycache__/cameras.cpython-311.pyc differ diff --git a/code/scene/__pycache__/dataset_readers.cpython-311.pyc b/code/scene/__pycache__/dataset_readers.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6e11b9b46cf22759b78d0bb684d1042b846b7734 Binary files /dev/null and b/code/scene/__pycache__/dataset_readers.cpython-311.pyc differ diff --git a/code/scene/__pycache__/gaussian_model.cpython-311.pyc b/code/scene/__pycache__/gaussian_model.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2b979185711c68e4c2235bcc1626aa4c9d205aff Binary files /dev/null and b/code/scene/__pycache__/gaussian_model.cpython-311.pyc differ diff --git a/code/scene/__pycache__/ground_model.cpython-311.pyc b/code/scene/__pycache__/ground_model.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56e37262223239179eb64a659c291dfd27aa7ec9 Binary files /dev/null and b/code/scene/__pycache__/ground_model.cpython-311.pyc differ diff --git a/code/scene/__pycache__/obj_model.cpython-311.pyc b/code/scene/__pycache__/obj_model.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0bd689427d9fad145e8cf497f24621cbeb32231a Binary files /dev/null and b/code/scene/__pycache__/obj_model.cpython-311.pyc differ diff --git a/code/scene/cameras.py b/code/scene/cameras.py new file mode 100644 index 0000000000000000000000000000000000000000..81fb43c3bfa3ab104f6d88e5056a6763764c9305 --- /dev/null +++ b/code/scene/cameras.py @@ -0,0 +1,71 @@ +import torch +from torch import nn + +class Camera(nn.Module): + def __init__(self, width, height, image, K, c2w, + image_name, data_device="cuda", + semantic2d=None, depth=None, mask=None, timestamp=-1, optical_image=None, dynamics={} + ): + super(Camera, self).__init__() + + try: + self.data_device = torch.device(data_device) + except Exception as e: + print(e) + print(f"[Warning] Custom device {data_device} failed, fallback to default cuda device" ) + self.data_device = torch.device("cuda") + + self.width = width + self.height = height + self.image_name = image_name + self.timestamp = timestamp + self.K = torch.from_numpy(K).float().cuda() + self.c2w = torch.from_numpy(c2w).float().cuda() + self.dynamics = dynamics + + self.original_image = torch.from_numpy(image).permute(2,0,1).float().clamp(0.0, 1.0).to(self.data_device) + if semantic2d is not None: + self.semantic2d = semantic2d.to(self.data_device) + else: + self.semantic2d = None + if depth is not None: + self.depth = depth.to(self.data_device) + else: + self.depth = None + if mask is not None: + self.mask = torch.from_numpy(mask).bool().to(self.data_device) + else: + self.mask = None + self.image_width = self.original_image.shape[2] + self.image_height = self.original_image.shape[1] + if optical_image is not None: + self.optical_gt = torch.from_numpy(optical_image).to(self.data_device) + else: + self.optical_gt = None + + +def loadCam(args, cam_info): + + if cam_info.semantic2d is not None: + semantic2d = torch.from_numpy(cam_info.semantic2d).long()[None, ...] + else: + semantic2d = None + + optical_image = cam_info.optical_image + mask = cam_info.mask + depth = cam_info.depth + + gt_image = cam_info.image[..., :3] / 255. + + return Camera(K=cam_info.K, c2w=cam_info.c2w, width=cam_info.width, height=cam_info.height, + image=gt_image, image_name=cam_info.image_name, data_device=args.model.data_device, + semantic2d=semantic2d, depth=depth, mask=mask, + timestamp=cam_info.timestamp, optical_image=optical_image, dynamics=cam_info.dynamics) + +def cameraList_from_camInfos(cam_infos, args): + camera_list = [] + + for c in cam_infos: + camera_list.append(loadCam(args, c)) + + return camera_list diff --git a/code/scene/dataset_readers.py b/code/scene/dataset_readers.py new file mode 100644 index 0000000000000000000000000000000000000000..a21f797953119b02562297b3052633a0cbecb7aa --- /dev/null +++ b/code/scene/dataset_readers.py @@ -0,0 +1,212 @@ +import os +from typing import NamedTuple +import numpy as np +import json +from plyfile import PlyData, PlyElement +from utils.sh_utils import SH2RGB +from scene.gaussian_model import BasicPointCloud +import torch.nn.functional as F +from imageio.v2 import imread +import torch + + +class CameraInfo(NamedTuple): + K: np.array + c2w: np.array + image: np.array + image_path: str + image_name: str + width: int + height: int + semantic2d: np.array + optical_image: np.array + depth: torch.tensor + mask: np.array + timestamp: int + dynamics: dict + +class SceneInfo(NamedTuple): + point_cloud: BasicPointCloud + train_cameras: list + test_cameras: list + nerf_normalization: dict + ply_path: str + verts: dict + +def getNerfppNorm(cam_info, data_type): + def get_center_and_diag(cam_centers): + cam_centers = np.hstack(cam_centers) + avg_cam_center = np.mean(cam_centers, axis=1, keepdims=True) + center = avg_cam_center + dist = np.linalg.norm(cam_centers - center, axis=0, keepdims=True) + diagonal = np.max(dist) + return center.flatten(), diagonal + + cam_centers = [] + for cam in cam_info: + cam_centers.append(cam.c2w[:3, 3:4]) # cam_centers in world coordinate + + radius = 10 + + return {'radius': radius} + +def fetchPly(path): + plydata = PlyData.read(path) + vertices = plydata['vertex'] + positions = np.vstack([vertices['x'], vertices['y'], vertices['z']]).T + if 'red' in vertices: + colors = np.vstack([vertices['red'], vertices['green'], vertices['blue']]).T / 255.0 + else: + print('Create random colors') + shs = np.ones((positions.shape[0], 3)) * 0.5 + colors = SH2RGB(shs) + normals = np.zeros((positions.shape[0], 3)) + return BasicPointCloud(points=positions, colors=colors, normals=normals) + +def storePly(path, xyz, rgb): + # Define the dtype for the structured array + dtype = [('x', 'f4'), ('y', 'f4'), ('z', 'f4'), + ('nx', 'f4'), ('ny', 'f4'), ('nz', 'f4'), + ('red', 'u1'), ('green', 'u1'), ('blue', 'u1')] + + normals = np.zeros_like(xyz) + + elements = np.empty(xyz.shape[0], dtype=dtype) + attributes = np.concatenate((xyz, normals, rgb), axis=1) + elements[:] = list(map(tuple, attributes)) + + # Create the PlyData object and write to file + vertex_element = PlyElement.describe(elements, 'vertex') + ply_data = PlyData([vertex_element]) + ply_data.write(path) + +def readHUGSIMCameras(path, data_type, ignore_dynamic): + train_cam_infos, test_cam_infos = [], [] + with open(os.path.join(path, 'meta_data.json')) as json_file: + meta_data = json.load(json_file) + + verts = {} + if 'verts' in meta_data and not ignore_dynamic: + verts_list = meta_data['verts'] + for k, v in verts_list.items(): + verts[k] = np.array(v) + + frames = meta_data['frames'] + for idx, frame in enumerate(frames): + c2w = np.array(frame['camtoworld']) + + rgb_path = os.path.join(path, frame['rgb_path'].replace('./', '')) + + rgb_split = rgb_path.split('/') + image_name = '_'.join([rgb_split[-2], rgb_split[-1][:-4]]) + image = imread(rgb_path) + + semantic_2d = None + semantic_pth = rgb_path.replace("images", "semantics").replace('.png', '.npy').replace('.jpg', '.npy') + if os.path.exists(semantic_pth): + semantic_2d = np.load(semantic_pth) + semantic_2d[(semantic_2d == 14) | (semantic_2d == 15)] = 13 + + optical_path = rgb_path.replace("images", "flow").replace('.png', '_flow.npy').replace('.jpg', '_flow.npy') + if os.path.exists(optical_path): + optical_image = np.load(optical_path) + else: + optical_image = None + + depth_path = rgb_path.replace("images", "depth").replace('.png', '.pt').replace('.jpg', '.pt') + if os.path.exists(depth_path): + depth = torch.load(depth_path, weights_only=True) + else: + depth = None + + mask = None + mask_path = rgb_path.replace("images", "masks").replace('.png', '.npy').replace('.jpg', '.npy') + if os.path.exists(mask_path): + mask = np.load(mask_path) + + timestamp = frame.get('timestamp', -1) + + intrinsic = np.array(frame['intrinsics']) + + dynamics = {} + if 'dynamics' in frame and not ignore_dynamic: + dynamics_list = frame['dynamics'] + for iid in dynamics_list.keys(): + dynamics[iid] = torch.tensor(dynamics_list[iid]).cuda() + + cam_info = CameraInfo(K=intrinsic, c2w=c2w, image=np.array(image), + image_path=rgb_path, image_name=image_name, height=image.shape[0], + width=image.shape[1], semantic2d=semantic_2d, + optical_image=optical_image, depth=depth, mask=mask, timestamp=timestamp, dynamics=dynamics) + + if data_type == 'kitti360': + if idx < 20: + train_cam_infos.append(cam_info) + elif idx % 20 < 16: + train_cam_infos.append(cam_info) + elif idx % 20 >= 16: + test_cam_infos.append(cam_info) + else: + continue + + elif data_type == 'kitti': + if idx < 10 or idx >= len(frames) - 4: + train_cam_infos.append(cam_info) + elif idx % 4 < 2: + train_cam_infos.append(cam_info) + elif idx % 4 == 2: + test_cam_infos.append(cam_info) + else: + continue + + elif data_type == "nuscenes": + if idx % 30 >= 24: + test_cam_infos.append(cam_info) + else: + train_cam_infos.append(cam_info) + + elif data_type == "waymo": + if idx % 15 >= 12: + test_cam_infos.append(cam_info) + else: + train_cam_infos.append(cam_info) + + elif data_type == "pandaset": + if idx > 30 and idx % 30 >= 24: + test_cam_infos.append(cam_info) + else: + train_cam_infos.append(cam_info) + + else: + raise NotImplementedError + + return train_cam_infos, test_cam_infos, verts + + +def readHUGSIMInfo(path, data_type, ignore_dynamic): + train_cam_infos, test_cam_infos, verts = readHUGSIMCameras(path, data_type, ignore_dynamic) + + print(f'Loaded {len(train_cam_infos)} train cameras and {len(test_cam_infos)} test cameras') + nerf_normalization = getNerfppNorm(train_cam_infos, data_type) + + ply_path = os.path.join(path, "points3d.ply") + if not os.path.exists(ply_path): + assert False, "Requires for initialize 3d points as inputs" + try: + pcd = fetchPly(ply_path) + except Exception as e: + print('When loading point clound, meet error:', e) + exit(0) + + scene_info = SceneInfo(point_cloud=pcd, + train_cameras=train_cam_infos, + test_cameras=test_cam_infos, + nerf_normalization=nerf_normalization, + ply_path=ply_path, + verts=verts) + return scene_info + + +sceneLoadTypeCallbacks = { + "HUGSIM": readHUGSIMInfo, +} \ No newline at end of file diff --git a/code/scene/gaussian_model.py b/code/scene/gaussian_model.py new file mode 100644 index 0000000000000000000000000000000000000000..bc89bf12b0a57a60c4662a4de464b9b8287fd82a --- /dev/null +++ b/code/scene/gaussian_model.py @@ -0,0 +1,636 @@ +import torch +import numpy as np +from utils.general_utils import inverse_sigmoid, get_expon_lr_func, build_rotation +from torch import nn +import os +from utils.system_utils import mkdir_p +from plyfile import PlyData, PlyElement +from utils.sh_utils import RGB2SH, SH2RGB +from simple_knn._C import distCUDA2 +from utils.graphics_utils import BasicPointCloud +from utils.general_utils import strip_symmetric, build_scaling_rotation +import open3d as o3d +import tinycudann as tcnn +from math import sqrt +from scene.ground_model import GroundModel +from io import BytesIO + + +class GaussianModel: + + def setup_functions(self): + def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation): + L = build_scaling_rotation(scaling_modifier * scaling, rotation) + actual_covariance = L @ L.transpose(1, 2) + symm = strip_symmetric(actual_covariance) + return symm + + self.scaling_activation = torch.exp + self.scaling_inverse_activation = torch.log + + self.covariance_activation = build_covariance_from_scaling_rotation + + self.opacity_activation = torch.sigmoid + self.inverse_opacity_activation = torch.logit + + self.rotation_activation = torch.nn.functional.normalize + + + def __init__(self, sh_degree : int, feat_mutable=True, affine=False, ground_args=None): + self.active_sh_degree = 0 + self.max_sh_degree = sh_degree + self._xyz = torch.empty(0) + self._features_dc = torch.empty(0) + self._features_rest = torch.empty(0) + self._feats3D = torch.empty(0) + self._scaling = torch.empty(0) + self._rotation = torch.empty(0) + self._opacity = torch.empty(0) + self.max_radii2D = torch.empty(0) + self.xyz_gradient_accum = torch.empty(0) + self.denom = torch.empty(0) + self.optimizer = None + self.percent_dense = 0 + self.spatial_lr_scale = 0 + self.feat_mutable = feat_mutable + self.setup_functions() + + self.pos_enc = tcnn.Encoding( + n_input_dims=3, + encoding_config={"otype": "Frequency", "n_frequencies": 2}, + ) + self.dir_enc = tcnn.Encoding( + n_input_dims=3, + encoding_config={ + "otype": "SphericalHarmonics", + "degree": 3, + }, + ) + + self.affine = affine + if affine: + self.appearance_model = tcnn.Network( + n_input_dims=self.pos_enc.n_output_dims + self.dir_enc.n_output_dims, + n_output_dims=12, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "None", + "n_neurons": 32, + "n_hidden_layers": 2, + } + ) + else: + self.appearance_model = None + + if ground_args: + self.ground_model = GroundModel(sh_degree, model_args=ground_args, finetune=True) + else: + self.ground_model = None + + def capture(self): + if self.ground_model is not None: + ground_model_params = self.ground_model.capture() + else: + ground_model_params = None + return ( + self.active_sh_degree, + self._xyz, + self._features_dc, + self._features_rest, + self._feats3D, + self._scaling, + self._rotation, + self._opacity, + self.spatial_lr_scale, + self.appearance_model.state_dict(), + ground_model_params, + ) + + def restore(self, model_args, training_args): + (self.active_sh_degree, + self._xyz, + self._features_dc, + self._features_rest, + self._feats3D, + self._scaling, + self._rotation, + self._opacity, + self.spatial_lr_scale, + appearance_state_dict, + ground_model_params, + ) = model_args + self.appearance_model.load_state_dict(appearance_state_dict, strict=False) + if training_args is not None: + self.training_setup(training_args) + if ground_model_params is not None: + self.ground_model = GroundModel(self.max_sh_degree, model_args=ground_model_params) + + @property + def get_scaling(self): + return self.scaling_activation(self._scaling) + + @property + def get_full_scaling(self): + assert self.ground_model is not None + return torch.cat([self.scaling_activation(self._scaling), self.ground_model.get_scaling]) + + @property + def get_rotation(self): + return self.rotation_activation(self._rotation) + + @property + def get_full_rotation(self): + assert self.ground_model is not None + return torch.cat([self.rotation_activation(self._rotation), self.ground_model.get_rotation]) + + @property + def get_xyz(self): + return self._xyz + + @property + def get_full_xyz(self): + assert self.ground_model is not None + return torch.cat([self._xyz, self.ground_model.get_xyz]) + + @property + def get_features(self): + features_dc = self._features_dc + features_rest = self._features_rest + return torch.cat((features_dc, features_rest), dim=1) + + @property + def get_full_features(self): + assert self.ground_model is not None + sh = torch.cat((self._features_dc, self._features_rest), dim=1) + return torch.cat([sh, self.ground_model.get_features]) + + @property + def get_3D_features(self): + return torch.softmax(self._feats3D, dim=-1) + + @property + def get_full_3D_features(self): + assert self.ground_model is not None + return torch.cat([torch.softmax(self._feats3D, dim=-1), self.ground_model.get_3D_features]) + + @property + def get_opacity(self): + return self.opacity_activation(self._opacity) + + @property + def get_full_opacity(self): + assert self.ground_model is not None + return torch.cat([self.opacity_activation(self._opacity), self.ground_model.get_opacity]) + + # def get_covariance(self, scaling_modifier = 1): + # return self.covariance_activation(self.get_scaling, scaling_modifier, self._rotation) + + def oneupSHdegree(self): + if self.active_sh_degree < self.max_sh_degree: + self.active_sh_degree += 1 + + def create_from_pcd(self, pcd : BasicPointCloud, spatial_lr_scale : float): + # self.spatial_lr_scale = 1 + self.spatial_lr_scale = spatial_lr_scale + fused_point_cloud = torch.tensor(np.asarray(pcd.points)).float().cuda() + fused_color = RGB2SH(torch.tensor(np.asarray(pcd.colors)).float().cuda()) + features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda() + features[:, :3, 0 ] = fused_color + features[:, 3:, 1:] = 0.0 + + if self.feat_mutable: + feats3D = torch.rand(fused_color.shape[0], 20).float().cuda() + self._feats3D = nn.Parameter(feats3D.requires_grad_(True)) + else: + feats3D = torch.zeros(fused_color.shape[0], 20).float().cuda() + feats3D[:, 13] = 1 + self._feats3D = feats3D + + print("Number of points at initialization : ", fused_point_cloud.shape[0]) + + dist2 = torch.clamp_min(distCUDA2(torch.from_numpy(np.asarray(pcd.points)).float().cuda()), 0.0000001) + scales = torch.log(torch.sqrt(dist2))[...,None].repeat(1, 3) + rots = torch.zeros((fused_point_cloud.shape[0], 4), device="cuda") + rots[:, 0] = 1 + + opacities = inverse_sigmoid(0.1 * torch.ones((fused_point_cloud.shape[0], 1), dtype=torch.float, device="cuda")) + + self._xyz = nn.Parameter(fused_point_cloud.requires_grad_(True)) + self._features_dc = nn.Parameter(features[:,:,0:1].transpose(1, 2).contiguous().requires_grad_(True)) + self._features_rest = nn.Parameter(features[:,:,1:].transpose(1, 2).contiguous().requires_grad_(True)) + self._scaling = nn.Parameter(scales.requires_grad_(True)) + self._rotation = nn.Parameter(rots.requires_grad_(True)) + self._opacity = nn.Parameter(opacities.requires_grad_(True)) + self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda") + + def training_setup(self, training_args): + self.percent_dense = training_args.percent_dense + self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda") + self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda") + + # self.spatial_lr_scale /= 3 + + l = [ + {'params': [self._xyz], 'lr': training_args.position_lr_init*self.spatial_lr_scale, "name": "xyz"}, + {'params': [self._features_dc], 'lr': training_args.feature_lr, "name": "f_dc"}, + {'params': [self._features_rest], 'lr': training_args.feature_lr / 20.0, "name": "f_rest"}, + {'params': [self._opacity], 'lr': training_args.opacity_lr, "name": "opacity"}, + {'params': [self._scaling], 'lr': training_args.scaling_lr*self.spatial_lr_scale, "name": "scaling"}, + {'params': [self._rotation], 'lr': training_args.rotation_lr, "name": "rotation"}, + ] + + if self.affine: + l.append({'params': [*self.appearance_model.parameters()], 'lr': 1e-3, "name": "appearance_model"}) + + if self.feat_mutable: + l.append({'params': [self._feats3D], 'lr': 1e-2, "name": "feats3D"}) + + if self.ground_model is not None: + self.ground_optimizer = self.ground_model.optimizer + else: + self.ground_optimizer = None + + self.optimizer = torch.optim.Adam(l, lr=0.0, eps=1e-15) + self.xyz_scheduler_args = get_expon_lr_func(lr_init=training_args.position_lr_init*self.spatial_lr_scale, + lr_final=training_args.position_lr_final*self.spatial_lr_scale, + lr_delay_mult=training_args.position_lr_delay_mult, + max_steps=training_args.position_lr_max_steps) + + def update_learning_rate(self, iteration): + ''' Learning rate scheduling per step ''' + for param_group in self.optimizer.param_groups: + if param_group["name"] == "xyz": + lr = self.xyz_scheduler_args(iteration) + param_group['lr'] = lr + return lr + + def construct_list_of_attributes(self): + l = ['x', 'y', 'z', 'nx', 'ny', 'nz'] + # All channels except the 3 DC + for i in range(self._features_dc.shape[1]*self._features_dc.shape[2]): + l.append('f_dc_{}'.format(i)) + for i in range(self._features_rest.shape[1]*self._features_rest.shape[2]): + l.append('f_rest_{}'.format(i)) + for i in range(self._feats3D.shape[1]): + l.append('semantic_{}'.format(i)) + l.append('opacity') + for i in range(self._scaling.shape[1]): + l.append('scale_{}'.format(i)) + for i in range(self._rotation.shape[1]): + l.append('rot_{}'.format(i)) + return l + + def save_ply(self, path=None): + mkdir_p(os.path.dirname(path)) + + if self.ground_model is not None: + xyz = self.get_full_xyz.detach().cpu().numpy() + normals = np.zeros_like(xyz) + f_dc = torch.cat([self._features_dc, self.ground_model._features_dc]).detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy() + f_rest = torch.cat([self._features_rest, self.ground_model._features_rest]).detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy() + feats3D = torch.cat([self._feats3D, self.ground_model._feats3D]).detach().cpu().numpy() + opacities = torch.cat([self._opacity, self.ground_model._opacity]).detach().cpu().numpy() + scale = self.scaling_inverse_activation(self.get_full_scaling).detach().cpu().numpy() + rotation = torch.cat([self._rotation, self.ground_model._rotation]).detach().cpu().numpy() + else: + xyz = self.get_xyz.detach().cpu().numpy() + normals = np.zeros_like(xyz) + f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy() + f_rest = self._features_rest.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy() + feats3D = self._feats3D.detach().cpu().numpy() + opacities = self._opacity.detach().cpu().numpy() + scale = self.scaling_inverse_activation(self.get_scaling).detach().cpu().numpy() + rotation = self._rotation.detach().cpu().numpy() + + dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()] + + elements = np.empty(xyz.shape[0], dtype=dtype_full) + attributes = np.concatenate((xyz, normals, f_dc, f_rest, feats3D, opacities, scale, rotation), axis=1) + elements[:] = list(map(tuple, attributes)) + el = PlyElement.describe(elements, 'vertex') + plydata = PlyData([el]) + if path is not None: + plydata.write(path) + return plydata + + def save_splat(self, ply_path, splat_path): + plydata = self.save_ply(ply_path) + vert = plydata["vertex"] + sorted_indices = np.argsort( + -np.exp(vert["scale_0"] + vert["scale_1"] + vert["scale_2"]) + / (1 + np.exp(-vert["opacity"])) + ) + buffer = BytesIO() + for idx in sorted_indices: + v = plydata["vertex"][idx] + position = np.array([v["x"], v["y"], v["z"]], dtype=np.float32) + scales = np.exp( + np.array( + [v["scale_0"], v["scale_1"], v["scale_2"]], + dtype=np.float32, + ) + ) + rot = np.array( + [v["rot_0"], v["rot_1"], v["rot_2"], v["rot_3"]], + dtype=np.float32, + ) + SH_C0 = 0.28209479177387814 + color = np.array( + [ + 0.5 + SH_C0 * v["f_dc_0"], + 0.5 + SH_C0 * v["f_dc_1"], + 0.5 + SH_C0 * v["f_dc_2"], + 1 / (1 + np.exp(-v["opacity"])), + ] + ) + buffer.write(position.tobytes()) + buffer.write(scales.tobytes()) + buffer.write((color * 255).clip(0, 255).astype(np.uint8).tobytes()) + buffer.write( + ((rot / np.linalg.norm(rot)) * 128 + 128) + .clip(0, 255) + .astype(np.uint8) + .tobytes() + ) + with open(splat_path, "wb") as f: + f.write(buffer.getvalue()) + + def save_semantic_pcd(self, path): + color_dict = { + 0: np.array([128, 64, 128]), # Road + 1: np.array([244, 35, 232]), # Sidewalk + 2: np.array([70, 70, 70]), # Building + 3: np.array([102, 102, 156]), # Wall + 4: np.array([190, 153, 153]), # Fence + 5: np.array([153, 153, 153]), # Pole + 6: np.array([250, 170, 30]), # Traffic Light + 7: np.array([220, 220, 0]), # Traffic Sign + 8: np.array([107, 142, 35]), # Vegetation + 9: np.array([152, 251, 152]), # Terrain + 10: np.array([0, 0, 0]), # Black (trainId 10) + 11: np.array([70, 130, 180]), # Sky + 12: np.array([220, 20, 60]), # Person + 13: np.array([255, 0, 0]), # Rider + 14: np.array([0, 0, 142]), # Car + 15: np.array([0, 0, 70]), # Truck + 16: np.array([0, 60, 100]), # Bus + 17: np.array([0, 80, 100]), # Train + 18: np.array([0, 0, 230]), # Motorcycle + 19: np.array([119, 11, 32]) # Bicycle + } + semantic_idx = torch.argmax(self.get_full_3D_features, dim=-1, keepdim=True) + opacities = self.get_full_opacity[:, 0] + mask = ((semantic_idx != 10)[:, 0]) & ((semantic_idx != 8)[:, 0]) & (opacities > 0.2) + + semantic_idx = semantic_idx[mask] + semantic_rgb = torch.zeros_like(semantic_idx).repeat(1, 3) + for idx in range(20): + rgb = torch.from_numpy(color_dict[idx]).to(semantic_rgb.device)[None, :] + semantic_rgb[(semantic_idx == idx)[:, 0], :] = rgb + semantic_rgb = semantic_rgb.float() / 255.0 + pcd_xyz = self.get_full_xyz[mask] + smt_pcd = o3d.geometry.PointCloud() + smt_pcd.points = o3d.utility.Vector3dVector(pcd_xyz.detach().cpu().numpy()) + smt_pcd.colors = o3d.utility.Vector3dVector(semantic_rgb.detach().cpu().numpy()) + o3d.io.write_point_cloud(path, smt_pcd) + + def save_vis_ply(self, path): + mkdir_p(os.path.dirname(path)) + xyz = self.get_xyz.detach().cpu().numpy() + if self.ground_model: + xyz = np.concatenate([xyz, self.ground_model.get_xyz.detach().cpu().numpy()]) + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(xyz) + colors = SH2RGB(self._features_dc[:, 0, :].detach().cpu().numpy()).clip(0, 1) + if self.ground_model: + ground_colors = SH2RGB(self.ground_model._features_dc[:, 0, :].detach().cpu().numpy()).clip(0, 1) + colors = np.concatenate([colors, ground_colors]) + pcd.colors = o3d.utility.Vector3dVector(colors) + o3d.io.write_point_cloud(path, pcd) + + def reset_opacity(self): + opacities_new = inverse_sigmoid(torch.min(self.get_opacity, torch.ones_like(self.get_opacity)*0.01)) + optimizable_tensors = self.replace_tensor_to_optimizer(opacities_new, "opacity") + self._opacity = optimizable_tensors["opacity"] + + def load_ply(self, path): + plydata = PlyData.read(path) + + xyz = np.stack((np.asarray(plydata.elements[0]["x"]), + np.asarray(plydata.elements[0]["y"]), + np.asarray(plydata.elements[0]["z"])), axis=1) + opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis] + + features_dc = np.zeros((xyz.shape[0], 3, 1)) + features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"]) + features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"]) + features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"]) + + extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")] + assert len(extra_f_names)==3*(self.max_sh_degree + 1) ** 2 - 3 + features_extra = np.zeros((xyz.shape[0], len(extra_f_names))) + for idx, attr_name in enumerate(extra_f_names): + features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name]) + # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC) + features_extra = features_extra.reshape((features_extra.shape[0], 3, (self.max_sh_degree + 1) ** 2 - 1)) + + scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")] + scales = np.zeros((xyz.shape[0], len(scale_names))) + for idx, attr_name in enumerate(scale_names): + scales[:, idx] = np.asarray(plydata.elements[0][attr_name]) + + rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")] + rots = np.zeros((xyz.shape[0], len(rot_names))) + for idx, attr_name in enumerate(rot_names): + rots[:, idx] = np.asarray(plydata.elements[0][attr_name]) + + self._xyz = nn.Parameter(torch.tensor(xyz, dtype=torch.float, device="cuda").requires_grad_(True)) + self._features_dc = nn.Parameter(torch.tensor(features_dc, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True)) + self._features_rest = nn.Parameter(torch.tensor(features_extra, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True)) + self._opacity = nn.Parameter(torch.tensor(opacities, dtype=torch.float, device="cuda").requires_grad_(True)) + self._scaling = nn.Parameter(torch.tensor(scales, dtype=torch.float, device="cuda").requires_grad_(True)) + self._rotation = nn.Parameter(torch.tensor(rots, dtype=torch.float, device="cuda").requires_grad_(True)) + + self.active_sh_degree = self.max_sh_degree + + def replace_tensor_to_optimizer(self, tensor, name): + optimizable_tensors = {} + for group in self.optimizer.param_groups: + if group["name"] == name: + stored_state = self.optimizer.state.get(group['params'][0], None) + if stored_state is not None: + stored_state["exp_avg"] = torch.zeros_like(tensor) + stored_state["exp_avg_sq"] = torch.zeros_like(tensor) + del self.optimizer.state[group['params'][0]] + group["params"][0] = nn.Parameter(tensor.requires_grad_(True)) + self.optimizer.state[group['params'][0]] = stored_state + optimizable_tensors[group["name"]] = group["params"][0] + else: + group["params"][0] = nn.Parameter(tensor.requires_grad_(True)) + optimizable_tensors[group["name"]] = group["params"][0] + return optimizable_tensors + + def _prune_optimizer(self, mask): + optimizable_tensors = {} + for group in self.optimizer.param_groups: + if group['name'] == 'appearance_model': + continue + stored_state = self.optimizer.state.get(group['params'][0], None) + if stored_state is not None: + stored_state["exp_avg"] = stored_state["exp_avg"][mask] + stored_state["exp_avg_sq"] = stored_state["exp_avg_sq"][mask] + + del self.optimizer.state[group['params'][0]] + group["params"][0] = nn.Parameter((group["params"][0][mask].requires_grad_(True))) + self.optimizer.state[group['params'][0]] = stored_state + + optimizable_tensors[group["name"]] = group["params"][0] + else: + group["params"][0] = nn.Parameter(group["params"][0][mask].requires_grad_(True)) + optimizable_tensors[group["name"]] = group["params"][0] + return optimizable_tensors + + def prune_points(self, mask): + valid_points_mask = ~mask + optimizable_tensors = self._prune_optimizer(valid_points_mask) + + self._xyz = optimizable_tensors["xyz"] + self._features_dc = optimizable_tensors["f_dc"] + self._features_rest = optimizable_tensors["f_rest"] + if self.feat_mutable: + self._feats3D = optimizable_tensors["feats3D"] + else: + self._feats3D = self._feats3D[1, :].repeat((self._xyz.shape[0], 1)) + self._opacity = optimizable_tensors["opacity"] + self._scaling = optimizable_tensors["scaling"] + self._rotation = optimizable_tensors["rotation"] + + self.xyz_gradient_accum = self.xyz_gradient_accum[valid_points_mask] + + self.denom = self.denom[valid_points_mask] + self.max_radii2D = self.max_radii2D[valid_points_mask] + + def cat_tensors_to_optimizer(self, tensors_dict): + optimizable_tensors = {} + for group in self.optimizer.param_groups: + if group['name'] not in tensors_dict: + continue + assert len(group["params"]) == 1 + extension_tensor = tensors_dict[group["name"]] + stored_state = self.optimizer.state.get(group["params"][0], None) + if stored_state is not None: + + stored_state["exp_avg"] = torch.cat((stored_state["exp_avg"], torch.zeros_like(extension_tensor)), dim=0) + stored_state["exp_avg_sq"] = torch.cat((stored_state["exp_avg_sq"], torch.zeros_like(extension_tensor)), dim=0) + + del self.optimizer.state[group["params"][0]] + group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True)) + self.optimizer.state[group["params"][0]] = stored_state + + optimizable_tensors[group["name"]] = group["params"][0] + else: + group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True)) + optimizable_tensors[group["name"]] = group["params"][0] + + return optimizable_tensors + + def densification_postfix(self, new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacities, new_scaling, new_rotation): + d = {"xyz": new_xyz, + "f_dc": new_features_dc, + "f_rest": new_features_rest, + "feats3D": new_feats3D, + "opacity": new_opacities, + "scaling" : new_scaling, + "rotation" : new_rotation} + + optimizable_tensors = self.cat_tensors_to_optimizer(d) + self._xyz = optimizable_tensors["xyz"] + self._features_dc = optimizable_tensors["f_dc"] + if self.feat_mutable: + self._feats3D = optimizable_tensors["feats3D"] + else: + self._feats3D = self._feats3D[1, :].repeat((self._xyz.shape[0], 1)) + self._features_rest = optimizable_tensors["f_rest"] + self._opacity = optimizable_tensors["opacity"] + self._scaling = optimizable_tensors["scaling"] + self._rotation = optimizable_tensors["rotation"] + + self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda") + self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda") + self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda") + + def densify_and_split(self, grads, grad_threshold, scene_extent, N=2): + n_init_points = self.get_xyz.shape[0] + # Extract points that satisfy the gradient condition + padded_grad = torch.zeros((n_init_points), device="cuda") + padded_grad[:grads.shape[0]] = grads.squeeze() + selected_pts_mask = torch.where(padded_grad >= grad_threshold, True, False) + selected_pts_mask = torch.logical_and(selected_pts_mask, + torch.max(self.get_scaling, dim=1).values > self.percent_dense*scene_extent) + + stds = self.get_scaling[selected_pts_mask].repeat(N,1) + means =torch.zeros((stds.size(0), 3),device="cuda") + samples = torch.normal(mean=means, std=stds) + rots = build_rotation(self._rotation[selected_pts_mask]).repeat(N,1,1) + new_xyz = torch.bmm(rots, samples.unsqueeze(-1)).squeeze(-1) + self.get_xyz[selected_pts_mask].repeat(N, 1) + new_scaling = self.scaling_inverse_activation(self.get_scaling[selected_pts_mask].repeat(N,1) / (0.8*N)) + new_rotation = self._rotation[selected_pts_mask].repeat(N,1) + new_features_dc = self._features_dc[selected_pts_mask].repeat(N,1,1) + new_features_rest = self._features_rest[selected_pts_mask].repeat(N,1,1) + new_feats3D = self._feats3D[selected_pts_mask].repeat(N,1) + new_opacity = self._opacity[selected_pts_mask].repeat(N,1) + + self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacity, new_scaling, new_rotation) + + prune_filter = torch.cat((selected_pts_mask, torch.zeros(N * selected_pts_mask.sum(), device="cuda", dtype=bool))) + self.prune_points(prune_filter) + + def densify_and_clone(self, grads, grad_threshold, scene_extent): + # Extract points that satisfy the gradient condition + selected_pts_mask = torch.where(torch.norm(grads, dim=-1) >= grad_threshold, True, False) + selected_pts_mask = torch.logical_and(selected_pts_mask, + torch.max(self.get_scaling, dim=1).values <= self.percent_dense*scene_extent) + + new_xyz = self._xyz[selected_pts_mask] + new_features_dc = self._features_dc[selected_pts_mask] + new_features_rest = self._features_rest[selected_pts_mask] + new_feats3D = self._feats3D[selected_pts_mask] + new_opacities = self._opacity[selected_pts_mask] + new_scaling = self._scaling[selected_pts_mask] + new_rotation = self._rotation[selected_pts_mask] + + self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacities, new_scaling, new_rotation) + + def densify_and_prune(self, max_grad, min_opacity, extent, max_screen_size, cam_pos=None): + grads = self.xyz_gradient_accum / self.denom + grads[grads.isnan()] = 0.0 + + self.densify_and_clone(grads, max_grad, extent) + self.densify_and_split(grads, max_grad, extent) + + prune_mask = (self.get_opacity < min_opacity).squeeze() + if max_screen_size: + big_points_vs = self.max_radii2D > max_screen_size + if cam_pos is not None: + # points_cam_dist = torch.abs(self.get_xyz[:, None, :] - cam_pos[None, ...]) + # points_cam_nearest_idx = torch.argmin(torch.norm(points_cam_dist, dim=-1), dim=1) + # points_cam_dist = points_cam_dist[torch.arange(points_cam_dist.shape[0]), points_cam_nearest_idx, :] + # near_mask1 = (points_cam_dist[:, 1] < 5) & (points_cam_dist[:, 0] < 10) & (points_cam_dist[:, 2] < 10) + # big_points_ws1 = near_mask1 & (self.get_scaling.max(dim=1).values > 1.0) + # near_mask2 = (points_cam_dist[:, 1] < 10) & (points_cam_dist[:, 0] < 20) & (points_cam_dist[:, 2] < 20) + # big_points_ws2 = near_mask2 & (self.get_scaling.max(dim=1).values > 3.0) + # big_points_ws = (self.get_scaling.max(dim=1).values > 10.0) | big_points_ws1 | big_points_ws2 + big_points_ws = self.get_scaling.max(dim=1).values > 10 + prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws) + else: + big_points_ws = self.get_scaling.max(dim=1).values > 5 + prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws) + self.prune_points(prune_mask) + + torch.cuda.empty_cache() + + def add_densification_stats_grad(self, tensor_grad, update_filter): + self.xyz_gradient_accum[update_filter] += torch.norm(tensor_grad[update_filter,:2], dim=-1, keepdim=True) + self.denom[update_filter] += 1 + \ No newline at end of file diff --git a/code/scene/ground_model.py b/code/scene/ground_model.py new file mode 100644 index 0000000000000000000000000000000000000000..fc0ce19ac94ffb4a45e6e57822ac6c0331d859d5 --- /dev/null +++ b/code/scene/ground_model.py @@ -0,0 +1,360 @@ +import torch +import numpy as np +from utils.general_utils import inverse_sigmoid, get_expon_lr_func, build_rotation +from torch import nn +import os +from utils.system_utils import mkdir_p +from plyfile import PlyData, PlyElement +from utils.sh_utils import RGB2SH, SH2RGB +from simple_knn._C import distCUDA2 +from utils.graphics_utils import BasicPointCloud +from utils.general_utils import strip_symmetric, build_scaling_rotation +import open3d as o3d +import math +from utils.graphics_utils import BasicPointCloud +from utils.sh_utils import RGB2SH + +class GroundModel: + + def setup_functions(self): + def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation): + L = build_scaling_rotation(scaling_modifier * scaling, rotation) + actual_covariance = L @ L.transpose(1, 2) + symm = strip_symmetric(actual_covariance) + return symm + + self.scaling_activation = torch.exp + self.scaling_inverse_activation = torch.log + + self.covariance_activation = build_covariance_from_scaling_rotation + + self.opacity_activation = torch.sigmoid + self.inverse_opacity_activation = torch.logit + + self.rotation_activation = torch.nn.functional.normalize + + + def __init__(self, sh_degree: int, ground_pcd: BasicPointCloud=None, model_args=None, finetune=False): + assert not ((ground_pcd is None) and (model_args is None)), "Need at least one way of initialization" + self.active_sh_degree = 0 + self.max_sh_degree = sh_degree + + self.scale = 0.1 + + if ground_pcd is not None: + self._xyz = nn.Parameter(torch.from_numpy(ground_pcd.points).float().cuda()) + fused_color = RGB2SH(torch.tensor(np.asarray(ground_pcd.colors)).float().cuda()) + features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda() + features[:, :3, 0 ] = fused_color + features[:, 3:, 1:] = 0.0 + self._features_dc = nn.Parameter(features[:,:,0:1].transpose(1, 2).contiguous().requires_grad_(True)) + self._features_rest = nn.Parameter(features[:,:,1:].transpose(1, 2).contiguous().requires_grad_(True)) + + self._feats3D = torch.zeros((self._xyz.shape[0], 20)).cuda() + self._feats3D[:, 1] = 1 + self._feats3D = nn.Parameter(self._feats3D) + self._rotation = torch.zeros((self._xyz.shape[0], 4)).cuda() + self._rotation[:, 0] = 1 + self._opacity = inverse_sigmoid(torch.ones((self._xyz.shape[0], 1)).cuda() * 0.99) + self._scaling = nn.Parameter(torch.ones((self._xyz.shape[0], 2)).float().cuda() * math.log(self.scale)) + + self.max_radii2D = torch.zeros((self._xyz.shape[0]), device="cuda") + self.percent_dense = 0.01 + self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda") + self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda") + else: + self.restore(model_args) + + if finetune: + self.param_groups = [ + {'params': [self._features_dc], 'lr': 2.5e-3, "name": "f_dc"}, + {'params': [self._features_rest], 'lr': 2.5e-3 / 20.0, "name": "f_rest"}, + {'params': [self._feats3D], 'lr': 1e-3, "name": "feats3D"}, + ] + else: + self.param_groups = [ + {'params': [self._xyz], 'lr': 1.6e-4, "name": "xyz"}, + {'params': [self._features_dc], 'lr': 2.5e-3, "name": "f_dc"}, + {'params': [self._features_rest], 'lr': 2.5e-3 / 20.0, "name": "f_rest"}, + {'params': [self._feats3D], 'lr': 1e-2, "name": "feats3D"}, + {'params': [self._opacity], 'lr': 0.05, "name": "opacity"}, + {'params': [self._scaling], 'lr': 1e-3, "name": "scaling"}, + ] + self.optimizer = torch.optim.Adam(self.param_groups, lr=0.0, eps=1e-15) + self.setup_functions() + + def capture(self): + return ( + self.active_sh_degree, + self._xyz, + # self._y, + # self._z, + self._features_dc, + self._features_rest, + self._feats3D, + self._scaling, + self._rotation, + self._opacity, + ) + + def restore(self, model_args): + (self.active_sh_degree, + self._xyz, + # self._y, + # self._z, + self._features_dc, + self._features_rest, + self._feats3D, + self._scaling, + self._rotation, + self._opacity) = model_args + + @property + def get_scaling(self): + scale_y = torch.ones_like(self._xyz[:, 0]) * math.log(0.001) + scaling = torch.stack((self._scaling[:, 0], scale_y, self._scaling[:, 1]), dim=1).cuda() + # scaling = torch.stack((self._scaling, scale_y, self._scaling), dim=1).cuda() + return self.scaling_activation(scaling) + + @property + def get_rotation(self): + return self.rotation_activation(self._rotation) + + @property + def get_xyz(self): + return self._xyz + + @property + def get_features(self): + features_dc = self._features_dc + features_rest = self._features_rest + return torch.cat((features_dc, features_rest), dim=1) + + @property + def get_3D_features(self): + return torch.softmax(self._feats3D, dim=-1) + + @property + def get_opacity(self): + return self.opacity_activation(self._opacity) + + def get_covariance(self, scaling_modifier = 1): + return self.covariance_activation(self.get_scaling, scaling_modifier, self._rotation) + + def oneupSHdegree(self): + if self.active_sh_degree < self.max_sh_degree: + self.active_sh_degree += 1 + + def construct_list_of_attributes(self): + l = ['x', 'y', 'z', 'nx', 'ny', 'nz'] + # All channels except the 3 DC + for i in range(self._features_dc.shape[1]*self._features_dc.shape[2]): + l.append('f_dc_{}'.format(i)) + for i in range(self._features_rest.shape[1]*self._features_rest.shape[2]): + l.append('f_rest_{}'.format(i)) + for i in range(self._feats3D.shape[1]): + l.append('semantic_{}'.format(i)) + l.append('opacity') + for i in range(self._scaling.shape[1]): + l.append('scale_{}'.format(i)) + for i in range(self._rotation.shape[1]): + l.append('rot_{}'.format(i)) + return l + + def save_ply(self, path): + mkdir_p(os.path.dirname(path)) + + xyz = self.get_xyz.detach().cpu().numpy() + normals = np.zeros_like(xyz) + f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy() + f_rest = self._features_rest.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy() + feats3D = self._feats3D.detach().cpu().numpy() + opacities = self._opacity.detach().cpu().numpy() + scale = self._scaling.detach().cpu().numpy() + rotation = self._rotation.detach().cpu().numpy() + + dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()] + + elements = np.empty(xyz.shape[0], dtype=dtype_full) + attributes = np.concatenate((xyz, normals, f_dc, f_rest, feats3D, opacities, scale, rotation), axis=1) + elements[:] = list(map(tuple, attributes)) + el = PlyElement.describe(elements, 'vertex') + PlyData([el]).write(path) + + def save_vis_ply(self, path): + mkdir_p(os.path.dirname(path)) + xyz = self.get_xyz.detach().cpu().numpy() + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(xyz) + colors = SH2RGB(self._features_dc[:, 0, :].detach().cpu().numpy()).clip(0, 1) + pcd.colors = o3d.utility.Vector3dVector(colors) + o3d.io.write_point_cloud(path, pcd) + + def reset_opacity(self): + opacities_new = inverse_sigmoid(torch.min(self.get_opacity, torch.ones_like(self.get_opacity)*0.01)) + optimizable_tensors = self.replace_tensor_to_optimizer(opacities_new, "opacity") + self._opacity = optimizable_tensors["opacity"] + + def replace_tensor_to_optimizer(self, tensor, name): + optimizable_tensors = {} + for group in self.optimizer.param_groups: + if group["name"] == name: + stored_state = self.optimizer.state.get(group['params'][0], None) + if stored_state is not None: + stored_state["exp_avg"] = torch.zeros_like(tensor) + stored_state["exp_avg_sq"] = torch.zeros_like(tensor) + del self.optimizer.state[group['params'][0]] + group["params"][0] = nn.Parameter(tensor.requires_grad_(True)) + self.optimizer.state[group['params'][0]] = stored_state + optimizable_tensors[group["name"]] = group["params"][0] + else: + group["params"][0] = nn.Parameter(tensor.requires_grad_(True)) + optimizable_tensors[group["name"]] = group["params"][0] + return optimizable_tensors + + def _prune_optimizer(self, mask): + optimizable_tensors = {} + for group in self.optimizer.param_groups: + if group['name'] == 'appearance_model': + continue + stored_state = self.optimizer.state.get(group['params'][0], None) + if stored_state is not None: + stored_state["exp_avg"] = stored_state["exp_avg"][mask] + stored_state["exp_avg_sq"] = stored_state["exp_avg_sq"][mask] + + del self.optimizer.state[group['params'][0]] + group["params"][0] = nn.Parameter((group["params"][0][mask].requires_grad_(True))) + self.optimizer.state[group['params'][0]] = stored_state + + optimizable_tensors[group["name"]] = group["params"][0] + else: + group["params"][0] = nn.Parameter(group["params"][0][mask].requires_grad_(True)) + optimizable_tensors[group["name"]] = group["params"][0] + return optimizable_tensors + + def prune_points(self, mask): + valid_points_mask = ~mask + optimizable_tensors = self._prune_optimizer(valid_points_mask) + + self._xyz = optimizable_tensors["xyz"] + self._features_dc = optimizable_tensors["f_dc"] + self._features_rest = optimizable_tensors["f_rest"] + self._feats3D = optimizable_tensors["feats3D"] + self._opacity = optimizable_tensors["opacity"] + self._scaling = optimizable_tensors["scaling"] + self._rotation = self._rotation[0, :].repeat((self._xyz.shape[0], 1)) + + self.xyz_gradient_accum = self.xyz_gradient_accum[valid_points_mask] + + self.denom = self.denom[valid_points_mask] + self.max_radii2D = self.max_radii2D[valid_points_mask] + + def cat_tensors_to_optimizer(self, tensors_dict): + optimizable_tensors = {} + for group in self.optimizer.param_groups: + if group['name'] not in tensors_dict: + continue + assert len(group["params"]) == 1 + extension_tensor = tensors_dict[group["name"]] + stored_state = self.optimizer.state.get(group["params"][0], None) + if stored_state is not None: + stored_state["exp_avg"] = torch.cat((stored_state["exp_avg"], torch.zeros_like(extension_tensor)), dim=0) + stored_state["exp_avg_sq"] = torch.cat((stored_state["exp_avg_sq"], torch.zeros_like(extension_tensor)), dim=0) + + del self.optimizer.state[group["params"][0]] + group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True)) + self.optimizer.state[group["params"][0]] = stored_state + + optimizable_tensors[group["name"]] = group["params"][0] + else: + group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True)) + optimizable_tensors[group["name"]] = group["params"][0] + + return optimizable_tensors + + def densification_postfix(self, new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacities, new_scaling, new_rotation): + d = {"xyz": new_xyz, + "f_dc": new_features_dc, + "f_rest": new_features_rest, + "feats3D": new_feats3D, + "opacity": new_opacities, + "scaling" : new_scaling} + + optimizable_tensors = self.cat_tensors_to_optimizer(d) + self._xyz = optimizable_tensors["xyz"] + self._features_dc = optimizable_tensors["f_dc"] + self._feats3D = optimizable_tensors["feats3D"] + self._features_rest = optimizable_tensors["f_rest"] + self._opacity = optimizable_tensors["opacity"] + self._scaling = optimizable_tensors["scaling"] + self._rotation = self._rotation[0, :].repeat((self._xyz.shape[0], 1)) + + self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda") + self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda") + self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda") + + def densify_and_split(self, grads, grad_threshold, scene_extent, N=2): + n_init_points = self.get_xyz.shape[0] + # Extract points that satisfy the gradient condition + padded_grad = torch.zeros((n_init_points), device="cuda") + padded_grad[:grads.shape[0]] = grads.squeeze() + selected_pts_mask = torch.where(padded_grad >= grad_threshold, True, False) + selected_pts_mask = torch.logical_and(selected_pts_mask, + torch.max(self.get_scaling, dim=1).values > self.percent_dense*scene_extent) + + stds = self.get_scaling[selected_pts_mask].repeat(N,1) + means =torch.zeros((stds.size(0), 3),device="cuda") + samples = torch.normal(mean=means, std=stds) + rots = build_rotation(self._rotation[selected_pts_mask]).repeat(N,1,1) + new_xyz = torch.bmm(rots, samples.unsqueeze(-1)).squeeze(-1) + self.get_xyz[selected_pts_mask].repeat(N, 1) + new_scaling = self.scaling_inverse_activation(self.get_scaling[selected_pts_mask].repeat(N,1) / (0.8*N))[:, [0,2]] + new_rotation = self._rotation[selected_pts_mask].repeat(N,1) + new_features_dc = self._features_dc[selected_pts_mask].repeat(N,1,1) + new_features_rest = self._features_rest[selected_pts_mask].repeat(N,1,1) + new_feats3D = self._feats3D[selected_pts_mask].repeat(N,1) + new_opacity = self._opacity[selected_pts_mask].repeat(N,1) + + self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacity, new_scaling, new_rotation) + + prune_filter = torch.cat((selected_pts_mask, torch.zeros(N * selected_pts_mask.sum(), device="cuda", dtype=bool))) + self.prune_points(prune_filter) + + def densify_and_clone(self, grads, grad_threshold, scene_extent): + # Extract points that satisfy the gradient condition + selected_pts_mask = torch.where(torch.norm(grads, dim=-1) >= grad_threshold, True, False) + selected_pts_mask = torch.logical_and(selected_pts_mask, + torch.max(self.get_scaling, dim=1).values <= self.percent_dense*scene_extent) + + new_xyz = self._xyz[selected_pts_mask] + new_features_dc = self._features_dc[selected_pts_mask] + new_features_rest = self._features_rest[selected_pts_mask] + new_feats3D = self._feats3D[selected_pts_mask] + new_opacities = self._opacity[selected_pts_mask] + new_scaling = self._scaling[selected_pts_mask] + new_rotation = self._rotation[selected_pts_mask] + + self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacities, new_scaling, new_rotation) + + def densify_and_prune(self, max_grad, min_opacity, extent, max_screen_size): + grads = self.xyz_gradient_accum / self.denom + grads[grads.isnan()] = 0.0 + + self.densify_and_clone(grads, max_grad, extent) + self.densify_and_split(grads, max_grad, extent) + + prune_mask = (self.get_opacity < min_opacity).squeeze() + if max_screen_size: + big_points_vs = self.max_radii2D > max_screen_size + big_points_ws = self.get_scaling.max(dim=1).values > 1.0 + prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws) + self.prune_points(prune_mask) + + torch.cuda.empty_cache() + + def add_densification_stats(self, viewspace_point_tensor, update_filter): + self.xyz_gradient_accum[update_filter] += torch.norm(viewspace_point_tensor.grad[update_filter,:2], dim=-1, keepdim=True) + self.denom[update_filter] += 1 + + def add_densification_stats_grad(self, tensor_grad, update_filter): + self.xyz_gradient_accum[update_filter] += torch.norm(tensor_grad[update_filter,:2], dim=-1, keepdim=True) + self.denom[update_filter] += 1 \ No newline at end of file diff --git a/code/scene/obj_model.py b/code/scene/obj_model.py new file mode 100644 index 0000000000000000000000000000000000000000..90ed72589e3c2b54d45c1135a424c0b40e608012 --- /dev/null +++ b/code/scene/obj_model.py @@ -0,0 +1,567 @@ +import torch +import numpy as np +from utils.general_utils import inverse_sigmoid, get_expon_lr_func, build_rotation +from torch import nn +import os +from utils.system_utils import mkdir_p +from plyfile import PlyData, PlyElement +from utils.sh_utils import RGB2SH, SH2RGB +from simple_knn._C import distCUDA2 +from utils.graphics_utils import BasicPointCloud +from utils.general_utils import strip_symmetric, build_scaling_rotation +import open3d as o3d +import tinycudann as tcnn +from io import BytesIO + + +class ObjModel: + + def setup_functions(self): + def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation): + L = build_scaling_rotation(scaling_modifier * scaling, rotation) + actual_covariance = L @ L.transpose(1, 2) + symm = strip_symmetric(actual_covariance) + return symm + + self.scaling_activation = torch.exp + self.scaling_inverse_activation = torch.log + + self.covariance_activation = build_covariance_from_scaling_rotation + + self.opacity_activation = torch.sigmoid + self.inverse_opacity_activation = torch.logit + + self.rotation_activation = torch.nn.functional.normalize + + + def __init__(self, sh_degree : int, feat_mutable=True, affine=False): + self.active_sh_degree = 0 + self.max_sh_degree = sh_degree + self._xyz = torch.empty(0) + self._features_dc = torch.empty(0) + self._features_rest = torch.empty(0) + self._feats3D = torch.empty(0) + self._scaling = torch.empty(0) + self._rotation = torch.empty(0) + self._opacity = torch.empty(0) + self.max_radii2D = torch.empty(0) + self.xyz_gradient_accum = torch.empty(0) + self.denom = torch.empty(0) + self.optimizer = None + self.percent_dense = 0 + self.spatial_lr_scale = 0 + self.feat_mutable = feat_mutable + self.setup_functions() + + self.pos_enc = tcnn.Encoding( + n_input_dims=3, + encoding_config={"otype": "Frequency", "n_frequencies": 2}, + ) + self.dir_enc = tcnn.Encoding( + n_input_dims=3, + encoding_config={ + "otype": "SphericalHarmonics", + "degree": 3, + }, + ) + + self.affine = affine + if affine: + self.appearance_model = tcnn.Network( + n_input_dims=self.pos_enc.n_output_dims + self.dir_enc.n_output_dims, + n_output_dims=12, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "None", + "n_neurons": 32, + "n_hidden_layers": 2, + } + ) + else: + self.appearance_model = None + + def capture(self): + return ( + self.active_sh_degree, + self._xyz, + self._features_dc, + self._features_rest, + self._feats3D, + self._scaling, + self._rotation, + self._opacity, + self.spatial_lr_scale, + ) + + def restore(self, model_args, training_args): + (self.active_sh_degree, + self._xyz, + self._features_dc, + self._features_rest, + self._feats3D, + self._scaling, + self._rotation, + self._opacity, + self.spatial_lr_scale, + ) = model_args + if training_args is not None: + self.training_setup(training_args) + + @property + def get_scaling(self): + return self.scaling_activation(self._scaling) + + @property + def get_rotation(self): + return self.rotation_activation(self._rotation) + + @property + def get_xyz(self): + return self._xyz + + @property + def get_features(self): + features_dc = self._features_dc + features_rest = self._features_rest + return torch.cat((features_dc, features_rest), dim=1) + + @property + def get_3D_features(self): + return torch.softmax(self._feats3D, dim=-1) + + @property + def get_opacity(self): + return self.opacity_activation(self._opacity) + + # def get_covariance(self, scaling_modifier = 1): + # return self.covariance_activation(self.get_scaling, scaling_modifier, self._rotation) + + def oneupSHdegree(self): + if self.active_sh_degree < self.max_sh_degree: + self.active_sh_degree += 1 + + def create_from_pcd(self, pcd : BasicPointCloud, spatial_lr_scale : float): + # self.spatial_lr_scale = 1 + self.spatial_lr_scale = spatial_lr_scale + fused_point_cloud = torch.tensor(np.asarray(pcd.points)).float().cuda() + fused_color = RGB2SH(torch.tensor(np.asarray(pcd.colors)).float().cuda()) + features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda() + features[:, :3, 0 ] = fused_color + features[:, 3:, 1:] = 0.0 + + if self.feat_mutable: + feats3D = torch.rand(fused_color.shape[0], 20).float().cuda() + self._feats3D = nn.Parameter(feats3D.requires_grad_(True)) + else: + feats3D = torch.zeros(fused_color.shape[0], 20).float().cuda() + feats3D[:, 13] = 1 + self._feats3D = feats3D + + print("Number of points at initialization : ", fused_point_cloud.shape[0]) + + dist2 = torch.clamp_min(distCUDA2(torch.from_numpy(np.asarray(pcd.points)).float().cuda()), 0.0000001) + scales = torch.log(torch.sqrt(dist2))[...,None].repeat(1, 3) + rots = torch.zeros((fused_point_cloud.shape[0], 4), device="cuda") + rots[:, 0] = 1 + + opacities = inverse_sigmoid(0.1 * torch.ones((fused_point_cloud.shape[0], 1), dtype=torch.float, device="cuda")) + + self._xyz = nn.Parameter(fused_point_cloud.requires_grad_(True)) + self._features_dc = nn.Parameter(features[:,:,0:1].transpose(1, 2).contiguous().requires_grad_(True)) + self._features_rest = nn.Parameter(features[:,:,1:].transpose(1, 2).contiguous().requires_grad_(True)) + self._scaling = nn.Parameter(scales.requires_grad_(True)) + self._rotation = nn.Parameter(rots.requires_grad_(True)) + self._opacity = nn.Parameter(opacities.requires_grad_(True)) + self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda") + + def training_setup(self, training_args): + self.percent_dense = training_args.percent_dense + self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda") + self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda") + + # self.spatial_lr_scale /= 3 + + l = [ + {'params': [self._xyz], 'lr': training_args.position_lr_init * 0.5, "name": "xyz"}, + {'params': [self._features_dc], 'lr': training_args.feature_lr, "name": "f_dc"}, + {'params': [self._features_rest], 'lr': training_args.feature_lr / 20.0, "name": "f_rest"}, + {'params': [self._opacity], 'lr': training_args.opacity_lr, "name": "opacity"}, + {'params': [self._scaling], 'lr': training_args.scaling_lr * 0.5, "name": "scaling"}, + {'params': [self._rotation], 'lr': training_args.rotation_lr, "name": "rotation"}, + ] + + if self.affine: + l.append({'params': [*self.appearance_model.parameters()], 'lr': 1e-3, "name": "appearance_model"}) + + if self.feat_mutable: + l.append({'params': [self._feats3D], 'lr': 1e-2, "name": "feats3D"}) + + self.optimizer = torch.optim.Adam(l, lr=0.0, eps=1e-15) + self.xyz_scheduler_args = get_expon_lr_func(lr_init=training_args.position_lr_init*self.spatial_lr_scale, + lr_final=training_args.position_lr_final*self.spatial_lr_scale, + lr_delay_mult=training_args.position_lr_delay_mult, + max_steps=training_args.position_lr_max_steps) + + def update_learning_rate(self, iteration): + ''' Learning rate scheduling per step ''' + for param_group in self.optimizer.param_groups: + if param_group["name"] == "xyz": + lr = self.xyz_scheduler_args(iteration) + param_group['lr'] = lr + return lr + + def construct_list_of_attributes(self): + l = ['x', 'y', 'z', 'nx', 'ny', 'nz'] + # All channels except the 3 DC + for i in range(self._features_dc.shape[1]*self._features_dc.shape[2]): + l.append('f_dc_{}'.format(i)) + for i in range(self._features_rest.shape[1]*self._features_rest.shape[2]): + l.append('f_rest_{}'.format(i)) + for i in range(self._feats3D.shape[1]): + l.append('semantic_{}'.format(i)) + l.append('opacity') + for i in range(self._scaling.shape[1]): + l.append('scale_{}'.format(i)) + for i in range(self._rotation.shape[1]): + l.append('rot_{}'.format(i)) + return l + + def save_ply(self, path=None): + mkdir_p(os.path.dirname(path)) + + xyz = self.get_xyz.detach().cpu().numpy() + normals = np.zeros_like(xyz) + f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy() + f_rest = self._features_rest.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy() + feats3D = self._feats3D.detach().cpu().numpy() + opacities = self._opacity.detach().cpu().numpy() + scale = self.scaling_inverse_activation(self.get_scaling).detach().cpu().numpy() + rotation = self._rotation.detach().cpu().numpy() + + dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()] + + elements = np.empty(xyz.shape[0], dtype=dtype_full) + attributes = np.concatenate((xyz, normals, f_dc, f_rest, feats3D, opacities, scale, rotation), axis=1) + elements[:] = list(map(tuple, attributes)) + el = PlyElement.describe(elements, 'vertex') + plydata = PlyData([el]) + if path is not None: + plydata.write(path) + return plydata + + def save_splat(self, ply_path, splat_path): + plydata = self.save_ply(ply_path) + vert = plydata["vertex"] + sorted_indices = np.argsort( + -np.exp(vert["scale_0"] + vert["scale_1"] + vert["scale_2"]) + / (1 + np.exp(-vert["opacity"])) + ) + buffer = BytesIO() + for idx in sorted_indices: + v = plydata["vertex"][idx] + position = np.array([v["x"], v["y"], v["z"]], dtype=np.float32) + scales = np.exp( + np.array( + [v["scale_0"], v["scale_1"], v["scale_2"]], + dtype=np.float32, + ) + ) + rot = np.array( + [v["rot_0"], v["rot_1"], v["rot_2"], v["rot_3"]], + dtype=np.float32, + ) + SH_C0 = 0.28209479177387814 + color = np.array( + [ + 0.5 + SH_C0 * v["f_dc_0"], + 0.5 + SH_C0 * v["f_dc_1"], + 0.5 + SH_C0 * v["f_dc_2"], + 1 / (1 + np.exp(-v["opacity"])), + ] + ) + buffer.write(position.tobytes()) + buffer.write(scales.tobytes()) + buffer.write((color * 255).clip(0, 255).astype(np.uint8).tobytes()) + buffer.write( + ((rot / np.linalg.norm(rot)) * 128 + 128) + .clip(0, 255) + .astype(np.uint8) + .tobytes() + ) + with open(splat_path, "wb") as f: + f.write(buffer.getvalue()) + + def save_semantic_pcd(self, path): + color_dict = { + 0: np.array([128, 64, 128]), # Road + 1: np.array([244, 35, 232]), # Sidewalk + 2: np.array([70, 70, 70]), # Building + 3: np.array([102, 102, 156]), # Wall + 4: np.array([190, 153, 153]), # Fence + 5: np.array([153, 153, 153]), # Pole + 6: np.array([250, 170, 30]), # Traffic Light + 7: np.array([220, 220, 0]), # Traffic Sign + 8: np.array([107, 142, 35]), # Vegetation + 9: np.array([152, 251, 152]), # Terrain + 10: np.array([0, 0, 0]), # Black (trainId 10) + 11: np.array([70, 130, 180]), # Sky + 12: np.array([220, 20, 60]), # Person + 13: np.array([255, 0, 0]), # Rider + 14: np.array([0, 0, 142]), # Car + 15: np.array([0, 0, 70]), # Truck + 16: np.array([0, 60, 100]), # Bus + 17: np.array([0, 80, 100]), # Train + 18: np.array([0, 0, 230]), # Motorcycle + 19: np.array([119, 11, 32]) # Bicycle + } + semantic_idx = torch.argmax(self.get_3D_features, dim=-1, keepdim=True) + opacities = self.get_opacity[:, 0] + mask = ((semantic_idx != 10)[:, 0]) & ((semantic_idx != 8)[:, 0]) & (opacities > 0.2) + + semantic_idx = semantic_idx[mask] + semantic_rgb = torch.zeros_like(semantic_idx).repeat(1, 3) + for idx in range(20): + rgb = torch.from_numpy(color_dict[idx]).to(semantic_rgb.device)[None, :] + semantic_rgb[(semantic_idx == idx)[:, 0], :] = rgb + semantic_rgb = semantic_rgb.float() / 255.0 + pcd_xyz = self.get_xyz[mask] + smt_pcd = o3d.geometry.PointCloud() + smt_pcd.points = o3d.utility.Vector3dVector(pcd_xyz.detach().cpu().numpy()) + smt_pcd.colors = o3d.utility.Vector3dVector(semantic_rgb.detach().cpu().numpy()) + o3d.io.write_point_cloud(path, smt_pcd) + + def save_vis_ply(self, path): + mkdir_p(os.path.dirname(path)) + xyz = self.get_xyz.detach().cpu().numpy() + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(xyz) + colors = SH2RGB(self._features_dc[:, 0, :].detach().cpu().numpy()).clip(0, 1) + pcd.colors = o3d.utility.Vector3dVector(colors) + o3d.io.write_point_cloud(path, pcd) + + def reset_opacity(self): + opacities_new = inverse_sigmoid(torch.min(self.get_opacity, torch.ones_like(self.get_opacity)*0.01)) + optimizable_tensors = self.replace_tensor_to_optimizer(opacities_new, "opacity") + self._opacity = optimizable_tensors["opacity"] + + def load_ply(self, path): + plydata = PlyData.read(path) + + xyz = np.stack((np.asarray(plydata.elements[0]["x"]), + np.asarray(plydata.elements[0]["y"]), + np.asarray(plydata.elements[0]["z"])), axis=1) + opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis] + + features_dc = np.zeros((xyz.shape[0], 3, 1)) + features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"]) + features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"]) + features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"]) + + extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")] + assert len(extra_f_names)==3*(self.max_sh_degree + 1) ** 2 - 3 + features_extra = np.zeros((xyz.shape[0], len(extra_f_names))) + for idx, attr_name in enumerate(extra_f_names): + features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name]) + # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC) + features_extra = features_extra.reshape((features_extra.shape[0], 3, (self.max_sh_degree + 1) ** 2 - 1)) + + scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")] + scales = np.zeros((xyz.shape[0], len(scale_names))) + for idx, attr_name in enumerate(scale_names): + scales[:, idx] = np.asarray(plydata.elements[0][attr_name]) + + rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")] + rots = np.zeros((xyz.shape[0], len(rot_names))) + for idx, attr_name in enumerate(rot_names): + rots[:, idx] = np.asarray(plydata.elements[0][attr_name]) + + self._xyz = nn.Parameter(torch.tensor(xyz, dtype=torch.float, device="cuda").requires_grad_(True)) + self._features_dc = nn.Parameter(torch.tensor(features_dc, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True)) + self._features_rest = nn.Parameter(torch.tensor(features_extra, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True)) + self._opacity = nn.Parameter(torch.tensor(opacities, dtype=torch.float, device="cuda").requires_grad_(True)) + self._scaling = nn.Parameter(torch.tensor(scales, dtype=torch.float, device="cuda").requires_grad_(True)) + self._rotation = nn.Parameter(torch.tensor(rots, dtype=torch.float, device="cuda").requires_grad_(True)) + + self.active_sh_degree = self.max_sh_degree + + def replace_tensor_to_optimizer(self, tensor, name): + optimizable_tensors = {} + for group in self.optimizer.param_groups: + if group["name"] == name: + stored_state = self.optimizer.state.get(group['params'][0], None) + if stored_state is not None: + stored_state["exp_avg"] = torch.zeros_like(tensor) + stored_state["exp_avg_sq"] = torch.zeros_like(tensor) + del self.optimizer.state[group['params'][0]] + group["params"][0] = nn.Parameter(tensor.requires_grad_(True)) + self.optimizer.state[group['params'][0]] = stored_state + optimizable_tensors[group["name"]] = group["params"][0] + else: + group["params"][0] = nn.Parameter(tensor.requires_grad_(True)) + optimizable_tensors[group["name"]] = group["params"][0] + return optimizable_tensors + + def _prune_optimizer(self, mask): + optimizable_tensors = {} + for group in self.optimizer.param_groups: + if group['name'] == 'appearance_model': + continue + stored_state = self.optimizer.state.get(group['params'][0], None) + if stored_state is not None: + stored_state["exp_avg"] = stored_state["exp_avg"][mask] + stored_state["exp_avg_sq"] = stored_state["exp_avg_sq"][mask] + + del self.optimizer.state[group['params'][0]] + group["params"][0] = nn.Parameter((group["params"][0][mask].requires_grad_(True))) + self.optimizer.state[group['params'][0]] = stored_state + + optimizable_tensors[group["name"]] = group["params"][0] + else: + group["params"][0] = nn.Parameter(group["params"][0][mask].requires_grad_(True)) + optimizable_tensors[group["name"]] = group["params"][0] + return optimizable_tensors + + def prune_points(self, mask): + valid_points_mask = ~mask + optimizable_tensors = self._prune_optimizer(valid_points_mask) + + self._xyz = optimizable_tensors["xyz"] + self._features_dc = optimizable_tensors["f_dc"] + self._features_rest = optimizable_tensors["f_rest"] + if self.feat_mutable: + self._feats3D = optimizable_tensors["feats3D"] + else: + self._feats3D = self._feats3D[1, :].repeat((self._xyz.shape[0], 1)) + self._opacity = optimizable_tensors["opacity"] + self._scaling = optimizable_tensors["scaling"] + self._rotation = optimizable_tensors["rotation"] + + self.xyz_gradient_accum = self.xyz_gradient_accum[valid_points_mask] + + self.denom = self.denom[valid_points_mask] + self.max_radii2D = self.max_radii2D[valid_points_mask] + + def cat_tensors_to_optimizer(self, tensors_dict): + optimizable_tensors = {} + for group in self.optimizer.param_groups: + if group['name'] not in tensors_dict: + continue + assert len(group["params"]) == 1 + extension_tensor = tensors_dict[group["name"]] + stored_state = self.optimizer.state.get(group["params"][0], None) + if stored_state is not None: + + stored_state["exp_avg"] = torch.cat((stored_state["exp_avg"], torch.zeros_like(extension_tensor)), dim=0) + stored_state["exp_avg_sq"] = torch.cat((stored_state["exp_avg_sq"], torch.zeros_like(extension_tensor)), dim=0) + + del self.optimizer.state[group["params"][0]] + group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True)) + self.optimizer.state[group["params"][0]] = stored_state + + optimizable_tensors[group["name"]] = group["params"][0] + else: + group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True)) + optimizable_tensors[group["name"]] = group["params"][0] + + return optimizable_tensors + + def densification_postfix(self, new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacities, new_scaling, new_rotation): + d = {"xyz": new_xyz, + "f_dc": new_features_dc, + "f_rest": new_features_rest, + "feats3D": new_feats3D, + "opacity": new_opacities, + "scaling" : new_scaling, + "rotation" : new_rotation} + + optimizable_tensors = self.cat_tensors_to_optimizer(d) + self._xyz = optimizable_tensors["xyz"] + self._features_dc = optimizable_tensors["f_dc"] + if self.feat_mutable: + self._feats3D = optimizable_tensors["feats3D"] + else: + self._feats3D = self._feats3D[1, :].repeat((self._xyz.shape[0], 1)) + self._features_rest = optimizable_tensors["f_rest"] + self._opacity = optimizable_tensors["opacity"] + self._scaling = optimizable_tensors["scaling"] + self._rotation = optimizable_tensors["rotation"] + + self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda") + self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda") + self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda") + + def densify_and_split(self, grads, grad_threshold, scene_extent, N=2): + n_init_points = self.get_xyz.shape[0] + # Extract points that satisfy the gradient condition + padded_grad = torch.zeros((n_init_points), device="cuda") + padded_grad[:grads.shape[0]] = grads.squeeze() + selected_pts_mask = torch.where(padded_grad >= grad_threshold, True, False) + selected_pts_mask = torch.logical_and(selected_pts_mask, + torch.max(self.get_scaling, dim=1).values > self.percent_dense*scene_extent) + + stds = self.get_scaling[selected_pts_mask].repeat(N,1) + means =torch.zeros((stds.size(0), 3),device="cuda") + samples = torch.normal(mean=means, std=stds) + rots = build_rotation(self._rotation[selected_pts_mask]).repeat(N,1,1) + new_xyz = torch.bmm(rots, samples.unsqueeze(-1)).squeeze(-1) + self.get_xyz[selected_pts_mask].repeat(N, 1) + new_scaling = self.scaling_inverse_activation(self.get_scaling[selected_pts_mask].repeat(N,1) / (0.8*N)) + new_rotation = self._rotation[selected_pts_mask].repeat(N,1) + new_features_dc = self._features_dc[selected_pts_mask].repeat(N,1,1) + new_features_rest = self._features_rest[selected_pts_mask].repeat(N,1,1) + new_feats3D = self._feats3D[selected_pts_mask].repeat(N,1) + new_opacity = self._opacity[selected_pts_mask].repeat(N,1) + + self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacity, new_scaling, new_rotation) + + prune_filter = torch.cat((selected_pts_mask, torch.zeros(N * selected_pts_mask.sum(), device="cuda", dtype=bool))) + self.prune_points(prune_filter) + + def densify_and_clone(self, grads, grad_threshold, scene_extent): + # Extract points that satisfy the gradient condition + selected_pts_mask = torch.where(torch.norm(grads, dim=-1) >= grad_threshold, True, False) + selected_pts_mask = torch.logical_and(selected_pts_mask, + torch.max(self.get_scaling, dim=1).values <= self.percent_dense*scene_extent) + + new_xyz = self._xyz[selected_pts_mask] + new_features_dc = self._features_dc[selected_pts_mask] + new_features_rest = self._features_rest[selected_pts_mask] + new_feats3D = self._feats3D[selected_pts_mask] + new_opacities = self._opacity[selected_pts_mask] + new_scaling = self._scaling[selected_pts_mask] + new_rotation = self._rotation[selected_pts_mask] + + self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacities, new_scaling, new_rotation) + + def densify_and_prune(self, max_grad, min_opacity, extent, max_screen_size, cam_pos=None): + grads = self.xyz_gradient_accum / self.denom + grads[grads.isnan()] = 0.0 + + self.densify_and_clone(grads, max_grad, extent) + self.densify_and_split(grads, max_grad, extent) + + prune_mask = (self.get_opacity < min_opacity).squeeze() + if max_screen_size: + big_points_vs = self.max_radii2D > max_screen_size + if cam_pos is not None: + # points_cam_dist = torch.abs(self.get_xyz[:, None, :] - cam_pos[None, ...]) + # points_cam_nearest_idx = torch.argmin(torch.norm(points_cam_dist, dim=-1), dim=1) + # points_cam_dist = points_cam_dist[torch.arange(points_cam_dist.shape[0]), points_cam_nearest_idx, :] + # near_mask1 = (points_cam_dist[:, 1] < 5) & (points_cam_dist[:, 0] < 10) & (points_cam_dist[:, 2] < 10) + # big_points_ws1 = near_mask1 & (self.get_scaling.max(dim=1).values > 1.0) + # near_mask2 = (points_cam_dist[:, 1] < 10) & (points_cam_dist[:, 0] < 20) & (points_cam_dist[:, 2] < 20) + # big_points_ws2 = near_mask2 & (self.get_scaling.max(dim=1).values > 3.0) + # big_points_ws = (self.get_scaling.max(dim=1).values > 10.0) | big_points_ws1 | big_points_ws2 + big_points_ws = self.get_scaling.max(dim=1).values > 10 + prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws) + else: + big_points_ws = self.get_scaling.max(dim=1).values > 5 + prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws) + self.prune_points(prune_mask) + + torch.cuda.empty_cache() + + def add_densification_stats_grad(self, tensor_grad, update_filter): + self.xyz_gradient_accum[update_filter] += torch.norm(tensor_grad[update_filter,:2], dim=-1, keepdim=True) + self.denom[update_filter] += 1 + \ No newline at end of file diff --git a/code/sim/hugsim_env.egg-info/PKG-INFO b/code/sim/hugsim_env.egg-info/PKG-INFO new file mode 100644 index 0000000000000000000000000000000000000000..baa4128d23be57803e61aeba0263da9c3ee26bbc --- /dev/null +++ b/code/sim/hugsim_env.egg-info/PKG-INFO @@ -0,0 +1,4 @@ +Metadata-Version: 2.4 +Name: hugsim-env +Version: 0.0.1 +Requires-Dist: gymnasium diff --git a/code/sim/hugsim_env.egg-info/SOURCES.txt b/code/sim/hugsim_env.egg-info/SOURCES.txt new file mode 100644 index 0000000000000000000000000000000000000000..fbf2ce7ba0fd8150f7fedc57be48fbb25e9e4396 --- /dev/null +++ b/code/sim/hugsim_env.egg-info/SOURCES.txt @@ -0,0 +1,10 @@ +pyproject.toml +setup.py +hugsim_env/__init__.py +hugsim_env.egg-info/PKG-INFO +hugsim_env.egg-info/SOURCES.txt +hugsim_env.egg-info/dependency_links.txt +hugsim_env.egg-info/requires.txt +hugsim_env.egg-info/top_level.txt +hugsim_env/envs/__init__.py +hugsim_env/envs/hug_sim.py \ No newline at end of file diff --git a/code/sim/hugsim_env.egg-info/dependency_links.txt b/code/sim/hugsim_env.egg-info/dependency_links.txt new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/code/sim/hugsim_env.egg-info/dependency_links.txt @@ -0,0 +1 @@ + diff --git a/code/sim/hugsim_env.egg-info/requires.txt b/code/sim/hugsim_env.egg-info/requires.txt new file mode 100644 index 0000000000000000000000000000000000000000..7aec64bcc5afba20d370404331125cba67102c95 --- /dev/null +++ b/code/sim/hugsim_env.egg-info/requires.txt @@ -0,0 +1 @@ +gymnasium diff --git a/code/sim/hugsim_env.egg-info/top_level.txt b/code/sim/hugsim_env.egg-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..d46c32cee37991dc71dae1118c8e5921dda29953 --- /dev/null +++ b/code/sim/hugsim_env.egg-info/top_level.txt @@ -0,0 +1 @@ +hugsim_env diff --git a/code/sim/hugsim_env/__init__.py b/code/sim/hugsim_env/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a614d31ac90f09236e2ae91e0048a1a77b454251 --- /dev/null +++ b/code/sim/hugsim_env/__init__.py @@ -0,0 +1,8 @@ +from gymnasium.envs.registration import register + + +register( + id="hugsim_env/HUGSim-v0", + entry_point="hugsim_env.envs:HUGSimEnv", + max_episode_steps=400, +) \ No newline at end of file diff --git a/code/sim/hugsim_env/__pycache__/__init__.cpython-311.pyc b/code/sim/hugsim_env/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cabe7c879560196ec39d5e6984ca6dd65140d196 Binary files /dev/null and b/code/sim/hugsim_env/__pycache__/__init__.cpython-311.pyc differ diff --git a/code/sim/hugsim_env/envs/__init__.py b/code/sim/hugsim_env/envs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d8f3cbdf0f544dc7f2456d2a0733ae6d31980ff7 --- /dev/null +++ b/code/sim/hugsim_env/envs/__init__.py @@ -0,0 +1 @@ +from hugsim_env.envs.hug_sim import HUGSimEnv \ No newline at end of file diff --git a/code/sim/hugsim_env/envs/__pycache__/__init__.cpython-311.pyc b/code/sim/hugsim_env/envs/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7d3394395b87efacc84084d5b802bf3ca976244e Binary files /dev/null and b/code/sim/hugsim_env/envs/__pycache__/__init__.cpython-311.pyc differ diff --git a/code/sim/hugsim_env/envs/__pycache__/hug_sim.cpython-311.pyc b/code/sim/hugsim_env/envs/__pycache__/hug_sim.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a2b38c4971553dcd5275960da3ff511d84d19fbd Binary files /dev/null and b/code/sim/hugsim_env/envs/__pycache__/hug_sim.cpython-311.pyc differ diff --git a/code/sim/hugsim_env/envs/hug_sim.py b/code/sim/hugsim_env/envs/hug_sim.py new file mode 100644 index 0000000000000000000000000000000000000000..e8c2bd04bf9d4260c52bb13c019d58e70505b7f5 --- /dev/null +++ b/code/sim/hugsim_env/envs/hug_sim.py @@ -0,0 +1,333 @@ +import torch +import numpy as np +from copy import deepcopy +import gymnasium +from gymnasium import spaces +from copy import deepcopy +from sim.utils.sim_utils import create_cam, rt2pose, pose2rt, load_camera_cfg, dense_cam_poses +from scipy.spatial.transform import Rotation as SCR +from sim.utils.score_calculator import create_rectangle, bg_collision_det +import os +import pickle +from sim.utils.plan import planner, UnifiedMap +from omegaconf import OmegaConf +import math +from gaussian_renderer import GaussianModel +from scene.obj_model import ObjModel +from gaussian_renderer import render +import open3d as o3d + + +def fg_collision_det(ego_box, objs): + ego_x, ego_y, _, ego_w, ego_l, ego_h, ego_yaw = ego_box + ego_poly = create_rectangle(ego_x, ego_y, ego_w, ego_l, ego_yaw) + for obs in objs: + obs_x, obs_y, _, obs_w, obs_l, _, obs_yaw = obs + obs_poly = create_rectangle( + obs_x, obs_y, obs_w, obs_l, obs_yaw) + if ego_poly.intersects(obs_poly): + return True + return False + +class HUGSimEnv(gymnasium.Env): + def __init__(self, cfg, output): + super().__init__() + + plan_list = cfg.scenario.plan_list + for control_param in plan_list: + control_param[5] = os.path.join(cfg.base.realcar_path, control_param[5]) + + # read ground infos + with open(os.path.join(cfg.model_path, 'ground_param.pkl'), 'rb') as f: + #numpy.ndarray, float, list + cam_poses, cam_heights, commands = pickle.load(f) + cam_poses, commands = dense_cam_poses(cam_poses, commands) + self.ground_model = (cam_poses, cam_heights, commands) + + if cfg.scenario.load_HD_map: + unified_map = UnifiedMap(cfg.base.HD_map.path, cfg.base.HD_map.version, cfg.scenario.scene_name) + else: + unified_map = None + + self.kinematic = OmegaConf.to_container(cfg.kinematic) + self.kinematic['min_steer'] = -math.radians(cfg.kinematic.min_steer) + self.kinematic['max_steer'] = math.radians(cfg.kinematic.max_steer) + self.kinematic['start_vr']= np.array(cfg.scenario.start_euler) / 180 * np.pi + self.kinematic['start_vab'] = np.array(cfg.scenario.start_ab) + self.kinematic['start_velo'] = cfg.scenario.start_velo + self.kinematic['start_steer'] = cfg.scenario.start_steer + + self.gaussians = GaussianModel(cfg.model.sh_degree, affine=cfg.affine) + + """ + plan_list: a, b, height, yaw, v, model_path, controller, params + Yaw is based on ego car's orientation. 0 means same direction as ego. + Right is positive and left is negative. + """ + + (model_params, iteration) = torch.load(os.path.join(cfg.model_path, "scene.pth"), weights_only=False) + self.gaussians.restore(model_params, None) + + dynamic_gaussians = {} + if len(plan_list) == 0: + self.planner = None + else: + self.planner = planner(plan_list, scene_path=cfg.model_path, unified_map=unified_map, ground=self.ground_model, dt=cfg.kinematic.dt) + for plan_id in self.planner.ckpts.keys(): + dynamic_gaussians[plan_id] = ObjModel(cfg.model.sh_degree, feat_mutable=False) + (model_params, iteration) = torch.load(self.planner.ckpts[plan_id], weights_only=False) + model_params = list(model_params) + dynamic_gaussians[plan_id].restore(model_params, None) + + semantic_idx = torch.argmax(self.gaussians.get_full_3D_features, dim=-1, keepdim=True) + ground_xyz = self.gaussians.get_full_xyz[(semantic_idx == 0)[:, 0]].detach().cpu().numpy() + scene_xyz = self.gaussians.get_full_xyz[((semantic_idx > 1) & (semantic_idx != 10))[:, 0]].detach().cpu().numpy() + ground_pcd = o3d.geometry.PointCloud() + ground_pcd.points = o3d.utility.Vector3dVector(ground_xyz.astype(float)) + o3d.io.write_point_cloud(os.path.join(output, 'ground.ply'), ground_pcd) + scene_pcd = o3d.geometry.PointCloud() + scene_pcd.points = o3d.utility.Vector3dVector(scene_xyz.astype(float)) + o3d.io.write_point_cloud(os.path.join(output, 'scene.ply'), scene_pcd) + + unicycles = {} + + if cfg.scenario.load_HD_map and self.planner is not None: + self.planner.update_agent_route() + + self.cam_params, cam_align, self.cam_rect = load_camera_cfg(cfg.camera) + + self.ego_verts = np.array([[0.5, 0, 0.5], [0.5, 0, -0.5], [0.5, 1.0, 0.5], [0.5, 1.0, -0.5], + [-0.5, 0, -0.5], [-0.5, 0, 0.5], [-0.5, 1.0, -0.5], [-0.5, 1.0, 0.5]]) + self.whl = np.array([1.6, 1.5, 3.0]) + self.ego_verts *= self.whl + self.data_type = cfg.data_type + + self.action_space = spaces.Dict( + { + "steer_rate": spaces.Box(self.kinematic['min_steer'], self.kinematic['max_steer'], dtype=float), + "acc": spaces.Box(self.kinematic['min_acc'], self.kinematic['max_acc'], dtype=float) + } + ) + self.observation_space = spaces.Dict( + { + 'rgb': spaces.Dict({ + cam_name: spaces.Box( + low=0, high=255, + shape=(params['intrinsic']['H'], params['intrinsic']['W'], 3), dtype=np.uint8 + ) for cam_name, params in self.cam_params.items() + }), + # 'semantic': spaces.Dict({ + # cam_name: spaces.Box( + # low=0, high=50, + # shape=(params['intrinsic']['H'], params['intrinsic']['W']), dtype=np.uint8 + # ) for cam_name, params in self.cam_params.items() + # }), + # 'depth': spaces.Dict({ + # cam_name: spaces.Box( + # low=0, high=1000, + # shape=(params['intrinsic']['H'], params['intrinsic']['W']), dtype=np.float32 + # ) for cam_name, params in self.cam_params.items() + # }), + } + ) + self.fric = self.kinematic['fric'] + + self.start_vr = self.kinematic['start_vr'] + self.start_vab = self.kinematic['start_vab'] + self.start_velo = self.kinematic['start_velo'] + self.vr = deepcopy(self.kinematic['start_vr']) + self.vab = deepcopy(self.kinematic['start_vab']) + self.velo = deepcopy(self.kinematic['start_velo']) + self.steer = deepcopy(self.kinematic['start_steer']) + self.dt = self.kinematic['dt'] + + bg_color = [1, 1, 1] if cfg.model.white_background else [0, 0, 0] + self.render_fn = render + self.render_kwargs = { + "pc": self.gaussians, + "bg_color": torch.tensor(bg_color, dtype=torch.float32, device="cuda"), + "dynamic_gaussians": dynamic_gaussians, + "unicycles": unicycles + } + gaussians = self.gaussians + semantic_idx = torch.argmax(gaussians.get_3D_features, dim=-1, keepdim=True) + opacities = gaussians.get_opacity[:, 0] + mask = ((semantic_idx > 1) & (semantic_idx != 10))[:, 0] & (opacities > 0.8) + self.points = gaussians.get_xyz[mask] + + self.last_accel = 0 + self.last_steer_rate = 0 + + self.timestamp = 0 + + def ground_height(self, u, v): + cam_poses, cam_height, _ = self.ground_model + cam_dist = np.sqrt( + (cam_poses[:, 0, 3] - u)**2 + (cam_poses[:, 2, 3] - v)**2 + ) + nearest_cam_idx = np.argmin(cam_dist, axis=0) + nearest_c2w = cam_poses[nearest_cam_idx] + + nearest_w2c = np.linalg.inv(nearest_c2w) + uhv_local = nearest_w2c[:3, :3] @ np.array([u, 0, v]) + nearest_w2c[:3, 3] + uhv_local[1] = 0 + uhv_world = nearest_c2w[:3, :3] @ uhv_local + nearest_c2w[:3, 3] + + return uhv_world[1] + + @property + def route_completion(self): + cam_poses, _, _ = self.ground_model + cam_dist = np.sqrt( + (cam_poses[:, 0, 3] - self.vab[0])**2 + (cam_poses[:, 2, 3] - self.vab[1])**2 + ) + nearest_cam_idx = np.argmin(cam_dist, axis=0) + return (nearest_cam_idx + 1) / (cam_poses.shape[0] * 0.9), cam_dist[nearest_cam_idx] + + + @property + def vt(self): + vt = np.zeros(3) + vt[[0, 2]] = self.vab + vt[1] = self.ground_height(self.vab[0], self.vab[1]) + return vt + + @property + def ego(self): + return rt2pose(self.vr, self.vt) + + @property + def ego_state(self): + return torch.tensor([self.vab[0], self.vab[1], self.vr[1], self.velo]) + + @property + def ego_box(self): + return [self.vt[2], -self.vt[0], -self.vt[1], self.whl[0], self.whl[2], self.whl[1], -self.vr[1]] + + @property + def objs_list(self): + obj_boxes = [] + objs = self.render_kwargs['planning'][0] + for obj_id, obj_b2w in objs.items(): + yaw = SCR.from_matrix(obj_b2w[:3, :3].detach().cpu().numpy()).as_euler('YXZ')[0] + # X, Y, Z in IMU, w, l, h + wlh = self.planner.wlhs[obj_id] + obj_boxes.append([obj_b2w[2, 3].item(), -obj_b2w[0, 3].item(), -obj_b2w[1, 3].item(), wlh[0], wlh[1], wlh[2], -yaw-0.5*np.pi]) + return obj_boxes + + def _get_obs(self): + rgbs, semantics, depths = {}, {}, {} + v2front = self.cam_params['CAM_FRONT']["v2c"] + for cam_name, params in self.cam_params.items(): + intrinsic, v2c = params['intrinsic'], params['v2c'] + c2front = v2front @ np.linalg.inv(v2c) @ self.cam_rect + c2w = self.ego @ c2front + viewpoint = create_cam(intrinsic, c2w) + with torch.no_grad(): + render_pkg = self.render_fn(viewpoint=viewpoint, prev_viewpoint=None, **self.render_kwargs) + rgb = (torch.permute(render_pkg['render'].clamp(0, 1), (1,2,0)).detach().cpu().numpy() * 255).astype(np.uint8) + smt = torch.argmax(render_pkg['feats'], dim=0).detach().cpu().numpy().astype(np.uint8) + depth = render_pkg['depth'][0].detach().cpu().numpy() + if (self.data_type == 'waymo' or self.data_type == 'kitti360') and 'BACK' in cam_name: + rgbs[cam_name] = np.zeros_like(rgb) + semantics[cam_name] = np.zeros_like(smt) + depths[cam_name] = np.zeros_like(depth) + else: + rgbs[cam_name] = rgb + semantics[cam_name] = smt + depths[cam_name] = depth + + return { + 'rgb': rgbs, + # 'semantic': semantics, + # 'depth': depths, + } + + def _get_info(self): + wego_r, wego_t = pose2rt(self.ego) + cam_poses, _, commands = self.ground_model + dist = np.sum((cam_poses[:, :3, 3] - self.vt) ** 2, axis=-1) + nearest_cam_idx = np.argmin(dist) + command = commands[nearest_cam_idx] + return { + 'ego_pos' : wego_t.tolist(), + 'ego_rot' : wego_r.tolist(), + 'ego_velo' : self.velo, + 'ego_steer': self.steer, + 'accelerate': self.last_accel, + 'steer_rate': self.last_steer_rate, + 'timestamp': self.timestamp, + 'command': command, + 'ego_box': self.ego_box, + 'obj_boxes': self.objs_list, + 'cam_params': self.cam_params, + # 'ego_verts': verts, + } + + def reset(self, seed=None, options=None): + self.vr = deepcopy(self.start_vr) + self.vab = deepcopy(self.start_vab) + self.velo = deepcopy(self.start_velo) + self.timestamp = 0 + + if self.planner is not None: + self.render_kwargs['planning'] = self.planner.plan_traj(self.timestamp, self.ego_state) + else: + self.render_kwargs['planning'] = [{}, {}] + + observation = self._get_obs() + info = self._get_info() + + return observation, info + + def step(self, action): + self.timestamp += self.dt + if self.planner is not None: + self.render_kwargs['planning'] = self.planner.plan_traj(self.timestamp, self.ego_state) + else: + self.render_kwargs['planning'] = [{}, {}] + steer_rate, acc = action['steer_rate'], action['acc'] + self.last_steer_rate, self.last_accel = steer_rate, acc + L = self.kinematic['Lr'] + self.kinematic['Lf'] + self.velo += acc * self.dt + self.steer += steer_rate * self.dt + theta = self.vr[1] + # print(theta / np.pi * 180, self.steer / np.pi * 180) + self.vab[0] = self.vab[0] + self.velo * np.sin(theta) * self.dt + self.vab[1] = self.vab[1] + self.velo * np.cos(theta) * self.dt + self.vr[1] = theta + self.velo * np.tan(self.steer) / L * self.dt + + terminated = False + reward = 0 + verts = (self.ego[:3, :3] @ self.ego_verts.T).T + self.ego[:3, 3] + verts = torch.from_numpy(verts.astype(np.float32)).cuda() + + bg_collision = bg_collision_det(self.points, verts) + if bg_collision: + terminated = True + print('Collision with background') + reward = -100 + + fg_collision = fg_collision_det(self.ego_box, self.objs_list) + if fg_collision: + terminated = True + print('Collision with foreground') + reward = -100 + + rc, dist = self.route_completion + if dist > 10: + terminated=True + print('Far from preset trajectory') + reward = -50 + + if rc >= 1: + terminated = True + print('Complete') + reward = 1000 + + observation = self._get_obs() + info = self._get_info() + info['rc'] = rc + info['collision'] = bg_collision or fg_collision + + return observation, reward, terminated, False, info \ No newline at end of file diff --git a/code/sim/ilqr/__pycache__/lqr.cpython-311.pyc b/code/sim/ilqr/__pycache__/lqr.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..69514602eff0ea70c11fe9027194293aeeb4b4e0 Binary files /dev/null and b/code/sim/ilqr/__pycache__/lqr.cpython-311.pyc differ diff --git a/code/sim/ilqr/__pycache__/lqr_solver.cpython-311.pyc b/code/sim/ilqr/__pycache__/lqr_solver.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e546ebb2696e7ec552a525cd11d99b37f40d039c Binary files /dev/null and b/code/sim/ilqr/__pycache__/lqr_solver.cpython-311.pyc differ diff --git a/code/sim/ilqr/__pycache__/utils.cpython-311.pyc b/code/sim/ilqr/__pycache__/utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8084f53110128e82974fda9195e0896215dd6732 Binary files /dev/null and b/code/sim/ilqr/__pycache__/utils.cpython-311.pyc differ diff --git a/code/sim/ilqr/lqr.py b/code/sim/ilqr/lqr.py new file mode 100644 index 0000000000000000000000000000000000000000..d65c2025de7723d7997acf501c7b2da4cd2c602b --- /dev/null +++ b/code/sim/ilqr/lqr.py @@ -0,0 +1,55 @@ +from sim.ilqr.lqr_solver import ILQRSolverParameters, ILQRWarmStartParameters, ILQRSolver +import numpy as np + +solver_params = ILQRSolverParameters( + discretization_time=0.5, + state_cost_diagonal_entries=[1.0, 1.0, 10.0, 0.0, 0.0], + input_cost_diagonal_entries=[1.0, 10.0], + state_trust_region_entries=[1.0] * 5, + input_trust_region_entries=[1.0] * 2, + max_ilqr_iterations=100, + convergence_threshold=1e-6, + max_solve_time=0.05, + max_acceleration=3.0, + max_steering_angle=np.pi / 3.0, + max_steering_angle_rate=0.4, + min_velocity_linearization=0.01, + wheelbase=2.7 +) + +warm_start_params = ILQRWarmStartParameters( + k_velocity_error_feedback=0.5, + k_steering_angle_error_feedback=0.05, + lookahead_distance_lateral_error=15.0, + k_lateral_error=0.1, + jerk_penalty_warm_start_fit=1e-4, + curvature_rate_penalty_warm_start_fit=1e-2, +) + +lqr = ILQRSolver(solver_params=solver_params, warm_start_params=warm_start_params) + +def plan2control(plan_traj, init_state): + current_state = init_state + solutions = lqr.solve(current_state, plan_traj) + optimal_inputs = solutions[-1].input_trajectory + accel_cmd = optimal_inputs[0, 0] + steering_rate_cmd = optimal_inputs[0, 1] + return accel_cmd, steering_rate_cmd + +if __name__ == '__main__': + # plan_traj = np.zeros((6,5)) + # plan_traj[:, 0] = 1 + # plan_traj[:, 1] = np.ones(6) + # plan_traj = np.cumsum(plan_traj, axis=0) + # print(plan_traj) + plan_traj = np.array([[-0.18724936, 2.29100776, 0., 0., 0., ], + [-0.29260731, 2.2971828 , 0., 0., 0. ], + [-0.46831554, 2.55596018, 0., 0., 0. ], + [-0.5859955 , 2.73183298, 0., 0., 0. ], + [-0.62684 , 2.84659386, 0., 0., 0. ], + [-0.67761713, 2.80647802, 0., 0., 0. ]]) + plan_traj = plan_traj[:, [1,0,2,3,4]] + init_state = np.array([0.00000000e+00, 3.46944695e-17, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00]) + print(plan_traj.shape, init_state.shape) + acc, steer = plan2control(plan_traj, init_state) + print(acc, steer) diff --git a/code/sim/ilqr/lqr_solver.py b/code/sim/ilqr/lqr_solver.py new file mode 100644 index 0000000000000000000000000000000000000000..f3a1bb93b9e5a0d1c24a95818fc887b1b3b71397 --- /dev/null +++ b/code/sim/ilqr/lqr_solver.py @@ -0,0 +1,689 @@ +""" +This provides an implementation of the iterative linear quadratic regulator (iLQR) algorithm for trajectory tracking. +It is specialized to the case with a discrete-time kinematic bicycle model and a quadratic trajectory tracking cost. + +Original (Nonlinear) Discrete Time System: + z_k = [x_k, y_k, theta_k, v_k, delta_k] + u_k = [a_k, phi_k] + + x_{k+1} = x_k + v_k * cos(theta_k) * dt + y_{k+1} = y_k + v_k * sin(theta_k) * dt + theta_{k+1} = theta_k + v_k * tan(delta_k) / L * dt + v_{k+1} = v_k + a_k * dt + delta_{k+1} = delta_k + phi_k * dt + + where (x_k, y_k, theta_k) is the pose at timestep k with time discretization dt, + v_k and a_k are velocity and acceleration, + delta_k and phi_k are steering angle and steering angle rate, + and L is the vehicle wheelbase. + +Quadratic Tracking Cost: + J = sum_{k=0}^{N-1} ||u_k||_2^{R_k} + + sum_{k=0}^N ||z_k - z_{ref,k}||_2^{Q_k} +For simplicity, we opt to use constant input cost matrices R_k = R and constant state cost matrices Q_k = Q. + +There are multiple improvements that can be done for this implementation, but omitted for simplicity of the code. +Some of these include: + * Handle constraints directly in the optimization (e.g. log-barrier / penalty method with quadratic cost estimate). + * Line search in the input policy update (feedforward term) to determine a good gradient step size. + +References Used: https://people.eecs.berkeley.edu/~pabbeel/cs287-fa19/slides/Lec5-LQR.pdf and + https://www.cs.cmu.edu/~rsalakhu/10703/Lectures/Lecture_trajectoryoptimization.pdf +""" + +import time +from dataclasses import dataclass, fields +from typing import List, Optional, Tuple + +import numpy as np +import numpy.typing as npt + +# from nuplan.common.actor_state.vehicle_parameters import get_pacifica_parameters +# from nuplan.common.geometry.compute import principal_value +# from nuplan.planning.simulation.controller.tracker.tracker_utils import ( +# complete_kinematic_state_and_inputs_from_poses, +# compute_steering_angle_feedback, +# ) +from sim.ilqr.utils import principal_value, complete_kinematic_state_and_inputs_from_poses, compute_steering_angle_feedback + +DoubleMatrix = npt.NDArray[np.float64] + + +@dataclass(frozen=True) +class ILQRSolverParameters: + """Parameters related to the solver implementation.""" + + discretization_time: float # [s] Time discretization used for integration. + + # Cost weights for state [x, y, heading, velocity, steering angle] and input variables [acceleration, steering rate]. + state_cost_diagonal_entries: List[float] + input_cost_diagonal_entries: List[float] + + # Trust region cost weights for state and input variables. Helps keep linearization error per update step bounded. + state_trust_region_entries: List[float] + input_trust_region_entries: List[float] + + # Parameters related to solver runtime / solution sub-optimality. + max_ilqr_iterations: int # Maximum number of iterations to run iLQR before timeout. + convergence_threshold: float # Threshold for delta inputs below which we can terminate iLQR early. + max_solve_time: Optional[ + float + ] # [s] If defined, sets a maximum time to run a solve call of iLQR before terminating. + + # Constraints for underlying dynamics model. + max_acceleration: float # [m/s^2] Absolute value threshold on acceleration input. + max_steering_angle: float # [rad] Absolute value threshold on steering angle state. + max_steering_angle_rate: float # [rad/s] Absolute value threshold on steering rate input. + + # Parameters for dynamics / linearization. + min_velocity_linearization: float # [m/s] Absolute value threshold below which linearization velocity is modified. + wheelbase: float # [m] Wheelbase length parameter for the vehicle. + + def __post_init__(self) -> None: + """Ensure entries lie in expected bounds and initialize wheelbase.""" + for entry in [ + "discretization_time", + "max_ilqr_iterations", + "convergence_threshold", + "max_acceleration", + "max_steering_angle", + "max_steering_angle_rate", + "min_velocity_linearization", + "wheelbase", + ]: + assert getattr(self, entry) > 0.0, f"Field {entry} should be positive." + + assert self.max_steering_angle < np.pi / 2.0, "Max steering angle should be less than 90 degrees." + + if isinstance(self.max_solve_time, float): + assert self.max_solve_time > 0.0, "The specified max solve time should be positive." + + assert np.all([x >= 0 for x in self.state_cost_diagonal_entries]), "Q matrix must be positive semidefinite." + assert np.all([x > 0 for x in self.input_cost_diagonal_entries]), "R matrix must be positive definite." + + assert np.all( + [x > 0 for x in self.state_trust_region_entries] + ), "State trust region cost matrix must be positive definite." + assert np.all( + [x > 0 for x in self.input_trust_region_entries] + ), "Input trust region cost matrix must be positive definite." + + +@dataclass(frozen=True) +class ILQRWarmStartParameters: + """Parameters related to generating a warm start trajectory for iLQR.""" + + k_velocity_error_feedback: float # Gain for initial velocity error for warm start acceleration. + k_steering_angle_error_feedback: float # Gain for initial steering angle error for warm start steering rate. + lookahead_distance_lateral_error: float # [m] Distance ahead for which we estimate lateral error. + k_lateral_error: float # Gain for lateral error to compute steering angle feedback. + jerk_penalty_warm_start_fit: float # Penalty for jerk in velocity profile estimation. + curvature_rate_penalty_warm_start_fit: float # Penalty for curvature rate in curvature profile estimation. + + def __post_init__(self) -> None: + """Ensure entries lie in expected bounds.""" + for entry in [ + "k_velocity_error_feedback", + "k_steering_angle_error_feedback", + "lookahead_distance_lateral_error", + "k_lateral_error", + "jerk_penalty_warm_start_fit", + "curvature_rate_penalty_warm_start_fit", + ]: + assert getattr(self, entry) > 0.0, f"Field {entry} should be positive." + + +@dataclass(frozen=True) +class ILQRIterate: + """Contains state, input, and associated Jacobian trajectories needed to perform an update step of iLQR.""" + + state_trajectory: DoubleMatrix + input_trajectory: DoubleMatrix + state_jacobian_trajectory: DoubleMatrix + input_jacobian_trajectory: DoubleMatrix + + def __post_init__(self) -> None: + """Check consistency of dimension across trajectory elements.""" + assert len(self.state_trajectory.shape) == 2, "Expect state trajectory to be a 2D matrix." + state_trajectory_length, state_dim = self.state_trajectory.shape + + assert len(self.input_trajectory.shape) == 2, "Expect input trajectory to be a 2D matrix." + input_trajectory_length, input_dim = self.input_trajectory.shape + + assert ( + input_trajectory_length == state_trajectory_length - 1 + ), "State trajectory should be 1 longer than the input trajectory." + assert self.state_jacobian_trajectory.shape == (input_trajectory_length, state_dim, state_dim) + assert self.input_jacobian_trajectory.shape == (input_trajectory_length, state_dim, input_dim) + + for field in fields(self): + # Make sure that we have no nan entries in our trajectory rollout prior to operating on this. + assert ~np.any(np.isnan(getattr(self, field.name))), f"{field.name} has unexpected nan values." + + +@dataclass(frozen=True) +class ILQRInputPolicy: + """Contains parameters for the perturbation input policy computed after performing LQR.""" + + state_feedback_matrices: DoubleMatrix + feedforward_inputs: DoubleMatrix + + def __post__init__(self) -> None: + """Check shape of policy parameters.""" + assert ( + len(self.state_feedback_matrices.shape) == 3 + ), "Expected state_feedback_matrices to have shape (n_horizon, n_inputs, n_states)" + + assert ( + len(self.feedforward_inputs.shape) == 2 + ), "Expected feedforward inputs to have shape (n_horizon, n_inputs)." + + assert ( + self.feedforward_inputs.shape == self.state_feedback_matrices.shape[:2] + ), "Inconsistent horizon or input dimension between feedforward inputs and state feedback matrices." + + for field in fields(self): + # Make sure that we have no nan entries in our policy parameters prior to using them. + assert ~np.any(np.isnan(getattr(self, field.name))), f"{field.name} has unexpected nan values." + + +@dataclass(frozen=True) +class ILQRSolution: + """Contains the iLQR solution with associated cost for consumption by the solver's client.""" + + state_trajectory: DoubleMatrix + input_trajectory: DoubleMatrix + + tracking_cost: float + + def __post_init__(self) -> None: + """Check consistency of dimension across trajectory elements and nonnegative cost.""" + assert len(self.state_trajectory.shape) == 2, "Expect state trajectory to be a 2D matrix." + state_trajectory_length, _ = self.state_trajectory.shape + + assert len(self.input_trajectory.shape) == 2, "Expect input trajectory to be a 2D matrix." + input_trajectory_length, _ = self.input_trajectory.shape + + assert ( + input_trajectory_length == state_trajectory_length - 1 + ), "State trajectory should be 1 longer than the input trajectory." + + assert self.tracking_cost >= 0.0, "Expect the tracking cost to be nonnegative." + + +class ILQRSolver: + """iLQR solver implementation, see module docstring for details.""" + + def __init__( + self, + solver_params: ILQRSolverParameters, + warm_start_params: ILQRWarmStartParameters, + ) -> None: + """ + Initialize solver parameters. + :param solver_params: Contains solver parameters for iLQR. + :param warm_start_params: Contains warm start parameters for iLQR. + """ + self._solver_params = solver_params + self._warm_start_params = warm_start_params + + self._n_states = 5 # state dimension + self._n_inputs = 2 # input dimension + + state_cost_diagonal_entries = self._solver_params.state_cost_diagonal_entries + assert ( + len(state_cost_diagonal_entries) == self._n_states + ), f"State cost matrix should have diagonal length {self._n_states}." + self._state_cost_matrix: DoubleMatrix = np.diag(state_cost_diagonal_entries) + + input_cost_diagonal_entries = self._solver_params.input_cost_diagonal_entries + assert ( + len(input_cost_diagonal_entries) == self._n_inputs + ), f"Input cost matrix should have diagonal length {self._n_inputs}." + self._input_cost_matrix: DoubleMatrix = np.diag(input_cost_diagonal_entries) + + state_trust_region_entries = self._solver_params.state_trust_region_entries + assert ( + len(state_trust_region_entries) == self._n_states + ), f"State trust region cost matrix should have diagonal length {self._n_states}." + self._state_trust_region_cost_matrix: DoubleMatrix = np.diag(state_trust_region_entries) + + input_trust_region_entries = self._solver_params.input_trust_region_entries + assert ( + len(input_trust_region_entries) == self._n_inputs + ), f"Input trust region cost matrix should have diagonal length {self._n_inputs}." + self._input_trust_region_cost_matrix: DoubleMatrix = np.diag(input_trust_region_entries) + + max_acceleration = self._solver_params.max_acceleration + max_steering_angle_rate = self._solver_params.max_steering_angle_rate + + # Define input clip limits once to avoid recomputation in _clip_inputs. + self._input_clip_min = (-max_acceleration, -max_steering_angle_rate) + self._input_clip_max = (max_acceleration, max_steering_angle_rate) + + def solve(self, current_state: DoubleMatrix, reference_trajectory: DoubleMatrix) -> List[ILQRSolution]: + """ + Run the main iLQR loop used to try to find (locally) optimal inputs to track the reference trajectory. + :param current_state: The initial state from which we apply inputs, z_0. + :param reference_trajectory: The state reference we'd like to track, inclusive of the initial timestep, + z_{r,k} for k in {0, ..., N}. + :return: A list of solution iterates after running the iLQR algorithm where the index is the iteration number. + """ + # Check that state parameter has the right shape. + assert current_state.shape == (self._n_states,), "Incorrect state shape." + + # Check that reference trajectory parameter has the right shape. + assert len(reference_trajectory.shape) == 2, "Reference trajectory should be a 2D matrix." + reference_trajectory_length, reference_trajectory_state_dimension = reference_trajectory.shape + assert reference_trajectory_length > 1, "The reference trajectory should be at least two timesteps long." + assert ( + reference_trajectory_state_dimension == self._n_states + ), "The reference trajectory should have a matching state dimension." + + # List of ILQRSolution results where the index corresponds to the iteration of iLQR. + solution_list: List[ILQRSolution] = [] + + # Get warm start input and state trajectory, as well as associated Jacobians. + current_iterate = self._input_warm_start(current_state, reference_trajectory) + + # Main iLQR Loop. + solve_start_time = time.perf_counter() + for _ in range(self._solver_params.max_ilqr_iterations): + # Determine the cost and store the associated solution object. + tracking_cost = self._compute_tracking_cost( + iterate=current_iterate, + reference_trajectory=reference_trajectory, + ) + solution_list.append( + ILQRSolution( + input_trajectory=current_iterate.input_trajectory, + state_trajectory=current_iterate.state_trajectory, + tracking_cost=tracking_cost, + ) + ) + + # Determine the LQR optimal perturbations to apply. + lqr_input_policy = self._run_lqr_backward_recursion( + current_iterate=current_iterate, + reference_trajectory=reference_trajectory, + ) + + # Apply the optimal perturbations to generate the next input trajectory iterate. + input_trajectory_next = self._update_inputs_with_policy( + current_iterate=current_iterate, + lqr_input_policy=lqr_input_policy, + ) + + # Check for convergence/timeout and terminate early if so. + # Else update the input_trajectory iterate and continue. + input_trajectory_norm_difference = np.linalg.norm(input_trajectory_next - current_iterate.input_trajectory) + + current_iterate = self._run_forward_dynamics(current_state, input_trajectory_next) + + if input_trajectory_norm_difference < self._solver_params.convergence_threshold: + break + + elapsed_time = time.perf_counter() - solve_start_time + if ( + isinstance(self._solver_params.max_solve_time, float) + and elapsed_time >= self._solver_params.max_solve_time + ): + break + + # Store the final iterate in the solution_dict. + tracking_cost = self._compute_tracking_cost( + iterate=current_iterate, + reference_trajectory=reference_trajectory, + ) + solution_list.append( + ILQRSolution( + input_trajectory=current_iterate.input_trajectory, + state_trajectory=current_iterate.state_trajectory, + tracking_cost=tracking_cost, + ) + ) + + return solution_list + + #################################################################################################################### + # Helper methods. + #################################################################################################################### + + def _compute_tracking_cost(self, iterate: ILQRIterate, reference_trajectory: DoubleMatrix) -> float: + """ + Compute the trajectory tracking cost given a candidate solution. + :param iterate: Contains the candidate state and input trajectory to evaluate. + :param reference_trajectory: The desired state reference trajectory with same length as state_trajectory. + :return: The tracking cost of the candidate state/input trajectory. + """ + input_trajectory = iterate.input_trajectory + state_trajectory = iterate.state_trajectory + + assert len(state_trajectory) == len( + reference_trajectory + ), "The state and reference trajectory should have the same length." + + error_state_trajectory = state_trajectory - reference_trajectory + error_state_trajectory[:, 2] = principal_value(error_state_trajectory[:, 2]) + + cost = np.sum([u.T @ self._input_cost_matrix @ u for u in input_trajectory]) + np.sum( + [e.T @ self._state_cost_matrix @ e for e in error_state_trajectory] + ) + + return float(cost) + + def _clip_inputs(self, inputs: DoubleMatrix) -> DoubleMatrix: + """ + Used to clip control inputs within constraints. + :param: inputs: The control inputs with shape (self._n_inputs,) to clip. + :return: Clipped version of the control inputs, unmodified if already within constraints. + """ + assert inputs.shape == (self._n_inputs,), f"The inputs should be a 1D vector with {self._n_inputs} elements." + + return np.clip(inputs, self._input_clip_min, self._input_clip_max) # type: ignore + + def _clip_steering_angle(self, steering_angle: float) -> float: + """ + Used to clip the steering angle state within bounds. + :param steering_angle: [rad] A steering angle (scalar) to clip. + :return: [rad] The clipped steering angle. + """ + steering_angle_sign = 1.0 if steering_angle >= 0 else -1.0 + steering_angle = steering_angle_sign * min(abs(steering_angle), self._solver_params.max_steering_angle) + return steering_angle + + def _input_warm_start(self, current_state: DoubleMatrix, reference_trajectory: DoubleMatrix) -> ILQRIterate: + """ + Given a reference trajectory, we generate the warm start (initial guess) by inferring the inputs applied based + on poses in the reference trajectory. + :param current_state: The initial state from which we apply inputs. + :param reference_trajectory: The reference trajectory we are trying to follow. + :return: The warm start iterate from which to start iLQR. + """ + reference_states_completed, reference_inputs_completed = complete_kinematic_state_and_inputs_from_poses( + discretization_time=self._solver_params.discretization_time, + wheel_base=self._solver_params.wheelbase, + poses=reference_trajectory[:, :3], + jerk_penalty=self._warm_start_params.jerk_penalty_warm_start_fit, + curvature_rate_penalty=self._warm_start_params.curvature_rate_penalty_warm_start_fit, + ) + + # We could just stop here and apply reference_inputs_completed (assuming it satisfies constraints). + # This could work if current_state = reference_states_completed[0,:] - i.e. no initial tracking error. + # We add feedback input terms for the first control input only to account for nonzero initial tracking error. + _, _, _, velocity_current, steering_angle_current = current_state + _, _, _, velocity_reference, steering_angle_reference = reference_states_completed[0, :] + + acceleration_feedback = -self._warm_start_params.k_velocity_error_feedback * ( + velocity_current - velocity_reference + ) + + steering_angle_feedback = compute_steering_angle_feedback( + pose_reference=current_state[:3], + pose_current=reference_states_completed[0, :3], + lookahead_distance=self._warm_start_params.lookahead_distance_lateral_error, + k_lateral_error=self._warm_start_params.k_lateral_error, + ) + steering_angle_desired = steering_angle_feedback + steering_angle_reference + steering_rate_feedback = -self._warm_start_params.k_steering_angle_error_feedback * ( + steering_angle_current - steering_angle_desired + ) + + reference_inputs_completed[0, 0] += acceleration_feedback + reference_inputs_completed[0, 1] += steering_rate_feedback + + # We rerun dynamics with constraints applied to make sure we have a feasible warm start for iLQR. + return self._run_forward_dynamics(current_state, reference_inputs_completed) + + #################################################################################################################### + # Dynamics and Jacobian. + #################################################################################################################### + + def _run_forward_dynamics(self, current_state: DoubleMatrix, input_trajectory: DoubleMatrix) -> ILQRIterate: + """ + Compute states and corresponding state/input Jacobian matrices using forward dynamics. + We additionally return the input since the dynamics may modify the input to ensure constraint satisfaction. + :param current_state: The initial state from which we apply inputs. Must be feasible given constraints. + :param input_trajectory: The input trajectory applied to the model. May be modified to ensure feasibility. + :return: A feasible iterate after applying dynamics with state/input trajectories and Jacobian matrices. + """ + # Store rollout as a set of numpy arrays, initialized as np.nan to ensure we correctly fill them in. + # The state trajectory includes the current_state, z_0, and is 1 element longer than the other arrays. + # The final_input_trajectory captures the applied input for the dynamics model satisfying constraints. + N = len(input_trajectory) + state_trajectory = np.nan * np.ones((N + 1, self._n_states), dtype=np.float64) + final_input_trajectory = np.nan * np.ones_like(input_trajectory, dtype=np.float64) + state_jacobian_trajectory = np.nan * np.ones((N, self._n_states, self._n_states), dtype=np.float64) + final_input_jacobian_trajectory = np.nan * np.ones((N, self._n_states, self._n_inputs), dtype=np.float64) + + state_trajectory[0] = current_state + + for idx_u, u in enumerate(input_trajectory): + state_next, final_input, state_jacobian, final_input_jacobian = self._dynamics_and_jacobian( + state_trajectory[idx_u], u + ) + + state_trajectory[idx_u + 1] = state_next + final_input_trajectory[idx_u] = final_input + state_jacobian_trajectory[idx_u] = state_jacobian + final_input_jacobian_trajectory[idx_u] = final_input_jacobian + + iterate = ILQRIterate( + state_trajectory=state_trajectory, # type: ignore + input_trajectory=final_input_trajectory, # type: ignore + state_jacobian_trajectory=state_jacobian_trajectory, # type: ignore + input_jacobian_trajectory=final_input_jacobian_trajectory, # type: ignore + ) + + return iterate + + def _dynamics_and_jacobian( + self, current_state: DoubleMatrix, current_input: DoubleMatrix + ) -> Tuple[DoubleMatrix, DoubleMatrix, DoubleMatrix, DoubleMatrix]: + """ + Propagates the state forward by one step and computes the corresponding state and input Jacobian matrices. + We also impose all constraints here to ensure the current input and next state are always feasible. + :param current_state: The current state z_k. + :param current_input: The applied input u_k. + :return: The next state z_{k+1}, (possibly modified) input u_k, and state (df/dz) and input (df/du) Jacobians. + """ + x, y, heading, velocity, steering_angle = current_state + + # Check steering angle is in expected range for valid Jacobian matrices. + assert ( + np.abs(steering_angle) < np.pi / 2.0 + ), f"The steering angle {steering_angle} is outside expected limits. There is a singularity at delta = np.pi/2." + + # Input constraints: clip inputs within bounds and then use. + current_input = self._clip_inputs(current_input) + acceleration, steering_rate = current_input + + # Euler integration of bicycle model. + discretization_time = self._solver_params.discretization_time + wheelbase = self._solver_params.wheelbase + + next_state: DoubleMatrix = np.copy(current_state) + next_state[0] += velocity * np.cos(heading) * discretization_time + next_state[1] += velocity * np.sin(heading) * discretization_time + next_state[2] += velocity * np.tan(steering_angle) / wheelbase * discretization_time + next_state[3] += acceleration * discretization_time + next_state[4] += steering_rate * discretization_time + + # Constrain heading angle to lie within +/- pi. + next_state[2] = principal_value(next_state[2]) + + # State constraints: clip the steering_angle within bounds and update steering_rate accordingly. + next_steering_angle = self._clip_steering_angle(next_state[4]) + applied_steering_rate = (next_steering_angle - steering_angle) / discretization_time + next_state[4] = next_steering_angle + current_input[1] = applied_steering_rate + + # Now we construct and populate the state and input Jacobians. + state_jacobian: DoubleMatrix = np.eye(self._n_states, dtype=np.float64) + input_jacobian: DoubleMatrix = np.zeros((self._n_states, self._n_inputs), dtype=np.float64) + + # Set a nonzero velocity to handle issues when linearizing at (near) zero velocity. + # This helps e.g. when the vehicle is stopped with zero steering angle and needs to accelerate/turn. + # Without this, the A matrix will indicate steering has no impact on heading due to Euler discretization. + # There will be a rank drop in the controllability matrix, so the discrete-time algebraic Riccati equation + # may not have a solution (uncontrollable subspace) or it may not be unique. + min_velocity_linearization = self._solver_params.min_velocity_linearization + if -min_velocity_linearization <= velocity and velocity <= min_velocity_linearization: + sign_velocity = 1.0 if velocity >= 0.0 else -1.0 + velocity = sign_velocity * min_velocity_linearization + + state_jacobian[0, 2] = -velocity * np.sin(heading) * discretization_time + state_jacobian[0, 3] = np.cos(heading) * discretization_time + + state_jacobian[1, 2] = velocity * np.cos(heading) * discretization_time + state_jacobian[1, 3] = np.sin(heading) * discretization_time + + state_jacobian[2, 3] = np.tan(steering_angle) / wheelbase * discretization_time + state_jacobian[2, 4] = velocity * discretization_time / (wheelbase * np.cos(steering_angle) ** 2) + + input_jacobian[3, 0] = discretization_time + input_jacobian[4, 1] = discretization_time + + return next_state, current_input, state_jacobian, input_jacobian + + #################################################################################################################### + # Core LQR implementation. + #################################################################################################################### + + def _run_lqr_backward_recursion( + self, + current_iterate: ILQRIterate, + reference_trajectory: DoubleMatrix, + ) -> ILQRInputPolicy: + """ + Computes the locally optimal affine state feedback policy by applying dynamic programming to linear perturbation + dynamics about a specified linearization trajectory. We include a trust region penalty as part of the cost. + :param current_iterate: Contains all relevant linearization information needed to compute LQR policy. + :param reference_trajectory: The desired state trajectory we are tracking. + :return: An affine state feedback policy - state feedback matrices and feedforward inputs found using LQR. + """ + state_trajectory = current_iterate.state_trajectory + input_trajectory = current_iterate.input_trajectory + state_jacobian_trajectory = current_iterate.state_jacobian_trajectory + input_jacobian_trajectory = current_iterate.input_jacobian_trajectory + + # Check reference matches the expected shape. + assert reference_trajectory.shape == state_trajectory.shape, "The reference trajectory has incorrect shape." + + # Compute nominal error trajectory. + error_state_trajectory = state_trajectory - reference_trajectory + error_state_trajectory[:, 2] = principal_value(error_state_trajectory[:, 2]) + + # The value function has the form V_k(\Delta z_k) = \Delta z_k^T P_k \Delta z_k + 2 \rho_k^T \Delta z_k. + # So p_current = P_k is related to the Hessian of the value function at the current timestep. + # And rho_current = rho_k is part of the linear cost term in the value function at the current timestep. + p_current = self._state_cost_matrix + self._state_trust_region_cost_matrix + rho_current = self._state_cost_matrix @ error_state_trajectory[-1] + + # The optimal LQR policy has the form \Delta u_k^* = K_k \Delta z_k + \kappa_k + # We refer to K_k as state_feedback_matrix and \kappa_k as feedforward input in the code below. + N = len(input_trajectory) + state_feedback_matrices = np.nan * np.ones((N, self._n_inputs, self._n_states), dtype=np.float64) + feedforward_inputs = np.nan * np.ones((N, self._n_inputs), dtype=np.float64) + + for i in reversed(range(N)): + A = state_jacobian_trajectory[i] + B = input_jacobian_trajectory[i] + u = input_trajectory[i] + error = error_state_trajectory[i] + + # Compute the optimal input policy for this timestep. + inverse_matrix_term = np.linalg.inv( + self._input_cost_matrix + self._input_trust_region_cost_matrix + B.T @ p_current @ B + ) # invertible since we checked input_cost / input_trust_region_cost are positive definite during creation. + state_feedback_matrix = -inverse_matrix_term @ B.T @ p_current @ A + feedforward_input = -inverse_matrix_term @ (self._input_cost_matrix @ u + B.T @ rho_current) + + # Compute the optimal value function for this timestep. + a_closed_loop = A + B @ state_feedback_matrix + + p_prior = ( + self._state_cost_matrix + + self._state_trust_region_cost_matrix + + state_feedback_matrix.T @ self._input_cost_matrix @ state_feedback_matrix + + state_feedback_matrix.T @ self._input_trust_region_cost_matrix @ state_feedback_matrix + + a_closed_loop.T @ p_current @ a_closed_loop + ) + + rho_prior = ( + self._state_cost_matrix @ error + + state_feedback_matrix.T @ self._input_cost_matrix @ (feedforward_input + u) + + state_feedback_matrix.T @ self._input_trust_region_cost_matrix @ feedforward_input + + a_closed_loop.T @ p_current @ B @ feedforward_input + + a_closed_loop.T @ rho_current + ) + + p_current = p_prior + rho_current = rho_prior + + state_feedback_matrices[i] = state_feedback_matrix + feedforward_inputs[i] = feedforward_input + + lqr_input_policy = ILQRInputPolicy( + state_feedback_matrices=state_feedback_matrices, # type: ignore + feedforward_inputs=feedforward_inputs, # type: ignore + ) + + return lqr_input_policy + + def _update_inputs_with_policy( + self, + current_iterate: ILQRIterate, + lqr_input_policy: ILQRInputPolicy, + ) -> DoubleMatrix: + """ + Used to update an iterate of iLQR by applying a perturbation input policy for local cost improvement. + :param current_iterate: Contains the state and input trajectory about which we linearized. + :param lqr_input_policy: Contains the LQR policy to apply. + :return: The next input trajectory found by applying the LQR policy. + """ + state_trajectory = current_iterate.state_trajectory + input_trajectory = current_iterate.input_trajectory + + # Trajectory of state perturbations while applying feedback policy. + # Starts with zero as the initial states match exactly, only later states might vary. + delta_state_trajectory = np.nan * np.ones((len(input_trajectory) + 1, self._n_states), dtype=np.float64) + delta_state_trajectory[0] = [0.0] * self._n_states + + # This is the updated input trajectory we will return after applying the input perturbations. + input_next_trajectory = np.nan * np.ones_like(input_trajectory, dtype=np.float64) + + zip_object = zip( + input_trajectory, + state_trajectory[:-1], + state_trajectory[1:], + lqr_input_policy.state_feedback_matrices, + lqr_input_policy.feedforward_inputs, + ) + + for input_idx, (input_lin, state_lin, state_lin_next, state_feedback_matrix, feedforward_input) in enumerate( + zip_object + ): + # Compute locally optimal input perturbation. + delta_state = delta_state_trajectory[input_idx] + delta_input = state_feedback_matrix @ delta_state + feedforward_input + + # Apply state and input perturbation. + input_perturbed = input_lin + delta_input + state_perturbed = state_lin + delta_state + state_perturbed[2] = principal_value(state_perturbed[2]) + + # Run dynamics with perturbed state/inputs to get next state. + # We get the actually applied input since it might have been clipped/modified to satisfy constraints. + state_perturbed_next, input_perturbed, _, _ = self._dynamics_and_jacobian(state_perturbed, input_perturbed) + + # Compute next state perturbation given next state. + delta_state_next = state_perturbed_next - state_lin_next + delta_state_next[2] = principal_value(delta_state_next[2]) + + delta_state_trajectory[input_idx + 1] = delta_state_next + input_next_trajectory[input_idx] = input_perturbed + + assert ~np.any(np.isnan(input_next_trajectory)), "All next inputs should be valid float values." + + return input_next_trajectory # type: ignore \ No newline at end of file diff --git a/code/sim/ilqr/utils.py b/code/sim/ilqr/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..90e4feda9cf2890ad4e96f223d47f64ca75bf287 --- /dev/null +++ b/code/sim/ilqr/utils.py @@ -0,0 +1,346 @@ +from typing import Tuple +import numpy as np +import numpy.typing as npt + +DoubleMatrix = npt.NDArray[np.float64] + +def principal_value(angle, min_=-np.pi): + """ + Wrap heading angle in to specified domain (multiples of 2 pi alias), + ensuring that the angle is between min_ and min_ + 2 pi. This function raises an error if the angle is infinite + :param angle: rad + :param min_: minimum domain for angle (rad) + :return angle wrapped to [min_, min_ + 2 pi). + """ + assert np.all(np.isfinite(angle)), "angle is not finite" + lhs = (angle - min_) % (2 * np.pi) + min_ + return lhs + +def compute_steering_angle_feedback( + pose_reference, pose_current, lookahead_distance, k_lateral_error +): + """ + Given pose information, determines the steering angle feedback value to address initial tracking error. + This is based on the feedback controller developed in Section 2.2 of the following paper: + https://ddl.stanford.edu/publications/design-feedback-feedforward-steering-controller-accurate-path-tracking-and-stability + :param pose_reference: Contains the reference pose at the current timestep. + :param pose_current: Contains the actual pose at the current timestep. + :param lookahead_distance: [m] Distance ahead for which we should estimate lateral error based on a linear fit. + :param k_lateral_error: Feedback gain for lateral error used to determine steering angle feedback. + :return: [rad] The steering angle feedback to apply. + """ + assert pose_reference.shape == (3,), "We expect a single reference pose." + assert pose_current.shape == (3,), "We expect a single current pose." + + assert lookahead_distance > 0.0, "Lookahead distance should be positive." + assert k_lateral_error > 0.0, "Feedback gain for lateral error should be positive." + + x_reference, y_reference, heading_reference = pose_reference + x_current, y_current, heading_current = pose_current + + x_error = x_current - x_reference + y_error = y_current - y_reference + heading_error = principal_value(heading_current - heading_reference) + + lateral_error = -x_error * np.sin(heading_reference) + y_error * np.cos(heading_reference) + + return float(-k_lateral_error * (lateral_error + lookahead_distance * heading_error)) + +def _convert_curvature_profile_to_steering_profile( + curvature_profile: DoubleMatrix, + discretization_time: float, + wheel_base: float, +) -> Tuple[DoubleMatrix, DoubleMatrix]: + """ + Converts from a curvature profile to the corresponding steering profile. + We assume a kinematic bicycle model where curvature = tan(steering_angle) / wheel_base. + For simplicity, we just use finite differences to determine steering rate. + :param curvature_profile: [rad] Curvature trajectory to convert. + :param discretization_time: [s] Time discretization used for integration. + :param wheel_base: [m] The vehicle wheelbase parameter required for conversion. + :return: The [rad] steering angle and [rad/s] steering rate (derivative) profiles. + """ + assert discretization_time > 0.0, "Discretization time must be positive." + assert wheel_base > 0.0, "The vehicle's wheelbase length must be positive." + + steering_angle_profile = np.arctan(wheel_base * curvature_profile) + steering_rate_profile = np.diff(steering_angle_profile) / discretization_time + + return steering_angle_profile, steering_rate_profile + + +def _get_xy_heading_displacements_from_poses(poses: DoubleMatrix) -> Tuple[DoubleMatrix, DoubleMatrix]: + """ + Returns position and heading displacements given a pose trajectory. + :param poses: A trajectory of poses (x, y, heading). + :return: Tuple of xy displacements with shape (num_poses-1, 2) and heading displacements with shape (num_poses-1,). + """ + assert len(poses.shape) == 2, "Expect a 2D matrix representing a trajectory of poses." + assert poses.shape[0] > 1, "Cannot get displacements given an empty or single element pose trajectory." + assert poses.shape[1] == 3, "Expect pose to have three elements (x, y, heading)." + + # Compute displacements that are used to complete the kinematic state and input. + pose_differences = np.diff(poses, axis=0) + xy_displacements = pose_differences[:, :2] + heading_displacements = principal_value(pose_differences[:, 2]) + + return xy_displacements, heading_displacements + + +def _make_banded_difference_matrix(number_rows: int) -> DoubleMatrix: + """ + Returns a banded difference matrix with specified number_rows. + When applied to a vector [x_1, ..., x_N], it returns [x_2 - x_1, ..., x_N - x_{N-1}]. + :param number_rows: The row dimension of the banded difference matrix (e.g. N-1 in the example above). + :return: A banded difference matrix with shape (number_rows, number_rows+1). + """ + banded_matrix: DoubleMatrix = -1.0 * np.eye(number_rows + 1, dtype=np.float64)[:-1, :] + for ind in range(len(banded_matrix)): + banded_matrix[ind, ind + 1] = 1.0 + + return banded_matrix + + + +def _fit_initial_velocity_and_acceleration_profile( + xy_displacements: DoubleMatrix, heading_profile: DoubleMatrix, discretization_time: float, jerk_penalty: float +) -> Tuple[float, DoubleMatrix]: + """ + Estimates initial velocity (v_0) and acceleration ({a_0, ...}) using least squares with jerk penalty regularization. + :param xy_displacements: [m] Deviations in x and y occurring between M+1 poses, a M by 2 matrix. + :param heading_profile: [rad] Headings associated to the starting timestamp for xy_displacements, a M-length vector. + :param discretization_time: [s] Time discretization used for integration. + :param jerk_penalty: A regularization parameter used to penalize acceleration differences. Should be positive. + :return: Least squares solution for initial velocity (v_0) and acceleration profile ({a_0, ..., a_M-1}) + for M displacement values. + """ + assert discretization_time > 0.0, "Discretization time must be positive." + assert jerk_penalty > 0, "Should have a positive jerk_penalty." + + assert len(xy_displacements.shape) == 2, "Expect xy_displacements to be a matrix." + assert xy_displacements.shape[1] == 2, "Expect xy_displacements to have 2 columns." + + num_displacements = len(xy_displacements) # aka M in the docstring + + assert heading_profile.shape == ( + num_displacements, + ), "Expect the length of heading_profile to match that of xy_displacements." + + # Core problem: minimize_x ||y-Ax||_2 + y = xy_displacements.flatten() # Flatten to a vector, [delta x_0, delta y_0, ...] + + A: DoubleMatrix = np.zeros((2 * num_displacements, num_displacements), dtype=np.float64) + for idx_timestep, heading in enumerate(heading_profile): + start_row = 2 * idx_timestep # Which row of A corresponds to x-coordinate information at timestep k. + + # Related to v_0, initial velocity - column 0. + # We fill in rows for measurements delta x_k, delta y_k. + A[start_row : (start_row + 2), 0] = np.array( + [ + np.cos(heading) * discretization_time, + np.sin(heading) * discretization_time, + ], + dtype=np.float64, + ) + + if idx_timestep > 0: + # Related to {a_0, ..., a_k-1}, acceleration profile - column 1 to k. + # We fill in rows for measurements delta x_k, delta y_k. + A[start_row : (start_row + 2), 1 : (1 + idx_timestep)] = np.array( + [ + [np.cos(heading) * discretization_time**2], + [np.sin(heading) * discretization_time**2], + ], + dtype=np.float64, + ) + + # Regularization using jerk penalty, i.e. difference of acceleration values. + # If there are M displacements, then we have M - 1 acceleration values. + # That means we have M - 2 jerk values, thus we make a banded difference matrix of that size. + banded_matrix = _make_banded_difference_matrix(num_displacements - 2) + R: DoubleMatrix = np.block([np.zeros((len(banded_matrix), 1)), banded_matrix]) + + # Compute regularized least squares solution. + x = np.linalg.pinv(A.T @ A + jerk_penalty * R.T @ R) @ A.T @ y + + # Extract profile from solution. + initial_velocity = x[0] + acceleration_profile = x[1:] + + return initial_velocity, acceleration_profile + + +def _generate_profile_from_initial_condition_and_derivatives( + initial_condition: float, derivatives: DoubleMatrix, discretization_time: float +) -> DoubleMatrix: + """ + Returns the corresponding profile (i.e. trajectory) given an initial condition and derivatives at + multiple timesteps by integration. + :param initial_condition: The value of the variable at the initial timestep. + :param derivatives: The trajectory of time derivatives of the variable at timesteps 0,..., N-1. + :param discretization_time: [s] Time discretization used for integration. + :return: The trajectory of the variable at timesteps 0,..., N. + """ + assert discretization_time > 0.0, "Discretization time must be positive." + + profile = initial_condition + np.insert(np.cumsum(derivatives * discretization_time), 0, 0.0) + + return profile # type: ignore + + +def _fit_initial_curvature_and_curvature_rate_profile( + heading_displacements: DoubleMatrix, + velocity_profile: DoubleMatrix, + discretization_time: float, + curvature_rate_penalty: float, + initial_curvature_penalty: float = 1e-10, +) -> Tuple[float, DoubleMatrix]: + """ + Estimates initial curvature (curvature_0) and curvature rate ({curvature_rate_0, ...}) + using least squares with curvature rate regularization. + :param heading_displacements: [rad] Angular deviations in heading occuring between timesteps. + :param velocity_profile: [m/s] Estimated or actual velocities at the timesteps matching displacements. + :param discretization_time: [s] Time discretization used for integration. + :param curvature_rate_penalty: A regularization parameter used to penalize curvature_rate. Should be positive. + :param initial_curvature_penalty: A regularization parameter to handle zero initial speed. Should be positive and small. + :return: Least squares solution for initial curvature (curvature_0) and curvature rate profile + (curvature_rate_0, ..., curvature_rate_{M-1}) for M heading displacement values. + """ + assert discretization_time > 0.0, "Discretization time must be positive." + assert curvature_rate_penalty > 0.0, "Should have a positive curvature_rate_penalty." + assert initial_curvature_penalty > 0.0, "Should have a positive initial_curvature_penalty." + + # Core problem: minimize_x ||y-Ax||_2 + y = heading_displacements + A: DoubleMatrix = np.tri(len(y), dtype=np.float64) # lower triangular matrix + A[:, 0] = velocity_profile * discretization_time + + for idx, velocity in enumerate(velocity_profile): + if idx == 0: + continue + A[idx, 1:] *= velocity * discretization_time**2 + + # Regularization on curvature rate. We add a small but nonzero weight on initial curvature too. + # This is since the corresponding row of the A matrix might be zero if initial speed is 0, leading to singularity. + # We guarantee that Q is positive definite such that the minimizer of the least squares problem is unique. + Q: DoubleMatrix = curvature_rate_penalty * np.eye(len(y)) + Q[0, 0] = initial_curvature_penalty + + # Compute regularized least squares solution. + x = np.linalg.pinv(A.T @ A + Q) @ A.T @ y + + # Extract profile from solution. + initial_curvature = x[0] + curvature_rate_profile = x[1:] + + return initial_curvature, curvature_rate_profile + + +def get_velocity_curvature_profiles_with_derivatives_from_poses( + discretization_time: float, + poses: DoubleMatrix, + jerk_penalty: float, + curvature_rate_penalty: float, +) -> Tuple[DoubleMatrix, DoubleMatrix, DoubleMatrix, DoubleMatrix]: + """ + Main function for joint estimation of velocity, acceleration, curvature, and curvature rate given N poses + sampled at discretization_time. This is done by solving two least squares problems with the given penalty weights. + :param discretization_time: [s] Time discretization used for integration. + :param poses: A trajectory of N poses (x, y, heading). + :param jerk_penalty: A regularization parameter used to penalize acceleration differences. Should be positive. + :param curvature_rate_penalty: A regularization parameter used to penalize curvature_rate. Should be positive. + :return: Profiles for velocity (N-1), acceleration (N-2), curvature (N-1), and curvature rate (N-2). + """ + xy_displacements, heading_displacements = _get_xy_heading_displacements_from_poses(poses) + + # Compute initial velocity + acceleration least squares solution and extract results. + # Note: If we have M displacements, we require the M associated heading values. + # Therefore, we exclude the last heading in the call below. + initial_velocity, acceleration_profile = _fit_initial_velocity_and_acceleration_profile( + xy_displacements=xy_displacements, + heading_profile=poses[:-1, 2], + discretization_time=discretization_time, + jerk_penalty=jerk_penalty, + ) + + velocity_profile = _generate_profile_from_initial_condition_and_derivatives( + initial_condition=initial_velocity, + derivatives=acceleration_profile, + discretization_time=discretization_time, + ) + + # Compute initial curvature + curvature rate least squares solution and extract results. It relies on velocity fit. + initial_curvature, curvature_rate_profile = _fit_initial_curvature_and_curvature_rate_profile( + heading_displacements=heading_displacements, + velocity_profile=velocity_profile, + discretization_time=discretization_time, + curvature_rate_penalty=curvature_rate_penalty, + ) + + curvature_profile = _generate_profile_from_initial_condition_and_derivatives( + initial_condition=initial_curvature, + derivatives=curvature_rate_profile, + discretization_time=discretization_time, + ) + + return velocity_profile, acceleration_profile, curvature_profile, curvature_rate_profile + + + +def complete_kinematic_state_and_inputs_from_poses( + discretization_time: float, + wheel_base: float, + poses: DoubleMatrix, + jerk_penalty: float, + curvature_rate_penalty: float, +) -> Tuple[DoubleMatrix, DoubleMatrix]: + """ + Main function for joint estimation of velocity, acceleration, steering angle, and steering rate given poses + sampled at discretization_time and the vehicle wheelbase parameter for curvature -> steering angle conversion. + One caveat is that we can only determine the first N-1 kinematic states and N-2 kinematic inputs given + N-1 displacement/difference values, so we need to extrapolate to match the length of poses provided. + This is handled by repeating the last input and extrapolating the motion model for the last state. + :param discretization_time: [s] Time discretization used for integration. + :param wheel_base: [m] The wheelbase length for the kinematic bicycle model being used. + :param poses: A trajectory of poses (x, y, heading). + :param jerk_penalty: A regularization parameter used to penalize acceleration differences. Should be positive. + :param curvature_rate_penalty: A regularization parameter used to penalize curvature_rate. Should be positive. + :return: kinematic_states (x, y, heading, velocity, steering_angle) and corresponding + kinematic_inputs (acceleration, steering_rate). + """ + ( + velocity_profile, + acceleration_profile, + curvature_profile, + curvature_rate_profile, + ) = get_velocity_curvature_profiles_with_derivatives_from_poses( + discretization_time=discretization_time, + poses=poses, + jerk_penalty=jerk_penalty, + curvature_rate_penalty=curvature_rate_penalty, + ) + + # Convert to steering angle given the wheelbase parameter. At this point, we don't need to worry about curvature. + steering_angle_profile, steering_rate_profile = _convert_curvature_profile_to_steering_profile( + curvature_profile=curvature_profile, + discretization_time=discretization_time, + wheel_base=wheel_base, + ) + + # Extend input fits with a repeated element and extrapolate state fits to match length of poses. + # This is since we fit with N-1 displacements but still have N poses at the end to deal with. + acceleration_profile = np.append(acceleration_profile, acceleration_profile[-1]) + steering_rate_profile = np.append(steering_rate_profile, steering_rate_profile[-1]) + + velocity_profile = np.append( + velocity_profile, velocity_profile[-1] + acceleration_profile[-1] * discretization_time + ) + steering_angle_profile = np.append( + steering_angle_profile, steering_angle_profile[-1] + steering_rate_profile[-1] * discretization_time + ) + + # Collect completed state and input in matrices. + kinematic_states: DoubleMatrix = np.column_stack((poses, velocity_profile, steering_angle_profile)) + kinematic_inputs: DoubleMatrix = np.column_stack((acceleration_profile, steering_rate_profile)) + + return kinematic_states, kinematic_inputs \ No newline at end of file diff --git a/code/sim/pyproject.toml b/code/sim/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..98223bc89abc472fa161f7f70fc6164f6f52b1a0 --- /dev/null +++ b/code/sim/pyproject.toml @@ -0,0 +1,9 @@ +[tool.hatch.build.targets.wheel] +packages = ["hugsim-env"] + +[project] +name = "hugsim-env" +version = "0.0.1" +dependencies = [ + "gymnasium", +] \ No newline at end of file diff --git a/code/sim/setup.py b/code/sim/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..ddba9e4afb8e6d7df0f7853d9d047626d5839701 --- /dev/null +++ b/code/sim/setup.py @@ -0,0 +1,7 @@ +from setuptools import setup, find_packages + +setup( + name="hugsim-env", + version="0.0.1", + packages=find_packages(), +) \ No newline at end of file diff --git a/code/sim/utils/__pycache__/agent_controller.cpython-311.pyc b/code/sim/utils/__pycache__/agent_controller.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..abf3805f5f92594d9e6f0cfc80e424c2b36eb974 Binary files /dev/null and b/code/sim/utils/__pycache__/agent_controller.cpython-311.pyc differ diff --git a/code/sim/utils/__pycache__/plan.cpython-311.pyc b/code/sim/utils/__pycache__/plan.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0c40ef369432b89c050b9565c8fb69b54a163bfb Binary files /dev/null and b/code/sim/utils/__pycache__/plan.cpython-311.pyc differ diff --git a/code/sim/utils/__pycache__/score_calculator.cpython-311.pyc b/code/sim/utils/__pycache__/score_calculator.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d896a29057922cce26ae726569f412783b6da16 Binary files /dev/null and b/code/sim/utils/__pycache__/score_calculator.cpython-311.pyc differ diff --git a/code/sim/utils/__pycache__/sim_utils.cpython-311.pyc b/code/sim/utils/__pycache__/sim_utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0a57a965714f8c366672bcb4f6d54e88497c0023 Binary files /dev/null and b/code/sim/utils/__pycache__/sim_utils.cpython-311.pyc differ diff --git a/code/sim/utils/agent_controller.py b/code/sim/utils/agent_controller.py new file mode 100644 index 0000000000000000000000000000000000000000..a7ff99354ae331d8fcde95a447630a327d69ad3d --- /dev/null +++ b/code/sim/utils/agent_controller.py @@ -0,0 +1,323 @@ +import math +import random +import numpy as np +from trajdata.maps import VectorMap +from submodules.Pplan.Sampling.spline_planner import SplinePlanner +import torch +import time +import math +from copy import deepcopy +from utils.dynamic_utils import unicycle + + +def constant_tracking(state, path, dt): + ''' + Args: + state: current state of the vehicle, of size [x, y, yaw, speed] + path: the path to follow, of size (N, [x, y, yaw]) + dt: time duration + ''' + + # find the nearest point in the path + dists = torch.norm(path[:, :2] - state[None, :2], dim=1) + nearest_index = torch.argmin(dists) + + # find the target point + lookahead_distance = state[3] * dt + target = path[-1] + is_end = True + for i in range(nearest_index + 1, len(path)): + if torch.norm(path[i, :2] - state[:2]) > lookahead_distance: + target = path[i] + is_end = False + break + + # compute the new state + target_distance = torch.norm(target[:2] - state[:2]) + ratio = lookahead_distance / target_distance.clamp(min=1e-6) + ratio = ratio.clamp(max=1.0) + + new_state = deepcopy(state) + new_state[:2] = state[:2] + ratio * (target[:2] - state[:2]) + new_state[2] = torch.atan2( + state[2].sin() + ratio * (target[2].sin() - state[2].sin()), + state[2].cos() + ratio * (target[2].cos() - state[2].cos()) + ) + if is_end: + new_state[3] = 0 + + return new_state + + +def constant_headaway(states, num_steps, dt): + ''' + Args: + states: current states of a batch of vehicles, of size (num_agents, [x, y, yaw, speed]) + num_steps: number of steps to move forward + dt: time duration + Return: + trajs: the trajectories of the vehicles, of size (num_agents, num_steps, [x, y, yaw, speed]) + ''' + + # state: [x, y, yaw, speed] + x = states[:, 0] + y = states[:, 1] + yaw = states[:, 2] + speed = states[:, 3] + + # Generate time steps + t_steps = torch.arange(num_steps) * dt + + # Calculate dx and dy for each step + dx = torch.outer(speed * torch.sin(yaw), t_steps) + dy = torch.outer(speed * torch.cos(yaw), t_steps) + + # Update x and y positions + x_traj = x.unsqueeze(1) + dx + y_traj = y.unsqueeze(1) + dy + + # Replicate the yaw and speed for each time step + yaw_traj = yaw.unsqueeze(1).repeat(1, num_steps) + speed_traj = speed.unsqueeze(1).repeat(1, num_steps) + + # Stack the x, y, yaw, and speed components to form the trajectory + trajs = torch.stack((x_traj, y_traj, yaw_traj, speed_traj), dim=-1) + + return trajs + + +class IDM: + def __init__( + self, v0=30.0, s0=5.0, T=2.0, a=2.0, b=4.0, delta=4.0, + lookahead_path_length=100, lead_distance_threshold=1.0 + ): + ''' + Args: + v0: desired speed + s0: minimum gap + T: safe time headway + a: max acceleration + b: comfortable deceleration + delta: acceleration exponent + lookahead_path_length: the length of path to look ahead + lead_distance_threshold: the distance to consider a vehicle as a lead vehicle + ''' + self.v0 = v0 + self.s0 = s0 + self.T = T + self.a = a + self.b = b + self.delta = delta + self.lookahead_path_length = lookahead_path_length + self.lead_distance_threshold = lead_distance_threshold + + def update(self, state, path, dt, neighbors): + ''' + Args: + state: current state of the vehicle, of size [x, y, yaw, speed] + path: the path to follow, of size (N, [x, y, yaw]) + dt: time duration + neighbors: the future states of the neighbors, of size (K, T, [x, y, yaw, speed]) + ''' + + if path is None: + return deepcopy(state) + + # find the nearest point in the path + dists = torch.norm(path[:, :2] - state[None, :2], dim=1) + nearest_index = torch.argmin(dists) + + + # lookahead_distance = state[3] * dt + # lookahead_targe = state[:2] + np.array([np.sin(state[2]) * lookahead_distance, np.cos(state[2]) * lookahead_distance]) + # # target = path[-1] + # is_end = False + # target_idx = torch.argmin(torch.norm(path[:, :2] - lookahead_targe, dim=-1)) + # target = path[target_idx] + + # find the target point + lookahead_distance = state[3] * dt + target = path[-1] + is_end = True + for i in range(nearest_index + 1, len(path)): + if torch.norm(path[i, :2] - state[:2]) > lookahead_distance: + target = path[i] + is_end = False + break + + # distance between neighbors and the path + lookahead_path = path[nearest_index + 1:][:self.lookahead_path_length] + lookahead_neighbors = neighbors[..., None, :].expand( + -1, -1, lookahead_path.shape[0], -1 + ) # (K, T, n, 4) + + dists_neighbors = torch.norm( + lookahead_neighbors[..., :2] - lookahead_path[None, None, :, :2], dim=-1 + ) # (K, T, n) + indices_neighbors = torch.arange( + lookahead_path.shape[0] + )[None, None].expand_as(dists_neighbors) + + # determine lead vehicles + is_lead = (dists_neighbors < self.lead_distance_threshold) + if is_lead.any(): + # compute lead distance + indices_lead = indices_neighbors[is_lead] # (num_lead) + lookahead_lengths = torch.cumsum(torch.norm( + lookahead_path[1:, :2] - lookahead_path[:-1, :2], dim=1 + ), dim=0) + lookahead_lengths = torch.cat([lookahead_lengths, lookahead_lengths[-1:]]) + lead_distance = lookahead_lengths[indices_lead] + + # compute lead speed + states_lead = lookahead_neighbors[is_lead] # (num_lead, 4) + ori_speed_lead = states_lead[:, 3] + yaw_lead = states_lead[:, 2] + yaw_path = lookahead_path[indices_lead, 2] + lead_speed = ori_speed_lead * (yaw_lead - yaw_path).cos() + + # compute acceleration + ego_speed = state[3] + delta_v = ego_speed - lead_speed + s_star = self.s0 + \ + (ego_speed * self.T + ego_speed * delta_v / (2 * math.sqrt(self.a * self.b))).clamp(min=0) + acceleration = self.a * (1 - (ego_speed / self.v0) ** self.delta - (s_star / lead_distance) ** 2) + acceleration = acceleration.min() + else: + acceleration = self.a * (1 - (state[3] / self.v0) ** self.delta) + + # compute the new state + target_distance = torch.norm(target[:2] - state[:2]) + ratio = lookahead_distance / target_distance.clamp(min=1e-6) + ratio = ratio.clamp(max=1.0) + + new_state = deepcopy(state) + new_state[:2] = state[:2] + ratio * (target[:2] - state[:2]) + new_state[2] = torch.atan2( + state[2].sin() + ratio * (target[2].sin() - state[2].sin()), + state[2].cos() + ratio * (target[2].cos() - state[2].cos()) + ) + if is_end: + new_state[3] = 0 + else: + new_state[3] = (state[3] + acceleration * dt).clamp(min=0) + + return new_state + + +class AttackPlanner: + def __init__(self, pred_steps=20, ATTACK_FREQ = 3, best_k=1, device='cpu'): + self.device = device + self.predict_steps = pred_steps + self.best_k = best_k + + self.planner = SplinePlanner( + device, + N_seg=self.predict_steps, + acce_grid=torch.linspace(-2, 5, 10).to(self.device), + acce_bound=[-6, 5], + vbound=[-2, 50] + ) + self.planner.psi_bound = [-math.pi * 2, math.pi * 2] + + self.exec_traj = None + self.exec_pointer = 1 + + def update( + self, state, unified_map, dt, + neighbors, attacked_states, + new_plan=True + ): + ''' + Args: + state: current state of the vehicle, of size [x, y, yaw, speed] + vector_map: the vector map + attacked_states: future states of the attacked agent, of size (T, [x, y, yaw, speed]) + neighbors: future states of the neighbors, of size (K, T, [x, y, yaw, speed]) + new_plan: whether to generate a new plan + ''' + assert self.exec_pointer > 0 + + # directly execute the current plan + if not new_plan: + if self.exec_traj is not None and \ + self.exec_pointer < self.exec_traj.shape[0]: + next_state = self.exec_traj[self.exec_pointer] + self.exec_pointer += 1 + return next_state + else: + new_plan = True + + assert attacked_states.shape[0] == self.predict_steps + + # state: [x, y, yaw, speed] + x, y, yaw, speed = state + + # query vector map to get lanes + query_xyzr = np.array([x, y, 0, yaw + np.pi / 2]) + # query_xyzr = unified_map.xyzr_local2world(np.array([x, y, 0, yaw])) + # lanes = unified_map.vector_map.get_lanes_within(query_xyzr[:3], dist=30) + # lanes = [unified_map.batch_xyzr_world2local(l.center.xyzh)[:, [0,1,3]] for l in lanes] + # lanes = [l.center.xyzh[:, [0,1,3]] for l in lanes] + lanes = None + + # for lane in lanes: + # plt.plot(lane[:, 0], lane[:, 1], 'k--', linewidth=0.5, alpha=0.5) + + # generate spline trajectories + x0 = torch.tensor([query_xyzr[0], query_xyzr[1], speed, query_xyzr[3]], device=self.device) + possible_trajs, xf_set = self.planner.gen_trajectories(x0, self.predict_steps * dt, lanes, + dyn_filter=True) # (num_trajs, T-1, [x, y, v, a, yaw, r, t]) + if possible_trajs.shape[0] == 0: + trajs = constant_headaway(state[None], self.predict_steps, dt) # (1, T, [x, y, yaw, speed]) + else: + trajs = torch.cat([ + state[None, None].expand(possible_trajs.shape[0], -1, -1), + possible_trajs[..., [0, 1, 4, 2]] + ], dim=1) + + # select the best trajectory + attack_distance = torch.norm(attacked_states[None, :, :2] - trajs[..., :2], dim=-1) + cost_attack = attack_distance.min(dim=1).values + cost_collision = ( + torch.norm(neighbors[None, ..., :2] - trajs[:, None, :, :2], dim=-1).min(dim=-1).values < 2.0).sum( + dim=-1) + cost = cost_attack + 0.1 * cost_collision + values, indices = torch.topk(cost, self.best_k, largest=False) + random_index = torch.randint(0, self.best_k, (1,)).item() + selected_index = indices[random_index] + traj_best = trajs[selected_index] + + # produce next state + self.exec_traj = traj_best + self.exec_traj[:, 2] -= np.pi / 2 + self.exec_pointer = 1 + next_state = self.exec_traj[self.exec_pointer] + # next_state[0] = -next_state[0] + self.exec_pointer += 1 + + return next_state + + +class ConstantPlanner: + def __init__(self): + return + + def update(self, state, dt): + a, b, yaw, v = state + a = a - v * np.sin(yaw) * dt + b = b + v * np.cos(yaw) * dt + return torch.tensor([a, b, yaw, v]) + + +class UnicyclePlanner: + def __init__(self, uc_path, speed=1.0): + self.uc_model = unicycle.restore(torch.load(uc_path, weights_only=False)) + self.t = 0 + self.speed = speed + + def update(self, dt): + self.t += dt * self.speed + a, b, v, pitchroll, yaw, h = self.uc_model.forward(self.t) + # return torch.tensor([a, b, yaw, v]), pitchroll.detach().cpu(), h.item() + return torch.tensor([a, b, yaw, v]) diff --git a/code/sim/utils/launch_ad.py b/code/sim/utils/launch_ad.py new file mode 100644 index 0000000000000000000000000000000000000000..b0a96c09b146a84b07ef739b5fc203beca864b03 --- /dev/null +++ b/code/sim/utils/launch_ad.py @@ -0,0 +1,30 @@ +import subprocess +import time +import os + + +def launch(shell_path, cuda_id, output): + os.makedirs(output, exist_ok=True) + print(os.path.join(output, 'output.txt')) + print(shell_path, cuda_id, output) + with open(os.path.join(output, 'output.txt'), 'w') as f: + process = subprocess.Popen( + ["zsh", shell_path, cuda_id, output], stdout=f, stderr=f + ) + return process + + +def check_alive(process, tolerant=100): + i = 0 + while i < tolerant: + return_code = process.poll() + if return_code is not None: + print(f"The AD algorithm completed with return code {return_code}.") + process.kill() + return + elif i % 5 == 0: + print(f"The AD algorithm is still running, remaining tolerant {tolerant - i}.") + time.sleep(1) + i += 1 + process.kill() + print("The AD algorithm process is killed.") \ No newline at end of file diff --git a/code/sim/utils/plan.py b/code/sim/utils/plan.py new file mode 100644 index 0000000000000000000000000000000000000000..1f8c0de276df47b052fca9f12a030eb54445f821 --- /dev/null +++ b/code/sim/utils/plan.py @@ -0,0 +1,238 @@ +import numpy as np +import torch +from scipy.spatial.transform import Rotation as SCR +import roma +from collections import namedtuple +from sim.utils.agent_controller import constant_headaway +from sim.utils import agent_controller +from collections import defaultdict +from trajdata import AgentType, UnifiedDataset +from trajdata.maps import MapAPI +from trajdata.simulation import SimulationScene +from sim.utils.sim_utils import rt2pose, pose2rt +from sim.utils.agent_controller import IDM, AttackPlanner, ConstantPlanner, UnicyclePlanner +import os +import json + +Model = namedtuple('Models', ['model_path', 'controller', 'controller_args']) + + +class planner: + def __init__(self, plan_list, scene_path=None, dt=0.2, unified_map=None, ground=None): + self.unified_map = unified_map + self.ground = ground + self.PREDICT_STEPS = 20 + self.NUM_NEIGHBORS = 3 + + self.rectify_angle = 0 + if self.unified_map is not None: + self.rectify_angle = self.unified_map.rectify_angle + + # plan_list: a, b, height, yaw, v, model_path, controller, controller_args: dict + self.stats, self.route, self.controller, self.ckpts, self.wlhs = {}, {}, {}, {}, {} + self.dt = dt + self.ATTACK_FREQ = 3 + for iid, args in enumerate(plan_list): + if args[6] == "UnicyclePlanner": + # self.ckpts[f"agent_{iid}"] = os.path.join(scene_path, "ckpts", f"dynamic_{args[7]}_chkpnt30000.pth") + # self.wlhs[f'agent_{iid}'] = [2.0, 4.0, 1.5] + self.ckpts[f'agent_{iid}'] = os.path.join(args[5], 'gs.pth') + with open(os.path.join(args[5], 'wlh.json')) as f: + self.wlhs[f'agent_{iid}'] = json.load(f) + uc_configs = args[7] + self.controller[f"agent_{iid}"] = UnicyclePlanner(os.path.join(scene_path, f"unicycle_{uc_configs['uc_id']}.pth"), speed=uc_configs['speed']) + a, b, v, pitchroll, yaw, h = self.controller[f"agent_{iid}"].uc_model.forward(0.0) + self.stats[f'agent_{iid}'] = torch.tensor([a, b, args[2], yaw, v]) + self.route[f'agent_{iid}'] = None + else: + model = Model(*args[5:]) + self.stats[f'agent_{iid}'] = torch.tensor(args[:5]) # a, b, height, yaw, v + self.stats[f'agent_{iid}'][3] += self.rectify_angle + self.route[f'agent_{iid}'] = None + self.ckpts[f'agent_{iid}'] = os.path.join(model.model_path, 'gs.pth') + with open(os.path.join(model.model_path, 'wlh.json')) as f: + self.wlhs[f'agent_{iid}'] = json.load(f) + self.controller[f'agent_{iid}'] = getattr(agent_controller, model.controller)(**model.controller_args) + if model.controller == "AttackPlanner": + self.ATTACK_FREQ = model.controller_args["ATTACK_FREQ"] + + def update_ground(self, ground): + self.ground = ground + + def update_agent_route(self): + assert self.unified_map is not None, "Map shouldn't be None to forecast agent path" + for iid, stat in self.stats.items(): + path = self.unified_map.get_route(stat) + if path is None: + print("path not found at ", self.stats) + if path is not None: + self.route[iid] = torch.from_numpy(np.hstack([path[:, :2], path[:, -1:]])) + + def ground_height(self, u, v): + cam_poses, cam_height, _ = self.ground + cam_poses = torch.from_numpy(cam_poses) + cam_dist = np.sqrt( + (cam_poses[:-1, 0, 3] - u) ** 2 + (cam_poses[:-1, 2, 3] - v) ** 2 + ) + nearest_cam_idx = np.argmin(cam_dist, axis=0) + nearest_c2w = cam_poses[nearest_cam_idx] + + nearest_w2c = np.linalg.inv(nearest_c2w) + uv_local = nearest_w2c[:3, :3] @ np.array([u, 0, v]) + nearest_w2c[:3, 3] + uv_local[1] = 0 + uv_world = nearest_c2w[:3, :3] @ uv_local + nearest_c2w[:3, 3] + + return uv_world[1] + cam_height + + def plan_traj(self, t, ego_stats): + all_stats = [ego_stats] + for iid, stat in self.stats.items(): + all_stats.append(stat[[0, 1, 3, 4]]) # a, b, yaw, v + all_stats = torch.stack(all_stats, dim=0) + future_states = constant_headaway(all_stats, num_steps=self.PREDICT_STEPS, dt=self.dt) + + b2ws = {} + for iid, stat in self.stats.items(): + # find closet neighbors + curr_xy_agents = all_stats[:, :2] + distance_agents = torch.norm(curr_xy_agents - stat[:2], dim=-1) + neighbor_idx = torch.argsort(distance_agents)[1:self.NUM_NEIGHBORS + 1] + neighbors = future_states[neighbor_idx] + + controller = self.controller[iid] + if type(controller) is IDM: + next_xyrv = controller.update(state=stat[[0, 1, 3, 4]], path=self.route[iid], dt=self.dt, + neighbors=neighbors) + elif type(controller) is AttackPlanner: + safe_neighbors = neighbors[1:, ...] + next_xyrv = controller.update(state=stat[[0, 1, 3, 4]], unified_map=self.unified_map, dt=0.1, + neighbors=safe_neighbors, attacked_states=future_states[0], + new_plan=((t // self.dt) % self.ATTACK_FREQ == 0)) + elif type(controller) is ConstantPlanner: + next_xyrv = controller.update(state=stat[[0, 1, 3, 4]], dt=self.dt) + elif type(controller) is UnicyclePlanner: + next_xyrv = controller.update(dt=self.dt) + else: + raise NotImplementedError + next_stat = torch.zeros_like(stat) + next_stat[[0, 1, 3, 4]] = next_xyrv.float() + next_stat[2] = stat[2] + self.stats[iid] = next_stat + + b2w = np.eye(4) + h = self.ground_height(next_xyrv[0].numpy(), next_xyrv[1].numpy()) + if type(controller) is UnicyclePlanner: + # b2w[:3, :3] = SCR.from_euler('xzy', [pitch_roll[0], pitch_roll[1], stat[3]]).as_matrix() + b2w[:3, :3] = SCR.from_euler('y', [-stat[3]]).as_matrix() + b2w[:3, 3] = np.array([next_stat[0], h + stat[2], next_stat[1]]) + else: + b2w[:3, :3] = SCR.from_euler('y', [-stat[3] - np.pi / 2 - self.rectify_angle]).as_matrix() + b2w[:3, 3] = np.array([next_stat[0], h + stat[2], next_stat[1]]) + b2ws[iid] = torch.tensor(b2w).float().cuda() + + return [b2ws, {}] + + +class UnifiedMap: + def __init__(self, datapath, version, scene_name): + self.datapath = datapath + self.version = version + + self.dataset = UnifiedDataset( + desired_data=[self.version], + data_dirs={ + self.version: self.datapath, + }, + cache_location="/app/app_datas/nusc_map_cache", + only_types=[AgentType.VEHICLE], + agent_interaction_distances=defaultdict(lambda: 50.0), + desired_dt=0.1, + num_workers=4, + verbose=True, + ) + + self.map_api = MapAPI(self.dataset.cache_path) + + self.scene = None + for scene in list(self.dataset.scenes()): + if scene.name == scene_name: + self.scene = scene + assert self.scene is not None, f"Can't find scene {scene_name}" + self.vector_map = self.map_api.get_map( + f"{self.version}:{self.scene.location}" + ) + self.ego_start_pos, self.ego_start_yaw = self.get_start_pose() + self.rectify_angle = 0 + if self.ego_start_yaw < 0: + self.ego_start_yaw += np.pi + self.rectify_angle = np.pi + self.PATH_LENGTH = 100 + + def get_start_pose(self): + sim_scene: SimulationScene = SimulationScene( + env_name=self.version, + scene_name=f"sim_scene", + scene=self.scene, + dataset=self.dataset, + init_timestep=0, + freeze_agents=True, + ) + obs = sim_scene.reset() + assert obs.agent_name[0] == 'ego', 'The first agent is not ego' + # We consider position of the first ego frame as origin + # This suppose is ok when the first frame front camera pose is set as origin + ego_start_pos = obs.curr_agent_state.position[0] + ego_start_yaw = obs.curr_agent_state.heading[0] + return ego_start_pos.numpy(), ego_start_yaw.item() + + def xyzr_local2world(self, stat): + alpha = np.arctan(stat[0] / stat[1]) + beta = self.ego_start_yaw - alpha + dist = np.linalg.norm(stat[:2]) + delta_x = dist * np.cos(beta) + delta_y = dist * np.sin(beta) + + world_stat = np.zeros(4) + world_stat[0] = delta_x + self.ego_start_pos[0] + world_stat[1] = delta_y + self.ego_start_pos[1] + world_stat[3] = stat[3] + self.ego_start_yaw + + return world_stat + + def batch_xyzr_world2local(self, stat): + beta = np.arctan((stat[:, 1] - self.ego_start_pos[1]) / (stat[:, 0] - self.ego_start_pos[0])) + alpha = self.ego_start_yaw - beta + dist = np.linalg.norm(stat[:, :2] - self.ego_start_pos, axis=1) + delta_x = dist * np.sin(alpha) + delta_y = dist * np.cos(alpha) + + local_stat = np.zeros_like(stat) + local_stat[:, 0] = delta_x + local_stat[:, 1] = delta_y + local_stat[:, 3] = stat[:, 3] - self.ego_start_yaw + + return local_stat + + def get_route(self, stat): + # stat: a, b, height, yaw, v + curr_xyzr = self.xyzr_local2world(stat[:4].numpy()) + + # lanes = self.vector_map.get_current_lane(curr_xyzr, max_dist=5, max_heading_error=np.pi/3) + lanes = self.vector_map.get_current_lane(curr_xyzr) + + if len(lanes) > 0: + curr_lane = lanes[0] + path = self.batch_xyzr_world2local(curr_lane.center.xyzh) + total_path_length = np.linalg.norm(curr_lane.center.xy[1:] - curr_lane.center.xy[:-1], axis=1).sum() + # random select next lanes until reach PATH_LENGTH + while total_path_length < self.PATH_LENGTH: + next_lanes = list(curr_lane.next_lanes) + if len(next_lanes) == 0: + break + next_lane = self.vector_map.get_road_lane(next_lanes[np.random.randint(len(next_lanes))]) + path = np.vstack([path, self.batch_xyzr_world2local(next_lane.center.xyzh)]) + total_path_length += np.linalg.norm(next_lane.center.xy[1:] - next_lane.center.xy[:-1], axis=1).sum() + curr_lane = next_lane + else: + path = None + return path diff --git a/code/sim/utils/score_calculator.py b/code/sim/utils/score_calculator.py new file mode 100644 index 0000000000000000000000000000000000000000..d93c5ae079f7a4be62fb71851d9c18e1ccd2aaf8 --- /dev/null +++ b/code/sim/utils/score_calculator.py @@ -0,0 +1,562 @@ +import pickle +from matplotlib import pyplot as plt +from matplotlib.patches import Rectangle, Polygon +from shapely.geometry import LineString, Point +import numpy as np +from shapely.geometry import Polygon as ShapelyPolygon +from shapely.geometry import Point +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor +import threading +import argparse +import os +import open3d as o3d +import torch +from scipy.spatial.transform import Rotation as SCR + +ego_verts_canonic = np.array([[0.5, 0.5, 0], [0.5, -0.5, 0], [0.5, 0.5, 1.0], [0.5, -0.5, 1.0], + [-0.5, -0.5, 0], [-0.5, 0.5, 0], [-0.5, -0.5, 1.0], [-0.5, 0.5, 1.0]]) + +# Define boundaries +boundaries = { + 'max_abs_lat_accel': 4.89, # [m/s^2] + 'max_lon_accel': 2.40, # [m/s^2] + 'min_lon_accel': -4.05, # [m/s^2] + 'max_abs_yaw_accel': 1.93, # [rad/s^2] + 'max_abs_lon_jerk': 8.37, # [m/s^3], + 'max_abs_yaw_rate': 0.95, # [rad/s] +} + +score_weight = { + 'ttc': 5, + 'c': 2, + 'ep': 5, +} + +def create_rectangle(center_x, center_y, width, length, yaw): + """Create a rectangle polygon.""" + cos_yaw = np.cos(yaw) + sin_yaw = np.sin(yaw) + + x_offs = [length/2, length/2, -length/2, -length/2] + y_offs = [width/2, -width/2, -width/2, width/2] + + x_pts = [center_x + x_off*cos_yaw - y_off * + sin_yaw for x_off, y_off in zip(x_offs, y_offs)] + y_pts = [center_y + x_off*sin_yaw + y_off * + cos_yaw for x_off, y_off in zip(x_offs, y_offs)] + + return ShapelyPolygon(zip(x_pts, y_pts)) + +def bg_collision_det(points, box): + O, A, B, C = box[0], box[1], box[2], box[5] + OA = A - O + OB = B - O + OC = C - O + POA, POB, POC = (points @ OA[..., None])[:, 0], (points @ OB[..., None])[:, 0], (points @ OC[..., None])[:, 0] + mask = (torch.dot(O, OA) < POA) & (POA < torch.dot(A, OA)) & \ + (torch.dot(O, OB) < POB) & (POB < torch.dot(B, OB)) & \ + (torch.dot(O, OC) < POC) & (POC < torch.dot(C, OC)) + return True if torch.sum(mask) > 100 else False + + +class ScoreCalculator: + def __init__(self, data): + self.data = data + + self.pdms = 0.0 + self.driving_score = None + pass + + def transform_to_ego_frame(self, traj, ego_box): + """ + Transform trajectory from global frame to ego-centric frame. + + :param traj: List of tuples (x, y, yaw) in global frame + :param ego_box: Tuple (x, y, z, w, l, h, yaw) of ego vehicle in global frame + :return: Numpy array of transformed trajectory + """ + ego_x, ego_y, _, _, _, _, ego_yaw = ego_box + + # Create rotation matrix + c, s = np.cos(-ego_yaw), np.sin(-ego_yaw) + R = np.array([[c, -s], [s, c]]) + + # Transform each point + transformed_traj = [] + for x, y, yaw in traj: + # Translate + x_translated, y_translated = x - ego_x, y - ego_y + + # Rotate + x_rotated, y_rotated = R @ np.array([x_translated, y_translated]) + + # Adjust yaw + yaw_adjusted = yaw - ego_yaw + + transformed_traj.append((x_rotated, y_rotated, yaw_adjusted)) + + return np.array(transformed_traj) + + def get_vehicle_corners(self, x, y, yaw, length, width): + """ + Calculate the corner points of the vehicle given its position, orientation, and dimensions. + + :param x: x-coordinate of the vehicle's center + :param y: y-coordinate of the vehicle's center + :param yaw: orientation of the vehicle in radians + :param length: length of the vehicle + :param width: width of the vehicle + :return: numpy array of corner coordinates (4x2) + """ + c, s = np.cos(yaw), np.sin(yaw) + front_left = np.array([x + c * length / 2 - s * width / 2, + y + s * length / 2 + c * width / 2]) + front_right = np.array([x + c * length / 2 + s * width / 2, + y + s * length / 2 - c * width / 2]) + rear_left = np.array([x - c * length / 2 - s * width / 2, + y - s * length / 2 + c * width / 2]) + rear_right = np.array([x - c * length / 2 + s * width / 2, + y - s * length / 2 - c * width / 2]) + return np.array([front_left, front_right, rear_right, rear_left]) + + def plot_trajectory_on_drivable_mask(self, drivable_mask, transformed_traj, vehicle_width, vehicle_length): + """ + Plot the transformed trajectory and vehicle bounding boxes on the drivable mask. + + :param drivable_mask: 2D numpy array representing the drivable area (200x200) + :param transformed_traj: Numpy array of transformed trajectory points + :param vehicle_width: Width of the vehicle in meters + :param vehicle_length: Length of the vehicle in meters + """ + + plt.figure(figsize=(10, 10)) + plt.imshow(drivable_mask, cmap='gray', extent=[-50, 50, -50, 50]) + + # Scale factor (200 pixels represent 100 meters) + scale_factor = 200 / 100 # pixels per meter + + # Plot trajectory + x_coords, y_coords, yaws = transformed_traj.T + plt.plot(x_coords, y_coords, 'r-', linewidth=2) + + # Plot vehicle bounding boxes + for x, y, yaw in transformed_traj: + corners = self.get_vehicle_corners( + x, y, yaw, vehicle_length, vehicle_width) + plt.gca().add_patch(Polygon(corners, fill=False, edgecolor='blue')) + + # Plot start and end points + plt.plot(x_coords[0], y_coords[0], 'go', markersize=10, label='Start') + plt.plot(x_coords[-1], y_coords[-1], 'bo', markersize=10, label='End') + + plt.title('Trajectory and Vehicle Bounding Boxes on Drivable Mask') + plt.legend() + plt.xlabel('x (meters)') + plt.ylabel('y (meters)') + plt.grid(True) + plt.tight_layout() + plt.show() + + def _calculate_drivable_area_compliance(self, ground, traj, vehicle_width, vehicle_length): + m, n = 2, 2 + dac = 1.0 + for traj_i, (x, y, yaw) in enumerate(traj): + cnt = 0 + c, s = np.cos(yaw), np.sin(yaw) + R = np.array([[c, -s], [s, c]]) + ground_in_ego = (np.linalg.inv(R) @ (ground + np.array([-x, -y])).T).T + x_bins = np.linspace(-vehicle_length/2, vehicle_length/2, m+1) + y_bins = np.linspace(-vehicle_width/2, vehicle_width/2, n+1) + for xi in range(m): + for yi in range(n): + min_x, max_x = x_bins[xi], x_bins[xi+1] + min_y, max_y = y_bins[yi], y_bins[yi+1] + ground_mask = (min_x < ground_in_ego[:, 0]) & (ground_in_ego[:, 0] < max_x) & \ + (min_y < ground_in_ego[:, 1]) & (ground_in_ego[:, 1] < max_y) + if ground_mask.sum() > 0: + cnt += 1 + drivable_ratio = cnt / (m*n) + if drivable_ratio < 0.3: + return 0 + elif drivable_ratio < 0.5: + dac = 0.5 + return dac + + def _calculate_progress(self, planned_traj, ref_taj): + def calculate_curve_length(points): + """Calculate the total length of a curve given by a set of points.""" + curve = LineString(points) + return curve.length + + def project_curve_onto_curve(curve_a, curve_b): + """Project curve_b onto curve_a and calculate the projected length.""" + projected_points = [] + for point in curve_b.coords: + projected_point = curve_a.interpolate( + curve_a.project(Point(point))) + projected_points.append(projected_point) + projected_curve = LineString(projected_points) + return projected_curve.length + + # Create Shapely LineString objects + plan_curve = LineString([(x, y) for x, y, _ in planned_traj]) + ref_curve = LineString([(x, y) for x, y, _ in ref_taj]) + + # Calculate lengths + plan_curve_length = calculate_curve_length(plan_curve) + ref_curve_length = calculate_curve_length(ref_curve) + projected_length = project_curve_onto_curve(ref_curve, plan_curve) + # print(f"plan_curve_length: {plan_curve_length}, ref_curve_length: {ref_curve_length}, project plan to ref_length: {projected_length}") + + ep = 0.0 + if max(plan_curve_length, ref_curve_length) < 5.0 or ref_curve_length < 1e-6: + ep = 1.0 + else: + ep = projected_length / ref_curve_length + return ep + + def _calculate_is_comfortable(self, traj, timestep): + """ + Check if all kinematic parameters of a trajectory are within specified boundaries. + + :param traj: List of tuples (x, y, yaw) representing the trajectory, in ego's local frame + :param timestep: Time interval between trajectory points in seconds + :return: 1.0 if all parameters are within boundaries, 0.0 otherwise + """ + + def calculate_trajectory_kinematics(traj, timestep): + """ + Calculate kinematic parameters for a given trajectory. + + :param traj: List of tuples (x, y, yaw) for each point in the trajectory + :param timestep: Time interval between each point in the trajectory + :return: Dictionary containing lists of calculated parameters + """ + # Convert trajectory to numpy array for easier calculations + x, y, yaw = zip(*traj) + x, y, yaw = np.array(x), np.array(y), np.array(yaw) + + # Calculate velocities + dx = np.diff(x) / timestep + dy = np.diff(y) / timestep + + # Calculate yaw rate + dyaw = np.diff(yaw) + dyaw = np.where(dyaw > np.pi, dyaw - 2*np.pi, dyaw) + dyaw = np.where(dyaw < -np.pi, dyaw + 2*np.pi, dyaw) + dyaw = dyaw / timestep + ddyaw = np.diff(dyaw) / timestep + + # Calculate speed + speed = np.sqrt(dx**2 + dy**2) + + # Calculate accelerations + accel = np.diff(speed) / timestep + jerk = np.diff(accel) / timestep + + # Calculate yaw rate (already calculated as dyaw) + yaw_rate = dyaw + # Calculate yaw acceleration + yaw_accel = ddyaw + + lon_accel = accel + lat_accel = np.zeros_like(lon_accel) + lon_jerk = jerk + + # Pad arrays to match the original trajectory length + yaw_rate = np.pad(yaw_rate, (0, 1), 'edge') + yaw_accel = np.pad(yaw_accel, (0, 2), 'edge') + lon_accel = np.pad(lon_accel, (0, 2), 'edge') + lat_accel = np.pad(lat_accel, (0, 2), 'edge') + lon_jerk = np.pad(lon_jerk, (0, 3), 'edge') + + return { + 'speed': speed, + 'yaw_rate': yaw_rate, + 'yaw_accel': yaw_accel, + 'lon_accel': lon_accel, + 'lat_accel': lat_accel, + 'lon_jerk': lon_jerk, + } + + # Calculate kinematic parameters + if len(traj) < 4: + return 1.0 + + kinematics = calculate_trajectory_kinematics(traj, timestep) + + # Check each parameter against its boundary + checks = [ + np.all(np.abs(kinematics['lat_accel']) <= + boundaries['max_abs_lat_accel']), + np.all(kinematics['lon_accel'] <= boundaries['max_abs_lat_accel']), + np.all(kinematics['lon_accel'] >= boundaries['min_lon_accel']), + np.all(np.abs(kinematics['lon_jerk']) <= + boundaries['max_abs_lon_jerk']), + np.all(np.abs(kinematics['yaw_accel']) <= + boundaries['max_abs_yaw_accel']), + np.all(np.abs(kinematics['yaw_rate']) <= + boundaries['max_abs_yaw_rate']) + ] + + # if not all(checks): + # print(traj) + # print(kinematics) + print(f"comfortable: {checks}") + + # Return 1.0 if all checks pass, 0.0 otherwise + return 1.0 if all(checks) else 0.0 + + def _calculate_no_collision(self, ego_box, planned_traj, obs_lists, scene_xyz): + ego_x, ego_y, z, ego_w, ego_l, ego_h, ego_yaw = ego_box + ego_verts_local = ego_verts_canonic * np.array([ego_l, ego_w, ego_h]) + for idx in range(planned_traj.shape[0]): + ego_x, ego_y, ego_yaw = planned_traj[idx] # ego_state= (x,y,yaw) + ego_trans_mat = np.eye(4) + ego_trans_mat[:3, :3] = SCR.from_euler('z', ego_yaw).as_matrix() + ego_trans_mat[:3, 3] = np.array([ego_x, ego_y, z]) + ego_verts_global = (ego_trans_mat[:3, :3] @ ego_verts_local.T).T + ego_trans_mat[:3, 3] + ego_verts_global = torch.from_numpy(ego_verts_global).float().cuda() + bk_collision = bg_collision_det(scene_xyz, ego_verts_global) + # scene_local = scene_xyz - np.array([ego_x, ego_y, z]) + # bk_collision = np.sum( + # (-ego_l/2 < scene_local[:, 0]) & (scene_local[:, 0] < ego_l/2) & \ + # (-ego_w/2 < scene_local[:, 1]) & (scene_local[:, 1] < ego_w/2) & \ + # (-ego_h/2 < scene_local[:, 2]) & (scene_local[:, 2] < ego_h/2) + # ) > 100 + if bk_collision: + print(f"collision with background detected! @ timestep{idx}") + return 0.0 + ego_poly = create_rectangle(ego_x, ego_y, ego_w, ego_l, ego_yaw) + obs_list = obs_lists[idx if idx < len(obs_lists) else -1] + for obs in obs_list: + # obs = (x,y,z,w,l,h,yaw) + obs_x, obs_y, _, obs_w, obs_l, _, obs_yaw = obs + obs_poly = create_rectangle( + obs_x, obs_y, obs_w, obs_l, obs_yaw) + if ego_poly.intersects(obs_poly): + print(f"collision with obstacle detected! @ timestep{idx}") + print( + f"ego_poly: {(ego_x, ego_y, ego_yaw,obs_w, obs_l)}, obs_poly: {(obs_x, obs_y, obs_yaw,obs_w, obs_l )}") + return 0.0 # Collision detected + return 1.0 + + def _calculate_time_to_collision(self, ego_box, planned_traj, obs_lists, scene_xyz, timestep): + # breakpoint() + t_list = [0.5, 1] # ttc time + + for t in t_list: + # Calculate velocities + velocities = np.diff(planned_traj[:, :2], axis=0) / timestep + + # Use the velocity of the second point for the first point + velocities = np.vstack([velocities[0], velocities]) + + # Calculate the displacement + displacement = velocities * t + + # Create the new trajectory + new_traj = planned_traj.copy() + new_traj[:, :2] += displacement + + is_collide_score = self._calculate_no_collision( + ego_box, new_traj, obs_lists, scene_xyz) + if is_collide_score == 0.0: + print(f" failed to pass ttc collision check, t={t}") + # breakpoint() + return 0.0 + + return 1.0 + + def calculate(self, ): + + print(f"current exp has {len(self.data['frames'])} frames") + if len(self.data['frames']) == 0: + return None + # todo: time_step need modify + score_list = {} + for i in range(0, len(self.data['frames']), 1): + frame = self.data['frames'][i] + if frame['is_key_frame'] == False: + continue + + print(f"frame {i} / {len(self.data['frames'])}") + timestamp = frame['time_stamp'] + planned_last_timestamp = timestamp + \ + len(frame['planned_traj']['traj']) * \ + frame['planned_traj']['timestep'] + ego_x, ego_y, _, ego_w, ego_l, _, ego_yaw = frame['ego_box'] + # frame['planned_traj']['traj'] + if len(frame['planned_traj']['traj'])<2: + continue + traj = frame['planned_traj']['traj'] + + planned_traj = np.concatenate(([np.array([ego_x, ego_y, ego_yaw])], traj), axis=0) + # print(planned_traj) + + # if the car is stopped, there may be error in the yaw of the planned trajectory + traj_distance = np.linalg.norm(planned_traj[-1, :2] - planned_traj[0, :2] ) + if traj_distance<1: + planned_traj[:, 2] = planned_traj[0, 2] # set all yaw to the first yaw + + current_timestamp = timestamp + current_frame_idx = i + obs_lists = [] + while current_timestamp <= planned_last_timestamp+1e-5: + if abs(current_timestamp - self.data['frames'][current_frame_idx]['time_stamp']) < 1e-5: + obs_list = [] + for idx, obj in enumerate(self.data['frames'][current_frame_idx]['obj_boxes']): + # obs_list.append(obj) + if self.data['frames'][current_frame_idx]['obj_names'][idx] == 'car': + obs_list.append(obj) + obs_lists.append(obs_list) + current_timestamp += frame['planned_traj']['timestep'] + + current_frame_idx += 1 + if current_frame_idx >= len(self.data['frames']): + break + + # breakpoint() + # plt.imshow(frame['drivable_mask'].astype(np.uint8)) + # plt.show() + # transformed_traj = self.transform_to_ego_frame(frame['planned_traj']['traj'], frame['ego_box']) + transformed_traj = self.transform_to_ego_frame( + planned_traj, frame['ego_box']) + # breakpoint() + + score_nc = self._calculate_no_collision( + frame['ego_box'], planned_traj, obs_lists, self.data['scene_xyz']) + # score_nc = 0.0 if frame['collision'] else 1.0 + score_dac = self._calculate_drivable_area_compliance( + self.data['ground_xy'], planned_traj, ego_w, ego_l) + score_ttc = self._calculate_time_to_collision( + frame['ego_box'], planned_traj, obs_lists, self.data['scene_xyz'], frame['planned_traj']['timestep']) + score_c = self._calculate_is_comfortable( + transformed_traj, frame['planned_traj']['timestep']) + # score_ep = self._calculate_progress( + # planned_traj, ref_traj) + + score_pdms = score_nc*score_dac*(score_weight['ttc']*score_ttc+score_weight['c']*score_c)/( + score_weight['ttc']+score_weight['c']) + print('nc, dac, ttc, com, pdms', [score_nc, score_dac, score_ttc, score_c, score_pdms]) + score_list[timestamp] = {'nc': score_nc, 'dac': score_dac, + 'ttc': score_ttc, 'c': score_c, 'pdms': score_pdms} + + totals = {metric: 0 for metric in next(iter(score_list.values()))} + for scores in score_list.values(): + for metric, value in scores.items(): + totals[metric] += value + + # avg scores + num_entries = len(score_list) + averages = {metric: total / num_entries for metric, + total in totals.items()} + + # writer.writerow(averages.values()) + + mean_score = averages['pdms'] + route_completion = max([f['rc'] for f in self.data['frames']]) + route_completion = route_completion if route_completion < 1 else 1.0 + driving_score = mean_score*route_completion + return mean_score, route_completion, driving_score, averages + + +def calculate(data): + print(f"this pkl file contains {len(data)} experiment records.") + # print(f"the first item metadata is {data[0]['metas']}.") + # breakpoint() + + def process_exp_data(exp_data): + score_calc = ScoreCalculator(exp_data) + score = score_calc.calculate() + print(f"The score of experiment is {score}.") + final_score_dict = score[3] + final_score_dict['rc'] = score[1] + final_score_dict['hdscore'] = score[2] + return final_score_dict + + def multi_threaded_process(data, max_workers=None): + all_averages = [] + + # Using thread locks for thread-safe append operations + lock = threading.Lock() + + def append_result(future): + result = future.result() + with lock: + all_averages.append(result) + + with ThreadPoolExecutor(max_workers=1) as executor: + futures = [executor.submit(process_exp_data, exp_data) + for exp_data in data] + for future in futures: + future.add_done_callback(append_result) + + return all_averages + + all_averages = multi_threaded_process(data) + + collected_values = defaultdict(list) + for averages in all_averages: + for key, value in averages.items(): + collected_values[key].append(value) + + # Calculation of mean and standard deviation for each indicator + results = {} + for key, values in collected_values.items(): + avg = np.mean(values) + # std = np.std(values) + results[key] = f"{avg:.4f}" + + # Output Results + print("=============================Results=============================") + for key, value in results.items(): + print(f"'{key}': {value}") + return results + +def parse_data(test_path): + data_file_name = os.path.join(test_path, "data.pkl") + ground_pcd_file_name = os.path.join(test_path, "ground.ply") + scene_pcd_file_name = os.path.join(test_path, "scene.ply") + + # Open the file and load the data + with open(data_file_name, 'rb') as f: + data = pickle.load(f) + + ground_pcd = o3d.io.read_point_cloud(ground_pcd_file_name) + ground_xyz = np.asarray(ground_pcd.points) # in camera coordinates + ground_xy = np.stack([ground_xyz[:, 2], -ground_xyz[:, 0]], axis=1) # in imu coordinates + + scene_pcd = o3d.io.read_point_cloud(scene_pcd_file_name) + scene_xyz = np.asarray(scene_pcd.points) # in camera coordinates + # in imu coordinates + scene_xyz = np.stack([scene_xyz[:, 2], -scene_xyz[:, 0], -scene_xyz[:, 1]], axis=1) + + data[0]['ground_xy'] = ground_xy + data[0]['scene_xyz'] = torch.from_numpy(scene_xyz).cuda() + # data[0]['scene_xyz'] = scene_xyz + return data + + +def hugsim_evaluate(test_data, ground_xyz, scene_xyz): + ground_xy = np.stack([ground_xyz[:, 2], -ground_xyz[:, 0]], axis=1) # in imu coordinates + scene_xyz = np.stack([scene_xyz[:, 2], -scene_xyz[:, 0], -scene_xyz[:, 1]], axis=1) + test_data[0]['ground_xy'] = ground_xy + test_data[0]['scene_xyz'] = torch.from_numpy(scene_xyz).float().cuda() + results = calculate(test_data) + return results + + +def get_opts(): + parser = argparse.ArgumentParser() + parser.add_argument('--test_path', type=str, required=True) + return parser.parse_args() + + +if __name__ == "__main__": + args = get_opts() + data = parse_data(args.test_path) + + # Call the main function with the loaded data + calculate(data) \ No newline at end of file diff --git a/code/sim/utils/sim_utils.py b/code/sim/utils/sim_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c291f42c39058beb13aaf49b24c2d9475433563b --- /dev/null +++ b/code/sim/utils/sim_utils.py @@ -0,0 +1,122 @@ +import numpy as np +from scipy.spatial.transform import Rotation as SCR +import math +from scene.cameras import Camera +from sim.ilqr.lqr import plan2control +from omegaconf import OmegaConf + +def rt2pose(r, t, degrees=False): + pose = np.eye(4) + pose[:3, :3] = SCR.from_euler('XYZ', r, degrees=degrees).as_matrix() + pose[:3, 3] = t + return pose + +def pose2rt(pose, degrees=False): + r = SCR.from_matrix(pose[:3, :3]).as_euler('XYZ', degrees=degrees) + t = pose[:3, 3] + return r, t + +def load_camera_cfg(cfg): + cam_params = {} + cams = OmegaConf.to_container(cfg.cams, resolve=True) + for cam_name, cam in cams.items(): + v2c = rt2pose(cam['extrinsics']['v2c_rot'], cam['extrinsics']['v2c_trans'], degrees=True) + l2c = rt2pose(cam['extrinsics']['l2c_rot'], cam['extrinsics']['l2c_trans'], degrees=True) + cam_intrin = cam['intrinsics'] + cam_intrin['fovx'] = cam_intrin['fovx'] / 180.0 * np.pi + cam_intrin['fovy'] = cam_intrin['fovy'] / 180.0 * np.pi + cam_params[cam_name] = {'intrinsic': cam_intrin, 'v2c': v2c, 'l2c': l2c} + + rect_mat = np.eye(4) + if 'cam_rect' in cfg: + rect_mat[:3, :3] = SCR.from_euler('XYZ', cfg.cam_rect.rot, degrees=True).as_matrix() + rect_mat[:3, 3] = np.array(cfg.cam_rect.trans) + + return cam_params, OmegaConf.to_container(cfg.cam_align, resolve=True), rect_mat + +def fov2focal(fov, pixels): + return pixels / (2 * math.tan(fov / 2)) + +def focal2fov(focal, pixels): + return 2*math.atan(pixels/(2*focal)) + +def create_cam(intrinsic, c2w): + fovx, fovy = intrinsic['fovx'], intrinsic['fovy'] + h, w = intrinsic['H'], intrinsic['W'] + K = np.eye(4) + K[0, 0], K[1, 1] = fov2focal(fovx, w), fov2focal(fovy, h) + K[0, 2], K[1, 2] = intrinsic['cx'], intrinsic['cy'] + cam = Camera(K=K, c2w=c2w, width=w, height=h, + image=np.zeros((h, w, 3)), image_name='', dynamics={}) + return cam + +def traj2control(plan_traj, info): + """ + The input plan trajectory is under lidar coordinates + x to right, y to forward and z to upward + """ + plan_traj_stats = np.zeros((plan_traj.shape[0]+1, 5)) + plan_traj_stats[1:, :2] = plan_traj[:, [1,0]] + prev_a, prev_b = 0, 0 + for i, (a, b) in enumerate(plan_traj): + rot = np.arctan((b - prev_b)/(a - prev_a)) + plan_traj_stats[i+1, 2] = rot + curr_stat = np.array( + [0, 0, 0, info['ego_velo'], info['ego_steer']] + ) + acc, steer_rate = plan2control(plan_traj_stats, curr_stat) + return acc, steer_rate + +def dense_cam_poses(cam_poses, cmds): + + for i in range(5): + dense_poses = [] + dense_cmds = [] + for i in range(cam_poses.shape[0]-1): + cam1 = cam_poses[i] + cam2 = cam_poses[i+1] + dense_poses.append(cam1) + dense_cmds.append(cmds[i]) + if np.linalg.norm(cam1[:3, 3]-cam2[:3, 3]) > 0.1: + euler1 = SCR.from_matrix(cam1[:3, :3]).as_euler("XYZ") + euler2 = SCR.from_matrix(cam2[:3, :3]).as_euler("XYZ") + interp_euler = (euler1 + euler2) / 2 + interp_trans = (cam1[:3, 3] + cam2[:3, 3]) / 2 + interp_pose = np.eye(4) + interp_pose[:3, :3] = SCR.from_euler("XYZ", interp_euler).as_matrix() + interp_pose[:3, 3] = interp_trans + dense_poses.append(interp_pose) + dense_cmds.append(cmds[i]) + dense_poses.append(cam_poses[-1]) + dense_poses = np.stack(dense_poses) + cam_poses = dense_poses + cmds = dense_cmds + + return cam_poses, cmds + +def traj_transform_to_global(traj, ego_box): + """ + Transform trajectory from ego-centeric frame to global frame + """ + ego_x, ego_y, _, _, _, _, ego_yaw = ego_box + global_points = [ + ( + ego_x + + px * math.cos(ego_yaw) + - py * math.sin(ego_yaw), + ego_y + + px * math.sin(ego_yaw) + + py * math.cos(ego_yaw), + ) + for px, py in traj + ] + global_trajs = [] + for i in range(1, len(global_points)): + x1, y1 = global_points[i - 1] + x2, y2 = global_points[i] + dx, dy = x2 - x1, y2 - y1 + # distance = math.sqrt(dx**2 + dy**2) + yaw = math.atan2(dy, dx) + global_trajs.append((x1, y1, yaw)) + return global_trajs + \ No newline at end of file diff --git a/code/submodules/Pplan/Policy/base.py b/code/submodules/Pplan/Policy/base.py new file mode 100644 index 0000000000000000000000000000000000000000..f55575c68c8a2765dc5d91afb4d203d54e96491f --- /dev/null +++ b/code/submodules/Pplan/Policy/base.py @@ -0,0 +1,16 @@ +import abc + + +class Policy(abc.ABC): + def __init__(self, device, *args, **kwargs): + self.device = device + + @abc.abstractmethod + def get_action(self, obs_dict, **kwargs): + """Predict an action based on the input observation """ + pass + + @abc.abstractmethod + def eval(self): + """Set the policy to evaluation mode""" + pass \ No newline at end of file diff --git a/code/submodules/Pplan/Policy/sampling_planner.py b/code/submodules/Pplan/Policy/sampling_planner.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/code/submodules/Pplan/Sampling/__init__.py b/code/submodules/Pplan/Sampling/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/code/submodules/Pplan/Sampling/__pycache__/__init__.cpython-311.pyc b/code/submodules/Pplan/Sampling/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8434af456fd17f1005fa474b0d1ea057e8952564 Binary files /dev/null and b/code/submodules/Pplan/Sampling/__pycache__/__init__.cpython-311.pyc differ diff --git a/code/submodules/Pplan/Sampling/__pycache__/forward_sampler.cpython-311.pyc b/code/submodules/Pplan/Sampling/__pycache__/forward_sampler.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..12de02131d72d12fd416b54f1282aa3948c90d9c Binary files /dev/null and b/code/submodules/Pplan/Sampling/__pycache__/forward_sampler.cpython-311.pyc differ diff --git a/code/submodules/Pplan/Sampling/__pycache__/spline_planner.cpython-311.pyc b/code/submodules/Pplan/Sampling/__pycache__/spline_planner.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea28552859a26c78de5fd09173c57fcbb8db3b1c Binary files /dev/null and b/code/submodules/Pplan/Sampling/__pycache__/spline_planner.cpython-311.pyc differ diff --git a/code/submodules/Pplan/Sampling/forward_sampler.py b/code/submodules/Pplan/Sampling/forward_sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..6508701cad30873c552587581e69e325580ea7f7 --- /dev/null +++ b/code/submodules/Pplan/Sampling/forward_sampler.py @@ -0,0 +1,141 @@ +from logging import raiseExceptions +import numpy as np +import torch +import pdb +from ..utils import geometry_utils as GeoUtils +import matplotlib.pyplot as plt +from scipy.interpolate import interp1d +import random +from typing import ( + Any, + Callable, + Dict, + Final, + Iterable, + List, + Optional, + Set, + Tuple, + Union, +) +class ForwardSampler(object): + def __init__(self,dt:float, acce_grid:list,dhm_grid:list, dhf_grid:list, max_rvel=8,max_steer=0.5, vbound=[-5.0, 30],device="cuda" if torch.cuda.is_available() else "cpu"): + self.device = device + self.accels = torch.tensor(acce_grid,device=self.device) + self.dhf_grid = torch.tensor(dhf_grid,device=self.device) + self.dhm_grid = torch.tensor(dhm_grid,device=self.device) + self.max_rvel = max_rvel + self.vbound = vbound + self.max_steer = max_steer + self.dt = dt + + + def velocity_plan(self,x0:torch.Tensor,T:int,acce:Optional[torch.Tensor]=None): + """plan velocity profile + + Args: + x0 (torch.Tensor): [B, 4], X,Y,v,heading + T (int): time horizon + acce (torch.Tensor): [B, N] + """ + bs = x0.shape[0] + if acce is None: + acce = self.accels[None,:].repeat_interleave(bs,0) + v0 = x0[...,2] # [B] + vdes = v0[:,None,None]+torch.arange(T,device=self.device)[None,None]*acce[:,:,None]*self.dt + vplan = torch.clip(vdes,min=self.vbound[0],max=self.vbound[1]) + return vplan # [B, N, T] + + def lateral_plan(self,x0:torch.Tensor,vplan:torch.Tensor,dhf:torch.Tensor,dhm:torch.Tensor,T:int,bangbang=True): + """plan lateral profile, + steering plan that ends with the desired heading change with mean heading change equal to dhm, if feasible + Args: + x0 (torch.Tensor): [B, 4], X,Y,v,heading + vplan (torch.Tensor): [B, N, T] velocity profile + dhf (torch.Tensor): [B, M] desired heading change at the end of the horizon + dhm (torch.Tensor): [B, M] mean heading change at the end of the horizon + T (int): horizon + """ + # using a linear steering profile + bs,M = dhf.shape + N = vplan.shape[1] + vplan = vplan[:,:,None] # [B, N, 1, T] + vl = torch.cat([x0[:,2].reshape(-1,1,1,1).repeat_interleave(N,1),vplan[...,:-1]],-1) + acce = vplan-vl + + c0 = torch.abs(vl) + c1 = torch.cumsum(c0*self.dt,-1) + c2 = torch.cumsum(c1*self.dt,-1) + c3 = torch.cumsum(c2*self.dt,-1) + + + # algebraic equation: c1[T]*a0+c2[T]*a1 = dhf, c2[T]*a0+c3[T]*a1 = dhm + + a0 = (c3[...,-1]*dhf.unsqueeze(1)-c2[...,-1]*dhm.unsqueeze(1))/(c1[...,-1]*c3[...,-1]-c2[...,-1]**2) # [B, N, M] + a1 = (dhf.unsqueeze(1)-c1[...,-1]*a0)/c2[...,-1] + + yawrate = a0[...,None]*c0+a1[...,None]*c1 + + if bangbang: + # turn into bang-bang control to reduce the peak steering value, but the mean heading value is not retained + pos_flag = (yawrate>0) + neg_flag = ~pos_flag + mean_pos_steering = (yawrate*pos_flag).sum(-1)/((c0*pos_flag).sum(-1)+1e-6) + mean_neg_steering = (yawrate*neg_flag).sum(-1)/((c0*neg_flag).sum(-1)+1e-6) + mean_pos_steering = torch.clip(mean_pos_steering,min=-self.max_steer,max=self.max_steer) + mean_neg_steering = torch.clip(mean_neg_steering,min=-self.max_steer,max=self.max_steer) + bb_yawrate = (mean_pos_steering[...,None]*pos_flag+mean_neg_steering[...,None]*neg_flag)*c0 + bb_yawrate = torch.clip(bb_yawrate,min=-self.max_rvel/c0,max=self.max_rvel/c0) + dh = torch.cumsum(bb_yawrate*self.dt,-1) + else: + yawrate = torch.clip(yawrate,min=-self.max_rvel/c0,max=self.max_rvel/c0) + yawrate = torch.clip(yawrate,min=-self.max_steer*c0,max=self.max_steer*c0) + dh = torch.cumsum(yawrate*self.dt,-1) + heading = x0[...,3,None,None,None]+dh + + vx = vplan*torch.cos(heading) + vy = vplan*torch.sin(heading) + traj = torch.stack([x0[:,None,None,None,0]+vx.cumsum(-1)*self.dt, + x0[:,None,None,None,1]+vy.cumsum(-1)*self.dt, + vplan.repeat_interleave(M,2), + heading],-1) + t = torch.arange(1,T+1,device=self.device)[None,None,None,:,None].repeat(bs,N,M,1,1)*self.dt + xyvaqrt = torch.cat([traj[...,:3],acce[...,None].repeat_interleave(M,2),traj[...,3:],yawrate[...,None],t],-1) + return xyvaqrt.reshape(bs,N*M,T,-1) # [B, N*M, T, 7] + + def sample_trajs(self,x0,T,bangbang=True): + # velocity sample + vplan = self.velocity_plan(x0,T) + bs = x0.shape[0] + dhf = self.dhf_grid + dhm = self.dhm_grid + Mf = dhf.shape[0] + Mm = dhm.shape[0] + dhm = dhm.repeat(Mf).unsqueeze(0).repeat_interleave(bs,0) + dhf = dhf.repeat_interleave(Mm,0).unsqueeze(0).repeat_interleave(bs,0)+dhm + return self.lateral_plan(x0,vplan,dhf,dhm,T,bangbang) + + + +def test(): + sampler = ForwardSampler(acce_grid=[-4,-2,0,2,4],dhm_grid=torch.linspace(-0.7,0.7,9),dhf_grid=[-0.4,0,0.4],dt=0.1) + x0 = torch.tensor([0,0,1.,0.],device="cuda").unsqueeze(0).repeat_interleave(3,0) + T = 10 + # vel_grid = sampler.velocity_plan(x0,T) + # dhf = torch.tensor([0.5,0,-0.5]).repeat(3).unsqueeze(0) + # dhm = torch.tensor([0.2,0,-0.2]).repeat_interleave(3,0).unsqueeze(0) + traj = sampler.sample_trajs(x0,T,bangbang=False) + traj = traj[0].reshape(-1,T,7).cpu().numpy() + import matplotlib.pyplot as plt + fig,ax = plt.subplots() + for i in range(traj.shape[0]): + ax.plot(traj[i,:,0],traj[i,:,1]) + plt.show() + + + +if __name__ == "__main__": + test() + + + \ No newline at end of file diff --git a/code/submodules/Pplan/Sampling/spline_planner.py b/code/submodules/Pplan/Sampling/spline_planner.py new file mode 100644 index 0000000000000000000000000000000000000000..fb1fa7aaa17a6d44e5c54e074f6a4d520abc906e --- /dev/null +++ b/code/submodules/Pplan/Sampling/spline_planner.py @@ -0,0 +1,385 @@ +from logging import raiseExceptions +import numpy as np +import torch +import pdb +from ..utils import geometry_utils as GeoUtils +import matplotlib.pyplot as plt +from scipy.interpolate import interp1d +import random + +from .forward_sampler import ForwardSampler + +STATE_INDEX = [0, 1, 2, 4] + +device = "cuda" if torch.cuda.is_available() else "cpu" + + +def mean_control_effort_coefficients(x0, dx0, xf, dxf): + """Returns `(c4, c3, c2)` corresponding to `c4 * tf**-4 + c3 * tf**-3 + c2 * tf**-2`.""" + return (12 * (x0 - xf) ** 2, 12 * (dx0 + dxf) * (x0 - xf), 4 * dx0 ** 2 + 4 * dx0 * dxf + 4 * dxf ** 2) + + +def cubic_spline_coefficients(x0, dx0, xf, dxf, tf): + return (x0, dx0, -2 * dx0 / tf - dxf / tf - 3 * x0 / tf ** 2 + 3 * xf / tf ** 2, + dx0 / tf ** 2 + dxf / tf ** 2 + 2 * x0 / tf ** 3 - 2 * xf / tf ** 3) + + +def compute_interpolating_spline(state_0, state_f, tf): + dx0, dy0 = state_0[..., 2] * \ + torch.cos(state_0[..., 3]), state_0[..., 2] * \ + torch.sin(state_0[..., 3]) + dxf, dyf = state_f[..., 2] * \ + torch.cos(state_f[..., 3]), state_f[..., 2] * \ + torch.sin(state_f[..., 3]) + tf = tf * torch.ones_like(state_0[..., 0]) + return ( + torch.stack(cubic_spline_coefficients( + state_0[..., 0], dx0, state_f[..., 0], dxf, tf), -1), + torch.stack(cubic_spline_coefficients( + state_0[..., 1], dy0, state_f[..., 1], dyf, tf), -1), + tf, + ) + + +def compute_spline_xyvaqrt(x_coefficients, y_coefficients, tf, N=10): + t = torch.arange(N).unsqueeze(0).to(tf.device) * tf.unsqueeze(-1) / (N - 1) + tp = t[..., None] ** torch.arange(4).to(tf.device) + dtp = t[..., None] ** torch.tensor([0, 0, 1, 2] + ).to(tf.device) * torch.arange(4).to(tf.device) + ddtp = t[..., None] ** torch.tensor([0, 0, 0, 1]).to( + tf.device) * torch.tensor([0, 0, 2, 6]).to(tf.device) + x_coefficients = x_coefficients.unsqueeze(-1) + y_coefficients = y_coefficients.unsqueeze(-1) + vx = dtp @ x_coefficients + vy = dtp @ y_coefficients + v = torch.hypot(vx, vy) + v_pos = torch.clip(v, min=1e-4) + ax = ddtp @ x_coefficients + ay = ddtp @ y_coefficients + a = (ax * vx + ay * vy) / v_pos + r = (-ax * vy + ay * vx) / (v_pos ** 2) + yaw = torch.atan2(vy, vx) + return torch.cat(( + tp @ x_coefficients, + tp @ y_coefficients, + v, + a, + yaw, + r, + t.unsqueeze(-1), + ), -1) + + +def patch_yaw_low_speed(traj): + idx = traj[...,] + + +def mean_control_effort_coefficients(x0, dx0, xf, dxf): + """Returns `(c4, c3, c2)` corresponding to `c4 * tf**-4 + c3 * tf**-3 + c2 * tf**-2`.""" + return (12 * (x0 - xf) ** 2, 12 * (dx0 + dxf) * (x0 - xf), 4 * dx0 ** 2 + 4 * dx0 * dxf + 4 * dxf ** 2) + + +class SplinePlanner(object): + def __init__(self, device, dx_grid=None, dy_grid=None, acce_grid=None, dyaw_grid=None, max_steer=0.5, max_rvel=8, + acce_bound=[-6, 4], vbound=[-2.0, 30], spline_order=3, N_seg=10, low_speed_threshold=2.0, seed=0): + self.spline_order = spline_order + self.device = device + assert spline_order == 3 + if dx_grid is None: + # self.dx_grid = torch.tensor([-4., 0, 4.]).to(self.device) + self.dx_grid = torch.tensor([0.]).to(self.device) + else: + self.dx_grid = torch.tensor(dx_grid).to(self.device) + if dy_grid is None: + self.dy_grid = torch.tensor([-3., -1.5, 0, 1.5, 3.]).to(self.device) + else: + self.dy_grid = torch.tensor(dy_grid).to(self.device) + self.dy_grid_lane = torch.tensor([-2., 0, 2., ]).to(self.device) + if acce_grid is None: + # self.acce_grid = torch.tensor([-1., -0.5, 0., 0.5, 1.]).to(self.device) + self.acce_grid = torch.tensor([-1., 0., 1.]).to(self.device) + else: + self.acce_grid = torch.tensor(acce_grid).to(self.device) + if dyaw_grid is None: + self.dyaw_grid = torch.tensor( + [-np.pi / 6, 0, np.pi / 6]).to(self.device) + else: + self.dyaw_grid = torch.tensor(dyaw_grid).to(self.device) + self.max_steer = max_steer + self.max_rvel = max_rvel + self.psi_bound = [-np.pi * 0.75, np.pi * 0.75] + self.acce_bound = acce_bound + self.vbound = vbound + self.N_seg = N_seg + self.low_speed_threshold = low_speed_threshold + self.forward_sampler = ForwardSampler(acce_grid=self.acce_grid, dhm_grid=torch.linspace(-0.7, 0.7, 9), + dhf_grid=[-0.4, 0, 0.4], dt=0.1, device=self.device) + torch.manual_seed(seed) + + def calc_trajectories(self, x0, tf, xf, N=None): + if N is None: + N = self.N_seg + if x0.ndim == 1: + x0_tile = x0.tile(xf.shape[0], 1) + xc, yc, tf = compute_interpolating_spline(x0_tile, xf, tf) + elif x0.ndim == xf.ndim: + xc, yc, tf = compute_interpolating_spline(x0, xf, tf) + else: + raise ValueError("wrong dimension for x0") + traj = compute_spline_xyvaqrt(xc, yc, tf, N) + return traj + + def gen_terminals_lane(self, x0, tf, lanes): + if lanes is None or len(lanes) == 0: + return self.gen_terminals(x0, tf) + + gs = [self.dx_grid.shape[0], self.dy_grid_lane.shape[0], self.acce_grid.shape[0]] + dx = self.dx_grid[:, None, None, None].repeat(1, gs[1], gs[2], 1).flatten() + dy = self.dy_grid_lane[None, :, None, None].repeat(gs[0], 1, gs[2], 1).flatten() + dv = self.acce_grid[None, None, :, None].repeat( + gs[0], gs[1], 1, 1).flatten() * tf + xf = list() + if x0.ndim == 1: + for lane in lanes: + f, p_start = lane + if isinstance(p_start, np.ndarray): + p_start = torch.from_numpy(p_start).to(x0.device) + elif isinstance(p_start, torch.Tensor): + p_start = p_start.to(x0.device) + offset = x0[:2] - p_start[:2] + s_offset = offset[0] * \ + torch.cos(p_start[2]) + offset[1] * torch.sin(p_start[2]) + ds = dx + dv / 2 * tf + x0[2:3] * tf + ss = ds + s_offset + xyyaw = torch.from_numpy(f(ss.cpu().numpy())).type( + torch.float).to(x0.device) + xyyaw[..., 1] += dy + xf.append( + torch.cat((xyyaw[:, :2], dv.reshape(-1, 1) + x0[2:3], xyyaw[:, 2:]), -1)) + # adding the end points not fixated on lane + xf_straight = torch.stack([ds, dy, dv + x0[2], x0[3].tile(ds.shape[0])], -1) + xf.append(xf_straight) + elif x0.ndim == 2: + for lane in lanes: + f, p_start = lane + if isinstance(p_start, np.ndarray): + p_start = torch.from_numpy(p_start).to(x0.device) + elif isinstance(p_start, torch.Tensor): + p_start = p_start.to(x0.device) + offset = x0[:, :2] - p_start[None, :2] + s_offset = offset[:, 0] * torch.cos(p_start[2]) + offset[:, 1] * torch.sin(p_start[2]) + + ds = (dx + dv / 2 * tf).unsqueeze(0) + x0[:, 2:3] * tf + ss = ds + s_offset.unsqueeze(-1) + xyyaw = torch.from_numpy(f(ss.cpu().numpy())).type( + torch.float).to(x0.device) + xyyaw[..., 1] += dy + xf.append(torch.cat((xyyaw[..., :2], dv.tile( + x0.shape[0], 1).unsqueeze(-1) + x0[:, None, 2:3], xyyaw[..., 2:]), -1)) + # adding the end points not fixated on lane + xf_straight = torch.stack([ds, dy.tile(x0.shape[0], 1), dv.tile(x0.shape[0], 1) + x0[:, None, 2], + x0[:, None, 3].tile(1, ds.shape[1])], -1) + xf.append(xf_straight) + else: + raise ValueError("x0 must have dimension 1 or 2") + xf = torch.cat(xf, -2) + return xf + + def gen_terminals(self, x0, tf): + gs = [self.dx_grid.shape[0], self.dy_grid.shape[0], + self.acce_grid.shape[0], self.dyaw_grid.shape[0]] + dx = self.dx_grid[:, None, None, None].repeat( + 1, gs[1], gs[2], gs[3]).flatten() + dy = self.dy_grid[None, :, None, None].repeat( + gs[0], 1, gs[2], gs[3]).flatten() + dv = tf * self.acce_grid[None, None, :, + None].repeat(gs[0], gs[1], 1, gs[3]).flatten() + dyaw = self.dyaw_grid[None, None, None, :].repeat( + gs[0], gs[1], gs[2], 1).flatten() + delta_x = torch.stack([dx, dy, dv, dyaw], -1) + + if x0.ndim == 1: + xy = torch.cat( + (delta_x[:, 0:1] + delta_x[:, 2:3] / 2 * tf + x0[2:3] * tf, delta_x[:, 1:2]), -1) + rotated_delta_xy = GeoUtils.batch_rotate_2D(xy, x0[3]) + refpsi = torch.arctan2(rotated_delta_xy[..., 1], rotated_delta_xy[..., 0]) + rotated_xy = rotated_delta_xy + x0[:2] + return torch.cat((rotated_xy, delta_x[:, 2:3] + x0[2:3], delta_x[:, 3:] + refpsi.unsqueeze(-1)), -1) + elif x0.ndim == 2: + delta_x = torch.tile(delta_x, [x0.shape[0], 1, 1]) + xy = torch.cat( + (delta_x[:, :, 0:1] + delta_x[:, :, 2:3] / 2 * tf + x0[:, None, 2:3] * tf, delta_x[:, :, 1:2]), -1) + rotated_delta_xy = GeoUtils.batch_rotate_2D(xy, x0[:, 3:4]) + refpsi = torch.arctan2(rotated_delta_xy[..., 1], rotated_delta_xy[..., 0]) + rotated_xy = rotated_delta_xy + x0[:, None, :2] + return torch.cat( + (rotated_xy, delta_x[:, :, 2:3] + x0[:, None, 2:3], delta_x[:, :, 3:] + refpsi.unsqueeze(-1)), -1) + else: + raise ValueError("x0 must have dimension 1 or 2") + + def feasible_flag(self, traj, xf): + diff = traj[..., -1, STATE_INDEX] - xf + + feas_flag = ((traj[..., 2] >= self.vbound[0]) & (traj[..., 2] < self.vbound[1]) & + (traj[..., 4] >= self.psi_bound[0]) & (traj[..., 4] < self.psi_bound[1]) & + (traj[..., 3] >= self.acce_bound[0]) & (traj[..., 3] <= self.acce_bound[1]) & + (torch.abs(traj[..., 5] * traj[..., 2]) <= self.max_rvel) & ( + torch.clip(torch.abs(traj[..., 2]), min=0.5) * self.max_steer >= torch.abs( + traj[..., 5]))).all(1) & ( + diff.abs() < 5e-3).all(-1) + + return feas_flag + + def gen_trajectories(self, x0, tf, lanes=None, dyn_filter=True, N=None, lane_only=False): + if N is None: + N = self.N_seg + if lanes is not None: + if isinstance(lanes, torch.Tensor): + lanes = lanes.cpu().numpy() + lane_interp = [GeoUtils.interp_lanes(lane) for lane in lanes] + xf_lane = self.gen_terminals_lane( + x0, tf, lane_interp) + else: + xf_lane = None + if lane_only: + assert xf_lane is not None + xf_set = xf_lane + else: + xf_set = self.gen_terminals(x0, tf) + if xf_lane is not None: + xf_set = torch.cat((xf_lane, xf_set), 0) + x0[..., 2] = torch.clip(x0[..., 2], min=1e-3) + xf_set[..., 2] = torch.clip(xf_set[..., 2], min=1e-3) + + # x, y, v, a, yaw,r, t + traj = self.calc_trajectories(x0, tf, xf_set, N) + if dyn_filter: + feas_flag = self.feasible_flag(traj, xf_set) + traj = traj[feas_flag] + xf = xf_set[feas_flag] + traj = traj[..., 1:, :] # remove the first time step + if x0[2] < self.low_speed_threshold: + # call forward sampler when velocity is low + extra_traj = self.forward_sampler.sample_trajs(x0.unsqueeze(0), int(tf / self.forward_sampler.dt)).squeeze( + 0) + f = interp1d(np.arange(1, extra_traj.shape[-2] + 1) * self.forward_sampler.dt, extra_traj.cpu().numpy(), + axis=-2) + extra_traj = torch.from_numpy(f(np.arange(1, N) * tf / N)).to(self.device) + traj = torch.cat((traj, extra_traj), 0) + return traj, traj[..., -1, STATE_INDEX] + + @staticmethod + def get_similarity_flag(x0, x1, thres=[2.0, 0.5, 2.0, np.pi / 12]): + thres = torch.tensor(thres, device=x0.device) + diff = x0.unsqueeze(-3) - x1.unsqueeze(-2) + flag = diff.abs() < thres + flag = flag.all(-1).any(-1) + return flag + + def gen_terminals_hardcoded(self, x0_set, tf): + X0, Y0, v0, psi0 = x0_set[..., 0:1], x0_set[..., 1:2], x0_set[..., 2:3], x0_set[..., 3:] + xf_set = list() + # drive straight + xf_straight = torch.cat((X0 + v0 * tf * torch.cos(psi0), Y0 + v0 * tf * torch.sin(psi0), v0, psi0), + -1).unsqueeze(1) + xf_set.append(xf_straight) + # hard brake + decel = torch.clip(-v0 / tf, min=self.acce_bound[0]) + xf_brake = torch.cat((X0 + (v0 + decel * 0.5 * tf) * tf * torch.cos(psi0), + Y0 + (v0 + decel * 0.5 * tf) * tf * torch.sin(psi0), v0 + decel * tf, psi0), + -1).unsqueeze(1) + xf_set.append(xf_brake) + xf_set = torch.cat(xf_set, 1) + return xf_set + + def gen_trajectory_batch(self, x0_set, tf, lanes=None, dyn_filter=True, N=None, max_children=None): + if N is None: + N = self.N_seg + device = x0_set.device + xf_set_sample = self.gen_terminals(x0_set, tf) + importance_score = torch.rand(xf_set_sample.shape[:2], device=device) + xf_set_hardcoded = self.gen_terminals_hardcoded(x0_set, tf) + xf_set = torch.cat((xf_set_sample, xf_set_hardcoded), 1) + importance_score = torch.cat((importance_score, 2 * torch.ones(xf_set_hardcoded.shape[:2], device=device)), 1) + if lanes is not None: + lane_interp = [GeoUtils.interp_lanes(lane) for lane in lanes] + xf_set_lane = self.gen_terminals_lane(x0_set, tf, lane_interp) + xf_set = torch.cat((xf_set, xf_set_lane), -2) + importance_score = torch.cat((importance_score, torch.ones(xf_set_lane.shape[:2], device=x0_set.device)), 1) + x0_set[..., 2] = torch.clip(x0_set[..., 2], min=1e-3) + xf_set[..., 2] = torch.clip(xf_set[..., 2], min=1e-3) + num_node = x0_set.shape[0] + num = xf_set.shape[1] + x0_tiled = x0_set.repeat_interleave(num, 0) + xf_tiled = xf_set.reshape(-1, xf_set.shape[-1]) + traj = self.calc_trajectories(x0_tiled, tf, xf_tiled, N) + if dyn_filter: + feas_flag = self.feasible_flag(traj, xf_tiled) + else: + feas_flag = torch.ones( + num * num_node, dtype=torch.bool).to(x0_set.device) + feas_flag = feas_flag.reshape(num_node, num) + traj = traj.reshape(num_node, num, *traj.shape[1:]) + if (x0_set[:, 2] < self.low_speed_threshold).any(): + extra_traj = self.forward_sampler.sample_trajs(x0_set, int(tf / self.forward_sampler.dt)) + f = interp1d(np.arange(1, extra_traj.shape[-2] + 1) * self.forward_sampler.dt, extra_traj.cpu().numpy(), + axis=-2, bounds_error=False, fill_value="extrapolate") + extra_traj = torch.from_numpy(f(np.arange(0, N) * tf / N)).to(self.device) + traj = torch.cat((traj, extra_traj), 1) + extra_importance_score = torch.rand(extra_traj.shape[:2], device=device) + importance_score = torch.cat((importance_score, extra_importance_score), 1) + feas_flag = torch.cat((feas_flag, torch.ones(extra_traj.shape[:2], device=device)), 1) + + importance_score = importance_score * feas_flag + chosen_idx = [torch.where(importance_score[i])[0].tolist() for i in range(num_node)] + if max_children is not None: + chosen_idx = [idx if len(idx) <= max_children else torch.topk(importance_score[i], max_children)[1] for + i, idx in enumerate(chosen_idx)] + traj_batch = [traj[i, chosen_idx[i], 1:] for i in range(num_node)] + return traj_batch + + def gen_trajectory_tree(self, x0, tf, n_layers, dyn_filter=True, N=None): + if N is None: + N = self.N_seg + trajs = list() + nodes = [x0[None, :]] + for i in range(n_layers): + xf = self.gen_terminals(nodes[i], tf) + x0i = torch.tile(nodes[i], [xf.shape[1], 1]) + xf = xf.reshape(-1, xf.shape[-1]) + + traj = self.calc_trajectories(x0i, tf, xf, N) + if dyn_filter: + feas_flag = self.feasible_flag(traj, xf) + traj = traj[feas_flag] + xf = xf[feas_flag] + + trajs.append(traj) + + nodes.append(xf.reshape(-1, xf.shape[-1])) + return trajs, nodes[1:] + + +if __name__ == "__main__": + planner = SplinePlanner("cuda") + x0 = torch.tensor([1., 2., 1., 0.]).cuda() + tf = 5 + traj, xf = planner.gen_trajectories(x0, tf) + trajs = planner.gen_trajectory_batch(xf, tf) + + # x, y, v, a, yaw,r, t = traj + msize = 12 + trajs, nodes = planner.gen_trajectory_tree(x0, tf, 2) + x0 = x0.cpu().numpy() + traj = traj.cpu().numpy() + plt.figure(figsize=(20, 10)) + plt.plot(x0[0], x0[1], marker="o", color="b", markersize=msize) + for node, traj in zip(nodes, trajs): + node = node.cpu().numpy() + traj = traj.cpu().numpy() + x = traj[..., 0] + y = traj[..., 1] + plt.plot(x.T, y.T, color="k") + for p in node: + plt.plot(p[0], p[1], marker="o", color="b", markersize=msize) + plt.show() diff --git a/code/submodules/Pplan/Sampling/trajectory_tree.py b/code/submodules/Pplan/Sampling/trajectory_tree.py new file mode 100644 index 0000000000000000000000000000000000000000..4f1c79ab54e8070ac3e45a90e259f672147aec55 --- /dev/null +++ b/code/submodules/Pplan/Sampling/trajectory_tree.py @@ -0,0 +1,156 @@ +import torch +import itertools +import pdb +import matplotlib.pyplot as plt +import numpy as np +from .tree import Tree +from torch.nn.utils.rnn import pad_sequence + +STATE_INDEX = [0, 1, 2, 4] + + + +class TrajTree(Tree): + + def __init__(self, traj, parent, depth): + self.traj = traj + self.state = traj[-1, STATE_INDEX] + self.children = list() + self.parent = parent + self.depth = depth + self.attribute = dict() + if parent is not None: + self.total_traj = torch.cat((parent.total_traj, traj), 0) + parent.expand(self) + else: + self.total_traj = traj + + def expand_children(self, expand_func): + trajs = expand_func(self.state) + children = [TrajTree(traj, self, self.depth + 1) for traj in trajs] + self.expand_set(children) + + + # @staticmethod + # def get_nodes_by_level(obj,depth,nodes=None): + # assert obj.depth<=depth + # if nodes is None: + # nodes = defaultdict(lambda: list()) + # if obj.depth==depth: + # nodes[depth].append(obj) + # return nodes, True + # else: + # if obj.isleaf(): + # return nodes, False + # else: + # flag = False + # for child in obj.children: + # nodes, child_flag = TrajTree.get_nodes_by_level(child,depth,nodes) + # flag = flag or child_flag + # if flag: + # nodes[obj.depth].append(obj) + # return nodes, flag + + @staticmethod + def get_children(obj): + if isinstance(obj, TrajTree): + return obj.children + elif isinstance(obj, list): + children = [node.children for node in obj] + children = list(itertools.chain.from_iterable(children)) + return children + else: + raise TypeError("obj must be a TrajTree or a list") + + def grow_tree(self, expand_func, n_steps): + assert len(self.children) == 0 + nodes = [self] + device = self.traj.device + with torch.no_grad(): + for i in range(n_steps): + nodes_state = torch.stack( + [node.state for node in nodes], 0).to(device) + trajs = expand_func(nodes_state) + new_nodes = list() + for node, traj in zip(nodes, trajs): + traj[..., -1] += node.traj[-1, -1] + node.children = [TrajTree(traj[i], node, node.depth + 1) + for i in range(traj.shape[0])] + new_nodes += node.children + nodes = new_nodes + if len(new_nodes) == 0: + return + + def plot_tree(self, ax=None, msize=12): + if ax is None: + fig, ax = plt.subplots(figsize=(20, 10)) + state = self.state.cpu().detach().numpy() + + ax.plot(state[0], state[1], marker="o", color="b", markersize=msize) + if self.traj.shape[0] > 1: + if self.parent is not None: + traj_l = torch.cat((self.parent.traj[-1:],self.traj),0) + traj = traj_l.cpu().detach().numpy() + else: + traj = self.traj.cpu().detach().numpy() + ax.plot(traj[:, 0], traj[:, 1], color="k") + for child in self.children: + child.plot_tree(ax) + return ax + + @staticmethod + def get_children_index_torch(nodes): + indices = dict() + for depth,nodes_d in nodes.items(): + if depth+1 in nodes: + childs_d = nodes[depth+1] + indices_d = list() + for node in nodes_d: + indices_d.append(torch.tensor([childs_d.index(child) for child in node.children])) + indices[depth] = pad_sequence(indices_d,batch_first=True,padding_value=-1) + return indices + + + +if __name__ == "__main__": + from Pplan.Sampling.spline_planner import SplinePlanner + from Pplan.utils.timer import Timer + + + + planner = SplinePlanner("cuda") + timer = Timer() + timer.tic() + with torch.no_grad(): + traj = torch.tensor([[4., 1., 5., 0., 0.1, 0., 0.]]).cuda() + + lane0 = np.stack( + [np.linspace(0, 100, 20), np.ones(20) * -3.6, np.zeros(20)]).T + lane1 = np.stack( + [np.linspace(0, 100, 20), np.ones(20) * 0, np.zeros(20)]).T + theta = np.linspace(0, np.pi/2, 20) + R = 50 + lane2 = np.stack( + [R*np.sin(theta), R*(1-np.cos(theta))+3.6, theta]).T + x0 = TrajTree(traj, None, 0) + tf = 5 + + def expand_func(x): return planner.gen_trajectory_batch( + x, tf, (lane0, lane1, lane2),max_children=10) + # def expand_func(x): return planner.gen_trajectory_batch( + # x, tf) + x0.grow_tree(expand_func, 2) + t = timer.toc() + nodes, _ = TrajTree.get_nodes_by_level(x0,2) + children_indices = TrajTree.get_children_index_torch(nodes) + import pdb + pdb.set_trace() + ax = x0.plot_tree() + ax.plot(lane0[:, 0], lane0[:, 1], + linestyle='dashed', linewidth=3, color="r") + ax.plot(lane1[:, 0], lane1[:, 1], + linestyle='dashed', linewidth=3, color="r") + ax.plot(lane2[:, 0], lane2[:, 1], + linestyle='dashed', linewidth=3, color="r") + plt.show() + print("execution time:", t) diff --git a/code/submodules/Pplan/Sampling/tree.py b/code/submodules/Pplan/Sampling/tree.py new file mode 100644 index 0000000000000000000000000000000000000000..e6cef83dc9d7e97f327d5696f5353236ed18cbc5 --- /dev/null +++ b/code/submodules/Pplan/Sampling/tree.py @@ -0,0 +1,72 @@ +import itertools +from collections import defaultdict +class Tree(object): + + def __init__(self, content, parent, depth): + self.content = content + self.children = list() + self.parent = parent + if parent is not None: + parent.expand(self) + self.depth = depth + self.attribute = dict() + + def expand(self, child): + self.children.append(child) + + def expand_set(self, children): + self.children += children + + def isroot(self): + return self.parent is None + + def isleaf(self): + return len(self.children) == 0 + + def get_subseq_trajs(self): + return [child.traj for child in self.children] + + + def get_all_leaves(self,leaf_set=[]): + if self.isleaf(): + leaf_set.append(self) + else: + for child in self.children: + leaf_set = child.get_all_leaves(leaf_set) + return leaf_set + + @staticmethod + def get_nodes_by_level(obj,depth,nodes=None,trim_short_branch=True): + assert obj.depth<=depth + if nodes is None: + nodes = defaultdict(lambda: list()) + if obj.depth==depth: + nodes[depth].append(obj) + return nodes, True + else: + if obj.isleaf(): + return nodes, False + + else: + flag = False + children_flags = dict() + for child in obj.children: + nodes, child_flag = Tree.get_nodes_by_level(child,depth,nodes) + children_flags[child] = child_flag + flag = flag or child_flag + if trim_short_branch: + obj.children = [child for child in obj.children if children_flags[child]] + if flag: + nodes[obj.depth].append(obj) + return nodes, flag + + @staticmethod + def get_children(obj): + if isinstance(obj, Tree): + return obj.children + elif isinstance(obj, list): + children = [node.children for node in obj] + children = list(itertools.chain.from_iterable(children)) + return children + else: + raise TypeError("obj must be a TrajTree or a list") \ No newline at end of file diff --git a/code/submodules/Pplan/__init__.py b/code/submodules/Pplan/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/code/submodules/Pplan/__pycache__/__init__.cpython-311.pyc b/code/submodules/Pplan/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..07e267f9366e62074b6f0f419a0522628f68c50b Binary files /dev/null and b/code/submodules/Pplan/__pycache__/__init__.cpython-311.pyc differ diff --git a/code/submodules/Pplan/utils/__pycache__/geometry_utils.cpython-311.pyc b/code/submodules/Pplan/utils/__pycache__/geometry_utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f9a41a342c5542224376d027d16c17aeb7989854 Binary files /dev/null and b/code/submodules/Pplan/utils/__pycache__/geometry_utils.cpython-311.pyc differ diff --git a/code/submodules/Pplan/utils/geometry_utils.py b/code/submodules/Pplan/utils/geometry_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..08f588c83e31b5071933342466b0277a33230c6d --- /dev/null +++ b/code/submodules/Pplan/utils/geometry_utils.py @@ -0,0 +1,448 @@ +import collections +import numpy as np +import torch +from enum import IntEnum +from scipy.interpolate import interp1d + + +def interp_lanes(lane): + """ generate interpolants for lanes + + Args: + lane (np.array()): [Nx3] + + Returns: + + """ + ds = np.cumsum( + np.hstack([0., np.linalg.norm(lane[1:, :2]-lane[:-1, :2], axis=-1)])) + return interp1d(ds, lane, fill_value="extrapolate", assume_sorted=True, axis=0), lane[0] + + +def batch_proj(x, line): + # x:[batch,3], line:[batch,N,3] + line_length = line.shape[-2] + batch_dim = x.ndim - 1 + if isinstance(x, torch.Tensor): + delta = line[..., 0:2] - torch.unsqueeze(x[..., 0:2], dim=-2).repeat( + *([1] * batch_dim), line_length, 1 + ) + dis = torch.linalg.norm(delta, axis=-1) + idx0 = torch.argmin(dis, dim=-1) + idx = idx0.view(*line.shape[:-2], 1, 1).repeat( + *([1] * (batch_dim + 1)), line.shape[-1] + ) + line_min = torch.squeeze(torch.gather(line, -2, idx), dim=-2) + dx = x[..., None, 0] - line[..., 0] + dy = x[..., None, 1] - line[..., 1] + delta_y = -dx * torch.sin(line_min[..., None, 2]) + dy * torch.cos( + line_min[..., None, 2] + ) + delta_x = dx * torch.cos(line_min[..., None, 2]) + dy * torch.sin( + line_min[..., None, 2] + ) + ref_pts = torch.stack( + [ + line_min[..., 0] + delta_x * torch.cos(line_min[..., 2]), + line_min[..., 1] + delta_x * torch.sin(line_min[..., 2]), + line_min[..., 2], + ], + dim=-1, + ) + delta_psi = round_2pi(x[..., 2] - line_min[..., 2]) + + return ( + delta_x, + delta_y, + torch.unsqueeze(delta_psi, dim=-1), + ref_pts, + ) + elif isinstance(x, np.ndarray): + delta = line[..., 0:2] - np.repeat( + x[..., np.newaxis, 0:2], line_length, axis=-2 + ) + dis = np.linalg.norm(delta, axis=-1) + idx0 = np.argmin(dis, axis=-1) + idx = idx0.reshape(*line.shape[:-2], 1, + 1).repeat(line.shape[-1], axis=-1) + line_min = np.squeeze(np.take_along_axis(line, idx, axis=-2), axis=-2) + dx = x[..., None, 0] - line[..., 0] + dy = x[..., None, 1] - line[..., 1] + delta_y = -dx * np.sin(line_min[..., None, 2]) + dy * np.cos( + line_min[..., None, 2] + ) + delta_x = dx * np.cos(line_min[..., None, 2]) + dy * np.sin( + line_min[..., None, 2] + ) + line_min[..., 0] += delta_x * np.cos(line_min[..., 2]) + line_min[..., 1] += delta_x * np.sin(line_min[..., 2]) + delta_psi = round_2pi(x[..., 2] - line_min[..., 2]) + return ( + delta_x, + delta_y, + np.expand_dims(delta_psi, axis=-1), + line_min, + ) + + +def round_2pi(x): + return (x + np.pi) % (2 * np.pi) - np.pi + + +def get_box_world_coords(pos, yaw, extent): + corners = (torch.tensor([[-1, -1], [-1, 1], [1, 1], [1, -1]]) * 0.5).to(pos.device) * ( + extent.unsqueeze(-2) + ) + s = torch.sin(yaw).unsqueeze(-1) + c = torch.cos(yaw).unsqueeze(-1) + rotM = torch.cat((torch.cat((c, s), dim=-1), + torch.cat((-s, c), dim=-1)), dim=-2) + rotated_corners = (corners + pos.unsqueeze(-2)) @ rotM + return rotated_corners + + +def get_upright_box(pos, extent): + yaws = torch.zeros(*pos.shape[:-1], 1).to(pos.device) + boxes = get_box_world_coords(pos, yaws, extent) + upright_boxes = boxes[..., [0, 2], :] + return upright_boxes + + +def batch_nd_transform_points(points, Mat): + ndim = Mat.shape[-1] - 1 + Mat = torch.transpose(Mat, -1, -2) + return (points.unsqueeze(-2) @ Mat[..., :ndim, :ndim]).squeeze(-2) + Mat[ + ..., -1:, :ndim + ].squeeze(-2) + + +def transform_points_tensor( + points: torch.Tensor, transf_matrix: torch.Tensor +) -> torch.Tensor: + """ + Transform a set of 2D/3D points using the given transformation matrix. + Assumes row major ordering of the input points. The transform function has 3 modes: + - points (N, F), transf_matrix (F+1, F+1) + all points are transformed using the matrix and the output points have shape (N, F). + - points (B, N, F), transf_matrix (F+1, F+1) + all sequences of points are transformed using the same matrix and the output points have shape (B, N, F). + transf_matrix is broadcasted. + - points (B, N, F), transf_matrix (B, F+1, F+1) + each sequence of points is transformed using its own matrix and the output points have shape (B, N, F). + Note this function assumes points.shape[-1] == matrix.shape[-1] - 1, which means that last + rows in the matrices do not influence the final results. + For 2D points only the first 2x3 parts of the matrices will be used. + + :param points: Input points of shape (N, F) or (B, N, F) + with F = 2 or 3 depending on input points are 2D or 3D points. + :param transf_matrix: Transformation matrix of shape (F+1, F+1) or (B, F+1, F+1) with F = 2 or 3. + :return: Transformed points of shape (N, F) or (B, N, F) depending on the dimensions of the input points. + """ + points_log = f" received points with shape {points.shape} " + matrix_log = f" received matrices with shape {transf_matrix.shape} " + + assert points.ndim in [ + 2, 3], f"points should have ndim in [2,3],{points_log}" + assert transf_matrix.ndim in [ + 2, + 3, + ], f"matrix should have ndim in [2,3],{matrix_log}" + assert ( + points.ndim >= transf_matrix.ndim + ), f"points ndim should be >= than matrix,{points_log},{matrix_log}" + + points_feat = points.shape[-1] + assert points_feat in [ + 2, 3], f"last points dimension must be 2 or 3,{points_log}" + assert ( + transf_matrix.shape[-1] == transf_matrix.shape[-2] + ), f"matrix should be a square matrix,{matrix_log}" + + matrix_feat = transf_matrix.shape[-1] + assert matrix_feat in [ + 3, 4], f"last matrix dimension must be 3 or 4,{matrix_log}" + assert ( + points_feat == matrix_feat - 1 + ), f"points last dim should be one less than matrix,{points_log},{matrix_log}" + + def _transform(points: torch.Tensor, transf_matrix: torch.Tensor) -> torch.Tensor: + num_dims = transf_matrix.shape[-1] - 1 + transf_matrix = torch.permute(transf_matrix, (0, 2, 1)) + return ( + points @ transf_matrix[:, :num_dims, :num_dims] + + transf_matrix[:, -1:, :num_dims] + ) + + if points.ndim == transf_matrix.ndim == 2: + points = torch.unsqueeze(points, 0) + transf_matrix = torch.unsqueeze(transf_matrix, 0) + return _transform(points, transf_matrix)[0] + + elif points.ndim == transf_matrix.ndim == 3: + return _transform(points, transf_matrix) + + elif points.ndim == 3 and transf_matrix.ndim == 2: + transf_matrix = torch.unsqueeze(transf_matrix, 0) + return _transform(points, transf_matrix) + else: + raise NotImplementedError( + f"unsupported case!{points_log},{matrix_log}") + + +def PED_PED_collision(p1, p2, S1, S2): + if isinstance(p1, torch.Tensor): + + return ( + torch.linalg.norm(p1[..., 0:2] - p2[..., 0:2], dim=-1) + - (S1[..., 0] + S2[..., 0]) / 2 + ) + + elif isinstance(p1, np.ndarray): + + return ( + np.linalg.norm(p1[..., 0:2] - p2[..., 0:2], axis=-1) + - (S1[..., 0] + S2[..., 0]) / 2 + ) + else: + raise NotImplementedError + + +def batch_rotate_2D(xy, theta): + if isinstance(xy, torch.Tensor): + x1 = xy[..., 0] * torch.cos(theta) - xy[..., 1] * torch.sin(theta) + y1 = xy[..., 1] * torch.cos(theta) + xy[..., 0] * torch.sin(theta) + return torch.stack([x1, y1], dim=-1) + elif isinstance(xy, np.ndarray): + x1 = xy[..., 0] * np.cos(theta) - xy[..., 1] * np.sin(theta) + y1 = xy[..., 1] * np.cos(theta) + xy[..., 0] * np.sin(theta) + return np.concatenate((x1[..., None], y1[..., None]), axis=-1) + + +def VEH_VEH_collision( + p1, p2, S1, S2, alpha=5, return_dis=False, offsetX=1.0, offsetY=0.3 +): + if isinstance(p1, torch.Tensor): + cornersX = torch.kron( + S1[..., 0] + + offsetX, torch.tensor([0.5, 0.5, -0.5, -0.5]).to(p1.device) + ) + cornersY = torch.kron( + S1[..., 1] + + offsetY, torch.tensor([0.5, -0.5, 0.5, -0.5]).to(p1.device) + ) + corners = torch.stack([cornersX, cornersY], dim=-1) + theta1 = p1[..., 2] + theta2 = p2[..., 2] + dx = (p1[..., 0:2] - p2[..., 0:2]).repeat_interleave(4, dim=-2) + delta_x1 = batch_rotate_2D( + corners, theta1.repeat_interleave(4, dim=-1)) + dx + delta_x2 = batch_rotate_2D( + delta_x1, -theta2.repeat_interleave(4, dim=-1)) + dis = torch.maximum( + torch.abs(delta_x2[..., 0]) - 0.5 * + S2[..., 0].repeat_interleave(4, dim=-1), + torch.abs(delta_x2[..., 1]) - 0.5 * + S2[..., 1].repeat_interleave(4, dim=-1), + ).view(*S1.shape[:-1], 4) + min_dis, _ = torch.min(dis, dim=-1) + + return min_dis + + elif isinstance(p1, np.ndarray): + cornersX = np.kron(S1[..., 0] + offsetX, + np.array([0.5, 0.5, -0.5, -0.5])) + cornersY = np.kron(S1[..., 1] + offsetY, + np.array([0.5, -0.5, 0.5, -0.5])) + corners = np.concatenate((cornersX, cornersY), axis=-1) + theta1 = p1[..., 2] + theta2 = p2[..., 2] + dx = (p1[..., 0:2] - p2[..., 0:2]).repeat(4, axis=-2) + delta_x1 = batch_rotate_2D(corners, theta1.repeat(4, axis=-1)) + dx + delta_x2 = batch_rotate_2D(delta_x1, -theta2.repeat(4, axis=-1)) + dis = np.maximum( + np.abs(delta_x2[..., 0]) - 0.5 * S2[..., 0].repeat(4, axis=-1), + np.abs(delta_x2[..., 1]) - 0.5 * S2[..., 1].repeat(4, axis=-1), + ).reshape(*S1.shape[:-1], 4) + min_dis = np.min(dis, axis=-1) + return min_dis + else: + raise NotImplementedError + + +def VEH_PED_collision(p1, p2, S1, S2): + if isinstance(p1, torch.Tensor): + + mask = torch.logical_or( + torch.abs(p1[..., 2]) > 0.1, torch.linalg.norm( + p2[..., 2:4], dim=-1) > 0.1 + ).detach() + theta = p1[..., 2] + dx = batch_rotate_2D(p2[..., 0:2] - p1[..., 0:2], -theta) + + return torch.maximum( + torch.abs(dx[..., 0]) - S1[..., 0] / 2 - S2[..., 0] / 2, + torch.abs(dx[..., 1]) - S1[..., 1] / 2 - S2[..., 0] / 2, + ) + elif isinstance(p1, np.ndarray): + + theta = p1[..., 2] + dx = batch_rotate_2D(p2[..., 0:2] - p1[..., 0:2], -theta) + return np.maximum( + np.abs(dx[..., 0]) - S1[..., 0] / 2 - S2[..., 0] / 2, + np.abs(dx[..., 1]) - S1[..., 1] / 2 - S2[..., 0] / 2, + ) + else: + raise NotImplementedError + + +def PED_VEH_collision(p1, p2, S1, S2): + return VEH_PED_collision(p2, p1, S2, S1) + + +def batch_proj(x, line): + # x:[batch,3], line:[batch,N,3] + line_length = line.shape[-2] + batch_dim = x.ndim - 1 + if isinstance(x, torch.Tensor): + delta = line[..., 0:2] - torch.unsqueeze(x[..., 0:2], dim=-2).repeat( + *([1] * batch_dim), line_length, 1 + ) + dis = torch.linalg.norm(delta, axis=-1) + idx0 = torch.argmin(dis, dim=-1) + idx = idx0.view(*line.shape[:-2], 1, 1).repeat( + *([1] * (batch_dim + 1)), line.shape[-1] + ) + line_min = torch.squeeze(torch.gather(line, -2, idx), dim=-2) + dx = x[..., None, 0] - line[..., 0] + dy = x[..., None, 1] - line[..., 1] + delta_y = -dx * torch.sin(line_min[..., None, 2]) + dy * torch.cos( + line_min[..., None, 2] + ) + delta_x = dx * torch.cos(line_min[..., None, 2]) + dy * torch.sin( + line_min[..., None, 2] + ) + # ref_pts = torch.stack( + # [ + # line_min[..., 0] + delta_x * torch.cos(line_min[..., 2]), + # line_min[..., 1] + delta_x * torch.sin(line_min[..., 2]), + # line_min[..., 2], + # ], + # dim=-1, + # ) + delta_psi = round_2pi(x[..., 2] - line_min[..., 2]) + + return ( + delta_x, + delta_y, + torch.unsqueeze(delta_psi, dim=-1), + ) + elif isinstance(x, np.ndarray): + delta = line[..., 0:2] - np.repeat( + x[..., np.newaxis, 0:2], line_length, axis=-2 + ) + dis = np.linalg.norm(delta, axis=-1) + idx0 = np.argmin(dis, axis=-1) + idx = idx0.reshape(*line.shape[:-2], 1, + 1).repeat(line.shape[-1], axis=-1) + line_min = np.squeeze(np.take_along_axis(line, idx, axis=-2), axis=-2) + dx = x[..., None, 0] - line[..., 0] + dy = x[..., None, 1] - line[..., 1] + delta_y = -dx * np.sin(line_min[..., None, 2]) + dy * np.cos( + line_min[..., None, 2] + ) + delta_x = dx * np.cos(line_min[..., None, 2]) + dy * np.sin( + line_min[..., None, 2] + ) + # line_min[..., 0] += delta_x * np.cos(line_min[..., 2]) + # line_min[..., 1] += delta_x * np.sin(line_min[..., 2]) + delta_psi = round_2pi(x[..., 2] - line_min[..., 2]) + return ( + delta_x, + delta_y, + np.expand_dims(delta_psi, axis=-1), + ) + + +class CollisionType(IntEnum): + """This enum defines the three types of collisions: front, rear and side.""" + FRONT = 0 + REAR = 1 + SIDE = 2 + + +def detect_collision( + ego_pos: np.ndarray, + ego_yaw: np.ndarray, + ego_extent: np.ndarray, + other_pos: np.ndarray, + other_yaw: np.ndarray, + other_extent: np.ndarray, +): + """ + Computes whether a collision occured between ego and any another agent. + Also computes the type of collision: rear, front, or side. + For this, we compute the intersection of ego's four sides with a target + agent and measure the length of this intersection. A collision + is classified into a class, if the corresponding length is maximal, + i.e. a front collision exhibits the longest intersection with + egos front edge. + + .. note:: please note that this funciton will stop upon finding the first + colision, so it won't return all collisions but only the first + one found. + + :param ego_pos: predicted centroid + :param ego_yaw: predicted yaw + :param ego_extent: predicted extent + :param other_pos: target agents + :return: None if not collision was found, and a tuple with the + collision type and the agent track_id + """ + from l5kit.planning import utils + ego_bbox = utils._get_bounding_box( + centroid=ego_pos, yaw=ego_yaw, extent=ego_extent) + # within_range_mask = utils.within_range(ego_pos, ego_extent, other_pos, other_extent) + for i in range(other_pos.shape[0]): + agent_bbox = utils._get_bounding_box( + other_pos[i], other_yaw[i], other_extent[i]) + + if ego_bbox.intersects(agent_bbox): + front_side, rear_side, left_side, right_side = utils._get_sides( + ego_bbox) + + intersection_length_per_side = np.asarray( + [ + agent_bbox.intersection(front_side).length, + agent_bbox.intersection(rear_side).length, + agent_bbox.intersection(left_side).length, + agent_bbox.intersection(right_side).length, + ] + ) + argmax_side = np.argmax(intersection_length_per_side) + + # Remap here is needed because there are two sides that are + # mapped to the same collision type CollisionType.SIDE + max_collision_types = max(CollisionType).value + remap_argmax = min(argmax_side, max_collision_types) + collision_type = CollisionType(remap_argmax) + return collision_type, i + return None + + +def calc_distance_map(road_flag, max_dis=10): + """mark the image with manhattan distance to the drivable area + + Args: + road_flag (torch.Tensor[B,W,H]): an image with 1 channel, 1 for drivable area, 0 for non-drivable area + max_dis (int, optional): maximum distance that the result saturates to. Defaults to 10. + """ + out = torch.zeros_like(road_flag, dtype=torch.float) + out[road_flag == 0] = max_dis + out[road_flag == 1] = 0 + for i in range(max_dis-1): + out[..., 1:, :] = torch.min(out[..., 1:, :], out[..., :-1, :]+1) + out[..., :-1, :] = torch.min(out[..., :-1, :], out[..., 1:, :]+1) + out[..., :, 1:] = torch.min(out[..., :, 1:], out[..., :, :-1]+1) + out[..., :, :-1] = torch.min(out[..., :, :-1], out[..., :, 1:]+1) + + return out diff --git a/code/submodules/Pplan/utils/timer.py b/code/submodules/Pplan/utils/timer.py new file mode 100644 index 0000000000000000000000000000000000000000..c0d299eec3e2a74a5e7aae972ae3dcccb4a20384 --- /dev/null +++ b/code/submodules/Pplan/utils/timer.py @@ -0,0 +1,65 @@ + +import time +import numpy as np +from contextlib import contextmanager + + +class Timer(object): + """A simple timer.""" + def __init__(self): + self.total_time = 0. + self.calls = 0 + self.start_time = 0. + self.diff = 0. + self.average_time = 0. + self.times = [] + + def recent_average_time(self, latest_n): + return np.mean(np.array(self.times)[-latest_n:]) + + def tic(self): + # using time.time instead of time.clock because time time.clock + # does not normalize for multithreading + self.start_time = time.time() + + def toc(self, average=True): + self.diff = time.time() - self.start_time + self.times.append(self.diff) + self.total_time += self.diff + self.calls += 1 + self.average_time = self.total_time / self.calls + if average: + return self.average_time + else: + return self.diff + + @contextmanager + def timed(self): + self.tic() + yield + self.toc() + + +class Timers(object): + def __init__(self): + self._timers = {} + + def tic(self, key): + if key not in self._timers: + self._timers[key] = Timer() + self._timers[key].tic() + + def toc(self, key): + self._timers[key].toc() + + @contextmanager + def timed(self, key): + self.tic(key) + yield + self.toc(key) + + def __str__(self): + msg = [] + for k, v in self._timers.items(): + msg.append('%s: %f' % (k, v.average_time)) + return ', '.join(msg) \ No newline at end of file diff --git a/code/submodules/simple-knn/ext.cpp b/code/submodules/simple-knn/ext.cpp new file mode 100644 index 0000000000000000000000000000000000000000..ae6cefe6ce61a38352a88d07b69a8e6cb9de5b31 --- /dev/null +++ b/code/submodules/simple-knn/ext.cpp @@ -0,0 +1,17 @@ +/* + * Copyright (C) 2023, Inria + * GRAPHDECO research group, https://team.inria.fr/graphdeco + * All rights reserved. + * + * This software is free for non-commercial, research and evaluation use + * under the terms of the LICENSE.md file. + * + * For inquiries contact george.drettakis@inria.fr + */ + +#include +#include "spatial.h" + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("distCUDA2", &distCUDA2); +} diff --git a/code/submodules/simple-knn/setup.py b/code/submodules/simple-knn/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..580d2bd8dc190ce642d87501d53de4f6d9d46c64 --- /dev/null +++ b/code/submodules/simple-knn/setup.py @@ -0,0 +1,35 @@ +# +# Copyright (C) 2023, Inria +# GRAPHDECO research group, https://team.inria.fr/graphdeco +# All rights reserved. +# +# This software is free for non-commercial, research and evaluation use +# under the terms of the LICENSE.md file. +# +# For inquiries contact george.drettakis@inria.fr +# + +from setuptools import setup +from torch.utils.cpp_extension import CUDAExtension, BuildExtension +import os + +cxx_compiler_flags = [] + +if os.name == 'nt': + cxx_compiler_flags.append("/wd4624") + +setup( + name="simple_knn", + ext_modules=[ + CUDAExtension( + name="simple_knn._C", + sources=[ + "spatial.cu", + "simple_knn.cu", + "ext.cpp"], + extra_compile_args={"nvcc": [], "cxx": cxx_compiler_flags}) + ], + cmdclass={ + 'build_ext': BuildExtension + } +) diff --git a/code/submodules/simple-knn/simple_knn.cu b/code/submodules/simple-knn/simple_knn.cu new file mode 100644 index 0000000000000000000000000000000000000000..e72e4c96ea9d161514835fc2fcee62b94954f2d9 --- /dev/null +++ b/code/submodules/simple-knn/simple_knn.cu @@ -0,0 +1,221 @@ +/* + * Copyright (C) 2023, Inria + * GRAPHDECO research group, https://team.inria.fr/graphdeco + * All rights reserved. + * + * This software is free for non-commercial, research and evaluation use + * under the terms of the LICENSE.md file. + * + * For inquiries contact george.drettakis@inria.fr + */ + +#define BOX_SIZE 1024 + +#include "cuda_runtime.h" +#include "device_launch_parameters.h" +#include "simple_knn.h" +#include +#include +#include +#include +#include +#include +#define __CUDACC__ +#include +#include + +namespace cg = cooperative_groups; + +struct CustomMin +{ + __device__ __forceinline__ + float3 operator()(const float3& a, const float3& b) const { + return { min(a.x, b.x), min(a.y, b.y), min(a.z, b.z) }; + } +}; + +struct CustomMax +{ + __device__ __forceinline__ + float3 operator()(const float3& a, const float3& b) const { + return { max(a.x, b.x), max(a.y, b.y), max(a.z, b.z) }; + } +}; + +__host__ __device__ uint32_t prepMorton(uint32_t x) +{ + x = (x | (x << 16)) & 0x030000FF; + x = (x | (x << 8)) & 0x0300F00F; + x = (x | (x << 4)) & 0x030C30C3; + x = (x | (x << 2)) & 0x09249249; + return x; +} + +__host__ __device__ uint32_t coord2Morton(float3 coord, float3 minn, float3 maxx) +{ + uint32_t x = prepMorton(((coord.x - minn.x) / (maxx.x - minn.x)) * ((1 << 10) - 1)); + uint32_t y = prepMorton(((coord.y - minn.y) / (maxx.y - minn.y)) * ((1 << 10) - 1)); + uint32_t z = prepMorton(((coord.z - minn.z) / (maxx.z - minn.z)) * ((1 << 10) - 1)); + + return x | (y << 1) | (z << 2); +} + +__global__ void coord2Morton(int P, const float3* points, float3 minn, float3 maxx, uint32_t* codes) +{ + auto idx = cg::this_grid().thread_rank(); + if (idx >= P) + return; + + codes[idx] = coord2Morton(points[idx], minn, maxx); +} + +struct MinMax +{ + float3 minn; + float3 maxx; +}; + +__global__ void boxMinMax(uint32_t P, float3* points, uint32_t* indices, MinMax* boxes) +{ + auto idx = cg::this_grid().thread_rank(); + + MinMax me; + if (idx < P) + { + me.minn = points[indices[idx]]; + me.maxx = points[indices[idx]]; + } + else + { + me.minn = { FLT_MAX, FLT_MAX, FLT_MAX }; + me.maxx = { -FLT_MAX,-FLT_MAX,-FLT_MAX }; + } + + __shared__ MinMax redResult[BOX_SIZE]; + + for (int off = BOX_SIZE / 2; off >= 1; off /= 2) + { + if (threadIdx.x < 2 * off) + redResult[threadIdx.x] = me; + __syncthreads(); + + if (threadIdx.x < off) + { + MinMax other = redResult[threadIdx.x + off]; + me.minn.x = min(me.minn.x, other.minn.x); + me.minn.y = min(me.minn.y, other.minn.y); + me.minn.z = min(me.minn.z, other.minn.z); + me.maxx.x = max(me.maxx.x, other.maxx.x); + me.maxx.y = max(me.maxx.y, other.maxx.y); + me.maxx.z = max(me.maxx.z, other.maxx.z); + } + __syncthreads(); + } + + if (threadIdx.x == 0) + boxes[blockIdx.x] = me; +} + +__device__ __host__ float distBoxPoint(const MinMax& box, const float3& p) +{ + float3 diff = { 0, 0, 0 }; + if (p.x < box.minn.x || p.x > box.maxx.x) + diff.x = min(abs(p.x - box.minn.x), abs(p.x - box.maxx.x)); + if (p.y < box.minn.y || p.y > box.maxx.y) + diff.y = min(abs(p.y - box.minn.y), abs(p.y - box.maxx.y)); + if (p.z < box.minn.z || p.z > box.maxx.z) + diff.z = min(abs(p.z - box.minn.z), abs(p.z - box.maxx.z)); + return diff.x * diff.x + diff.y * diff.y + diff.z * diff.z; +} + +template +__device__ void updateKBest(const float3& ref, const float3& point, float* knn) +{ + float3 d = { point.x - ref.x, point.y - ref.y, point.z - ref.z }; + float dist = d.x * d.x + d.y * d.y + d.z * d.z; + for (int j = 0; j < K; j++) + { + if (knn[j] > dist) + { + float t = knn[j]; + knn[j] = dist; + dist = t; + } + } +} + +__global__ void boxMeanDist(uint32_t P, float3* points, uint32_t* indices, MinMax* boxes, float* dists) +{ + int idx = cg::this_grid().thread_rank(); + if (idx >= P) + return; + + float3 point = points[indices[idx]]; + float best[3] = { FLT_MAX, FLT_MAX, FLT_MAX }; + + for (int i = max(0, idx - 3); i <= min(P - 1, idx + 3); i++) + { + if (i == idx) + continue; + updateKBest<3>(point, points[indices[i]], best); + } + + float reject = best[2]; + best[0] = FLT_MAX; + best[1] = FLT_MAX; + best[2] = FLT_MAX; + + for (int b = 0; b < (P + BOX_SIZE - 1) / BOX_SIZE; b++) + { + MinMax box = boxes[b]; + float dist = distBoxPoint(box, point); + if (dist > reject || dist > best[2]) + continue; + + for (int i = b * BOX_SIZE; i < min(P, (b + 1) * BOX_SIZE); i++) + { + if (i == idx) + continue; + updateKBest<3>(point, points[indices[i]], best); + } + } + dists[indices[idx]] = (best[0] + best[1] + best[2]) / 3.0f; +} + +void SimpleKNN::knn(int P, float3* points, float* meanDists) +{ + float3* result; + cudaMalloc(&result, sizeof(float3)); + size_t temp_storage_bytes; + + float3 init = { 0, 0, 0 }, minn, maxx; + + cub::DeviceReduce::Reduce(nullptr, temp_storage_bytes, points, result, P, CustomMin(), init); + thrust::device_vector temp_storage(temp_storage_bytes); + + cub::DeviceReduce::Reduce(temp_storage.data().get(), temp_storage_bytes, points, result, P, CustomMin(), init); + cudaMemcpy(&minn, result, sizeof(float3), cudaMemcpyDeviceToHost); + + cub::DeviceReduce::Reduce(temp_storage.data().get(), temp_storage_bytes, points, result, P, CustomMax(), init); + cudaMemcpy(&maxx, result, sizeof(float3), cudaMemcpyDeviceToHost); + + thrust::device_vector morton(P); + thrust::device_vector morton_sorted(P); + coord2Morton << <(P + 255) / 256, 256 >> > (P, points, minn, maxx, morton.data().get()); + + thrust::device_vector indices(P); + thrust::sequence(indices.begin(), indices.end()); + thrust::device_vector indices_sorted(P); + + cub::DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes, morton.data().get(), morton_sorted.data().get(), indices.data().get(), indices_sorted.data().get(), P); + temp_storage.resize(temp_storage_bytes); + + cub::DeviceRadixSort::SortPairs(temp_storage.data().get(), temp_storage_bytes, morton.data().get(), morton_sorted.data().get(), indices.data().get(), indices_sorted.data().get(), P); + + uint32_t num_boxes = (P + BOX_SIZE - 1) / BOX_SIZE; + thrust::device_vector boxes(num_boxes); + boxMinMax << > > (P, points, indices_sorted.data().get(), boxes.data().get()); + boxMeanDist << > > (P, points, indices_sorted.data().get(), boxes.data().get(), meanDists); + + cudaFree(result); +} \ No newline at end of file diff --git a/code/submodules/simple-knn/simple_knn.egg-info/PKG-INFO b/code/submodules/simple-knn/simple_knn.egg-info/PKG-INFO new file mode 100644 index 0000000000000000000000000000000000000000..5f75033b9d43932227e1f3210fa1f2a9d0105e81 --- /dev/null +++ b/code/submodules/simple-knn/simple_knn.egg-info/PKG-INFO @@ -0,0 +1,3 @@ +Metadata-Version: 2.4 +Name: simple_knn +Version: 0.0.0 diff --git a/code/submodules/simple-knn/simple_knn.egg-info/SOURCES.txt b/code/submodules/simple-knn/simple_knn.egg-info/SOURCES.txt new file mode 100644 index 0000000000000000000000000000000000000000..6d197584eea85619778013921e053c37e36d058d --- /dev/null +++ b/code/submodules/simple-knn/simple_knn.egg-info/SOURCES.txt @@ -0,0 +1,8 @@ +ext.cpp +setup.py +simple_knn.cu +spatial.cu +simple_knn.egg-info/PKG-INFO +simple_knn.egg-info/SOURCES.txt +simple_knn.egg-info/dependency_links.txt +simple_knn.egg-info/top_level.txt \ No newline at end of file diff --git a/code/submodules/simple-knn/simple_knn.egg-info/dependency_links.txt b/code/submodules/simple-knn/simple_knn.egg-info/dependency_links.txt new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/code/submodules/simple-knn/simple_knn.egg-info/dependency_links.txt @@ -0,0 +1 @@ + diff --git a/code/submodules/simple-knn/simple_knn.egg-info/top_level.txt b/code/submodules/simple-knn/simple_knn.egg-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..bae7cb8c8a9fdfcc9e62de3793a7b9a5b1f7b4b9 --- /dev/null +++ b/code/submodules/simple-knn/simple_knn.egg-info/top_level.txt @@ -0,0 +1 @@ +simple_knn diff --git a/code/submodules/simple-knn/simple_knn.h b/code/submodules/simple-knn/simple_knn.h new file mode 100644 index 0000000000000000000000000000000000000000..3fcfdb87c53faaadc1fd820d5deeb1b2b5c21a86 --- /dev/null +++ b/code/submodules/simple-knn/simple_knn.h @@ -0,0 +1,21 @@ +/* + * Copyright (C) 2023, Inria + * GRAPHDECO research group, https://team.inria.fr/graphdeco + * All rights reserved. + * + * This software is free for non-commercial, research and evaluation use + * under the terms of the LICENSE.md file. + * + * For inquiries contact george.drettakis@inria.fr + */ + +#ifndef SIMPLEKNN_H_INCLUDED +#define SIMPLEKNN_H_INCLUDED + +class SimpleKNN +{ +public: + static void knn(int P, float3* points, float* meanDists); +}; + +#endif \ No newline at end of file diff --git a/code/submodules/simple-knn/simple_knn/.gitkeep b/code/submodules/simple-knn/simple_knn/.gitkeep new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/code/submodules/simple-knn/spatial.cu b/code/submodules/simple-knn/spatial.cu new file mode 100644 index 0000000000000000000000000000000000000000..1a6a654ba6f8c6a1856a40d14fb7a53c96602bad --- /dev/null +++ b/code/submodules/simple-knn/spatial.cu @@ -0,0 +1,26 @@ +/* + * Copyright (C) 2023, Inria + * GRAPHDECO research group, https://team.inria.fr/graphdeco + * All rights reserved. + * + * This software is free for non-commercial, research and evaluation use + * under the terms of the LICENSE.md file. + * + * For inquiries contact george.drettakis@inria.fr + */ + +#include "spatial.h" +#include "simple_knn.h" + +torch::Tensor +distCUDA2(const torch::Tensor& points) +{ + const int P = points.size(0); + + auto float_opts = points.options().dtype(torch::kFloat32); + torch::Tensor means = torch::full({P}, 0.0, float_opts); + + SimpleKNN::knn(P, (float3*)points.contiguous().data(), means.contiguous().data()); + + return means; +} \ No newline at end of file diff --git a/code/submodules/simple-knn/spatial.h b/code/submodules/simple-knn/spatial.h new file mode 100644 index 0000000000000000000000000000000000000000..280c953a0321a769e433a43535fd36c251b730f0 --- /dev/null +++ b/code/submodules/simple-knn/spatial.h @@ -0,0 +1,14 @@ +/* + * Copyright (C) 2023, Inria + * GRAPHDECO research group, https://team.inria.fr/graphdeco + * All rights reserved. + * + * This software is free for non-commercial, research and evaluation use + * under the terms of the LICENSE.md file. + * + * For inquiries contact george.drettakis@inria.fr + */ + +#include + +torch::Tensor distCUDA2(const torch::Tensor& points); \ No newline at end of file diff --git a/code/utils/__pycache__/dynamic_utils.cpython-311.pyc b/code/utils/__pycache__/dynamic_utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6e42adc9c3063b0eaba04281f65138defb3b5b34 Binary files /dev/null and b/code/utils/__pycache__/dynamic_utils.cpython-311.pyc differ diff --git a/code/utils/__pycache__/general_utils.cpython-311.pyc b/code/utils/__pycache__/general_utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8f9116b2161873be5541e2dc991a8f5dbe133e86 Binary files /dev/null and b/code/utils/__pycache__/general_utils.cpython-311.pyc differ diff --git a/code/utils/__pycache__/graphics_utils.cpython-311.pyc b/code/utils/__pycache__/graphics_utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..36fe1577276fa18b2bed06d181389c1d7f4cab5d Binary files /dev/null and b/code/utils/__pycache__/graphics_utils.cpython-311.pyc differ diff --git a/code/utils/__pycache__/sh_utils.cpython-311.pyc b/code/utils/__pycache__/sh_utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..28ab253cdafd82f991a82a01ee32fdc916b20dfe Binary files /dev/null and b/code/utils/__pycache__/sh_utils.cpython-311.pyc differ diff --git a/code/utils/__pycache__/system_utils.cpython-311.pyc b/code/utils/__pycache__/system_utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..81769a5188b4373135e970b0aff5b6ca2bfbb37f Binary files /dev/null and b/code/utils/__pycache__/system_utils.cpython-311.pyc differ diff --git a/code/utils/cmap.py b/code/utils/cmap.py new file mode 100644 index 0000000000000000000000000000000000000000..cf14e2333fd2452f77d3450e18734fba57daa1b6 --- /dev/null +++ b/code/utils/cmap.py @@ -0,0 +1,90 @@ +import numpy as np + +_color_map_errors = np.array([ + [149, 54, 49], #0: log2(x) = -infinity + [180, 117, 69], #0.0625: log2(x) = -4 + [209, 173, 116], #0.125: log2(x) = -3 + [233, 217, 171], #0.25: log2(x) = -2 + [248, 243, 224], #0.5: log2(x) = -1 + [144, 224, 254], #1.0: log2(x) = 0 + [97, 174, 253], #2.0: log2(x) = 1 + [67, 109, 244], #4.0: log2(x) = 2 + [39, 48, 215], #8.0: log2(x) = 3 + [38, 0, 165], #16.0: log2(x) = 4 + [38, 0, 165] #inf: log2(x) = inf +]).astype(float) + +def color_error_image(errors, scale=1, mask=None, BGR=True): + """ + Color an input error map. + + Arguments: + errors -- HxW numpy array of errors + [scale=1] -- scaling the error map (color change at unit error) + [mask=None] -- zero-pixels are masked white in the result + [BGR=True] -- toggle between BGR and RGB + + Returns: + colored_errors -- HxWx3 numpy array visualizing the errors + """ + + errors_flat = errors.flatten() + errors_color_indices = np.clip(np.log2(errors_flat / scale + 1e-5) + 5, 0, 9) + i0 = np.floor(errors_color_indices).astype(int) + f1 = errors_color_indices - i0.astype(float) + colored_errors_flat = _color_map_errors[i0, :] * (1-f1).reshape(-1,1) + _color_map_errors[i0+1, :] * f1.reshape(-1,1) + + if mask is not None: + colored_errors_flat[mask.flatten() == 0] = 255 + + if not BGR: + colored_errors_flat = colored_errors_flat[:,[2,1,0]] + + return colored_errors_flat.reshape(errors.shape[0], errors.shape[1], 3).astype(np.int) + +_color_map_depths = np.array([ + [0, 0, 0], # 0.000 + [0, 0, 255], # 0.114 + [255, 0, 0], # 0.299 + [255, 0, 255], # 0.413 + [0, 255, 0], # 0.587 + [0, 255, 255], # 0.701 + [255, 255, 0], # 0.886 + [255, 255, 255], # 1.000 + [255, 255, 255], # 1.000 +]).astype(float) +_color_map_bincenters = np.array([ + 0.0, + 0.114, + 0.299, + 0.413, + 0.587, + 0.701, + 0.886, + 1.000, + 2.000, # doesn't make a difference, just strictly higher than 1 +]) + +def color_depth_map(depths, scale=None): + """ + Color an input depth map. + + Arguments: + depths -- HxW numpy array of depths + [scale=None] -- scaling the values (defaults to the maximum depth) + + Returns: + colored_depths -- HxWx3 numpy array visualizing the depths + """ + + # if scale is None: + # scale = depths.max() / 1.5 + scale = 50 + values = np.clip(depths.flatten() / scale, 0, 1) + # for each value, figure out where they fit in in the bincenters: what is the last bincenter smaller than this value? + lower_bin = ((values.reshape(-1, 1) >= _color_map_bincenters.reshape(1,-1)) * np.arange(0,9)).max(axis=1) + lower_bin_value = _color_map_bincenters[lower_bin] + higher_bin_value = _color_map_bincenters[lower_bin + 1] + alphas = (values - lower_bin_value) / (higher_bin_value - lower_bin_value) + colors = _color_map_depths[lower_bin] * (1-alphas).reshape(-1,1) + _color_map_depths[lower_bin + 1] * alphas.reshape(-1,1) + return colors.reshape(depths.shape[0], depths.shape[1], 3).astype(np.uint8) \ No newline at end of file diff --git a/code/utils/dataset.py b/code/utils/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..954995ab809072256e138ff69d42aa73e07e63ae --- /dev/null +++ b/code/utils/dataset.py @@ -0,0 +1,56 @@ +from torch.utils.data import Dataset + +class HUGSIM_dataset(Dataset): + def __init__(self, views, data_type): + super().__init__() + self.views = views + self.data_type = data_type + if data_type == 'kitti360': + self.gap = 4 + elif data_type == 'waymo': + self.gap = 3 + elif data_type == 'kitti': + self.gap = 2 + else: + self.gap = 6 + + def __getitem__(self, index): + if index - self.gap >= 0: + prev_index = index-self.gap + else: + prev_index = -1 + + viewpoint_cam = self.views[index] + + gt_image = viewpoint_cam.original_image + if viewpoint_cam.semantic2d is not None: + gt_semantic = viewpoint_cam.semantic2d + else: + gt_semantic = None + if viewpoint_cam.optical_gt is not None: + gt_optical = viewpoint_cam.optical_gt + else: + gt_optical = None + if viewpoint_cam.depth is not None: + gt_depth = viewpoint_cam.depth + else: + gt_depth = None + if viewpoint_cam.mask is not None: + mask = viewpoint_cam.mask + else: + mask = None + + return index, prev_index, gt_image, gt_semantic, gt_optical, gt_depth, mask + + def __len__(self): + return len(self.views) + +def tocuda(ans): + if ans is None: + return None + else: + return ans.cuda() + +def hugsim_collate(data): + assert len(data) == 1 + return data[0] \ No newline at end of file diff --git a/code/utils/dynamic_utils.py b/code/utils/dynamic_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..de188772b39567a89d6f8bdc9835a70d1635d31e --- /dev/null +++ b/code/utils/dynamic_utils.py @@ -0,0 +1,230 @@ +import numpy as np +import torch +from torch import optim +from torch import nn +from tqdm import tqdm +from matplotlib import pyplot as plt +import torch.nn.functional as F +from collections import defaultdict +import os +import roma + + +class unicycle(torch.nn.Module): + + def __init__(self, train_timestamp, centers, eulers, heights=None): + super(unicycle, self).__init__() + self.train_timestamp = train_timestamp + self.delta = torch.diff(self.train_timestamp) + + self.input_a = centers[:, 0].clone() + self.input_b = centers[:, 1].clone() + + self.a = nn.Parameter(centers[:, 0]) + self.b = nn.Parameter(centers[:, 1]) + + diff_a = torch.diff(centers[:, 0]) / self.delta + diff_b = torch.diff(centers[:, 1]) / self.delta + v = torch.sqrt(diff_a ** 2 + diff_b**2) + + self.v = nn.Parameter(F.pad(v, (0, 1), 'constant', v[-1].item())) + self.pitchroll = eulers[:, :2] + self.yaw = nn.Parameter(eulers[:, -1]) + + if heights is None: + self.h = torch.zeros_like(train_timestamp).float() + else: + self.h = heights + + def acc_omega(self): + acc = torch.diff(self.v) / self.delta + omega = torch.diff(self.yaw) / self.delta + acc = F.pad(acc, (0, 1), 'constant', acc[-1].item()) + omega = F.pad(omega, (0, 1), 'constant', omega[-1].item()) + return acc, omega + + def forward(self, timestamp): + if timestamp < self.train_timestamp[0]: + delta_t = self.train_timestamp[0] - timestamp + a = self.a[0] - delta_t * torch.cos(self.yaw[0]) * self.v[0] + b = self.b[0] - delta_t * torch.sin(self.yaw[0]) * self.v[0] + return a, b, self.v[0], self.pitchroll[0], self.yaw[0], self.h[0] + elif timestamp > self.train_timestamp[-1]: + delta_t = timestamp - self.train_timestamp[-1] + a = self.a[-1] + delta_t * torch.cos(self.yaw[-1]) * self.v[-1] + b = self.b[-1] + delta_t * torch.sin(self.yaw[-1]) * self.v[-1] + return a, b, self.v[-1], self.pitchroll[-1], self.yaw[-1], self.h[-1] + idx = torch.searchsorted(self.train_timestamp, timestamp, side='left') + if self.train_timestamp[idx] == timestamp: + return self.a[idx], self.b[idx], self.v[idx], self.pitchroll[idx], self.yaw[idx], self.h[idx] + else: + prev_timestamps = self.train_timestamp[idx-1] + delta_t = timestamp - prev_timestamps + prev_a, prev_b = self.a[idx-1], self.b[idx-1] + prev_v, prev_yaw = self.v[idx-1], self.yaw[idx-1] + acc, omega = self.acc_omega() + v = prev_v + acc[idx-1] * delta_t + yaw = prev_yaw + omega[idx-1] * delta_t + a = prev_a + prev_v * ((torch.sin(yaw) - torch.sin(prev_yaw)) / (omega[idx-1] + 1e-6)) + b = prev_b - prev_v * ((torch.cos(yaw) - torch.cos(prev_yaw)) / (omega[idx-1] + 1e-6)) + h = self.h[idx-1] + return a, b, v, self.pitchroll[idx-1], yaw, h + + def capture(self): + return ( + self.a, + self.b, + self.v, + self.pitchroll, + self.yaw, + self.h, + self.train_timestamp, + self.delta + ) + + @classmethod + def restore(cls, model_args): + ( + a, + b, + v, + pitchroll, + yaw, + h, + train_timestamp, + delta + ) = model_args + model = cls(train_timestamp, + torch.stack([a.clone().detach(),b.clone().detach()], dim=-1), + torch.concat([pitchroll.clone().detach(), yaw.clone().detach()[:, None]], dim=-1), + h.clone().detach()) + model.a = a + model.b = b + model.v = v + model.pitchroll = pitchroll + model.yaw = yaw + model.h = h + model.train_timestamp = train_timestamp + model.delta = delta + return model + + def visualize(self, save_path, noise_centers=None, gt_centers=None): + a = self.a.detach().cpu().numpy() + b = self.b.detach().cpu().numpy() + yaw = self.yaw.detach().cpu().numpy() + plt.scatter(a, b, marker='x', color='b') + plt.quiver(a, b, np.ones_like(a) * np.cos(yaw), np.ones_like(b) * np.sin(yaw), scale=20, width=0.005) + if noise_centers is not None: + noise_centers = noise_centers.detach().cpu().numpy() + plt.scatter(noise_centers[:, 0], noise_centers[:, 1], marker='o', color='gray') + if gt_centers is not None: + gt_centers = gt_centers.detach().cpu().numpy() + plt.scatter(gt_centers[:, 0], gt_centers[:, 1], marker='v', color='g') + plt.axis('equal') + plt.savefig(save_path) + plt.close() + + def reg_loss(self): + reg = 0 + acc, omega = self.acc_omega() + reg += torch.mean(torch.abs(torch.diff(acc))) * 0.01 + reg += torch.mean(torch.abs(torch.diff(omega))) * 0.1 + reg_a_motion = self.v[:-1] * ((torch.sin(self.yaw[1:]) - torch.sin(self.yaw[:-1])) / (omega[:-1] + 1e-6)) + reg_b_motion = -self.v[:-1] * ((torch.cos(self.yaw[1:]) - torch.cos(self.yaw[:-1])) / (omega[:-1] + 1e-6)) + reg_a = self.a[:-1] + reg_a_motion + reg_b = self.b[:-1] + reg_b_motion + reg += torch.mean((reg_a - self.a[1:])**2 + (reg_b - self.b[1:])**2) * 1 + return reg + + def pos_loss(self): + return torch.mean((self.a - self.input_a) ** 2 + (self.b - self.input_b) ** 2) * 10 + + +def create_unicycle_model(train_cams, model_path, opt_iter=0, opt_pos=False, data_type='kitti'): + unicycle_models = {} + if data_type == 'kitti': + cameras = [cam for cam in train_cams if 'cam_0' in cam.image_name] + elif data_type == 'waymo': + cameras = [cam for cam in train_cams if 'cam_1' in cam.image_name] + elif data_type == 'nuscenes': + cameras = [cam for cam in train_cams if (('CAM_FRONT' in cam.image_name) and ('LEFT' not in cam.image_name) and ('RIGHT' not in cam.image_name))] + elif data_type == 'pandaset': + cameras = [cam for cam in train_cams if 'front_camera' in cam.image_name] + else: + raise NotImplementedError + + cameras = sorted(cameras, key=lambda x: x.timestamp) + + os.makedirs(os.path.join(model_path, "unicycle"), exist_ok=True) + + start_time = cameras[0].timestamp + all_centers, all_heights, all_eulers, all_timestamps = defaultdict(list), defaultdict(list), defaultdict(list), defaultdict(list) + for cam in cameras: + t = cam.timestamp - start_time + for track_id, b2w in cam.dynamics.items(): + all_centers[track_id].append(b2w[[0, 2], 3]) + all_heights[track_id].append(b2w[1, 3]) + eulers = roma.rotmat_to_euler('xzy', b2w[:3, :3]) + all_eulers[track_id].append(-eulers + torch.pi / 2) + # all_eulers[track_id].append(eulers) + all_timestamps[track_id].append(t) + + for track_id in all_centers.keys(): + centers = torch.stack(all_centers[track_id], dim=0).cuda() + timestamps = torch.tensor(all_timestamps[track_id]).cuda() + heights = torch.tensor(all_heights[track_id]).cuda() + eulers = torch.stack(all_eulers[track_id]).cuda() + + model = unicycle(timestamps, centers.clone(), eulers.clone(), heights.clone()) + model.visualize(os.path.join(model_path, "unicycle", f"{track_id}_init.png")) + l = [ + {'params': [model.v], 'lr': 1e-3, "name": "v"}, + {'params': [model.yaw], 'lr': 1e-4, "name": "yaw"}, + ] + + if opt_pos: + l.extend([ + {'params': [model.a], 'lr': 1e-3, "name": "a"}, + {'params': [model.b], 'lr': 1e-3, "name": "b"}, + ]) + + optimizer = optim.Adam(l, lr=0.0) + + t_range = tqdm(range(opt_iter), desc=f"Fitting {track_id}") + for iter in t_range: + loss = 5e-3 * model.reg_loss() + 1e-3 * model.pos_loss() + t_range.set_postfix({'loss': loss.item()}) + optimizer.zero_grad() + loss.backward() + optimizer.step() + + unicycle_models[track_id] = {'model': model, + 'optimizer': optimizer, + 'input_centers': centers} + + model.visualize(os.path.join(model_path, "unicycle", f"{track_id}_iter0.png")) + torch.save(model.capture(), os.path.join(model_path, f"ckpts/unicycle_{track_id}.pth")) + + return unicycle_models + + +if __name__ == "__main__": + from scene import Scene, GaussianModel + from omegaconf import OmegaConf + from argparse import ArgumentParser + + parser = ArgumentParser(description="Training script parameters") + parser.add_argument("--base_cfg", type=str, default="./configs/gs_base.yaml") + parser.add_argument("--data_cfg", type=str, default="./configs/nusc.yaml") + parser.add_argument("--source_path", type=str, default="") + parser.add_argument("--model_path", type=str, default="") + args = parser.parse_args() + cfg = OmegaConf.merge(OmegaConf.load(args.base_cfg), OmegaConf.load(args.data_cfg)) + if len(args.source_path) > 0: + cfg.source_path = args.source_path + if len(args.model_path) > 0: + cfg.model_path = args.model_path + gaussians = GaussianModel(3, feat_mutable=True) + print("loading scene...") + scene = Scene(cfg, gaussians, data_type=cfg.data_type) + create_unicycle_model(scene.getTrainCameras(), cfg.model_path, 500, False, cfg.data_type) \ No newline at end of file diff --git a/code/utils/general_utils.py b/code/utils/general_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..aafde9e5db8d668764034d577002ec808b0d29db --- /dev/null +++ b/code/utils/general_utils.py @@ -0,0 +1,174 @@ +# +# Copyright (C) 2023, Inria +# GRAPHDECO research group, https://team.inria.fr/graphdeco +# All rights reserved. +# +# This software is free for non-commercial, research and evaluation use +# under the terms of the LICENSE.md file. +# +# For inquiries contact george.drettakis@inria.fr +# + +import torch +import sys +from datetime import datetime +import numpy as np +import random +import os +import cv2 + +def inverse_sigmoid(x): + return torch.log(x/(1-x)) + +def PILtoTorch(pil_image, resolution): + resized_image_PIL = pil_image.resize(resolution) + resized_image = torch.from_numpy(np.array(resized_image_PIL)) / 255.0 + if len(resized_image.shape) == 3: + return resized_image.permute(2, 0, 1) + else: + return resized_image.unsqueeze(dim=-1).permute(2, 0, 1) + +def PIL2toTorch(pil_image, resolution): + resized_image_PIL = pil_image.resize(resolution) + resized_image = torch.from_numpy(np.array(resized_image_PIL)) / 255.0 * (2.0 ** 16 - 1.0) + return resized_image + +def decode_op(optical_png): + # use 'PIL Image.Open' to READ + "Convert from .png (h, w, 3-rgb) -> (h,w,2)(flow_x, flow_y) .. float32 array" + optical_png = optical_png[..., [2, 1, 0]] # bgr -> rgb + h, w, _c = optical_png.shape + assert optical_png.dtype == np.uint16 and _c == 3 + "invalid flow flag: b == 0 for sky or other invalid flow" + invalid_points = np.where(optical_png[..., 2] == 0) + out_flow = torch.empty((h, w, 2)) + decoded = 2.0 / (2**16 - 1.0) * optical_png.astype('f4') - 1 + out_flow[..., 0] = torch.tensor(decoded[:, :, 0] * (w - 1)) # (pixel) delta_x : R + out_flow[..., 1] = torch.tensor(decoded[:, :, 1] * (h - 1)) # delta_y : G + out_flow[invalid_points[0], invalid_points[1], :] = 0 # B=0 for invalid flow + return out_flow + +def get_expon_lr_func( + lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000 +): + """ + Copied from Plenoxels + + Continuous learning rate decay function. Adapted from JaxNeRF + The returned rate is lr_init when step=0 and lr_final when step=max_steps, and + is log-linearly interpolated elsewhere (equivalent to exponential decay). + If lr_delay_steps>0 then the learning rate will be scaled by some smooth + function of lr_delay_mult, such that the initial learning rate is + lr_init*lr_delay_mult at the beginning of optimization but will be eased back + to the normal learning rate when steps>lr_delay_steps. + :param conf: config subtree 'lr' or similar + :param max_steps: int, the number of steps during optimization. + :return HoF which takes step as input + """ + + def helper(step): + if step < 0 or (lr_init == 0.0 and lr_final == 0.0): + # Disable this parameter + return 0.0 + if lr_delay_steps > 0: + # A kind of reverse cosine decay. + delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin( + 0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1) + ) + else: + delay_rate = 1.0 + t = np.clip(step / max_steps, 0, 1) + log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t) + return delay_rate * log_lerp + + return helper + +def strip_lowerdiag(L): + uncertainty = torch.zeros((L.shape[0], 6), dtype=torch.float, device="cuda") + + uncertainty[:, 0] = L[:, 0, 0] + uncertainty[:, 1] = L[:, 0, 1] + uncertainty[:, 2] = L[:, 0, 2] + uncertainty[:, 3] = L[:, 1, 1] + uncertainty[:, 4] = L[:, 1, 2] + uncertainty[:, 5] = L[:, 2, 2] + return uncertainty + +def strip_symmetric(sym): + return strip_lowerdiag(sym) + +def build_rotation(r): + norm = torch.sqrt(r[:,0]*r[:,0] + r[:,1]*r[:,1] + r[:,2]*r[:,2] + r[:,3]*r[:,3]) + + q = r / norm[:, None] + + R = torch.zeros((q.size(0), 3, 3), device='cuda') + + r = q[:, 0] + x = q[:, 1] + y = q[:, 2] + z = q[:, 3] + + R[:, 0, 0] = 1 - 2 * (y*y + z*z) + R[:, 0, 1] = 2 * (x*y - r*z) + R[:, 0, 2] = 2 * (x*z + r*y) + R[:, 1, 0] = 2 * (x*y + r*z) + R[:, 1, 1] = 1 - 2 * (x*x + z*z) + R[:, 1, 2] = 2 * (y*z - r*x) + R[:, 2, 0] = 2 * (x*z - r*y) + R[:, 2, 1] = 2 * (y*z + r*x) + R[:, 2, 2] = 1 - 2 * (x*x + y*y) + return R + +def build_scaling_rotation(s, r): + L = torch.zeros((s.shape[0], 3, 3), dtype=torch.float, device="cuda") + R = build_rotation(r) + + L[:,0,0] = s[:,0] + L[:,1,1] = s[:,1] + L[:,2,2] = s[:,2] + + L = R @ L + return L + +DEFAULT_RANDOM_SEED = 0 + +def seedBasic(seed=DEFAULT_RANDOM_SEED): + random.seed(seed) + os.environ['PYTHONHASHSEED'] = str(seed) + np.random.seed(seed) + +def seedTorch(seed=DEFAULT_RANDOM_SEED): + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + +# basic + tensorflow + torch +def seedEverything(seed=DEFAULT_RANDOM_SEED): + seedBasic(seed) + seedTorch(seed) + +def safe_state(silent): + old_f = sys.stdout + class F: + def __init__(self, silent): + self.silent = silent + + def write(self, x): + if not self.silent: + if x.endswith("\n"): + old_f.write(x.replace("\n", " [{}]\n".format(str(datetime.now().strftime("%d/%m %H:%M:%S"))))) + else: + old_f.write(x) + + def flush(self): + old_f.flush() + + sys.stdout = F(silent) + + random.seed(DEFAULT_RANDOM_SEED) + np.random.seed(DEFAULT_RANDOM_SEED) + torch.manual_seed(DEFAULT_RANDOM_SEED) + torch.cuda.set_device(torch.device("cuda:0")) + # sys.stdout = old_f diff --git a/code/utils/graphics_utils.py b/code/utils/graphics_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..35cb0a2c46f6569cf593165ee3cc69995d0a335f --- /dev/null +++ b/code/utils/graphics_utils.py @@ -0,0 +1,78 @@ +# +# Copyright (C) 2023, Inria +# GRAPHDECO research group, https://team.inria.fr/graphdeco +# All rights reserved. +# +# This software is free for non-commercial, research and evaluation use +# under the terms of the LICENSE.md file. +# +# For inquiries contact george.drettakis@inria.fr +# + +import torch +import math +import numpy as np +from typing import NamedTuple + +class BasicPointCloud(NamedTuple): + points : np.array + colors : np.array + normals : np.array + +def geom_transform_points(points, transf_matrix): + P, _ = points.shape + ones = torch.ones(P, 1, dtype=points.dtype, device=points.device) + points_hom = torch.cat([points, ones], dim=1) + points_out = torch.matmul(points_hom, transf_matrix.unsqueeze(0)) + + denom = points_out[..., 3:] + 0.0000001 + return (points_out[..., :3] / denom).squeeze(dim=0) + +def getWorld2View(R, t): + Rt = np.zeros((4, 4)) + Rt[:3, :3] = R.transpose() + Rt[:3, 3] = t + Rt[3, 3] = 1.0 + return np.float32(Rt) + +def getWorld2View2(R, t, translate=np.array([.0, .0, .0]), scale=1.0): + Rt = np.zeros((4, 4)) + Rt[:3, :3] = R.transpose() + Rt[:3, 3] = t + Rt[3, 3] = 1.0 + + C2W = np.linalg.inv(Rt) + cam_center = C2W[:3, 3] + cam_center = (cam_center + translate) * scale + C2W[:3, 3] = cam_center + Rt = np.linalg.inv(C2W) + return np.float32(Rt) + +def getProjectionMatrix(znear, zfar, fovX, fovY, cx_ratio, cy_ratio): + tanHalfFovY = math.tan((fovY / 2)) + tanHalfFovX = math.tan((fovX / 2)) + + top = tanHalfFovY * znear + bottom = -top + right = tanHalfFovX * znear + left = -right + + P = torch.zeros(4, 4) + + z_sign = 1.0 + + P[0, 0] = 2.0 * znear / (right - left) + P[1, 1] = 2.0 * znear / (top - bottom) + P[0, 2] = (right + left) / (right - left) - 1 + cx_ratio + P[1, 2] = (top + bottom) / (top - bottom) - 1 + cy_ratio + P[3, 2] = z_sign + P[2, 2] = z_sign * (zfar + znear) / (zfar - znear) + P[2, 3] = -(2 * zfar * znear) / (zfar - znear) + + return P + +def fov2focal(fov, pixels): + return pixels / (2 * math.tan(fov / 2)) + +def focal2fov(focal, pixels): + return 2*math.atan(pixels/(2*focal)) \ No newline at end of file diff --git a/code/utils/image_utils.py b/code/utils/image_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..cdeaa1b6d250e549181ab165070f82ccd31b3eb9 --- /dev/null +++ b/code/utils/image_utils.py @@ -0,0 +1,19 @@ +# +# Copyright (C) 2023, Inria +# GRAPHDECO research group, https://team.inria.fr/graphdeco +# All rights reserved. +# +# This software is free for non-commercial, research and evaluation use +# under the terms of the LICENSE.md file. +# +# For inquiries contact george.drettakis@inria.fr +# + +import torch + +def mse(img1, img2): + return (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True) + +def psnr(img1, img2): + mse = (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True) + return 20 * torch.log10(1.0 / torch.sqrt(mse)) diff --git a/code/utils/iou_utils.py b/code/utils/iou_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..cff657484de352bae65ae5990836d347de341132 --- /dev/null +++ b/code/utils/iou_utils.py @@ -0,0 +1,159 @@ +# 3D IoU caculate code for 3D object detection +# Kent 2018/12 + +import numpy as np +from scipy.spatial import ConvexHull +from numpy import * + +def polygon_clip(subjectPolygon, clipPolygon): + """ Clip a polygon with another polygon. + + Ref: https://rosettacode.org/wiki/Sutherland-Hodgman_polygon_clipping#Python + + Args: + subjectPolygon: a list of (x,y) 2d points, any polygon. + clipPolygon: a list of (x,y) 2d points, has to be *convex* + Note: + **points have to be counter-clockwise ordered** + + Return: + a list of (x,y) vertex point for the intersection polygon. + """ + def inside(p): + return(cp2[0]-cp1[0])*(p[1]-cp1[1]) > (cp2[1]-cp1[1])*(p[0]-cp1[0]) + + def computeIntersection(): + dc = [ cp1[0] - cp2[0], cp1[1] - cp2[1] ] + dp = [ s[0] - e[0], s[1] - e[1] ] + n1 = cp1[0] * cp2[1] - cp1[1] * cp2[0] + n2 = s[0] * e[1] - s[1] * e[0] + n3 = 1.0 / (dc[0] * dp[1] - dc[1] * dp[0]) + return [(n1*dp[0] - n2*dc[0]) * n3, (n1*dp[1] - n2*dc[1]) * n3] + + outputList = subjectPolygon + cp1 = clipPolygon[-1] + + for clipVertex in clipPolygon: + cp2 = clipVertex + inputList = outputList + outputList = [] + s = inputList[-1] + + for subjectVertex in inputList: + e = subjectVertex + if inside(e): + if not inside(s): + outputList.append(computeIntersection()) + outputList.append(e) + elif inside(s): + outputList.append(computeIntersection()) + s = e + cp1 = cp2 + if len(outputList) == 0: + return None + return(outputList) + +def poly_area(x,y): + """ Ref: http://stackoverflow.com/questions/24467972/calculate-area-of-polygon-given-x-y-coordinates """ + return 0.5*np.abs(np.dot(x,np.roll(y,1))-np.dot(y,np.roll(x,1))) + +def convex_hull_intersection(p1, p2): + """ Compute area of two convex hull's intersection area. + p1,p2 are a list of (x,y) tuples of hull vertices. + return a list of (x,y) for the intersection and its volume + """ + inter_p = polygon_clip(p1,p2) + if inter_p is not None: + hull_inter = ConvexHull(inter_p) + return inter_p, hull_inter.volume + else: + return None, 0.0 + +def box3d_vol(corners): + ''' corners: (8,3) no assumption on axis direction ''' + a = np.sqrt(np.sum((corners[0,:] - corners[1,:])**2)) + b = np.sqrt(np.sum((corners[1,:] - corners[2,:])**2)) + c = np.sqrt(np.sum((corners[0,:] - corners[4,:])**2)) + return a*b*c + +def is_clockwise(p): + x = p[:,0] + y = p[:,1] + return np.dot(x,np.roll(y,1))-np.dot(y,np.roll(x,1)) > 0 + +def box3d_iou(corners1, corners2): + ''' Compute 3D bounding box IoU. + + Input: + corners1: numpy array (8,3), assume up direction is negative Y + corners2: numpy array (8,3), assume up direction is negative Y + Output: + iou: 3D bounding box IoU + iou_2d: bird's eye view 2D bounding box IoU + + todo (kent): add more description on corner points' orders. + ''' + # corner points are in counter clockwise order + rect1 = [(corners1[i,0], corners1[i,2]) for i in [4,5,1,0]] + rect2 = [(corners2[i,0], corners2[i,2]) for i in [4,5,1,0]] + + area1 = poly_area(np.array(rect1)[:,0], np.array(rect1)[:,1]) + area2 = poly_area(np.array(rect2)[:,0], np.array(rect2)[:,1]) + + inter, inter_area = convex_hull_intersection(rect1, rect2) + iou_2d = inter_area/(area1+area2-inter_area) + # if iou_2d < 0: + # print(inter_area, area1, area2) + # ymax = min(corners1[0,1], corners2[0,1]) + # ymin = max(corners1[4,1], corners2[4,1]) + + # inter_vol = inter_area * max(0.0, ymax-ymin) + + # vol1 = box3d_vol(corners1) + # vol2 = box3d_vol(corners2) + # iou = inter_vol / (vol1 + vol2 - inter_vol) + # return iou, iou_2d + return 0, iou_2d + +# ---------------------------------- +# Helper functions for evaluation +# ---------------------------------- + +def get_3d_box(box_size, heading_angle, center): + ''' Calculate 3D bounding box corners from its parameterization. + + Input: + box_size: tuple of (length,wide,height) + heading_angle: rad scalar, clockwise from pos x axis + center: tuple of (x,y,z) + Output: + corners_3d: numpy array of shape (8,3) for 3D box cornders + ''' + def roty(t): + c = np.cos(t) + s = np.sin(t) + return np.array([[c, 0, s], + [0, 1, 0], + [-s, 0, c]]) + + R = roty(heading_angle) + l,w,h = box_size + x_corners = [l/2,l/2,-l/2,-l/2,l/2,l/2,-l/2,-l/2]; + y_corners = [h/2,h/2,h/2,h/2,-h/2,-h/2,-h/2,-h/2]; + z_corners = [w/2,-w/2,-w/2,w/2,w/2,-w/2,-w/2,w/2]; + corners_3d = np.dot(R, np.vstack([x_corners,y_corners,z_corners])) + corners_3d[0,:] = corners_3d[0,:] + center[0]; + corners_3d[1,:] = corners_3d[1,:] + center[1]; + corners_3d[2,:] = corners_3d[2,:] + center[2]; + corners_3d = np.transpose(corners_3d) + return corners_3d + + +if __name__=='__main__': + print('------------------') + # get_3d_box(box_size, heading_angle, center) + corners_3d_ground = get_3d_box((1.497255,1.644981, 3.628938), -1.531692, (2.882992 ,1.698800 ,20.785644)) + corners_3d_predict = get_3d_box((1.458242, 1.604773, 3.707947), -1.549553, (2.756923, 1.661275, 20.943280 )) + (IOU_3d,IOU_2d)=box3d_iou(corners_3d_predict,corners_3d_ground) + print (IOU_3d,IOU_2d) #3d IoU/ 2d IoU of BEV(bird eye's view) + \ No newline at end of file diff --git a/code/utils/loss_utils.py b/code/utils/loss_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..742577f3d8c9bb42d636c14d1506753b56ae216e --- /dev/null +++ b/code/utils/loss_utils.py @@ -0,0 +1,347 @@ +# +# Copyright (C) 2023, Inria +# GRAPHDECO research group, https://team.inria.fr/graphdeco +# All rights reserved. +# +# This software is free for non-commercial, research and evaluation use +# under the terms of the LICENSE.md file. +# +# For inquiries contact george.drettakis@inria.fr +# + +import torch +import torch.nn.functional as F +from torch.autograd import Variable +from math import exp +from torch import Tensor, nn +from typing import Dict, Literal, Optional, Tuple, cast +from jaxtyping import Bool, Float + +L1Loss = nn.L1Loss +MSELoss = nn.MSELoss + +def l1_loss(network_output, gt, mask=None): + l1 = torch.abs((network_output - gt)) + if mask is not None: + l1 = l1[:, mask] + return l1.mean() + +def l2_loss(network_output, gt): + return ((network_output - gt) ** 2).mean() + +def gaussian(window_size, sigma): + gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + +def create_window(window_size, channel): + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + +def ssim(img1, img2, window_size=11, size_average=True): + channel = img1.size(-3) + window = create_window(window_size, channel) + + if img1.is_cuda: + window = window.cuda(img1.get_device()) + window = window.type_as(img1) + + return _ssim(img1, img2, window, window_size, channel, size_average) + +def _ssim(img1, img2, window, window_size, channel, size_average=True): + mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel) + + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2 + + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + + ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + + if size_average: + return ssim_map.mean() + else: + return ssim_map.mean(1).mean(1).mean(1) + +def ssim_loss(img1, img2, window_size=11, size_average=True, mask=None): + channel = img1.size(-3) + window = create_window(window_size, channel) + + if img1.is_cuda: + window = window.cuda(img1.get_device()) + window = window.type_as(img1) + + return _ssim_loss(img1, img2, window, window_size, channel, size_average, mask) + +def _ssim_loss(img1, img2, window, window_size, channel, size_average=True, mask=None): + mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel) + + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2 + + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + + ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + ssim_map = 1 - ssim_map + + if mask is not None: + ssim_map = ssim_map[:, mask] + if size_average: + return ssim_map.mean() + else: + return ssim_map.mean(1).mean(1).mean(1) + + +def masked_reduction( + input_tensor: Float[Tensor, "1 32 mult"], + mask: Bool[Tensor, "1 32 mult"], + reduction_type: Literal["image", "batch"], +) -> Tensor: + """ + Whether to consolidate the input_tensor across the batch or across the image + Args: + input_tensor: input tensor + mask: mask tensor + reduction_type: either "batch" or "image" + Returns: + input_tensor: reduced input_tensor + """ + if reduction_type == "batch": + # avoid division by 0 (if sum(M) = sum(sum(mask)) = 0: sum(image_loss) = 0) + divisor = torch.sum(mask) + if divisor == 0: + return torch.tensor(0, device=input_tensor.device) + input_tensor = torch.sum(input_tensor) / divisor + elif reduction_type == "image": + # avoid division by 0 (if M = sum(mask) = 0: image_loss = 0) + valid = mask.nonzero() + + input_tensor[valid] = input_tensor[valid] / mask[valid] + input_tensor = torch.mean(input_tensor) + return input_tensor + + +def normalized_depth_scale_and_shift( + prediction: Float[Tensor, "1 32 mult"], target: Float[Tensor, "1 32 mult"], mask: Bool[Tensor, "1 32 mult"] +): + """ + More info here: https://arxiv.org/pdf/2206.00665.pdf supplementary section A2 Depth Consistency Loss + This function computes scale/shift required to normalizes predicted depth map, + to allow for using normalized depth maps as input from monocular depth estimation networks. + These networks are trained such that they predict normalized depth maps. + + Solves for scale/shift using a least squares approach with a closed form solution: + Based on: + https://github.com/autonomousvision/monosdf/blob/d9619e948bf3d85c6adec1a643f679e2e8e84d4b/code/model/loss.py#L7 + Args: + prediction: predicted depth map + target: ground truth depth map + mask: mask of valid pixels + Returns: + scale and shift for depth prediction + """ + # system matrix: A = [[a_00, a_01], [a_10, a_11]] + a_00 = torch.sum(mask * prediction * prediction, (1, 2)) + a_01 = torch.sum(mask * prediction, (1, 2)) + a_11 = torch.sum(mask, (1, 2)) + + # right hand side: b = [b_0, b_1] + b_0 = torch.sum(mask * prediction * target, (1, 2)) + b_1 = torch.sum(mask * target, (1, 2)) + + # solution: x = A^-1 . b = [[a_11, -a_01], [-a_10, a_00]] / (a_00 * a_11 - a_01 * a_10) . b + scale = torch.zeros_like(b_0) + shift = torch.zeros_like(b_1) + + det = a_00 * a_11 - a_01 * a_01 + valid = det.nonzero() + + scale[valid] = (a_11[valid] * b_0[valid] - a_01[valid] * b_1[valid]) / det[valid] + shift[valid] = (-a_01[valid] * b_0[valid] + a_00[valid] * b_1[valid]) / det[valid] + + return scale, shift + + +class MiDaSMSELoss(nn.Module): + """ + data term from MiDaS paper + """ + + def __init__(self, reduction_type: Literal["image", "batch"] = "batch"): + super().__init__() + + self.reduction_type: Literal["image", "batch"] = reduction_type + # reduction here is different from the image/batch-based reduction. This is either "mean" or "sum" + self.mse_loss = MSELoss(reduction="none") + + def forward( + self, + prediction: Float[Tensor, "1 32 mult"], + target: Float[Tensor, "1 32 mult"], + mask: Bool[Tensor, "1 32 mult"], + ) -> Float[Tensor, "0"]: + """ + Args: + prediction: predicted depth map + target: ground truth depth map + mask: mask of valid pixels + Returns: + mse loss based on reduction function + """ + summed_mask = torch.sum(mask, (1, 2)) + image_loss = torch.sum(self.mse_loss(prediction, target) * mask, (1, 2)) + # multiply by 2 magic number? + image_loss = masked_reduction(image_loss, 2 * summed_mask, self.reduction_type) + + return image_loss + + +class GradientLoss(nn.Module): + """ + multiscale, scale-invariant gradient matching term to the disparity space. + This term biases discontinuities to be sharp and to coincide with discontinuities in the ground truth + More info here https://arxiv.org/pdf/1907.01341.pdf Equation 11 + """ + + def __init__(self, scales: int = 4, reduction_type: Literal["image", "batch"] = "batch"): + """ + Args: + scales: number of scales to use + reduction_type: either "batch" or "image" + """ + super().__init__() + self.reduction_type: Literal["image", "batch"] = reduction_type + self.__scales = scales + + def forward( + self, + prediction: Float[Tensor, "1 32 mult"], + target: Float[Tensor, "1 32 mult"], + mask: Bool[Tensor, "1 32 mult"], + ) -> Float[Tensor, "0"]: + """ + Args: + prediction: predicted depth map + target: ground truth depth map + mask: mask of valid pixels + Returns: + gradient loss based on reduction function + """ + assert self.__scales >= 1 + total = 0.0 + + for scale in range(self.__scales): + step = pow(2, scale) + + grad_loss = self.gradient_loss( + prediction[:, ::step, ::step], + target[:, ::step, ::step], + mask[:, ::step, ::step], + ) + total += grad_loss + + assert isinstance(total, Tensor) + return total + + def gradient_loss( + self, + prediction: Float[Tensor, "1 32 mult"], + target: Float[Tensor, "1 32 mult"], + mask: Bool[Tensor, "1 32 mult"], + ) -> Float[Tensor, "0"]: + """ + multiscale, scale-invariant gradient matching term to the disparity space. + This term biases discontinuities to be sharp and to coincide with discontinuities in the ground truth + More info here https://arxiv.org/pdf/1907.01341.pdf Equation 11 + Args: + prediction: predicted depth map + target: ground truth depth map + reduction: reduction function, either reduction_batch_based or reduction_image_based + Returns: + gradient loss based on reduction function + """ + summed_mask = torch.sum(mask, (1, 2)) + diff = prediction - target + diff = torch.mul(mask, diff) + + grad_x = torch.abs(diff[:, :, 1:] - diff[:, :, :-1]) + mask_x = torch.mul(mask[:, :, 1:], mask[:, :, :-1]) + grad_x = torch.mul(mask_x, grad_x) + + grad_y = torch.abs(diff[:, 1:, :] - diff[:, :-1, :]) + mask_y = torch.mul(mask[:, 1:, :], mask[:, :-1, :]) + grad_y = torch.mul(mask_y, grad_y) + + image_loss = torch.sum(grad_x, (1, 2)) + torch.sum(grad_y, (1, 2)) + image_loss = masked_reduction(image_loss, summed_mask, self.reduction_type) + + return image_loss + + +class ScaleAndShiftInvariantLoss(nn.Module): + """ + Scale and shift invariant loss as described in + "Towards Robust Monocular Depth Estimation: Mixing Datasets for Zero-shot Cross-dataset Transfer" + https://arxiv.org/pdf/1907.01341.pdf + """ + + def __init__(self, alpha: float = 0.5, scales: int = 4, reduction_type: Literal["image", "batch"] = "batch"): + """ + Args: + alpha: weight of the regularization term + scales: number of scales to use + reduction_type: either "batch" or "image" + """ + super().__init__() + self.__data_loss = MiDaSMSELoss(reduction_type=reduction_type) + self.__regularization_loss = GradientLoss(scales=scales, reduction_type=reduction_type) + self.__alpha = alpha + + self.__prediction_ssi = None + + def forward( + self, + prediction: Float[Tensor, "1 32 mult"], + target: Float[Tensor, "1 32 mult"], + mask: Bool[Tensor, "1 32 mult"], + ) -> Float[Tensor, "0"]: + """ + Args: + prediction: predicted depth map (unnormalized) + target: ground truth depth map (normalized) + mask: mask of valid pixels + Returns: + scale and shift invariant loss + """ + scale, shift = normalized_depth_scale_and_shift(prediction, target, mask) + self.__prediction_ssi = scale.view(-1, 1, 1) * prediction + shift.view(-1, 1, 1) + + total = self.__data_loss(self.__prediction_ssi, target, mask) + # if self.__alpha > 0: + # total += self.__alpha * self.__regularization_loss(self.__prediction_ssi, target, mask) + + return total + + def __get_prediction_ssi(self): + """ + scale and shift invariant prediction + from https://arxiv.org/pdf/1907.01341.pdf equation 1 + """ + return self.__prediction_ssi + + prediction_ssi = property(__get_prediction_ssi) \ No newline at end of file diff --git a/code/utils/semantic_utils.py b/code/utils/semantic_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c05154eae2221da344aa38f49a5488ee85bc35ef --- /dev/null +++ b/code/utils/semantic_utils.py @@ -0,0 +1,222 @@ +#!/usr/bin/python +# +# KITTI-360 labels +# + +from collections import namedtuple +from PIL import Image +import numpy as np + + +#-------------------------------------------------------------------------------- +# Definitions +#-------------------------------------------------------------------------------- + +# a label and all meta information +Label = namedtuple( 'Label' , [ + + 'name' , # The identifier of this label, e.g. 'car', 'person', ... . + # We use them to uniquely name a class + + 'id' , # An integer ID that is associated with this label. + # The IDs are used to represent the label in ground truth images + # An ID of -1 means that this label does not have an ID and thus + # is ignored when creating ground truth images (e.g. license plate). + # Do not modify these IDs, since exactly these IDs are expected by the + # evaluation server. + + 'trainId' , # Feel free to modify these IDs as suitable for your method. Then create + # ground truth images with train IDs, using the tools provided in the + # 'preparation' folder. However, make sure to validate or submit results + # to our evaluation server using the regular IDs above! + # For trainIds, multiple labels might have the same ID. Then, these labels + # are mapped to the same class in the ground truth images. For the inverse + # mapping, we use the label that is defined first in the list below. + # For example, mapping all void-type classes to the same ID in training, + # might make sense for some approaches. + # Max value is 255! + + 'category' , # The name of the category that this label belongs to + + 'categoryId' , # The ID of this category. Used to create ground truth images + # on category level. + + 'hasInstances', # Whether this label distinguishes between single instances or not + + 'ignoreInEval', # Whether pixels having this class as ground truth label are ignored + # during evaluations or not + + 'color' , # The color of this label + ] ) + + +#-------------------------------------------------------------------------------- +# A list of all labels +#-------------------------------------------------------------------------------- + +# Please adapt the train IDs as appropriate for your approach. +# Note that you might want to ignore labels with ID 255 during training. +# Further note that the current train IDs are only a suggestion. You can use whatever you like. +# Make sure to provide your results using the original IDs and not the training IDs. +# Note that many IDs are ignored in evaluation and thus you never need to predict these! + +labels = [ + # name id trainId category catId hasInstances ignoreInEval color + Label( 'unlabeled' , 0 , 255 , 'void' , 0 , False , True , ( 0, 0, 0) ), + Label( 'ego vehicle' , 1 , 255 , 'void' , 0 , False , True , ( 0, 0, 0) ), + Label( 'rectification border' , 2 , 255 , 'void' , 0 , False , True , ( 0, 0, 0) ), + Label( 'out of roi' , 3 , 255 , 'void' , 0 , False , True , ( 0, 0, 0) ), + Label( 'static' , 4 , 255 , 'void' , 0 , False , True , ( 0, 0, 0) ), + Label( 'dynamic' , 5 , 255 , 'void' , 0 , False , True , (111, 74, 0) ), + Label( 'ground' , 6 , 255 , 'void' , 0 , False , True , ( 81, 0, 81) ), + Label( 'road' , 7 , 0 , 'flat' , 1 , False , False , (128, 64,128) ), + Label( 'sidewalk' , 8 , 1 , 'flat' , 1 , False , False , (244, 35,232) ), + Label( 'parking' , 9 , 255 , 'flat' , 1 , False , True , (250,170,160) ), + Label( 'rail track' , 10 , 255 , 'flat' , 1 , False , True , (230,150,140) ), + Label( 'building' , 11 , 2 , 'construction' , 2 , False , False , ( 70, 70, 70) ), + Label( 'wall' , 12 , 3 , 'construction' , 2 , False , False , (102,102,156) ), + Label( 'fence' , 13 , 4 , 'construction' , 2 , False , False , (190,153,153) ), + Label( 'guard rail' , 14 , 255 , 'construction' , 2 , False , True , (180,165,180) ), + Label( 'bridge' , 15 , 255 , 'construction' , 2 , False , True , (150,100,100) ), + Label( 'tunnel' , 16 , 255 , 'construction' , 2 , False , True , (150,120, 90) ), + Label( 'pole' , 17 , 5 , 'object' , 3 , False , False , (153,153,153) ), + Label( 'polegroup' , 18 , 255 , 'object' , 3 , False , True , (153,153,153) ), + Label( 'traffic light' , 19 , 6 , 'object' , 3 , False , False , (250,170, 30) ), + Label( 'traffic sign' , 20 , 7 , 'object' , 3 , False , False , (220,220, 0) ), + Label( 'vegetation' , 21 , 8 , 'nature' , 4 , False , False , (107,142, 35) ), + Label( 'terrain' , 22 , 9 , 'nature' , 4 , False , False , (152,251,152) ), + Label( 'sky' , 23 , 10 , 'sky' , 5 , False , False , ( 70,130,180) ), + Label( 'person' , 24 , 11 , 'human' , 6 , True , False , (220, 20, 60) ), + Label( 'rider' , 25 , 12 , 'human' , 6 , True , False , (255, 0, 0) ), + Label( 'car' , 26 , 13 , 'vehicle' , 7 , True , False , ( 0, 0,142) ), + Label( 'truck' , 27 , 14 , 'vehicle' , 7 , True , False , ( 0, 0, 70) ), + Label( 'bus' , 28 , 15 , 'vehicle' , 7 , True , False , ( 0, 60,100) ), + Label( 'caravan' , 29 , 255 , 'vehicle' , 7 , True , True , ( 0, 0, 90) ), + Label( 'trailer' , 30 , 255 , 'vehicle' , 7 , True , True , ( 0, 0,110) ), + Label( 'train' , 31 , 16 , 'vehicle' , 7 , True , False , ( 0, 80,100) ), + Label( 'motorcycle' , 32 , 17 , 'vehicle' , 7 , True , False , ( 0, 0,230) ), + Label( 'bicycle' , 33 , 18 , 'vehicle' , 7 , True , False , (119, 11, 32) ), + Label( 'license plate' , -1 , -1 , 'vehicle' , 7 , False , True , ( 0, 0,142) ), +] + + +#-------------------------------------------------------------------------------- +# Create dictionaries for a fast lookup +#-------------------------------------------------------------------------------- + +# Please refer to the main method below for example usages! + +# name to label object +name2label = { label.name : label for label in labels } +# id to label object +id2label = { label.id : label for label in labels } +# trainId to label object +trainId2label = { label.trainId : label for label in reversed(labels) } +# label2trainid +label2trainid = { label.id : label.trainId for label in labels } +# trainId to label object +trainId2name = { label.trainId : label.name for label in labels } +trainId2color = { label.trainId : label.color for label in labels } +# category to list of label objects +category2labels = {} +for label in labels: + category = label.category + if category in category2labels: + category2labels[category].append(label) + else: + category2labels[category] = [label] + +#-------------------------------------------------------------------------------- +# color mapping +#-------------------------------------------------------------------------------- + +palette = [128, 64, 128, + 244, 35, 232, + 70, 70, 70, + 102, 102, 156, + 190, 153, 153, + 153, 153, 153, + 250, 170, 30, + 220, 220, 0, + 107, 142, 35, + 152, 251, 152, + 70, 130, 180, + 220, 20, 60, + 255, 0, 0, + 0, 0, 142, + 0, 0, 70, + 0, 60, 100, + 0, 80, 100, + 0, 0, 230, + 119, 11, 32] +zero_pad = 256 * 3 - len(palette) +for i in range(zero_pad): + palette.append(0) +color_mapping = palette + +def colorize(image_array): + new_mask = Image.fromarray(image_array.astype(np.uint8)).convert('P') + new_mask.putpalette(color_mapping) + return new_mask + +#-------------------------------------------------------------------------------- +# Assure single instance name +#-------------------------------------------------------------------------------- + +# returns the label name that describes a single instance (if possible) +# e.g. input | output +# ---------------------- +# car | car +# cargroup | car +# foo | None +# foogroup | None +# skygroup | None +def assureSingleInstanceName( name ): + # if the name is known, it is not a group + if name in name2label: + return name + # test if the name actually denotes a group + if not name.endswith("group"): + return None + # remove group + name = name[:-len("group")] + # test if the new name exists + if not name in name2label: + return None + # test if the new name denotes a label that actually has instances + if not name2label[name].hasInstances: + return None + # all good then + return name + +#-------------------------------------------------------------------------------- +# Main for testing +#-------------------------------------------------------------------------------- + +# just a dummy main +if __name__ == "__main__": + # Print all the labels + print("List of KITTI-360 labels:") + print("") + print(" {:>21} | {:>3} | {:>7} | {:>14} | {:>10} | {:>12} | {:>12}".format( 'name', 'id', 'trainId', 'category', 'categoryId', 'hasInstances', 'ignoreInEval' )) + print(" " + ('-' * 98)) + for label in labels: + # print(" {:>21} | {:>3} | {:>7} | {:>14} | {:>10} | {:>12} | {:>12}".format( label.name, label.id, label.trainId, label.category, label.categoryId, label.hasInstances, label.ignoreInEval )) + print(" \"{:}\"".format(label.name)) + print("") + + print("Example usages:") + + # Map from name to label + name = 'car' + id = name2label[name].id + print("ID of label '{name}': {id}".format( name=name, id=id )) + + # Map from ID to label + category = id2label[id].category + print("Category of label with ID '{id}': {category}".format( id=id, category=category )) + + # Map from trainID to label + trainId = 0 + name = trainId2label[trainId].name + print("Name of label with trainID '{id}': {name}".format( id=trainId, name=name )) \ No newline at end of file diff --git a/code/utils/sh_utils.py b/code/utils/sh_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..bbca7d192aa3a7edf8c5b2d24dee535eac765785 --- /dev/null +++ b/code/utils/sh_utils.py @@ -0,0 +1,118 @@ +# Copyright 2021 The PlenOctree Authors. +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. + +import torch + +C0 = 0.28209479177387814 +C1 = 0.4886025119029199 +C2 = [ + 1.0925484305920792, + -1.0925484305920792, + 0.31539156525252005, + -1.0925484305920792, + 0.5462742152960396 +] +C3 = [ + -0.5900435899266435, + 2.890611442640554, + -0.4570457994644658, + 0.3731763325901154, + -0.4570457994644658, + 1.445305721320277, + -0.5900435899266435 +] +C4 = [ + 2.5033429417967046, + -1.7701307697799304, + 0.9461746957575601, + -0.6690465435572892, + 0.10578554691520431, + -0.6690465435572892, + 0.47308734787878004, + -1.7701307697799304, + 0.6258357354491761, +] + + +def eval_sh(deg, sh, dirs): + """ + Evaluate spherical harmonics at unit directions + using hardcoded SH polynomials. + Works with torch/np/jnp. + ... Can be 0 or more batch dimensions. + Args: + deg: int SH deg. Currently, 0-3 supported + sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2] + dirs: jnp.ndarray unit directions [..., 3] + Returns: + [..., C] + """ + assert deg <= 4 and deg >= 0 + coeff = (deg + 1) ** 2 + assert sh.shape[-1] >= coeff + + result = C0 * sh[..., 0] + if deg > 0: + x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3] + result = (result - + C1 * y * sh[..., 1] + + C1 * z * sh[..., 2] - + C1 * x * sh[..., 3]) + + if deg > 1: + xx, yy, zz = x * x, y * y, z * z + xy, yz, xz = x * y, y * z, x * z + result = (result + + C2[0] * xy * sh[..., 4] + + C2[1] * yz * sh[..., 5] + + C2[2] * (2.0 * zz - xx - yy) * sh[..., 6] + + C2[3] * xz * sh[..., 7] + + C2[4] * (xx - yy) * sh[..., 8]) + + if deg > 2: + result = (result + + C3[0] * y * (3 * xx - yy) * sh[..., 9] + + C3[1] * xy * z * sh[..., 10] + + C3[2] * y * (4 * zz - xx - yy)* sh[..., 11] + + C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12] + + C3[4] * x * (4 * zz - xx - yy) * sh[..., 13] + + C3[5] * z * (xx - yy) * sh[..., 14] + + C3[6] * x * (xx - 3 * yy) * sh[..., 15]) + + if deg > 3: + result = (result + C4[0] * xy * (xx - yy) * sh[..., 16] + + C4[1] * yz * (3 * xx - yy) * sh[..., 17] + + C4[2] * xy * (7 * zz - 1) * sh[..., 18] + + C4[3] * yz * (7 * zz - 3) * sh[..., 19] + + C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20] + + C4[5] * xz * (7 * zz - 3) * sh[..., 21] + + C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22] + + C4[7] * xz * (xx - 3 * yy) * sh[..., 23] + + C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) * sh[..., 24]) + return result + +def RGB2SH(rgb): + return (rgb - 0.5) / C0 + +def SH2RGB(sh): + return sh * C0 + 0.5 \ No newline at end of file diff --git a/code/utils/system_utils.py b/code/utils/system_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..533a99a14e99626abe78cd901dba170142709628 --- /dev/null +++ b/code/utils/system_utils.py @@ -0,0 +1,28 @@ +# +# Copyright (C) 2023, Inria +# GRAPHDECO research group, https://team.inria.fr/graphdeco +# All rights reserved. +# +# This software is free for non-commercial, research and evaluation use +# under the terms of the LICENSE.md file. +# +# For inquiries contact george.drettakis@inria.fr +# + +from errno import EEXIST +from os import makedirs, path +import os + +def mkdir_p(folder_path): + # Creates a directory. equivalent to using mkdir -p on the command line + try: + makedirs(folder_path) + except OSError as exc: # Python >2.5 + if exc.errno == EEXIST and path.isdir(folder_path): + pass + else: + raise + +def searchForMaxIteration(folder): + saved_iters = [int(fname.split("_")[-1]) for fname in os.listdir(folder)] + return max(saved_iters) \ No newline at end of file diff --git a/code/web_server_config/kinematic.yaml b/code/web_server_config/kinematic.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8d2eb2866568aa0d5e63471d18da5d882e11e254 --- /dev/null +++ b/code/web_server_config/kinematic.yaml @@ -0,0 +1,9 @@ +min_steer: 15 +max_steer: 15 +min_acc: -2.0 +max_acc: 2.0 +fric: 0.2 +Lr: 1.35 +Lf: 1.35 +max_speed: 0.5 +dt: 0.25 \ No newline at end of file diff --git a/code/web_server_config/kitti360_base.yaml b/code/web_server_config/kitti360_base.yaml new file mode 100644 index 0000000000000000000000000000000000000000..08449b626d042c006e14cf2d34c72b4c94eb2832 --- /dev/null +++ b/code/web_server_config/kitti360_base.yaml @@ -0,0 +1,3 @@ +realcar_path: /app/app_datas/3DRealCar +model_base: /app/app_datas/ss/scenes/kitti360 +output_dir: /app/app_datas/outputs/benchmark/out/kitti360_ \ No newline at end of file diff --git a/code/web_server_config/kitti360_camera.yaml b/code/web_server_config/kitti360_camera.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d87ad9e36c42dff640138dfd38405ce56b3b6bab --- /dev/null +++ b/code/web_server_config/kitti360_camera.yaml @@ -0,0 +1,166 @@ +cams: + CAM_BACK: + extrinsics: + l2c_rot: + - 92.25568216952541 + - 0.49511589981751 + - -179.5254516169236 + l2c_trans: + - 0.0008379812089021367 + - -0.27867326994260555 + - -0.9287049615140774 + v2c_rot: + - 89.60442534979997 + - 0.6418405371350101 + - -89.47195526949999 + v2c_trans: + - -0.028860307553156085 + - 1.5680612963127507 + - 0.0443194618131316 + intrinsics: + H: 450 + W: 800 + cx: 428.8887163431848 + cy: 238.44244942037076 + fovx: 90.22309438895985 + fovy: 58.9063564284116 + CAM_BACK_LEFT: + extrinsics: + l2c_rot: + - 89.82295180299249 + - 2.3155403060452167 + - -108.62593888549094 + l2c_trans: + - -0.21405828773135127 + - -0.28654507725108536 + - -0.43206949836662323 + v2c_rot: + - 88.81655348310001 + - -0.14162822446102205 + - -18.57965602679999 + v2c_trans: + - -1.1439217599288067 + - 1.5620432465039058 + - -0.15609448627566294 + intrinsics: + H: 450 + W: 800 + cx: 414.78846668154955 + cy: 233.58402809319935 + fovx: 65.03154053626905 + fovy: 39.45259285521372 + CAM_BACK_RIGHT: + extrinsics: + l2c_rot: + - 90.1950785524308 + - -2.6262664650943512 + - 112.51452117672765 + l2c_trans: + - 0.24641784245418819 + - -0.2948173619480623 + - -0.3999756259757112 + v2c_rot: + - 89.33706258359997 + - -0.11343093253300289 + - -157.478809169 + v2c_trans: + - 1.1606735483944377 + - 1.5478119044158865 + - -0.043687868897601415 + intrinsics: + H: 450 + W: 800 + cx: 412.6884022687992 + cy: 231.27408192854378 + fovx: 65.24002919568859 + fovy: 39.59883582992881 + CAM_FRONT: + extrinsics: + l2c_rot: + - 86.86166850342303 + - -0.5023434125361063 + - -0.6429117094019539 + l2c_trans: + - -0.007022340992421169 + - -0.3515059821513211 + - -0.7332408994883366 + v2c_rot: + - 89.50957761160001 + - -0.700717347395009 + - 89.4120791145 + v2c_trans: + - 0.00536785836469495 + - 1.4801653630391234 + - -1.7346966604269407 + intrinsics: + H: 450 + W: 800 + cx: 413.294057390699 + cy: 234.99233131122904 + fovx: 65.12158487224077 + fovy: 39.515728607991626 + CAM_FRONT_LEFT: + extrinsics: + l2c_rot: + - 87.64723817274033 + - 2.0461415281358453 + - -55.34666714311071 + l2c_trans: + - -0.20156170366177595 + - -0.34522197648384223 + - -0.7410957430382523 + v2c_rot: + - 89.0143474659 + - -0.23019817163998718 + - 34.65885052869996 + v2c_trans: + - -1.0050284941954049 + - 1.4882946049903183 + - -1.3333331421595005 + intrinsics: + H: 450 + W: 800 + cx: 413.6205315547843 + cy: 225.457749102887 + fovx: 64.91268361861167 + fovy: 39.3693134188315 + CAM_FRONT_RIGHT: + extrinsics: + l2c_rot: + - 88.10298835347594 + - -3.051668588782139 + - 57.49277888679618 + l2c_trans: + - 0.22950989625543705 + - -0.33063487372538725 + - -0.7384322069296176 + v2c_rot: + - 89.66956817789999 + - -0.9071318410000028 + - 147.574088459 + v2c_trans: + - 1.0906326784058709 + - 1.4930721654490449 + - -1.2775493224769925 + intrinsics: + H: 450 + W: 800 + cx: 408.8943785479856 + cy: 225.97708900475635 + fovx: 64.95866942785656 + fovy: 39.40152637651543 +cam_rect: + rot: + - 5.0 + - 0.0 + - 0.0 + trans: + - 0.0 + - 0.0 + - 0.0 + +cam_align: + [ + ['CAM_FRONT_LEFT', 'CAM_FRONT', 'CAM_FRONT_RIGHT'], + ['CAM_BACK_RIGHT', 'CAM_BACK', 'CAM_BACK_LEFT'], + ] \ No newline at end of file diff --git a/code/web_server_config/nuscenes_base.yaml b/code/web_server_config/nuscenes_base.yaml new file mode 100644 index 0000000000000000000000000000000000000000..216a0a2dca73d664590e1e529eb7d454796cc4c3 --- /dev/null +++ b/code/web_server_config/nuscenes_base.yaml @@ -0,0 +1,6 @@ +realcar_path: /app/app_datas/3DRealCar +model_base: /app/app_datas/ss/scenes/nuscenes +output_dir: /app/app_datas/outputs/benchmark/out/nusc_ +HD_map: + path: /app/app_datas/datasets/nuScenes/raw/Trainval + version: nusc_trainval diff --git a/code/web_server_config/nuscenes_camera.yaml b/code/web_server_config/nuscenes_camera.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a116c1a3302a9021ae66480012f5510d5f57a92f --- /dev/null +++ b/code/web_server_config/nuscenes_camera.yaml @@ -0,0 +1,166 @@ +cams: + CAM_BACK: + extrinsics: + l2c_rot: + - 92.25568216952541 + - 0.49511589981751 + - -179.5254516169236 + l2c_trans: + - 0.0008379812089021367 + - -0.27867326994260555 + - -0.9287049615140774 + v2c_rot: + - 89.60442534979997 + - 0.6418405371350101 + - -89.47195526949999 + v2c_trans: + - -0.028860307553156085 + - 1.5680612963127507 + - 0.0443194618131316 + intrinsics: + H: 450 + W: 800 + cx: 428.8887163431848 + cy: 238.44244942037076 + fovx: 90.22309438895985 + fovy: 58.9063564284116 + CAM_BACK_LEFT: + extrinsics: + l2c_rot: + - 89.82295180299249 + - 2.3155403060452167 + - -108.62593888549094 + l2c_trans: + - -0.21405828773135127 + - -0.28654507725108536 + - -0.43206949836662323 + v2c_rot: + - 88.81655348310001 + - -0.14162822446102205 + - -18.57965602679999 + v2c_trans: + - -1.1439217599288067 + - 1.5620432465039058 + - -0.15609448627566294 + intrinsics: + H: 450 + W: 800 + cx: 414.78846668154955 + cy: 233.58402809319935 + fovx: 65.03154053626905 + fovy: 39.45259285521372 + CAM_BACK_RIGHT: + extrinsics: + l2c_rot: + - 90.1950785524308 + - -2.6262664650943512 + - 112.51452117672765 + l2c_trans: + - 0.24641784245418819 + - -0.2948173619480623 + - -0.3999756259757112 + v2c_rot: + - 89.33706258359997 + - -0.11343093253300289 + - -157.478809169 + v2c_trans: + - 1.1606735483944377 + - 1.5478119044158865 + - -0.043687868897601415 + intrinsics: + H: 450 + W: 800 + cx: 412.6884022687992 + cy: 231.27408192854378 + fovx: 65.24002919568859 + fovy: 39.59883582992881 + CAM_FRONT: + extrinsics: + l2c_rot: + - 86.86166850342303 + - -0.5023434125361063 + - -0.6429117094019539 + l2c_trans: + - -0.007022340992421169 + - -0.3515059821513211 + - -0.7332408994883366 + v2c_rot: + - 89.50957761160001 + - -0.700717347395009 + - 89.4120791145 + v2c_trans: + - 0.00536785836469495 + - 1.4801653630391234 + - -1.7346966604269407 + intrinsics: + H: 450 + W: 800 + cx: 413.294057390699 + cy: 234.99233131122904 + fovx: 65.12158487224077 + fovy: 39.515728607991626 + CAM_FRONT_LEFT: + extrinsics: + l2c_rot: + - 87.64723817274033 + - 2.0461415281358453 + - -55.34666714311071 + l2c_trans: + - -0.20156170366177595 + - -0.34522197648384223 + - -0.7410957430382523 + v2c_rot: + - 89.0143474659 + - -0.23019817163998718 + - 34.65885052869996 + v2c_trans: + - -1.0050284941954049 + - 1.4882946049903183 + - -1.3333331421595005 + intrinsics: + H: 450 + W: 800 + cx: 413.6205315547843 + cy: 225.457749102887 + fovx: 64.91268361861167 + fovy: 39.3693134188315 + CAM_FRONT_RIGHT: + extrinsics: + l2c_rot: + - 88.10298835347594 + - -3.051668588782139 + - 57.49277888679618 + l2c_trans: + - 0.22950989625543705 + - -0.33063487372538725 + - -0.7384322069296176 + v2c_rot: + - 89.66956817789999 + - -0.9071318410000028 + - 147.574088459 + v2c_trans: + - 1.0906326784058709 + - 1.4930721654490449 + - -1.2775493224769925 + intrinsics: + H: 450 + W: 800 + cx: 408.8943785479856 + cy: 225.97708900475635 + fovx: 64.95866942785656 + fovy: 39.40152637651543 +cam_rect: + rot: + - 0.0 + - 0.0 + - 0.0 + trans: + - 0.0 + - 0.3 + - 0.0 + +cam_align: + [ + ['CAM_FRONT_LEFT', 'CAM_FRONT', 'CAM_FRONT_RIGHT'], + ['CAM_BACK_RIGHT', 'CAM_BACK', 'CAM_BACK_LEFT'], + ] \ No newline at end of file diff --git a/code/web_server_config/pandaset_base.yaml b/code/web_server_config/pandaset_base.yaml new file mode 100644 index 0000000000000000000000000000000000000000..eb0e773ec4d909e8383bfd2ee775cc6611da700a --- /dev/null +++ b/code/web_server_config/pandaset_base.yaml @@ -0,0 +1,3 @@ +realcar_path: /app/app_datas/3DRealCar +model_base: /app/app_datas/ss/scenes/pandaset +output_dir: /app/app_datas/outputs/benchmark/out/pandaset_ \ No newline at end of file diff --git a/code/web_server_config/pandaset_camera.yaml b/code/web_server_config/pandaset_camera.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e5eff0fef0c72c0e54464428a2018ad8faadd741 --- /dev/null +++ b/code/web_server_config/pandaset_camera.yaml @@ -0,0 +1,166 @@ +cams: + CAM_BACK: + extrinsics: + l2c_rot: + - 92.25568216952541 + - 0.49511589981751 + - -179.5254516169236 + l2c_trans: + - 0.0008379812089021367 + - -0.27867326994260555 + - -0.9287049615140774 + v2c_rot: + - 89.60442534979997 + - 0.6418405371350101 + - -89.47195526949999 + v2c_trans: + - -0.028860307553156085 + - 1.5680612963127507 + - 0.0443194618131316 + intrinsics: + H: 450 + W: 800 + cx: 428.8887163431848 + cy: 238.44244942037076 + fovx: 80.22309438895985 + fovy: 50.9063564284116 + CAM_BACK_LEFT: + extrinsics: + l2c_rot: + - 89.82295180299249 + - 2.3155403060452167 + - -108.62593888549094 + l2c_trans: + - -0.21405828773135127 + - -0.28654507725108536 + - -0.43206949836662323 + v2c_rot: + - 88.81655348310001 + - -0.14162822446102205 + - -18.57965602679999 + v2c_trans: + - -1.1439217599288067 + - 1.5620432465039058 + - -0.15609448627566294 + intrinsics: + H: 450 + W: 800 + cx: 414.78846668154955 + cy: 233.58402809319935 + fovx: 58.03154053626905 + fovy: 32.45259285521372 + CAM_BACK_RIGHT: + extrinsics: + l2c_rot: + - 90.1950785524308 + - -2.6262664650943512 + - 112.51452117672765 + l2c_trans: + - 0.24641784245418819 + - -0.2948173619480623 + - -0.3999756259757112 + v2c_rot: + - 89.33706258359997 + - -0.11343093253300289 + - -157.478809169 + v2c_trans: + - 1.1606735483944377 + - 1.5478119044158865 + - -0.043687868897601415 + intrinsics: + H: 450 + W: 800 + cx: 412.6884022687992 + cy: 231.27408192854378 + fovx: 58.24002919568859 + fovy: 32.59883582992881 + CAM_FRONT: + extrinsics: + l2c_rot: + - 86.86166850342303 + - -0.5023434125361063 + - -0.6429117094019539 + l2c_trans: + - -0.007022340992421169 + - -0.3515059821513211 + - -0.7332408994883366 + v2c_rot: + - 89.50957761160001 + - -0.700717347395009 + - 89.4120791145 + v2c_trans: + - 0.00536785836469495 + - 1.4801653630391234 + - -1.7346966604269407 + intrinsics: + H: 450 + W: 800 + cx: 413.294057390699 + cy: 234.99233131122904 + fovx: 58.12158487224077 + fovy: 32.515728607991626 + CAM_FRONT_LEFT: + extrinsics: + l2c_rot: + - 87.64723817274033 + - 2.0461415281358453 + - -55.34666714311071 + l2c_trans: + - -0.20156170366177595 + - -0.34522197648384223 + - -0.7410957430382523 + v2c_rot: + - 89.0143474659 + - -0.23019817163998718 + - 34.65885052869996 + v2c_trans: + - -1.0050284941954049 + - 1.4882946049903183 + - -1.3333331421595005 + intrinsics: + H: 450 + W: 800 + cx: 413.6205315547843 + cy: 225.457749102887 + fovx: 58.91268361861167 + fovy: 32.3693134188315 + CAM_FRONT_RIGHT: + extrinsics: + l2c_rot: + - 88.10298835347594 + - -3.051668588782139 + - 57.49277888679618 + l2c_trans: + - 0.22950989625543705 + - -0.33063487372538725 + - -0.7384322069296176 + v2c_rot: + - 89.66956817789999 + - -0.9071318410000028 + - 147.574088459 + v2c_trans: + - 1.0906326784058709 + - 1.4930721654490449 + - -1.2775493224769925 + intrinsics: + H: 450 + W: 800 + cx: 408.8943785479856 + cy: 225.97708900475635 + fovx: 58.95866942785656 + fovy: 32.40152637651543 +cam_rect: + rot: + - 0.0 + - 0.0 + - 0.0 + trans: + - 0.0 + - 0.7 + - 0.0 + +cam_align: + [ + ['CAM_FRONT_LEFT', 'CAM_FRONT', 'CAM_FRONT_RIGHT'], + ['CAM_BACK_RIGHT', 'CAM_BACK', 'CAM_BACK_LEFT'], + ] \ No newline at end of file diff --git a/docker/web_server_config/scene-0383-medium-00.yaml b/code/web_server_config/scene-0383-medium-00.yaml similarity index 100% rename from docker/web_server_config/scene-0383-medium-00.yaml rename to code/web_server_config/scene-0383-medium-00.yaml diff --git a/code/web_server_config/waymo_base.yaml b/code/web_server_config/waymo_base.yaml new file mode 100644 index 0000000000000000000000000000000000000000..945d938d5b0f4a19f7c70c4e49f026dc404f1fd0 --- /dev/null +++ b/code/web_server_config/waymo_base.yaml @@ -0,0 +1,3 @@ +realcar_path: /app/app_datas/3DRealCar +model_base: /app/app_datas/ss/scenes/waymo +output_dir: /app/app_datas/outputs/benchmark/out/waymo_ \ No newline at end of file diff --git a/code/web_server_config/waymo_camera.yaml b/code/web_server_config/waymo_camera.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f03fcd947943b301703a6098d240195afa7909d2 --- /dev/null +++ b/code/web_server_config/waymo_camera.yaml @@ -0,0 +1,166 @@ +cams: + CAM_BACK: + extrinsics: + l2c_rot: + - 92.25568216952541 + - 0.49511589981751 + - -179.5254516169236 + l2c_trans: + - 0.0008379812089021367 + - -0.27867326994260555 + - -0.9287049615140774 + v2c_rot: + - 89.60442534979997 + - 0.6418405371350101 + - -89.47195526949999 + v2c_trans: + - -0.028860307553156085 + - 1.5680612963127507 + - 0.0443194618131316 + intrinsics: + H: 450 + W: 800 + cx: 428.8887163431848 + cy: 238.44244942037076 + fovx: 82.22309438895985 + fovy: 50.9063564284116 + CAM_BACK_LEFT: + extrinsics: + l2c_rot: + - 89.82295180299249 + - 2.3155403060452167 + - -108.62593888549094 + l2c_trans: + - -0.21405828773135127 + - -0.28654507725108536 + - -0.43206949836662323 + v2c_rot: + - 88.81655348310001 + - -0.14162822446102205 + - -18.57965602679999 + v2c_trans: + - -1.1439217599288067 + - 1.5620432465039058 + - -0.15609448627566294 + intrinsics: + H: 450 + W: 800 + cx: 414.78846668154955 + cy: 233.58402809319935 + fovx: 58.03154053626905 + fovy: 32.45259285521372 + CAM_BACK_RIGHT: + extrinsics: + l2c_rot: + - 90.1950785524308 + - -2.6262664650943512 + - 112.51452117672765 + l2c_trans: + - 0.24641784245418819 + - -0.2948173619480623 + - -0.3999756259757112 + v2c_rot: + - 89.33706258359997 + - -0.11343093253300289 + - -157.478809169 + v2c_trans: + - 1.1606735483944377 + - 1.5478119044158865 + - -0.043687868897601415 + intrinsics: + H: 450 + W: 800 + cx: 412.6884022687992 + cy: 231.27408192854378 + fovx: 58.24002919568859 + fovy: 32.59883582992881 + CAM_FRONT: + extrinsics: + l2c_rot: + - 86.86166850342303 + - -0.5023434125361063 + - -0.6429117094019539 + l2c_trans: + - -0.007022340992421169 + - -0.3515059821513211 + - -0.7332408994883366 + v2c_rot: + - 89.50957761160001 + - -0.700717347395009 + - 89.4120791145 + v2c_trans: + - 0.00536785836469495 + - 1.4801653630391234 + - -1.7346966604269407 + intrinsics: + H: 450 + W: 800 + cx: 413.294057390699 + cy: 234.99233131122904 + fovx: 58.12158487224077 + fovy: 32.515728607991626 + CAM_FRONT_LEFT: + extrinsics: + l2c_rot: + - 87.64723817274033 + - 2.0461415281358453 + - -55.34666714311071 + l2c_trans: + - -0.20156170366177595 + - -0.34522197648384223 + - -0.7410957430382523 + v2c_rot: + - 89.0143474659 + - -0.23019817163998718 + - 34.65885052869996 + v2c_trans: + - -1.0050284941954049 + - 1.4882946049903183 + - -1.3333331421595005 + intrinsics: + H: 450 + W: 800 + cx: 413.6205315547843 + cy: 225.457749102887 + fovx: 58.91268361861167 + fovy: 32.3693134188315 + CAM_FRONT_RIGHT: + extrinsics: + l2c_rot: + - 88.10298835347594 + - -3.051668588782139 + - 57.49277888679618 + l2c_trans: + - 0.22950989625543705 + - -0.33063487372538725 + - -0.7384322069296176 + v2c_rot: + - 89.66956817789999 + - -0.9071318410000028 + - 147.574088459 + v2c_trans: + - 1.0906326784058709 + - 1.4930721654490449 + - -1.2775493224769925 + intrinsics: + H: 450 + W: 800 + cx: 408.8943785479856 + cy: 225.97708900475635 + fovx: 58.95866942785656 + fovy: 32.40152637651543 +cam_rect: + rot: + - 0.0 + - 0.0 + - 0.0 + trans: + - 0.0 + - 0.3 + - 0.0 + +cam_align: + [ + ['CAM_FRONT_LEFT', 'CAM_FRONT', 'CAM_FRONT_RIGHT'], + ['CAM_BACK_RIGHT', 'CAM_BACK', 'CAM_BACK_LEFT'], + ] \ No newline at end of file diff --git a/download_pre_datas.py b/download_pre_datas.py index 0a7e207d814e89d517dab0930484424419e98758..507fbfcf66a5513ffd971695a3abca67a0f03a41 100644 --- a/download_pre_datas.py +++ b/download_pre_datas.py @@ -4,11 +4,16 @@ from pathlib import Path from huggingface_hub import snapshot_download -snapshot_download(repo_id='XDimLab/HUGSIM',revision='main',local_dir='/app/app_datas/PAMI2024/release/',local_dir_use_symlinks=False,allow_patterns=['3DRealCar/**'],repo_type='dataset') -snapshot_download(repo_id='XDimLab/HUGSIM',revision='main',local_dir='/app/app_datas/PAMI2024/release/ss',local_dir_use_symlinks=False,allow_patterns=['scenes/nuscenes/scene-0383.zip'], repo_type='dataset') - -for file in Path("/app/app_datas/PAMI2024/release/ss/scenes/nuscenes").rglob("*.zip"): +def unzip(file): file_path = file.as_posix() dir_path = file.parent.as_posix() with zipfile.ZipFile(file_path, 'r') as zip_ref: zip_ref.extractall(dir_path) + +snapshot_download(repo_id='XDimLab/HUGSIM',revision='main',local_dir='/app/app_datas/',local_dir_use_symlinks=False,allow_patterns=['3DRealCar/**'],repo_type='dataset') +snapshot_download(repo_id='XDimLab/HUGSIM',revision='main',local_dir='/app/app_datas/',local_dir_use_symlinks=False,allow_patterns=['nusc_map_cache.zip'],repo_type='dataset') +snapshot_download(repo_id='XDimLab/HUGSIM_private',revision='main',local_dir='/app/app_datas/ss/',local_dir_use_symlinks=False,allow_patterns=['scenes/**'], repo_type='dataset') +snapshot_download(repo_id='XDimLab/HUGSIM_private',revision='main',local_dir='/app/app_datas/ss/',local_dir_use_symlinks=False,allow_patterns=['scenarios.zip'], repo_type='dataset') + +unzip("/app/app_datas/nusc_map_cache.zip") +unzip("/app/app_datas/ss/scenarios.zip") \ No newline at end of file diff --git a/web_server.py b/web_server.py index c0e3e186cd8c9e5af1c90d63e64dde1cec4b30be..6fbd87a4982114d86e27911e30c755ea2a7fbe19 100644 --- a/web_server.py +++ b/web_server.py @@ -26,6 +26,7 @@ import gymnasium import uvicorn import psutil import torch +from glob import glob from sim.utils.sim_utils import traj2control, traj_transform_to_global from sim.utils.score_calculator import hugsim_evaluate @@ -204,12 +205,16 @@ class EnvHandler: if scene_index < 0 or scene_index >= len(self.scene_list): raise ValueError("Invalid scene index.") + self.close() self.cur_scene_index = scene_index scene_config = self.scene_list[scene_index] - self.close() - self.cur_otuput = os.path.join(self.base_output, scene_config.name) - os.makedirs(self.cur_otuput, exist_ok=True) - self.env = gymnasium.make('hugsim_env/HUGSim-v0', cfg=scene_config.cfg, output=self.cur_otuput) + self._log(f"Switch to scene: {scene_config.name}_{scene_config.cfg.scenario.mode}") + print(f"Switch to scene: {scene_config.name}_{scene_config.cfg.scenario.mode}") + + self.cur_output = os.path.join(self.base_output, + f"{scene_config.name}_{scene_config.cfg.scenario.mode}") + os.makedirs(self.cur_output, exist_ok=True) + self.env = gymnasium.make('hugsim_env/HUGSim-v0', cfg=scene_config.cfg, output=self.cur_output) self._scene_cnt = 0 self._scene_done = False self._save_data = {'type': 'closeloop', 'frames': []} @@ -383,31 +388,35 @@ class EnvHandlerManager: Returns: List[SceneConfig]: A list of scene configurations. """ - base_path = os.path.join(os.path.dirname(__file__), 'docker', "web_server_config", 'nuscenes_base.yaml') - scenario_path = os.path.join(os.path.dirname(__file__), 'docker', "web_server_config", 'scene-0383-medium-00.yaml') - camera_path = os.path.join(os.path.dirname(__file__), 'docker', "web_server_config", 'nuscenes_camera.yaml') - kinematic_path = os.path.join(os.path.dirname(__file__), 'docker', "web_server_config", 'kinematic.yaml') - - scenario_config = OmegaConf.load(scenario_path) - base_config = OmegaConf.load(base_path) - camera_config = OmegaConf.load(camera_path) - kinematic_config = OmegaConf.load(kinematic_path) - cfg = OmegaConf.merge( - {"scenario": scenario_config}, - {"base": base_config}, - {"camera": camera_config}, - {"kinematic": kinematic_config} - ) - - model_path = os.path.join(cfg.base.model_base, cfg.scenario.scene_name) - model_config = OmegaConf.load(os.path.join(model_path, 'cfg.yaml')) - model_config.update({"model_path": "/app/app_datas/PAMI2024/release/ss/scenes/nuscenes/scene-0383"}) - cfg.update(model_config) - cfg.base.output_dir = base_output - return [ - SceneConfig(name=cfg.scenario.scene_name, cfg=cfg), - SceneConfig(name="temp_mock_scene", cfg=cfg), - ] + scene_list = [] + for data_type in ['kitti360', 'waymo', 'nuscenes', 'pandaset']: + base_path = os.path.join(os.path.dirname(__file__), "web_server_config", f'{data_type}_base.yaml') + camera_path = os.path.join(os.path.dirname(__file__), "web_server_config", f'{data_type}_camera.yaml') + kinematic_path = os.path.join(os.path.dirname(__file__), "web_server_config", 'kinematic.yaml') + + base_config = OmegaConf.load(base_path) + camera_config = OmegaConf.load(camera_path) + kinematic_config = OmegaConf.load(kinematic_path) + + scenarios_list = glob(f"/app/app_datas/ss/scenarios/{data_type}/*.yaml") + + for scenario_path in scenarios_list: + scenario_config = OmegaConf.load(scenario_path) + cfg = OmegaConf.merge( + {"scenario": scenario_config}, + {"base": base_config}, + {"camera": camera_config}, + {"kinematic": kinematic_config} + ) + + model_path = os.path.join(cfg.base.model_base, cfg.scenario.scene_name) + model_config = OmegaConf.load(os.path.join(model_path, 'cfg.yaml')) + model_config.update({"model_path": f"/app/app_datas/ss/scenes/{data_type}/{cfg.scenario.scene_name}"}) + cfg.update(model_config) + cfg.base.output_dir = base_output + scene_list.append(SceneConfig(name=cfg.scenario.scene_name, cfg=cfg)) + + return scene_list def _generate_env_handler(self, env_id: str): base_output = "/app/app_datas/env_output"