[ { "path": ".gitignore", "content": "*.pyc\n.vscode\noutput\nbuild\ndiff_rasterization/diff_rast.egg-info\ndiff_rasterization/dist\ntensorboard_3d\nscreenshots\n*.egg-info\nexternal\nshell\n.DS_Store" }, { "path": ".gitmodules", "content": "[submodule \"submodules/simple-knn\"]\n\tpath = submodules/simple-knn\n\turl = https://gitlab.inria.fr/bkerbl/simple-knn.git\n[submodule \"submodules/hugs-rasterization\"]\n\tpath = submodules/hugs-rasterization\n\turl = https://github.com/hyzhou404/hugs-rasterization\n" }, { "path": "LICENSE.md", "content": "HUGS License \n=========================== \n\n**Zhejiang University** hold all the ownership rights on the *Software* named **HUGS**. \n\nThe *Software* is still being developed by the *Licensor*. \n\n*Licensor*'s goal is to allow the research community to use, test and evaluate \nthe *Software*. \n\n## 1. Definitions \n\n*Licensee* means any person or entity that uses the *Software* and distributes \nits *Work*. \n\n*Licensor* means the owners of the *Software*, i.e Zhejiang University\n\n*Software* means the original work of authorship made available under this \nLicense ie HUGS. \n\n*Work* means the *Software* and any additions to or derivative works of the \n*Software* that are made available under this License. \n\n\n## 2. Purpose \nThis license is intended to define the rights granted to the *Licensee* by \nLicensors under the *Software*. \n\n## 3. Rights granted \n\nFor the above reasons Licensors have decided to distribute the *Software*. \nLicensors grant non-exclusive rights to use the *Software* for research purposes \nto research users (both academic and industrial), free of charge, without right \nto sublicense.. The *Software* may be used \"non-commercially\", i.e., for research \nand/or evaluation purposes only. \n\nSubject to the terms and conditions of this License, you are granted a \nnon-exclusive, royalty-free, license to reproduce, prepare derivative works of, \npublicly display, publicly perform and distribute its *Work* and any resulting \nderivative works in any form. \n\n## 4. Limitations \n\n**4.1 Redistribution.** You may reproduce or distribute the *Work* only if (a) you do \nso under this License, (b) you include a complete copy of this License with \nyour distribution, and (c) you retain without modification any copyright, \npatent, trademark, or attribution notices that are present in the *Work*. \n\n**4.2 Derivative Works.** You may specify that additional or different terms apply \nto the use, reproduction, and distribution of your derivative works of the *Work* \n(\"Your Terms\") only if (a) Your Terms provide that the use limitation in \nSection 2 applies to your derivative works, and (b) you identify the specific \nderivative works that are subject to Your Terms. Notwithstanding Your Terms, \nthis License (including the redistribution requirements in Section 3.1) will \ncontinue to apply to the *Work* itself. \n\n**4.3** Any other use without of prior consent of Licensors is prohibited. Research \nusers explicitly acknowledge having received from Licensors all information \nallowing to appreciate the adequacy between of the *Software* and their needs and \nto undertake all necessary precautions for its execution and use. \n\n**4.4** The *Software* is provided both as a compiled library file and as source \ncode. In case of using the *Software* for a publication or other results obtained \nthrough the use of the *Software*, users are strongly encouraged to cite the \ncorresponding publications as explained in the documentation of the *Software*. \n\n## 5. Disclaimer \n\nTHE USER CANNOT USE, EXPLOIT OR DISTRIBUTE THE *SOFTWARE* FOR COMMERCIAL PURPOSES \nWITHOUT PRIOR AND EXPLICIT CONSENT OF LICENSORS. YOU MUST CONTACT Zhejiang University FOR ANY \nUNAUTHORIZED USE: yiyi.liao@zju.edu.cn. ANY SUCH ACTION WILL \nCONSTITUTE A FORGERY. THIS *SOFTWARE* IS PROVIDED \"AS IS\" WITHOUT ANY WARRANTIES \nOF ANY NATURE AND ANY EXPRESS OR IMPLIED WARRANTIES, WITH REGARDS TO COMMERCIAL \nUSE, PROFESSIONNAL USE, LEGAL OR NOT, OR OTHER, OR COMMERCIALISATION OR \nADAPTATION. UNLESS EXPLICITLY PROVIDED BY LAW, IN NO EVENT, SHALL Zhejiang University OR THE \nAUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR \nCONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE \nGOODS OR SERVICES, LOSS OF USE, DATA, OR PROFITS OR BUSINESS INTERRUPTION) \nHOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT \nLIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING FROM, OUT OF OR \nIN CONNECTION WITH THE *SOFTWARE* OR THE USE OR OTHER DEALINGS IN THE *SOFTWARE*. \n" }, { "path": "README.md", "content": "# HUGS: Holistic Urban 3D Scene Understanding via Gaussian Splatting\n\n[Hongyu Zhou](https://github.com/hyzhou404), [Jiahao Shao](https://jhaoshao.github.io/), Lu Xu, Dongfeng Bai, [Weichao Qiu](https://weichaoqiu.com/), Bingbing Liu, [Yue Wang](https://ywang-zju.github.io/), [Andreas Geiger](https://www.cvlibs.net/) , [Yiyi Liao](https://yiyiliao.github.io/)
\n\n| [Webpage](https://xdimlab.github.io/hugs_website/) | [Full Paper](https://openaccess.thecvf.com/content/CVPR2024/html/Zhou_HUGS_Holistic_Urban_3D_Scene_Understanding_via_Gaussian_Splatting_CVPR_2024_paper.html) | [Video](https://www.youtube.com/watch?v=DmPhL-8FeT4)\n\nThis repository contains the official authors implementation associated with the paper \"HUGS: Holistic Urban 3D Scene Understanding via Gaussian Splatting\", which can be found [here](https://xdimlab.github.io/hugs_website/). \n\n![image teaser](./assets/teaser.png)\n\nAbstract: *Holistic understanding of urban scenes based on RGB images is a challenging yet important problem. It encompasses understanding both the geometry and appearance to enable novel view synthesis, parsing semantic labels, and tracking moving objects. Despite considerable progress, existing approaches often focus on specific aspects of this task and require additional inputs such as LiDAR scans or manually annotated 3D bounding boxes. In this paper, we introduce a novel pipeline that utilizes 3D Gaussian Splatting for holistic urban scene understanding. Our main idea involves the joint optimization of geometry, appearance, semantics, and motion using a combination of static and dynamic 3D Gaussians, where moving object poses are regularized via physical constraints. Our approach offers the ability to render new viewpoints in real-time, yielding 2D and 3D semantic information with high accuracy, and reconstruct dynamic scenes, even in scenarios where 3D bounding box detection are highly noisy. Experimental results on KITTI, KITTI-360, and Virtual KITTI 2 demonstrate the effectiveness of our approach.*\n\n\n\n## Cloning the Repository\n\nThe repository contains submodules, thus please check it out with \n```shell\n# SSH\ngit clone git@github.com:hyzhou404/hugs.git --recursive\n```\nor\n```shell\n# HTTPS\ngit clone https://github.com/hyzhou404/hugs --recursive\n```\n\n\n\n## Prepare Enviroments\n\nCreate conda environment:\n\n```shell\nconda create -n hugs python=3.10 -y\n```\n\nPlease install [PyTorch](https://pytorch.org/), [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn), [pytorch3d](https://github.com/facebookresearch/pytorch3d/tree/main) and [flow-vis-torch](https://github.com/ChristophReich1996/Optical-Flow-Visualization-PyTorch) by following official instructions.\n\nInstall submodules by running:\n\n```shell\npip install submodules/simple-knn\npip install submodules/hugs-rasterization\n```\n\nInstall remaining packages by running:\n ```shell\n pip install -r requirements.txt\n ```\n\n\n## Data & Checkpoints Download\n\nwe have made available two sequences from KITTI as indicated in our paper. Furthermore, three sequences from KITTI-360 and one sequence from Waymo has also been provided.\n\nDownload checkpoints from [here](https://huggingface.co/datasets/hyzhou404/hugs_release).\n\n```python\nunzip ${sequence}.zip\n```\n\n\n\n## Rendering\n\nRender test views by running:\n\n```shell\npython render.py -m ${checkpoint_path} --data_type ${dataset_type} --iteration 30000 --affine \n```\n\nThe variable **dataset_type** is a string, and its value can be one of the following: **kitti**, **kitti360**, or **waymo**.\n\n\n## Evaluation\n\n```\npython metrics.py -m ${checkpoint_path}\n```\n\n## Training\nThis repository only includes the inference code of HUGS. The code for training will be released in future work.\n\n\n
\n
\n

BibTeX

\n 
@InProceedings{Zhou_2024_CVPR,\n    author    = {Zhou, Hongyu and Shao, Jiahao and Xu, Lu and Bai, Dongfeng and Qiu, Weichao and Liu, Bingbing and Wang, Yue and Geiger, Andreas and Liao, Yiyi},\n    title     = {HUGS: Holistic Urban 3D Scene Understanding via Gaussian Splatting},\n    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},\n    month     = {June},\n    year      = {2024},\n    pages     = {21336-21345}\n    }
\n
\n
\n" }, { "path": "arguments/__init__.py", "content": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# This software is free for non-commercial, research and evaluation use \n# under the terms of the LICENSE.md file.\n#\n# For inquiries contact george.drettakis@inria.fr\n#\n\nfrom argparse import ArgumentParser, Namespace\nimport sys\nimport os\n\nclass GroupParams:\n pass\n\nclass ParamGroup:\n def __init__(self, parser: ArgumentParser, name : str, fill_none = False):\n group = parser.add_argument_group(name)\n for key, value in vars(self).items():\n shorthand = False\n if key.startswith(\"_\"):\n shorthand = True\n key = key[1:]\n t = type(value)\n value = value if not fill_none else None \n if shorthand:\n if t == bool:\n group.add_argument(\"--\" + key, (\"-\" + key[0:1]), default=value, action=\"store_true\")\n else:\n group.add_argument(\"--\" + key, (\"-\" + key[0:1]), default=value, type=t)\n else:\n if t == bool:\n group.add_argument(\"--\" + key, default=value, action=\"store_true\")\n else:\n group.add_argument(\"--\" + key, default=value, type=t)\n\n def extract(self, args):\n group = GroupParams()\n for arg in vars(args).items():\n if arg[0] in vars(self) or (\"_\" + arg[0]) in vars(self):\n setattr(group, arg[0], arg[1])\n return group\n\nclass ModelParams(ParamGroup): \n def __init__(self, parser, sentinel=False):\n self.sh_degree = 3\n self._source_path = \"\"\n self._model_path = \"\"\n self._images = \"images\"\n self._resolution = -1\n self._white_background = False\n self.data_device = \"cpu\"\n self.eval = False\n super().__init__(parser, \"Loading Parameters\", sentinel)\n\n def extract(self, args):\n g = super().extract(args)\n g.source_path = os.path.abspath(g.source_path)\n return g\n\nclass PipelineParams(ParamGroup):\n def __init__(self, parser):\n self.convert_SHs_python = False\n self.compute_cov3D_python = False\n self.debug = False\n super().__init__(parser, \"Pipeline Parameters\")\n\nclass OptimizationParams(ParamGroup):\n def __init__(self, parser):\n self.iterations = 30_000\n self.position_lr_init = 0.00016\n self.position_lr_final = 0.0000016\n self.position_lr_delay_mult = 0.01\n self.position_lr_max_steps = 30_000\n self.feature_lr = 0.0025\n self.opacity_lr = 0.05\n self.scaling_lr = 0.001\n self.rotation_lr = 0.001\n self.percent_dense = 0.001\n self.lambda_dssim = 0.2\n self.densification_interval = 100\n self.opacity_reset_interval = 3000\n self.densify_from_iter = 500\n self.densify_until_iter = 15_000\n self.densify_grad_threshold = 0.0002\n super().__init__(parser, \"Optimization Parameters\")\n\ndef get_combined_args(parser : ArgumentParser):\n cmdlne_string = sys.argv[1:]\n cfgfile_string = \"Namespace()\"\n args_cmdline = parser.parse_args(cmdlne_string)\n\n try:\n cfgfilepath = os.path.join(args_cmdline.model_path, \"cfg_args\")\n print(\"Looking for config file in\", cfgfilepath)\n with open(cfgfilepath) as cfg_file:\n print(\"Config file found: {}\".format(cfgfilepath))\n cfgfile_string = cfg_file.read()\n except TypeError:\n print(\"Config file not found at\")\n pass\n args_cfgfile = eval(cfgfile_string)\n\n merged_dict = vars(args_cfgfile).copy()\n for k,v in vars(args_cmdline).items():\n if v != None:\n merged_dict[k] = v\n return Namespace(**merged_dict)\n" }, { "path": "environment.yml", "content": "name: gaussian_splatting\nchannels:\n - pytorch\n - conda-forge\n - defaults\ndependencies:\n - cudatoolkit=11.6\n - plyfile=0.8.1\n - python=3.7.13\n - pip=22.3.1\n - pytorch=1.12.1\n - torchaudio=0.12.1\n - torchvision=0.13.1\n - tqdm\n - pip:\n - submodules/diff-gaussian-rasterization\n - submodules/simple-knn" }, { "path": "gaussian_renderer/__init__.py", "content": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# This software is free for non-commercial, research and evaluation use \n# under the terms of the LICENSE.md file.\n#\n# For inquiries contact george.drettakis@inria.fr\n#\n\nimport torch\nimport math\nfrom diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer\nfrom scene.gaussian_model import GaussianModel\nfrom utils.sh_utils import eval_sh, RGB2SH\nfrom pytorch3d.transforms import quaternion_to_matrix, matrix_to_quaternion\n\ndef euler2matrix(yaw):\n cos = torch.cos(-yaw)\n sin = torch.sin(-yaw)\n rot = torch.eye(3).float().cuda()\n rot[0,0] = cos\n rot[0,2] = sin\n rot[2,0] = -sin\n rot[2,2] = cos\n return rot\n\ndef cat_bgfg(bg, fg, only_dynamic=False, only_xyz=False):\n if only_xyz:\n bg_feats = [bg.get_xyz]\n else:\n bg_feats = [bg.get_xyz, bg.get_opacity, bg.get_scaling, bg.get_rotation, bg.get_features, bg.get_3D_features]\n \n output = []\n for fg_feat, bg_feat in zip(fg, bg_feats):\n if fg_feat is None:\n output.append(bg_feat)\n elif only_dynamic:\n output.append(fg_feat)\n else:\n output.append(torch.cat((bg_feat, fg_feat), dim=0))\n \n return output\n\n\ndef cat_all_fg(all_fg, next_fg):\n output = []\n for feat, next_feat in zip(all_fg, next_fg):\n if feat is None:\n feat = next_feat\n else:\n feat = torch.cat((feat, next_feat), dim=0)\n output.append(feat)\n return output\n\n\ndef proj_uv(xyz, cam):\n device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n intr = torch.as_tensor(cam.K[:3, :3]).float().to(device) # (3, 3)\n w2c = torch.tensor(cam.w2c).float().to(device)[:3, :] # (3, 4)\n\n c_xyz = (w2c[:3, :3] @ xyz.T).T + w2c[:3, 3]\n i_xyz = (intr @ c_xyz.mT).mT # (N, 3)\n uv = i_xyz[:, :2] / i_xyz[:, -1:].clip(1e-3) # (N, 2)\n return uv\n\n\ndef unicycle_b2w(timestamp, model):\n # model = unicycle_models[track_id]['model']\n pred = model(timestamp)\n if pred is None:\n return None\n pred_a, pred_b, pred_v, pred_phi, pred_h = pred\n # r = euler_angles_to_matrix(torch.tensor([0, pred_phi-torch.pi, 0]), 'XYZ')\n rt = torch.eye(4).float().cuda()\n rt[:3,:3] = euler2matrix(pred_phi)\n rt[1, 3], rt[0, 3], rt[2, 3] = pred_h, pred_a, pred_b\n return rt\n\ndef render(viewpoint_camera, prev_viewpoint_camera, pc : GaussianModel, dynamic_gaussians : dict, \n unicycles : dict, pipe, bg_color : torch.Tensor, \n render_optical=False, scaling_modifier = 1.0, only_dynamic=False):\n \"\"\"\n Render the scene. \n \n Background tensor (bg_color) must be on GPU!\n \"\"\"\n timestamp = viewpoint_camera.timestamp\n\n all_fg = [None, None, None, None, None, None]\n prev_all_fg = [None]\n\n if len(unicycles) == 0:\n track_dict = viewpoint_camera.dynamics\n if prev_viewpoint_camera is not None:\n prev_track_dict = prev_viewpoint_camera.dynamics\n else:\n track_dict, prev_track_dict = {}, {}\n for track_id, uni_model in unicycles.items():\n B2W = unicycle_b2w(timestamp, uni_model['model'])\n track_dict[track_id] = B2W\n if prev_viewpoint_camera is not None:\n prev_B2W = unicycle_b2w(prev_viewpoint_camera.timestamp, uni_model['model'])\n prev_track_dict[track_id] = prev_B2W\n\n for track_id, B2W in track_dict.items():\n w_dxyz = (B2W[:3, :3] @ dynamic_gaussians[track_id].get_xyz.T).T + B2W[:3, 3]\n drot = quaternion_to_matrix(dynamic_gaussians[track_id].get_rotation)\n w_drot = matrix_to_quaternion(B2W[:3, :3] @ drot)\n next_fg = [w_dxyz, \n dynamic_gaussians[track_id].get_opacity, \n dynamic_gaussians[track_id].get_scaling, \n w_drot,\n dynamic_gaussians[track_id].get_features,\n dynamic_gaussians[track_id].get_3D_features]\n # next_fg = get_next_fg(dynamic_gaussians[track_id], B2W)\n # w_dxyz = next_fg[0]\n all_fg = cat_all_fg(all_fg, next_fg)\n\n if render_optical and prev_viewpoint_camera is not None:\n if track_id in prev_track_dict:\n prev_B2W = prev_track_dict[track_id]\n prev_w_dxyz = torch.mm(prev_B2W[:3, :3], dynamic_gaussians[track_id].get_xyz.T).T + prev_B2W[:3, 3]\n prev_all_fg = cat_all_fg(prev_all_fg, [prev_w_dxyz])\n else:\n prev_all_fg = cat_all_fg(prev_all_fg, [w_dxyz])\n \n xyz, opacity, scales, rotations, shs, feats3D = cat_bgfg(pc, all_fg)\n if render_optical and prev_viewpoint_camera is not None:\n prev_xyz = cat_bgfg(pc, prev_all_fg, only_xyz=True)[0]\n uv = proj_uv(xyz, viewpoint_camera)\n prev_uv = proj_uv(prev_xyz, prev_viewpoint_camera)\n delta_uv = uv - prev_uv\n delta_uv = torch.cat([delta_uv, torch.ones_like(delta_uv[:, :1], device=delta_uv.device)], dim=-1)\n else:\n delta_uv = torch.zeros_like(xyz)\n\n # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means\n screenspace_points = torch.zeros_like(xyz, dtype=xyz.dtype, requires_grad=True, device=\"cuda\") + 0\n try:\n screenspace_points.retain_grad()\n except:\n pass\n\n # Set up rasterization configuration\n tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)\n tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)\n\n if pc.affine:\n cam_xyz, cam_dir = viewpoint_camera.c2w[:3, 3].cuda(), viewpoint_camera.c2w[:3, 2].cuda()\n o_enc = pc.pos_enc(cam_xyz[None, :] / 60)\n d_enc = pc.dir_enc(cam_dir[None, :])\n appearance = pc.appearance_model(torch.cat([o_enc, d_enc], dim=1)) * 1e-1\n affine_weight, affine_bias = appearance[:, :9].view(3, 3), appearance[:, -3:]\n affine_weight = affine_weight + torch.eye(3, device=appearance.device)\n\n # bg_img = pc.sky_model(enc).view(*rays_d.shape).permute(2, 0, 1).float()\n\n raster_settings = GaussianRasterizationSettings(\n image_height=int(viewpoint_camera.image_height),\n image_width=int(viewpoint_camera.image_width),\n tanfovx=tanfovx,\n tanfovy=tanfovy,\n bg=bg_color,\n scale_modifier=scaling_modifier,\n viewmatrix=viewpoint_camera.world_view_transform,\n projmatrix=viewpoint_camera.full_proj_transform,\n sh_degree=pc.active_sh_degree,\n campos=viewpoint_camera.camera_center,\n prefiltered=False,\n debug=pipe.debug\n )\n\n rasterizer = GaussianRasterizer(raster_settings=raster_settings)\n \n means3D = xyz\n means2D = screenspace_points\n\n cov3D_precomp = None\n colors_precomp = None\n\n # Rasterize visible Gaussians to image, obtain their radii (on screen). \n rendered_image, radii, feats, depth, flow = rasterizer(\n means3D = means3D,\n means2D = means2D,\n shs = shs,\n colors_precomp = colors_precomp,\n opacities = opacity,\n scales = scales,\n rotations = rotations,\n cov3D_precomp = cov3D_precomp,\n feats3D = feats3D,\n delta = delta_uv)\n \n if pc.affine:\n colors = rendered_image.view(3, -1).permute(1, 0) # (H*W, 3)\n refined_image = (colors @ affine_weight + affine_bias).clip(0, 1).permute(1, 0).view(*rendered_image.shape)\n else:\n refined_image = rendered_image\n\n # Those Gaussians that were frustum culled or had a radius of 0 were not visible.\n # They will be excluded from value updates used in the splitting criteria.\n return {\"render\": refined_image,\n \"feats\": feats,\n \"depth\": depth,\n \"opticalflow\": flow,\n \"viewspace_points\": screenspace_points,\n \"visibility_filter\" : radii > 0,\n \"radii\": radii}\n" }, { "path": "lpipsPyTorch/__init__.py", "content": "import torch\n\nfrom .modules.lpips import LPIPS\n\n\ndef lpips(x: torch.Tensor,\n y: torch.Tensor,\n net_type: str = 'alex',\n version: str = '0.1'):\n r\"\"\"Function that measures\n Learned Perceptual Image Patch Similarity (LPIPS).\n\n Arguments:\n x, y (torch.Tensor): the input tensors to compare.\n net_type (str): the network type to compare the features: \n 'alex' | 'squeeze' | 'vgg'. Default: 'alex'.\n version (str): the version of LPIPS. Default: 0.1.\n \"\"\"\n device = x.device\n criterion = LPIPS(net_type, version).to(device)\n return criterion(x, y)\n" }, { "path": "lpipsPyTorch/modules/lpips.py", "content": "import torch\nimport torch.nn as nn\n\nfrom .networks import get_network, LinLayers\nfrom .utils import get_state_dict\n\n\nclass LPIPS(nn.Module):\n r\"\"\"Creates a criterion that measures\n Learned Perceptual Image Patch Similarity (LPIPS).\n\n Arguments:\n net_type (str): the network type to compare the features: \n 'alex' | 'squeeze' | 'vgg'. Default: 'alex'.\n version (str): the version of LPIPS. Default: 0.1.\n \"\"\"\n def __init__(self, net_type: str = 'alex', version: str = '0.1'):\n\n assert version in ['0.1'], 'v0.1 is only supported now'\n\n super(LPIPS, self).__init__()\n\n # pretrained network\n self.net = get_network(net_type)\n\n # linear layers\n self.lin = LinLayers(self.net.n_channels_list)\n self.lin.load_state_dict(get_state_dict(net_type, version))\n\n def forward(self, x: torch.Tensor, y: torch.Tensor):\n feat_x, feat_y = self.net(x), self.net(y)\n\n diff = [(fx - fy) ** 2 for fx, fy in zip(feat_x, feat_y)]\n res = [l(d).mean((2, 3), True) for d, l in zip(diff, self.lin)]\n\n return torch.sum(torch.cat(res, 0), 0, True)\n" }, { "path": "lpipsPyTorch/modules/networks.py", "content": "from typing import Sequence\n\nfrom itertools import chain\n\nimport torch\nimport torch.nn as nn\nfrom torchvision import models\n\nfrom .utils import normalize_activation\n\n\ndef get_network(net_type: str):\n if net_type == 'alex':\n return AlexNet()\n elif net_type == 'squeeze':\n return SqueezeNet()\n elif net_type == 'vgg':\n return VGG16()\n else:\n raise NotImplementedError('choose net_type from [alex, squeeze, vgg].')\n\n\nclass LinLayers(nn.ModuleList):\n def __init__(self, n_channels_list: Sequence[int]):\n super(LinLayers, self).__init__([\n nn.Sequential(\n nn.Identity(),\n nn.Conv2d(nc, 1, 1, 1, 0, bias=False)\n ) for nc in n_channels_list\n ])\n\n for param in self.parameters():\n param.requires_grad = False\n\n\nclass BaseNet(nn.Module):\n def __init__(self):\n super(BaseNet, self).__init__()\n\n # register buffer\n self.register_buffer(\n 'mean', torch.Tensor([-.030, -.088, -.188])[None, :, None, None])\n self.register_buffer(\n 'std', torch.Tensor([.458, .448, .450])[None, :, None, None])\n\n def set_requires_grad(self, state: bool):\n for param in chain(self.parameters(), self.buffers()):\n param.requires_grad = state\n\n def z_score(self, x: torch.Tensor):\n return (x - self.mean) / self.std\n\n def forward(self, x: torch.Tensor):\n x = self.z_score(x)\n\n output = []\n for i, (_, layer) in enumerate(self.layers._modules.items(), 1):\n x = layer(x)\n if i in self.target_layers:\n output.append(normalize_activation(x))\n if len(output) == len(self.target_layers):\n break\n return output\n\n\nclass SqueezeNet(BaseNet):\n def __init__(self):\n super(SqueezeNet, self).__init__()\n\n self.layers = models.squeezenet1_1(True).features\n self.target_layers = [2, 5, 8, 10, 11, 12, 13]\n self.n_channels_list = [64, 128, 256, 384, 384, 512, 512]\n\n self.set_requires_grad(False)\n\n\nclass AlexNet(BaseNet):\n def __init__(self):\n super(AlexNet, self).__init__()\n\n self.layers = models.alexnet(True).features\n self.target_layers = [2, 5, 8, 10, 12]\n self.n_channels_list = [64, 192, 384, 256, 256]\n\n self.set_requires_grad(False)\n\n\nclass VGG16(BaseNet):\n def __init__(self):\n super(VGG16, self).__init__()\n\n self.layers = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1).features\n self.target_layers = [4, 9, 16, 23, 30]\n self.n_channels_list = [64, 128, 256, 512, 512]\n\n self.set_requires_grad(False)\n" }, { "path": "lpipsPyTorch/modules/utils.py", "content": "from collections import OrderedDict\n\nimport torch\n\n\ndef normalize_activation(x, eps=1e-10):\n norm_factor = torch.sqrt(torch.sum(x ** 2, dim=1, keepdim=True))\n return x / (norm_factor + eps)\n\n\ndef get_state_dict(net_type: str = 'alex', version: str = '0.1'):\n # build url\n url = 'https://raw.githubusercontent.com/richzhang/PerceptualSimilarity/' \\\n + f'master/lpips/weights/v{version}/{net_type}.pth'\n\n # download\n old_state_dict = torch.hub.load_state_dict_from_url(\n url, progress=True,\n map_location=None if torch.cuda.is_available() else torch.device('cpu')\n )\n\n # rename keys\n new_state_dict = OrderedDict()\n for key, val in old_state_dict.items():\n new_key = key\n new_key = new_key.replace('lin', '')\n new_key = new_key.replace('model.', '')\n new_state_dict[new_key] = val\n\n return new_state_dict\n" }, { "path": "metrics.py", "content": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# This software is free for non-commercial, research and evaluation use \n# under the terms of the LICENSE.md file.\n#\n# For inquiries contact george.drettakis@inria.fr\n#\n\nfrom pathlib import Path\nimport os\nfrom PIL import Image\nimport torch\nimport torchvision.transforms.functional as tf\nfrom utils.loss_utils import ssim\nfrom lpipsPyTorch import lpips\nimport json\nfrom tqdm import tqdm\nfrom utils.image_utils import psnr\nfrom argparse import ArgumentParser\nfrom collections import OrderedDict\n\ndef readImages(renders_dir, gt_dir):\n renders = []\n gts = []\n image_names = []\n for fname in os.listdir(renders_dir):\n render = Image.open(renders_dir / fname)\n gt = Image.open(gt_dir / fname)\n renders.append(tf.to_tensor(render).unsqueeze(0)[:, :3, :, :].cuda())\n gts.append(tf.to_tensor(gt).unsqueeze(0)[:, :3, :, :].cuda())\n image_names.append(fname)\n return renders, gts, image_names\n\ndef evaluate(model_paths, write):\n # import ipdb; ipdb.set_trace()\n full_dict = {}\n per_view_dict = {}\n full_dict_polytopeonly = {}\n per_view_dict_polytopeonly = {}\n print(\"\")\n\n scene_dir = model_paths[0]\n\n print(\"Scene:\", scene_dir)\n\n for splits in ['test', 'train']:\n full_dict[splits] = {}\n per_view_dict[splits] = {}\n dir_path = Path(scene_dir) / splits\n for method in os.listdir(dir_path):\n print(\"Method:\", method)\n full_dict[splits][method] = {}\n per_view_dict[splits][method] = {}\n\n method_dir = dir_path / method\n gt_dir = method_dir/ \"gt\"\n renders_dir = method_dir / \"renders\"\n renders, gts, image_names = readImages(renders_dir, gt_dir)\n\n ssims = []\n psnrs = []\n lpipss = []\n\n for idx in tqdm(range(len(renders)), desc=\"Metric evaluation progress\"):\n ssims.append(ssim(renders[idx], gts[idx]))\n psnrs.append(psnr(renders[idx], gts[idx]))\n lpipss.append(lpips(renders[idx], gts[idx], net_type='alex'))\n\n print(\" SSIM : {:>12.7f}\".format(torch.tensor(ssims).mean(), \".5\"))\n print(\" PSNR : {:>12.7f}\".format(torch.tensor(psnrs).mean(), \".5\"))\n print(\" LPIPS: {:>12.7f}\".format(torch.tensor(lpipss).mean(), \".5\"))\n print(\"\")\n\n full_dict[splits][method].update({\"SSIM\": torch.tensor(ssims).mean().item(),\n \"PSNR\": torch.tensor(psnrs).mean().item(),\n \"LPIPS\": torch.tensor(lpipss).mean().item()})\n per_view_dict[splits][method].update({\n \"SSIM\": OrderedDict(sorted({name: ssim for ssim, name in zip(torch.tensor(ssims).tolist(), image_names)}.items())),\n \"PSNR\": OrderedDict(sorted({name: psnr for psnr, name in zip(torch.tensor(psnrs).tolist(), image_names)}.items())),\n \"LPIPS\": OrderedDict(sorted({name: lp for lp, name in zip(torch.tensor(lpipss).tolist(), image_names)}.items()))\n })\n if write:\n with open(scene_dir + \"/metric_results.json\", 'w') as fp:\n json.dump(full_dict, fp, indent=True)\n with open(scene_dir + \"/per_view.json\", 'w') as fp:\n json.dump(per_view_dict, fp, indent=True)\n\nif __name__ == \"__main__\":\n device = torch.device(\"cuda:0\")\n torch.cuda.set_device(device)\n\n # Set up command line argument parser\n parser = ArgumentParser(description=\"Training script parameters\")\n parser.add_argument('--model_paths', '-m', required=True, nargs=\"+\", type=str, default=[])\n parser.add_argument('--write', action='store_false', default=True)\n args = parser.parse_args()\n evaluate(args.model_paths, args.write)\n" }, { "path": "render.py", "content": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# This software is free for non-commercial, research and evaluation use \n# under the terms of the LICENSE.md file.\n#\n# For inquiries contact george.drettakis@inria.fr\n#\n\nimport torch\nfrom scene import Scene\nimport os\nfrom tqdm import tqdm\nfrom os import makedirs\nfrom gaussian_renderer import render\nimport torchvision\nfrom utils.general_utils import safe_state\nfrom argparse import ArgumentParser\nfrom arguments import ModelParams, PipelineParams, get_combined_args\nfrom gaussian_renderer import GaussianModel\nimport numpy as np\nfrom copy import deepcopy\nfrom torchmetrics.functional import structural_similarity_index_measure as ssim\nimport matplotlib.pyplot as plt\nfrom mpl_toolkits.axes_grid1 import make_axes_locatable\nfrom matplotlib import cm\nfrom utils.semantic_utils import colorize\nimport flow_vis_torch\nfrom utils.cmap import color_depth_map\nfrom imageio.v2 import imwrite\n\ndef to4x4(R, T):\n RT = np.eye(4,4)\n RT[:3, :3] = R\n RT[:3, 3] = T\n return RT\n\ndef apply_colormap(image, cmap=\"viridis\"):\n colormap = cm.get_cmap(cmap)\n colormap = torch.tensor(colormap.colors).to(image.device) # type: ignore\n image_long = (image * 255).long()\n image_long_min = torch.min(image_long)\n image_long_max = torch.max(image_long)\n assert image_long_min >= 0, f\"the min value is {image_long_min}\"\n assert image_long_max <= 255, f\"the max value is {image_long_max}\"\n return colormap[image_long[0, ...]].permute(2, 0, 1)\n\n\ndef apply_depth_colormap(depth, near_plane=None, far_plane=None, cmap=\"turbo\"):\n near_plane = near_plane or float(torch.min(depth))\n far_plane = far_plane or float(torch.max(depth))\n depth = (depth - near_plane) / (far_plane - near_plane + 1e-10)\n depth = torch.clip(depth, 0, 1)\n\n colored_image = apply_colormap(depth, cmap=cmap)\n return colored_image\n\n\ndef render_set(model_path, name, iteration, views, scene, pipeline, background, data_type):\n render_path = os.path.join(model_path, name, \"ours_{}\".format(iteration), \"renders\")\n semantic_path = os.path.join(model_path, name, \"ours_{}\".format(iteration), \"semantic\")\n optical_path = os.path.join(model_path, name, \"ours_{}\".format(iteration), \"optical\")\n gts_path = os.path.join(model_path, name, \"ours_{}\".format(iteration), \"gt\")\n error_path = os.path.join(model_path, name, \"ours_{}\".format(iteration), \"error_map\")\n depth_path = os.path.join(model_path, name, \"ours_{}\".format(iteration), \"depth\")\n\n makedirs(render_path, exist_ok=True)\n makedirs(semantic_path, exist_ok=True)\n makedirs(optical_path, exist_ok=True)\n makedirs(gts_path, exist_ok=True)\n makedirs(error_path, exist_ok=True)\n makedirs(depth_path, exist_ok=True)\n\n for idx, view in enumerate(tqdm(views, desc=\"Rendering progress\")):\n \n if data_type == 'kitti':\n gap = 2\n elif data_type == 'kitti360':\n gap = 4\n elif data_type == 'waymo':\n gap = 1\n elif data_type == 'nuscenes' or data_type == 'pandaset':\n gap = 6\n\n if idx - gap < 0:\n prev_view = None\n else:\n prev_view = views[idx-4]\n render_pkg = render(\n view, prev_view, scene.gaussians, scene.dynamic_gaussians, scene.unicycles, pipeline, background, True\n )\n rendering = render_pkg['render'].detach().cpu()\n semantic = render_pkg['feats'].detach().cpu()\n semantic = torch.argmax(semantic, dim=0)\n semantic_rgb = colorize(semantic.detach().cpu().numpy())\n depth = render_pkg['depth']\n color_depth = color_depth_map(depth[0].detach().cpu().numpy())\n color_depth[semantic == 10] = np.array([255.0, 255.0, 255.0])\n gt = view.original_image[0:3, :, :]\n \n # _, ssim_map = ssim(rendering[None, ...], gt[None, ...], return_full_image=True)\n # ssim_map = torch.mean(ssim_map[0], dim=0).clip(0, 1)[None, ...]\n # error_map = 1 - ssim_maps\n error_map = torch.mean((rendering - gt) ** 2, dim=0)[None, ...]\n\n fig = plt.figure(frameon=False)\n fig.set_size_inches(1.408, 0.376)\n ax = plt.Axes(fig, [0., 0., 1., 1.])\n ax.set_axis_off()\n fig.add_axes(ax)\n ax.imshow((error_map.detach().cpu().numpy().transpose(1,2,0)), cmap='jet')\n plt.savefig(os.path.join(error_path, view.image_name + \".png\"), dpi=1000)\n plt.close('all')\n\n torchvision.utils.save_image(rendering, os.path.join(render_path, view.image_name + \".png\"))\n torchvision.utils.save_image(gt, os.path.join(gts_path, view.image_name + \".png\"))\n semantic_rgb.save(os.path.join(semantic_path, view.image_name + \".png\"))\n imwrite(os.path.join(depth_path, view.image_name + \".png\"), color_depth)\n \n opticalflow = render_pkg[\"opticalflow\"]\n opticalflow = opticalflow.permute(1,2,0)\n opticalflow = opticalflow[..., :2]\n pytorch_optic_rgb = flow_vis_torch.flow_to_color(opticalflow.permute(2, 0, 1)) # (2, h, w)\n torchvision.utils.save_image(pytorch_optic_rgb.float(), os.path.join(optical_path, view.image_name + \".png\"), normalize=True)\n # torchvision.utils.save_image(error_map, os.path.join(error_path, '{0:05d}'.format(idx) + \".png\"))\n\ndef render_sets(dataset : ModelParams, iteration : int, pipeline : PipelineParams, \n skip_train : bool, skip_test : bool, data_type, affine, ignore_dynamic):\n with torch.no_grad():\n gaussians = GaussianModel(dataset.sh_degree, affine=affine)\n scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False, data_type=data_type, ignore_dynamic=ignore_dynamic)\n\n bg_color = [1,1,1] if dataset.white_background else [0, 0, 0]\n background = torch.tensor(bg_color, dtype=torch.float32, device=\"cuda\")\n\n if not skip_train:\n render_set(dataset.model_path, \"train\", scene.loaded_iter, scene.getTrainCameras(), scene, pipeline, background, data_type)\n\n if not skip_test:\n render_set(dataset.model_path, \"test\", scene.loaded_iter, scene.getTestCameras(), scene, pipeline, background, data_type)\n\n\nif __name__ == \"__main__\":\n # Set up command line argument parser\n parser = ArgumentParser(description=\"Testing script parameters\")\n model = ModelParams(parser, sentinel=True)\n pipeline = PipelineParams(parser)\n parser.add_argument(\"--iteration\", default=-1, type=int)\n parser.add_argument(\"--data_type\", default='kitti360', type=str)\n parser.add_argument(\"--affine\", action=\"store_true\")\n parser.add_argument(\"--ignore_dynamic\", action=\"store_true\")\n parser.add_argument(\"--skip_train\", action=\"store_true\")\n parser.add_argument(\"--skip_test\", action=\"store_true\")\n parser.add_argument(\"--quiet\", action=\"store_true\")\n args = get_combined_args(parser)\n print(\"Rendering \" + args.model_path)\n args.source_path = os.path.join(args.model_path, 'data')\n\n # Initialize system state (RNG)\n # safe_state(args.quiet)\n\n render_sets(model.extract(args), args.iteration, pipeline.extract(args), \n args.skip_train, args.skip_test, args.data_type, args.affine, args.ignore_dynamic)" }, { "path": "requirements.txt", "content": "config==0.5.1\ndatasets==2.19.2\n# flow_vis_torch==0.1\nimageio==2.34.1\nmatplotlib==3.9.0\nnetwork==0.1\nnumpy==1.26.4\nopen3d==0.18.0\nopencv_python==4.10.0.82\nPillow==10.3.0\nplyfile==1.0.3\n# pytorch3d==0.7.4\nrunx==0.0.11\nscipy==1.13.1\nsetuptools==69.5.1\n# torch==2.3.1+cu118\ntorchmetrics==1.4.0.post0\n# torchvision==0.18.1+cu118\ntqdm==4.66.4" }, { "path": "scene/__init__.py", "content": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# This software is free for non-commercial, research and evaluation use \n# under the terms of the LICENSE.md file.\n#\n# For inquiries contact george.drettakis@inria.fr\n#\n\nimport os\nimport random\nimport json\nfrom utils.system_utils import searchForMaxIteration\nfrom scene.dataset_readers import sceneLoadTypeCallbacks\nfrom scene.gaussian_model import GaussianModel\nfrom arguments import ModelParams\nfrom utils.camera_utils import cameraList_from_camInfos, camera_to_JSON\nimport torch\nimport open3d as o3d\nimport numpy as np\nfrom utils.dynamic_utils import create_unicycle_model\nimport shutil\n\nclass Scene:\n\n gaussians : GaussianModel\n\n def __init__(self, args : ModelParams, gaussians : GaussianModel, load_iteration=None, shuffle=True, \n unicycle=False, uc_fit_iter=0, resolution_scales=[1.0], data_type='kitti360', ignore_dynamic=False):\n \"\"\"b\n :param path: Path to colmap scene main folder.\n \"\"\"\n self.model_path = args.model_path\n self.loaded_iter = None\n self.gaussians = gaussians\n\n if load_iteration:\n if load_iteration == -1:\n self.loaded_iter = searchForMaxIteration(os.path.join(self.model_path, \"ckpts\"))\n else:\n self.loaded_iter = load_iteration\n print(\"Loading trained model at iteration {}\".format(self.loaded_iter))\n\n self.train_cameras = {}\n self.test_cameras = {}\n if os.path.exists(os.path.join(args.source_path, \"sparse\")):\n # scene_info = sceneLoadTypeCallbacks[\"Colmap\"](args.source_path, args.images, args.eval)\n raise NotImplementedError\n elif os.path.exists(os.path.join(args.source_path, \"transforms_train.json\")):\n print(\"Found transforms_train.json file, assuming Blender data set!\")\n # scene_info = sceneLoadTypeCallbacks[\"Blender\"](args.source_path, args.white_background, args.eval)\n raise NotImplementedError\n elif os.path.exists(os.path.join(args.source_path, \"meta_data.json\")):\n print(\"Found meta_data.json file, assuming Studio data set!\")\n scene_info = sceneLoadTypeCallbacks['Studio'](args.source_path, args.white_background, args.eval, data_type, ignore_dynamic)\n else:\n assert False, \"Could not recognize scene type!\"\n\n self.dynamic_verts = scene_info.verts\n self.dynamic_gaussians = {}\n for track_id in scene_info.verts:\n self.dynamic_gaussians[track_id] = GaussianModel(args.sh_degree, feat_mutable=False)\n \n if unicycle:\n self.unicycles = create_unicycle_model(scene_info.train_cameras, self.model_path, uc_fit_iter, data_type)\n else:\n self.unicycles = {}\n\n if not self.loaded_iter:\n with open(scene_info.ply_path, 'rb') as src_file, open(os.path.join(self.model_path, \"input.ply\") , 'wb') as dest_file:\n dest_file.write(src_file.read())\n json_cams = []\n camlist = []\n if scene_info.test_cameras:\n camlist.extend(scene_info.test_cameras)\n if scene_info.train_cameras:\n camlist.extend(scene_info.train_cameras)\n for id, cam in enumerate(camlist):\n json_cams.append(camera_to_JSON(id, cam))\n with open(os.path.join(self.model_path, \"cameras.json\"), 'w') as file:\n json.dump(json_cams, file)\n shutil.copyfile(os.path.join(args.source_path, 'meta_data.json'), os.path.join(self.model_path, 'meta_data.json'))\n\n if shuffle:\n random.shuffle(scene_info.train_cameras) # Multi-res consistent random shuffling\n random.shuffle(scene_info.test_cameras) # Multi-res consistent random shuffling\n\n self.cameras_extent = scene_info.nerf_normalization[\"radius\"]\n\n for resolution_scale in resolution_scales:\n print(\"Loading Training Cameras\")\n self.train_cameras[resolution_scale] = cameraList_from_camInfos(scene_info.train_cameras, resolution_scale, args)\n print(\"Loading Test Cameras\")\n self.test_cameras[resolution_scale] = cameraList_from_camInfos(scene_info.test_cameras, resolution_scale, args)\n\n if self.loaded_iter:\n (model_params, first_iter) = torch.load(os.path.join(self.model_path, \"ckpts\", f\"chkpnt{self.loaded_iter}.pth\"))\n gaussians.restore(model_params, None)\n for iid, dynamic_gaussian in self.dynamic_gaussians.items():\n (model_params, first_iter) = torch.load(os.path.join(self.model_path, \"ckpts\", f\"dynamic_{iid}_chkpnt{self.loaded_iter}.pth\"))\n dynamic_gaussian.restore(model_params, None)\n for iid, unicycle_pkg in self.unicycles.items():\n model_params = torch.load(os.path.join(self.model_path, \"ckpts\", f\"unicycle_{iid}_chkpnt{self.loaded_iter}.pth\"))\n unicycle_pkg['model'].restore(model_params)\n else:\n self.gaussians.create_from_pcd(scene_info.point_cloud, self.cameras_extent)\n for track_id in self.dynamic_gaussians.keys():\n vertices = scene_info.verts[track_id]\n\n # init from template\n l, h, w = vertices[:, 0].max() - vertices[:, 0].min(), vertices[:, 1].max() - vertices[:, 1].min(), vertices[:, 2].max() - vertices[:, 2].min()\n pcd = o3d.io.read_point_cloud(f\"utils/vehicle_template/benz_{data_type}.ply\")\n points = np.array(pcd.points) * np.array([l, h, w])\n pcd.points = o3d.utility.Vector3dVector(points)\n pcd.colors = o3d.utility.Vector3dVector(np.ones_like(points) * 0.5)\n\n self.dynamic_gaussians[track_id].create_from_pcd(pcd, self.cameras_extent)\n\n def save(self, iteration):\n # self.gaussians.save_ply(os.path.join(point_cloud_path, \"point_cloud.ply\"))\n point_cloud_vis_path = os.path.join(self.model_path, \"point_cloud_vis/iteration_{}\".format(iteration))\n self.gaussians.save_vis_ply(os.path.join(point_cloud_vis_path, \"point.ply\"))\n for iid, dynamic_gaussian in self.dynamic_gaussians.items():\n dynamic_gaussian.save_vis_ply(os.path.join(point_cloud_vis_path, f\"dynamic_{iid}.ply\"))\n\n def getTrainCameras(self, scale=1.0):\n return self.train_cameras[scale]\n\n def getTestCameras(self, scale=1.0):\n return self.test_cameras[scale]" }, { "path": "scene/cameras.py", "content": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# This software is free for non-commercial, research and evaluation use \n# under the terms of the LICENSE.md file.\n#\n# For inquiries contact george.drettakis@inria.fr\n#\n\nimport torch\nfrom torch import nn\nimport numpy as np\nfrom utils.graphics_utils import getWorld2View2, getProjectionMatrix, fov2focal\nfrom utils.general_utils import decode_op\n\nclass Camera(nn.Module):\n def __init__(self, colmap_id, R, T, K, FoVx, FoVy, image,\n image_name, uid,\n trans=np.array([0.0, 0.0, 0.0]), scale=1.0, data_device=\"cuda\", \n cx_ratio=None, cy_ratio=None, semantic2d=None, mask=None, timestamp=-1, optical_image=None, dynamics={}\n ):\n super(Camera, self).__init__()\n self.uid = uid\n self.colmap_id = colmap_id\n self.R = R\n self.T = T\n self.K = K\n self.FoVx = FoVx\n self.FoVy = FoVy\n self.image_name = image_name\n self.cx_ratio = cx_ratio\n self.cy_ratio = cy_ratio\n self.timestamp = timestamp\n _, self.H, self.W = image.shape\n self.w2c = np.eye(4)\n self.w2c[:3, :3] = self.R.T\n self.w2c[:3, 3] = self.T\n self.c2w = torch.from_numpy(np.linalg.inv(self.w2c)).cuda()\n self.fx = fov2focal(self.FoVx, self.W)\n self.fy = fov2focal(self.FoVy, self.H)\n self.dynamics = dynamics\n\n try:\n self.data_device = torch.device(data_device)\n except Exception as e:\n print(e)\n print(f\"[Warning] Custom device {data_device} failed, fallback to default cuda device\" )\n self.data_device = torch.device(\"cuda\")\n\n self.original_image = image.clamp(0.0, 1.0).to(self.data_device)\n if semantic2d is not None:\n self.semantic2d = semantic2d.to(self.data_device)\n else:\n self.semantic2d = None\n if mask is not None:\n self.mask = torch.from_numpy(mask).bool().to(self.data_device)\n else:\n self.mask = None\n self.image_width = self.original_image.shape[2]\n self.image_height = self.original_image.shape[1]\n if optical_image is not None:\n self.optical_gt = torch.from_numpy(optical_image).to(self.data_device)\n else:\n self.optical_gt = None\n\n self.zfar = 100.0\n self.znear = 0.01\n\n self.trans = trans\n self.scale = scale\n\n self.world_view_transform = torch.tensor(getWorld2View2(R, T, trans, scale)).transpose(0, 1).cuda()\n self.projection_matrix = getProjectionMatrix(znear=self.znear, zfar=self.zfar, \n fovX=self.FoVx, fovY=self.FoVy, cx_ratio=cx_ratio, cy_ratio=cy_ratio).transpose(0,1).cuda()\n self.full_proj_transform = (self.world_view_transform.unsqueeze(0).bmm(self.projection_matrix.unsqueeze(0))).squeeze(0)\n self.camera_center = self.world_view_transform.inverse()[3, :3]\n\n def get_rays(self):\n i, j = torch.meshgrid(torch.linspace(0, self.W-1, self.W), \n torch.linspace(0, self.H-1, self.H)) # pytorch's meshgrid has indexing='ij'\n i = i.t()\n j = j.t()\n dirs = torch.stack([(i-self.cx_ratio)/self.fx, -(j-self.cy_ratio)/self.fy, -torch.ones_like(i)], -1)\n rays_d = torch.sum(dirs[..., np.newaxis, :] * self.c2w[:3,:3], -1).to(self.data_device)\n rays_o = self.c2w[:3,-1].expand(rays_d.shape).to(self.data_device)\n rays_d = torch.nn.functional.normalize(rays_d, dim=-1)\n return rays_o.permute(2,0,1), rays_d.permute(2,0,1)\n\nclass MiniCam:\n def __init__(self, width, height, fovy, fovx, znear, zfar, world_view_transform, full_proj_transform):\n self.image_width = width\n self.image_height = height \n self.FoVy = fovy\n self.FoVx = fovx\n self.znear = znear\n self.zfar = zfar\n self.world_view_transform = world_view_transform\n self.full_proj_transform = full_proj_transform\n view_inv = torch.inverse(self.world_view_transform)\n self.camera_center = view_inv[3][:3]\n\n" }, { "path": "scene/dataset_readers.py", "content": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# This software is free for non-commercial, research and evaluation use \n# under the terms of the LICENSE.md file.\n#\n# For inquiries contact george.drettakis@inria.fr\n#\n\nimport os\nimport sys\nfrom PIL import Image\nfrom typing import NamedTuple\nfrom utils.graphics_utils import getWorld2View2, focal2fov, fov2focal\nimport numpy as np\nimport json\nfrom pathlib import Path\nfrom plyfile import PlyData, PlyElement\nfrom utils.sh_utils import SH2RGB\nfrom scene.gaussian_model import BasicPointCloud\nimport torch.nn.functional as F\nfrom imageio.v2 import imread\nimport torch\nimport random\n\n\nclass CameraInfo(NamedTuple):\n uid: int\n R: np.array\n T: np.array\n K: np.array\n FovY: np.array\n FovX: np.array\n image: np.array\n image_path: str\n image_name: str\n width: int\n height: int\n cx_ratio: float\n cy_ratio: float\n semantic2d: np.array\n optical_image: np.array\n mask: np.array\n timestamp: int\n dynamics: dict\n\nclass SceneInfo(NamedTuple):\n point_cloud: BasicPointCloud\n train_cameras: list\n test_cameras: list\n nerf_normalization: dict\n ply_path: str\n verts: dict\n\ndef getNerfppNorm(cam_info):\n def get_center_and_diag(cam_centers):\n cam_centers = np.hstack(cam_centers)\n avg_cam_center = np.mean(cam_centers, axis=1, keepdims=True)\n center = avg_cam_center\n dist = np.linalg.norm(cam_centers - center, axis=0, keepdims=True)\n diagonal = np.max(dist)\n return center.flatten(), diagonal\n\n cam_centers = []\n\n for cam in cam_info:\n W2C = getWorld2View2(cam.R, cam.T)\n C2W = np.linalg.inv(W2C)\n cam_centers.append(C2W[:3, 3:4]) # cam_centers in world coordinate\n\n center, diagonal = get_center_and_diag(cam_centers)\n # radius = diagonal * 1.1 + 30\n radius = 10\n\n translate = -center\n\n return {\"translate\": translate, \"radius\": radius}\n\ndef fetchPly(path):\n plydata = PlyData.read(path)\n vertices = plydata['vertex']\n positions = np.vstack([vertices['x'], vertices['y'], vertices['z']]).T\n if 'red' in vertices:\n colors = np.vstack([vertices['red'], vertices['green'], vertices['blue']]).T / 255.0\n else:\n print('Create random colors')\n # shs = np.random.random((positions.shape[0], 3)) / 255.0\n shs = np.ones((positions.shape[0], 3)) * 0.5\n colors = SH2RGB(shs)\n # shs = np.ones((positions.shape[0], 3)) * 0.5\n # colors = SH2RGB(shs)\n normals = np.zeros((positions.shape[0], 3))\n return BasicPointCloud(points=positions, colors=colors, normals=normals)\n\ndef storePly(path, xyz, rgb):\n # Define the dtype for the structured array\n dtype = [('x', 'f4'), ('y', 'f4'), ('z', 'f4'),\n ('nx', 'f4'), ('ny', 'f4'), ('nz', 'f4'),\n ('red', 'u1'), ('green', 'u1'), ('blue', 'u1')]\n \n normals = np.zeros_like(xyz)\n\n elements = np.empty(xyz.shape[0], dtype=dtype)\n attributes = np.concatenate((xyz, normals, rgb), axis=1)\n elements[:] = list(map(tuple, attributes))\n\n # Create the PlyData object and write to file\n vertex_element = PlyElement.describe(elements, 'vertex')\n ply_data = PlyData([vertex_element])\n ply_data.write(path)\n\ndef readStudioCameras(path, white_background, data_type, ignore_dynamic):\n train_cam_infos, test_cam_infos = [], []\n with open(os.path.join(path, 'meta_data.json')) as json_file:\n meta_data = json.load(json_file)\n\n verts = {}\n if 'verts' in meta_data and not ignore_dynamic:\n verts_list = meta_data['verts']\n for k, v in verts_list.items():\n verts[k] = np.array(v)\n\n frames = meta_data['frames']\n for idx, frame in enumerate(frames):\n matrix = np.linalg.inv(np.array(frame['camtoworld']))\n R = matrix[:3, :3]\n T = matrix[:3, 3]\n R = np.transpose(R)\n\n rgb_path = os.path.join(path, frame['rgb_path'].replace('./', ''))\n\n rgb_split = rgb_path.split('/')\n image_name = '_'.join([rgb_split[-2], rgb_split[-1][:-4]])\n image = Image.open(rgb_path)\n\n semantic_2d = None\n semantic_pth = rgb_path.replace(\"images\", \"semantics\").replace('.png', '.npy').replace('.jpg', '.npy')\n if os.path.exists(semantic_pth):\n semantic_2d = np.load(semantic_pth)\n semantic_2d[(semantic_2d == 14) | (semantic_2d == 15)] = 13\n\n optical_path = rgb_path.replace(\"images\", \"flow\").replace('.png', '_flow.npy').replace('.jpg', '_flow.npy')\n if os.path.exists(optical_path):\n optical_image = np.load(optical_path)\n else:\n optical_image = None\n\n mask = None\n mask_path = rgb_path.replace(\"images\", \"masks\").replace('.png', '.npy').replace('.jpg', '.npy')\n if os.path.exists(mask_path):\n mask = np.load(mask_path)\n\n timestamp = frame.get('timestamp', -1)\n\n intrinsic = np.array(frame['intrinsics'])\n FovX = focal2fov(intrinsic[0, 0], image.size[0])\n FovY = focal2fov(intrinsic[1, 1], image.size[1])\n cx, cy = intrinsic[0, 2], intrinsic[1, 2]\n w, h = image.size\n \n dynamics = {}\n if 'dynamics' in frame and not ignore_dynamic:\n dynamics_list = frame['dynamics']\n for iid in dynamics_list.keys():\n dynamics[iid] = torch.tensor(dynamics_list[iid]).cuda()\n \n cam_info = CameraInfo(uid=idx, R=R, T=T, K=intrinsic, FovY=FovY, FovX=FovX, image=image,\n image_path=rgb_path, image_name=image_name, width=image.size[0],\n height=image.size[1], cx_ratio=2*cx/w, cy_ratio=2*cy/h, semantic2d=semantic_2d, \n optical_image=optical_image, mask=mask, timestamp=timestamp, dynamics=dynamics)\n \n # kitti360\n if data_type == 'kitti360':\n # if 'cam_2' in cam_info.image_name or 'cam_3' in cam_info.image_name:\n # train_cam_infos.append(cam_info)\n # # continue\n if idx < 20:\n train_cam_infos.append(cam_info)\n elif idx % 8 < 4:\n train_cam_infos.append(cam_info)\n elif idx % 8 >= 4:\n test_cam_infos.append(cam_info)\n else:\n continue\n\n elif data_type == 'kitti':\n if idx < 10 or idx >= len(frames) - 4:\n train_cam_infos.append(cam_info)\n elif idx % 4 < 2:\n train_cam_infos.append(cam_info)\n elif idx % 4 == 2:\n test_cam_infos.append(cam_info)\n else:\n continue\n\n elif data_type == \"nuscenes\":\n if idx < 600 or idx >= 1200:\n continue\n elif idx % 30 >= 24:\n # print('test', cam_info.image_name)\n test_cam_infos.append(cam_info)\n else:\n # print('train', cam_info.image_name)\n train_cam_infos.append(cam_info)\n\n elif data_type == \"waymo\":\n if idx > 10 and idx % 10 >= 9:\n test_cam_infos.append(cam_info)\n else:\n train_cam_infos.append(cam_info)\n\n elif data_type == \"pandaset\":\n # if idx >= 360:\n # continue\n if idx > 30 and idx % 30 >= 24:\n test_cam_infos.append(cam_info)\n else:\n train_cam_infos.append(cam_info)\n \n else:\n raise NotImplementedError\n return train_cam_infos, test_cam_infos, verts\n\n\ndef readStudioInfo(path, white_background, eval, data_type, ignore_dynamic):\n train_cam_infos, test_cam_infos, verts = readStudioCameras(path, white_background, data_type, ignore_dynamic)\n\n print(f'Loaded {len(train_cam_infos)} train cameras and {len(test_cam_infos)} test cameras')\n nerf_normalization = getNerfppNorm(train_cam_infos)\n\n ply_path = os.path.join(path, \"points3d.ply\")\n # ply_path = os.path.join(path, 'lidar', 'cat.ply')\n if not os.path.exists(ply_path):\n # Since this data set has no colmap data, we start with random points\n num_pts = 500_000\n print(f\"Generating random point cloud ({num_pts})...\")\n\n # We create random points inside the bounds of the synthetic Blender scenes\n AABB = [[-20, -25, -20], [20, 5, 80]]\n xyz = np.random.uniform(AABB[0], AABB[1], (500000, 3))\n # xyz = np.load(os.path.join(path, 'lidar_point.npy'))\n num_pts = xyz.shape[0]\n shs = np.ones((num_pts, 3)) * 0.5\n pcd = BasicPointCloud(points=xyz, colors=SH2RGB(shs), normals=np.zeros((num_pts, 3)))\n\n storePly(ply_path, xyz, SH2RGB(shs) * 255)\n try:\n pcd = fetchPly(ply_path)\n except Exception as e:\n print('When loading point clound, meet error:', e)\n exit(0)\n\n scene_info = SceneInfo(point_cloud=pcd,\n train_cameras=train_cam_infos,\n test_cameras=test_cam_infos,\n nerf_normalization=nerf_normalization,\n ply_path=ply_path,\n verts=verts)\n return scene_info\n\n\nsceneLoadTypeCallbacks = {\n \"Studio\": readStudioInfo,\n}" }, { "path": "scene/gaussian_model.py", "content": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# This software is free for non-commercial, research and evaluation use \n# under the terms of the LICENSE.md file.\n#\n# For inquiries contact george.drettakis@inria.fr\n#\n\nimport torch\nimport numpy as np\nfrom utils.general_utils import inverse_sigmoid, get_expon_lr_func, build_rotation\nfrom torch import nn\nimport os\nfrom utils.system_utils import mkdir_p\nfrom plyfile import PlyData, PlyElement\nfrom utils.sh_utils import RGB2SH, SH2RGB\nfrom simple_knn._C import distCUDA2\nfrom utils.graphics_utils import BasicPointCloud\nfrom utils.general_utils import strip_symmetric, build_scaling_rotation\nimport open3d as o3d\nimport tinycudann as tcnn\nfrom math import sqrt\n\nclass CustomAdam(torch.optim.Optimizer):\n def __init__(self, params, lr=0.001, betas=(0.9, 0.999), eps=1e-8):\n defaults = dict(lr=lr, betas=betas, eps=eps)\n super(CustomAdam, self).__init__(params, defaults)\n \n def step(self, custom_lr=None, name=None):\n for group in self.param_groups:\n for p in group['params']:\n if p.grad is None:\n continue\n\n grad = p.grad.data\n if grad.is_sparse:\n raise RuntimeError('Adam does not support sparse gradients')\n\n state = self.state[p]\n\n # State initialization\n if len(state) == 0:\n state['step'] = 0\n # Exponential moving averages of gradient values\n state['exp_avg'] = torch.zeros_like(p.data)\n # Exponential moving averages of squared gradient values\n state['exp_avg_sq'] = torch.zeros_like(p.data)\n\n exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n beta1, beta2 = group['betas']\n\n # Add op to update moving averages\n state['step'] += 1\n exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)\n exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)\n\n denom = exp_avg_sq.sqrt().add_(group['eps'])\n \n bias_correction1 = 1.0 - beta1 ** state['step']\n bias_correction2 = 1.0 - beta2 ** state['step']\n\n if (custom_lr is not None) and (name is not None) and (group['name'] in name):\n step_size = custom_lr[:, None] * group['lr'] * (sqrt(bias_correction2) / bias_correction1)\n else:\n step_size = group['lr'] * (sqrt(bias_correction2) / bias_correction1)\n \n p.data -= step_size * exp_avg / denom\n \nclass GaussianModel:\n\n def setup_functions(self):\n def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation):\n L = build_scaling_rotation(scaling_modifier * scaling, rotation)\n actual_covariance = L @ L.transpose(1, 2)\n symm = strip_symmetric(actual_covariance)\n return symm\n \n self.scaling_activation = torch.exp\n self.scaling_inverse_activation = torch.log\n\n self.covariance_activation = build_covariance_from_scaling_rotation\n\n self.opacity_activation = torch.sigmoid\n self.inverse_opacity_activation = inverse_sigmoid\n\n self.rotation_activation = torch.nn.functional.normalize\n\n\n def __init__(self, sh_degree : int, feat_mutable=True, affine=False):\n self.active_sh_degree = 0\n self.max_sh_degree = sh_degree \n self._xyz = torch.empty(0)\n self._features_dc = torch.empty(0)\n self._features_rest = torch.empty(0)\n self._feats3D = torch.empty(0)\n self._scaling = torch.empty(0)\n self._rotation = torch.empty(0)\n self._opacity = torch.empty(0)\n self.max_radii2D = torch.empty(0)\n self.xyz_gradient_accum = torch.empty(0)\n self.denom = torch.empty(0)\n self.optimizer = None\n self.percent_dense = 0\n self.spatial_lr_scale = 0\n self.feat_mutable = feat_mutable\n self.setup_functions()\n\n self.pos_enc = tcnn.Encoding(\n n_input_dims=3,\n encoding_config={\"otype\": \"Frequency\", \"n_frequencies\": 2},\n )\n self.dir_enc = tcnn.Encoding(\n n_input_dims=3,\n encoding_config={\n \"otype\": \"SphericalHarmonics\",\n \"degree\": 3,\n },\n )\n\n self.affine = affine\n if affine:\n self.appearance_model = tcnn.Network(\n n_input_dims=self.pos_enc.n_output_dims + self.dir_enc.n_output_dims,\n n_output_dims=12,\n network_config={\n \"otype\": \"FullyFusedMLP\",\n \"activation\": \"ReLU\",\n \"output_activation\": \"None\",\n \"n_neurons\": 32,\n \"n_hidden_layers\": 2,\n }\n )\n else:\n self.appearance_model = None\n\n def capture(self):\n return (\n self.active_sh_degree,\n self._xyz,\n self._features_dc,\n self._features_rest,\n self._feats3D,\n self._scaling,\n self._rotation,\n self._opacity,\n self.max_radii2D,\n self.xyz_gradient_accum,\n self.denom,\n self.optimizer.state_dict(),\n self.spatial_lr_scale,\n self.appearance_model,\n )\n \n def restore(self, model_args, training_args):\n (self.active_sh_degree, \n self._xyz, \n self._features_dc, \n self._features_rest,\n self._feats3D,\n self._scaling, \n self._rotation, \n self._opacity,\n self.max_radii2D, \n xyz_gradient_accum, \n denom,\n opt_dict, \n self.spatial_lr_scale,\n self.appearance_model,) = model_args\n self.xyz_gradient_accum = xyz_gradient_accum\n self.denom = denom\n if training_args is not None:\n self.training_setup(training_args)\n self.optimizer.load_state_dict(opt_dict)\n\n @property\n def get_scaling(self):\n return self.scaling_activation(self._scaling)\n \n @property\n def get_rotation(self):\n return self.rotation_activation(self._rotation)\n \n # TODO add get_xyz for dynamic car\n @property\n def get_xyz(self):\n return self._xyz\n \n @property\n def get_features(self):\n features_dc = self._features_dc\n features_rest = self._features_rest\n return torch.cat((features_dc, features_rest), dim=1)\n \n @property\n def get_3D_features(self):\n return torch.softmax(self._feats3D, dim=-1)\n \n @property\n def get_opacity(self):\n return self.opacity_activation(self._opacity)\n \n def get_covariance(self, scaling_modifier = 1):\n return self.covariance_activation(self.get_scaling, scaling_modifier, self._rotation)\n\n def oneupSHdegree(self):\n if self.active_sh_degree < self.max_sh_degree:\n self.active_sh_degree += 1\n\n def create_from_pcd(self, pcd : BasicPointCloud, spatial_lr_scale : float):\n # self.spatial_lr_scale = 1\n self.spatial_lr_scale = spatial_lr_scale\n fused_point_cloud = torch.tensor(np.asarray(pcd.points)).float().cuda()\n fused_color = RGB2SH(torch.tensor(np.asarray(pcd.colors)).float().cuda())\n features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda()\n features[:, :3, 0 ] = fused_color\n features[:, 3:, 1:] = 0.0\n\n if self.feat_mutable:\n feats3D = torch.zeros(fused_color.shape[0], 20).float().cuda()\n self._feats3D = nn.Parameter(feats3D.requires_grad_(True))\n else:\n feats3D = torch.zeros(fused_color.shape[0], 20).float().cuda()\n feats3D[:, 13] = 1\n self._feats3D = nn.Parameter(feats3D.requires_grad_(True))\n\n print(\"Number of points at initialisation : \", fused_point_cloud.shape[0])\n\n dist2 = torch.clamp_min(distCUDA2(torch.from_numpy(np.asarray(pcd.points)).float().cuda()), 0.0000001)\n scales = torch.log(torch.sqrt(dist2))[...,None].repeat(1, 3)\n rots = torch.zeros((fused_point_cloud.shape[0], 4), device=\"cuda\")\n rots[:, 0] = 1\n\n opacities = inverse_sigmoid(0.1 * torch.ones((fused_point_cloud.shape[0], 1), dtype=torch.float, device=\"cuda\"))\n\n self._xyz = nn.Parameter(fused_point_cloud.requires_grad_(True))\n self._features_dc = nn.Parameter(features[:,:,0:1].transpose(1, 2).contiguous().requires_grad_(True))\n self._features_rest = nn.Parameter(features[:,:,1:].transpose(1, 2).contiguous().requires_grad_(True))\n self._scaling = nn.Parameter(scales.requires_grad_(True))\n self._rotation = nn.Parameter(rots.requires_grad_(True))\n self._opacity = nn.Parameter(opacities.requires_grad_(True))\n self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device=\"cuda\")\n\n def training_setup(self, training_args):\n self.percent_dense = training_args.percent_dense\n self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device=\"cuda\")\n self.denom = torch.zeros((self.get_xyz.shape[0], 1), device=\"cuda\")\n\n # self.spatial_lr_scale /= 3\n\n l = [\n {'params': [self._xyz], 'lr': training_args.position_lr_init*self.spatial_lr_scale, \"name\": \"xyz\"},\n {'params': [self._features_dc], 'lr': training_args.feature_lr, \"name\": \"f_dc\"},\n {'params': [self._features_rest], 'lr': training_args.feature_lr / 20.0, \"name\": \"f_rest\"},\n {'params': [self._opacity], 'lr': training_args.opacity_lr, \"name\": \"opacity\"},\n {'params': [self._scaling], 'lr': training_args.scaling_lr*self.spatial_lr_scale, \"name\": \"scaling\"},\n {'params': [self._rotation], 'lr': training_args.rotation_lr, \"name\": \"rotation\"},\n ]\n\n if self.affine:\n l.append({'params': [*self.appearance_model.parameters()], 'lr': 1e-3, \"name\": \"appearance_model\"})\n \n if self.feat_mutable:\n l.append({'params': [self._feats3D], 'lr': 1e-2, \"name\": \"feats3D\"})\n\n self.optimizer = torch.optim.Adam(l, lr=0.0, eps=1e-15)\n # self.optimizer = CustomAdam(l, lr=0.0, eps=1e-15)\n self.xyz_scheduler_args = get_expon_lr_func(lr_init=training_args.position_lr_init*self.spatial_lr_scale,\n lr_final=training_args.position_lr_final*self.spatial_lr_scale,\n lr_delay_mult=training_args.position_lr_delay_mult,\n max_steps=training_args.position_lr_max_steps)\n\n def update_learning_rate(self, iteration):\n ''' Learning rate scheduling per step '''\n for param_group in self.optimizer.param_groups:\n if param_group[\"name\"] == \"xyz\":\n lr = self.xyz_scheduler_args(iteration)\n param_group['lr'] = lr\n return lr\n\n def construct_list_of_attributes(self):\n l = ['x', 'y', 'z', 'nx', 'ny', 'nz']\n # All channels except the 3 DC\n for i in range(self._features_dc.shape[1]*self._features_dc.shape[2]):\n l.append('f_dc_{}'.format(i))\n for i in range(self._features_rest.shape[1]*self._features_rest.shape[2]):\n l.append('f_rest_{}'.format(i))\n for i in range(self._feats3D.shape[1]):\n l.append('semantic_{}'.format(i))\n l.append('opacity')\n for i in range(self._scaling.shape[1]):\n l.append('scale_{}'.format(i))\n for i in range(self._rotation.shape[1]):\n l.append('rot_{}'.format(i))\n return l\n\n def save_ply(self, path):\n mkdir_p(os.path.dirname(path))\n\n xyz = self._xyz.detach().cpu().numpy()\n normals = np.zeros_like(xyz)\n f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()\n f_rest = self._features_rest.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()\n feats3D = self._feats3D.detach().cpu().numpy()\n opacities = self._opacity.detach().cpu().numpy()\n scale = self._scaling.detach().cpu().numpy()\n rotation = self._rotation.detach().cpu().numpy()\n\n dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()]\n\n elements = np.empty(xyz.shape[0], dtype=dtype_full)\n attributes = np.concatenate((xyz, normals, f_dc, f_rest, feats3D, opacities, scale, rotation), axis=1)\n elements[:] = list(map(tuple, attributes))\n el = PlyElement.describe(elements, 'vertex')\n PlyData([el]).write(path)\n\n def save_vis_ply(self, path):\n mkdir_p(os.path.dirname(path))\n xyz = self.get_xyz.detach().cpu().numpy()\n pcd = o3d.geometry.PointCloud()\n pcd.points = o3d.utility.Vector3dVector(xyz)\n colors = SH2RGB(self._features_dc[:, 0, :].detach().cpu().numpy()).clip(0, 1)\n pcd.colors = o3d.utility.Vector3dVector(colors)\n o3d.io.write_point_cloud(path, pcd)\n\n def reset_opacity(self):\n opacities_new = inverse_sigmoid(torch.min(self.get_opacity, torch.ones_like(self.get_opacity)*0.01))\n optimizable_tensors = self.replace_tensor_to_optimizer(opacities_new, \"opacity\")\n self._opacity = optimizable_tensors[\"opacity\"]\n\n def load_ply(self, path):\n plydata = PlyData.read(path)\n\n xyz = np.stack((np.asarray(plydata.elements[0][\"x\"]),\n np.asarray(plydata.elements[0][\"y\"]),\n np.asarray(plydata.elements[0][\"z\"])), axis=1)\n opacities = np.asarray(plydata.elements[0][\"opacity\"])[..., np.newaxis]\n\n features_dc = np.zeros((xyz.shape[0], 3, 1))\n features_dc[:, 0, 0] = np.asarray(plydata.elements[0][\"f_dc_0\"])\n features_dc[:, 1, 0] = np.asarray(plydata.elements[0][\"f_dc_1\"])\n features_dc[:, 2, 0] = np.asarray(plydata.elements[0][\"f_dc_2\"])\n\n extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith(\"f_rest_\")]\n assert len(extra_f_names)==3*(self.max_sh_degree + 1) ** 2 - 3\n features_extra = np.zeros((xyz.shape[0], len(extra_f_names)))\n for idx, attr_name in enumerate(extra_f_names):\n features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name])\n # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC)\n features_extra = features_extra.reshape((features_extra.shape[0], 3, (self.max_sh_degree + 1) ** 2 - 1))\n\n scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith(\"scale_\")]\n scales = np.zeros((xyz.shape[0], len(scale_names)))\n for idx, attr_name in enumerate(scale_names):\n scales[:, idx] = np.asarray(plydata.elements[0][attr_name])\n\n rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith(\"rot\")]\n rots = np.zeros((xyz.shape[0], len(rot_names)))\n for idx, attr_name in enumerate(rot_names):\n rots[:, idx] = np.asarray(plydata.elements[0][attr_name])\n\n self._xyz = nn.Parameter(torch.tensor(xyz, dtype=torch.float, device=\"cuda\").requires_grad_(True))\n self._features_dc = nn.Parameter(torch.tensor(features_dc, dtype=torch.float, device=\"cuda\").transpose(1, 2).contiguous().requires_grad_(True))\n self._features_rest = nn.Parameter(torch.tensor(features_extra, dtype=torch.float, device=\"cuda\").transpose(1, 2).contiguous().requires_grad_(True))\n self._opacity = nn.Parameter(torch.tensor(opacities, dtype=torch.float, device=\"cuda\").requires_grad_(True))\n self._scaling = nn.Parameter(torch.tensor(scales, dtype=torch.float, device=\"cuda\").requires_grad_(True))\n self._rotation = nn.Parameter(torch.tensor(rots, dtype=torch.float, device=\"cuda\").requires_grad_(True))\n\n self.active_sh_degree = self.max_sh_degree\n\n def replace_tensor_to_optimizer(self, tensor, name):\n optimizable_tensors = {}\n for group in self.optimizer.param_groups:\n if group[\"name\"] == name:\n stored_state = self.optimizer.state.get(group['params'][0], None)\n stored_state[\"exp_avg\"] = torch.zeros_like(tensor)\n stored_state[\"exp_avg_sq\"] = torch.zeros_like(tensor)\n\n del self.optimizer.state[group['params'][0]]\n group[\"params\"][0] = nn.Parameter(tensor.requires_grad_(True))\n self.optimizer.state[group['params'][0]] = stored_state\n\n optimizable_tensors[group[\"name\"]] = group[\"params\"][0]\n return optimizable_tensors\n\n def _prune_optimizer(self, mask):\n optimizable_tensors = {}\n for group in self.optimizer.param_groups:\n if group['name'] == 'appearance_model':\n continue\n stored_state = self.optimizer.state.get(group['params'][0], None)\n if stored_state is not None:\n stored_state[\"exp_avg\"] = stored_state[\"exp_avg\"][mask]\n stored_state[\"exp_avg_sq\"] = stored_state[\"exp_avg_sq\"][mask]\n\n del self.optimizer.state[group['params'][0]]\n group[\"params\"][0] = nn.Parameter((group[\"params\"][0][mask].requires_grad_(True)))\n self.optimizer.state[group['params'][0]] = stored_state\n\n optimizable_tensors[group[\"name\"]] = group[\"params\"][0]\n else:\n group[\"params\"][0] = nn.Parameter(group[\"params\"][0][mask].requires_grad_(True))\n optimizable_tensors[group[\"name\"]] = group[\"params\"][0]\n return optimizable_tensors\n\n def prune_points(self, mask):\n valid_points_mask = ~mask\n optimizable_tensors = self._prune_optimizer(valid_points_mask)\n\n self._xyz = optimizable_tensors[\"xyz\"]\n self._features_dc = optimizable_tensors[\"f_dc\"]\n self._features_rest = optimizable_tensors[\"f_rest\"]\n if self.feat_mutable:\n self._feats3D = optimizable_tensors[\"feats3D\"]\n else:\n self._feats3D = self._feats3D[1, :].repeat((self._xyz.shape[0], 1))\n self._opacity = optimizable_tensors[\"opacity\"]\n self._scaling = optimizable_tensors[\"scaling\"]\n self._rotation = optimizable_tensors[\"rotation\"]\n\n self.xyz_gradient_accum = self.xyz_gradient_accum[valid_points_mask]\n\n self.denom = self.denom[valid_points_mask]\n self.max_radii2D = self.max_radii2D[valid_points_mask]\n\n def cat_tensors_to_optimizer(self, tensors_dict):\n optimizable_tensors = {}\n for group in self.optimizer.param_groups:\n if group['name'] not in tensors_dict:\n continue\n assert len(group[\"params\"]) == 1\n extension_tensor = tensors_dict[group[\"name\"]]\n stored_state = self.optimizer.state.get(group[\"params\"][0], None)\n if stored_state is not None:\n\n stored_state[\"exp_avg\"] = torch.cat((stored_state[\"exp_avg\"], torch.zeros_like(extension_tensor)), dim=0)\n stored_state[\"exp_avg_sq\"] = torch.cat((stored_state[\"exp_avg_sq\"], torch.zeros_like(extension_tensor)), dim=0)\n\n del self.optimizer.state[group[\"params\"][0]]\n group[\"params\"][0] = nn.Parameter(torch.cat((group[\"params\"][0], extension_tensor), dim=0).requires_grad_(True))\n self.optimizer.state[group[\"params\"][0]] = stored_state\n\n optimizable_tensors[group[\"name\"]] = group[\"params\"][0]\n else:\n group[\"params\"][0] = nn.Parameter(torch.cat((group[\"params\"][0], extension_tensor), dim=0).requires_grad_(True))\n optimizable_tensors[group[\"name\"]] = group[\"params\"][0]\n\n return optimizable_tensors\n\n def densification_postfix(self, new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacities, new_scaling, new_rotation):\n d = {\"xyz\": new_xyz,\n \"f_dc\": new_features_dc,\n \"f_rest\": new_features_rest,\n \"feats3D\": new_feats3D,\n \"opacity\": new_opacities,\n \"scaling\" : new_scaling,\n \"rotation\" : new_rotation}\n\n optimizable_tensors = self.cat_tensors_to_optimizer(d)\n self._xyz = optimizable_tensors[\"xyz\"]\n self._features_dc = optimizable_tensors[\"f_dc\"]\n if self.feat_mutable:\n self._feats3D = optimizable_tensors[\"feats3D\"]\n else:\n self._feats3D = self._feats3D[1, :].repeat((self._xyz.shape[0], 1))\n self._features_rest = optimizable_tensors[\"f_rest\"]\n self._opacity = optimizable_tensors[\"opacity\"]\n self._scaling = optimizable_tensors[\"scaling\"]\n self._rotation = optimizable_tensors[\"rotation\"]\n\n self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device=\"cuda\")\n self.denom = torch.zeros((self.get_xyz.shape[0], 1), device=\"cuda\")\n self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device=\"cuda\")\n\n def densify_and_split(self, grads, grad_threshold, scene_extent, N=2):\n n_init_points = self.get_xyz.shape[0]\n # Extract points that satisfy the gradient condition\n padded_grad = torch.zeros((n_init_points), device=\"cuda\")\n padded_grad[:grads.shape[0]] = grads.squeeze()\n selected_pts_mask = torch.where(padded_grad >= grad_threshold, True, False)\n selected_pts_mask = torch.logical_and(selected_pts_mask,\n torch.max(self.get_scaling, dim=1).values > self.percent_dense*scene_extent)\n\n stds = self.get_scaling[selected_pts_mask].repeat(N,1)\n means =torch.zeros((stds.size(0), 3),device=\"cuda\")\n samples = torch.normal(mean=means, std=stds)\n rots = build_rotation(self._rotation[selected_pts_mask]).repeat(N,1,1)\n new_xyz = torch.bmm(rots, samples.unsqueeze(-1)).squeeze(-1) + self.get_xyz[selected_pts_mask].repeat(N, 1)\n new_scaling = self.scaling_inverse_activation(self.get_scaling[selected_pts_mask].repeat(N,1) / (0.8*N))\n new_rotation = self._rotation[selected_pts_mask].repeat(N,1)\n new_features_dc = self._features_dc[selected_pts_mask].repeat(N,1,1)\n new_features_rest = self._features_rest[selected_pts_mask].repeat(N,1,1)\n new_feats3D = self._feats3D[selected_pts_mask].repeat(N,1)\n new_opacity = self._opacity[selected_pts_mask].repeat(N,1)\n\n self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacity, new_scaling, new_rotation)\n\n prune_filter = torch.cat((selected_pts_mask, torch.zeros(N * selected_pts_mask.sum(), device=\"cuda\", dtype=bool)))\n self.prune_points(prune_filter)\n\n def densify_and_clone(self, grads, grad_threshold, scene_extent):\n # Extract points that satisfy the gradient condition\n selected_pts_mask = torch.where(torch.norm(grads, dim=-1) >= grad_threshold, True, False)\n selected_pts_mask = torch.logical_and(selected_pts_mask,\n torch.max(self.get_scaling, dim=1).values <= self.percent_dense*scene_extent)\n \n new_xyz = self._xyz[selected_pts_mask]\n new_features_dc = self._features_dc[selected_pts_mask]\n new_features_rest = self._features_rest[selected_pts_mask]\n new_feats3D = self._feats3D[selected_pts_mask]\n new_opacities = self._opacity[selected_pts_mask]\n new_scaling = self._scaling[selected_pts_mask]\n new_rotation = self._rotation[selected_pts_mask]\n\n self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_feats3D, new_opacities, new_scaling, new_rotation)\n\n def densify_and_prune(self, max_grad, min_opacity, extent, max_screen_size):\n grads = self.xyz_gradient_accum / self.denom\n grads[grads.isnan()] = 0.0\n\n self.densify_and_clone(grads, max_grad, extent)\n self.densify_and_split(grads, max_grad, extent)\n\n prune_mask = (self.get_opacity < min_opacity).squeeze()\n if max_screen_size:\n big_points_vs = self.max_radii2D > max_screen_size\n big_points_ws = self.get_scaling.max(dim=1).values > 0.1 * extent * 10\n prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws)\n self.prune_points(prune_mask)\n\n torch.cuda.empty_cache()\n\n def add_densification_stats(self, viewspace_point_tensor, update_filter):\n self.xyz_gradient_accum[update_filter] += torch.norm(viewspace_point_tensor.grad[update_filter,:2], dim=-1, keepdim=True)\n self.denom[update_filter] += 1\n \n def add_densification_stats_grad(self, tensor_grad, update_filter):\n self.xyz_gradient_accum[update_filter] += torch.norm(tensor_grad[update_filter,:2], dim=-1, keepdim=True)\n self.denom[update_filter] += 1" }, { "path": "submodules/simple-knn/ext.cpp", "content": "/*\n * Copyright (C) 2023, Inria\n * GRAPHDECO research group, https://team.inria.fr/graphdeco\n * All rights reserved.\n *\n * This software is free for non-commercial, research and evaluation use \n * under the terms of the LICENSE.md file.\n *\n * For inquiries contact george.drettakis@inria.fr\n */\n\n#include \n#include \"spatial.h\"\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"distCUDA2\", &distCUDA2);\n}\n" }, { "path": "submodules/simple-knn/setup.py", "content": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# This software is free for non-commercial, research and evaluation use \n# under the terms of the LICENSE.md file.\n#\n# For inquiries contact george.drettakis@inria.fr\n#\n\nfrom setuptools import setup\nfrom torch.utils.cpp_extension import CUDAExtension, BuildExtension\nimport os\n\ncxx_compiler_flags = []\n\nif os.name == 'nt':\n cxx_compiler_flags.append(\"/wd4624\")\n\nsetup(\n name=\"simple_knn\",\n ext_modules=[\n CUDAExtension(\n name=\"simple_knn._C\",\n sources=[\n \"spatial.cu\", \n \"simple_knn.cu\",\n \"ext.cpp\"],\n extra_compile_args={\"nvcc\": [], \"cxx\": cxx_compiler_flags})\n ],\n cmdclass={\n 'build_ext': BuildExtension\n }\n)\n" }, { "path": "submodules/simple-knn/simple_knn/.gitkeep", "content": "" }, { "path": "submodules/simple-knn/simple_knn.cu", "content": "/*\n * Copyright (C) 2023, Inria\n * GRAPHDECO research group, https://team.inria.fr/graphdeco\n * All rights reserved.\n *\n * This software is free for non-commercial, research and evaluation use \n * under the terms of the LICENSE.md file.\n *\n * For inquiries contact george.drettakis@inria.fr\n */\n\n#define BOX_SIZE 1024\n\n#include \"cuda_runtime.h\"\n#include \"device_launch_parameters.h\"\n#include \"simple_knn.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#define __CUDACC__\n#include \n#include \n\nnamespace cg = cooperative_groups;\n\nstruct CustomMin\n{\n\t__device__ __forceinline__\n\t\tfloat3 operator()(const float3& a, const float3& b) const {\n\t\treturn { min(a.x, b.x), min(a.y, b.y), min(a.z, b.z) };\n\t}\n};\n\nstruct CustomMax\n{\n\t__device__ __forceinline__\n\t\tfloat3 operator()(const float3& a, const float3& b) const {\n\t\treturn { max(a.x, b.x), max(a.y, b.y), max(a.z, b.z) };\n\t}\n};\n\n__host__ __device__ uint32_t prepMorton(uint32_t x)\n{\n\tx = (x | (x << 16)) & 0x030000FF;\n\tx = (x | (x << 8)) & 0x0300F00F;\n\tx = (x | (x << 4)) & 0x030C30C3;\n\tx = (x | (x << 2)) & 0x09249249;\n\treturn x;\n}\n\n__host__ __device__ uint32_t coord2Morton(float3 coord, float3 minn, float3 maxx)\n{\n\tuint32_t x = prepMorton(((coord.x - minn.x) / (maxx.x - minn.x)) * ((1 << 10) - 1));\n\tuint32_t y = prepMorton(((coord.y - minn.y) / (maxx.y - minn.y)) * ((1 << 10) - 1));\n\tuint32_t z = prepMorton(((coord.z - minn.z) / (maxx.z - minn.z)) * ((1 << 10) - 1));\n\n\treturn x | (y << 1) | (z << 2);\n}\n\n__global__ void coord2Morton(int P, const float3* points, float3 minn, float3 maxx, uint32_t* codes)\n{\n\tauto idx = cg::this_grid().thread_rank();\n\tif (idx >= P)\n\t\treturn;\n\n\tcodes[idx] = coord2Morton(points[idx], minn, maxx);\n}\n\nstruct MinMax\n{\n\tfloat3 minn;\n\tfloat3 maxx;\n};\n\n__global__ void boxMinMax(uint32_t P, float3* points, uint32_t* indices, MinMax* boxes)\n{\n\tauto idx = cg::this_grid().thread_rank();\n\n\tMinMax me;\n\tif (idx < P)\n\t{\n\t\tme.minn = points[indices[idx]];\n\t\tme.maxx = points[indices[idx]];\n\t}\n\telse\n\t{\n\t\tme.minn = { FLT_MAX, FLT_MAX, FLT_MAX };\n\t\tme.maxx = { -FLT_MAX,-FLT_MAX,-FLT_MAX };\n\t}\n\n\t__shared__ MinMax redResult[BOX_SIZE];\n\n\tfor (int off = BOX_SIZE / 2; off >= 1; off /= 2)\n\t{\n\t\tif (threadIdx.x < 2 * off)\n\t\t\tredResult[threadIdx.x] = me;\n\t\t__syncthreads();\n\n\t\tif (threadIdx.x < off)\n\t\t{\n\t\t\tMinMax other = redResult[threadIdx.x + off];\n\t\t\tme.minn.x = min(me.minn.x, other.minn.x);\n\t\t\tme.minn.y = min(me.minn.y, other.minn.y);\n\t\t\tme.minn.z = min(me.minn.z, other.minn.z);\n\t\t\tme.maxx.x = max(me.maxx.x, other.maxx.x);\n\t\t\tme.maxx.y = max(me.maxx.y, other.maxx.y);\n\t\t\tme.maxx.z = max(me.maxx.z, other.maxx.z);\n\t\t}\n\t\t__syncthreads();\n\t}\n\n\tif (threadIdx.x == 0)\n\t\tboxes[blockIdx.x] = me;\n}\n\n__device__ __host__ float distBoxPoint(const MinMax& box, const float3& p)\n{\n\tfloat3 diff = { 0, 0, 0 };\n\tif (p.x < box.minn.x || p.x > box.maxx.x)\n\t\tdiff.x = min(abs(p.x - box.minn.x), abs(p.x - box.maxx.x));\n\tif (p.y < box.minn.y || p.y > box.maxx.y)\n\t\tdiff.y = min(abs(p.y - box.minn.y), abs(p.y - box.maxx.y));\n\tif (p.z < box.minn.z || p.z > box.maxx.z)\n\t\tdiff.z = min(abs(p.z - box.minn.z), abs(p.z - box.maxx.z));\n\treturn diff.x * diff.x + diff.y * diff.y + diff.z * diff.z;\n}\n\ntemplate\n__device__ void updateKBest(const float3& ref, const float3& point, float* knn)\n{\n\tfloat3 d = { point.x - ref.x, point.y - ref.y, point.z - ref.z };\n\tfloat dist = d.x * d.x + d.y * d.y + d.z * d.z;\n\tfor (int j = 0; j < K; j++)\n\t{\n\t\tif (knn[j] > dist)\n\t\t{\n\t\t\tfloat t = knn[j];\n\t\t\tknn[j] = dist;\n\t\t\tdist = t;\n\t\t}\n\t}\n}\n\n__global__ void boxMeanDist(uint32_t P, float3* points, uint32_t* indices, MinMax* boxes, float* dists)\n{\n\tint idx = cg::this_grid().thread_rank();\n\tif (idx >= P)\n\t\treturn;\n\n\tfloat3 point = points[indices[idx]];\n\tfloat best[3] = { FLT_MAX, FLT_MAX, FLT_MAX };\n\n\tfor (int i = max(0, idx - 3); i <= min(P - 1, idx + 3); i++)\n\t{\n\t\tif (i == idx)\n\t\t\tcontinue;\n\t\tupdateKBest<3>(point, points[indices[i]], best);\n\t}\n\n\tfloat reject = best[2];\n\tbest[0] = FLT_MAX;\n\tbest[1] = FLT_MAX;\n\tbest[2] = FLT_MAX;\n\n\tfor (int b = 0; b < (P + BOX_SIZE - 1) / BOX_SIZE; b++)\n\t{\n\t\tMinMax box = boxes[b];\n\t\tfloat dist = distBoxPoint(box, point);\n\t\tif (dist > reject || dist > best[2])\n\t\t\tcontinue;\n\n\t\tfor (int i = b * BOX_SIZE; i < min(P, (b + 1) * BOX_SIZE); i++)\n\t\t{\n\t\t\tif (i == idx)\n\t\t\t\tcontinue;\n\t\t\tupdateKBest<3>(point, points[indices[i]], best);\n\t\t}\n\t}\n\tdists[indices[idx]] = (best[0] + best[1] + best[2]) / 3.0f;\n}\n\nvoid SimpleKNN::knn(int P, float3* points, float* meanDists)\n{\n\tfloat3* result;\n\tcudaMalloc(&result, sizeof(float3));\n\tsize_t temp_storage_bytes;\n\n\tfloat3 init = { 0, 0, 0 }, minn, maxx;\n\n\tcub::DeviceReduce::Reduce(nullptr, temp_storage_bytes, points, result, P, CustomMin(), init);\n\tthrust::device_vector temp_storage(temp_storage_bytes);\n\n\tcub::DeviceReduce::Reduce(temp_storage.data().get(), temp_storage_bytes, points, result, P, CustomMin(), init);\n\tcudaMemcpy(&minn, result, sizeof(float3), cudaMemcpyDeviceToHost);\n\n\tcub::DeviceReduce::Reduce(temp_storage.data().get(), temp_storage_bytes, points, result, P, CustomMax(), init);\n\tcudaMemcpy(&maxx, result, sizeof(float3), cudaMemcpyDeviceToHost);\n\n\tthrust::device_vector morton(P);\n\tthrust::device_vector morton_sorted(P);\n\tcoord2Morton << <(P + 255) / 256, 256 >> > (P, points, minn, maxx, morton.data().get());\n\n\tthrust::device_vector indices(P);\n\tthrust::sequence(indices.begin(), indices.end());\n\tthrust::device_vector indices_sorted(P);\n\n\tcub::DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes, morton.data().get(), morton_sorted.data().get(), indices.data().get(), indices_sorted.data().get(), P);\n\ttemp_storage.resize(temp_storage_bytes);\n\n\tcub::DeviceRadixSort::SortPairs(temp_storage.data().get(), temp_storage_bytes, morton.data().get(), morton_sorted.data().get(), indices.data().get(), indices_sorted.data().get(), P);\n\n\tuint32_t num_boxes = (P + BOX_SIZE - 1) / BOX_SIZE;\n\tthrust::device_vector boxes(num_boxes);\n\tboxMinMax << > > (P, points, indices_sorted.data().get(), boxes.data().get());\n\tboxMeanDist << > > (P, points, indices_sorted.data().get(), boxes.data().get(), meanDists);\n\n\tcudaFree(result);\n}" }, { "path": "submodules/simple-knn/simple_knn.h", "content": "/*\n * Copyright (C) 2023, Inria\n * GRAPHDECO research group, https://team.inria.fr/graphdeco\n * All rights reserved.\n *\n * This software is free for non-commercial, research and evaluation use \n * under the terms of the LICENSE.md file.\n *\n * For inquiries contact george.drettakis@inria.fr\n */\n\n#ifndef SIMPLEKNN_H_INCLUDED\n#define SIMPLEKNN_H_INCLUDED\n\nclass SimpleKNN\n{\npublic:\n\tstatic void knn(int P, float3* points, float* meanDists);\n};\n\n#endif" }, { "path": "submodules/simple-knn/spatial.cu", "content": "/*\n * Copyright (C) 2023, Inria\n * GRAPHDECO research group, https://team.inria.fr/graphdeco\n * All rights reserved.\n *\n * This software is free for non-commercial, research and evaluation use \n * under the terms of the LICENSE.md file.\n *\n * For inquiries contact george.drettakis@inria.fr\n */\n\n#include \"spatial.h\"\n#include \"simple_knn.h\"\n\ntorch::Tensor\ndistCUDA2(const torch::Tensor& points)\n{\n const int P = points.size(0);\n\n auto float_opts = points.options().dtype(torch::kFloat32);\n torch::Tensor means = torch::full({P}, 0.0, float_opts);\n \n SimpleKNN::knn(P, (float3*)points.contiguous().data(), means.contiguous().data());\n\n return means;\n}" }, { "path": "submodules/simple-knn/spatial.h", "content": "/*\n * Copyright (C) 2023, Inria\n * GRAPHDECO research group, https://team.inria.fr/graphdeco\n * All rights reserved.\n *\n * This software is free for non-commercial, research and evaluation use \n * under the terms of the LICENSE.md file.\n *\n * For inquiries contact george.drettakis@inria.fr\n */\n\n#include \n\ntorch::Tensor distCUDA2(const torch::Tensor& points);" }, { "path": "utils/camera_utils.py", "content": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# This software is free for non-commercial, research and evaluation use \n# under the terms of the LICENSE.md file.\n#\n# For inquiries contact george.drettakis@inria.fr\n#\n\nfrom scene.cameras import Camera\nimport numpy as np\nfrom utils.general_utils import PILtoTorch, PIL2toTorch\nfrom utils.graphics_utils import fov2focal\nimport torch\n\nWARNED = False\n\ndef loadCam(args, id, cam_info, resolution_scale):\n orig_w, orig_h = cam_info.image.size\n\n if args.resolution in [1, 2, 4, 8]:\n resolution = round(orig_w/(resolution_scale * args.resolution)), round(orig_h/(resolution_scale * args.resolution))\n else: # should be a type that converts to float\n if args.resolution == -1:\n if orig_w > 1600:\n global WARNED\n if not WARNED:\n print(\"[ INFO ] Encountered quite large input images (>1.6K pixels width), rescaling to 1.6K.\\n \"\n \"If this is not desired, please explicitly specify '--resolution/-r' as 1\")\n WARNED = True\n global_down = orig_w / 1600\n else:\n global_down = 1\n else:\n global_down = orig_w / args.resolution\n\n scale = float(global_down) * float(resolution_scale)\n resolution = (int(orig_w / scale), int(orig_h / scale))\n\n resized_image_rgb = PILtoTorch(cam_info.image, resolution)\n\n if cam_info.semantic2d is not None:\n semantic2d = torch.from_numpy(cam_info.semantic2d).long()[None, ...]\n else:\n semantic2d = None\n\n optical_image = cam_info.optical_image\n mask = cam_info.mask\n\n gt_image = resized_image_rgb[:3, ...]\n\n return Camera(colmap_id=cam_info.uid, R=cam_info.R, T=cam_info.T, K=cam_info.K,\n FoVx=cam_info.FovX, FoVy=cam_info.FovY, \n image=gt_image, image_name=cam_info.image_name, uid=id, data_device=args.data_device,\n cx_ratio=cam_info.cx_ratio, cy_ratio=cam_info.cy_ratio, semantic2d=semantic2d, mask=mask,\n timestamp=cam_info.timestamp, optical_image=optical_image, dynamics=cam_info.dynamics)\n\ndef cameraList_from_camInfos(cam_infos, resolution_scale, args):\n camera_list = []\n\n for id, c in enumerate(cam_infos):\n camera_list.append(loadCam(args, id, c, resolution_scale))\n\n return camera_list\n\ndef camera_to_JSON(id, camera : Camera):\n Rt = np.zeros((4, 4))\n Rt[:3, :3] = camera.R.transpose()\n Rt[:3, 3] = camera.T\n Rt[3, 3] = 1.0\n\n W2C = np.linalg.inv(Rt)\n pos = W2C[:3, 3]\n rot = W2C[:3, :3]\n serializable_array_2d = [x.tolist() for x in rot]\n camera_entry = {\n 'id' : id,\n 'img_name' : camera.image_name,\n 'width' : camera.width,\n 'height' : camera.height,\n 'position': pos.tolist(),\n 'rotation': serializable_array_2d,\n 'fy' : fov2focal(camera.FovY, camera.height),\n 'fx' : fov2focal(camera.FovX, camera.width),\n }\n return camera_entry\n" }, { "path": "utils/cmap.py", "content": "import numpy as np\n\n_color_map_errors = np.array([\n [149, 54, 49], #0: log2(x) = -infinity\n [180, 117, 69], #0.0625: log2(x) = -4\n [209, 173, 116], #0.125: log2(x) = -3\n [233, 217, 171], #0.25: log2(x) = -2\n [248, 243, 224], #0.5: log2(x) = -1\n [144, 224, 254], #1.0: log2(x) = 0\n [97, 174, 253], #2.0: log2(x) = 1\n [67, 109, 244], #4.0: log2(x) = 2\n [39, 48, 215], #8.0: log2(x) = 3\n [38, 0, 165], #16.0: log2(x) = 4\n [38, 0, 165] #inf: log2(x) = inf\n]).astype(float)\n\ndef color_error_image(errors, scale=1, mask=None, BGR=True):\n \"\"\"\n Color an input error map.\n \n Arguments:\n errors -- HxW numpy array of errors\n [scale=1] -- scaling the error map (color change at unit error)\n [mask=None] -- zero-pixels are masked white in the result\n [BGR=True] -- toggle between BGR and RGB\n\n Returns:\n colored_errors -- HxWx3 numpy array visualizing the errors\n \"\"\"\n \n errors_flat = errors.flatten()\n errors_color_indices = np.clip(np.log2(errors_flat / scale + 1e-5) + 5, 0, 9)\n i0 = np.floor(errors_color_indices).astype(int)\n f1 = errors_color_indices - i0.astype(float)\n colored_errors_flat = _color_map_errors[i0, :] * (1-f1).reshape(-1,1) + _color_map_errors[i0+1, :] * f1.reshape(-1,1)\n\n if mask is not None:\n colored_errors_flat[mask.flatten() == 0] = 255\n\n if not BGR:\n colored_errors_flat = colored_errors_flat[:,[2,1,0]]\n\n return colored_errors_flat.reshape(errors.shape[0], errors.shape[1], 3).astype(np.int)\n\n_color_map_depths = np.array([\n [0, 0, 0], # 0.000\n [0, 0, 255], # 0.114\n [255, 0, 0], # 0.299\n [255, 0, 255], # 0.413\n [0, 255, 0], # 0.587\n [0, 255, 255], # 0.701\n [255, 255, 0], # 0.886\n [255, 255, 255], # 1.000\n [255, 255, 255], # 1.000\n]).astype(float)\n_color_map_bincenters = np.array([\n 0.0,\n 0.114,\n 0.299,\n 0.413,\n 0.587,\n 0.701,\n 0.886,\n 1.000,\n 2.000, # doesn't make a difference, just strictly higher than 1\n])\n\ndef color_depth_map(depths, scale=None):\n \"\"\"\n Color an input depth map.\n \n Arguments:\n depths -- HxW numpy array of depths\n [scale=None] -- scaling the values (defaults to the maximum depth)\n\n Returns:\n colored_depths -- HxWx3 numpy array visualizing the depths\n \"\"\"\n\n # if scale is None:\n # scale = depths.max() / 1.5\n scale = 50\n values = np.clip(depths.flatten() / scale, 0, 1)\n # for each value, figure out where they fit in in the bincenters: what is the last bincenter smaller than this value?\n lower_bin = ((values.reshape(-1, 1) >= _color_map_bincenters.reshape(1,-1)) * np.arange(0,9)).max(axis=1)\n lower_bin_value = _color_map_bincenters[lower_bin]\n higher_bin_value = _color_map_bincenters[lower_bin + 1]\n alphas = (values - lower_bin_value) / (higher_bin_value - lower_bin_value)\n colors = _color_map_depths[lower_bin] * (1-alphas).reshape(-1,1) + _color_map_depths[lower_bin + 1] * alphas.reshape(-1,1)\n return colors.reshape(depths.shape[0], depths.shape[1], 3).astype(np.uint8)" }, { "path": "utils/dynamic_utils.py", "content": "import numpy as np\nimport torch\nfrom torch import optim\nfrom torch import nn\nfrom tqdm import tqdm\nfrom matplotlib import pyplot as plt\nimport torch.nn.functional as F\nfrom collections import defaultdict\nimport os\n\n\ndef rot2Euler(R):\n sy = torch.sqrt(R[0,0] * R[0,0] + R[1,0] * R[1,0])\n singular = sy < 1e-6\n\n if not singular:\n x = torch.atan2(R[2,1] , R[2,2])\n y = torch.atan2(-R[2,0], sy)\n z = torch.atan2(R[1,0], R[0,0])\n else:\n x = torch.atan2(-R[1,2], R[1,1])\n y = torch.atan2(-R[2,0], sy)\n z = 0\n\n return torch.stack([x,y,z])\n\nclass unicycle(torch.nn.Module):\n\n def __init__(self, train_timestamp, centers=None, heights=None, phis=None):\n super(unicycle, self).__init__()\n self.train_timestamp = train_timestamp\n self.delta = torch.diff(self.train_timestamp)\n\n self.input_a = centers[:, 0].clone()\n self.input_b = centers[:, 1].clone()\n\n if centers is None:\n self.a = nn.Parameter(torch.zeros_like(train_timestamp).float())\n self.b = nn.Parameter(torch.zeros_like(train_timestamp).float())\n else:\n self.a = nn.Parameter(centers[:, 0])\n self.b = nn.Parameter(centers[:, 1])\n \n diff_a = torch.diff(centers[:, 0]) / self.delta\n diff_b = torch.diff(centers[:, 1]) / self.delta\n v = torch.sqrt(diff_a ** 2 + diff_b**2)\n self.v = nn.Parameter(F.pad(v, (0, 1), 'constant', v[-1].item()))\n self.phi = nn.Parameter(phis)\n\n if heights is None:\n self.h = nn.Parameter(torch.zeros_like(train_timestamp).float())\n else:\n self.h = nn.Parameter(heights)\n\n def acc_omega(self):\n acc = torch.diff(self.v) / self.delta\n omega = torch.diff(self.phi) / self.delta\n acc = F.pad(acc, (0, 1), 'constant', acc[-1].item())\n omega = F.pad(omega, (0, 1), 'constant', omega[-1].item())\n return acc, omega\n\n def forward(self, timestamps):\n idx = torch.searchsorted(self.train_timestamp, timestamps, side='left')\n invalid = (idx == self.train_timestamp.shape[0])\n idx[invalid] -= 1\n idx[self.train_timestamp[idx] != timestamps] -= 1\n idx[invalid] += 1\n prev_timestamps = self.train_timestamp[idx]\n delta_t = timestamps - prev_timestamps\n prev_a, prev_b = self.a[idx], self.b[idx]\n prev_v, prev_phi = self.v[idx], self.phi[idx]\n \n acc, omega = self.acc_omega()\n v = prev_v + acc[idx] * delta_t\n phi = prev_phi + omega[idx] * delta_t\n a = prev_a + prev_v * ((torch.sin(phi) - torch.sin(prev_phi)) / (omega[idx] + 1e-6))\n b = prev_b - prev_v * ((torch.cos(phi) - torch.cos(prev_phi)) / (omega[idx] + 1e-6))\n h = self.h[idx]\n return a, b, v, phi, h\n\n def capture(self):\n return (\n self.a,\n self.b,\n self.v,\n self.phi,\n self.h,\n self.train_timestamp,\n self.delta\n )\n \n def restore(self, model_args):\n (\n self.a,\n self.b,\n self.v,\n self.phi,\n self.h,\n self.train_timestamp,\n self.delta\n ) = model_args\n\n def visualize(self, save_path, noise_centers=None, gt_centers=None):\n a, b, _, phi, _ = self.forward(self.train_timestamp)\n a = a.detach().cpu().numpy()\n b = b.detach().cpu().numpy()\n phi = phi.detach().cpu().numpy()\n plt.scatter(a, b, marker='x', color='b')\n plt.quiver(a, b, np.ones_like(a) * np.cos(phi), np.ones_like(b) * np.sin(phi), scale=20, width=0.005)\n if noise_centers is not None:\n noise_centers = noise_centers.detach().cpu().numpy()\n plt.scatter(noise_centers[:, 0], noise_centers[:, 1], marker='o', color='gray')\n if gt_centers is not None:\n gt_centers = gt_centers.detach().cpu().numpy()\n plt.scatter(gt_centers[:, 0], gt_centers[:, 1], marker='v', color='g')\n plt.axis('equal')\n plt.savefig(save_path)\n plt.close()\n\n def reg_loss(self):\n reg = 0\n acc, omega = self.acc_omega()\n reg += torch.mean(torch.abs(torch.diff(acc))) * 1\n reg += torch.mean(torch.abs(torch.diff(omega))) * 1\n reg_a_motion = self.v[:-1] * ((torch.sin(self.phi[1:]) - torch.sin(self.phi[:-1])) / (omega[:-1] + 1e-6)) \n reg_b_motion = -self.v[:-1] * ((torch.cos(self.phi[1:]) - torch.cos(self.phi[:-1])) / (omega[:-1] + 1e-6))\n reg_a = self.a[:-1] + reg_a_motion\n reg_b = self.b[:-1] + reg_b_motion\n reg += torch.mean((reg_a - self.a[1:])**2 + (reg_b - self.b[1:])**2) * 1\n return reg\n \n def pos_loss(self):\n # a, b, _, _, _ = self.forward(self.train_timestamp)\n return torch.mean((self.a - self.input_a) ** 2 + (self.b - self.input_b) ** 2) * 10\n \n\ndef create_unicycle_model(train_cams, model_path, opt_iter=0, data_type='kitti'):\n unicycle_models = {}\n if data_type == 'kitti':\n cameras = [cam for cam in train_cams if 'cam_0' in cam.image_name]\n elif data_type == 'waymo':\n cameras = [cam for cam in train_cams if 'cam_1' in cam.image_name]\n else:\n raise NotImplementedError \n\n all_centers, all_heights, all_phis, all_timestamps = defaultdict(list), defaultdict(list), defaultdict(list), defaultdict(list)\n seq_timestamps = []\n for cam in cameras:\n t = cam.timestamp\n seq_timestamps.append(t)\n for track_id, b2w in cam.dynamics.items():\n all_centers[track_id].append(b2w[[0, 2], 3])\n all_heights[track_id].append(b2w[1, 3])\n eulers = rot2Euler(b2w[:3, :3])\n all_phis[track_id].append(eulers[1])\n all_timestamps[track_id].append(t)\n\n for track_id in all_centers.keys():\n centers = torch.stack(all_centers[track_id], dim=0).cuda()\n timestamps = torch.tensor(all_timestamps[track_id]).cuda()\n heights = torch.tensor(all_heights[track_id]).cuda()\n phis = torch.tensor(all_phis[track_id]).cuda() + torch.pi\n model = unicycle(timestamps, centers.clone(), heights.clone(), phis.clone())\n l = [\n {'params': [model.a], 'lr': 1e-2, \"name\": \"a\"},\n {'params': [model.b], 'lr': 1e-2, \"name\": \"b\"},\n {'params': [model.v], 'lr': 1e-3, \"name\": \"v\"},\n {'params': [model.phi], 'lr': 1e-4, \"name\": \"phi\"},\n {'params': [model.h], 'lr': 0, \"name\": \"h\"}\n ]\n\n optimizer = optim.Adam(l, lr=0.0)\n\n t_range = tqdm(range(opt_iter), desc=f\"Fitting {track_id}\")\n for iter in t_range:\n loss = 0.2 * model.pos_loss() + model.reg_loss()\n t_range.set_postfix({'loss': loss.item()})\n optimizer.zero_grad()\n loss.backward()\n optimizer.step()\n\n unicycle_models[track_id] = {'model': model, \n 'optimizer': optimizer,\n 'input_centers': centers}\n \n os.makedirs(os.path.join(model_path, \"unicycle\"), exist_ok=True)\n for track_id, unicycle_pkg in unicycle_models.items():\n model = unicycle_pkg['model']\n optimizer = unicycle_pkg['optimizer']\n \n model.visualize(os.path.join(model_path, \"unicycle\", f\"{track_id}_init.png\"),\n # noise_centers=unicycle_pkg['input_centers']\n )\n # gt_centers=gt_centers)\n\n return unicycle_models" }, { "path": "utils/general_utils.py", "content": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# This software is free for non-commercial, research and evaluation use \n# under the terms of the LICENSE.md file.\n#\n# For inquiries contact george.drettakis@inria.fr\n#\n\nimport torch\nimport sys\nfrom datetime import datetime\nimport numpy as np\nimport random\nimport os\nimport cv2\n\ndef inverse_sigmoid(x):\n return torch.log(x/(1-x))\n\ndef PILtoTorch(pil_image, resolution):\n resized_image_PIL = pil_image.resize(resolution)\n resized_image = torch.from_numpy(np.array(resized_image_PIL)) / 255.0\n if len(resized_image.shape) == 3:\n return resized_image.permute(2, 0, 1)\n else:\n return resized_image.unsqueeze(dim=-1).permute(2, 0, 1)\n \ndef PIL2toTorch(pil_image, resolution):\n resized_image_PIL = pil_image.resize(resolution) \n resized_image = torch.from_numpy(np.array(resized_image_PIL)) / 255.0 * (2.0 ** 16 - 1.0)\n return resized_image\n\ndef decode_op(optical_png):\n # use 'PIL Image.Open' to READ \n \"Convert from .png (h, w, 3-rgb) -> (h,w,2)(flow_x, flow_y) .. float32 array\"\n optical_png = optical_png[..., [2, 1, 0]] # bgr -> rgb\n h, w, _c = optical_png.shape\n assert optical_png.dtype == np.uint16 and _c == 3\n \"invalid flow flag: b == 0 for sky or other invalid flow\"\n invalid_points = np.where(optical_png[..., 2] == 0)\n out_flow = torch.empty((h, w, 2))\n decoded = 2.0 / (2**16 - 1.0) * optical_png.astype('f4') - 1\n out_flow[..., 0] = torch.tensor(decoded[:, :, 0] * (w - 1)) # (pixel) delta_x : R\n out_flow[..., 1] = torch.tensor(decoded[:, :, 1] * (h - 1)) # delta_y : G\n out_flow[invalid_points[0], invalid_points[1], :] = 0 # B=0 for invalid flow\n return out_flow\n\ndef get_expon_lr_func(\n lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000\n):\n \"\"\"\n Copied from Plenoxels\n\n Continuous learning rate decay function. Adapted from JaxNeRF\n The returned rate is lr_init when step=0 and lr_final when step=max_steps, and\n is log-linearly interpolated elsewhere (equivalent to exponential decay).\n If lr_delay_steps>0 then the learning rate will be scaled by some smooth\n function of lr_delay_mult, such that the initial learning rate is\n lr_init*lr_delay_mult at the beginning of optimization but will be eased back\n to the normal learning rate when steps>lr_delay_steps.\n :param conf: config subtree 'lr' or similar\n :param max_steps: int, the number of steps during optimization.\n :return HoF which takes step as input\n \"\"\"\n\n def helper(step):\n if step < 0 or (lr_init == 0.0 and lr_final == 0.0):\n # Disable this parameter\n return 0.0\n if lr_delay_steps > 0:\n # A kind of reverse cosine decay.\n delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin(\n 0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1)\n )\n else:\n delay_rate = 1.0\n t = np.clip(step / max_steps, 0, 1)\n log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t)\n return delay_rate * log_lerp\n\n return helper\n\ndef strip_lowerdiag(L):\n uncertainty = torch.zeros((L.shape[0], 6), dtype=torch.float, device=\"cuda\")\n\n uncertainty[:, 0] = L[:, 0, 0]\n uncertainty[:, 1] = L[:, 0, 1]\n uncertainty[:, 2] = L[:, 0, 2]\n uncertainty[:, 3] = L[:, 1, 1]\n uncertainty[:, 4] = L[:, 1, 2]\n uncertainty[:, 5] = L[:, 2, 2]\n return uncertainty\n\ndef strip_symmetric(sym):\n return strip_lowerdiag(sym)\n\ndef build_rotation(r):\n norm = torch.sqrt(r[:,0]*r[:,0] + r[:,1]*r[:,1] + r[:,2]*r[:,2] + r[:,3]*r[:,3])\n\n q = r / norm[:, None]\n\n R = torch.zeros((q.size(0), 3, 3), device='cuda')\n\n r = q[:, 0]\n x = q[:, 1]\n y = q[:, 2]\n z = q[:, 3]\n\n R[:, 0, 0] = 1 - 2 * (y*y + z*z)\n R[:, 0, 1] = 2 * (x*y - r*z)\n R[:, 0, 2] = 2 * (x*z + r*y)\n R[:, 1, 0] = 2 * (x*y + r*z)\n R[:, 1, 1] = 1 - 2 * (x*x + z*z)\n R[:, 1, 2] = 2 * (y*z - r*x)\n R[:, 2, 0] = 2 * (x*z - r*y)\n R[:, 2, 1] = 2 * (y*z + r*x)\n R[:, 2, 2] = 1 - 2 * (x*x + y*y)\n return R\n\ndef build_scaling_rotation(s, r):\n L = torch.zeros((s.shape[0], 3, 3), dtype=torch.float, device=\"cuda\")\n R = build_rotation(r)\n\n L[:,0,0] = s[:,0]\n L[:,1,1] = s[:,1]\n L[:,2,2] = s[:,2]\n\n L = R @ L\n return L\n\nDEFAULT_RANDOM_SEED = 0\n\ndef seedBasic(seed=DEFAULT_RANDOM_SEED):\n random.seed(seed)\n os.environ['PYTHONHASHSEED'] = str(seed)\n np.random.seed(seed)\n\ndef seedTorch(seed=DEFAULT_RANDOM_SEED):\n torch.manual_seed(seed)\n torch.cuda.manual_seed(seed)\n torch.backends.cudnn.deterministic = True\n torch.backends.cudnn.benchmark = False\n \n# basic + tensorflow + torch \ndef seedEverything(seed=DEFAULT_RANDOM_SEED):\n seedBasic(seed)\n seedTorch(seed)\n\ndef safe_state(silent):\n old_f = sys.stdout\n class F:\n def __init__(self, silent):\n self.silent = silent\n\n def write(self, x):\n if not self.silent:\n if x.endswith(\"\\n\"):\n old_f.write(x.replace(\"\\n\", \" [{}]\\n\".format(str(datetime.now().strftime(\"%d/%m %H:%M:%S\")))))\n else:\n old_f.write(x)\n\n def flush(self):\n old_f.flush()\n\n sys.stdout = F(silent)\n\n random.seed(DEFAULT_RANDOM_SEED)\n np.random.seed(DEFAULT_RANDOM_SEED)\n torch.manual_seed(DEFAULT_RANDOM_SEED)\n torch.cuda.set_device(torch.device(\"cuda:0\"))\n # sys.stdout = old_f\n" }, { "path": "utils/graphics_utils.py", "content": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# This software is free for non-commercial, research and evaluation use \n# under the terms of the LICENSE.md file.\n#\n# For inquiries contact george.drettakis@inria.fr\n#\n\nimport torch\nimport math\nimport numpy as np\nfrom typing import NamedTuple\n\nclass BasicPointCloud(NamedTuple):\n points : np.array\n colors : np.array\n normals : np.array\n # feats3D : np.array\n\ndef geom_transform_points(points, transf_matrix):\n P, _ = points.shape\n ones = torch.ones(P, 1, dtype=points.dtype, device=points.device)\n points_hom = torch.cat([points, ones], dim=1)\n points_out = torch.matmul(points_hom, transf_matrix.unsqueeze(0))\n\n denom = points_out[..., 3:] + 0.0000001\n return (points_out[..., :3] / denom).squeeze(dim=0)\n\ndef getWorld2View(R, t):\n Rt = np.zeros((4, 4))\n Rt[:3, :3] = R.transpose()\n Rt[:3, 3] = t\n Rt[3, 3] = 1.0\n return np.float32(Rt)\n\ndef getWorld2View2(R, t, translate=np.array([.0, .0, .0]), scale=1.0):\n Rt = np.zeros((4, 4))\n Rt[:3, :3] = R.transpose()\n Rt[:3, 3] = t\n Rt[3, 3] = 1.0\n\n C2W = np.linalg.inv(Rt)\n cam_center = C2W[:3, 3]\n cam_center = (cam_center + translate) * scale\n C2W[:3, 3] = cam_center\n Rt = np.linalg.inv(C2W)\n return np.float32(Rt)\n\ndef getProjectionMatrix(znear, zfar, fovX, fovY, cx_ratio, cy_ratio):\n tanHalfFovY = math.tan((fovY / 2))\n tanHalfFovX = math.tan((fovX / 2))\n\n top = tanHalfFovY * znear\n bottom = -top\n right = tanHalfFovX * znear\n left = -right\n\n P = torch.zeros(4, 4)\n\n z_sign = 1.0\n\n P[0, 0] = 2.0 * znear / (right - left)\n P[1, 1] = 2.0 * znear / (top - bottom)\n P[0, 2] = (right + left) / (right - left) - 1 + cx_ratio\n P[1, 2] = (top + bottom) / (top - bottom) - 1 + cy_ratio\n P[3, 2] = z_sign\n P[2, 2] = z_sign * (zfar + znear) / (zfar - znear)\n P[2, 3] = -(2 * zfar * znear) / (zfar - znear)\n\n # P[0, 0] = 2.0 * znear / (right - left)\n # P[1, 1] = 2.0 * znear / (top - bottom)\n # P[0, 2] = (right + left) / (right - left)\n # P[1, 2] = (top + bottom) / (top - bottom)\n # P[3, 2] = z_sign\n # P[2, 2] = z_sign * zfar / (zfar - znear)\n # P[2, 3] = -(zfar * znear) / (zfar - znear)\n return P\n\ndef fov2focal(fov, pixels):\n return pixels / (2 * math.tan(fov / 2))\n\ndef focal2fov(focal, pixels):\n return 2*math.atan(pixels/(2*focal))" }, { "path": "utils/image_utils.py", "content": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# This software is free for non-commercial, research and evaluation use \n# under the terms of the LICENSE.md file.\n#\n# For inquiries contact george.drettakis@inria.fr\n#\n\nimport torch\n\ndef mse(img1, img2):\n return (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True)\n\ndef psnr(img1, img2):\n mse = (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True)\n return 20 * torch.log10(1.0 / torch.sqrt(mse))\n" }, { "path": "utils/iou_utils.py", "content": "# 3D IoU caculate code for 3D object detection \n# Kent 2018/12\n\nimport numpy as np\nfrom scipy.spatial import ConvexHull\nfrom numpy import *\n\ndef polygon_clip(subjectPolygon, clipPolygon):\n \"\"\" Clip a polygon with another polygon.\n\n Ref: https://rosettacode.org/wiki/Sutherland-Hodgman_polygon_clipping#Python\n\n Args:\n subjectPolygon: a list of (x,y) 2d points, any polygon.\n clipPolygon: a list of (x,y) 2d points, has to be *convex*\n Note:\n **points have to be counter-clockwise ordered**\n\n Return:\n a list of (x,y) vertex point for the intersection polygon.\n \"\"\"\n def inside(p):\n return(cp2[0]-cp1[0])*(p[1]-cp1[1]) > (cp2[1]-cp1[1])*(p[0]-cp1[0])\n \n def computeIntersection():\n dc = [ cp1[0] - cp2[0], cp1[1] - cp2[1] ]\n dp = [ s[0] - e[0], s[1] - e[1] ]\n n1 = cp1[0] * cp2[1] - cp1[1] * cp2[0]\n n2 = s[0] * e[1] - s[1] * e[0] \n n3 = 1.0 / (dc[0] * dp[1] - dc[1] * dp[0])\n return [(n1*dp[0] - n2*dc[0]) * n3, (n1*dp[1] - n2*dc[1]) * n3]\n \n outputList = subjectPolygon\n cp1 = clipPolygon[-1]\n \n for clipVertex in clipPolygon:\n cp2 = clipVertex\n inputList = outputList\n outputList = []\n s = inputList[-1]\n \n for subjectVertex in inputList:\n e = subjectVertex\n if inside(e):\n if not inside(s):\n outputList.append(computeIntersection())\n outputList.append(e)\n elif inside(s):\n outputList.append(computeIntersection())\n s = e\n cp1 = cp2\n if len(outputList) == 0:\n return None\n return(outputList)\n\ndef poly_area(x,y):\n \"\"\" Ref: http://stackoverflow.com/questions/24467972/calculate-area-of-polygon-given-x-y-coordinates \"\"\"\n return 0.5*np.abs(np.dot(x,np.roll(y,1))-np.dot(y,np.roll(x,1)))\n\ndef convex_hull_intersection(p1, p2):\n \"\"\" Compute area of two convex hull's intersection area.\n p1,p2 are a list of (x,y) tuples of hull vertices.\n return a list of (x,y) for the intersection and its volume\n \"\"\"\n inter_p = polygon_clip(p1,p2)\n if inter_p is not None:\n hull_inter = ConvexHull(inter_p)\n return inter_p, hull_inter.volume\n else:\n return None, 0.0 \n\ndef box3d_vol(corners):\n ''' corners: (8,3) no assumption on axis direction '''\n a = np.sqrt(np.sum((corners[0,:] - corners[1,:])**2))\n b = np.sqrt(np.sum((corners[1,:] - corners[2,:])**2))\n c = np.sqrt(np.sum((corners[0,:] - corners[4,:])**2))\n return a*b*c\n\ndef is_clockwise(p):\n x = p[:,0]\n y = p[:,1]\n return np.dot(x,np.roll(y,1))-np.dot(y,np.roll(x,1)) > 0\n\ndef box3d_iou(corners1, corners2):\n ''' Compute 3D bounding box IoU.\n\n Input:\n corners1: numpy array (8,3), assume up direction is negative Y\n corners2: numpy array (8,3), assume up direction is negative Y\n Output:\n iou: 3D bounding box IoU\n iou_2d: bird's eye view 2D bounding box IoU\n\n todo (kent): add more description on corner points' orders.\n '''\n # corner points are in counter clockwise order\n rect1 = [(corners1[i,0], corners1[i,2]) for i in [4,5,1,0]]\n rect2 = [(corners2[i,0], corners2[i,2]) for i in [4,5,1,0]] \n \n area1 = poly_area(np.array(rect1)[:,0], np.array(rect1)[:,1])\n area2 = poly_area(np.array(rect2)[:,0], np.array(rect2)[:,1])\n \n inter, inter_area = convex_hull_intersection(rect1, rect2)\n iou_2d = inter_area/(area1+area2-inter_area)\n # if iou_2d < 0:\n # print(inter_area, area1, area2)\n # ymax = min(corners1[0,1], corners2[0,1])\n # ymin = max(corners1[4,1], corners2[4,1])\n\n # inter_vol = inter_area * max(0.0, ymax-ymin)\n \n # vol1 = box3d_vol(corners1)\n # vol2 = box3d_vol(corners2)\n # iou = inter_vol / (vol1 + vol2 - inter_vol)\n # return iou, iou_2d\n return 0, iou_2d\n\n# ----------------------------------\n# Helper functions for evaluation\n# ----------------------------------\n\ndef get_3d_box(box_size, heading_angle, center):\n ''' Calculate 3D bounding box corners from its parameterization.\n\n Input:\n box_size: tuple of (length,wide,height)\n heading_angle: rad scalar, clockwise from pos x axis\n center: tuple of (x,y,z)\n Output:\n corners_3d: numpy array of shape (8,3) for 3D box cornders\n '''\n def roty(t):\n c = np.cos(t)\n s = np.sin(t)\n return np.array([[c, 0, s],\n [0, 1, 0],\n [-s, 0, c]])\n\n R = roty(heading_angle)\n l,w,h = box_size\n x_corners = [l/2,l/2,-l/2,-l/2,l/2,l/2,-l/2,-l/2];\n y_corners = [h/2,h/2,h/2,h/2,-h/2,-h/2,-h/2,-h/2];\n z_corners = [w/2,-w/2,-w/2,w/2,w/2,-w/2,-w/2,w/2];\n corners_3d = np.dot(R, np.vstack([x_corners,y_corners,z_corners]))\n corners_3d[0,:] = corners_3d[0,:] + center[0];\n corners_3d[1,:] = corners_3d[1,:] + center[1];\n corners_3d[2,:] = corners_3d[2,:] + center[2];\n corners_3d = np.transpose(corners_3d)\n return corners_3d\n\n \nif __name__=='__main__':\n print('------------------')\n # get_3d_box(box_size, heading_angle, center)\n corners_3d_ground = get_3d_box((1.497255,1.644981, 3.628938), -1.531692, (2.882992 ,1.698800 ,20.785644)) \n corners_3d_predict = get_3d_box((1.458242, 1.604773, 3.707947), -1.549553, (2.756923, 1.661275, 20.943280 ))\n (IOU_3d,IOU_2d)=box3d_iou(corners_3d_predict,corners_3d_ground)\n print (IOU_3d,IOU_2d) #3d IoU/ 2d IoU of BEV(bird eye's view)\n " }, { "path": "utils/loss_utils.py", "content": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# This software is free for non-commercial, research and evaluation use \n# under the terms of the LICENSE.md file.\n#\n# For inquiries contact george.drettakis@inria.fr\n#\n\nimport torch\nimport torch.nn.functional as F\nfrom torch.autograd import Variable\nfrom math import exp\n\ndef l1_loss(network_output, gt, mask=None):\n l1 = torch.abs((network_output - gt))\n if mask is not None:\n l1 = l1[:, mask]\n return l1.mean()\n\ndef l2_loss(network_output, gt):\n return ((network_output - gt) ** 2).mean()\n\ndef gaussian(window_size, sigma):\n gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)])\n return gauss / gauss.sum()\n\ndef create_window(window_size, channel):\n _1D_window = gaussian(window_size, 1.5).unsqueeze(1)\n _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)\n window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())\n return window\n\ndef ssim(img1, img2, window_size=11, size_average=True):\n channel = img1.size(-3)\n window = create_window(window_size, channel)\n\n if img1.is_cuda:\n window = window.cuda(img1.get_device())\n window = window.type_as(img1)\n\n return _ssim(img1, img2, window, window_size, channel, size_average)\n\ndef _ssim(img1, img2, window, window_size, channel, size_average=True):\n mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel)\n mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel)\n\n mu1_sq = mu1.pow(2)\n mu2_sq = mu2.pow(2)\n mu1_mu2 = mu1 * mu2\n\n sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq\n sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq\n sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2\n\n C1 = 0.01 ** 2\n C2 = 0.03 ** 2\n\n ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))\n\n if size_average:\n return ssim_map.mean()\n else:\n return ssim_map.mean(1).mean(1).mean(1)\n\ndef ssim_loss(img1, img2, window_size=11, size_average=True, mask=None):\n channel = img1.size(-3)\n window = create_window(window_size, channel)\n\n if img1.is_cuda:\n window = window.cuda(img1.get_device())\n window = window.type_as(img1)\n\n return _ssim_loss(img1, img2, window, window_size, channel, size_average, mask)\n\ndef _ssim_loss(img1, img2, window, window_size, channel, size_average=True, mask=None):\n mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel)\n mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel)\n\n mu1_sq = mu1.pow(2)\n mu2_sq = mu2.pow(2)\n mu1_mu2 = mu1 * mu2\n\n sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq\n sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq\n sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2\n\n C1 = 0.01 ** 2\n C2 = 0.03 ** 2\n\n ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))\n ssim_map = 1 - ssim_map\n\n if mask is not None:\n ssim_map = ssim_map[:, mask]\n if size_average:\n return ssim_map.mean()\n else:\n return ssim_map.mean(1).mean(1).mean(1)" }, { "path": "utils/nvseg_utils.py", "content": "import sys\nsys.path.append(\"/data0/hyzhou/workspace/nv_seg\")\nfrom network import get_model\nfrom config import cfg, torch_version_float\nfrom datasets.cityscapes import Loader as dataset_cls\nfrom runx.logx import logx\nimport cv2\nimport torch\nfrom imageio.v2 import imread, imwrite\nimport os\nimport numpy as np\nfrom glob import glob\nfrom tqdm import tqdm\nfrom torchvision.utils import save_image\n\ndef restore_net(net, checkpoint):\n assert 'state_dict' in checkpoint, 'cant find state_dict in checkpoint'\n forgiving_state_restore(net, checkpoint['state_dict'])\n\n\ndef forgiving_state_restore(net, loaded_dict):\n \"\"\"\n Handle partial loading when some tensors don't match up in size.\n Because we want to use models that were trained off a different\n number of classes.\n \"\"\"\n\n net_state_dict = net.state_dict()\n new_loaded_dict = {}\n for k in net_state_dict:\n new_k = k\n if new_k in loaded_dict and net_state_dict[k].size() == loaded_dict[new_k].size():\n new_loaded_dict[k] = loaded_dict[new_k]\n else: \n logx.msg(\"Skipped loading parameter {}\".format(k))\n net_state_dict.update(new_loaded_dict)\n net.load_state_dict(net_state_dict)\n return net\n\ndef get_nvseg_model():\n logx.initialize(logdir=\"./results\",\n global_rank=0)\n\n cfg.immutable(False)\n cfg.DATASET.NUM_CLASSES = dataset_cls.num_classes\n cfg.DATASET.IGNORE_LABEL = dataset_cls.ignore_label\n cfg.MODEL.MSCALE = True\n cfg.MODEL.N_SCALES = [0.5,1.0,2.0]\n cfg.MODEL.BNFUNC = torch.nn.BatchNorm2d\n cfg.OPTIONS.TORCH_VERSION = torch_version_float()\n cfg.DATASET_INST = dataset_cls('folder')\n cfg.immutable(True)\n colorize_mask_fn = cfg.DATASET_INST.colorize_mask\n\n net = get_model(network='network.ocrnet.HRNet_Mscale',\n num_classes=cfg.DATASET.NUM_CLASSES,\n criterion=None)\n\n snapshot = \"ASSETS_PATH/seg_weights/cityscapes_trainval_ocr.HRNet_Mscale_nimble-chihuahua.pth\".replace('ASSETS_PATH', cfg.ASSETS_PATH)\n checkpoint = torch.load(snapshot, map_location=torch.device('cpu'))\n renamed_ckpt = {'state_dict': {}}\n for k, v in checkpoint['state_dict'].items():\n renamed_ckpt['state_dict'][k.replace('module.', '')] = v\n restore_net(net, renamed_ckpt)\n net = net.eval().cuda()\n return net" }, { "path": "utils/semantic_utils.py", "content": "#!/usr/bin/python\n#\n# KITTI-360 labels\n#\n\nfrom collections import namedtuple\nfrom PIL import Image\nimport numpy as np\n\n\n#--------------------------------------------------------------------------------\n# Definitions\n#--------------------------------------------------------------------------------\n\n# a label and all meta information\nLabel = namedtuple( 'Label' , [\n\n 'name' , # The identifier of this label, e.g. 'car', 'person', ... .\n # We use them to uniquely name a class\n\n 'id' , # An integer ID that is associated with this label.\n # The IDs are used to represent the label in ground truth images\n # An ID of -1 means that this label does not have an ID and thus\n # is ignored when creating ground truth images (e.g. license plate).\n # Do not modify these IDs, since exactly these IDs are expected by the\n # evaluation server.\n\n 'trainId' , # Feel free to modify these IDs as suitable for your method. Then create\n # ground truth images with train IDs, using the tools provided in the\n # 'preparation' folder. However, make sure to validate or submit results\n # to our evaluation server using the regular IDs above!\n # For trainIds, multiple labels might have the same ID. Then, these labels\n # are mapped to the same class in the ground truth images. For the inverse\n # mapping, we use the label that is defined first in the list below.\n # For example, mapping all void-type classes to the same ID in training,\n # might make sense for some approaches.\n # Max value is 255!\n\n 'category' , # The name of the category that this label belongs to\n\n 'categoryId' , # The ID of this category. Used to create ground truth images\n # on category level.\n\n 'hasInstances', # Whether this label distinguishes between single instances or not\n\n 'ignoreInEval', # Whether pixels having this class as ground truth label are ignored\n # during evaluations or not\n\n 'color' , # The color of this label\n ] )\n\n\n#--------------------------------------------------------------------------------\n# A list of all labels\n#--------------------------------------------------------------------------------\n\n# Please adapt the train IDs as appropriate for your approach.\n# Note that you might want to ignore labels with ID 255 during training.\n# Further note that the current train IDs are only a suggestion. You can use whatever you like.\n# Make sure to provide your results using the original IDs and not the training IDs.\n# Note that many IDs are ignored in evaluation and thus you never need to predict these!\n\nlabels = [\n # name id trainId category catId hasInstances ignoreInEval color\n Label( 'unlabeled' , 0 , 255 , 'void' , 0 , False , True , ( 0, 0, 0) ),\n Label( 'ego vehicle' , 1 , 255 , 'void' , 0 , False , True , ( 0, 0, 0) ),\n Label( 'rectification border' , 2 , 255 , 'void' , 0 , False , True , ( 0, 0, 0) ),\n Label( 'out of roi' , 3 , 255 , 'void' , 0 , False , True , ( 0, 0, 0) ),\n Label( 'static' , 4 , 255 , 'void' , 0 , False , True , ( 0, 0, 0) ),\n Label( 'dynamic' , 5 , 255 , 'void' , 0 , False , True , (111, 74, 0) ),\n Label( 'ground' , 6 , 255 , 'void' , 0 , False , True , ( 81, 0, 81) ),\n Label( 'road' , 7 , 0 , 'flat' , 1 , False , False , (128, 64,128) ),\n Label( 'sidewalk' , 8 , 1 , 'flat' , 1 , False , False , (244, 35,232) ),\n Label( 'parking' , 9 , 255 , 'flat' , 1 , False , True , (250,170,160) ),\n Label( 'rail track' , 10 , 255 , 'flat' , 1 , False , True , (230,150,140) ),\n Label( 'building' , 11 , 2 , 'construction' , 2 , False , False , ( 70, 70, 70) ),\n Label( 'wall' , 12 , 3 , 'construction' , 2 , False , False , (102,102,156) ),\n Label( 'fence' , 13 , 4 , 'construction' , 2 , False , False , (190,153,153) ),\n Label( 'guard rail' , 14 , 255 , 'construction' , 2 , False , True , (180,165,180) ),\n Label( 'bridge' , 15 , 255 , 'construction' , 2 , False , True , (150,100,100) ),\n Label( 'tunnel' , 16 , 255 , 'construction' , 2 , False , True , (150,120, 90) ),\n Label( 'pole' , 17 , 5 , 'object' , 3 , False , False , (153,153,153) ),\n Label( 'polegroup' , 18 , 255 , 'object' , 3 , False , True , (153,153,153) ),\n Label( 'traffic light' , 19 , 6 , 'object' , 3 , False , False , (250,170, 30) ),\n Label( 'traffic sign' , 20 , 7 , 'object' , 3 , False , False , (220,220, 0) ),\n Label( 'vegetation' , 21 , 8 , 'nature' , 4 , False , False , (107,142, 35) ),\n Label( 'terrain' , 22 , 9 , 'nature' , 4 , False , False , (152,251,152) ),\n Label( 'sky' , 23 , 10 , 'sky' , 5 , False , False , ( 70,130,180) ),\n Label( 'person' , 24 , 11 , 'human' , 6 , True , False , (220, 20, 60) ),\n Label( 'rider' , 25 , 12 , 'human' , 6 , True , False , (255, 0, 0) ),\n Label( 'car' , 26 , 13 , 'vehicle' , 7 , True , False , ( 0, 0,142) ),\n Label( 'truck' , 27 , 14 , 'vehicle' , 7 , True , False , ( 0, 0, 70) ),\n Label( 'bus' , 28 , 15 , 'vehicle' , 7 , True , False , ( 0, 60,100) ),\n Label( 'caravan' , 29 , 255 , 'vehicle' , 7 , True , True , ( 0, 0, 90) ),\n Label( 'trailer' , 30 , 255 , 'vehicle' , 7 , True , True , ( 0, 0,110) ),\n Label( 'train' , 31 , 16 , 'vehicle' , 7 , True , False , ( 0, 80,100) ),\n Label( 'motorcycle' , 32 , 17 , 'vehicle' , 7 , True , False , ( 0, 0,230) ),\n Label( 'bicycle' , 33 , 18 , 'vehicle' , 7 , True , False , (119, 11, 32) ),\n Label( 'license plate' , -1 , -1 , 'vehicle' , 7 , False , True , ( 0, 0,142) ),\n]\n\n\n#--------------------------------------------------------------------------------\n# Create dictionaries for a fast lookup\n#--------------------------------------------------------------------------------\n\n# Please refer to the main method below for example usages!\n\n# name to label object\nname2label = { label.name : label for label in labels }\n# id to label object\nid2label = { label.id : label for label in labels }\n# trainId to label object\ntrainId2label = { label.trainId : label for label in reversed(labels) }\n# label2trainid\nlabel2trainid = { label.id : label.trainId for label in labels }\n# trainId to label object\ntrainId2name = { label.trainId : label.name for label in labels }\ntrainId2color = { label.trainId : label.color for label in labels }\n# category to list of label objects\ncategory2labels = {}\nfor label in labels:\n category = label.category\n if category in category2labels:\n category2labels[category].append(label)\n else:\n category2labels[category] = [label]\n\n#--------------------------------------------------------------------------------\n# color mapping\n#--------------------------------------------------------------------------------\n\npalette = [128, 64, 128,\n 244, 35, 232,\n 70, 70, 70,\n 102, 102, 156,\n 190, 153, 153,\n 153, 153, 153,\n 250, 170, 30,\n 220, 220, 0,\n 107, 142, 35,\n 152, 251, 152,\n 70, 130, 180,\n 220, 20, 60,\n 255, 0, 0,\n 0, 0, 142,\n 0, 0, 70,\n 0, 60, 100,\n 0, 80, 100,\n 0, 0, 230,\n 119, 11, 32]\nzero_pad = 256 * 3 - len(palette)\nfor i in range(zero_pad):\n palette.append(0)\ncolor_mapping = palette\n\ndef colorize(image_array):\n new_mask = Image.fromarray(image_array.astype(np.uint8)).convert('P')\n new_mask.putpalette(color_mapping)\n return new_mask\n\n#--------------------------------------------------------------------------------\n# Assure single instance name\n#--------------------------------------------------------------------------------\n\n# returns the label name that describes a single instance (if possible)\n# e.g. input | output\n# ----------------------\n# car | car\n# cargroup | car\n# foo | None\n# foogroup | None\n# skygroup | None\ndef assureSingleInstanceName( name ):\n # if the name is known, it is not a group\n if name in name2label:\n return name\n # test if the name actually denotes a group\n if not name.endswith(\"group\"):\n return None\n # remove group\n name = name[:-len(\"group\")]\n # test if the new name exists\n if not name in name2label:\n return None\n # test if the new name denotes a label that actually has instances\n if not name2label[name].hasInstances:\n return None\n # all good then\n return name\n\n#--------------------------------------------------------------------------------\n# Main for testing\n#--------------------------------------------------------------------------------\n\n# just a dummy main\nif __name__ == \"__main__\":\n # Print all the labels\n print(\"List of KITTI-360 labels:\")\n print(\"\")\n print(\" {:>21} | {:>3} | {:>7} | {:>14} | {:>10} | {:>12} | {:>12}\".format( 'name', 'id', 'trainId', 'category', 'categoryId', 'hasInstances', 'ignoreInEval' ))\n print(\" \" + ('-' * 98))\n for label in labels:\n # print(\" {:>21} | {:>3} | {:>7} | {:>14} | {:>10} | {:>12} | {:>12}\".format( label.name, label.id, label.trainId, label.category, label.categoryId, label.hasInstances, label.ignoreInEval ))\n print(\" \\\"{:}\\\"\".format(label.name))\n print(\"\")\n\n print(\"Example usages:\")\n\n # Map from name to label\n name = 'car'\n id = name2label[name].id\n print(\"ID of label '{name}': {id}\".format( name=name, id=id ))\n\n # Map from ID to label\n category = id2label[id].category\n print(\"Category of label with ID '{id}': {category}\".format( id=id, category=category ))\n\n # Map from trainID to label\n trainId = 0\n name = trainId2label[trainId].name\n print(\"Name of label with trainID '{id}': {name}\".format( id=trainId, name=name ))" }, { "path": "utils/sh_utils.py", "content": "# Copyright 2021 The PlenOctree Authors.\n# Redistribution and use in source and binary forms, with or without\n# modification, are permitted provided that the following conditions are met:\n#\n# 1. Redistributions of source code must retain the above copyright notice,\n# this list of conditions and the following disclaimer.\n#\n# 2. Redistributions in binary form must reproduce the above copyright notice,\n# this list of conditions and the following disclaimer in the documentation\n# and/or other materials provided with the distribution.\n#\n# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE\n# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE\n# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR\n# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF\n# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS\n# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN\n# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)\n# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE\n# POSSIBILITY OF SUCH DAMAGE.\n\nimport torch\n\nC0 = 0.28209479177387814\nC1 = 0.4886025119029199\nC2 = [\n 1.0925484305920792,\n -1.0925484305920792,\n 0.31539156525252005,\n -1.0925484305920792,\n 0.5462742152960396\n]\nC3 = [\n -0.5900435899266435,\n 2.890611442640554,\n -0.4570457994644658,\n 0.3731763325901154,\n -0.4570457994644658,\n 1.445305721320277,\n -0.5900435899266435\n]\nC4 = [\n 2.5033429417967046,\n -1.7701307697799304,\n 0.9461746957575601,\n -0.6690465435572892,\n 0.10578554691520431,\n -0.6690465435572892,\n 0.47308734787878004,\n -1.7701307697799304,\n 0.6258357354491761,\n] \n\n\ndef eval_sh(deg, sh, dirs):\n \"\"\"\n Evaluate spherical harmonics at unit directions\n using hardcoded SH polynomials.\n Works with torch/np/jnp.\n ... Can be 0 or more batch dimensions.\n Args:\n deg: int SH deg. Currently, 0-3 supported\n sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2]\n dirs: jnp.ndarray unit directions [..., 3]\n Returns:\n [..., C]\n \"\"\"\n assert deg <= 4 and deg >= 0\n coeff = (deg + 1) ** 2\n assert sh.shape[-1] >= coeff\n\n result = C0 * sh[..., 0]\n if deg > 0:\n x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3]\n result = (result -\n C1 * y * sh[..., 1] +\n C1 * z * sh[..., 2] -\n C1 * x * sh[..., 3])\n\n if deg > 1:\n xx, yy, zz = x * x, y * y, z * z\n xy, yz, xz = x * y, y * z, x * z\n result = (result +\n C2[0] * xy * sh[..., 4] +\n C2[1] * yz * sh[..., 5] +\n C2[2] * (2.0 * zz - xx - yy) * sh[..., 6] +\n C2[3] * xz * sh[..., 7] +\n C2[4] * (xx - yy) * sh[..., 8])\n\n if deg > 2:\n result = (result +\n C3[0] * y * (3 * xx - yy) * sh[..., 9] +\n C3[1] * xy * z * sh[..., 10] +\n C3[2] * y * (4 * zz - xx - yy)* sh[..., 11] +\n C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12] +\n C3[4] * x * (4 * zz - xx - yy) * sh[..., 13] +\n C3[5] * z * (xx - yy) * sh[..., 14] +\n C3[6] * x * (xx - 3 * yy) * sh[..., 15])\n\n if deg > 3:\n result = (result + C4[0] * xy * (xx - yy) * sh[..., 16] +\n C4[1] * yz * (3 * xx - yy) * sh[..., 17] +\n C4[2] * xy * (7 * zz - 1) * sh[..., 18] +\n C4[3] * yz * (7 * zz - 3) * sh[..., 19] +\n C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20] +\n C4[5] * xz * (7 * zz - 3) * sh[..., 21] +\n C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22] +\n C4[7] * xz * (xx - 3 * yy) * sh[..., 23] +\n C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) * sh[..., 24])\n return result\n\ndef RGB2SH(rgb):\n return (rgb - 0.5) / C0\n\ndef SH2RGB(sh):\n return sh * C0 + 0.5" }, { "path": "utils/system_utils.py", "content": "#\n# Copyright (C) 2023, Inria\n# GRAPHDECO research group, https://team.inria.fr/graphdeco\n# All rights reserved.\n#\n# This software is free for non-commercial, research and evaluation use \n# under the terms of the LICENSE.md file.\n#\n# For inquiries contact george.drettakis@inria.fr\n#\n\nfrom errno import EEXIST\nfrom os import makedirs, path\nimport os\n\ndef mkdir_p(folder_path):\n # Creates a directory. equivalent to using mkdir -p on the command line\n try:\n makedirs(folder_path)\n except OSError as exc: # Python >2.5\n if exc.errno == EEXIST and path.isdir(folder_path):\n pass\n else:\n raise\n\ndef searchForMaxIteration(folder):\n saved_iters = [int(fname.split(\"_\")[-1]) for fname in os.listdir(folder)]\n return max(saved_iters)" } ]