├── .gitignore ├── .gitmodules ├── LICENSE ├── README.md ├── VERSION ├── app ├── anim │ ├── __init__.py │ ├── anim.py │ ├── file_anim.py │ └── reverse_ego_anim.py ├── loss │ ├── __init__.py │ ├── clearance.py │ ├── color_lipshitz.py │ ├── conditional.py │ ├── eikonal.py │ ├── flow.py │ ├── lidar.py │ ├── mahattan.py │ ├── mask.py │ ├── mask_entropy.py │ ├── mono.py │ ├── perceptual.py │ ├── photometric.py │ ├── ray_vw_entropy.py │ ├── sdf_curvature.py │ ├── sparsity.py │ └── weight_reg.py ├── models │ ├── __init__.py │ ├── asset_base.py │ ├── env │ │ ├── __init__.py │ │ └── sky.py │ ├── large │ │ ├── __init__.py │ │ └── neus.py │ ├── misc │ │ ├── __init__.py │ │ └── color_transform.py │ ├── scene │ │ ├── __init__.py │ │ ├── image_embeddings.py │ │ └── learnable_params.py │ ├── shared │ │ ├── __init__.py │ │ ├── batched_dynamic_neus.py │ │ ├── batched_neus.py │ │ └── utils.py │ └── single │ │ ├── __init__.py │ │ ├── dynamic_nerf.py │ │ ├── dynamic_neus.py │ │ ├── nerf.py │ │ └── neus.py ├── renderers │ ├── __init__.py │ ├── buffer_compose_renderer.py │ ├── render_parallel.py │ ├── repr_compose_renderer.py │ ├── single_volume_renderer.py │ └── utils.py ├── resources │ ├── __init__.py │ ├── asset_bank.py │ ├── nodes.py │ ├── observers │ │ ├── __init__.py │ │ ├── cameras.py │ │ ├── fisheye.py │ │ ├── lidars.py │ │ ├── orth_camera.py │ │ └── panaroma.py │ ├── scene_bank.py │ ├── scenes.py │ └── utils.py ├── space_builder │ ├── __init__.py │ └── grid_space_builder.py ├── visible_grid.py └── visualizer │ ├── gui_runner_forest.py │ ├── gui_runner_single_cuboid.py │ ├── utils.py │ ├── visualize_scene.py │ └── visualize_visible_grid.py ├── code_multi ├── README.md ├── configs │ ├── exps │ │ ├── bg_neus=multi_block │ │ │ ├── fg=hyperlotd_no_occ.221224.yaml │ │ │ └── fg=permuto_occ.240212.yaml │ │ ├── fg_emernerf │ │ │ ├── no_gtbox_no_flow.240208.yaml │ │ │ └── no_gtbox_with_flow.240208.yaml │ │ ├── fg_neus=hyper_lotd │ │ │ └── no_fg_occ.221218.yaml │ │ ├── fg_neus=permuto │ │ │ ├── all_occ.240201.yaml │ │ │ └── all_occ.with_normals.240201.yaml │ │ └── fg_neus=permuto_emernerf │ │ │ └── with_gtbox.240212.yaml │ └── waymo │ │ └── dataset_dynamic_81.yaml └── tools │ ├── demo_category.py │ ├── demo_lidar_sim.sh │ ├── demo_lidar_sim_anim.sh │ ├── demo_manipulate.sh │ ├── eval.py │ ├── eval_directory.py │ ├── extract_mesh.py │ ├── extract_visible_grid.py │ ├── manipulate.py │ ├── manipulate_directory.py │ ├── nvs.py │ ├── render.py │ ├── render_anim.py │ ├── render_topdown.py │ ├── render_waymo_top_lidar.sh │ ├── run.py │ ├── train.py │ ├── utils.py │ └── vis_anno.py ├── code_single ├── README.md ├── configs │ ├── indoor │ │ └── lotd_neus.replica.230814.yaml │ ├── object_centric │ │ ├── lotd_neus.bmvs.230814.yaml │ │ ├── lotd_neus.dtu.230814.yaml │ │ └── permuto_neus.bmvs.230814.yaml │ └── waymo │ │ ├── dataset_static_new8.yaml │ │ ├── ngp_withlidar.230814.yaml │ │ ├── streetsurf.230814 │ │ ├── lidaronly_filterobj.230814.yaml │ │ ├── nomask_withlidar.230814.yaml │ │ ├── withmask_nolidar.230814.yaml │ │ ├── withmask_withlidar.230814.yaml │ │ └── withmask_withlidar_withnormal.230814.yaml │ │ └── streetsurf │ │ ├── lidaronly_filterobj.240219.yaml │ │ ├── nomask_withlidar_joint.240219.yaml │ │ ├── withmask_nolidar.240219.yaml │ │ ├── withmask_withlidar_joint.240219.yaml │ │ └── withmask_withlidar_withnormal.240219.yaml └── tools │ ├── assets │ ├── dashAgGridFunctions.js │ └── dashAgGridStyle.css │ ├── demo_lidar_sim.sh │ ├── eval.py │ ├── eval_directory.py │ ├── eval_lidar.py │ ├── eval_lidar_directory.py │ ├── extract_mesh.py │ ├── extract_mesh_directory.sh │ ├── extract_occgrid.py │ ├── extract_occgrid_directory.sh │ ├── inspect_rendering.py │ ├── render.py │ ├── render_directory.sh │ ├── render_lidar_directory.sh │ ├── run.py │ ├── train.py │ └── visualize_slice.py ├── dataio ├── __init__.py ├── autonomous_driving │ ├── __init__.py │ ├── custom │ │ ├── README.md │ │ ├── custom_autodrive_dataset.py │ │ └── filter_dynamic.py │ ├── kitti │ │ └── kitti_dataset.py │ ├── nuscenes │ │ └── nuscenes_dataset.py │ ├── pandaset │ │ ├── extract_masks.py │ │ ├── pandaset_dataset.py │ │ └── preprocess.py │ ├── waymo │ │ ├── README.md │ │ ├── __init__.py │ │ ├── build_asset_dataset_v1_panoptic.py │ │ ├── build_asset_dataset_v2_asset.py │ │ ├── copy_selected_seqs.py │ │ ├── download_waymo.sh │ │ ├── download_waymo_v2.sh │ │ ├── env_backup.yml │ │ ├── experimental │ │ │ ├── make_vid_asset_v1.py │ │ │ ├── make_vid_asset_v2.py │ │ │ ├── make_vid_camera_box_v1.py │ │ │ ├── make_vid_camera_box_v2.py │ │ │ ├── make_vid_camera_box_v2_try1.py │ │ │ ├── make_vid_camera_box_v2_try2.py │ │ │ ├── make_vid_camera_seg_v1.py │ │ │ └── unit_test.py │ │ ├── extract_masks.py │ │ ├── extract_masks_vit_adapter.py │ │ ├── extract_mono_cues.py │ │ ├── filter_demo.py │ │ ├── filter_dynamic.py │ │ ├── generate_multiseq_config.py │ │ ├── preprocess.py │ │ ├── train_and_eval_multiple.py │ │ ├── train_multi_and_eval_multiple.py │ │ ├── waymo_dataset.py │ │ ├── waymo_dynamic_81.lst │ │ ├── waymo_dynamic_81_dbg.lst │ │ ├── waymo_filereader.py │ │ ├── waymo_static_32.lst │ │ ├── waymo_static_new8.lst │ │ ├── waymo_unanno.lst │ │ └── zip_selected_seqs.py │ └── zod │ │ └── zod_dataset.py ├── block_nerf │ ├── README.md │ ├── __init__.py │ └── block_nerf_dataset.py ├── bmvs │ ├── README.md │ ├── __init__.py │ ├── bmvs_dataset.py │ ├── normalize_bmvs.py │ └── normalized.lst ├── colmap │ ├── README.md │ ├── __init__.py │ ├── colmap_dataset.py │ ├── colmap_loader.py │ └── extract_mono_cues.py ├── custom_old │ ├── __init__.py │ └── custom_dataset.py ├── data_loader │ ├── __init__.py │ ├── base_loader.py │ ├── image_loader.py │ ├── lidar_loader.py │ ├── patch_sampler.py │ ├── pixel_loader.py │ ├── sampler.py │ ├── unit_test.py │ └── view_loader.py ├── dtu │ ├── README.md │ ├── __init__.py │ └── dtu_dataset.py ├── gtav_nerf │ └── gtav_nerf_dataset.py ├── mega_nerf │ ├── __init__.py │ └── mega_nerf_dataset.py ├── monosdf │ ├── README.md │ ├── __init__.py │ └── monosdf_dataset.py ├── nerf │ └── nerf_dataset.py ├── ners │ ├── __init__.py │ ├── instance_dataset.py │ └── instance_dataset_cropped.py ├── neural_recon_w │ ├── neural_recon_w_dataset.py │ └── preprocess.py ├── scannet │ └── download.py ├── scene_dataset.py └── utils.py ├── docs ├── data │ ├── autonomous_driving.md │ └── autonomous_driving_cn.md ├── exps │ ├── exp_lipshitz_3d.py │ ├── exp_lipshitz_3d_modulated.py │ ├── exp_permuto_2d_modulated.py │ ├── exp_permuto_3d_modulated.py │ ├── exp_sky_oneframe.py │ ├── mll_enc_videos.py │ └── permuto_enc_video.py ├── methods │ ├── neuralsim.md │ ├── neus_in_10_minutes.md │ ├── ngp_lidar.md │ ├── street_gaussian.md │ ├── streetsurf.md │ ├── streetsurf │ │ ├── cuboid_space_sdf.png │ │ ├── raymarching_with_eikonal_on_occ.png │ │ ├── raymarching_without_eikonal_on_occ.png │ │ ├── sdf_nablas_norm_relu.png │ │ ├── sdf_slice_with_eikonal_on_occ.png │ │ └── sdf_slice_without_eikonal_on_occ.png │ └── streetsurf_unisim.md └── tutorials │ ├── appearace_factorization.md │ ├── compare_pcl_mesh.py │ └── semantic_field.md ├── media ├── github_assets.txt ├── logo_blue.png ├── logo_blue.svg ├── multi_object_volume_render.png ├── occ_grid_batched.jpg ├── occ_grid_batched_dynamic.jpg ├── occ_grid_dynamic.jpg ├── scene_graph.png ├── vis_frustum_culling.jpeg └── vis_scene_graph.jpeg └── set_env.sh /.gitignore: -------------------------------------------------------------------------------- 1 | .spyproject/ 2 | .vscode/ 3 | 4 | **/__pycache__ 5 | *.pyc 6 | 7 | *.egg 8 | build/ 9 | dist/ 10 | *.egg-info/ 11 | 12 | /debug/ 13 | /logs/ 14 | /data/ 15 | /data 16 | /dev_test/ 17 | /out/ 18 | /trained_models/ 19 | *.bk 20 | 21 | dbg_*/** 22 | 23 | pyrightconfig.json 24 | *.lprof 25 | 26 | search_*/ 27 | bk_search_*/ 28 | /dev_tools/debug_hidden 29 | /code_single/configs/backups/ 30 | 31 | /app/resources/observers/RisleyPrismCsvData/*.csv 32 | **/internal/ 33 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "nr3d_lib"] 2 | path = nr3d_lib 3 | url = https://github.com/pjlab-ADG/nr3d_lib 4 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 ADG@PJLab 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /VERSION: -------------------------------------------------------------------------------- 1 | 0.6.0 -------------------------------------------------------------------------------- /app/anim/__init__.py: -------------------------------------------------------------------------------- 1 | from app.resources import Scene 2 | 3 | from .anim import Anim 4 | from .file_anim import FileAnim 5 | from .reverse_ego_anim import ReverseEgoAnim 6 | 7 | 8 | def create_anim(anim_file_or_type: str, scene: Scene) -> Anim: 9 | if anim_file_or_type == "reverse_ego": 10 | return ReverseEgoAnim(scene) 11 | else: 12 | return FileAnim(anim_file_or_type, scene) 13 | -------------------------------------------------------------------------------- /app/anim/anim.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from app.resources import Scene 3 | 4 | 5 | class Anim: 6 | scene: Scene 7 | device: torch.device 8 | 9 | def __init__(self, scene: Scene) -> None: 10 | self.scene = scene 11 | self.device = scene.device 12 | 13 | def slice_at(self, global_frame: int): 14 | raise NotImplementedError() 15 | -------------------------------------------------------------------------------- /app/anim/file_anim.py: -------------------------------------------------------------------------------- 1 | import json 2 | import re 3 | import torch 4 | from typing import Dict 5 | 6 | from nr3d_lib.models.attributes import * 7 | 8 | from app.resources import Scene, SceneNode 9 | from app.anim.anim import Anim 10 | 11 | 12 | class FileAnim(Anim): 13 | n_frames: int 14 | time_step: float 15 | clip_range: range 16 | start_at_scene_frame: int 17 | pause_scene_anim: bool 18 | class_cfg: Dict[str, Any] 19 | _anim: Dict[str, Dict[str, torch.Tensor]] 20 | 21 | def __init__(self, file: str, scene: Scene): 22 | super().__init__(scene) 23 | with open(file, "r") as fp: 24 | json_data: dict = json.load(fp) 25 | 26 | self.time_step = json_data["time_step"] 27 | self.n_frames = len(json_data["frames"]) 28 | self.clip_range = range(json_data.get("clip_start", 0), 29 | json_data.get("clip_stop", self.n_frames)) 30 | self.start_at_scene_frame = json_data.get("start_at_scene_frame", 0) 31 | self.pause_scene_anim = json_data.get("pause_scene_anim", 0) 32 | if self.pause_scene_anim is True: 33 | self.pause_scene_anim = self.clip_range[1] - self.clip_range[0] 34 | self.class_cfg = json_data["class_cfg"] 35 | 36 | self._anim = {} 37 | for i, frame in enumerate(json_data["frames"]): 38 | for key, value in frame.items(): 39 | if isinstance(value, list): 40 | if ":" not in key: 41 | key = key + ":0" 42 | id = self._find_or_create_node(scene, *key.split(":")) 43 | self._add(i, id, value) 44 | else: 45 | for short_id, trs in value.items(): 46 | id = self._find_or_create_node(scene, key, short_id) 47 | self._add(i, id, trs) 48 | 49 | def slice_at(self, global_frame: int): 50 | anim_frame = global_frame - self.start_at_scene_frame 51 | scene_frame = global_frame - max(min(self.pause_scene_anim, anim_frame), 0) 52 | scene_frame = min(scene_frame, len(self.scene) - 1) 53 | self.scene.slice_at(scene_frame) 54 | 55 | for class_name, class_anim_cfg in self.class_cfg.items(): 56 | if class_anim_cfg["merge_mode"] == "replace": 57 | for node in self.scene.all_nodes_by_class_name[class_name]: 58 | node.i_valid = False 59 | 60 | if anim_frame >= 0 and anim_frame < self.clip_range.stop - self.clip_range.start: 61 | anim_frame_1 = anim_frame + self.clip_range.start 62 | for class_name, class_anim_cfg in self.class_cfg.items(): 63 | for node in self.scene.all_nodes_by_class_name[class_name]: 64 | if node.id in self._anim: 65 | node.i_valid = True 66 | node.i_valid_flags = self._anim[node.id]["valid_flag"][anim_frame_1] 67 | node.transform = TransformMat4x4(self._anim[node.id]["transforms"][anim_frame_1]) 68 | if self._anim[node.id]["scales"][anim_frame_1, 0] > 0: 69 | node.scale = Scale(self._anim[node.id]["scales"][anim_frame_1]) 70 | 71 | self.scene.root.update() 72 | 73 | def _add(self, frame_ind: int, id: str, trs: List[float]): 74 | if id not in self._anim: 75 | self._anim[id] = { 76 | "valid_flag": torch.zeros(self.n_frames, dtype=torch.bool, device=self.device), 77 | "transforms": torch.zeros(self.n_frames, 4, 4, device=self.device), 78 | "scales": torch.zeros(self.n_frames, 3, device=self.device) 79 | } 80 | self._anim[id]["valid_flag"][frame_ind] = True 81 | self._anim[id]["transforms"][frame_ind] = torch.tensor( 82 | trs[:12] + [0., 0., 0., 1.], device=self.device).reshape(4, 4) 83 | if len(trs) == 15: 84 | self._anim[id]["scales"][frame_ind] = torch.tensor(trs[12:], device=self.device) 85 | 86 | def _find_or_create_node(self, scene: Scene, class_name: str, short_id: str): 87 | if re.match("^\d+$", short_id): 88 | return scene.all_nodes_by_class_name[class_name][int(short_id)].id 89 | for node in scene.all_nodes_by_class_name.get(class_name, []): 90 | if node.id.endswith(short_id): 91 | return node.id 92 | for scene_id, obj_id in scene.asset_bank.drawable_shared_map.get(class_name, []): 93 | if obj_id.endswith(short_id): 94 | new_node = SceneNode(obj_id, class_name, scene.device, scene.dtype) 95 | new_node_model_id = scene.asset_bank.asset_compute_id(obj=new_node, scene=scene, class_name=class_name) 96 | new_node.model = scene.asset_bank[new_node_model_id] 97 | scene.add_node(new_node) 98 | return obj_id 99 | raise ValueError(f"{short_id} not exists in asset bank") 100 | -------------------------------------------------------------------------------- /app/anim/reverse_ego_anim.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from typing import Dict 3 | 4 | from nr3d_lib.models.attributes import * 5 | 6 | from app.resources import Scene 7 | from app.anim.anim import Anim 8 | 9 | class ReverseEgoAnim(Anim): 10 | n_frames: int 11 | time_step: float 12 | clip_range: range 13 | start_at_scene_frame: int 14 | pause_scene_anim: bool 15 | class_cfg: Dict[str, Any] 16 | _anim: Dict[str, Dict[str, torch.Tensor]] 17 | 18 | def __init__(self, scene: Scene): 19 | super().__init__(scene) 20 | 21 | def slice_at(self, global_frame: int): 22 | scene_frame = min(global_frame, len(self.scene) - 1) 23 | self.scene.slice_at(scene_frame) 24 | 25 | if "EgoVehicle" in self.scene.all_nodes_by_class_name: 26 | ego_node = self.scene.all_nodes_by_class_name["EgoVehicle"][0] 27 | ego_node._slice_at(len(self.scene) - 1 - scene_frame) 28 | ego_node.transform.tensor[..., :2] *= -1. # Flip x and y 29 | ego_node.update() 30 | else: 31 | raise NotImplementedError("Currently only support scene containing ego-vehicle node") -------------------------------------------------------------------------------- /app/loss/__init__.py: -------------------------------------------------------------------------------- 1 | from .clearance import ClearanceLoss 2 | from .conditional import DeformationLoss, LatentLoss 3 | from .eikonal import EikonalLoss 4 | from .lidar import LidarLoss 5 | from .mask import MaskOccupancyLoss 6 | from .photometric import PhotometricLoss 7 | from .sparsity import SparsityLoss 8 | from .mono import MonoDepthLoss, MonoNormalLoss 9 | from .weight_reg import WeightRegLoss 10 | from .color_lipshitz import ColorLipshitzRegLoss 11 | from .sdf_curvature import SDFCurvatureRegLoss 12 | from .ray_vw_entropy import RayVisWeightEntropyRegLoss 13 | from .mask_entropy import MaskEntropyRegLoss 14 | from .mahattan import RoadNormalLoss 15 | from .perceptual import PerceptualLoss, S3IMLoss 16 | from .flow import FlowLoss 17 | -------------------------------------------------------------------------------- /app/loss/clearance.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file clearance.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief SDF near field clearance regularization loss, to prevent too small or even minus near field SDFs (minus near SDF = camera inside shape) 5 | """ 6 | 7 | from copy import deepcopy 8 | from numbers import Number 9 | from typing import Dict, List, Literal, Union 10 | 11 | import torch 12 | import torch.nn as nn 13 | 14 | from nr3d_lib.config import ConfigDict 15 | from nr3d_lib.models.annealers import get_anneal_val 16 | 17 | from app.resources import Scene, SceneNode 18 | 19 | class ClearanceLoss(nn.Module): 20 | def __init__( 21 | self, 22 | class_name_cfgs: Union[ConfigDict, float], 23 | drawable_class_names: List[str]) -> None: 24 | """ 25 | Near-field SDF clearance regularization. 26 | This regularization penalizes small or negative near-SDF values to avoid the common local minimum problem of camera-inside-geometry. 27 | It also aids in eliminating near-field in-the-air artifacts in street views. 28 | 29 | Args: 30 | class_name_cfgs (Union[ConfigDict, float]): Regularization configuration for each corresponding model class_name. 31 | Each configuration has the following format (for example): 32 | "Street": 33 | { 34 | 'w': 1.0, # Loss weight, 35 | 'anenal': ..., # Optional weight annealing configuration 36 | 'thresh': 0.01, # Only penalizes near-SDF values below this threshold. 37 | 'beta': 10.0, # Growth factor for the penalty as the near-SDF value falls further below the threshold. 38 | } 39 | drawable_class_names (List[str]): List of all potential class_names. Defaults to []. 40 | """ 41 | 42 | super().__init__() 43 | if isinstance(class_name_cfgs, Number): 44 | class_name_cfgs = {class_name: {'w': class_name_cfgs} for class_name in drawable_class_names} 45 | else: 46 | for k, v in class_name_cfgs.items(): 47 | if isinstance(v, Number): 48 | class_name_cfgs[k] = {'w' : v} 49 | self.class_name_cfgs: Dict[str, ConfigDict] = class_name_cfgs 50 | 51 | def fn_penalty_sdf(self, near_sdf: torch.Tensor, beta: float = 1.0, thresh: float = 0.01): 52 | mask_in = near_sdf < thresh 53 | num_pen_pts = mask_in.sum().item() 54 | # penalty = torch.sigmoid(-beta * near_sdf[mask_in]).mean() if num_pen_pts > 0 else near_sdf.new_zeros([1,]) 55 | penalty = (torch.exp(-beta * (near_sdf[mask_in]-thresh)).sum() / mask_in.numel()) if num_pen_pts > 0 else near_sdf.new_zeros([1,]) 56 | return num_pen_pts, penalty 57 | 58 | def fn_penalty_density(self, near_density: torch.Tensor, ): 59 | raise NotImplementedError 60 | 61 | def forward( 62 | self, 63 | scene: Scene, ret: dict, uniform_samples: dict, sample: dict, ground_truth: dict, it: int, 64 | mode: Literal['pixel', 'lidar', 'image_patch'] = ...) -> Dict[str, torch.Tensor]: 65 | 66 | ret_losses = {} 67 | for _, obj_raw_ret in ret['raw_per_obj_model'].items(): 68 | if obj_raw_ret['volume_buffer']['type'] == 'empty': 69 | continue # Skip not rendered models 70 | class_name = obj_raw_ret['class_name'] 71 | model_id = obj_raw_ret['model_id'] 72 | model = scene.asset_bank[model_id] 73 | if class_name not in self.class_name_cfgs.keys(): 74 | continue 75 | 76 | config = deepcopy(self.class_name_cfgs[class_name]) 77 | w = config.pop('w', None) 78 | if (anneal_cfg:=config.pop('anneal', None)) is not None: 79 | w = get_anneal_val(it=it, **anneal_cfg) 80 | assert w is not None, f"Can not get w for {self.__class__.__name__}.{class_name}" 81 | 82 | if obj_raw_ret['volume_buffer']['type'] == 'empty': 83 | continue 84 | 85 | if 'near_sdf' in obj_raw_ret['details']: 86 | near_sdf = obj_raw_ret['details']['near_sdf'] 87 | _, penalty = self.fn_penalty_sdf(near_sdf, **config) 88 | elif 'near_sigma' in obj_raw_ret['details']: 89 | near_sigma = obj_raw_ret['details']['near_sdf'] 90 | _, penalty = self.fn_penalty_density(near_sigma, **config) 91 | else: 92 | raise RuntimeError(f"Can not find 'near_sdf' or 'near_density' in details for {class_name}") 93 | ret_losses[f'loss_clearance.{class_name}'] = w * penalty 94 | 95 | return ret_losses 96 | -------------------------------------------------------------------------------- /app/loss/color_lipshitz.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file color_lipshitz.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Color regularization loss 5 | Borrowed from 6 | PermutoSDF: Fast Multi-View Reconstruction with Implicit Surfaces using Permutohedral Lattices, 7 | Radu Alexandru Rosu and Sven Behnke 8 | """ 9 | 10 | from copy import deepcopy 11 | from numbers import Number 12 | from typing import Dict, List, Union 13 | 14 | import torch 15 | import torch.nn as nn 16 | 17 | from nr3d_lib.config import ConfigDict 18 | from nr3d_lib.models.loss.recon import * 19 | from nr3d_lib.models.annealers import get_anneal_val 20 | 21 | from app.resources import Scene, SceneNode 22 | 23 | class ColorLipshitzRegLoss(nn.Module): 24 | def __init__( 25 | self, 26 | class_name_cfgs: Union[ConfigDict, float], 27 | drawable_class_names: List[str], 28 | enable_after: int = 0, 29 | ) -> None: 30 | super().__init__() 31 | 32 | self.enable_after = enable_after 33 | if isinstance(class_name_cfgs, Number): 34 | class_name_cfgs = {class_name: {'w': class_name_cfgs} for class_name in drawable_class_names} 35 | else: 36 | for k, v in class_name_cfgs.items(): 37 | if isinstance(v, Number): 38 | class_name_cfgs[k] = {'w' : v} 39 | self.class_name_cfgs: Dict[str, ConfigDict] = class_name_cfgs 40 | 41 | def forward(self, scene: Scene, ret: dict, sample: dict, ground_truth: dict, it: int) -> Dict[str, torch.Tensor]: 42 | if it < self.enable_after: 43 | return {} 44 | ret_losses = {} 45 | for _, obj_raw_ret in ret['raw_per_obj_model'].items(): 46 | if obj_raw_ret['volume_buffer']['type'] == 'empty': 47 | continue # Skip not rendered models 48 | class_name = obj_raw_ret['class_name'] 49 | model_id = obj_raw_ret['model_id'] 50 | model = scene.asset_bank[model_id] 51 | if class_name not in self.class_name_cfgs.keys(): 52 | continue 53 | 54 | config = deepcopy(self.class_name_cfgs[class_name]) 55 | w = config.pop('w', None) 56 | if (anneal_cfg:=config.get('anneal', None)) is not None: 57 | w = get_anneal_val(it=it, **anneal_cfg) 58 | assert w is not None, f"Can not get w for {self.__class__.__name__}.{class_name}" 59 | 60 | assert hasattr(model, 'get_color_lipshitz_bound'), f"{model.id} has no get_color_lipshitz_bound" 61 | loss = config.w * model.get_color_lipshitz_bound() 62 | ret_losses[f"loss_color_reg.{class_name}"] = loss 63 | return ret_losses 64 | 65 | -------------------------------------------------------------------------------- /app/loss/conditional.py: -------------------------------------------------------------------------------- 1 | 2 | from numbers import Number 3 | import functools 4 | from typing import Dict, Union 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | 10 | from nr3d_lib.config import ConfigDict 11 | from nr3d_lib.models.loss.recon import l2_loss, relative_l2_loss, l1_loss 12 | from nr3d_lib.models.loss.safe import safe_binary_cross_entropy 13 | from nr3d_lib.models.annealers import get_annealer 14 | 15 | from app.resources import Scene 16 | 17 | class LatentLoss(nn.Module): 18 | pass 19 | 20 | class DeformationLoss(nn.Module): 21 | pass 22 | 23 | class ConditionalLosses(nn.Module): 24 | def __init__(self) -> None: 25 | super().__init__() 26 | 27 | def forward(self) -> Dict[str, torch.Tensor]: 28 | pass -------------------------------------------------------------------------------- /app/loss/mahattan.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file mahattan.py 3 | @brief Loss with extracted mono cues (mono depth prior and mono normals prior) 4 | Modified from https://github.com/autonomousvision/monosdf 5 | """ 6 | 7 | import functools 8 | from typing import List, Literal, Union 9 | 10 | import torch 11 | import torch.nn as nn 12 | import torch.nn.functional as F 13 | 14 | from nr3d_lib.config import ConfigDict 15 | from nr3d_lib.models.loss.recon import * 16 | 17 | from app.resources import Scene 18 | from app.resources.observers import Camera 19 | 20 | class RoadNormalLoss(nn.Module): 21 | """ 22 | Mahattan assumption for street views: constraints to road normals 23 | """ 24 | def __init__( 25 | self, 26 | w_l1: float = 0.03, w_cos: float = 0.03, 27 | distant_mode: Literal['crdv', 'cr_only'] = 'crdv', 28 | mask_pred_thresh: float = 0.95, 29 | enable_after: int = 0, 30 | apply_in_pixel_train_step=False, 31 | detach_mean = True, 32 | use_l1 = True, 33 | ): 34 | super().__init__() 35 | self.w_l1 = w_l1 36 | self.w_cos = w_cos 37 | self.mask_pred_thresh = mask_pred_thresh 38 | self.distant_mode = distant_mode 39 | self.enable_after = enable_after 40 | self.detach_mean = detach_mean 41 | self.use_l1 = use_l1 42 | self.apply_in_pixel_train_step = apply_in_pixel_train_step 43 | 44 | if self.distant_mode == 'cr_only': 45 | self.requires_render_per_class = True 46 | else: 47 | self.requires_render_per_class = False 48 | 49 | def forward(self, scene: Scene, cam: Camera, ret: dict, sample: dict, ground_truth: dict, it: int): 50 | if it < self.enable_after: 51 | return {} 52 | 53 | device = scene.device 54 | 55 | if self.distant_mode == 'crdv': 56 | rendered = ret['rendered'] 57 | elif self.distant_mode == 'cr_only': 58 | rendered = ret['rendered_per_obj_in_scene']['street'] 59 | else: 60 | raise RuntimeError(f"Invalid distant_mode={self.distant_mode}") 61 | 62 | normal_pred = rendered['normals_volume'] 63 | mask_pred = (rendered['mask_volume'].data > self.mask_pred_thresh) # detached 64 | 65 | #---- Road 66 | assert 'image_road_mask' in ground_truth 67 | road_mask = ground_truth['image_road_mask'].view(mask_pred.shape) 68 | 69 | mask = mask_pred & road_mask 70 | if self.detach_mean: 71 | road_normals_pred = F.normalize(normal_pred.data[mask], dim=-1) # Must be [N,3] 72 | else: 73 | road_normals_pred = F.normalize(normal_pred[mask], dim=-1) # Must be [N,3] 74 | 75 | if road_normals_pred.numel() > 0: 76 | # Encourage road to have same normals 77 | # (penalize difference with their mean) 78 | mean = F.normalize(road_normals_pred.mean(dim=0), dim=-1) # Detached ? 79 | if self.use_l1: 80 | loss_l1 = torch.abs(road_normals_pred - mean).sum(dim=-1).mean() 81 | else: 82 | loss_l1 = torch.tensor([0.], device=device) 83 | loss_cos = (1. - torch.sum(road_normals_pred * mean, dim = -1)).mean() 84 | else: 85 | loss_l1 = torch.tensor([0.], device=device) 86 | loss_cos = torch.tensor([0.], device=device) 87 | 88 | ret_losses = { 89 | 'loss_road_normal.l1': self.w_l1 * loss_l1, 90 | 'loss_road_normal.cos': self.w_cos * loss_cos 91 | } 92 | return ret_losses 93 | 94 | class MahattanLoss(nn.Module): 95 | """ 96 | The original mahattan assumption on indoor datasets 97 | """ 98 | pass -------------------------------------------------------------------------------- /app/loss/ray_vw_entropy.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file ray_vw_entropy.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Entropy regularization 5 | """ 6 | 7 | from typing import Dict, List 8 | 9 | import torch 10 | import torch.nn as nn 11 | 12 | from nr3d_lib.config import ConfigDict 13 | from nr3d_lib.graphics.pack_ops import packed_mean 14 | from nr3d_lib.models.annealers import get_annealer 15 | 16 | from app.resources import Scene 17 | 18 | class RayVisWeightEntropyRegLoss(nn.Module): 19 | def __init__( 20 | self, 21 | w: float, anneal: dict = None, mode: str='total', 22 | drawable_class_names: List[str] = []) -> None: 23 | super().__init__() 24 | self.w = w 25 | self.w_fn = None if anneal is None else get_annealer(**anneal) 26 | self.mode = mode 27 | 28 | def fn(self, volume_buffer: dict): 29 | vw = volume_buffer['vw'] 30 | entropy = -vw*torch.log(vw+1e-8) 31 | if (buffer_type:=volume_buffer['type']) == 'packed': 32 | loss = packed_mean(entropy, volume_buffer['pack_infos_collect']).mean() 33 | elif buffer_type == 'batched': 34 | loss = entropy.mean() 35 | elif buffer_type == 'emtpy': 36 | loss = 0 37 | return loss 38 | 39 | def fn_in_total(self, volume_buffer: dict): 40 | if volume_buffer['type'] == 'empty': 41 | return 0 42 | else: 43 | vw = volume_buffer['vw_in_total'] 44 | entropy = -vw*torch.log(vw+1e-8) 45 | return packed_mean(entropy, volume_buffer['pack_infos_collect']).mean() 46 | 47 | def forward_code_single(self, scene: Scene, ret: dict, sample: dict, ground_truth: dict, it: int) -> Dict[str, torch.Tensor]: 48 | w = self.w if self.w_fn is None else self.w_fn(it=it) 49 | 50 | ret_losses = dict() 51 | raw_per_obj_model = ret['raw_per_obj_model'] 52 | if 'total' in self.mode: 53 | ret_losses['loss_entropy'] = w * self.fn(ret['volume_buffer']) 54 | if 'cr' in self.mode: 55 | main_class_name = scene.main_class_name 56 | cr_obj_id = scene.drawable_groups_by_class_name[main_class_name][0].id 57 | ret_losses[f'loss_entropy.{main_class_name}'] = w * self.fn_in_total(raw_per_obj_model[cr_obj_id]['volume_buffer']) 58 | if 'dv' in self.mode: 59 | dv_class_name = 'Distant' 60 | dv_obj_id = scene.drawable_groups_by_class_name[dv_class_name][0].id 61 | ret_losses[f'loss_entropy.{dv_class_name}'] = w * self.fn_in_total(raw_per_obj_model[dv_obj_id]['volume_buffer']) 62 | return ret_losses 63 | -------------------------------------------------------------------------------- /app/loss/sdf_curvature.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file sdf_curvature.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief SDF curvature smoothness loss. 5 | Borrowed from 6 | PermutoSDF: Fast Multi-View Reconstruction with Implicit Surfaces using Permutohedral Lattices, 7 | Radu Alexandru Rosu and Sven Behnke 8 | """ 9 | 10 | from copy import deepcopy 11 | from numbers import Number 12 | from typing import Dict, List, Union 13 | 14 | import torch 15 | import torch.nn as nn 16 | import torch.nn.functional as F 17 | 18 | from nr3d_lib.config import ConfigDict 19 | from nr3d_lib.models.loss.safe import safe_mse_loss 20 | from nr3d_lib.models.annealers import get_anneal_val 21 | 22 | from app.resources import Scene, SceneNode 23 | 24 | class SDFCurvatureRegLoss(nn.Module): 25 | def __init__( 26 | self, 27 | class_name_cfgs: Union[ConfigDict, float], drawable_class_names: List[str], 28 | on_uniform_samples=True, eps=1.0e-4, enable_after: int = 0) -> None: 29 | super().__init__() 30 | if isinstance(class_name_cfgs, Number): 31 | class_name_cfgs = {class_name: {'w': class_name_cfgs} for class_name in drawable_class_names} 32 | else: 33 | for k, v in class_name_cfgs.items(): 34 | if isinstance(v, Number): 35 | class_name_cfgs[k] = {'w' : v} 36 | self.class_name_cfgs: Dict[str, ConfigDict] = class_name_cfgs 37 | self.on_uniform_samples = on_uniform_samples 38 | self.eps = eps 39 | self.enable_after = enable_after 40 | 41 | def fn(self, curvature: torch.Tensor): 42 | return curvature.clamp_max_(0.5).abs().mean() 43 | 44 | def forward(self, scene: Scene, ret: dict, uniform_samples: dict, sample: dict, ground_truth: dict, it: int) -> Dict[str, torch.Tensor]: 45 | if it < self.enable_after: 46 | return {} 47 | 48 | ret_losses = {} 49 | for _, obj_raw_ret in ret['raw_per_obj_model'].items(): 50 | if obj_raw_ret['volume_buffer']['type'] == 'empty': 51 | continue # Skip not rendered models 52 | class_name = obj_raw_ret['class_name'] 53 | model_id = obj_raw_ret['model_id'] 54 | model = scene.asset_bank[model_id] 55 | if class_name not in self.class_name_cfgs.keys(): 56 | continue 57 | 58 | config = deepcopy(self.class_name_cfgs[class_name]) 59 | w = config.pop('w', None) 60 | if (anneal_cfg:=config.get('anneal', None)) is not None: 61 | w = get_anneal_val(it=it, **anneal_cfg) 62 | assert w is not None, f"Can not get w for {self.__class__.__name__}.{class_name}" 63 | if w <= 0: 64 | continue 65 | 66 | alpha_uniform = 1.0 67 | if self.on_uniform_samples: 68 | assert class_name in uniform_samples.keys(), f"uniform_samples should contain {class_name}" 69 | loss_on_uniform = self.fn(model.get_sdf_curvature_1d(uniform_samples['net_x'], uniform_samples['nablas'], eps=self.eps)) 70 | ret_losses[f"loss_sdf_curvature_reg.{class_name}.uniform"] = w * alpha_uniform * loss_on_uniform 71 | 72 | alpha_render = config.get('alpha_loss_on_render', 0) 73 | if (alpha_render > 0) and (obj_raw_ret['volume_buffer']['type'] != 'empty'): 74 | volume_buffer = obj_raw_ret['volume_buffer'] 75 | loss_on_render = self.fn(model.get_sdf_curvature_1d(volume_buffer['net_x'], volume_buffer['nablas'], eps=self.eps)) 76 | ret_losses[f"loss_sdf_curvature_reg.{class_name}.render"] = w * alpha_render * loss_on_render 77 | 78 | return ret_losses 79 | -------------------------------------------------------------------------------- /app/loss/sparsity.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file sparsity.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Geometry sparsity regularization loss 5 | """ 6 | 7 | from copy import deepcopy 8 | from numbers import Number 9 | from typing import Dict, List, Literal, Union 10 | 11 | import torch 12 | import torch.nn as nn 13 | 14 | from nr3d_lib.config import ConfigDict 15 | from nr3d_lib.models.annealers import get_anneal_val 16 | from nr3d_lib.maths import normalized_logistic_density 17 | 18 | from app.resources import Scene, SceneNode 19 | 20 | class SparsityLoss(nn.Module): 21 | def __init__( 22 | self, 23 | class_name_cfgs: Union[ConfigDict, float] = {}, 24 | drawable_class_names: List[str] = [], 25 | enable_after: int = 0, 26 | ) -> None: 27 | """ Sparsity regularization to encourage free spaces in unobserved regions. 28 | 29 | Args: 30 | class_name_cfgs (Union[ConfigDict, float], optional): Sparsity loss configuration for each corresponding model class_name. 31 | Each configuration has the following format (for example): 32 | "Street": 33 | { 34 | 'w': 1.0, # Loss weight, 35 | 'anenal': ..., # Optional weight annealing configuration 36 | 'key': 'sdf, # The key to query the `uniform_samples` dict to get the sampled geometry values 37 | 'type': 'normalized_logistic_density', # The type of function to map the queried raw geometry values to loss values 38 | } 39 | drawable_class_names (List[str], optional): List of all possible class_names. Defaults to []. 40 | enable_after (int, optional): Enable this loss after this iteration. Defaults to 0. 41 | """ 42 | 43 | super().__init__() 44 | self.class_name_cfgs: Dict[str, ConfigDict] = class_name_cfgs 45 | self.enable_after = enable_after 46 | 47 | def fn_normal(self, x: torch.Tensor, std: float = 1.0): 48 | return torch.exp(-x**std).mean() 49 | 50 | def fn_nld(self, x: torch.Tensor, inv_scale: float = 16.0): 51 | return normalized_logistic_density(x,inv_s=inv_scale).mean() 52 | 53 | def fn_density_reg(self, x: torch.Tensor, lamb: float = 0.05): 54 | return (1 - (-lamb * x).exp()).abs().mean() 55 | 56 | def forward(self, scene: Scene, ret: dict, uniform_samples: dict, sample: dict, ground_truth: dict, it: int) -> Dict[str, torch.Tensor]: 57 | if it < self.enable_after: 58 | return {} 59 | 60 | ret_losses = {} 61 | for _, obj_raw_ret in ret['raw_per_obj_model'].items(): 62 | if obj_raw_ret['volume_buffer']['type'] == 'empty': 63 | continue # Skip not rendered models 64 | class_name = obj_raw_ret['class_name'] 65 | model_id = obj_raw_ret['model_id'] 66 | model = scene.asset_bank[model_id] 67 | if class_name not in self.class_name_cfgs.keys(): 68 | continue 69 | 70 | config = deepcopy(self.class_name_cfgs[class_name]) 71 | assert class_name in uniform_samples.keys(), f"uniform_samples should contain {class_name}" 72 | w = config.pop('w', None) 73 | if (anneal_cfg:=config.pop('anneal', None)) is not None: 74 | w = get_anneal_val(it=it, **anneal_cfg) 75 | assert w is not None, f"Can not get w for {self.__class__.__name__}.{class_name}" 76 | 77 | val = uniform_samples[class_name][config.pop('key', 'sdf')] 78 | fn_type = config.pop('type', 'normalized_logistic_density') 79 | if fn_type == 'normal': 80 | loss = self.fn_normal(val, **config) 81 | elif fn_type == 'normalized_logistic_density': 82 | loss = self.fn_nld(val, **config) 83 | elif fn_type == 'density_reg': 84 | loss = self.fn_density_reg(val, **config) 85 | else: 86 | raise RuntimeError(f"Invalid type={fn_type}") 87 | ret_losses[f"loss_sparsity.{class_name}"] = w * loss 88 | return ret_losses 89 | -------------------------------------------------------------------------------- /app/loss/weight_reg.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file weight_reg.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Network weight regularization loss 5 | """ 6 | 7 | from copy import deepcopy 8 | from numbers import Number 9 | from typing import Dict, List, Union 10 | 11 | import torch 12 | import torch.nn as nn 13 | 14 | from nr3d_lib.config import ConfigDict 15 | from nr3d_lib.models.loss.recon import * 16 | from nr3d_lib.models.annealers import get_anneal_val 17 | 18 | from app.resources import Scene, SceneNode 19 | 20 | class WeightRegLoss(nn.Module): 21 | def __init__( 22 | self, 23 | class_name_cfgs: Union[ConfigDict, float], 24 | drawable_class_names: List[str], 25 | ) -> None: 26 | """ Parameter regularization to prevent diversion 27 | NOTE: Will invoke each configured model's get_weight_reg() function to get a flattenned tensor of all params' norm. 28 | 29 | Args: 30 | class_name_cfgs (Union[ConfigDict, float], optional): Regularization configuration for each corresponding model class_name. 31 | Defaults to {}. 32 | drawable_class_names (List[str], optional): List of all possible class_names. Defaults to []. 33 | enable_after (int, optional): Enable this loss after this iteration. Defaults to 0. 34 | """ 35 | 36 | super().__init__() 37 | 38 | if isinstance(class_name_cfgs, Number): 39 | class_name_cfgs = {class_name: {'w': class_name_cfgs} for class_name in drawable_class_names} 40 | else: 41 | for k, v in class_name_cfgs.items(): 42 | if isinstance(v, Number): 43 | class_name_cfgs[k] = {'w' : v} 44 | self.class_name_cfgs: Dict[str, ConfigDict] = class_name_cfgs 45 | 46 | def forward( 47 | self, scene: Scene, ret: dict, sample: dict, ground_truth: dict, it: int 48 | ) -> Dict[str, torch.Tensor]: 49 | ret_losses = {} 50 | 51 | for _, obj_raw_ret in ret['raw_per_obj_model'].items(): 52 | if obj_raw_ret['volume_buffer']['type'] == 'empty': 53 | continue # Skip not rendered models 54 | class_name = obj_raw_ret['class_name'] 55 | model_id = obj_raw_ret['model_id'] 56 | model = scene.asset_bank[model_id] 57 | if class_name not in self.class_name_cfgs.keys(): 58 | continue 59 | 60 | config = deepcopy(self.class_name_cfgs[class_name]) 61 | w = config.pop('w', None) 62 | if (anneal_cfg:=config.pop('anneal', None)) is not None: 63 | w = get_anneal_val(it=it, **anneal_cfg) 64 | assert w is not None, f"Can not get w for {self.__class__.__name__}.{class_name}" 65 | 66 | assert hasattr(model, 'get_weight_reg'), f"{model_id} has no get_weight_reg" 67 | loss = model.get_weight_reg(**config).sum() 68 | ret_losses[f"loss_weight_reg.{class_name}"] = w * loss 69 | return ret_losses -------------------------------------------------------------------------------- /app/models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/app/models/__init__.py -------------------------------------------------------------------------------- /app/models/env/__init__.py: -------------------------------------------------------------------------------- 1 | from .sky import * -------------------------------------------------------------------------------- /app/models/env/sky.py: -------------------------------------------------------------------------------- 1 | __all__ = [ 2 | "SimpleSky", 3 | "PureColorSky" 4 | ] 5 | 6 | from typing import List 7 | 8 | import torch 9 | import torch.nn as nn 10 | from app.resources import Scene, SceneNode 11 | 12 | from nr3d_lib.models.blocks import get_blocks 13 | from nr3d_lib.models.embedders import get_embedder 14 | 15 | from app.models.asset_base import AssetAssignment, AssetModelMixin 16 | 17 | class SimpleSky(AssetModelMixin, nn.Module): 18 | """ 19 | A simple sky model represented by directional MLPs 20 | """ 21 | assigned_to = AssetAssignment.OBJECT # NOTE: For now, sky is an unique object; might be changed to scene-level property 22 | is_ray_query_supported = False 23 | is_batched_query_supported = False 24 | def __init__( 25 | self, 26 | dir_embed_cfg:dict={'type':'spherical', 'degree': 4}, 27 | D=3, W=64, skips=[], activation='relu', output_activation='sigmoid', 28 | n_appear_embedding: int = 0, appear_embed_cfg:dict={'type':'identity'}, 29 | weight_norm=False, dtype=torch.float, device=None, use_tcnn_backend=False): 30 | super().__init__() 31 | 32 | dir_embed_cfg.setdefault('use_tcnn_backend', use_tcnn_backend) 33 | self.embed_fn_view, input_ch_views = get_embedder(dir_embed_cfg, 3) 34 | 35 | self.use_h_appear = n_appear_embedding > 0 36 | # h_appear 37 | if self.use_h_appear: 38 | self.embed_fn_appear, input_ch_h_appear = get_embedder(appear_embed_cfg, n_appear_embedding) 39 | else: 40 | input_ch_h_appear = 0 41 | 42 | self.blocks = get_blocks( 43 | input_ch_views + input_ch_h_appear, 3, 44 | D=D, W=W, skips=skips, activation=activation, output_activation=output_activation, 45 | dtype=dtype, device=device, weight_norm=weight_norm, use_tcnn_backend=use_tcnn_backend) 46 | 47 | def forward(self, v: torch.Tensor, *, h_appear: torch.Tensor = None): 48 | network_input = self.embed_fn_view(v) 49 | if self.use_h_appear > 0: 50 | network_input = torch.cat([network_input, h_appear], dim=-1) 51 | return self.blocks(network_input) 52 | 53 | @classmethod 54 | def asset_compute_id(cls, scene: Scene = None, obj: SceneNode = None, class_name: str = None) -> str: 55 | return f"{cls.__name__}#{class_name or obj.class_name}#{scene.id}#{obj.id}" 56 | 57 | class PureColorSky(AssetModelMixin, nn.Module): 58 | """ 59 | A dummy pure color sky model 60 | """ 61 | assigned_to = AssetAssignment.OBJECT 62 | is_ray_query_supported = False 63 | is_batched_query_supported = False 64 | def __init__(self, RGB: List[int]=[255,255,255]) -> None: 65 | super().__init__() 66 | self.register_buffer('RGB', torch.tensor(RGB, dtype=torch.float)/255., persistent=False) 67 | def forward(self, v: torch.Tensor, x: torch.Tensor=None): 68 | prefix = v.shape[:-1] 69 | return self.RGB.tile([*prefix,1]) 70 | 71 | @classmethod 72 | def asset_compute_id(cls, scene: Scene = None, obj: SceneNode = None, class_name: str = None) -> str: 73 | return f"{cls.__name__}#{class_name or obj.class_name}#{scene.id}#{obj.id}" -------------------------------------------------------------------------------- /app/models/large/__init__.py: -------------------------------------------------------------------------------- 1 | from .neus import * -------------------------------------------------------------------------------- /app/models/misc/__init__.py: -------------------------------------------------------------------------------- 1 | from .color_transform import * -------------------------------------------------------------------------------- /app/models/misc/color_transform.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file color_transform.py 3 | @author Nianchen Deng, Shanghai AI Lab 4 | @brief Learnable pixel transform module for image postprocessing 5 | """ 6 | 7 | __all__ = [ 8 | 'ColorTransform' 9 | ] 10 | 11 | import torch 12 | import torch.nn as nn 13 | 14 | from nr3d_lib.utils import torch_dtype 15 | from nr3d_lib.models.blocks import get_blocks 16 | 17 | from app.resources import Scene, SceneNode 18 | from app.models.asset_base import AssetAssignment, AssetModelMixin 19 | 20 | class ColorTransform(AssetModelMixin, nn.Module): 21 | assigned_to = AssetAssignment.MISC 22 | def __init__(self, embedding_dim: int, mode: str, dtype=torch.float, **block_params) -> None: 23 | super().__init__() 24 | self.dtype = torch_dtype(dtype) 25 | self.mode = mode 26 | 27 | if self.mode == "exposure" or self.mode == "exposure+brightness": 28 | return 29 | 30 | color_chns = 3 31 | coord_chns = 2 32 | n_affine_elements = 12 33 | if self.mode == "direct": 34 | decoder_input_chns = embedding_dim + color_chns + coord_chns 35 | decoder_output_chns = color_chns 36 | elif self.mode == "pixel_affine": 37 | decoder_input_chns = embedding_dim + coord_chns 38 | decoder_output_chns = n_affine_elements 39 | elif self.mode == "global_affine": 40 | decoder_input_chns = embedding_dim 41 | decoder_output_chns = n_affine_elements 42 | else: 43 | raise ValueError("Invalid value for argument \"mode\"") 44 | self.decoder = get_blocks(decoder_input_chns, decoder_output_chns, dtype=self.dtype, **block_params) 45 | 46 | def forward(self, h: torch.Tensor, xy: torch.Tensor, rgbs: torch.Tensor) -> torch.Tensor: 47 | """ Transform input colors by the affine matrices decoded from image embeddings. 48 | 49 | Args: 50 | h (torch.Tensor): [N, D] Features 51 | xy (torch.Tensor): [N, 2] pixel coordinates in range [0,1] 52 | rgbs (torch.Tensor): [N, 3] rgb colors in range [0,1] 53 | 54 | Returns: 55 | torch.Tensor: [N, 3] transformed rgb colors 56 | """ 57 | if self.mode == "exposure": 58 | return rgbs * torch.pow(2., h) 59 | elif self.mode == "exposure+brightness": 60 | return rgbs * torch.pow(2., h[:, 0:1]) + h[:, 1:2] 61 | 62 | if self.mode == "direct": 63 | decoder_input = torch.cat([h, xy, rgbs], dim=-1) 64 | elif self.mode == "pixel_affine": 65 | decoder_input = torch.cat([h, xy], dim=-1) 66 | elif self.mode == "global_affine": 67 | decoder_input = h 68 | 69 | decoder_output = self.decoder(decoder_input) 70 | 71 | if self.mode == "direct": 72 | rgbs_output = decoder_output 73 | elif self.mode == "pixel_affine" or self.mode == "global_affine": 74 | affine_trs = decoder_output.reshape(-1, 3, 4) # (N, 3, 4) 75 | rgbs_output = (affine_trs[:, :3, :3] @ rgbs[..., None] + affine_trs[:, :3, 3:])[..., 0] 76 | return rgbs_output.to(rgbs.dtype) 77 | 78 | @classmethod 79 | def asset_compute_id(cls, scene: Scene = None, obj: SceneNode = None, class_name: str = None) -> str: 80 | return f"{cls.__name__}#{scene.id}" -------------------------------------------------------------------------------- /app/models/scene/__init__.py: -------------------------------------------------------------------------------- 1 | from .image_embeddings import * 2 | from .learnable_params import * -------------------------------------------------------------------------------- /app/models/scene/image_embeddings.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file image_embeddings.py 3 | @author Nianchen Deng, Shanghai AI Lab 4 | @brief Learnable image embeddings 5 | """ 6 | 7 | __all__ = [ 8 | 'ImageEmbeddings' 9 | ] 10 | 11 | import torch 12 | import torch.nn as nn 13 | from nr3d_lib.config import ConfigDict 14 | from nr3d_lib.models.embeddings import SeqEmbedding 15 | 16 | from nr3d_lib.utils import torch_dtype 17 | from nr3d_lib.fmt import log 18 | 19 | from app.resources import Scene, SceneNode 20 | from app.models.asset_base import AssetAssignment, AssetModelMixin 21 | 22 | 23 | class ImageEmbeddings(AssetModelMixin, nn.ModuleDict): 24 | assigned_to = AssetAssignment.SCENE 25 | def __init__( 26 | self, 27 | dims: int, 28 | ego_node_id: str=None, ego_class_name: str="Camera", 29 | weight_init: str="normal", weight_init_std: float=1.0, 30 | dtype=torch.float, device=None) -> None: 31 | super().__init__() 32 | 33 | self.dims = dims 34 | self.ego_node_id = ego_node_id 35 | self.ego_class_name = ego_class_name 36 | self.weight_init = weight_init 37 | self.weight_init_std = weight_init_std 38 | self.dtype = torch_dtype(dtype) 39 | self.set_device = device 40 | 41 | @property 42 | def device(self) -> torch.device: 43 | return next(self.parameters()).device 44 | 45 | def asset_populate( 46 | self, 47 | scene: Scene = None, obj: SceneNode = None, config: dict = None, 48 | device=None, **kwargs): 49 | 50 | device = device or self.set_device 51 | self.set_device = device 52 | 53 | self.scene = scene 54 | if self.ego_node_id is not None: 55 | ego_node_list = [self.scene.all_nodes[self.ego_node_id]] 56 | elif self.ego_class_name is not None: 57 | ego_node_list = self.scene.all_nodes_by_class_name[self.ego_class_name] 58 | else: 59 | raise RuntimeError( 60 | f"Invalid combination of arguments ego_node_id={self.ego_node_id}, " 61 | f"ego_class_name={self.ego_class_name}") 62 | 63 | self.exposures = nn.ModuleDict() 64 | for ego_node in ego_node_list: 65 | # NOTE: Different nodes might have different frame lengths 66 | embedding_weight = torch.empty(len(ego_node.frame_global_ts), self.dims, dtype=self.dtype) 67 | if self.weight_init == "uniform": 68 | embedding_weight.uniform_(-self.weight_init_std, self.weight_init_std) 69 | elif self.weight_init == "normal": 70 | embedding_weight.normal_(0., self.weight_init_std) 71 | elif self.weight_init == "zero": 72 | embedding_weight.zero_() 73 | else: 74 | raise ValueError("Unknown weight initial method") 75 | self[ego_node.id] = SeqEmbedding(ego_node.frame_global_ts, v_keyframes=embedding_weight, dim=self.dims, dtype=self.dtype, device=device) 76 | log.info(f"{self.scene.id} create image embeddings for {[node.id for node in ego_node_list]}") 77 | 78 | @classmethod 79 | def asset_compute_id(cls, scene: Scene = None, obj: SceneNode = None, class_name: str = None) -> str: 80 | return f"{cls.__name__}#{scene.id}" -------------------------------------------------------------------------------- /app/models/shared/__init__.py: -------------------------------------------------------------------------------- 1 | from .batched_neus import * 2 | from .batched_dynamic_neus import * 3 | -------------------------------------------------------------------------------- /app/models/shared/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/app/models/shared/utils.py -------------------------------------------------------------------------------- /app/models/single/__init__.py: -------------------------------------------------------------------------------- 1 | from .nerf import * 2 | from .neus import * 3 | from .dynamic_neus import * 4 | from .dynamic_nerf import * -------------------------------------------------------------------------------- /app/models/single/dynamic_neus.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file neus.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief neuralsim's API for NeuS models. 5 | """ 6 | 7 | __all__ = [ 8 | 'DynamicPermutoConcatNeuSObj', 9 | ] 10 | 11 | 12 | import numpy as np 13 | from typing import List 14 | 15 | import torch 16 | import torch.nn as nn 17 | 18 | from nr3d_lib.logger import Logger 19 | from nr3d_lib.config import ConfigDict 20 | from nr3d_lib.models.embeddings import SeqEmbedding 21 | from nr3d_lib.models.accelerations import get_accel_class, accel_types_dynamic 22 | from nr3d_lib.models.fields_dynamic.neus import DynamicPermutoConcatNeuSModel 23 | 24 | from app.models.asset_base import AssetAssignment, AssetMixin 25 | from app.resources import Scene, SceneNode 26 | from app.resources.observers import Camera 27 | 28 | class DynamicPermutoConcatNeuSObj(AssetMixin, DynamicPermutoConcatNeuSModel): 29 | """ 30 | MRO: 31 | -> AssetMixin 32 | -> DynamicPermutoConcatNeuSModel 33 | -> NeusRendererMixinDynamic 34 | -> DynamicPermutoConcatNeuS 35 | -> ModelMixin 36 | -> nn.Module 37 | """ 38 | assigned_to = AssetAssignment.OBJECT 39 | is_ray_query_supported: bool = True 40 | use_ts: bool = True 41 | 42 | """ Asset functions """ 43 | @classmethod 44 | def asset_compute_id(cls, scene: Scene = None, obj: SceneNode = None, class_name: str = None) -> str: 45 | return f"{cls.__name__}#{class_name or obj.class_name}#{scene.id}#{obj.id}" 46 | 47 | def asset_populate( 48 | self, 49 | scene: Scene = None, obj: SceneNode = None, 50 | config: dict = None, device=None, **kwargs): 51 | assert isinstance(obj, list) and isinstance(obj[0], SceneNode), f"Input `obj` for populate should be a list of SceneNode" 52 | 53 | ts_keyframes = obj.frame_global_ts.data.clone() # No gradients 54 | z_time_all = SeqEmbedding(ts_keyframes, **self.latents_cfg['z_time'], dtype=torch.float, device=device) 55 | self.z_time_all = z_time_all 56 | self.z_time_single = None 57 | 58 | #---- Dynamic Accel 59 | if self.accel_cfg is not None: 60 | accel_cls = get_accel_class(self.accel_cfg['type']) 61 | # NOTE: `config` is from `self.populate_cfg` 62 | accel_n_jump_frames = int(config.get('accel_n_jump_frames', 2)) 63 | if accel_cls in accel_types_dynamic: 64 | self.accel_cfg.update(ts_keyframes=ts_keyframes[::accel_n_jump_frames].contiguous()) 65 | 66 | #---- Model network's populate 67 | super().populate(config=config, device=device, **kwargs) 68 | 69 | def asset_training_initialize(self, scene: Scene, obj: SceneNode, config: dict, logger: Logger=None, log_prefix: str=None): 70 | # self.grad_guard_when_render.logger = logger 71 | # self.grad_guard_when_uniform.logger = logger 72 | return super().training_initialize(config, logger=logger, log_prefix=log_prefix) 73 | 74 | """ New define or model functions overwrite """ 75 | def set_z_time(self, z_time_single: torch.Tensor): 76 | self.z_time_single = z_time_single 77 | assert z_time_single.dim() == 1, "Only support manually specifying one z_time" 78 | 79 | def clean_z_time(self): 80 | self.z_time_single = None 81 | 82 | def _check_or_get_z_time_per_x( 83 | self, x: torch.Tensor, ts: torch.Tensor = None, z_time: torch.Tensor = None 84 | ) -> torch.Tensor: 85 | x_prefix = [*x.shape[:-1]] 86 | z_time = self.z_time_all.get_z_per_input(x_prefix, ts_per_input=ts, z_single=self.z_time_single, z_per_input=z_time) 87 | return z_time 88 | 89 | def query_sdf( 90 | self, x: torch.Tensor, *, z_time: torch.Tensor = None, 91 | ts: torch.Tensor = None, **kwargs): 92 | z_time = self._check_or_get_z_time_per_x(x, ts=ts, z_time=z_time) 93 | return super().query_sdf(x, z_time=z_time, ts=ts, **kwargs) 94 | 95 | def forward_sdf( 96 | self, x: torch.Tensor, *, z_time: torch.Tensor = None, 97 | ts: torch.Tensor = None, **kwargs): 98 | z_time = self._check_or_get_z_time_per_x(x, ts=ts, z_time=z_time) 99 | return super().forward_sdf(x, z_time=z_time, ts=ts, **kwargs) 100 | 101 | def forward_sdf_nablas( 102 | self, x: torch.Tensor, *, z_time: torch.Tensor = None, 103 | ts: torch.Tensor = None, **kwargs): 104 | z_time = self._check_or_get_z_time_per_x(x, ts=ts, z_time=z_time) 105 | return super().forward_sdf_nablas(x, z_time=z_time, ts=ts, **kwargs) 106 | 107 | def forward( 108 | self, x: torch.Tensor, *, z_time: torch.Tensor = None, 109 | ts: torch.Tensor = None, **kwargs): 110 | z_time = self._check_or_get_z_time_per_x(x, ts=ts, z_time=z_time) 111 | return super().forward(x, z_time=z_time, ts=ts, **kwargs) 112 | 113 | if __name__ == "__main__": 114 | def unit_test(): 115 | pass 116 | -------------------------------------------------------------------------------- /app/renderers/__init__.py: -------------------------------------------------------------------------------- 1 | from .buffer_compose_renderer import BufferComposeRenderer 2 | from .repr_compose_renderer import ReprComposeRenderer 3 | from .single_volume_renderer import SingleVolumeRenderer 4 | 5 | from .render_parallel import * 6 | from .utils import * -------------------------------------------------------------------------------- /app/renderers/repr_compose_renderer.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file repr_compose_renderer.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Joint rendering of multiple objects in one scene. 5 | Implements a renderer that: 6 | first composes one holistic representation from each object representation, 7 | then applies ray tracing or rasterization to volume render. 8 | 9 | FEATURES: 10 | 11 | """ 12 | 13 | import itertools 14 | import functools 15 | import numpy as np 16 | from typing import Any, Dict, Tuple, Union, List 17 | 18 | import torch 19 | import torch.nn as nn 20 | 21 | from nr3d_lib.profile import profile 22 | 23 | from app.models.asset_base import AssetAssignment, AssetModelMixin 24 | from app.resources import AssetBank, Scene, SceneNode, namedtuple_ind_id_obj 25 | from app.renderers.utils import rotate_volume_buffer_nablas, prepare_empty_rendered 26 | from app.resources.observers import Camera, MultiCamBundle, Lidar, RaysLidar, MultiRaysLidarBundle 27 | from app.renderers.render_parallel import render_parallel, render_parallel_with_replicas, EvalParallelWrapper 28 | 29 | class ReprComposeRenderer(nn.Module): 30 | """ 31 | Joint rendering of multiple objects in one scene. 32 | Implements a renderer that: 33 | first composes one holistic representation from each object representation, 34 | then applies ray tracing or rasterization to volume render. 35 | """ 36 | 37 | @profile 38 | def compose_repr(self, drawables: List[SceneNode]): 39 | """ 40 | Compose a holistic representation from all the seperate mode reprs from each node. 41 | The underlying representation can be anything that is "composable" (i.e. has the definition of compose() method) 42 | For example: 43 | - Composing multiple object's feature grids together into one holistic feature grid. e.g. unisim / GIRAFFE 44 | - Composing multiple objects' 3D gaussians or other types of kernels info one holistic group of kernels. 45 | """ 46 | pass 47 | 48 | @profile 49 | def view_query( 50 | self, 51 | observer: Camera, 52 | #---- Keyword arguments 53 | drawable_ids: List[str] = None, 54 | scene: Scene = ..., 55 | #---- Keyword arguments (View query configs) 56 | ): 57 | """ 58 | Rasterize the composed holistic representation. 59 | """ 60 | assert isinstance(observer, Camera), "view_query() only supports observer type=Camera" 61 | assert scene.i_is_single, "view_query() requires the scene to be frozen at single frame index / timestamp" 62 | 63 | if drawable_ids is None: 64 | drawables = observer.filter_drawable_groups(scene.get_drawables()) 65 | else: 66 | drawables = scene.drawables[drawable_ids] 67 | 68 | """ 69 | Consider the behavior here that falls back to buffer compose -> various composed representations. 70 | """ 71 | 72 | 73 | @profile 74 | def ray_query( 75 | self, 76 | #---- Tensor inputs 77 | rays_o: torch.Tensor, 78 | rays_d: torch.Tensor, 79 | rays_ts: torch.Tensor = None, 80 | rays_pix: torch.Tensor = None, 81 | *, 82 | #---- Keyword arguments 83 | scene: Scene, 84 | observer: Union[Camera, MultiCamBundle, Lidar, RaysLidar, MultiRaysLidarBundle]=None, 85 | ): 86 | """ 87 | Ray trace the composed holistic representation. 88 | """ 89 | 90 | -------------------------------------------------------------------------------- /app/renderers/utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file utils.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Utilities for volume renderers 5 | """ 6 | 7 | __all__ = [ 8 | 'rotate_volume_buffer_nablas' 9 | ] 10 | 11 | from typing import Tuple 12 | 13 | import torch 14 | 15 | from nr3d_lib.graphics.pack_ops import packed_matmul 16 | 17 | def rotate_volume_buffer_nablas(rotation: torch.Tensor, nablas: torch.Tensor, volume_buffer: dict = None): 18 | if rotation.dim() == 2: # Single rotation 19 | nablas_in_world = (rotation * nablas.unsqueeze(-2)).sum(-1) 20 | return nablas_in_world 21 | else: # Multi-frame rotation (When obj is frozen at multi-frame) 22 | assert volume_buffer is not None, "Requires volume_buffer's `rays_inds_hit` and `pack_infos_hit` for rotating nablas" 23 | rotation_on_hitrays = rotation[volume_buffer['rays_inds_hit']] 24 | if (buffer_type:=volume_buffer['type']) == 'packed': 25 | nablas_in_world = packed_matmul(nablas, rotation_on_hitrays, volume_buffer['pack_infos_hit']) 26 | elif buffer_type == 'batched': 27 | nablas_in_world = (rotation_on_hitrays * nablas.unsqueeze(-2)).sum(-1) 28 | return nablas_in_world 29 | 30 | def prepare_empty_rendered( 31 | prefix: Tuple[int], dtype=torch.float, device=None, 32 | with_rgb=True, with_normal=True, with_feature_dim: int = 0): 33 | rendered = dict( 34 | mask_volume = torch.zeros(prefix, dtype=dtype, device=device), 35 | depth_volume = torch.zeros(prefix, dtype=dtype, device=device) 36 | ) 37 | if with_rgb: 38 | rendered['rgb_volume'] = torch.zeros([*prefix, 3], dtype=dtype, device=device) 39 | if with_normal: 40 | rendered['normals_volume'] = torch.zeros([*prefix, 3], dtype=dtype, device=device) 41 | if with_feature_dim: 42 | rendered['feature_volume'] = torch.zeros([*prefix, with_feature_dim], dtype=dtype, device=device) 43 | return rendered 44 | 45 | -------------------------------------------------------------------------------- /app/resources/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from .nodes import * 3 | from .scenes import * 4 | from .scene_bank import * 5 | from .asset_bank import * 6 | from .utils import * -------------------------------------------------------------------------------- /app/resources/observers/__init__.py: -------------------------------------------------------------------------------- 1 | from typing import Union 2 | from .cameras import * 3 | from .lidars import * 4 | from .orth_camera import * 5 | from .panaroma import * 6 | 7 | OBSERVER_CLASS_NAMES = CAMERA_CLASS_NAMES + LIDAR_CLASS_NAMES 8 | CAMERA_TYPE = Camera 9 | LIDAR_TYPE = Union[Lidar, RaysLidar] 10 | OBSERVER_TYPE = Union[Camera, Lidar, RaysLidar] -------------------------------------------------------------------------------- /app/resources/observers/fisheye.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file fisheye.py 3 | @author Xinyu Cai, Shanghai AI Lab 4 | @brief 5 | """ 6 | 7 | import math 8 | import torch 9 | import numpy as np 10 | 11 | def fish_eye_distort(src: torch.Tensor): 12 | # K = cam.intr.mat_3x3().cpu().numpy() 13 | h,w = src.shape[0:2] 14 | fx = 318.44998905930794 / 1280 * w 15 | fy = 317.8314899911656 / 960 * h 16 | cx = 636.2089399611955 / 1280 * w 17 | cy = 481.71423781914115 / 960 * h 18 | K=np.array([[fx, 0.0, cx], 19 | [0.0, fy, cy], 20 | [0.0, 0.0, 1.0]]) 21 | D=np.array([0.18198802503702904, -0.04198598106075817, 0.010013633995507613, -0.0025294664427881705]) 22 | IK = np.linalg.inv(K) 23 | fx = K[0][0] 24 | fy = K[1][1] 25 | cx = K[0][2] 26 | cy = K[1][2] 27 | import math 28 | dstImg = np.zeros(src.shape, np.uint8) 29 | for i in range(h): 30 | for j in range(w): 31 | _x = i*IK[0][1] + IK[0][2] + j*IK[0][0] 32 | _y = i*IK[1][1] + IK[1][2] + j*IK[1][0] 33 | _w = i*IK[2][1] + IK[2][2] + j*IK[2][0] 34 | a = _x / _w 35 | b = _y / _w 36 | r = math.sqrt(a*a + b*b) 37 | theta = math.atan(r) 38 | theta2 = theta * theta 39 | theta4 = theta2 * theta2 40 | theta6 = theta4 * theta2 41 | theta8 = theta4 * theta4 42 | theta_d = theta * (1 + D[0]*theta2 + D[1]*theta4 + D[2]*theta6 + D[3]*theta8) 43 | scale = 1.0 if r == 0 else theta_d / r 44 | u = int(fx * a * scale + cx) 45 | v = int(fy * b * scale + cy) 46 | dstImg[v][u] = src[i][j] 47 | return dstImg -------------------------------------------------------------------------------- /app/resources/observers/orth_camera.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file orth_camera.py 3 | @author Nianchen Deng, Shanghai AI Lab 4 | @brief Orthogonal camera 5 | """ 6 | 7 | __all__ = [ 8 | "OrthogonalCamera" 9 | ] 10 | 11 | from typing import NamedTuple, List, Tuple 12 | 13 | import torch 14 | import torch.nn.functional as F 15 | 16 | from nr3d_lib.models.attributes import OrthoCameraIntrinsics 17 | from nr3d_lib.graphics.cameras import pinhole_view_frustum 18 | 19 | from app.resources import SceneNode 20 | from app.resources.observers.cameras import Camera, CameraBase, namedtuple_mask_nuvd, namedtuple_mask_niuvd 21 | 22 | class OrthogonalCamera(Camera): 23 | def __init__(self, unique_id: str, scene=..., device=None, dtype=torch.float): 24 | super().__init__(unique_id=unique_id, scene=scene, device=device, dtype=dtype) 25 | # Additional attributes 26 | self.intr = OrthoCameraIntrinsics(device=device) 27 | 28 | def get_all_rays(self, return_ts=False) -> List[torch.Tensor]: 29 | """ 30 | - support single frame: ✓ 31 | - support batched frames: o should only be used when H or W is the same across differerent batches. 32 | """ 33 | H, W, device = self.intr.H, self.intr.W, self.device 34 | i, j = torch.meshgrid(torch.linspace(0, W-1, W, device=device), torch.linspace(0, H-1, H, device=device), indexing='ij') 35 | i, j = i.t().reshape(H*W)+0.5, j.t().reshape(H*W)+0.5 # Pixel centers 36 | rays_o = self.world_transform(self.intr.lift(i, j, torch.full_like(i, self.near or 0.))) 37 | rays_d = self.world_transform.mat_3x4()[..., :, 2].expand_as(rays_o) 38 | 39 | ret = [rays_o, rays_d] 40 | return ret 41 | 42 | def build_view_frustum(self): 43 | pass 44 | 45 | # def check_spheres_inside_frustum(self, sphere_center_radius: torch.Tensor, holistic_body=False): 46 | # pass 47 | 48 | @torch.no_grad() 49 | def get_view_frustum_pts(self, near=None, far=None): 50 | """ 51 | - support single frame: ✓ 52 | - support batched frames: ✓ `...` means arbitary prefix-dims (i.e. `self.i_prefix`) 53 | """ 54 | near_clip = self.near or near or 0. 55 | far_clip = self.far or far or 100. 56 | _W, _H = self.intr.wh().movedim(-1, 0) 57 | _0, _near, _far = _W.new_zeros(_W.shape), _W.new_full(_W.shape, near_clip), \ 58 | _W.new_full(_W.shape, far_clip) 59 | # [..., 8] 60 | u, v, d = torch.stack([ 61 | torch.stack([_0, _0, _near], 0), 62 | torch.stack([_W, _0, _near], 0), 63 | torch.stack([_W, _H, _near], 0), 64 | torch.stack([_0, _H, _near], 0), 65 | torch.stack([_0, _0, _far], 0), 66 | torch.stack([_W, _0, _far], 0), 67 | torch.stack([_W, _H, _far], 0), 68 | torch.stack([_0, _H, _far], 0), 69 | ], dim=-1) 70 | # [..., 8, 3] 71 | pts = self.world_transform(self.intr.lift(u, v, d)[..., :3]) 72 | return pts 73 | 74 | # @profile 75 | def filter_drawable_groups(self, drawables: List[SceneNode], draw_self=False): 76 | """ 77 | Frustum culling to filter drawables by checking whether drawables' bounding spheres has intersection with view frustums. 78 | - support single frame: ✓ 79 | - support batched frames: ✓ `...` means arbitary prefix-dims (i.e. `self.i_prefix`) 80 | """ 81 | # For now, we don't consider observer's self shading 82 | if self.model is not None and not draw_self: 83 | drawables = list(filter(lambda obj: obj.id != self.id, drawables)) 84 | 85 | if len(collected := [[obj, obj.model_bounding_sphere, obj.i_valid_flags] for obj in drawables if obj.model_bounding_sphere is not None]) > 0: 86 | drawables_with_bound, model_bounding_spheres, obj_valids = zip(*collected) 87 | else: 88 | drawables_with_bound = [] 89 | drawables_no_bound = [obj for obj in drawables if obj.model_bounding_sphere is None] 90 | 91 | # No frustum Culling 92 | 93 | return list(drawables_with_bound) + drawables_no_bound 94 | -------------------------------------------------------------------------------- /app/resources/utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file utils.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Utility functions for scene resources 5 | """ 6 | 7 | __all__ = ['load_scenes_and_assets'] 8 | 9 | import os 10 | import torch 11 | from typing import Tuple, Dict, Any 12 | 13 | from nr3d_lib.fmt import log 14 | from nr3d_lib.checkpoint import sorted_ckpts 15 | from nr3d_lib.config import ConfigDict 16 | from nr3d_lib.utils import IDListedDict 17 | 18 | from app.resources.asset_bank import AssetBank 19 | from app.resources.scenes import Scene 20 | from app.resources.scene_bank import load_scene_bank 21 | 22 | 23 | def load_scenes_and_assets(exp_dir: str, assetbank_cfg: ConfigDict, training: ConfigDict, *, 24 | load_pt: str = None, device=None, **kwargs) -> Tuple[IDListedDict[Scene], AssetBank, Dict[str, Any]]: 25 | # --------------------------------------------- 26 | # ----------- Load Checkpoint -------------- 27 | # --------------------------------------------- 28 | # Automatically load 'final_xxx.pt' or 'latest.pt' if argument load_pt is not specified 29 | ckpt_file = load_pt or sorted_ckpts(os.path.join(exp_dir, 'ckpts'))[-1] 30 | log.info("=> Use ckpt:" + str(ckpt_file)) 31 | state_dict = torch.load(ckpt_file, map_location=device) 32 | 33 | # --------------------------------------------- 34 | # ----------- Scene Bank -------------- 35 | # --------------------------------------------- 36 | scene_bank, _ = load_scene_bank(os.path.join(exp_dir, 'scenarios'), device=device) 37 | 38 | # --------------------------------------------- 39 | # ----------- Asset Bank -------------- 40 | # --------------------------------------------- 41 | asset_bank = AssetBank(assetbank_cfg) 42 | asset_bank.create_asset_bank(scene_bank, load_state_dict=state_dict['asset_bank'], device=device) 43 | asset_bank.training_before_per_step(training.num_iters) 44 | asset_bank.eval() 45 | 46 | # --------------------------------------------- 47 | # ---- Load assets to scene objects ---- 48 | # --------------------------------------------- 49 | for scene in scene_bank: 50 | scene.load_assets(asset_bank) 51 | 52 | return scene_bank, asset_bank, state_dict 53 | -------------------------------------------------------------------------------- /app/space_builder/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/app/space_builder/__init__.py -------------------------------------------------------------------------------- /app/visualizer/utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file utils.py 3 | @author Nianchen Deng, Shanghai AI Lab 4 | @brief 5 | """ 6 | import torch 7 | import open3d as o3d 8 | import open3d.visualization.gui as gui 9 | 10 | from nr3d_lib.config import ConfigDict 11 | from nr3d_lib.models.attributes import * 12 | 13 | 14 | def initialize_visualizer(title: str, width: int = 1024, height: int = 768, has_sidebar: bool = False): 15 | 16 | app = gui.Application.instance 17 | app.initialize() 18 | w = app.create_window(title, width, height) 19 | em = w.theme.font_size 20 | 21 | if has_sidebar: 22 | sidebar = gui.Vert(.5 * em, gui.Margins(em, em, em, em)) 23 | w.add_child(sidebar) 24 | 25 | # Create scene widget 26 | widget3d = gui.SceneWidget() 27 | widget3d.scene = o3d.visualization.rendering.Open3DScene(w.renderer) 28 | w.add_child(widget3d) 29 | 30 | # Register layout callback 31 | def on_layout_callback(layout_context): 32 | win_rect = w.content_rect 33 | if has_sidebar: 34 | sidebar_preferred_size = sidebar.calc_preferred_size( 35 | layout_context, gui.Widget.Constraints()) 36 | sidebar_width = max(sidebar_preferred_size.width, 15 * em) 37 | sidebar.frame = gui.Rect(win_rect.x, win_rect.y, sidebar_width, win_rect.height) 38 | widget3d.frame = gui.Rect(win_rect.x + sidebar_width, win_rect.y, 39 | win_rect.width - sidebar_width, win_rect.height) 40 | else: 41 | widget3d.frame = win_rect 42 | 43 | w.set_on_layout(on_layout_callback) 44 | 45 | ret = ConfigDict(app=app, window=w, widget3d=widget3d) 46 | if has_sidebar: 47 | ret["sidebar"] = sidebar 48 | return ret 49 | 50 | 51 | def create_voxels_geometry(min_corners: torch.Tensor, max_corners: torch.Tensor, 52 | colors: Union[torch.Tensor, List[float]], mesh: bool = False): 53 | n_voxels = min_corners.shape[0] 54 | weights = min_corners.new_tensor([ 55 | [0., 0., 0.], [0., 0., 1.], [0., 1., 0.], [0., 1., 1.], 56 | [1., 0., 0.], [1., 0., 1.], [1., 1., 0.], [1., 1., 1.] 57 | ]) # (8, 3) 58 | corners = min_corners[:, None] * (1. - weights) + max_corners[:, None] * weights 59 | corners = corners.reshape(-1, 3) 60 | offsets = torch.arange(0, n_voxels * 8, 8, device=corners.device)[:, None, None] 61 | if mesh: 62 | voxel_triangles = min_corners.new_tensor([ 63 | [0, 1, 2], [5, 4, 7], [0, 4, 1], [6, 2, 7], [0, 2, 4], [3, 1, 7], 64 | [3, 2, 1], [6, 7, 4], [5, 1, 4], [3, 7, 2], [6, 4, 2], [5, 7, 1] 65 | ], dtype=torch.long) # (12, 3) 66 | triangles = voxel_triangles + offsets # (N, 12, 3) 67 | triangles = triangles.reshape(-1, 3) 68 | geo = o3d.geometry.TriangleMesh( 69 | o3d.utility.Vector3dVector(corners.cpu().numpy()), 70 | o3d.utility.Vector3iVector(triangles.cpu().numpy())) 71 | geo.compute_triangle_normals() 72 | else: 73 | voxel_lines = min_corners.new_tensor([ 74 | (0, 1), (1, 3), (3, 2), (2, 0), 75 | (4, 5), (5, 7), (7, 6), (6, 4), 76 | (0, 4), (1, 5), (2, 6), (3, 7) 77 | ], dtype=torch.long) # (12, 2) 78 | lines = voxel_lines + offsets # (N, 12, 2) 79 | lines = lines.reshape(-1, 2) 80 | geo = o3d.geometry.LineSet(o3d.utility.Vector3dVector(corners.cpu().numpy()), 81 | o3d.utility.Vector2iVector(lines.cpu().numpy())) 82 | if isinstance(colors, torch.Tensor): 83 | colors = colors[:, None].expand(-1, 8, -1).reshape(-1, 3) 84 | geo.colors = colors.cpu().numpy() 85 | else: 86 | geo.paint_uniform_color(colors) 87 | return geo 88 | -------------------------------------------------------------------------------- /code_multi/README.md: -------------------------------------------------------------------------------- 1 | ## Major settings 2 | 3 | For multi-object datasets. FG stands for foreground, and BG stands for background. 4 | 5 | | Methods | Official / Un-official | Get started | Notes, major difference from paper, etc. | 6 | | ------------------------------------------------------------ | ---------------------- | -------------------------------------------- | ------------------------------------------------------------ | 7 | | Neuralsim | Official | [readme](../docs/methods/neuralsim.md) | | 8 | 9 | ## General guide 10 | 11 | ```shell 12 | python code_multi/tools/run.py train,render --config code_multi/configs/exps/fg_neus=permuto/all_occ.with_normals.240201.yaml --render.downscale=4 13 | ``` -------------------------------------------------------------------------------- /code_multi/tools/demo_lidar_sim.sh: -------------------------------------------------------------------------------- 1 | # export EXP_DIR=/home/guojianfei/neuralsim_results/seg767010_exp108_splitblock_netv4_lv3_far=240.0_nr8192_cstep=0.6_step=0.05 2 | export EXP_DIR=/home/guojianfei/neuralsim_results/seg767010_exp111_15_splitblock_netv4_lv3_vi=10.0_far=280.0_nr8192_cstep=0. 3 | 4 | # original rerender 5 | python code_multi/tools/replay.py --resume_dir $EXP_DIR --render_lidar --start_frame 41 --stop_frame 95 dataset_cfg.param.root=/data1/waymo/processed --no_cam --downscale 4 --forward_inv_s 64000 --lidar_model=original_reren --lidar_id=lidar_TOP --dirname demo_lidar_original_reren --lidar_vis_vmin=-2. --lidar_vis_vmax=9. assetbank_cfg.class_name_cfgs.Vehicle.model_params.framework.model_params.ray_query_cfg.forward_inv_s=6400 6 | 7 | # horizon 8 | python code_multi/tools/replay.py --resume_dir $EXP_DIR --render_lidar --start_frame 41 --stop_frame 95 dataset_cfg.param.root=/data1/waymo/processed --no_cam --downscale 4 --forward_inv_s 64000 --lidar_model=Risley_prism --lidar_id=horizon --dirname demo_lidar_horizon --lidar_vis_vmin=-2. --lidar_vis_vmax=9. assetbank_cfg.class_name_cfgs.Vehicle.model_params.framework.model_params.ray_query_cfg.forward_inv_s=6400 9 | 10 | # panda qt 11 | python code_multi/tools/replay.py --resume_dir $EXP_DIR --render_lidar --start_frame 41 --stop_frame 95 dataset_cfg.param.root=/data1/waymo/processed --no_cam --downscale 4 --forward_inv_s 64000 --lidar_model=Surround --lidar_id=pandar_qt --dirname demo_lidar_pandar_qt --lidar_vis_vmin=-2. --lidar_vis_width=960 --lidar_vis_vmax=9. assetbank_cfg.class_name_cfgs.Vehicle.model_params.framework.model_params.ray_query_cfg.forward_inv_s=6400 12 | 13 | # panda 128 14 | python code_multi/tools/replay.py --resume_dir $EXP_DIR --render_lidar --start_frame 41 --stop_frame 95 dataset_cfg.param.root=/data1/waymo/processed --no_cam --downscale 4 --forward_inv_s 64000 --lidar_model=Surround --lidar_id=pandar128 --dirname demo_lidar_pandar_128 --lidar_vis_vmin=-2. --lidar_vis_vmax=9. assetbank_cfg.class_name_cfgs.Vehicle.model_params.framework.model_params.ray_query_cfg.forward_inv_s=6400 15 | 16 | # panda xt 17 | python code_multi/tools/replay.py --resume_dir $EXP_DIR --render_lidar --start_frame 41 --stop_frame 95 dataset_cfg.param.root=/data1/waymo/processed --no_cam --downscale 4 --forward_inv_s 64000 --lidar_model=Surround --lidar_id=pandar_xt --dirname demo_lidar_pandar_xt --lidar_vis_vmin=-2. --lidar_vis_vmax=9. assetbank_cfg.class_name_cfgs.Vehicle.model_params.framework.model_params.ray_query_cfg.forward_inv_s=6400 18 | 19 | # hdl64 20 | python code_multi/tools/replay.py --resume_dir $EXP_DIR --render_lidar --start_frame 41 --stop_frame 95 dataset_cfg.param.root=/data1/waymo/processed --no_cam --downscale 4 --forward_inv_s 64000 --lidar_model=Surround --lidar_id=hdl64 --dirname demo_lidar_hdl64 --lidar_vis_vmin=-2. --lidar_vis_vmax=9. assetbank_cfg.class_name_cfgs.Vehicle.model_params.framework.model_params.ray_query_cfg.forward_inv_s=6400 21 | 22 | # vlp16 23 | python code_multi/tools/replay.py --resume_dir $EXP_DIR --render_lidar --start_frame 41 --stop_frame 95 dataset_cfg.param.root=/data1/waymo/processed --no_cam --downscale 4 --forward_inv_s 64000 --lidar_model=Surround --lidar_id=vlp16 --dirname demo_lidar_vlp16 --lidar_vis_vmin=-2. --lidar_vis_vmax=9. assetbank_cfg.class_name_cfgs.Vehicle.model_params.framework.model_params.ray_query_cfg.forward_inv_s=6400 24 | 25 | # bpearl 26 | python code_multi/tools/replay.py --resume_dir $EXP_DIR --render_lidar --start_frame 41 --stop_frame 95 dataset_cfg.param.root=/data1/waymo/processed --no_cam --downscale 4 --forward_inv_s 64000 --lidar_model=Surround --lidar_id=bpearl --dirname demo_lidar_bpearl --lidar_vis_vmin=-2. --lidar_vis_vmax=9. --lidar_vis_width=960 assetbank_cfg.class_name_cfgs.Vehicle.model_params.framework.model_params.ray_query_cfg.forward_inv_s=6400 27 | 28 | # rs_m1 29 | python code_multi/tools/replay.py --resume_dir $EXP_DIR --render_lidar --start_frame 41 --stop_frame 95 dataset_cfg.param.root=/data1/waymo/processed --no_cam --downscale 4 --forward_inv_s 64000 --lidar_model=Solid_state --lidar_id=rs_m1 --dirname demo_lidar_rs_m1 --lidar_vis_vmin=-2. --lidar_vis_vmax=9. --lidar_vis_width=960 assetbank_cfg.class_name_cfgs.Vehicle.model_params.framework.model_params.ray_query_cfg.forward_inv_s=6400 -------------------------------------------------------------------------------- /code_multi/tools/demo_lidar_sim_anim.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/bash 2 | EXP_DIR=/home/dengnianchen/Work/neuralsim/logs/seg745856_merged_exp220 3 | ANIM=/home/dengnianchen/Work/neuralsim/logs/seg745856_merged_exp220/vehicle_trajectory1.json 4 | DATASET_ROOT=/data/waymo 5 | START_FRAME=52 6 | STOP_FRAME=100 7 | FPS=18 8 | LIDAR_VIS_WIDTH=480 9 | 10 | EXTRA_ARGS=( 11 | "--lidar_vis_view" "TBDnDBg3lh6aQuGsMgqpgQ@45,70" 12 | #"--draw_anno" "TBDnDBg3lh6aQuGsMgqpgQ" "--anno_color" "0" 13 | ) 14 | 15 | if [ -n "$ANIM" ] 16 | then 17 | EXTRA_ARGS+=("--anim" "$ANIM") 18 | fi 19 | 20 | LIDAR_MODELS=(original_reren Risley_prism Surround Surround Surround Surround Surround Surround Solid_state) 21 | LIDAR_NAMES=(lidar_TOP horizon pandar_qt pandar128 pandar_xt hdl64 vlp16 bpearl rs_m1) 22 | 23 | for ((i = 0; i < ${#LIDAR_MODELS[@]}; i++)); do 24 | python code_multi/tools/render_anim.py --resume_dir "$EXP_DIR" dataset_cfg.param.root="$DATASET_ROOT" \ 25 | --start_frame $START_FRAME --stop_frame $STOP_FRAME --fps $FPS --no_gt --no_cam --render_lidar \ 26 | --forward_inv_s 64000 assetbank_cfg.class_name_cfgs.Vehicle.model_params.framework.model_params.ray_query_cfg.forward_inv_s=6400 \ 27 | --lidar_model="${LIDAR_MODELS[i]}" --lidar_id="${LIDAR_NAMES[i]}" \ 28 | --lidar_vis_vmin=-2. --lidar_vis_vmax=9. --lidar_vis_width $LIDAR_VIS_WIDTH \ 29 | ${EXTRA_ARGS[*]} 30 | done 31 | -------------------------------------------------------------------------------- /code_multi/tools/demo_manipulate.sh: -------------------------------------------------------------------------------- 1 | PY_ARGS=${@:1} # --start_frame 80 --downscale 4 2 | 3 | # source /etc/profile.d/conda.sh 4 | # conda activate /cpfs2/user/guojianfei/ml/ 5 | # cd /cpfs2/user/guojianfei/ai_ws/neuralsim/ 6 | export PYTHONPATH="${DIR}":$PYTHONPATH 7 | 8 | python code_multi/tools/manipulate.py --only_all --mode random ${PY_ARGS} --start_frame 60 --num_frames 60 --fix_gt --render_lidar --only_all 9 | python code_multi/tools/manipulate.py --only_all --mode replay_random ${PY_ARGS} --start_frame 0 --num_frames 120 --render_lidar --only_all 10 | # python code_multi/tools/manipulate.py --only_all --mode replay_rotation n_rots=5 --stop_frame -1 ${PY_ARGS} 11 | # python code_multi/tools/manipulate.py --only_all --mode replay_translation --stop_frame -1 ${PY_ARGS} 12 | # python code_multi/tools/manipulate.py --only_all --mode replay_scale --stop_frame -1 ${PY_ARGS} 13 | # python code_multi/tools/manipulate.py --only_all --mode rotation --num_frames 48 --fix_gt n_rots=1.0 ${PY_ARGS} 14 | # python code_multi/tools/manipulate.py --only_all --mode translation --fix_gt ${PY_ARGS} 15 | # python code_multi/tools/manipulate.py --only_all --mode scale --num_frames 48 --fix_gt ${PY_ARGS} 16 | # python code_multi/tools/manipulate.py --only_all --mode thanos --stop_frame -1 ${PY_ARGS} 17 | 18 | python code_multi/tools/manipulate.py --only_all --mode self_fly ${PY_ARGS} --num_frames 48 --only_all 19 | python code_multi/tools/manipulate.py --only_all --mode self_zoom_out_fix_obj ${PY_ARGS} --fix_gt --num_frames 48 --only_all 20 | python code_multi/tools/manipulate.py --only_all --mode self_rotate ${PY_ARGS} --num_frames 120 --render_lidar --only_all 21 | # python code_multi/tools/manipulate.py --only_all --mode self_trans ${PY_ARGS} 22 | 23 | # python code_multi/tools/manipulate.py --only_all --mode clone --stop_frame -1 ${PY_ARGS} 24 | -------------------------------------------------------------------------------- /code_multi/tools/eval_directory.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file eval_directory.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Appearance evaluation for all exps in a specified directory (--overall_dir) 5 | """ 6 | import os 7 | import sys 8 | def set_env(depth: int): 9 | # Add project root to sys.path 10 | current_file_path = os.path.abspath(__file__) 11 | project_root_path = os.path.dirname(current_file_path) 12 | for _ in range(depth): 13 | project_root_path = os.path.dirname(project_root_path) 14 | if project_root_path not in sys.path: 15 | sys.path.append(project_root_path) 16 | print(f"Added {project_root_path} to sys.path") 17 | set_env(2) 18 | 19 | import os 20 | from tqdm import tqdm 21 | 22 | from nr3d_lib.utils import is_file_being_written 23 | from nr3d_lib.config import OmegaConf, ConfigDict 24 | 25 | from code_multi.tools.eval import main_function as eval_main, make_parser as eval_parser 26 | 27 | def main_function(args: ConfigDict): 28 | sdir_list = list(sorted(os.listdir(args.overall_dir))) 29 | for sdirname in tqdm(sdir_list, 'evaluating...'): 30 | full_sdir = os.path.join(args.overall_dir, sdirname) 31 | if os.path.isdir(full_sdir): 32 | # NOTE: Only eval finished exps 33 | ckpt_dir = os.path.join(full_sdir, 'ckpts') 34 | finished = False 35 | if os.path.exists(ckpt_dir): 36 | ckpt_list = os.listdir(ckpt_dir) 37 | for ckpt in ckpt_list: 38 | if 'final' in ckpt and not is_file_being_written(ckpt): 39 | finished = True 40 | break 41 | if not finished: 42 | continue 43 | 44 | eval_dir = os.path.join(full_sdir, args.dirname) 45 | if not args.no_ignore_existing and os.path.exists(eval_dir): 46 | continue 47 | 48 | sargs = OmegaConf.create(args.to_dict()) 49 | sargs.exp_dir = full_sdir 50 | detail_cfg = OmegaConf.load(os.path.join(full_sdir, 'config.yaml')) 51 | 52 | sargs = OmegaConf.merge(detail_cfg, sargs) 53 | sargs = ConfigDict(OmegaConf.to_container(sargs, resolve=True)) 54 | try: 55 | eval_main(sargs) 56 | except FileExistsError: 57 | pass 58 | 59 | if __name__ == "__main__": 60 | bc = eval_parser() 61 | bc.parser.add_argument("--overall_dir", type=str, required=True, help="Specifies the overall directory.") 62 | bc.parser.add_argument("--no_ignore_existing", action='store_true', help="If set, existing evaluations will NOT be ignored.") 63 | args = bc.parse(stand_alone=False) 64 | main_function(args) -------------------------------------------------------------------------------- /code_multi/tools/extract_mesh.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file extract_mesh.py 3 | @author Jianfei Guo, Shanghai AI Lab & Nianchen Deng, Shanghai AI Lab 4 | @brief Extract meshes of objects in a trained scene. 5 | """ 6 | import os 7 | import sys 8 | def set_env(depth: int): 9 | # Add project root to sys.path 10 | current_file_path = os.path.abspath(__file__) 11 | project_root_path = os.path.dirname(current_file_path) 12 | for _ in range(depth): 13 | project_root_path = os.path.dirname(project_root_path) 14 | if project_root_path not in sys.path: 15 | sys.path.append(project_root_path) 16 | print(f"Added {project_root_path} to sys.path") 17 | set_env(2) 18 | 19 | import os 20 | import torch 21 | import numpy as np 22 | from omegaconf import OmegaConf 23 | 24 | from nr3d_lib.fmt import log 25 | from nr3d_lib.config import BaseConfig 26 | from nr3d_lib.graphics.trianglemesh import extract_mesh 27 | from nr3d_lib.models.spatial import AABBSpace, ForestBlockSpace, BatchedBlockSpace 28 | 29 | from app.resources import load_scenes_and_assets 30 | from app.visible_grid import VisibleGrid 31 | from app.models.asset_base import AssetAssignment 32 | 33 | def main(args, device=torch.device('cuda')): 34 | scene_bank, *_ = load_scenes_and_assets(**args, device=device) 35 | output_root = os.path.join(args.exp_dir, args.dirname) 36 | os.makedirs(output_root, exist_ok=True) 37 | 38 | #--------------------------------------------- 39 | #---------- Mesh Extraction ---------- 40 | #--------------------------------------------- 41 | reconstruct_cfg = OmegaConf.load(args.reconstruct_cfg) 42 | 43 | log.info(f"Start [extract_mesh] in {args.exp_dir}") 44 | with torch.no_grad(): 45 | for scene in scene_bank: 46 | scene.frozen_at_global_frame(0) 47 | for class_name, class_cfg in reconstruct_cfg.items(): 48 | if args.extract_classes and class_name not in args.extract_classes: 49 | continue 50 | for i, obj in enumerate(scene.all_nodes_by_class_name[class_name]): 51 | model = obj.model 52 | if model.assigned_to in [AssetAssignment.MULTI_OBJ, AssetAssignment.MULTI_OBJ_ONE_SCENE]: 53 | model.set_condition({"keys": [obj.full_unique_id]}) 54 | 55 | if "visible_grid" in class_cfg: 56 | visible_grid = VisibleGrid.load(class_cfg.visible_grid, model.space) 57 | visible_grid.build_accel() 58 | visible_grid.postprocess("close2") 59 | model.accel = visible_grid.accel 60 | 61 | mesh_file = os.path.join(output_root, f"{class_name}#{i}#{obj.id}_N={class_cfg['resolution']}.ply") 62 | if isinstance(model.space, (ForestBlockSpace, AABBSpace)): 63 | diameter3d = np.ones((3,)) 64 | aabb = class_cfg.get("aabb") or model.space.aabb.flatten().tolist() 65 | query_sdf_fn = lambda x: model.forward_sdf(x, input_normalized=False)['sdf'] 66 | query_color_fn = lambda x, v: model.forward(x, v, with_rgb=True, with_normal=False, input_normalized=False)['rgb'] 67 | elif isinstance(model.space, BatchedBlockSpace): 68 | diameter3d = model.space.radius3d * 2 69 | aabb = class_cfg.get("aabb") or [-1., -1., -1., 1., 1., 1.] 70 | query_sdf_fn = lambda x: model.forward_sdf(x, bidx=torch.zeros_like(x[..., 0]), input_normalized=True)['sdf'] 71 | query_color_fn = lambda x, v: model.forward(x, v, bidx=torch.zeros_like(x[..., 0]), 72 | input_normalized=True, with_rgb=True, with_normal=False)['rgb'] 73 | extract_mesh(query_sdf_fn, query_color_fn, bmin=aabb[:3], bmax=aabb[3:], 74 | N=class_cfg["resolution"], include_color=True, filepath=mesh_file, 75 | show_progress=True, chunk=args.chunk, scale=diameter3d) 76 | scene.unfrozen() 77 | 78 | 79 | if __name__ == "__main__": 80 | bc = BaseConfig() 81 | bc.parser.add_argument("--load_pt", type=str, default=None, help="Typically unnecessary as the final or latest ckpt is loaded automatically. \n"\ 82 | "Only specify the ckpt file path if indeed a non-final or non-latest ckpt needs to be loaded.") 83 | bc.parser.add_argument("--chunk", type=int, default=16*1024, 84 | help='net chunk when querying the network. change for smaller GPU memory.') 85 | bc.parser.add_argument("--dirname", type=str, default='meshes', 86 | help="Sets the output directory to /path/to/exp_dir/${dirname}`.") 87 | bc.parser.add_argument("--reconstruct_cfg", type=str, default="reconstruct.yaml", 88 | help="Path to the reconstruct configuration file. Defaults to \"reconstruct.yaml\" under `resume_dir`") 89 | bc.parser.add_argument("--extract_classes", type=str, nargs="+", 90 | help="Only extract meshes of specified classes. Defaults to extract all classes defined in reconstruct configuration.") 91 | main(bc.parse()) -------------------------------------------------------------------------------- /code_multi/tools/manipulate_directory.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file manipulate_directory.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Scene editing for all exps in a specified directory (--overall_dir) 5 | """ 6 | import os 7 | import sys 8 | def set_env(depth: int): 9 | # Add project root to sys.path 10 | current_file_path = os.path.abspath(__file__) 11 | project_root_path = os.path.dirname(current_file_path) 12 | for _ in range(depth): 13 | project_root_path = os.path.dirname(project_root_path) 14 | if project_root_path not in sys.path: 15 | sys.path.append(project_root_path) 16 | print(f"Added {project_root_path} to sys.path") 17 | set_env(2) 18 | 19 | import os 20 | from tqdm import tqdm 21 | 22 | from nr3d_lib.config import OmegaConf, ConfigDict 23 | 24 | from code_multi.tools.manipulate import main_function as manipulate_main, make_parser as manipulate_parser 25 | from nr3d_lib.utils import is_file_being_written 26 | 27 | def main_function(args: ConfigDict): 28 | sdir_list = list(sorted(os.listdir(args.overall_dir))) 29 | for sdirname in tqdm(sdir_list, 'evaluating...'): 30 | full_sdir = os.path.join(args.overall_dir, sdirname) 31 | if os.path.isdir(full_sdir): 32 | # NOTE: Only eval finished exps 33 | ckpt_dir = os.path.join(full_sdir, 'ckpts') 34 | finished = False 35 | if os.path.exists(ckpt_dir): 36 | ckpt_list = os.listdir(ckpt_dir) 37 | for ckpt in ckpt_list: 38 | if 'final' in ckpt and not is_file_being_written(ckpt): 39 | finished = True 40 | break 41 | if not finished: 42 | continue 43 | 44 | eval_dir = os.path.join(full_sdir, args.dirname) 45 | if not args.no_ignore_existing and os.path.exists(eval_dir): 46 | continue 47 | 48 | sargs = OmegaConf.create(args.to_dict()) 49 | sargs.exp_dir = full_sdir 50 | detail_cfg = OmegaConf.load(os.path.join(full_sdir, 'config.yaml')) 51 | 52 | sargs = OmegaConf.merge(detail_cfg, sargs) 53 | sargs = ConfigDict(OmegaConf.to_container(sargs, resolve=True)) 54 | try: 55 | manipulate_main(sargs) 56 | except FileExistsError: 57 | pass 58 | 59 | if __name__ == "__main__": 60 | bc = manipulate_parser() 61 | bc.parser.add_argument("--overall_dir", type=str, required=True, help="Specifies the overall directory.") 62 | bc.parser.add_argument("--no_ignore_existing", action='store_true', help="If set, existing evaluations will NOT be ignored.") 63 | args = bc.parse(stand_alone=False) 64 | main_function(args) -------------------------------------------------------------------------------- /code_multi/tools/render_topdown.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file render_topdown.py 3 | @author Nianchen Deng, Shanghai AI Lab 4 | @brief 5 | """ 6 | import os 7 | import sys 8 | def set_env(depth: int): 9 | # Add project root to sys.path 10 | current_file_path = os.path.abspath(__file__) 11 | project_root_path = os.path.dirname(current_file_path) 12 | for _ in range(depth): 13 | project_root_path = os.path.dirname(project_root_path) 14 | if project_root_path not in sys.path: 15 | sys.path.append(project_root_path) 16 | print(f"Added {project_root_path} to sys.path") 17 | set_env(2) 18 | 19 | import os 20 | import imageio 21 | import numpy as np 22 | from datetime import datetime 23 | 24 | import torch 25 | 26 | from nr3d_lib.models.attributes import * 27 | 28 | from app.resources import Scene, load_scenes_and_assets 29 | from app.resources.observers import OrthogonalCamera, OrthoCameraIntrinsics 30 | from app.renderers import BufferComposeRenderer 31 | 32 | DEBUG_LIDAR = False 33 | 34 | 35 | def main_function(args, device=torch.device('cuda')): 36 | # --------------------------------------------- 37 | # -------------- Load ----------------- 38 | # --------------------------------------------- 39 | scene_bank, asset_bank, *_ = load_scenes_and_assets(**args, device=device, class_name_list=args.class_name) 40 | asset_bank.eval() 41 | 42 | observer = OrthogonalCamera("topdown_observer", device=device) 43 | observer.intr = OrthoCameraIntrinsics( 44 | phyW=torch.tensor([60.], device=device), 45 | phyH=torch.tensor([180.], device=device), 46 | W=torch.tensor([1600], device=device), 47 | H=torch.tensor([4800], device=device)) 48 | observer.transform = TransformMat4x4(np.array([ 49 | [1., 0., 0., 20.], 50 | [0., -1., 0., -90.], 51 | [0., 0., -1., 8.], 52 | [0., 0., 0., 1.] 53 | ]), device=observer.device, dtype=observer.dtype) 54 | scene.add_node(observer, scene.root) 55 | 56 | # --------------------------------------------- 57 | # ------------ Renderer --------------- 58 | # --------------------------------------------- 59 | renderer = BufferComposeRenderer(args.renderer) 60 | renderer.populate(asset_bank) 61 | renderer.eval() 62 | asset_bank.eval() 63 | renderer.config.rayschunk = args.rayschunk 64 | renderer.config.with_normal = False 65 | for scene in scene_bank: 66 | for obs in scene.get_observers(False): 67 | obs.near = renderer.config.near 68 | obs.far = renderer.config.far 69 | 70 | with torch.no_grad(): 71 | scene: Scene = scene_bank[0] 72 | scene.frozen_at_global_frame(0) 73 | ret = renderer.render(scene, observer=observer, render_per_obj_individual=True, show_progress=args.progress) 74 | 75 | def to_img(tensor): 76 | return (tensor * 255).clamp(0., 255.).to(torch.uint8)\ 77 | .reshape(observer.intr.H, observer.intr.W, -1).cpu().numpy() 78 | 79 | rgb_volume = to_img(ret['rendered']['rgb_volume']) 80 | 81 | 82 | #------------- Background 83 | bg_obj_id = scene.get_drawable_groups_by_class_name(args.class_name)[0].id 84 | bg_rendered = ret['rendered_per_obj'][bg_obj_id] 85 | bgrgb_volume = to_img(bg_rendered['rgb_volume']) 86 | imageio.imwrite(os.path.join(args.exp_dir, "topdown.png"), bgrgb_volume) 87 | 88 | 89 | if __name__ == "__main__": 90 | from nr3d_lib.config import BaseConfig 91 | bc = BaseConfig() 92 | bc.parser.add_argument("--class_name", type=str, default='Street', help="The class_name of the object you want to operate with.") 93 | bc.parser.add_argument("--progress", action='store_true', help="If set, shows per frame progress.") 94 | bc.parser.add_argument("--load_pt", type=str, default=None, help="Typically unnecessary as the final or latest ckpt is loaded automatically. \n"\ 95 | "Only specify the ckpt file path if indeed a non-final or non-latest ckpt needs to be loaded.") 96 | bc.parser.add_argument("--quality", type=int, default=None, help="Sets the quality for imageio.mimwrite (range: 0-10; 10 is the highest; default is 5).") 97 | bc.parser.add_argument("--outbase", type=str, default=datetime.now().strftime("%Y_%m_%d_%H_%M_%S"), 98 | help="Sets the basename of the output file (without extension).") 99 | bc.parser.add_argument("--rayschunk", type=int, default=4096) 100 | 101 | main_function(bc.parse()) 102 | -------------------------------------------------------------------------------- /code_multi/tools/render_waymo_top_lidar.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | exp_dir=$1 4 | PY_ARGS=${@:2} 5 | 6 | python code_multi/tools/replay.py --resume_dir $exp_dir --render_lidar --no_cam --downscale 4 --forward_inv_s 64000 --lidar_model=original_reren --lidar_id=lidar_TOP_reren --dirname demo_lidar_original_reren --lidar_vis_vmin=-2. --lidar_vis_vmax=9. assetbank_cfg.class_name_cfgs.Vehicle.model_params.framework.model_params.ray_query_cfg.forward_inv_s=6400 ${PY_ARGS} -------------------------------------------------------------------------------- /code_single/configs/waymo/dataset_static_new8.yaml: -------------------------------------------------------------------------------- 1 | 2 | scenebank_cfg: 3 | # NOTE: scene_id[,start_frame[,n_frames]] 4 | scenarios: 5 | - segment-5328596138024684667_2180_000_2200_000_with_camera_labels 6 | - segment-8345535260120974350_1980_000_2000_000_with_camera_labels 7 | - segment-8494653877777333091_540_000_560_000_with_camera_labels 8 | - segment-10017090168044687777_6380_000_6400_000_with_camera_labels, 0, 180 9 | - segment-10096619443888687526_2820_000_2840_000_with_camera_labels 10 | - segment-13667377240304615855_500_000_520_000_with_camera_labels 11 | - segment-14766384747691229841_6315_730_6335_730_with_camera_labels 12 | - segment-17330200445788773877_2700_000_2720_000_with_camera_labels 13 | observer_cfgs: 14 | Camera: 15 | list: ${camera_list} 16 | RaysLidar: 17 | list: ${lidar_list} 18 | object_cfgs: 19 | Vehicle: 20 | dynamic_only: true 21 | no_objects: true # Set to true to skip loading foreground objects into scene graph 22 | align_orientation: true 23 | consider_distortion: true 24 | scene_graph_has_ego_car: true -------------------------------------------------------------------------------- /code_single/tools/assets/dashAgGridFunctions.js: -------------------------------------------------------------------------------- 1 | var dagfuncs = window.dashAgGridFunctions = window.dashAgGridFunctions || {}; 2 | 3 | dagfuncs.getDataPath = function (data) { 4 | return data.hierarchy; 5 | } 6 | 7 | dagfuncs.labelFormatter = function({ value }) { 8 | return `${value}%` 9 | } 10 | 11 | dagfuncs.barFormatter = function(params) { 12 | const { yValue } = params; 13 | return { 14 | fill: yValue <= 20 ? '#4fa2d9' : yValue < 60 ? '#277cb5' : '#195176', 15 | }; 16 | } -------------------------------------------------------------------------------- /code_single/tools/assets/dashAgGridStyle.css: -------------------------------------------------------------------------------- 1 | .ag-theme-alpine { 2 | --ag-font-size: 4px !important; 3 | --ag-font-family: monospace !important; 4 | --ag-grid-size: 4px !important; 5 | } -------------------------------------------------------------------------------- /code_single/tools/demo_lidar_sim.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | EXP_DIR=$1 4 | PY_ARGS=${@:2} 5 | 6 | LIDAR_MODELS=(original original_reren Risley_prism Surround Surround Surround Surround Surround Surround Solid_state) 7 | LIDAR_NAMES=(lidar_TOP lidar_TOP horizon pandar_qt pandar128 pandar_xt hdl64 vlp16 bpearl rs_m1) 8 | LIDAR_RADIUS=(2 2 1 4 2 2 2 2 4 4) 9 | 10 | for ((i = 0; i < ${#LIDAR_MODELS[@]}; i++)); do 11 | python code_single/tools/render.py --resume_dir $EXP_DIR --render_lidar --fps=24 --no_cam --dirname demo_lidar_sim --lidar_forward_inv_s 64000 --lidar_id="${LIDAR_NAMES[i]}" --lidar_model="${LIDAR_MODELS[i]}" --lidar_vis_radius="${LIDAR_RADIUS[i]}" --lidar_vis_vmin=-2. --lidar_vis_vmax=9. --lidar_vis_rgb_choice=height ${PY_ARGS} 12 | done -------------------------------------------------------------------------------- /code_single/tools/eval_directory.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file eval_directory.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Appearance evaluation for all exps in a specified directory (--overall_dir) 5 | """ 6 | import os 7 | import sys 8 | def set_env(depth: int): 9 | # Add project root to sys.path 10 | current_file_path = os.path.abspath(__file__) 11 | project_root_path = os.path.dirname(current_file_path) 12 | for _ in range(depth): 13 | project_root_path = os.path.dirname(project_root_path) 14 | if project_root_path not in sys.path: 15 | sys.path.append(project_root_path) 16 | print(f"Added {project_root_path} to sys.path") 17 | set_env(2) 18 | 19 | import os 20 | from tqdm import tqdm 21 | 22 | from nr3d_lib.config import OmegaConf, ConfigDict 23 | 24 | from code_single.tools.eval import main_function as eval_main, make_parser as eval_parser 25 | from nr3d_lib.utils import is_file_being_written 26 | 27 | def main_function(args: ConfigDict): 28 | sdir_list = list(sorted(os.listdir(args.overall_dir))) 29 | for sdirname in tqdm(sdir_list, 'evaluating...'): 30 | full_sdir = os.path.join(args.overall_dir, sdirname) 31 | if os.path.isdir(full_sdir): 32 | # NOTE: Only eval finished exps 33 | ckpt_dir = os.path.join(full_sdir, 'ckpts') 34 | finished = False 35 | if os.path.exists(ckpt_dir): 36 | ckpt_list = os.listdir(ckpt_dir) 37 | for ckpt in ckpt_list: 38 | if 'final' in ckpt and not is_file_being_written(ckpt): 39 | finished = True 40 | break 41 | if not finished: 42 | continue 43 | 44 | eval_dir = os.path.join(full_sdir, args.dirname) 45 | if not args.no_ignore_existing and os.path.exists(eval_dir): 46 | continue 47 | 48 | sargs = OmegaConf.create(args.to_dict()) 49 | sargs.exp_dir = full_sdir 50 | detail_cfg = OmegaConf.load(os.path.join(full_sdir, 'config.yaml')) 51 | 52 | sargs = OmegaConf.merge(detail_cfg, sargs) 53 | sargs = ConfigDict(OmegaConf.to_container(sargs, resolve=True)) 54 | try: 55 | eval_main(sargs) 56 | except FileExistsError: 57 | pass 58 | 59 | if __name__ == "__main__": 60 | bc = eval_parser() 61 | bc.parser.add_argument("--overall_dir", type=str, required=True, help="Specifies the overall directory.") 62 | bc.parser.add_argument("--no_ignore_existing", action='store_true', help="If set, existing evaluations will NOT be ignored.") 63 | args = bc.parse(stand_alone=False) 64 | main_function(args) -------------------------------------------------------------------------------- /code_single/tools/eval_lidar_directory.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file eval_lidar_directory.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief LiDAR evaluation for all exps in a specified directory (--overall_dir) 5 | """ 6 | import os 7 | import sys 8 | def set_env(depth: int): 9 | # Add project root to sys.path 10 | current_file_path = os.path.abspath(__file__) 11 | project_root_path = os.path.dirname(current_file_path) 12 | for _ in range(depth): 13 | project_root_path = os.path.dirname(project_root_path) 14 | if project_root_path not in sys.path: 15 | sys.path.append(project_root_path) 16 | print(f"Added {project_root_path} to sys.path") 17 | set_env(2) 18 | 19 | import os 20 | from tqdm import tqdm 21 | 22 | from nr3d_lib.config import OmegaConf, ConfigDict 23 | 24 | from code_single.tools.eval_lidar import main_function as eval_lidar_main, make_parser as eval_lidar_parser 25 | from nr3d_lib.utils import is_file_being_written 26 | 27 | def main_function(args: ConfigDict): 28 | sdir_list = list(sorted(os.listdir(args.overall_dir))) 29 | for sdirname in tqdm(sdir_list, 'evaluating...'): 30 | full_sdir = os.path.join(args.overall_dir, sdirname) 31 | if os.path.isdir(full_sdir): 32 | # NOTE: Only eval finished exps 33 | ckpt_dir = os.path.join(full_sdir, 'ckpts') 34 | finished = False 35 | if os.path.exists(ckpt_dir): 36 | ckpt_list = os.listdir(ckpt_dir) 37 | for ckpt in ckpt_list: 38 | if 'final' in ckpt and not is_file_being_written(ckpt): 39 | finished = True 40 | break 41 | if not finished: 42 | continue 43 | 44 | eval_dir = os.path.join(full_sdir, args.dirname) 45 | if not args.no_ignore_existing and os.path.exists(eval_dir): 46 | continue 47 | 48 | sargs = OmegaConf.create(args.to_dict()) 49 | sargs.exp_dir = full_sdir 50 | detail_cfg = OmegaConf.load(os.path.join(full_sdir, 'config.yaml')) 51 | 52 | sargs = OmegaConf.merge(detail_cfg, sargs) 53 | sargs = ConfigDict(OmegaConf.to_container(sargs, resolve=True)) 54 | try: 55 | eval_lidar_main(sargs) 56 | except FileExistsError: 57 | pass 58 | 59 | if __name__ == "__main__": 60 | bc = eval_lidar_parser() 61 | bc.parser.add_argument("--overall_dir", type=str, required=True, help="Specifies the overall directory.") 62 | bc.parser.add_argument("--no_ignore_existing", action='store_true', help="If set, existing evaluations will NOT be ignored.") 63 | args = bc.parse(stand_alone=False) 64 | main_function(args) -------------------------------------------------------------------------------- /code_single/tools/extract_mesh_directory.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Mesh extraction for all exps in a specified directory ($1) 4 | 5 | # $1: overall_dir 6 | # the rest: other params 7 | 8 | overall_dir=$1 9 | PY_ARGS=${@:2} 10 | 11 | shopt -s nullglob 12 | 13 | for expdir in $overall_dir/* 14 | do 15 | expname=$(basename $expdir) 16 | # echo $expdir "->" $expname 17 | 18 | # files=(${expdir}/meshes/*0.2.ply) 19 | # if [ ${#files[@]} -gt 0 ]; then 20 | # echo "exist: $expdir" 21 | # else 22 | # echo "no: $expdir" 23 | # python code_single/tools/extract_mesh.py --resume_dir $expdir ${PY_ARGS} 24 | # fi 25 | 26 | python code_single/tools/extract_mesh.py --resume_dir $expdir ${PY_ARGS} 27 | done 28 | 29 | echo "Done extract_mesh_directory.sh in dir ${overall_dir}" -------------------------------------------------------------------------------- /code_single/tools/extract_occgrid_directory.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | overall_dir=$1 4 | PY_ARGS=${@:2} 5 | 6 | for expdir in $overall_dir/* 7 | do 8 | expname=$(basename $expdir) 9 | echo $expdir "->" $expname 10 | python code_single/tools/extract_occgrid.py --resume_dir $expdir ${PY_ARGS} 11 | done 12 | 13 | echo "Done extract_occgrid_directory.sh in dir ${overall_dir}" -------------------------------------------------------------------------------- /code_single/tools/render_directory.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Render all experiments in a specified directory ($1) 4 | 5 | overall_dir=$1 6 | PY_ARGS=${@:2} 7 | 8 | for expdir in $overall_dir/* 9 | do 10 | expname=$(basename $expdir) 11 | echo $expdir "->" $expname 12 | python code_single/tools/replay.py --resume_dir $expdir ${PY_ARGS} 13 | done 14 | 15 | echo "Done replay_directory.sh in dir ${overall_dir}" -------------------------------------------------------------------------------- /code_single/tools/render_lidar_directory.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | overall_dir=$1 4 | lidar_model=$2 # original_reren 5 | lidar_id=$3 # lidar_TOP 6 | PY_ARGS=${@:4} 7 | 8 | for expdir in $overall_dir/* 9 | do 10 | expname=$(basename $expdir) 11 | echo $expdir "->" $expname 12 | python code_single/tools/replay.py --resume_dir $expdir --no_cam --render_lidar --forward_inv_s 64000 --lidar_model=${lidar_model} --lidar_id=${lidar_id} --dirname=render_${lidar_model}_${lidar_id} --lidar_vis_vmin=-2. --lidar_vis_vmax=9. ${PY_ARGS} 13 | done 14 | 15 | echo "Done render_lidar_directory.sh in dir ${overall_dir}" -------------------------------------------------------------------------------- /dataio/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/dataio/__init__.py -------------------------------------------------------------------------------- /dataio/autonomous_driving/__init__.py: -------------------------------------------------------------------------------- 1 | from .waymo.waymo_dataset import WaymoDataset 2 | # from .kitti.kitti_dataset import KITTIDataset 3 | # from .nuscenes.nuscenes_dataset import NuScenesDataset 4 | from .custom.custom_autodrive_dataset import CustomAutoDriveDataset -------------------------------------------------------------------------------- /dataio/autonomous_driving/custom/README.md: -------------------------------------------------------------------------------- 1 | ### Preprocess 2 | 3 | We have standardized the preprocessed datasets into **universal formats applicable to all autonomous driving datasets**. For more details on this universal format, please refer to [docs/data/autonomous_driving.md](../../../docs/data/autonomous_driving.md). Regardless of the type of autonomous driving dataset you are using, once converted into this preprocessed format, it can be directly loaded by [custom_autodrive_dataset.py](custom_autodrive_dataset.py). 4 | 5 | ### Prior extraction - monocular cues and masks 6 | 7 | For extraction of monocular depths and surface normals cues, please refer to [Extract monocular normals & depths priors](../waymo/README.md#extract-monocular-normals--depths-priors). 8 | 9 | For extraction of semantic masks for sky and dynamic objects, please refer to [Extract mask priors - for sky, pedestrian etc](../waymo/README.md#extract-mask-priors----for-sky-pedestrian-etc). 10 | 11 | -------------------------------------------------------------------------------- /dataio/autonomous_driving/custom/filter_dynamic.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file filter_dynamic.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Filter dynamic objects 5 | """ 6 | 7 | import argparse 8 | import numpy as np 9 | from typing import List 10 | 11 | def collect_loc_motion(scenario: dict, loc_eps=0.03): 12 | """ 13 | return path: {id: path_xyz} 14 | path_xyz: path_x: [] (3*length) 15 | """ 16 | categ_stats = {} 17 | for oid, odict in scenario['objects'].items(): 18 | class_name = odict['class_name'] 19 | # Location at world coordinate 20 | loc_diff_norms = [] 21 | for seg in odict['segments']: 22 | locations = seg['data']['transform'][:, :3, 3] 23 | loc_diff = np.diff(locations, axis=0) 24 | loc_diff_norm = np.linalg.norm(loc_diff, axis=-1) 25 | loc_diff_norms.append(loc_diff_norm) 26 | categ_stats.setdefault(class_name, {})[oid] = loc_diff_norms 27 | return categ_stats 28 | 29 | def stat_dynamic_objects(scenario: dict, loc_eps=0.03, speed_eps=None, all_class_name: List[str] = None): 30 | if all_class_name is None: 31 | all_class_name = [odict['class_name'] for oid, odict in scenario['objects'].items()] 32 | all_class_name = list(set(all_class_name)) 33 | stats = {cls_name:{'n_dynamic': 0, 'is_dynamic':[], 'by_speed': [], 'by_loc': []} for cls_name in all_class_name} 34 | #------------------------------------------------ 35 | # Filter according to center_x and center_y 36 | loc_motion_stats = collect_loc_motion(scenario) 37 | for cls_name, cls_dict in loc_motion_stats.items(): 38 | by_loc = [] 39 | for oid, loc_diff_norms in cls_dict.items(): 40 | loc_diff_norms = np.concatenate(loc_diff_norms) 41 | if len(loc_diff_norms) == 0: 42 | continue 43 | # print(oid, loc_diff_norms.max()) 44 | if loc_diff_norms.max() > loc_eps: 45 | by_loc.append(oid) 46 | stats[cls_name]['by_loc'] = by_loc 47 | #------------------------------------------------ 48 | # Filter according to speed_x and speed_y 49 | pass 50 | #------------------------------------------------ 51 | # Collect results from box_speed and loc_motion 52 | for cls_name, cls_dict in stats.items(): 53 | # li_dyna = list(set(cls_dict['by_speed']) | set(cls_dict['by_loc'])) 54 | li_dyna = cls_dict['by_loc'] 55 | cls_dict['is_dynamic'] = li_dyna 56 | cls_dict['n_dynamic'] = len(li_dyna) 57 | 58 | return stats -------------------------------------------------------------------------------- /dataio/autonomous_driving/kitti/kitti_dataset.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file kitti_dataset.py 3 | @author Jianfei Guo, Shanghai AI Lab. 4 | @brief Dataset IO for KITTI datasets 5 | 6 | To be merged. 7 | """ 8 | import numpy as np 9 | from typing import Dict, Any 10 | 11 | from nr3d_lib.config import ConfigDict 12 | 13 | from dataio.scene_dataset import SceneDataset 14 | 15 | class KITTIDataset(SceneDataset): 16 | def __init__(self, config: dict) -> None: 17 | config = config.copy() 18 | 19 | def get_scenario(self, scene_id: str, **kwargs) -> Dict[str, Any]: 20 | pass 21 | 22 | def get_image(self, scene_id: str, camera_id: str, frame_index: int) -> np.ndarray: 23 | pass 24 | 25 | def get_image_occupancy_mask(self, scene_id: str, camera_id: str, frame_index: int) -> np.ndarray: 26 | pass 27 | 28 | def get_image_semantic_mask_all(self, scene_id: str, camera_id: str, frame_index: int, *, compress=True) -> np.ndarray: 29 | # Integer semantic mask on RGB image. 30 | raise NotImplementedError 31 | 32 | def get_lidar(self, scene_id: str, lidar_id: str, frame_index: int) -> Dict[str, np.ndarray]: 33 | pass -------------------------------------------------------------------------------- /dataio/autonomous_driving/nuscenes/nuscenes_dataset.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file nuscenes_dataset.py 3 | @author Chenjing Ding, Sensetime 4 | @brief Dataset IO for NuScenes dataset 5 | 6 | To be merged. 7 | TODO: https://github.com/PJLab-ADG/neuralsim_dev/blob/tsdf_sample/dataio/nuscense_io.py 8 | """ 9 | import numpy as np 10 | from typing import Dict, Any 11 | 12 | from nr3d_lib.config import ConfigDict 13 | 14 | from dataio.scene_dataset import SceneDataset 15 | 16 | class NuScenesDataset(SceneDataset): 17 | def __init__(self, config: dict) -> None: 18 | config = config.copy() 19 | 20 | def get_scenario(self, scene_id: str, **kwargs) -> Dict[str, Any]: 21 | pass 22 | 23 | def get_image(self, scene_id: str, camera_id: str, frame_index: int) -> np.ndarray: 24 | pass 25 | 26 | def get_image_occupancy_mask(self, scene_id: str, camera_id: str, frame_index: int) -> np.ndarray: 27 | pass 28 | 29 | def get_image_semantic_mask_all(self, scene_id: str, camera_id: str, frame_index: int, *, compress=True) -> np.ndarray: 30 | # Integer semantic mask on RGB image. 31 | raise NotImplementedError 32 | 33 | def get_lidar(self, scene_id: str, lidar_id: str, frame_index: int) -> Dict[str, np.ndarray]: 34 | pass -------------------------------------------------------------------------------- /dataio/autonomous_driving/pandaset/pandaset_dataset.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | """ 4 | import os 5 | import json 6 | import pickle 7 | import numpy as np 8 | import pandas as pd 9 | from glob import glob 10 | import transforms3d as t3d # pip install transforms3d 11 | from typing import Any, Dict 12 | from operator import itemgetter 13 | 14 | from nr3d_lib.config import ConfigDict 15 | from nr3d_lib.utils import get_image_size 16 | 17 | from dataio.scene_dataset import SceneDataset 18 | 19 | def idx_to_frame_str(frame_index): 20 | return f'{frame_index:02d}' 21 | 22 | class PandarsetDataset(SceneDataset): 23 | def __init__(self, config: dict) -> None: 24 | self.config = config 25 | self.populate(**config) 26 | 27 | def populate( 28 | self, root: str, 29 | rgb_dirname: str = 'camera', 30 | lidar_dirname: str = 'lidar', 31 | mask_dirname: str = 'masks', 32 | ): 33 | self.root = root 34 | self.rgb_dirname = rgb_dirname 35 | self.lidar_dirname = lidar_dirname 36 | 37 | def get_scenario(self, scene_id: str, **kwargs) -> Dict[str, Any]: 38 | sequence_dir = os.path.join(self.root, scene_id) 39 | 40 | 41 | 42 | def group_vis_gps(): 43 | pass 44 | 45 | if __name__ == '__main__': 46 | pass -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/dataio/autonomous_driving/waymo/__init__.py -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/build_asset_dataset_v1_panoptic.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file build_asset_dataset_v1_panoptic.py 3 | @brief Build categorical asset dataset using sparse panoptic annotation provided by WOD v1.4.2 4 | @author Fengrui Tian, Shanghai AI Lab && Jianfei Guo, Shanghai AI Lab 5 | """ 6 | 7 | def build_asset_dataset( 8 | 9 | ): 10 | pass -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/build_asset_dataset_v2_asset.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file build_asset_dataset_v2_asset.py 3 | @brief Build categorical asset dataset using WOD v2 API & asset annotations 4 | @author Jianfei Guo, Shanghai AI Lab 5 | """ 6 | 7 | -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/copy_selected_seqs.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copy processed data of selected waymo sequences to a target parent directory (for data transfering between devices) 3 | """ 4 | 5 | import os 6 | from tqdm import tqdm 7 | from glob import glob 8 | from pathlib import Path 9 | 10 | import shutil 11 | from os import PathLike 12 | from typing import List, Union 13 | 14 | def copy_files( 15 | out_root_dir: str, root_dir: Union[str, PathLike], select_scene_ids: List[str], 16 | fields=['masks', 'normals', 'depths']): 17 | src_root = Path(root_dir).expanduser().resolve(strict=True) 18 | out_root = Path(out_root_dir).expanduser().resolve(strict=True) 19 | for i, scene_id in enumerate(tqdm(select_scene_ids, f'copying...')): 20 | for field in fields: 21 | if field == 'scenarios': 22 | shutil.copy(src_root.joinpath(scene_id, 'scenario.pt'), out_root.joinpath(scene_id, 'scenario.pt')) 23 | else: 24 | shutil.copytree(src_root.joinpath(scene_id, field), out_root.joinpath(scene_id, field)) 25 | 26 | def copy_raw( 27 | out_root_dir: str, root_dir: str, select_scene_ids: List[str], 28 | ): 29 | for i, scene_id in enumerate(tqdm(select_scene_ids, f'copying raw...')): 30 | shutil.copyfile(os.path.join(root_dir, f"{scene_id}.tfrecord"), os.path.join(out_root_dir, f"{scene_id}.tfrecord")) 31 | 32 | if __name__ == "__main__": 33 | import argparse 34 | parser = argparse.ArgumentParser() 35 | parser.add_argument("--seq_list", type=str, help='specify --seq_list if you want to limit the list of seqs', default="dataio/autonomous_driving/waymo/waymo_static_32.lst") 36 | parser.add_argument("--data_root", type=str, default="/data1/waymo/processed") 37 | parser.add_argument("--out_root", help='output root directory') 38 | args = parser.parse_args() 39 | 40 | assert args.out_root is not None 41 | os.makedirs(args.out_root, exist_ok=True) 42 | 43 | with open(args.seq_list, 'r') as f: 44 | seq_list = f.read().splitlines() 45 | select_scene_ids = [s.split(',')[0].rstrip(".tfrecord") for s in seq_list] 46 | 47 | copy_files(args.out_root, args.data_root, select_scene_ids) 48 | # copy_files(args.out_root, args.data_root, select_scene_ids, ['masks_vit_adapter']) 49 | # copy_raw(args.out_root, args.data_root, select_scene_ids) 50 | -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/download_waymo.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # NOTE: Before proceeding, you need to fill out the Waymo terms of use and complete `gcloud auth login`. 4 | 5 | lst=$1 # dataio/autonomous_driving/waymo/waymo_static_32.lst 6 | dest=$2 # /data1/waymo/training/ 7 | source=gs://waymo_open_dataset_v_1_4_2/individual_files/training 8 | 9 | mkdir -p $dest 10 | 11 | # Get the total number of filenames 12 | total_files=$(wc -l < $lst) 13 | counter=0 14 | 15 | # Read filenames from the .lst file and process them one by one 16 | while IFS= read -r filename; do 17 | counter=$((counter + 1)) 18 | echo "[${counter}/${total_files}] Dowloading $filename ..." 19 | gsutil cp -n ${source}/${filename}.tfrecord ${dest} 20 | done < $lst 21 | -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/download_waymo_v2.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # NOTE: Before proceeding, you need to fill out the Waymo terms of use and complete `gcloud auth login`. 4 | 5 | source=gs://waymo_open_dataset_v_2_0_0/training 6 | dest=/data1/waymo/v2 7 | 8 | tags=(\ 9 | stats \ 10 | camera_image \ 11 | camera_box \ 12 | camera_hkp \ 13 | camera_to_lidar_box_association \ 14 | lidar \ 15 | lidar_box \ 16 | veh_asset_camera_sensor \ 17 | ped_asset_camera_sensor \ 18 | veh_asset_lidar_sensor \ 19 | ped_asset_lidar_sensor \ 20 | veh_asset_auto_label \ 21 | ped_asset_auto_label \ 22 | veh_asset_ray \ 23 | ped_asset_ray \ 24 | veh_asset_ray_compressed \ 25 | ped_asset_ray_compressed \ 26 | ) 27 | 28 | # tags=(\ 29 | # camera_box \ 30 | # lidar_box \ 31 | # ) 32 | 33 | seqs=(\ 34 | 10023947602400723454_1120_000_1140_000 \ 35 | 7670103006580549715_360_000_380_000 \ 36 | ) 37 | 38 | for ((i = 0; i < ${#tags[@]}; i++)); do 39 | mkdir -p ${dest}/${tags[i]} 40 | for ((j = 0; j < ${#seqs[@]}; j++)); do 41 | gsutil cp -n ${source}/${tags[i]}/${seqs[j]}.parquet ${dest}/${tags[i]}/ 42 | done 43 | done 44 | 45 | # codec file for ray decompression 46 | mkdir -p ${dest}/veh_asset_ray_compressed/ 47 | mkdir -p ${dest}/ped_asset_ray_compressed/ 48 | gsutil cp -n ${source}/veh_asset_ray_compressed/codec_config.json ${dest}/veh_asset_ray_compressed/ 49 | gsutil cp -n ${source}/ped_asset_ray_compressed/codec_config.json ${dest}/ped_asset_ray_compressed/ -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/experimental/make_vid_asset_v1.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/dataio/autonomous_driving/waymo/experimental/make_vid_asset_v1.py -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/experimental/make_vid_camera_box_v1.py: -------------------------------------------------------------------------------- 1 | """ 2 | - Annotations of camera box and id for each frame 3 | - Existence of association 4 | """ 5 | 6 | 7 | import io 8 | import os 9 | import cv2 10 | import pickle 11 | import imageio # NOTE: Also needs pip install imageio-ffmpeg 12 | import functools 13 | import numpy as np 14 | from tqdm import tqdm 15 | from PIL import Image 16 | from typing import List 17 | from numbers import Number 18 | import matplotlib.pyplot as plt 19 | 20 | import tensorflow as tf 21 | 22 | gpus = tf.config.experimental.list_physical_devices('GPU') 23 | for gpu in gpus: 24 | tf.config.experimental.set_memory_growth(gpu, True) 25 | 26 | from waymo_open_dataset import dataset_pb2 27 | from waymo_open_dataset.utils import frame_utils, transform_utils, range_image_utils 28 | from dataio.autonomous_driving.waymo.waymo_filereader import WaymoDataFileReader 29 | from dataio.autonomous_driving.waymo.filter_dynamic import stat_dynamic_objects 30 | from dataio.autonomous_driving.waymo.waymo_dataset import * 31 | 32 | from nr3d_lib.utils import pad_images_to_same_size, image_downscale 33 | from nr3d_lib.plot import draw_2dbox_on_im 34 | 35 | def make_video(sequence_file, vid_base: str, save_perframe=True, downscale: int = 1): 36 | if save_perframe: 37 | os.makedirs(vid_base, exist_ok=True) 38 | 39 | dataset = WaymoDataFileReader(str(sequence_file)) 40 | 41 | im_per_frame = [] 42 | 43 | # frame = next(iter(dataset)) 44 | for find, frame in enumerate(tqdm(dataset, "processing...")): 45 | im_per_cam = [] 46 | 47 | for camera_name in ['SIDE_LEFT', 'FRONT_LEFT', 'FRONT', 'FRONT_RIGHT', 'SIDE_RIGHT']: 48 | _j = WAYMO_CAMERAS.index(camera_name) 49 | camera_image = frame.images[_j] 50 | for c in frame.context.camera_calibrations: 51 | if c.name == camera_image.name: 52 | break 53 | for cl in frame.camera_labels: 54 | if cl.name == camera_image.name: 55 | break 56 | 57 | im = Image.open(io.BytesIO(camera_image.image)) 58 | im = np.asarray(im) # uint8, 0~255, [H, W, 3] 59 | im = image_downscale(im, downscale=downscale) 60 | im = (im * 255).clip(0, 255).astype(np.uint8) 61 | 62 | for l in cl.labels: 63 | box = l.box 64 | 65 | cam_box_label = f"cid: {l.id[:6]}" 66 | lidar_box_label = None 67 | if l.HasField('association'): 68 | lidar_box_label = l.association.laser_object_id 69 | lidar_box_label = f"lid: {lidar_box_label[:6]}" 70 | 71 | im = draw_2dbox_on_im( 72 | im, 73 | box.center_x / downscale, box.center_y / downscale, 74 | width=box.length / downscale, height=box.width / downscale, 75 | label=cam_box_label, label2=lidar_box_label, fontscale=1./downscale, 76 | color=(0, 128, 64), fillalpha=0.2, 77 | ) 78 | 79 | im_per_cam.append(im) 80 | 81 | im_per_cam = pad_images_to_same_size(im_per_cam, padding='top_left') 82 | im_per_cam = np.concatenate(im_per_cam, axis=1) 83 | 84 | if save_perframe: 85 | imageio.imwrite(os.path.join(vid_base, f"{find:04d}.jpg"), im_per_cam) 86 | im_per_frame.append(im_per_cam) 87 | 88 | vid_path = vid_base + ".mp4" 89 | imageio.mimwrite(vid_path, im_per_frame, fps=24) 90 | 91 | if __name__ == "__main__": 92 | sequence_file = "/data2/neuralsim_open/dataset/waymo_dynamic_81/segment-7670103006580549715_360_000_380_000_with_camera_labels.tfrecord" 93 | make_video(sequence_file, "./dev_test/767010_cam_box", save_perframe=True, downscale=2) -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/experimental/make_vid_camera_seg_v1.py: -------------------------------------------------------------------------------- 1 | """ 2 | - Camera segmentation 3 | - Existence of association 4 | """ 5 | 6 | 7 | import io 8 | import os 9 | import cv2 10 | import pickle 11 | import imageio # NOTE: Also needs pip install imageio-ffmpeg 12 | import functools 13 | import numpy as np 14 | from tqdm import tqdm 15 | from PIL import Image 16 | from typing import List 17 | from numbers import Number 18 | import matplotlib.pyplot as plt 19 | 20 | import tensorflow as tf 21 | 22 | gpus = tf.config.experimental.list_physical_devices('GPU') 23 | for gpu in gpus: 24 | tf.config.experimental.set_memory_growth(gpu, True) 25 | 26 | from waymo_open_dataset import dataset_pb2 27 | from waymo_open_dataset.utils import frame_utils, transform_utils, range_image_utils 28 | from dataio.autonomous_driving.waymo.waymo_filereader import WaymoDataFileReader 29 | from dataio.autonomous_driving.waymo.filter_dynamic import stat_dynamic_objects 30 | from dataio.autonomous_driving.waymo.waymo_dataset import * 31 | 32 | from nr3d_lib.utils import pad_images_to_same_size, image_downscale 33 | from nr3d_lib.plot import draw_2dbox_on_im 34 | 35 | def make_video(sequence_file, vid_base: str, save_perframe=True, downscale: int = 1): 36 | if save_perframe: 37 | os.makedirs(vid_base, exist_ok=True) 38 | 39 | dataset = WaymoDataFileReader(str(sequence_file)) 40 | 41 | im_per_frame = [] 42 | 43 | # frame = next(iter(dataset)) 44 | for find, frame in enumerate(tqdm(dataset, "processing...")): 45 | im_per_cam = [] 46 | 47 | for camera_name in ['SIDE_LEFT', 'FRONT_LEFT', 'FRONT', 'FRONT_RIGHT', 'SIDE_RIGHT']: 48 | _j = WAYMO_CAMERAS.index(camera_name) 49 | camera_image = frame.images[_j] 50 | for c in frame.context.camera_calibrations: 51 | if c.name == camera_image.name: 52 | break 53 | for cl in frame.camera_labels: 54 | if cl.name == camera_image.name: 55 | break 56 | 57 | im = Image.open(io.BytesIO(camera_image.image)) 58 | im = np.asarray(im) # uint8, 0~255, [H, W, 3] 59 | im = image_downscale(im, downscale=downscale) 60 | im = (im * 255).clip(0, 255).astype(np.uint8) 61 | 62 | for l in cl.labels: 63 | box = l.box 64 | 65 | cam_box_label = f"cid: {l.id[:6]}" 66 | lidar_box_label = None 67 | if l.HasField('association'): 68 | lidar_box_label = l.association.laser_object_id 69 | lidar_box_label = f"lid: {lidar_box_label[:6]}" 70 | 71 | im = draw_2dbox_on_im( 72 | im, 73 | box.center_x / downscale, box.center_y / downscale, 74 | width=box.length / downscale, height=box.width / downscale, 75 | label=cam_box_label, label2=lidar_box_label, fontscale=1./downscale, 76 | color=(0, 128, 64), fillalpha=0.2, 77 | ) 78 | 79 | im_per_cam.append(im) 80 | 81 | im_per_cam = pad_images_to_same_size(im_per_cam, padding='top_left') 82 | im_per_cam = np.concatenate(im_per_cam, axis=1) 83 | 84 | if save_perframe: 85 | imageio.imwrite(os.path.join(vid_base, f"{find:04d}.jpg"), im_per_cam) 86 | im_per_frame.append(im_per_cam) 87 | 88 | vid_path = vid_base + ".mp4" 89 | imageio.mimwrite(vid_path, im_per_frame, fps=24) 90 | 91 | if __name__ == "__main__": 92 | sequence_file = "/data2/neuralsim_open/dataset/waymo_dynamic_81/segment-7670103006580549715_360_000_380_000_with_camera_labels.tfrecord" 93 | make_video(sequence_file, "./dev_test/767010_cam_box", save_perframe=True, downscale=2) -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/filter_demo.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file filter_demo.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Try to filter some nice waymo sequence for making demos. 5 | """ 6 | 7 | import os 8 | import json 9 | import argparse 10 | import functools 11 | import numpy as np 12 | from glob import glob 13 | from tqdm import tqdm 14 | 15 | def main_function(args): 16 | from dataio.autonomous_driving.waymo.waymo_filereader import WaymoDataFileReader 17 | from dataio.autonomous_driving.waymo.waymo_dataset import parse_seq_file_list, WAYMO_CLASSES, file_to_scene_id 18 | import concurrent.futures as futures 19 | 20 | def process_single_sequence(seq_fpath: str, out_root: str): 21 | scene_id = file_to_scene_id(seq_fpath) 22 | os.makedirs(os.path.join(out_root, scene_id), exist_ok=True) 23 | 24 | dataset = WaymoDataFileReader(str(seq_fpath)) 25 | frame0 = next(iter(dataset)) 26 | stats = frame0.context.stats 27 | 28 | meta_fpath = os.path.join(out_root, scene_id, "metadata.json") 29 | dic = { 30 | 'stats': { 31 | 'name': frame0.context.name, 32 | 'time_of_day': stats.time_of_day, 33 | 'weather': stats.weather, 34 | 'location': stats.location 35 | } 36 | } 37 | if os.path.exists(meta_fpath): 38 | with open(meta_fpath, 'r') as f: 39 | dic2 = json.load(f) 40 | dic.update(dic2) 41 | with open(meta_fpath, 'w') as f: 42 | json.dump(dic, f) 43 | print(f"=> File saved to {meta_fpath}") 44 | 45 | os.makedirs(args.out_root, exist_ok=True) 46 | seq_fpath_list = sorted(parse_seq_file_list(args.root, args.seq_list)) 47 | num_workers = min(args.j, len(seq_fpath_list)) 48 | 49 | if args.generate_meta: 50 | print("=> Generating metadata...") 51 | if num_workers > 1: 52 | with futures.ThreadPoolExecutor(num_workers) as executor: 53 | executor.map(functools.partial(process_single_sequence, out_root=args.out_root), seq_fpath_list) 54 | else: 55 | for seq_fpath in tqdm(seq_fpath_list): 56 | process_single_sequence(seq_fpath, out_root=args.out_root) 57 | 58 | seq_list = [] 59 | for meta_fpath in tqdm(sorted(glob(os.path.join(args.out_root, "*", 'metadata.json'))), desc="=> Filtering..."): 60 | with open(meta_fpath, 'r') as f: 61 | meta = json.load(f) 62 | 63 | dyna_stats = meta['dynamic_stats'] 64 | # n_dyna = stats['Vehicle']['n_dynamic'] 65 | n_dyna = len( list(set(dyna_stats['Vehicle']['by_speed']) & set(dyna_stats['Vehicle']['by_loc'])) ) 66 | n_total = dyna_stats['Vehicle']['n_total'] 67 | 68 | stats = meta['stats'] 69 | 70 | egomotion = meta['egomotion'] 71 | linear = np.array(egomotion['linear']).max() 72 | angular = np.array(egomotion['angular']).max() 73 | 74 | if linear < 8. and angular < 15.: 75 | seq_list.append([n_dyna, n_total, linear, angular, stats, file_to_scene_id(meta_fpath)]) 76 | 77 | output_list = os.path.join(args.out_root, 'vehicle_debug.seq_list') 78 | seq_list = sorted(seq_list, key=(lambda it: it[0]), reverse=True) 79 | with open(output_list, 'w') as f: 80 | f.writelines('\n'.join([f"{n_dyna:04d}, {n_total:04d}, {linear}, {angular}, {stats['time_of_day']:10s}, {stats['weather']:16s}, {stats['location']:16s}, {scene_id}" for (n_dyna, n_total, linear, angular, stats, scene_id) in seq_list])) 81 | print(f"=> File saved to {output_list}") 82 | 83 | if __name__ == "__main__": 84 | parser = argparse.ArgumentParser() 85 | parser.add_argument( 86 | "--root", type=str, default="/media/guojianfei/DataBank0/dataset/waymo/training", required=True, 87 | help="Root directory of raw .tfrecords") 88 | parser.add_argument( 89 | "--seq_list", type=str, default=None, 90 | help="Optional specify subset of sequences. If None, will process all sequences contained in args.root") 91 | parser.add_argument( 92 | "--out_root", type=str, default="/data1/waymo/processed", required=True, 93 | help="Output root directory") 94 | parser.add_argument('-j', type=int, default=4, help='max num workers') 95 | parser.add_argument('--generate_meta', action='store_true', help='whether to generate_meta meta info. ') 96 | args = parser.parse_args() 97 | 98 | main_function(args) -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/generate_multiseq_config.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | def set_env(depth: int): 4 | # Add project root to sys.path 5 | current_file_path = os.path.abspath(__file__) 6 | project_root_path = os.path.dirname(current_file_path) 7 | for _ in range(depth): 8 | project_root_path = os.path.dirname(project_root_path) 9 | if project_root_path not in sys.path: 10 | sys.path.append(project_root_path) 11 | print(f"Added {project_root_path} to sys.path") 12 | set_env(3) 13 | 14 | import os 15 | from glob import glob 16 | from nr3d_lib.config import ConfigDict, save_config 17 | 18 | import argparse 19 | parser = argparse.ArgumentParser() 20 | parser.add_argument("--data_root", type=str, default="/data1/waymo/processed/") 21 | parser.add_argument("--start", type=int, default=0) 22 | parser.add_argument("--stop", type=int, default=None) 23 | parser.add_argument("--out", type=str, required=True, default="dataio/autonomous_driving/waymo/all.yaml") 24 | args = parser.parse_args() 25 | 26 | scenario_root = os.path.join(os.path.normpath(args.data_root), 'scenarios') 27 | scenario_file_list = list(sorted(glob(os.path.join(scenario_root, "*.pt")))) 28 | scenario_list = [os.path.splitext(os.path.basename(s))[0] for s in scenario_file_list] 29 | 30 | if args.stop is None: 31 | args.stop = len(scenario_list) 32 | scenario_list = scenario_list[args.start:args.stop] 33 | 34 | config = ConfigDict(seq_list=scenario_list) 35 | 36 | save_config(config, args.out, ignore_fields=[]) 37 | print(f"=> Seq config saved to {args.out}") 38 | -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/train_and_eval_multiple.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | from tqdm import tqdm 4 | 5 | from app.resources import parse_scene_bank_cfg 6 | from code_single.tools.train import main_function as train_main 7 | from code_single.tools.eval import main_function as eval_main, make_parser as eval_parser 8 | 9 | if __name__ == "__main__": 10 | parser = eval_parser() 11 | args = parser.parse(stand_alone=False) 12 | scenario_cfg_list = args.scenebank_cfg.pop('scenarios') 13 | exp_parent_dir = args.exp_parent_dir 14 | for sce_cfg in tqdm(scenario_cfg_list): 15 | try: 16 | sce_id, _, _ = parse_scene_bank_cfg(sce_cfg) 17 | waymo_short_id = sce_id[8:14] 18 | 19 | # Make local config 20 | sce_args = args.deepcopy() 21 | sce_args.scenebank_cfg.scenarios = [sce_cfg] 22 | if 'test_scenebank_cfg' in sce_args: 23 | sce_args.test_scenebank_cfg.scenarios = [sce_cfg] 24 | 25 | exp_dir = sce_args.exp_dir = os.path.join(exp_parent_dir, "seg"+waymo_short_id) 26 | ckpt_dir = os.path.join(exp_dir, 'ckpts') 27 | eval_dir = os.path.join(exp_dir, args.dirname) 28 | if not os.path.exists(exp_dir): 29 | train_main(sce_args) 30 | if not os.path.exists(eval_dir): 31 | finished = False 32 | if os.path.exists(ckpt_dir): 33 | ckpt_list = os.listdir(ckpt_dir) 34 | for ckpt in ckpt_list: 35 | if 'final' in ckpt: 36 | finished = True 37 | break 38 | # NOTE: Only eval finished exps 39 | if finished: 40 | eval_main(sce_args) 41 | except FileExistsError: 42 | pass 43 | -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/train_multi_and_eval_multiple.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | from tqdm import tqdm 4 | 5 | from app.resources import parse_scene_bank_cfg 6 | from code_multi.tools.train import main_function as train_main 7 | from code_multi.tools.eval import main_function as eval_main, make_parser as eval_parser 8 | 9 | if __name__ == "__main__": 10 | parser = eval_parser() 11 | args = parser.parse(stand_alone=False) 12 | scenario_cfg_list = args.scenebank_cfg.pop('scenarios') 13 | exp_parent_dir = args.exp_parent_dir 14 | for sce_cfg in tqdm(scenario_cfg_list): 15 | try: 16 | sce_id, _, _ = parse_scene_bank_cfg(sce_cfg) 17 | waymo_short_id = sce_id[8:14] 18 | 19 | # Make local config 20 | sce_args = args.deepcopy() 21 | sce_args.scenebank_cfg.scenarios = [sce_cfg] 22 | if 'test_scenebank_cfg' in sce_args: 23 | sce_args.test_scenebank_cfg.scenarios = [sce_cfg] 24 | 25 | exp_dir = sce_args.exp_dir = os.path.join(exp_parent_dir, "seg"+waymo_short_id) 26 | ckpt_dir = os.path.join(exp_dir, 'ckpts') 27 | eval_dir = os.path.join(exp_dir, args.dirname) 28 | if not os.path.exists(exp_dir): 29 | train_main(sce_args) 30 | if not os.path.exists(eval_dir): 31 | finished = False 32 | if os.path.exists(ckpt_dir): 33 | ckpt_list = os.listdir(ckpt_dir) 34 | for ckpt in ckpt_list: 35 | if 'final' in ckpt: 36 | finished = True 37 | break 38 | # NOTE: Only eval finished exps 39 | if finished: 40 | eval_main(sce_args) 41 | except FileExistsError: 42 | pass 43 | -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/waymo_dynamic_81.lst: -------------------------------------------------------------------------------- 1 | segment-7670103006580549715_360_000_380_000_with_camera_labels 2 | segment-1191788760630624072_3880_000_3900_000_with_camera_labels 3 | segment-1730266523558914470_305_260_325_260_with_camera_labels 4 | segment-1758724094753801109_1251_037_1271_037_with_camera_labels 5 | segment-1773696223367475365_1060_000_1080_000_with_camera_labels 6 | segment-1887497421568128425_94_000_114_000_with_camera_labels 7 | segment-1891390218766838725_4980_000_5000_000_with_camera_labels 8 | segment-1918764220984209654_5680_000_5700_000_with_camera_labels 9 | segment-1926967104529174124_5214_780_5234_780_with_camera_labels 10 | segment-2922309829144504838_1840_000_1860_000_with_camera_labels 11 | segment-2935377810101940676_300_000_320_000_with_camera_labels 12 | segment-2961247865039433386_920_000_940_000_with_camera_labels 13 | segment-3195159706851203049_2763_790_2783_790_with_camera_labels 14 | segment-3461228720457810721_4511_120_4531_120_with_camera_labels 15 | segment-3490810581309970603_11125_000_11145_000_with_camera_labels 16 | segment-3591015878717398163_1381_280_1401_280_with_camera_labels 17 | segment-3644145307034257093_3000_400_3020_400_with_camera_labels 18 | segment-3657581213864582252_340_000_360_000_with_camera_labels 19 | segment-3919438171935923501_280_000_300_000_with_camera_labels 20 | segment-4164064449185492261_400_000_420_000_with_camera_labels 21 | segment-4414235478445376689_2020_000_2040_000_with_camera_labels 22 | segment-4468278022208380281_455_820_475_820_with_camera_labels 23 | segment-4537254579383578009_3820_000_3840_000_with_camera_labels 24 | segment-4604173119409817302_2820_000_2840_000_with_camera_labels 25 | segment-4808842546020773462_2310_000_2330_000_with_camera_labels 26 | segment-4960194482476803293_4575_960_4595_960_with_camera_labels 27 | segment-5451442719480728410_5660_000_5680_000_with_camera_labels 28 | segment-5495302100265783181_80_000_100_000_with_camera_labels 29 | segment-6234738900256277070_320_000_340_000_with_camera_labels 30 | segment-6242822583398487496_73_000_93_000_with_camera_labels 31 | segment-6390847454531723238_6000_000_6020_000_with_camera_labels 32 | segment-6417523992887712896_1180_000_1200_000_with_camera_labels 33 | segment-6792191642931213648_1522_000_1542_000_with_camera_labels 34 | segment-6814918034011049245_134_170_154_170_with_camera_labels 35 | segment-7000927478052605119_1052_330_1072_330_with_camera_labels 36 | segment-7313718849795510302_280_000_300_000_with_camera_labels 37 | segment-7458568461947999548_700_000_720_000_with_camera_labels 38 | segment-7554208726220851641_380_000_400_000_with_camera_labels 39 | segment-7643597152739318064_3979_000_3999_000_with_camera_labels 40 | segment-7799671367768576481_260_000_280_000_with_camera_labels 41 | segment-7885161619764516373_289_280_309_280_with_camera_labels 42 | segment-7912728502266478772_1202_200_1222_200_with_camera_labels 43 | segment-7940496892864900543_4783_540_4803_540_with_camera_labels 44 | segment-7996500550445322129_2333_304_2353_304_with_camera_labels 45 | segment-8700094808505895018_7272_488_7292_488_with_camera_labels 46 | segment-8938046348067069210_3800_000_3820_000_with_camera_labels 47 | segment-9058545212382992974_5236_200_5256_200_with_camera_labels 48 | segment-9179922063516210200_157_000_177_000_with_camera_labels 49 | segment-9653249092275997647_980_000_1000_000_with_camera_labels 50 | segment-10072231702153043603_5725_000_5745_000_with_camera_labels 51 | segment-10082223140073588526_6140_000_6160_000_with_camera_labels 52 | segment-10391312872392849784_4099_400_4119_400_with_camera_labels 53 | segment-10517728057304349900_3360_000_3380_000_with_camera_labels 54 | segment-10526338824408452410_5714_660_5734_660_with_camera_labels 55 | segment-10588771936253546636_2300_000_2320_000_with_camera_labels 56 | segment-10750135302241325253_180_000_200_000_with_camera_labels 57 | segment-11017034898130016754_697_830_717_830_with_camera_labels 58 | segment-11454085070345530663_1905_000_1925_000_with_camera_labels 59 | segment-11566385337103696871_5740_000_5760_000_with_camera_labels 60 | segment-12208410199966712301_4480_000_4500_000_with_camera_labels 61 | segment-12212767626682531382_2100_150_2120_150_with_camera_labels 62 | segment-12337317986514501583_5346_260_5366_260_with_camera_labels 63 | segment-13679757109245957439_4167_170_4187_170_with_camera_labels 64 | segment-14369250836076988112_7249_040_7269_040_with_camera_labels 65 | segment-14430914081327266277_6480_000_6500_000_with_camera_labels 66 | segment-14734824171146590110_880_000_900_000_with_camera_labels 67 | segment-14763701469114129880_2260_000_2280_000_with_camera_labels 68 | segment-14777753086917826209_4147_000_4167_000_with_camera_labels 69 | segment-15053781258223091665_3192_117_3212_117_with_camera_labels 70 | segment-15166409572599113654_808_000_828_000_with_camera_labels 71 | segment-16262849101474060261_3459_585_3479_585_with_camera_labels 72 | segment-16801666784196221098_2480_000_2500_000_with_camera_labels 73 | segment-16911037681440249335_700_000_720_000_with_camera_labels 74 | segment-17216329305659006368_4800_000_4820_000_with_camera_labels 75 | segment-17407069523496279950_4354_900_4374_900_with_camera_labels 76 | segment-17437352085580560526_2120_000_2140_000_with_camera_labels 77 | segment-17601040886987343289_472_000_492_000_with_camera_labels 78 | segment-17778522338768131809_5920_000_5940_000_with_camera_labels 79 | segment-17885096890374683162_755_580_775_580_with_camera_labels 80 | segment-17987556068410436875_520_610_540_610_with_camera_labels 81 | segment-18141076662151909970_2755_710_2775_710_with_camera_labels -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/waymo_dynamic_81_dbg.lst: -------------------------------------------------------------------------------- 1 | segment-9653249092275997647_980_000_1000_000_with_camera_labels -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/waymo_filereader.py: -------------------------------------------------------------------------------- 1 | """ 2 | Modified from https://github.com/gdlg/simple-waymo-open-dataset-reader 3 | """ 4 | 5 | # Copyright (c) 2019, Grégoire Payen de La Garanderie, Durham University 6 | # 7 | # Licensed under the Apache License, Version 2.0 (the "License"); 8 | # you may not use this file except in compliance with the License. 9 | # You may obtain a copy of the License at 10 | # 11 | # http://www.apache.org/licenses/LICENSE-2.0 12 | # 13 | # Unless required by applicable law or agreed to in writing, software 14 | # distributed under the License is distributed on an "AS IS" BASIS, 15 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 16 | # See the License for the specific language governing permissions and 17 | # limitations under the License. 18 | # ============================================================================== 19 | 20 | import struct 21 | from waymo_open_dataset import dataset_pb2 22 | 23 | def read_record(f, header_only = False): 24 | header = f.read(12) 25 | if header == b'': 26 | raise StopIteration() 27 | length, lengthcrc = struct.unpack("QI", header) 28 | if header_only: 29 | # Skip length+4 bytes ahead 30 | f.seek(length+4,1) 31 | return None 32 | else: 33 | data = f.read(length) 34 | datacrc = struct.unpack("I",f.read(4)) 35 | 36 | frame = dataset_pb2.Frame() 37 | frame.ParseFromString(data) 38 | return frame 39 | 40 | class WaymoDataFileReader: 41 | def __init__(self, filename): 42 | self.filename = filename 43 | with open(self.filename, 'rb') as f: 44 | f.seek(0,0) 45 | table = [] 46 | while f: 47 | offset = f.tell() 48 | try: 49 | read_record(f, header_only=True) 50 | table.append(offset) 51 | except StopIteration: 52 | break 53 | f.seek(0,0) 54 | 55 | self.table = table 56 | 57 | def __len__(self): 58 | return len(self.table) 59 | 60 | def __iter__(self): 61 | with open(self.filename, 'rb') as f: 62 | f.seek(0,0) 63 | while f: 64 | try: 65 | yield read_record(f) 66 | except StopIteration: 67 | break -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/waymo_static_32.lst: -------------------------------------------------------------------------------- 1 | segment-1172406780360799916_1660_000_1680_000_with_camera_labels 2 | segment-4058410353286511411_3980_000_4000_000_with_camera_labels 3 | segment-10061305430875486848_1080_000_1100_000_with_camera_labels 4 | segment-14869732972903148657_2420_000_2440_000_with_camera_labels 5 | segment-16646360389507147817_3320_000_3340_000_with_camera_labels 6 | segment-15062351272945542584_5921_360_5941_360_with_camera_labels 7 | segment-13238419657658219864_4630_850_4650_850_with_camera_labels 8 | segment-13476374534576730229_240_000_260_000_with_camera_labels 9 | segment-14424804287031718399_1281_030_1301_030_with_camera_labels 10 | segment-15270638100874320175_2720_000_2740_000_with_camera_labels 11 | segment-15349503153813328111_2160_000_2180_000_with_camera_labels 12 | segment-15868625208244306149_4340_000_4360_000_with_camera_labels 13 | segment-16608525782988721413_100_000_120_000_with_camera_labels 14 | segment-17761959194352517553_5448_420_5468_420_with_camera_labels 15 | segment-3224923476345749285_4480_000_4500_000_with_camera_labels 16 | segment-3425716115468765803_977_756_997_756_with_camera_labels 17 | segment-3988957004231180266_5566_500_5586_500_with_camera_labels 18 | segment-9385013624094020582_2547_650_2567_650_with_camera_labels 19 | segment-8811210064692949185_3066_770_3086_770_with_camera_labels 20 | segment-10275144660749673822_5755_561_5775_561_with_camera_labels 21 | segment-10676267326664322837_311_180_331_180_with_camera_labels 22 | segment-12879640240483815315_5852_605_5872_605_with_camera_labels 23 | segment-13142190313715360621_3888_090_3908_090_with_camera_labels 24 | segment-13196796799137805454_3036_940_3056_940_with_camera_labels 25 | segment-14348136031422182645_3360_000_3380_000_with_camera_labels 26 | segment-15365821471737026848_1160_000_1180_000_with_camera_labels 27 | segment-16470190748368943792_4369_490_4389_490_with_camera_labels 28 | segment-11379226583756500423_6230_810_6250_810_with_camera_labels 29 | segment-13085453465864374565_2040_000_2060_000_with_camera_labels 30 | segment-14004546003548947884_2331_861_2351_861_with_camera_labels 31 | segment-15221704733958986648_1400_000_1420_000_with_camera_labels 32 | segment-16345319168590318167_1420_000_1440_000_with_camera_labels 33 | -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/waymo_static_new8.lst: -------------------------------------------------------------------------------- 1 | segment-5328596138024684667_2180_000_2200_000_with_camera_labels 2 | segment-8345535260120974350_1980_000_2000_000_with_camera_labels 3 | segment-8494653877777333091_540_000_560_000_with_camera_labels 4 | segment-10017090168044687777_6380_000_6400_000_with_camera_labels 5 | segment-10096619443888687526_2820_000_2840_000_with_camera_labels 6 | segment-13667377240304615855_500_000_520_000_with_camera_labels 7 | segment-14766384747691229841_6315_730_6335_730_with_camera_labels 8 | segment-17330200445788773877_2700_000_2720_000_with_camera_labels -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/waymo_unanno.lst: -------------------------------------------------------------------------------- 1 | segment-2711351338963414257_1360_000_1380_000_with_camera_labels 2 | segment-1231623110026745648_480_000_500_000_with_camera_labels -------------------------------------------------------------------------------- /dataio/autonomous_driving/waymo/zip_selected_seqs.py: -------------------------------------------------------------------------------- 1 | """ 2 | Compress processed data of selected waymo sequences to a zip file (for data transfering between devices) 3 | """ 4 | 5 | from tqdm import tqdm 6 | from tqdm.contrib.concurrent import thread_map 7 | from pathlib import Path 8 | from zipfile import ZIP_DEFLATED, ZipFile 9 | 10 | from os import PathLike 11 | from typing import List, Union 12 | 13 | def make_zip( 14 | zip_name: str, root_dir: Union[str, PathLike], select_scene_ids: List[str], 15 | fields=['images', 'lidars', 'masks', 'normals', 'depths', 'scenarios']): 16 | src_path = Path(root_dir).expanduser().resolve(strict=True) 17 | with ZipFile(zip_name, 'w', ZIP_DEFLATED) as zf: 18 | def recursive_append_to_zip(path: Path): 19 | # for file in path.rglob('*.npz'): 20 | for file in path.rglob('*'): 21 | zf.write(file, file.relative_to(src_path)) 22 | for i, scene_id in enumerate(tqdm(select_scene_ids, f'compressing ...')): 23 | for field in fields: 24 | if field == 'scenarios': 25 | scenario_file = src_path.joinpath(scene_id, f"scenario.pt") 26 | zf.write(scenario_file, scenario_file.relative_to(src_path)) 27 | else: 28 | recursive_append_to_zip(src_path.joinpath(scene_id, field)) 29 | 30 | def make_multi_zip( 31 | out_root: Union[str, PathLike], root: Union[str, PathLike], select_scene_ids: List[str], 32 | fields=['images', 'lidars', 'masks', 'normals', 'depths', 'scenarios']): 33 | src_path = Path(root).expanduser().resolve(strict=True) 34 | dst_path = Path(out_root).expanduser().resolve(strict=True) 35 | 36 | def make_single_zip(scene_id: str): 37 | zip_name = dst_path.joinpath(f"{scene_id}.zip") 38 | with ZipFile(zip_name, 'w', ZIP_DEFLATED) as zf: 39 | def recursive_append_to_zip(path: Path): 40 | for file in path.rglob('*'): 41 | zf.write(file, file.relative_to(src_path)) 42 | for field in fields: 43 | if field == 'scenarios': 44 | scenario_file = src_path.joinpath(scene_id, f"scenario.pt") 45 | zf.write(scenario_file, scenario_file.relative_to(src_path)) 46 | else: 47 | recursive_append_to_zip(src_path.joinpath(scene_id, field)) 48 | 49 | thread_map(make_single_zip, select_scene_ids, desc='Compressing') 50 | 51 | if __name__ == "__main__": 52 | import argparse 53 | parser = argparse.ArgumentParser() 54 | parser.add_argument("--seq_list", type=str, help='specify --seq_list if you want to limit the list of seqs', default="dataio/autonomous_driving/waymo/seq_list_select_20221214.lst") 55 | parser.add_argument("--data_root", type=str, default="/data1/waymo/processed") 56 | parser.add_argument("--out", help='output zip file path', default=None) 57 | parser.add_argument("--out_root", help='output zip files root dir', default=None) 58 | args = parser.parse_args() 59 | 60 | with open(args.seq_list, 'r') as f: 61 | seq_list = f.read().splitlines() 62 | select_scene_ids = [s.split(',')[0].rstrip(".tfrecord") for s in seq_list] 63 | 64 | if args.out is not None: 65 | make_zip(args.out, args.data_root, select_scene_ids) 66 | # make_zip(args.out, args.data_root, select_scene_ids, ['masks_vit_adapter']) 67 | elif args.out_root is not None: 68 | make_multi_zip(args.out_root, args.data_root, select_scene_ids) 69 | -------------------------------------------------------------------------------- /dataio/autonomous_driving/zod/zod_dataset.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/dataio/autonomous_driving/zod/zod_dataset.py -------------------------------------------------------------------------------- /dataio/block_nerf/README.md: -------------------------------------------------------------------------------- 1 | ## (WIP) -------------------------------------------------------------------------------- /dataio/block_nerf/__init__.py: -------------------------------------------------------------------------------- 1 | from .block_nerf_dataset import BlockNeRFDataset -------------------------------------------------------------------------------- /dataio/bmvs/README.md: -------------------------------------------------------------------------------- 1 | # [neuralsim] BlendedMVS dataset 2 | 3 | ## Usage 4 | 5 | ### Download 6 | 7 | #### > Tested small split (187MB) (already normalized) 8 | 9 | [Google Drive](https://drive.google.com/file/d/13hpXTYtjXNDu1HyFZSy66-ZF1H6wFPYW/view?usp=drive_link) 10 | 11 | #### > Full (27.5GB) 12 | 13 | Download the [BlendedMVS](https://github.com/YoYo000/BlendedMVS) dataset via [this link](https://github.com/YoYo000/BlendedMVS#download). 14 | 15 | For now, we have tested on the `low-res set` of the `BlendedMVS` split (27.5GB). 16 | 17 | ### Normalization 18 | 19 | > This is only for the original BMVS dataset. We have provided a small tested split with already normalized cams. See links above. 20 | 21 | The camera poses in [BlendedMVS](https://github.com/YoYo000/BlendedMVS) dataset have random centers and random scales. Therefore, it is necessary to first normalize the camera centers to the origin and standardize the scale of the camera distances. 22 | 23 | In order to achieve this, we have developed a simple script to find the focus center of **multiple object-centric views**, by minimizing the average projected pixel distance w.r.t. the image centers across all frames. This script only requires camera intrinsics and poses. 24 | 25 | ```shell 26 | cd /path/to/neuralsim 27 | python dataio/bmvs/normalize_bmvs.py --root /path/to/bmvs 28 | ``` 29 | 30 | It will generate `cameras_sphere.npz` files in the instance directories: 31 | 32 | ``` 33 | /path/to/bmvs 34 | ├── 5aa0f9d7a9efce63548c69a1 35 | │ ├── cameras_sphere.npz 36 | │ ├── blended_image 37 | │ │ ├── 00000000.jpg 38 | │ │ ├── 00000001.jpg 39 | │ │ ├── ... 40 | │ ├── cams 41 | │ │ ├── 00000000_cam.txt 42 | │ │ ├── 00000001_cam.txt 43 | │ │ ├── ... 44 | ├── 5aa235f64a17b335eeaf9609 45 | │ ├── cameras_sphere.npz 46 | │ ├── ... 47 | ├── ... 48 | ``` 49 | 50 | ## DEBUG 51 | 52 | ```shell 53 | cd /path/to/neuralsim 54 | source set_env.sh 55 | python dataio/bmvs/bmvs_dataset.py 56 | ``` 57 | -------------------------------------------------------------------------------- /dataio/bmvs/__init__.py: -------------------------------------------------------------------------------- 1 | from .bmvs_dataset import BMVSDataset -------------------------------------------------------------------------------- /dataio/bmvs/bmvs_dataset.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file bmvs_dataset.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Dataset IO for IDR/NeuS/DTU format datasets. 5 | """ 6 | 7 | import os 8 | import numpy as np 9 | from typing import Any, Dict, List, Tuple, Union 10 | 11 | from nr3d_lib.config import ConfigDict 12 | from nr3d_lib.utils import load_rgb, glob_imgs, get_image_size 13 | from nr3d_lib.graphics.cameras import decompose_intr_c2w_from_proj_np 14 | 15 | from dataio.scene_dataset import SceneDataset 16 | 17 | class BMVSDataset(SceneDataset): 18 | def __init__(self, config: dict) -> None: 19 | self.config = config 20 | self.populate(**config) 21 | 22 | def populate( 23 | self, 24 | root: str, 25 | instance_id: str, 26 | cam_file='cameras_sphere.npz', 27 | scale_radius=-1): 28 | 29 | self.main_class_name = "Main" 30 | self.instance_id = instance_id 31 | self.instance_dir = os.path.join(root, instance_id) 32 | 33 | assert os.path.exists(self.instance_dir), f"Not exist: {self.instance_dir}" 34 | 35 | self.image_paths = list(sorted(glob_imgs(os.path.join(self.instance_dir, 'blended_images')))) 36 | self.image_paths = [p for p in self.image_paths if not "masked" in p] 37 | self.n_images = len(self.image_paths) 38 | 39 | self.cam_file = os.path.join(self.instance_dir, cam_file) 40 | 41 | camera_dict = np.load(self.cam_file) 42 | scale_mats = [camera_dict['scale_mat_%d' % idx].astype(np.float32) for idx in range(self.n_images)] 43 | world_mats = [camera_dict['world_mat_%d' % idx].astype(np.float32) for idx in range(self.n_images)] 44 | 45 | intrs_all = [] 46 | c2ws_all = [] 47 | cam_center_norms = [] 48 | for scale_mat, world_mat in zip(scale_mats, world_mats): 49 | P = world_mat @ scale_mat 50 | P = P[:3, :4] 51 | intrinsics, pose = decompose_intr_c2w_from_proj_np(P) 52 | cam_center_norms.append(np.linalg.norm(pose[:3,3])) 53 | intrs_all.append(intrinsics.astype(np.float32)) 54 | c2ws_all.append(pose.astype(np.float32)) 55 | 56 | self.intrs_all = np.array(intrs_all) 57 | self.c2ws_all = np.array(c2ws_all) 58 | max_cam_norm = max(cam_center_norms) 59 | if scale_radius > 0: 60 | for i in range(len(self.c2ws_all)): 61 | self.c2ws_all[i][:3, 3] *= (scale_radius / max_cam_norm / 1.1) 62 | 63 | hw = [] 64 | for image_path in self.image_paths: 65 | W, H = get_image_size(image_path) 66 | hw.append([H, W]) 67 | hw = np.array(hw) 68 | self.hws_all = hw 69 | self.scale_mat = scale_mats[0] 70 | 71 | def get_scenario(self, scene_id: str, **kwargs) -> Dict[str, Any]: 72 | metas = dict( 73 | n_frames=self.n_images, 74 | main_class_name=self.main_class_name 75 | ) 76 | cam = dict( 77 | id='camera', 78 | class_name='Camera', 79 | n_frames=self.n_images, 80 | data=dict( 81 | hw=self.hws_all, 82 | intr=self.intrs_all, 83 | transform=self.c2ws_all, 84 | global_frame_inds=np.arange(self.n_images) 85 | ) 86 | ) 87 | obj = dict( 88 | id=self.instance_id, 89 | class_name=self.main_class_name, 90 | # Has no recorded data. 91 | ) 92 | scenario = dict( 93 | scene_id=f"BMVS-{self.instance_id}", 94 | metas=metas, 95 | objects={obj['id']: obj}, 96 | observers={cam['id']: cam} 97 | ) 98 | return scenario 99 | 100 | def get_image_wh(self, scene_id: str, camera_id: str, frame_index: Union[int, List[int]]): 101 | return self.hws_all[frame_index][::-1] 102 | 103 | def get_image(self, scene_id: str, camera_id: str, frame_index: int) -> np.ndarray: 104 | fpath = self.image_paths[frame_index] 105 | return load_rgb(fpath) 106 | 107 | if __name__ == '__main__': 108 | def unit_test_cams(): 109 | eg_config = ConfigDict( 110 | root="/data1/bmvs", 111 | # instance_id="5aa515e613d42d091d29d300", 112 | # instance_id="5a4a38dad38c8a075495b5d2", 113 | instance_id="5a8315f624b8e938486e0bd8", 114 | cam_file="cameras_sphere.npz" 115 | ) 116 | instance = BMVSDataset(eg_config) 117 | 118 | [print(p) for p in instance.image_paths] 119 | 120 | # from nr3d_lib.plot import vis_camera_o3d_from_arrays 121 | # vis_camera_o3d_from_arrays( 122 | # instance.intrs_all, instance.c2ws_all, instance.hws_all[..., 0], instance.hws_all[..., 1], 123 | # cam_size=0.03) 124 | 125 | import matplotlib.pyplot as plt 126 | from nr3d_lib.plot import vis_camera_mplot 127 | fig = plt.figure() 128 | ax = plt.axes(projection='3d') 129 | vis_camera_mplot(ax, instance.intrs_all, instance.c2ws_all, instance.hws_all[0, 0], instance.hws_all[0, 1], 130 | cam_size=0.03, annotation=True, per_cam_axis=False) 131 | plt.show() 132 | 133 | unit_test_cams() 134 | -------------------------------------------------------------------------------- /dataio/bmvs/normalized.lst: -------------------------------------------------------------------------------- 1 | 58c4bb4f4a69c55606122be4 2 | 58cf4771d0f5fb221defe6da 3 | 58d36897f387231e6c929903 4 | 58eaf1513353456af3a1682a 5 | 58f7f7299f5b5647873cb110 6 | 59056e6760bb961de55f3501 7 | 59338e76772c3e6384afbb15 8 | 59350ca084b7f26bf5ce6eb8 9 | 5947719bf1b45630bd096665 10 | 5947b62af1b45630bd0c2a02 11 | 59817e4a1bd4b175e7038d19 12 | 599aa591d5b41f366fed0d58 13 | 59d2657f82ca7774b1ec081d 14 | 59e75a2ca9e91f2c5526005d 15 | 59e864b2a9e91f2c5529325f 16 | 59ecfd02e225f6492d20fcc9 17 | 59f363a8b45be22330016cad 18 | 59f87d0bfa6280566fb38c9a 19 | 5a0271884e62597cdee0d0eb 20 | 5a3ca9cb270f0e3f14d0eddb 21 | 5a3cb4e4270f0e3f14d12f43 22 | 5a3f4aba5889373fbbc5d3b5 23 | 5a489fb1c7dab83a7d7b1070 24 | 5a48ba95c7dab83a7d7b44ed 25 | 5a48c4e9c7dab83a7d7b5cc7 26 | 5a48d4b2c7dab83a7d7b9851 27 | 5a4a38dad38c8a075495b5d2 28 | 5a563183425d0f5186314855 29 | 5a572fd9fc597b0478a81d14 30 | 5a57542f333d180827dfc132 31 | 5a588a8193ac3d233f77fbca 32 | 5a618c72784780334bc1972d 33 | 5a6400933d809f1d8200af15 34 | 5a6464143d809f1d8208c43c 35 | 5a69c47d0d5d0a7f3b2e9752 36 | 5a7d3db14989e929563eb153 37 | 5a8315f624b8e938486e0bd8 38 | 5a8aa0fab18050187cbe060e 39 | 5a969eea91dfc339a9a3ad2c 40 | 5aa0f9d7a9efce63548c69a1 41 | 5aa235f64a17b335eeaf9609 42 | 5aa515e613d42d091d29d300 43 | 5ab85f1dac4291329b17cb50 44 | 5ab8713ba3799a1d138bd69a 45 | 5ab8b8e029f5351f7f2ccf59 46 | 5abc2506b53b042ead637d86 47 | 5acf8ca0f3d8a750097e4b15 48 | 5adc6bd52430a05ecb2ffb85 49 | 5ae2e9c5fe405c5076abc6b2 50 | 5af28cea59bc705737003253 51 | 5afacb69ab00705d0cefdd5b 52 | 5b192eb2170cf166458ff886 53 | 5b21e18c58e2823a67a10dd8 54 | 5b22269758e2823a67a3bd03 55 | 5b2c67b5e0878c381608b8d8 56 | 5b3b353d8d46a939f93524b9 57 | 5b4933abf2b5f44e95de482a 58 | 5b6e716d67b396324c2d77cb 59 | 5b6eff8b67b396324c5b2672 60 | 5b78e57afc8fcf6781d0c3ba 61 | 5b908d3dc6ab78485f3d24a9 62 | 5b950c71608de421b1e7318f 63 | 5ba19a8a360c7c30c1c169df 64 | 5ba75d79d76ffa2c86cf2f05 65 | 5bb7a08aea1cfa39f1a947ab 66 | 5bc5f0e896b66a2cd8f9bd36 67 | 5bccd6beca24970bce448134 68 | 5bce7ac9ca24970bce4934b6 69 | 5bcf979a6d5f586b95c258cd 70 | 5bd43b4ba6b28b1ee86b92dd 71 | 5be3a5fb8cfdd56947f6b67c 72 | 5be3ae47f44e235bdbbc9771 73 | 5be47bf9b18881428d8fbc1d 74 | 5be4ab93870d330ff2dce134 75 | 5be883a4f98cee15019d5b83 76 | 5bea87f4abd34c35e1860ab5 77 | 5beb6e66abd34c35e18e66b9 78 | 5bf03590d4392319481971dc 79 | 5bf17c0fd439231948355385 80 | 5bf18642c50e6f7f8bdbd492 81 | 5bf21799d43923194842c001 82 | 5bf3a82cd439231948877aed 83 | 5bf7d63575c26f32dbf7413b 84 | 5bfc9d5aec61ca1dd69132a2 85 | 5bfd0f32ec61ca1dd69dc77b 86 | 5bff3c5cfe0ea555e6bcbf3a 87 | 5c062d84a96e33018ff6f0a6 88 | 5c0d13b795da9479e12e2ee9 89 | 5c1892f726173c3a09ea9aeb 90 | 5c189f2326173c3a09ed7ef3 91 | 5c1af2e2bee9a723c963d019 92 | 5c1b1500bee9a723c96c3e78 93 | 5c1dbf200843bc542d8ef8c4 94 | 5c20ca3a0843bc542d94e3e2 95 | 5c2b3ed5e611832e8aed46bf 96 | 5c34300a73a8df509add216d 97 | 5c34529873a8df509ae57b58 98 | -------------------------------------------------------------------------------- /dataio/colmap/README.md: -------------------------------------------------------------------------------- 1 | 2 | # [neuralsim] COLMAP-like dataset 3 | 4 | This SceneDataset can be used for any dataset collected with COLMAP. 5 | 6 | The directory file structure should be as follows: 7 | 8 | ``` 9 | data_dir 10 | ├── images 11 | │ ├── 00000000.jpg 12 | │ ├── 00000001.jpg 13 | │ ├── ... 14 | └── sparse 15 | └── 0 16 | ├── cameras.bin (or cameras.txt) 17 | ├── images.bin (or images.txt) 18 | └── points3D.bin 19 | ``` 20 | 21 | ## Extract monocular normals & depths priors 22 | 23 | NOTE: 24 | 25 | - Normal cues are generally more important than depth cues. In most cases, using only normal cues is sufficient. 26 | - The scale and shift of monocular depths have no correlation with real-world depths. They can only be indirectly used as weak hints. 27 | - The inferred normals are in range `[-1,1]`. The inferred depths are typically in range `[0,1]`. 28 | 29 | ### Setup `omnidata` 30 | 31 | Clone `omnidata` and install requirements. 32 | 33 | ```shell 34 | # Clone omnidata into your favorite directory 35 | git clone https://github.com/EPFL-VILAB/omnidata 36 | 37 | # Install dependencies 38 | pip install einops joblib pandas h5py scipy seaborn kornia timm pytorch-lightning 39 | ``` 40 | 41 | Download pretrained models following [download weights and code](https://github.com/EPFL-VILAB/omnidata/tree/main/omnidata_tools/torch#pretrained-models). 42 | 43 | NOTE: 44 | 45 | - If you encounter `gdown` error - access denied, try this answer: https://github.com/EPFL-VILAB/omnidata/issues/52#issuecomment-1627969898 46 | 47 | ### Extract mono cues 48 | 49 | ```shell 50 | cd /path/to/neuralsim/dataio/colmap/ 51 | 52 | # Extract depth 53 | python extract_mono_cues.py --task=depth --data_dir=/path/to/data_dir --omnidata_path=/path/to/omnidata/omnidata_tools/torch/ 54 | 55 | # Extract normals 56 | python extract_mono_cues.py --task=normal --data_dir=/path/to/data_dir --omnidata_path=/path/to/omnidata/omnidata_tools/torch/ 57 | ``` 58 | 59 | NOTE: You can pass `--verbose` if needed. -------------------------------------------------------------------------------- /dataio/colmap/__init__.py: -------------------------------------------------------------------------------- 1 | from .colmap_dataset import COLMAPDataset -------------------------------------------------------------------------------- /dataio/custom_old/__init__.py: -------------------------------------------------------------------------------- 1 | from .custom_dataset import CustomDataset -------------------------------------------------------------------------------- /dataio/data_loader/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | @author Jianfei Guo, Shanghai AI Lab 3 | @brief This module implements a generic dataloader containing scene graph, images, LiDAR data, and annotations, \ 4 | and defines different specialized datasets for pixel, image, and LiDAR sampling. 5 | 6 | FEATURES: 7 | - `SceneDataLoader` 8 | - Defines the basic and common IO logic for the dataset. 9 | - Loads and caches common APIs used in dataio/dataset_impl.py, including scenarios, images, LiDAR data, and annotations. 10 | - Performs image downsampling. 11 | 12 | - Defines 4 different specialized dataloader based on `SceneDataLoader`: 13 | - `PixelDataset` / `JointFramePixelDataset`: Sampling returns individual rays (pixels); supports importance sampling using error_map; 14 | - `PixelDataset`: in one sampled batch, rays originate from the same frame of image. 15 | - `JointFramePixelDataset`: in one sampled batch, rays could originate from different frames of image, \ 16 | since frame index is jointly sampled along with ray index in importance sampling 17 | - `ImageDataset`: Sampling returns the full (optionally downsampled) image; 18 | - `ImagePatchDataset`: Sampling extracts a patch from the image according to certain scaling and shifting rules; 19 | - `LidarDataset`: Sampling returns individual rays from the lidar. 20 | 21 | NOTE: 22 | (Coding conventions) 23 | - Avoid the use of "scene.slice_at" or any other operations related to the scene graph or nodes, or use them strictly within a "no_grad()" context. 24 | - The reason for this is that the scene graph may involve propagation of pose gradients across nodes, 25 | and we do not expect the dataloader to provide any gradients. 26 | These gradients should only be present in the forward process of the trainer. 27 | - In particular, calculations of camera rays should not be performed here, as they may require pose gradients. 28 | Instead, this code outputs `rays_xy` (the sampled pixel location in [0,1]) and `rays_sel` (an optional selector). 29 | for the "cam.get_selected_rays" function in the trainer's forward method. 30 | - LiDAR's merging and filter_in_cams require the scene graph, but should be performed within a "no_grad()" context. 31 | """ 32 | 33 | from .base_loader import * 34 | from .pixel_loader import * 35 | from .lidar_loader import * 36 | from .image_loader import * -------------------------------------------------------------------------------- /dataio/data_loader/unit_test.py: -------------------------------------------------------------------------------- 1 | from dataio.data_loader import * 2 | 3 | if __name__ == "__main__": 4 | import torch 5 | def unit_test(device=torch.device('cuda')): 6 | import numpy as np 7 | from icecream import ic 8 | from torch.utils.data.dataloader import DataLoader 9 | from nr3d_lib.utils import import_str 10 | from nr3d_lib.config import load_config, ConfigDict 11 | from app.resources import create_scene_bank, AssetBank 12 | from dataio.scene_dataset import SceneDataset 13 | 14 | dataset_cfg = ConfigDict( 15 | target='dataio.autonomous_driving.WaymoDataset', 16 | param=ConfigDict( 17 | root='/data1/waymo/processed', 18 | rgb_dirname="images", 19 | lidar_dirname="lidars" 20 | ) 21 | ) 22 | 23 | scenebank_cfg = ConfigDict( 24 | scenarios=['segment-7670103006580549715_360_000_380_000_with_camera_labels, 15'], 25 | observer_cfgs=ConfigDict( 26 | Camera=ConfigDict( 27 | list=['camera_FRONT', 'camera_FRONT_LEFT', 'camera_FRONT_RIGHT', 28 | 'camera_SIDE_LEFT', 'camera_SIDE_RIGHT'] 29 | ), 30 | RaysLidar=ConfigDict( 31 | list=['lidar_TOP', 'lidar_FRONT', 'lidar_REAR', 32 | 'lidar_SIDE_LEFT', 'lidar_SIDE_RIGHT'] 33 | ), 34 | ), 35 | object_cfgs=ConfigDict( 36 | Vehicle=ConfigDict( 37 | dynamic_only=False 38 | ), 39 | Pedestrian=ConfigDict( 40 | dynamic_only=False 41 | ) 42 | ), 43 | no_objects=False, 44 | align_orientation=False, 45 | aabb_extend=120., 46 | consider_distortion=True 47 | ) 48 | 49 | dataset_impl: SceneDataset = import_str(dataset_cfg.target)(dataset_cfg.param) 50 | 51 | scene_bank, _ = create_scene_bank( 52 | dataset=dataset_impl, device=device, 53 | scenebank_cfg=scenebank_cfg, 54 | drawable_class_names=['Vehicle', 'Pedestrian', 'Street', 'Distant'], 55 | misc_node_class_names=['node', 'EgoVehicle', 'EgoDrone'], 56 | ) 57 | 58 | scene = scene_bank[0] 59 | 60 | scene_dataloader = SceneDataLoader( 61 | scene_bank, dataset_impl, device=device, 62 | config=ConfigDict( 63 | preload=False, 64 | tags=ConfigDict( 65 | camera=ConfigDict( 66 | downscale=2, 67 | list=scenebank_cfg.observer_cfgs.Camera.list 68 | ), 69 | lidar=ConfigDict( 70 | list=scenebank_cfg.observer_cfgs.RaysLidar.list, 71 | multi_lidar_merge=True, 72 | filter_kwargs=ConfigDict( 73 | filter_valid=True, 74 | filter_in_cams=False, 75 | filter_out_objs=False 76 | ), 77 | downsample_cfg=ConfigDict( 78 | lidar_TOP=4 79 | ) 80 | ) 81 | ) 82 | )) 83 | 84 | cam0 = scene.observers[scene_dataloader.cam_id_list[0]] 85 | lidar0 = scene.observers[scene_dataloader.lidar_id_list[0]] 86 | 87 | multi_cam_weight = np.random.rand(len(scene_dataloader.cam_id_list)) 88 | multi_cam_weight /= multi_cam_weight.sum() 89 | pixel_dataset = PixelDataset( 90 | scene_dataloader, 91 | equal_mode='ray_batch', num_rays=4096, 92 | camera_sample_mode='weighted', multi_cam_weight=multi_cam_weight, 93 | frame_sample_mode='error_map', 94 | pixel_sample_mode='error_map', error_map_hw=(128,128)) 95 | 96 | pixel_dataloader = DataLoader(pixel_dataset, batch_size=None, sampler=pixel_dataset.get_random_sampler()) 97 | for sample, ground_truth in pixel_dataloader: 98 | print(sample, ground_truth) 99 | break 100 | 101 | pixel_dataloader = DataLoader(pixel_dataset, batch_size=7, sampler=pixel_dataset.get_random_sampler()) 102 | for sample, ground_truth in pixel_dataloader: 103 | print(sample, ground_truth) 104 | break 105 | 106 | lidar_dataset = LidarDataset( 107 | scene_dataloader, 108 | equal_mode= 'ray_batch', num_rays=8192, 109 | frame_sample_mode = 'uniform', 110 | lidar_sample_mode = 'merged_weighted', 111 | multi_lidar_weight = [0.5, 0.1, 0.1, 0.1], 112 | ) 113 | # lidar_dataset.sample_merged(scene.id, 83, 8192) 114 | # lidar_dataset.sample_merged(scene.id, 91, 8192) 115 | lidar_dataset.sample_merged(scene.id, 132, 8192) 116 | # lidar_dataset.sample_merged(scene.id, 135, 8192) 117 | # lidar_dataset.sample_merged(scene.id, 186, 8192) 118 | # lidar_dataset.sample_merged(scene.id, 180, 8192) 119 | 120 | unit_test() -------------------------------------------------------------------------------- /dataio/data_loader/view_loader.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/dataio/data_loader/view_loader.py -------------------------------------------------------------------------------- /dataio/dtu/README.md: -------------------------------------------------------------------------------- 1 | # [neuralsim] DTU / IDR dataset 2 | 3 | This SceneDataset can be used for any dataset collected using the [NeuS](https://github.com/Totoro97/NeuS) / [IDR](https://github.com/lioryariv/idr) data format. 4 | 5 | Currently, the most typical example we have tested is the [DTU](https://roboimagedata.compute.dtu.dk/?page_id=36) and [BlendedMVS](https://github.com/YoYo000/BlendedMVS) dataset provided by [NeuS](https://github.com/Totoro97/NeuS). 6 | 7 | ## Usage 8 | 9 | ### Download 10 | 11 | Go to the [NeuS](https://github.com/Totoro97/NeuS) repo and download their data. 12 | 13 | ## DEBUG 14 | 15 | ```shell 16 | cd /path/to/neuralsim 17 | source set_env.sh 18 | python dataio/dtu/dtu_dataset.py 19 | ``` 20 | 21 | -------------------------------------------------------------------------------- /dataio/dtu/__init__.py: -------------------------------------------------------------------------------- 1 | from .dtu_dataset import DTUDataset -------------------------------------------------------------------------------- /dataio/gtav_nerf/gtav_nerf_dataset.py: -------------------------------------------------------------------------------- 1 | """ 2 | WIP 3 | """ -------------------------------------------------------------------------------- /dataio/mega_nerf/__init__.py: -------------------------------------------------------------------------------- 1 | from .mega_nerf_dataset import MegaNeRFDataset -------------------------------------------------------------------------------- /dataio/monosdf/README.md: -------------------------------------------------------------------------------- 1 | # [neuralsim] MonoSDF Dataset 2 | 3 | ## Usage 4 | 5 | ### Download 6 | 7 | Follow [this link](https://github.com/autonomousvision/monosdf#dataset) to download the MonoSDF's preprocessed data of [Replica](https://github.com/facebookresearch/Replica-Dataset) / [scannet](http://www.scan-net.org/) indoor datasets. -------------------------------------------------------------------------------- /dataio/monosdf/__init__.py: -------------------------------------------------------------------------------- 1 | from .monosdf_dataset import MonoSDFDataset -------------------------------------------------------------------------------- /dataio/nerf/nerf_dataset.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file nerf_dataset.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Dataset IO for NeRF-standard datasets 5 | """ 6 | 7 | import os 8 | import json 9 | import numpy as np 10 | from typing import Any, Dict, Literal 11 | 12 | from nr3d_lib.config import ConfigDict 13 | from nr3d_lib.utils import get_image_size, load_rgb 14 | 15 | from dataio.scene_dataset import SceneDataset 16 | 17 | class NeRFDataset(SceneDataset): 18 | def __init__(self, config: dict) -> None: 19 | self.config = config 20 | self.populate(**config) 21 | 22 | def populate( 23 | self, 24 | datadir: str, ): 25 | pass -------------------------------------------------------------------------------- /dataio/ners/__init__.py: -------------------------------------------------------------------------------- 1 | from .instance_dataset import MVMCNeRSInstanceDataset 2 | from .instance_dataset_cropped import MVMCNeRSInstanceDatasetCropped -------------------------------------------------------------------------------- /dataio/ners/instance_dataset.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file instance_dataset.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Dataset IO for NeRS's MVMC dataset; 5 | Single instance loader (with original image and its corresponding intrinsics). 6 | """ 7 | 8 | import os 9 | import json 10 | import numpy as np 11 | from typing import Any, Dict, Literal 12 | 13 | from nr3d_lib.config import ConfigDict 14 | from nr3d_lib.utils import get_image_size, load_rgb 15 | 16 | from dataio.scene_dataset import SceneDataset 17 | 18 | def rle_to_binary_mask(rle): 19 | """ 20 | rle should be coco format: {"counts": [], "size": []} 21 | """ 22 | if isinstance(rle, list): 23 | return np.stack([rle_to_binary_mask(r) for r in rle]) 24 | counts = rle["counts"] 25 | if isinstance(counts, str): 26 | counts = list(map(int, counts.split(" "))) 27 | mask = np.zeros(np.prod(rle["size"]), dtype=bool) 28 | running_length = 0 29 | for start, length in zip(counts[::2], counts[1::2]): 30 | running_length += start 31 | mask[running_length : running_length + length] = 1 32 | running_length += length 33 | return mask.reshape(rle["size"], order="F") 34 | 35 | class MVMCNeRSInstanceDataset(SceneDataset): 36 | def __init__(self, config: dict) -> None: 37 | self.config = config 38 | self.populate(**config) 39 | 40 | def populate( 41 | self, 42 | root: str, 43 | instance_id: str, 44 | camera_type: Literal['camera_optimized', 'camera_pretrained'] = 'camera_optimized' 45 | ): 46 | instance_id = str(instance_id) 47 | 48 | self.main_class_name = "Main" 49 | self.instance_id = instance_id 50 | self.instance_dir = os.path.join(root, instance_id) 51 | 52 | annotations_json = os.path.join(self.instance_dir, "annotations.json") 53 | with open(annotations_json) as f: 54 | annotations = json.load(f) 55 | 56 | hws_all = [] 57 | image_paths = [] 58 | Rs = [] 59 | Ts = [] 60 | fovs = [] 61 | masks = [] 62 | for annotation in annotations["annotations"]: 63 | image_path = os.path.join(self.instance_dir, "images", annotation["filename"]) 64 | image_paths.append(image_path) 65 | 66 | W, H = get_image_size(image_path) 67 | hws_all.append([H, W]) 68 | 69 | Rs.append(annotation[camera_type]["R"]) 70 | Ts.append(annotation[camera_type]["T"]) 71 | fovs.append(annotation[camera_type]["fov"]) 72 | 73 | mask = rle_to_binary_mask(annotation["mask"]) 74 | masks.append(mask) 75 | 76 | Rs, Ts, fovs = np.array(Rs), np.array(Ts), np.array(fovs) 77 | hws_all = np.array(hws_all) 78 | self.hws_all = hws_all 79 | masks = np.array(masks) 80 | self.masks = masks 81 | self.image_paths = image_paths 82 | self.n_images = len(self.image_paths) 83 | 84 | #------ Intrinsics 85 | fovs = np.deg2rad(fovs) 86 | focal_length = np.abs(1 / np.tan(fovs / 2)) 87 | focal_length_px = focal_length[..., None] * hws_all[..., [1,0]]/2. 88 | 89 | intrs_all = np.zeros([self.n_images, 3, 3]) 90 | intrs_all[..., 0, 0] = focal_length_px[..., 0] 91 | intrs_all[..., 1, 1] = focal_length_px[..., 1] 92 | intrs_all[..., 0, 2] = hws_all[..., 1] / 2. 93 | intrs_all[..., 1, 2] = hws_all[..., 0] / 2. 94 | intrs_all[..., 2, 2] = 1. 95 | self.intrs_all = intrs_all 96 | 97 | #------ Optimized camera poses from NeRS 98 | # From pytorch3d's w2c to OpenCV's c2w 99 | c2ws = np.eye(4)[None, ...].repeat(self.n_images, 0) 100 | c2ws[..., :3, :3] = Rs 101 | c2ws[..., :3, 3] = -np.einsum('...ij,...j->...i', Rs, Ts) 102 | c2ws = c2ws @ np.diagflat([-1,-1,1,1])[None,...] 103 | self.c2ws_all = c2ws 104 | 105 | def get_scenario(self, scene_id: str, **kwargs) -> Dict[str, Any]: 106 | metas = dict( 107 | n_frames=self.n_images, 108 | main_class_name=self.main_class_name 109 | ) 110 | cam = dict( 111 | id='camera', 112 | class_name='Camera', 113 | n_frames=self.n_images, 114 | data=dict( 115 | hw=self.hws_all, 116 | intr=self.intrs_all, 117 | transform=self.c2ws_all, 118 | global_frame_inds=np.arange(self.n_images) 119 | ) 120 | ) 121 | obj = dict( 122 | id=self.instance_id, 123 | class_name=self.main_class_name, 124 | # Has no recorded data. 125 | ) 126 | scenario = dict( 127 | scene_id=f"MVMC-{self.instance_id}", 128 | metas=metas, 129 | objects={obj['id']: obj}, 130 | observers={cam['id']: cam} 131 | ) 132 | return scenario 133 | 134 | def get_image(self, scene_id: str, camera_id: str, frame_index: int) -> np.ndarray: 135 | return load_rgb(self.image_paths[frame_index]) 136 | 137 | def get_image_occupancy_mask(self, scene_id: str, camera_id: str, frame_index: int) -> np.ndarray: 138 | return self.masks[frame_index] -------------------------------------------------------------------------------- /dataio/neural_recon_w/neural_recon_w_dataset.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file neural_recon_w_dataset.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Dataset IO for Neural Reconstruction in the wild datasets. 5 | """ 6 | 7 | import os 8 | import numpy as np 9 | from typing import Any, Dict 10 | 11 | from nr3d_lib.config import ConfigDict 12 | from nr3d_lib.utils import load_mask, load_rgb, glob_imgs, get_image_size 13 | 14 | from dataio.scene_dataset import SceneDataset 15 | 16 | class NeuralReconWDataset(SceneDataset): 17 | def __init__(self, config: dict) -> None: 18 | self.config = config 19 | self.populate(**config) 20 | 21 | def populate( 22 | self, 23 | root: str, 24 | 25 | ): 26 | pass -------------------------------------------------------------------------------- /dataio/neural_recon_w/preprocess.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | """ 4 | 5 | import cv2 6 | import numpy as np 7 | from tqdm import tqdm 8 | 9 | import torch 10 | -------------------------------------------------------------------------------- /dataio/scene_dataset.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file scene_dataset.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Dataset implementation abstract interfaces. 5 | """ 6 | 7 | import numpy as np 8 | from typing import Any, Dict, List, Literal, Tuple, Union 9 | from abc import ABC, abstractmethod 10 | 11 | from nr3d_lib.config import ConfigDict 12 | 13 | class SceneDataset(ABC): 14 | @abstractmethod # NOTE: This is the only method that must be implemented. 15 | def get_scenario(self, scene_id: str, **kwargs) -> Dict[str, Any]: 16 | # NOTE: Must be implemented to enable scene loading and the whole training framework. 17 | raise NotImplementedError 18 | 19 | @property 20 | def up_vec(self) -> np.ndarray: 21 | raise NotImplementedError 22 | 23 | @property 24 | def forward_vec(self) -> np.ndarray: 25 | raise NotImplementedError 26 | 27 | @property 28 | def right_vec(self) -> np.ndarray: 29 | raise NotImplementedError 30 | 31 | def get_all_available_scenarios(self) -> List[str]: 32 | # All available scene id list of current dataset 33 | raise NotImplementedError 34 | 35 | def get_metadata(self, scene_id: str) -> Dict[str, Any]: 36 | raise NotImplementedError 37 | 38 | def get_image_wh(self, scene_id: str, camera_id: str, frame_index: Union[int, List[int]]) -> np.ndarray: 39 | raise NotImplementedError 40 | 41 | def get_image(self, scene_id: str, camera_id: str, frame_index: int) -> np.ndarray: 42 | # [H, W, 3], np.float32, range [0,1] 43 | raise NotImplementedError 44 | 45 | def get_image_mono_depth(self, scene_id: str, camera_id: str, frame_index: int) -> np.ndarray: 46 | # [H, W], np.float32, range [0, inf] 47 | raise NotImplementedError 48 | 49 | def get_image_mono_normals(self, scene_id: str, camera_id: str, frame_index: int) -> np.ndarray: 50 | # [H, W, 3], np.float32, range [-1,1] 51 | raise NotImplementedError 52 | 53 | def get_image_occupancy_mask(self, scene_id: str, camera_id: str, frame_index: int) -> np.ndarray: 54 | # [H, W], bool, binary occupancy mask on RGB image. 1 for occpied, 0 for not. 55 | raise NotImplementedError 56 | 57 | def get_image_semantic_mask_by_type( 58 | self, scene_id: str, camera_id: str, 59 | sem_type: Literal['dynamic', 'human', 'road', 'anno_dontcare'], 60 | frame_index: int) -> np.ndarray: 61 | # Binary semantic mask on RGB image. 1 for matched, 0 for not. 62 | raise NotImplementedError 63 | 64 | def get_image_semantic_mask_all(self, scene_id: str, camera_id: str, frame_index: int) -> np.ndarray: 65 | # [H, W], int16, integer semantic mask on RGB image. 66 | raise NotImplementedError 67 | 68 | def get_lidar(self, scene_id: str, lidar_id: str, frame_index: int) -> Dict[str, np.ndarray]: 69 | # { 70 | # 'rays_o': [..., 3], np.float32, lidar beam's starting points 71 | # 'rays_d': [..., 3], np.float32, lidar beam's direction vectors 72 | # 'ranges': [...], np.float32, lidar beam's termination depth (usally the first return) 73 | # } 74 | raise NotImplementedError 75 | 76 | # def get_aabb(self, scene_id: str, lidar_id: str, frame_index: int) -> Dict[str, np.ndarray]: 77 | # raise NotImplementedError 78 | 79 | # def create_dataset(self, source_data_cfg: dict, j=8): 80 | # raise NotImplementedError -------------------------------------------------------------------------------- /dataio/utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | @file utils.py 3 | @author Jianfei Guo, Shanghai AI Lab 4 | @brief Common utilities for processing datasets 5 | """ 6 | 7 | import numpy as np 8 | 9 | def clip_node_data(odict: dict, start, stop): 10 | """ 11 | In-place clip odict['data'][...], odict['start_frame'], odict['n_frames'] 12 | 13 | start/stop: [int] start/stop frame ind, or [float] start/stop timestamp 14 | both under start <= ... < stop convention, i.e. item \in [start,stop) 15 | """ 16 | assert 'data' in odict, "Only works with node with data" 17 | if (start is None and stop is None) or ('global_frame_inds' not in odict['data'] and 'global_timestamps' not in odict['data']): 18 | t = None 19 | else: 20 | if 'global_frame_inds' in odict['data']: 21 | assert (start is None or isinstance(start, int)) and (stop is None or isinstance(stop, int)), f"Please use frame ind clipping." 22 | t = odict['data']['global_frame_inds'] 23 | else: 24 | assert (start is None or isinstance(start, float)) and (stop is None or isinstance(stop, float)), f"Please use timestamp clipping." 25 | t = odict['data']['global_timestamps'] 26 | 27 | if t is not None: 28 | if start is None: 29 | mask = t < stop 30 | elif stop is None: 31 | mask = t >= start 32 | else: 33 | mask = (t >= start) & (t < stop) 34 | 35 | if ~mask.any(): 36 | # NOTE: DO NOT let any useful nodes to be not-used parent's children, 37 | # because all children and descendants of a not-used node will be ignored. 38 | return 39 | 40 | # Modified keys: data{}, start_frame, n_frames 41 | for k, v in odict['data'].items(): 42 | odict['data'][k] = v[mask] 43 | odict['n_frames'] = np.sum(mask) 44 | if 'global_frame_inds' in odict['data']: 45 | if start is not None: 46 | # TODO: This only supports `start`, `stop` being frame inds, not ts 47 | odict['data']['global_frame_inds'] -= start 48 | odict['start_frame'] = np.min(odict['data']['global_frame_inds']) 49 | return odict 50 | 51 | def clip_node_segments(odict: dict, start, stop): 52 | """ 53 | In-place clip odict['segments'] (seg['data'][...], seg['start_frame'], seg['n_frames']) 54 | """ 55 | assert 'segments' in odict, "Only works with node with segments" 56 | old_segs = odict.pop('segments') 57 | new_segs = [] 58 | for seg in old_segs: 59 | if (start is None and stop is None) or ('global_frame_inds' not in seg['data'] and 'global_timestamps' not in seg['data']): 60 | new_segs.append(seg) 61 | continue 62 | else: 63 | if 'global_frame_inds' in seg['data']: 64 | assert (start is None or isinstance(start, int)) and (stop is None or isinstance(stop, int)), f"Please use frame ind clipping." 65 | t = seg['data']['global_frame_inds'] 66 | else: 67 | assert (start is None or isinstance(start, float)) and (stop is None or isinstance(stop, float)), f"Please use timestamp clipping." 68 | t = seg['data']['global_timestamps'] 69 | 70 | if start is None: 71 | mask = t < stop 72 | elif stop is None: 73 | mask = t >= start 74 | else: 75 | mask = (t >= start) & (t < stop) 76 | 77 | if ~mask.any(): 78 | continue 79 | 80 | # Modified keys: data{}, start_frame, n_frames 81 | for k, v in seg['data'].items(): 82 | seg['data'][k] = v[mask] 83 | seg['n_frames'] = np.sum(mask) 84 | if 'global_frame_inds' in seg['data']: 85 | if start is not None: 86 | # TODO: This only supports `start`, `stop` being frame inds, not ts 87 | seg['data']['global_frame_inds'] -= start 88 | seg['start_frame'] = np.min(seg['data']['global_frame_inds']) 89 | new_segs.append(seg) 90 | 91 | odict['segments'] = new_segs 92 | return odict -------------------------------------------------------------------------------- /docs/exps/exp_lipshitz_3d.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | def set_env(depth: int): 4 | # Add project root to sys.path 5 | current_file_path = os.path.abspath(__file__) 6 | project_root_path = os.path.dirname(current_file_path) 7 | for _ in range(depth): 8 | project_root_path = os.path.dirname(project_root_path) 9 | if project_root_path not in sys.path: 10 | sys.path.append(project_root_path) 11 | print(f"Added {project_root_path} to sys.path") 12 | set_env(1) 13 | 14 | import os 15 | import numpy as np 16 | from tqdm import tqdm 17 | from functools import partial 18 | 19 | import torch 20 | import torch.nn as nn 21 | import torch.nn.functional as F 22 | from torch.optim import Adam 23 | 24 | from nr3d_lib.utils import cond_mkdir 25 | from nr3d_lib.models.blocks import LipshitzMLP, MLP 26 | from nr3d_lib.graphics.trianglemesh import extract_mesh 27 | 28 | import argparse 29 | parser = argparse.ArgumentParser() 30 | parser.add_argument("--lr", type=float, default=1.0e-3) 31 | # NOTE: Best for now: 3k -> w=3.0e-5; 30k -> w=1.0e-6 32 | parser.add_argument("--lipshitz", action='store_true') 33 | parser.add_argument("--w_lip", type=float, default=3.0e-5) 34 | parser.add_argument("--num_iters", type=int, default=3000) 35 | args = parser.parse_args() 36 | 37 | device = torch.device('cuda') 38 | dtype = torch.float32 39 | 40 | n_latent_dim = 8 41 | n_input_dim = 3 42 | 43 | exp_dir = "./dev_test/lip/3d/" 44 | cond_mkdir(exp_dir) 45 | 46 | if args.lipshitz: 47 | m = LipshitzMLP(n_input_dim + n_latent_dim, 1, D=4, W=256, c_init_factor=1.0, dtype=dtype, device=device) 48 | else: 49 | m = MLP(n_input_dim + n_latent_dim, 1, D=4, W=256, dtype=dtype, device=device) 50 | print(m) 51 | 52 | # https://iquilezles.org/articles/distfunctions/ 53 | def sdf_gt_cube(x: torch.Tensor, r=0.5): 54 | q = x.abs() - r 55 | return q.clamp_min(0.).norm(dim=-1) + q.max(dim=-1).values.clamp_max(0.) 56 | 57 | def sdf_gt_cube_frame(x: torch.Tensor, r=0.5, e=0.1): 58 | p = x.abs() - r 59 | q = (p+e).abs() - e 60 | v1 = torch.stack([p[..., 0], q[..., 1], q[..., 2]], dim=-1) 61 | v2 = torch.stack([q[..., 0], p[..., 1], q[..., 2]], dim=-1) 62 | v3 = torch.stack([q[..., 0], q[..., 1], p[..., 2]], dim=-1) 63 | vs = torch.stack([v1, v2, v3], dim=0) 64 | return (vs.clamp_min(0.).norm(dim=-1) + vs.max(dim=-1).values.clamp_max(0.)).min(dim=0).values 65 | 66 | def sdf_gt_sphere(x: torch.Tensor, r=0.5): 67 | return x.norm(dim=-1) - r 68 | 69 | sdf_gt_fn1 = partial(sdf_gt_sphere, r=0.7) 70 | sdf_gt_fn2 = partial(sdf_gt_cube_frame, r=0.5) 71 | sdf_gt_fn3 = partial(sdf_gt_cube, r=0.8) 72 | 73 | extract_mesh(sdf_gt_fn1, None, N=128, filepath=os.path.join(exp_dir, 'gt1.ply')) 74 | extract_mesh(sdf_gt_fn2, None, N=128, filepath=os.path.join(exp_dir, 'gt2.ply')) 75 | extract_mesh(sdf_gt_fn3, None, N=128, filepath=os.path.join(exp_dir, 'gt3.ply')) 76 | 77 | l1 = nn.Parameter(torch.empty([n_latent_dim], device=device).normal_(std=1.), requires_grad=True) 78 | l2 = nn.Parameter(torch.empty([n_latent_dim], device=device).normal_(std=1.), requires_grad=True) 79 | l3 = nn.Parameter(torch.empty([n_latent_dim], device=device).normal_(std=1.), requires_grad=True) 80 | 81 | def sdf_pred_fn(x: torch.Tensor, l: torch.Tensor): 82 | h = torch.cat([x, l.unsqueeze(0).expand(x.shape[0], -1)], dim=-1) 83 | return m.forward(h).squeeze_(-1) 84 | 85 | optimzer = Adam(list(m.parameters())+[l1,l2,l3], lr=1.0e-3) 86 | 87 | with tqdm(range(args.num_iters)) as pbar: 88 | for _ in pbar: 89 | optimzer.zero_grad() 90 | 91 | # [-1, 1] 92 | x = torch.rand([5000, n_input_dim], device=device) * 2 - 1 93 | 94 | sdf_pred1 = sdf_pred_fn(x, l1) 95 | sdf_pred2 = sdf_pred_fn(x, l2) 96 | sdf_pred3 = sdf_pred_fn(x, l3) 97 | 98 | sdf_gt1 = sdf_gt_fn1(x) 99 | sdf_gt2 = sdf_gt_fn2(x) 100 | sdf_gt3 = sdf_gt_fn3(x) 101 | 102 | loss = F.l1_loss(sdf_pred1, sdf_gt1) \ 103 | + F.l1_loss(sdf_pred2, sdf_gt2) \ 104 | + F.l1_loss(sdf_pred3, sdf_gt3) 105 | loss = loss.mean() 106 | 107 | if args.lipshitz: 108 | loss += args.w_lip * m.lipshitz_bound_full() 109 | 110 | loss.backward() 111 | optimzer.step() 112 | 113 | pbar.set_postfix(loss=loss.item()) 114 | 115 | with torch.no_grad(): 116 | extract_mesh(lambda x: sdf_pred_fn(x, l1), None, N=128, filepath=os.path.join(exp_dir, 'pred1.ply')) 117 | extract_mesh(lambda x: sdf_pred_fn(x, l2), None, N=128, filepath=os.path.join(exp_dir, 'pred2.ply')) 118 | extract_mesh(lambda x: sdf_pred_fn(x, l2), None, N=128, filepath=os.path.join(exp_dir, 'pred3.ply')) 119 | 120 | alphas = np.linspace(0, 1, endpoint=True, num=10).tolist() 121 | mesh_files = [] 122 | for alpha in alphas: 123 | l = l1 * (1-alpha) + l2 * alpha 124 | mesh_file = os.path.join(exp_dir, f'alpha={alpha:.2f}.ply') 125 | try: 126 | extract_mesh(lambda x: sdf_pred_fn(x, l), None, N=128, filepath=mesh_file) 127 | mesh_files.append(mesh_file) 128 | except: 129 | mesh_files.append(None) 130 | 131 | import open3d as o3d 132 | things_to_draw = [] 133 | for i, (alpha, mesh_file) in enumerate(zip(alphas, mesh_files)): 134 | if mesh_file is None: 135 | continue 136 | mesh = o3d.io.read_triangle_mesh(mesh_file) 137 | mesh.compute_vertex_normals() 138 | mesh.translate([-2*i, 0, 0]) 139 | things_to_draw.append(mesh) 140 | o3d.visualization.draw_geometries(things_to_draw) -------------------------------------------------------------------------------- /docs/exps/exp_permuto_2d_modulated.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/docs/exps/exp_permuto_2d_modulated.py -------------------------------------------------------------------------------- /docs/exps/exp_sky_oneframe.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | def set_env(depth: int): 4 | # Add project root to sys.path 5 | current_file_path = os.path.abspath(__file__) 6 | project_root_path = os.path.dirname(current_file_path) 7 | for _ in range(depth): 8 | project_root_path = os.path.dirname(project_root_path) 9 | if project_root_path not in sys.path: 10 | sys.path.append(project_root_path) 11 | print(f"Added {project_root_path} to sys.path") 12 | set_env(1) 13 | 14 | import pickle 15 | import numpy as np 16 | from tqdm import tqdm 17 | 18 | import torch 19 | from torch.optim import Adam 20 | import torch.nn.functional as F 21 | 22 | from nr3d_lib.graphics.cameras import pinhole_get_rays 23 | from nr3d_lib.config import ConfigDict 24 | from nr3d_lib.utils import check_to_torch, load_rgb 25 | from nr3d_lib.models.utils import batchify_query, get_scheduler 26 | from nr3d_lib.models.loss.recon import mse_loss 27 | 28 | from app.models.env.neural_sky import SimpleSky 29 | 30 | NUM_ITERS = 30000 31 | LR = 1.0e-4 32 | 33 | #---------------- Cityscapes semantic segmentation 34 | cityscapes_classes = [ 35 | 'road', 'sidewalk', 'building', 'wall', 'fence', 'pole', 36 | 'traffic light', 'traffic sign', 'vegetation', 'terrain', 'sky', 37 | 'person', 'rider', 'car', 'truck', 'bus', 'train', 'motorcycle', 38 | 'bicycle' 39 | ] 40 | cityscapes_classes_ind_map = {cn: i for i, cn in enumerate(cityscapes_classes)} 41 | 42 | device = torch.device('cuda') 43 | 44 | img = load_rgb('/data1/waymo/processed/images/segment-7670103006580549715_360_000_380_000_with_camera_labels/camera_FRONT/00000000.jpg') 45 | mask_all = np.load('/data1/waymo/processed/masks/segment-7670103006580549715_360_000_380_000_with_camera_labels/camera_FRONT/00000000.npz')['arr_0'] 46 | sky_mask = (mask_all==cityscapes_classes_ind_map['sky']) 47 | sky_mask = check_to_torch(sky_mask, device=device).flatten(0) 48 | 49 | H, W, *_ = img.shape 50 | rgb_gt = check_to_torch(img, device=device, dtype=torch.float).flatten(0,-2) 51 | 52 | # sky = SimpleSky({'type':'spherical', 'degree':8}, D=2, W=256, dtype=torch.float, device=device) 53 | sky = SimpleSky({'type':'sinusoidal', 'n_frequencies':10}, D=2, W=256, dtype=torch.float, device=device) 54 | optimzer = Adam(sky.parameters(), lr=LR) 55 | scheduler = get_scheduler(ConfigDict(type='exponential', num_iters=NUM_ITERS, min_factor=0.03), optimzer) 56 | 57 | with open('/data1/waymo/processed/scenarios/segment-7670103006580549715_360_000_380_000_with_camera_labels.pt', 'rb') as f: 58 | scenario = pickle.load(f) 59 | 60 | obs_data = scenario['observers']['camera_FRONT']['data'] 61 | intr = check_to_torch(obs_data['intr'], device=device, dtype=torch.float)[0] 62 | c2w = check_to_torch(obs_data['c2w'], device=device, dtype=torch.float)[0] 63 | 64 | rays_o_all, rays_d_all = pinhole_get_rays(c2w, intr, H, W) 65 | 66 | with tqdm(range(NUM_ITERS)) as pbar: 67 | for _ in pbar: 68 | optimzer.zero_grad() 69 | inds = torch.randint(H*W, [4096], device=device) 70 | rays_d = rays_d_all[inds] 71 | 72 | pred = sky.forward(F.normalize(rays_d, dim=-1)) 73 | gt = rgb_gt[inds] 74 | 75 | loss = mse_loss(pred, gt, sky_mask[inds].unsqueeze(-1)) 76 | loss.backward() 77 | 78 | pbar.set_postfix(loss=loss.item()) 79 | 80 | optimzer.step() 81 | scheduler.step() 82 | 83 | pred_img: torch.Tensor = batchify_query(sky.forward, rays_d_all, chunk=65536, show_progress=True) 84 | pred_img = pred_img.reshape(H, W, 3).data.cpu().numpy() -------------------------------------------------------------------------------- /docs/exps/permuto_enc_video.py: -------------------------------------------------------------------------------- 1 | import os 2 | import imageio 3 | import numpy as np 4 | from glob import glob 5 | from tqdm import tqdm 6 | from typing import Tuple 7 | import matplotlib.pyplot as plt 8 | 9 | import torch 10 | import torch.nn as nn 11 | from torch import optim 12 | 13 | from nr3d_lib.utils import load_rgb 14 | from nr3d_lib.config import ConfigDict 15 | from nr3d_lib.models.utils import get_scheduler 16 | from nr3d_lib.models.loss.recon import huber_loss 17 | from nr3d_lib.models.grid_encodings.permuto import GenerativePermutoConcat 18 | 19 | n_latent_dim = 1 20 | num_iters = 10000 21 | num_pixels = 100000 22 | lr = 0.01 23 | """ 24 | 0.03 ~ 0.1 barely allows the all-zero-initialized latent to exhibit a monotonically increasing behavior 25 | Below 0.01 is still not quite enough 26 | """ 27 | wreg = 0.03 28 | use_scheduler = True 29 | min_factor = 0.2 30 | # param_dtype = torch.float16 31 | param_dtype = torch.float32 32 | device = torch.device('cuda') 33 | 34 | m = GenerativePermutoConcat( 35 | n_latent_dim, 2, 3, 36 | encoding_cfg=ConfigDict(permuto_auto_compute_cfg=ConfigDict( 37 | type='multi_res', coarsest_res=4.0, finest_res=2000.0, n_levels=16, n_feats=4, 38 | log2_hashmap_size=18, apply_random_shifts_per_level=True)), 39 | decoder_cfg=ConfigDict(type='mlp', D=2, W=64, output_activation='sigmoid'), dtype=param_dtype, device=device 40 | ) 41 | 42 | print(m) 43 | 44 | data_root = '/data1/video_dataset/london_cut_0014/raw_images' 45 | imgs = [load_rgb(p) for p in sorted(glob(os.path.join(data_root, '*.png')))] 46 | imgs = [torch.tensor(im, dtype=torch.float) for im in imgs] 47 | 48 | latents_init = torch.zeros([len(imgs), n_latent_dim], dtype=torch.float, device=device) 49 | # latents_init = torch.linspace(0, 1, len(imgs), dtype=torch.float, device=device) 50 | # latents_init = torch.zeros([len(imgs), n_latent_dim], dtype=torch.float, device=device).uniform_(-0.1, 0.1) 51 | latents = nn.Embedding(len(imgs), n_latent_dim, dtype=torch.float, device=device, _weight=latents_init.clone()) 52 | 53 | optimer = optim.Adam(list(latents.parameters())+list(m.parameters()), lr=lr, eps=1.0e-15, betas=(0.9, 0.99)) 54 | if use_scheduler: 55 | scheduler = get_scheduler( 56 | ConfigDict(type='exponential', min_factor=min_factor, warmup_steps=500, num_iters=num_iters, ), 57 | optimizer=optimer) 58 | 59 | logging_loss = [] 60 | 61 | with tqdm(range(num_iters)) as pbar: 62 | for _ in pbar: 63 | im_ind = np.random.randint(0, len(imgs)) 64 | im = imgs[im_ind] 65 | H, W, _ = im.shape 66 | xy = torch.rand([num_pixels, 2], dtype=torch.float, device=device) 67 | wh = (xy * xy.new_tensor([W,H])).long() 68 | wh.clamp_(wh.new_tensor([0]), wh.new_tensor([W-1, H-1])) 69 | 70 | z = latents(torch.tensor([[im_ind]], dtype=torch.long, device=device)) 71 | rgb_pred = m.forward(xy.unsqueeze(0), z=z/2+0.5)['output'] 72 | rgb_gt = im[wh[:, 1], wh[:, 0]].to(device).unsqueeze(0) 73 | 74 | loss = huber_loss(rgb_pred, rgb_gt) 75 | # latent close by 76 | loss += wreg * latents.weight.diff(dim=0).norm(dim=-1).mean() 77 | 78 | optimer.zero_grad() 79 | loss.backward() 80 | optimer.step() 81 | if use_scheduler: 82 | scheduler.step() 83 | 84 | logging_loss.append(loss.item()) 85 | pbar.set_postfix(lr=optimer.param_groups[0]['lr'], loss=loss.item()) 86 | 87 | @torch.no_grad() 88 | def pred_im(z: torch.Tensor, HW: Tuple[int,int]): 89 | z = z.view(-1, n_latent_dim) 90 | B = z.size(0) 91 | H, W = HW 92 | i, j = torch.meshgrid(torch.linspace(0, W-1, W, device=device), torch.linspace(0, H-1, H, device=device), indexing='xy') 93 | wh = torch.stack([i+0.5, j+0.5], dim=-1) 94 | xy = wh / wh.new_tensor([W,H]) 95 | xy = xy.tile(B,1,1,1) # BHW2 96 | im = m.forward(xy, z=z/2+0.5)['output'] # BHW3 97 | return im 98 | 99 | z_test = latents(torch.tensor([[0]], dtype=torch.long, device=device)) 100 | im_test = pred_im(z_test, (800, 1200)) 101 | 102 | plt.plot(logging_loss) 103 | plt.show() 104 | 105 | plt.imshow(im_test[0].data.float().cpu().numpy()) 106 | plt.show() 107 | 108 | latents_new = latents.weight.data.clone() 109 | z_vid = torch.linspace(-latents_new.min().item(),-latents_new.max().item(),24,dtype=torch.float, device=device).unsqueeze(-1).tile(1,n_latent_dim) 110 | vid = pred_im(z_vid, (800, 1200)).data.float().cpu().numpy() 111 | vid = (vid*255.).clip(0,255).astype(np.uint8) 112 | vid_pth = 'dev_test/test_permuto_vid.mp4' 113 | imageio.mimwrite(vid_pth, vid) 114 | print(f"Video saved to {vid_pth}") 115 | -------------------------------------------------------------------------------- /docs/methods/neuralsim.md: -------------------------------------------------------------------------------- 1 | # Neuralsim 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /docs/methods/neus_in_10_minutes.md: -------------------------------------------------------------------------------- 1 | # [neuralsim] NeuS in 10 minutes 2 | 3 | [website](https://lingjie0206.github.io/papers/NeuS/) | [arxiv](https://arxiv.org/abs/2106.10689) | [official_repo](https://github.com/Totoro97/NeuS) | 4 | 5 | An **unofficial** and improved implementation of "NeuS: Learning Neural Implicit Surfaces by Volume Rendering for Multi-view Reconstruction". 6 | 7 | ```bibtex 8 | @inproceedings{wang2021neus, 9 | title={NeuS: Learning Neural Implicit Surfaces by Volume Rendering for Multi-view Reconstruction}, 10 | author={Wang, Peng and Liu, Lingjie and Liu, Yuan and Theobalt, Christian and Komura, Taku and Wang, Wenping}, 11 | booktitle={Proc. Advances in Neural Information Processing Systems (NeurIPS)}, 12 | volume={34}, 13 | pages={27171--27183}, 14 | year={2021} 15 | } 16 | ``` 17 | 18 | ## [NVS Demo] Rendered & Depth & Surface normals 19 | 20 | https://github.com/PJLab-ADG/neuralsim/assets/25529198/ce6ec6fc-2d0e-4c2b-9d91-d1b992d13ff4 21 | 22 | https://github.com/PJLab-ADG/neuralsim/assets/25529198/32d6fe6f-39a1-403d-8b12-16d26e092375 23 | 24 | https://github.com/PJLab-ADG/neuralsim/assets/25529198/29fff4d9-c70d-4097-8f12-0bc307d339f3 25 | 26 | 27 | ## Highlights (demo coming soon!) 28 | 29 | - Stable training within 10 minutes without necessarily needing mask 30 | - Worried about your camera pose accuracy ? We can refine them ! 31 | - Worried about your footage quality & consistency ? We have in the wild image embeddings ! 32 | - Worried about geometric distortions like depressions or bulges ? We opt to use monocular normal priors ! 33 | - Object-centric, indoor or outdoors ? We can cover them all ! 34 | 35 | ## Object-centric datasets 36 | 37 | ### Requirements 38 | 39 | - <10 mins training time on single RTX3090 40 | - 6 GiB GPU Mem 41 | 42 | ### Major settings 43 | 44 | #### > Without mask 45 | 46 | | Dataset | Config file | 47 | | ------------------------------------------------------------ | ------------------------------------------------------------ | 48 | | COLMAP dataset + Apperance Embeddings | [lotd_neus.colmap.230826.yaml](../../code_single/configs/exps/lotd_neus.colmap.230826.yaml) | 49 | | COLMAP dataset + Pose refinement + Apperance Embeddings | [lotd_neus.colmap_refine.230826.yaml](../../code_single/configs/exps/lotd_neus.colmap_refine.230826.yaml) | 50 | | [BlendedMVS dataset preparation](../../dataio/bmvs/README.md) | [lotd_neus.bmvs.230814.yaml](../../code_single/configs/object_centric/lotd_neus.bmvs.230814.yaml) | 51 | | [NeuS/DTU dataset preparation](../../dataio/dtu/README.md) | | 52 | 53 | #### > With mask (WIP) 54 | 55 | ### Instructions 56 | 57 | For detailed instructions, please refer to the [general guide](../../code_single/README.md#general-usage) section in `code_single`. 58 | 59 | ## Indoor datasets 60 | 61 | Can be viewed as an **unofficial** and improved implementation of "MonoSDF: Exploring Monocular Geometric Cues for Neural Implicit Surface Reconstruction". 62 | 63 | [website](https://niujinshuchong.github.io/monosdf/) | [arxiv](https://arxiv.org/abs/2206.00665) | [offcial_repo](https://github.com/autonomousvision/monosdf) | :warning: Unofficial implementation :warning: 64 | 65 | ```bibtex 66 | @inproceedings{Yu2022MonoSDF, 67 | author = {Yu, Zehao and Peng, Songyou and Niemeyer, Michael and Sattler, Torsten and Geiger, Andreas}, 68 | title = {MonoSDF: Exploring Monocular Geometric Cues for Neural Implicit Surface Reconstruction}, 69 | booktitle={Proc. Advances in Neural Information Processing Systems (NeurIPS)}, 70 | year = {2022}, 71 | } 72 | ``` 73 | 74 | ### Requirements 75 | 76 | - <10 mins training time on single RTX3090 77 | - 6 GiB GPU Mem 78 | 79 | ### Dataset preparation 80 | 81 | Follow [this link](https://github.com/autonomousvision/monosdf#dataset) to download the MonoSDF's preprocessed data of [Replica](https://github.com/facebookresearch/Replica-Dataset) / [scannet](http://www.scan-net.org/) indoor datasets. 82 | 83 | ### Major settings 84 | 85 | | Settings / Dataset | Config file | 86 | | ------------------------------------------------------------ | ------------------------------------------------------------ | 87 | | Replica dataset (processed by [MonoSDF](https://github.com/autonomousvision/monosdf)) | [lotd_neus.replica.230814.yaml](../../code_single/configs/indoor/lotd_neus.replica.230814.yaml) | 88 | | Scan net dataset (processed by [MonoSDF](https://github.com/autonomousvision/monosdf)) | WIP | 89 | 90 | ### Instructions 91 | 92 | For detailed instructions, please refer to the [general guide](../../code_single/README.md#general-usage) section in `code_single`. 93 | -------------------------------------------------------------------------------- /docs/methods/ngp_lidar.md: -------------------------------------------------------------------------------- 1 | # [neuralsim] InstantNGP + UrbanNeRF 2 | 3 | [Instant-NGP](https://github.com/NVlabs/instant-ngp) | [UrbanNeRF](https://urban-radiance-fields.github.io/) 4 | 5 | :warning: Unofficial implementation :warning: 6 | 7 | An unofficial implementation (combination) of "Instant Neural Graphics Primitives with a Multiresolution Hash Encoding" and "Urban Radiance Fields" 8 | 9 | ```bibtex 10 | @article{muller2022instantngp, 11 | title={Instant neural graphics primitives with a multiresolution hash encoding}, 12 | author={M{\"u}ller, Thomas and Evans, Alex and Schied, Christoph and Keller, Alexander}, 13 | journal={ACM Transactions on Graphics (ToG)}, 14 | volume={41}, 15 | number={4}, 16 | pages={1--15}, 17 | year={2022} 18 | } 19 | ``` 20 | 21 | ```bibtex 22 | @inproceedings{rematas2022urban, 23 | title={Urban radiance fields}, 24 | author={Rematas, Konstantinos and Liu, Andrew and Srinivasan, Pratul P and Barron, Jonathan T and Tagliasacchi, Andrea and Funkhouser, Thomas and Ferrari, Vittorio}, 25 | booktitle={Proc. {IEEE/CVF} Conference on Computer Vision and Pattern Recognition (CVPR)}, 26 | pages={12932--12942}, 27 | year={2022} 28 | } 29 | ``` 30 | 31 | ## Usage 32 | 33 | ### Requirements 34 | 35 | - ~45 mins training time on single RTX3090 36 | 37 | - ~11 GiB GPU Mem 38 | - \>20 GiB CPU Mem (Caching data to speed up) 39 | 40 | ### Dataset preparation 41 | 42 | - Waymo Open Dataset - Perception 43 | 44 | - [README](../../dataio/autonomous_driving/waymo/README.md) 45 | 46 | - split file: [waymo_static_32.lst](../../dataio/autonomous_driving/waymo/waymo_static_32.lst) 47 | 48 | 49 | ### Major settings 50 | 51 | | Settings | Config file | 52 | | ------------------------------------ | ------------------------------------------------------------ | 53 | | Multi-view reconstruction with LiDAR | [ngp_withlidar.230814.yaml](../../code_single/configs/waymo/ngp_withlidar.230814.yaml) | 54 | 55 | ## Instructions 56 | 57 | For detailed instructions, please refer to the [general guide](../../code_single/README.md#general-usage) section in `code_single`. 58 | -------------------------------------------------------------------------------- /docs/methods/street_gaussian.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/docs/methods/street_gaussian.md -------------------------------------------------------------------------------- /docs/methods/streetsurf/cuboid_space_sdf.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/docs/methods/streetsurf/cuboid_space_sdf.png -------------------------------------------------------------------------------- /docs/methods/streetsurf/raymarching_with_eikonal_on_occ.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/docs/methods/streetsurf/raymarching_with_eikonal_on_occ.png -------------------------------------------------------------------------------- /docs/methods/streetsurf/raymarching_without_eikonal_on_occ.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/docs/methods/streetsurf/raymarching_without_eikonal_on_occ.png -------------------------------------------------------------------------------- /docs/methods/streetsurf/sdf_nablas_norm_relu.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/docs/methods/streetsurf/sdf_nablas_norm_relu.png -------------------------------------------------------------------------------- /docs/methods/streetsurf/sdf_slice_with_eikonal_on_occ.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/docs/methods/streetsurf/sdf_slice_with_eikonal_on_occ.png -------------------------------------------------------------------------------- /docs/methods/streetsurf/sdf_slice_without_eikonal_on_occ.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/docs/methods/streetsurf/sdf_slice_without_eikonal_on_occ.png -------------------------------------------------------------------------------- /docs/methods/streetsurf_unisim.md: -------------------------------------------------------------------------------- 1 | # (WIP) 2 | 3 | -------------------------------------------------------------------------------- /docs/tutorials/appearace_factorization.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/docs/tutorials/appearace_factorization.md -------------------------------------------------------------------------------- /docs/tutorials/semantic_field.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/docs/tutorials/semantic_field.md -------------------------------------------------------------------------------- /media/github_assets.txt: -------------------------------------------------------------------------------- 1 | # A storage of the automatically uploaded assets in github 2 | 3 | # 3 of 360v2 dataset 4 | https://github.com/PJLab-ADG/neuralsim/assets/25529198/ce6ec6fc-2d0e-4c2b-9d91-d1b992d13ff4 5 | https://github.com/PJLab-ADG/neuralsim/assets/25529198/32d6fe6f-39a1-403d-8b12-16d26e092375 6 | https://github.com/PJLab-ADG/neuralsim/assets/25529198/29fff4d9-c70d-4097-8f12-0bc307d339f3 7 | 8 | # gundam 9 | teaser_training_bmvs_gundam 10 | 11 | # village house 12 | teaser_training_bmvs_village_house 13 | 14 | # replica scan1 (new) 15 | 16 | 17 | # teaser_categorical.mp4 18 | 19 | 20 | # teaser seg100613 21 | 22 | 23 | # teaser sensor simulation 24 | 25 | 26 | # teaser mix 27 | 28 | 29 | # teaser manipulate 30 | 31 | 32 | # teaser multi-verse 33 | 34 | 35 | # teaser_seg767010_style 36 | 37 | 38 | # img seg100613_ds=4_withmesh 39 | seg100613_ds=4_withmesh 40 | 41 | # [visualization] video mesh & cam & lidar 42 | https://github.com/PJLab-ADG/neuralsim/assets/25529198/42080c6e-6bde-4ecf-8d42-55b5fddf7be9 43 | 44 | # [visualization] video lidar pcl 45 | https://github.com/PJLab-ADG/neuralsim/assets/25529198/b502d393-dcf1-4a87-ba0f-bfa7eae26eee 46 | 47 | # [streetsurf] waymo static_32 48 | streetsurf_waymo_static_32 49 | 50 | # [streetsurf] lidar_only all waymo 51 | https://github.com/PJLab-ADG/neuralsim/assets/25529198/583800ff-42e0-4b9d-ace7-dae43229aa35 52 | 53 | # [streetsurf] An example of rendering with mesh 54 | https://github.com/PJLab-ADG/neuralsim/assets/25529198/7c1767c6-6eae-41c4-a60d-2eb74ff4c4c9 55 | 56 | # [neuralsim] render video seg965324 57 | https://github.com/PJLab-ADG/neuralsim/assets/25529198/dbd7854e-a9db-48be-a23c-61e3592f8faf 58 | 59 | # [neuralsim] lens dirt seg1009661 60 | https://github.com/PJLab-ADG/neuralsim/assets/25529198/62d10c4a-e8c1-4c22-be79-ca5622d6e0f1 61 | -> screeshot 62 | ![lens_dirt_1009661](https://github.com/PJLab-ADG/neuralsim/assets/25529198/91cd9908-9a14-4766-8b5a-983d55ba671b) 63 | 64 | # [neuralsim] lens flare seg938501 65 | https://github.com/PJLab-ADG/neuralsim/assets/25529198/5e88aed0-d271-4d63-bfea-08e700cd8047 66 | -> screenshot 67 | ![lens_flare_9385013](https://github.com/PJLab-ADG/neuralsim/assets/25529198/c626b905-e50b-4977-830a-83f255a12cac) 68 | 69 | #---- Deprecated: stored in private `neuralsim_dev` repo. 70 | # gundam 71 | https://github.com/PJLab-ADG/neuralsim_dev/assets/25529198/6755fb78-4103-41c0-b25e-74e78f11ee00 72 | # gundam large text 73 | https://github.com/PJLab-ADG/neuralsim_dev/assets/25529198/e46a21db-f564-461e-906e-af0a9a77d7ad 74 | # statue 75 | https://github.com/PJLab-ADG/neuralsim_dev/assets/25529198/87ff0d2c-4cd1-416e-b125-e7b5fd72f221 76 | # village house 77 | https://github.com/PJLab-ADG/neuralsim_dev/assets/25529198/5a922e08-5a45-400a-aa53-b3c08e87b1a8 78 | # village house large text 79 | https://github.com/PJLab-ADG/neuralsim_dev/assets/25529198/a671cf1e-ece0-4ee3-9131-6417cfb8e77e 80 | # archway 81 | https://github.com/PJLab-ADG/neuralsim_dev/assets/25529198/0e335e09-31db-4e06-abad-34f3738088cc 82 | # replica scan1 83 | https://github.com/PJLab-ADG/neuralsim_dev/assets/25529198/f41d541d-c047-4aaf-8b4c-54c6e61862fc 84 | # replica scan1 (new) 85 | https://github.com/PJLab-ADG/neuralsim_dev/assets/25529198/2791e214-451f-4244-8594-18adbff62edf 86 | # teaser_categorical.mp4 87 | https://github.com/PJLab-ADG/neuralsim_dev/assets/25529198/f3d99126-d228-4964-97ad-c12563998df7 88 | # teaser seg100613 89 | https://github.com/PJLab-ADG/neuralsim_dev/assets/25529198/294bbdd3-e1e0-4ca7-8779-44e36f806f4a 90 | # teaser 767 mix 91 | https://github.com/PJLab-ADG/neuralsim_dev/assets/25529198/334d521d-1266-49a7-aa04-034f76a8c730 92 | # teaser manipulate 93 | https://github.com/PJLab-ADG/neuralsim_dev/assets/25529198/6e7aad96-a9c4-4618-833a-19b564edbfd5 94 | -------------------------------------------------------------------------------- /media/logo_blue.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/media/logo_blue.png -------------------------------------------------------------------------------- /media/multi_object_volume_render.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/media/multi_object_volume_render.png -------------------------------------------------------------------------------- /media/occ_grid_batched.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/media/occ_grid_batched.jpg -------------------------------------------------------------------------------- /media/occ_grid_batched_dynamic.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/media/occ_grid_batched_dynamic.jpg -------------------------------------------------------------------------------- /media/occ_grid_dynamic.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/media/occ_grid_dynamic.jpg -------------------------------------------------------------------------------- /media/scene_graph.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/media/scene_graph.png -------------------------------------------------------------------------------- /media/vis_frustum_culling.jpeg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/media/vis_frustum_culling.jpeg -------------------------------------------------------------------------------- /media/vis_scene_graph.jpeg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PJLab-ADG/neuralsim/faba099e0feb11ea0089490a5e87565e25bc4a2c/media/vis_scene_graph.jpeg -------------------------------------------------------------------------------- /set_env.sh: -------------------------------------------------------------------------------- 1 | # NOTE: Insert project directory into PYTHONPATH 2 | # Usage: source set_env.sh 3 | 4 | script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )" 5 | script_dir=$(realpath $script_dir) 6 | export PYTHONPATH="${script_dir}":$PYTHONPATH 7 | echo "Added $script_dir to PYTHONPATH" 8 | --------------------------------------------------------------------------------