{"id":20144936,"url":"https://github.com/blurgyy/jaxngp","last_synced_at":"2025-04-09T19:17:58.752Z","repository":{"id":171185497,"uuid":"619811733","full_name":"blurgyy/jaxngp","owner":"blurgyy","description":"JAX implementation of instant-ngp (NeRF part)","archived":false,"fork":false,"pushed_at":"2024-05-25T09:49:10.000Z","size":927,"stargazers_count":32,"open_issues_count":2,"forks_count":3,"subscribers_count":4,"default_branch":"main","last_synced_at":"2025-04-09T19:17:53.052Z","etag":null,"topics":["cuda","hashgrid","instant-ngp","jax","nerf","neural-radiance-field","nix","python","wsl"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/blurgyy.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":"CITATION.cff","codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2023-03-27T13:33:49.000Z","updated_at":"2024-12-03T12:04:08.000Z","dependencies_parsed_at":null,"dependency_job_id":"2f3aed0e-185e-469c-8588-692f0875f508","html_url":"https://github.com/blurgyy/jaxngp","commit_stats":null,"previous_names":["blurgyy/jaxngp"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/blurgyy%2Fjaxngp","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/blurgyy%2Fjaxngp/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/blurgyy%2Fjaxngp/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/blurgyy%2Fjaxngp/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/blurgyy","download_url":"https://codeload.github.com/blurgyy/jaxngp/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248094989,"owners_count":21046770,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["cuda","hashgrid","instant-ngp","jax","nerf","neural-radiance-field","nix","python","wsl"],"created_at":"2024-11-13T22:12:50.708Z","updated_at":"2025-04-09T19:17:58.723Z","avatar_url":"https://github.com/blurgyy.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"jaxngp\n======\n\nThis repository contains [JAX] implementations of:\n* a multiresolution hash encoder (JAX)\n* an accelerated volume renderer for fast training of NeRFs (CUDA + JAX), with\n  * occupancy grid pruning during ray marching\n  * early stop during ray color integration\n* an inference-time renderer for real-time rendering of NeRFs (CUDA + JAX)\n* a GUI for visualizing \\\u0026 interacting \\\u0026 exploring NeRFs **[@seimeicyx]**\n\nBenchmarks\n----------\n\n### [NeRF-synthetic]\n\n|                                            | mic   | ficus | chair | hotdog | materials | drums | ship  | lego  | average |\n|:--- |:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| _@33.7k steps_ (this codebase) | 37.04 | 33.14 | 35.10 | 37.21  | 29.50     | 25.85 | 30.93 | 35.95 | 33.09   |\n| _@51.2k steps_ (this codebase) | 37.07 | 33.17 | 35.16 | 37.26  | 29.50     | 25.86 | 30.94 | 36.03 | 33.124  |\n| paper ([instant-ngp])          | 36.22 | 33.51 | 35.00 | 37.40  | 29.78     | 26.02 | 31.10 | 36.39 | 33.176  |\n\n\u003csup\u003e\nFor each scene, the network is trained on 100 training images (800x800 each) for 30k steps with\ndefault parameters, reported PSNR is averaged across 200 test images.\n\u003c/sup\u003e\n\nEnvironment Setup\n-----------------\n\njaxngp manages environments with Nix, but it's also possible to setup the environment with any other package manager (e.g. Conda).\n\n### With Nix (recommended)\n\n1. Install Nix with the [official installer](https://nixos.org/download) or the [nix-installer](https://github.com/DeterminateSystems/nix-installer/releases).\n2. With the `nix` executable available, clone this repository and setup environment:\n   ```bash\n   $ git clone https://github.com/blurgyy/jaxngp.git\n   $ cd jaxngp/\n   $ NIXPKGS_ALLOW_UNFREE=1 nix develop --impure\n   ```\n   This will download (or build if necessary) all the dependencies, and opens a new shell with all the dependencies configured.\n   \u003e **Note**: to avoid the built environment being garbage collected when `nix gc` or `nix-collect-garbage` is called, append a `--profile \u003cPATH\u003e` argument:\n   \u003e ```bash\n   \u003e $ NIXPKGS_ALLOW_UNFREE=1 nix develop --impure --profile .git/devshell.profile\n   \u003e ```\n\n### With Conda\n\nTODO\n\nRunning\n-------\n\n\u003c!-- \u003e **Note**: All the commands below are run after the environment has been setup. --\u003e\n\nThe program's entrance is at `python3 -m app.nerf`.  It provides three subcommands: [`train`](./app/nerf/train.py), [`test`](./app/nerf/test.py), and [`gui`](./app/nerf/gui.py).  Pass `-h|--help` to any of the subcommand to see its usage, e.g.:\n\n\u003cdetails\u003e\n\u003csummary\u003e\n  \u003ccode\u003epython3 -m app.nerf train --help\u003c/code\u003e\n\u003c/summary\u003e\n\n```markdown\nusage: __main__.py train [-h] --exp-dir PATH [--raymarch.diagonal-n-steps INT]\n                         [--raymarch.perturb | --raymarch.no-perturb]\n                         [--raymarch.density-grid-res INT] [--render.bg FLOAT FLOAT FLOAT]\n                         [--render.random-bg | --render.no-random-bg]\n                         [--scene.sharpness-threshold FLOAT] [--scene.world-scale FLOAT]\n                         [--scene.resolution-scale FLOAT] [--scene.camera-near FLOAT]\n                         [--logging {DEBUG,INFO,WARN,WARNING,ERROR,CRITICAL}] [--seed INT]\n                         [--summary | --no-summary] [--frames-val PATH [PATH ...]]\n                         [--ckpt {None}|PATH] [--lr FLOAT] [--tv-scale FLOAT] [--bs INT]\n                         [--n-epochs INT] [--n-batches INT] [--data-loop INT] [--validate-every INT]\n                         [--keep INT] [--keep-every {None}|INT]\n                         [--raymarch-eval.diagonal-n-steps INT]\n                         [--raymarch-eval.perturb | --raymarch-eval.no-perturb]\n                         [--raymarch-eval.density-grid-res INT] [--render-eval.bg FLOAT FLOAT FLOAT]\n                         [--render-eval.random-bg | --render-eval.no-random-bg]\n                         PATH [PATH ...]\n\n╭─ positional arguments ───────────────────────────────────────────────────────────────────────────╮\n│ PATH [PATH ...]         directories or transform.json files containing data for training         │\n│                         (required)                                                               │\n╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\n╭─ arguments ──────────────────────────────────────────────────────────────────────────────────────╮\n│ -h, --help              show this help message and exit                                          │\n│ --exp-dir PATH          experiment artifacts are saved under this directory (required)           │\n│ --frames-val PATH [PATH ...]                                                                     │\n│                         directories or transform.json files containing data for validation       │\n│                         (default: )                                                              │\n│ --ckpt {None}|PATH      if specified, continue training from this checkpoint (default: None)     │\n╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\n╭─ raymarch arguments ─────────────────────────────────────────────────────────────────────────────╮\n│ raymarching/rendering options during training                                                    │\n│ ──────────────────────────────────────────────────────────────────────────────────────────────── │\n│ --raymarch.diagonal-n-steps INT                                                                  │\n│                         for calculating the length of a minimal ray marching step, the NGP paper │\n│                         uses 1024 (appendix E.1) (default: 1024)                                 │\n│ --raymarch.perturb, --raymarch.no-perturb                                                        │\n│                         whether to fluctuate the first sample along the ray with a tiny          │\n│                         perturbation (default: True)                                             │\n│ --raymarch.density-grid-res INT                                                                  │\n│                         resolution for the auxiliary density/occupancy grid, the NGP paper uses  │\n│                         128 (appendix E.2) (default: 128)                                        │\n╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\n╭─ render arguments ───────────────────────────────────────────────────────────────────────────────╮\n│ raymarching/rendering options during training                                                    │\n│ ──────────────────────────────────────────────────────────────────────────────────────────────── │\n│ --render.bg FLOAT FLOAT FLOAT                                                                    │\n│                         background color for transparent parts of the image, has no effect if    │\n│                         `random_bg` is True (default: 1.0 1.0 1.0)                               │\n│ --render.random-bg, --render.no-random-bg                                                        │\n│                         ignore `bg` specification and use random color for transparent parts of  │\n│                         the image (default: True)                                                │\n╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\n╭─ scene arguments ────────────────────────────────────────────────────────────────────────────────╮\n│ raymarching/rendering options during training                                                    │\n│ ──────────────────────────────────────────────────────────────────────────────────────────────── │\n│ --scene.sharpness-threshold FLOAT                                                                │\n│                         images with sharpness lower than this value will be discarded (default:  │\n│                         -1.0)                                                                    │\n│ --scene.world-scale FLOAT                                                                        │\n│                         scale both the scene's camera positions and bounding box with this       │\n│                         factor (default: 1.0)                                                    │\n│ --scene.resolution-scale FLOAT                                                                   │\n│                         scale input images in case they are too large, camera intrinsics are     │\n│                         also scaled to match the updated image resolution. (default: 1.0)        │\n│ --scene.camera-near FLOAT                                                                        │\n│                         (default: 0.3)                                                           │\n╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\n╭─ common arguments ───────────────────────────────────────────────────────────────────────────────╮\n│ --logging {DEBUG,INFO,WARN,WARNING,ERROR,CRITICAL}                                               │\n│                         log level (default: INFO)                                                │\n│ --seed INT              random seed (default: 1000000007)                                        │\n│ --summary, --no-summary                                                                          │\n│                         display model information after model init (default: False)              │\n╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\n╭─ train arguments ────────────────────────────────────────────────────────────────────────────────╮\n│ training hyper parameters                                                                        │\n│ ──────────────────────────────────────────────────────────────────────────────────────────────── │\n│ --lr FLOAT              learning rate (default: 0.01)                                            │\n│ --tv-scale FLOAT        scalar multiplied to total variation loss, set this to a positive value  │\n│                         to enable calculation of TV loss (default: 0.0)                          │\n│ --bs INT                batch size (default: 1048576)                                            │\n│ --n-epochs INT          training epochs (default: 50)                                            │\n│ --n-batches INT         batches per epoch (default: 1024)                                        │\n│ --data-loop INT         loop within training data for this number of iterations, this helps      │\n│                         reduce the effective dataloader overhead. (default: 1)                   │\n│ --validate-every INT    will validate every `validate_every` epochs, set this to a large value   │\n│                         to disable validation (default: 10)                                      │\n│ --keep INT              number of latest checkpoints to keep (default: 1)                        │\n│ --keep-every {None}|INT                                                                          │\n│                         how many epochs should a new checkpoint to be kept (in addition to       │\n│                         keeping the last `keep` checkpoints) (default: 8)                        │\n╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\n╭─ raymarch-eval arguments ────────────────────────────────────────────────────────────────────────╮\n│ raymarching/rendering options for validating during training                                     │\n│ ──────────────────────────────────────────────────────────────────────────────────────────────── │\n│ --raymarch-eval.diagonal-n-steps INT                                                             │\n│                         for calculating the length of a minimal ray marching step, the NGP paper │\n│                         uses 1024 (appendix E.1) (default: 1024)                                 │\n│ --raymarch-eval.perturb, --raymarch-eval.no-perturb                                              │\n│                         whether to fluctuate the first sample along the ray with a tiny          │\n│                         perturbation (default: False)                                            │\n│ --raymarch-eval.density-grid-res INT                                                             │\n│                         resolution for the auxiliary density/occupancy grid, the NGP paper uses  │\n│                         128 (appendix E.2) (default: 128)                                        │\n╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\n╭─ render-eval arguments ──────────────────────────────────────────────────────────────────────────╮\n│ raymarching/rendering options for validating during training                                     │\n│ ──────────────────────────────────────────────────────────────────────────────────────────────── │\n│ --render-eval.bg FLOAT FLOAT FLOAT                                                               │\n│                         background color for transparent parts of the image, has no effect if    │\n│                         `random_bg` is True (default: 0.0 0.0 0.0)                               │\n│ --render-eval.random-bg, --render-eval.no-random-bg                                              │\n│                         ignore `bg` specification and use random color for transparent parts of  │\n│                         the image (default: False)                                               │\n╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\n```\n\n\u003csup\u003e\n  Above is just an example and might not reflect the state of the latest codebase.\n\u003c/sup\u003e\n\u003c/details\u003e\n\n### Examples\n\n* `train`\n  * Just train with the default parameters on the `lego` scene from the [NeRF-synthetic] dataset:\n    ```bash\n    $ python3 -m app.nerf train data/nerf_synthetic/lego/transforms_train.json --exp-dir=logs/lego\n    ```\n  * Train for 10 epochs, with a batch size of 262144, on all the 400 (100\\*train + 100\\*validation + 200\\*test) images from the `lego` scene:\n    ```bash\n    $ python3 -m app.nerf train data/nerf_synthetic/lego --exp-dir=logs/lego-trainvaltest --{n-epochs=10,bs=262144}\n    ```\n  * Train on the training and validation splits of the `drums` scene, with a weight of 1e-5 on the Total Variation (TV) loss (by default this weight is 0):\n    ```bash\n    $ python3 -m app.nerf train data/nerf_synthetic/drums/transforms_{train,val}.json --exp-dir=logs/drums-trainval --tv-scale=1e-5\n    ```\n  * Train on the training split of the `mic` scene, validate with the validation split, validate after every epoch:\n    ```bash\n    $ python3 -m app.nerf train data/nerf_synthetic/mic/transforms_train.json --frames-val=data/nerf_synthetic/mic/transforms_val.json --exp-dir=logs/mic --validate-every=1\n    ```\n    \u003e **Note**: The validated images are logged to [tensorboard], located under `--exp-dir`'s `logs/` directory.  View it in browser with:\n    \u003e ```bash\n    \u003e $ tensorboard serve --logdir logs/mic/logs/ --bind_all\n    \u003e TensorBoard 2.10.0 at http://localhost:6006/ (Press CTRL+C to quit)\n    \u003e ```\n* `test`\n  * Test using the latest checkpoint under `logs/mic` directory, with the camera intrinsics and extrinsics of the `mic` scene's test split\n    ```bash\n    $ python3 -m app.nerf test data/nerf_synthetic/mic/transforms_test.json --ckpt=logs/mic/ --exp-dir=output\n    ```\n  * Test on the `mic` scene with given camera extrinsics, but override the camera's resolution to 1920x1080, and use white as background color:\n    ```bash\n    $ python3 -m app.nerf test data/nerf_synthetic/mic/transforms_test.json --ckpt=logs/mic/ --exp-dir=output --camera-override.{width=1920,height=1080} --render.bg 1 1 1\n    ```\n  * Test with a generated orbiting trajectory (see [Demos] for an example) on the `mic` scene, with resolution 1920x1080:\n    ```bash\n    $ python3 -m app.nerf test data/nerf_synthetic/mic --trajectory=orbit --ckpt=logs/mic/ --exp-dir=output --camera-override.{width=1920,height=1080}\n    ```\n* `gui`\n  \u003e **Note**: The `gui` subcommand accepts all the parameters of the `train` subcommand, and additionally a `--viewport` parameter (but the default values of `--viewport` are sane enough to leave as-is).\n  * Train on all the 400 images from the `lego` scene, with 1280x720 as the default size of the rendering viewport:\n    ```bash\n    $ python3 -m app.nerf gui data/nerf_synthetic/lego --exp-dir=logs/gui-lego --viewport.{W=1280,H=720}\n    ```\n\nRunning on Custom Data\n----------------------\n\nA helper CLI (just a [colmap] wrapper via [pycolmap]) is provided for creating an Instant-NGP-compatible scene from a casually captured video or a directory of images:\n\n```bash\n$ python3 -m utils create-scene --help  # see full usage\nusage: ...\n$\n$ # capture or download a video\n$ mkdir -p data/_src\n$ curl https://github.com/blurgyy/jaxngp/assets/44701880/022a7b3c-344d-418f-aba0-0ccb9bfeb374 -Lo data/_src/gundam.mp4\n$\n$ # create a scene from the video, set scene bound to 16, with a background color model\n$ python3 -m utils create-scene data/_src/gundam.mp4 --root-dir=data/gundam --matcher=Sequential --fps=5 --bound=16 --bg\n[...]\n```\n\nAfter the scene has been created, the rest (training/validating/testing) are the same:\n\n```bash\n$ # train on all the registered images\n$ python3 -m utils train data/gundam --exp-dir=logs/gundam\n[...]\n$\n$ # Render novel views, with a resolution of 1920x1080, save results as images and a video (video shown below in the Demo section)\n$ python3 -m app.nerf test data/gundam --{ckpt,exp-dir}=logs/gundam --trajectory=orbit --camera-override.{width=1920,height=1080,no-distortion} --orbit.high=1 --save-as=\"video and images\"\n```\n\n### Demos\n\n* Orbiting camera trajectory:\n\n  [gundam-nvs]\n\n  \u003csup\u003e[Source video][gundam-source] courtesy of **[@lzyronrico]** | Rendering time: \\~18min for 288 frames | Resolution: 1920x1080 | Hardware: 1 GTX 1080Ti\u003c/sup\u003e\n\n* Interactive rendering in the GUI:\n\n  [gui-teaser]\n\n  \u003csup\u003e[Source video][metasun-source] courtesy of **[@filiptronicek]** | Hardware: 1 GTX 1080Ti\u003c/sup\u003e\n\n\n[JAX]: \u003chttps://github.com/google/jax\u003e\n[instant-ngp]: \u003chttps://github.com/NVLabs/instant-ngp\u003e\n[colmap]: \u003chttps://github.com/colmap/colmap\u003e\n[pycolmap]: \u003chttps://github.com/colmap/pycolmap\u003e\n[tensorboard]: \u003chttps://github.com/tensorflow/tensorboard\u003e\n\n[@seimeicyx]: https://github.com/seimeicyx\n[@lzyronrico]: https://github.com/lzyronrico\n[@filiptronicek]: https://github.com/filiptronicek\n\n[NeRF-synthetic]: \u003chttps://drive.google.com/drive/folders/1JDdLGDruGNXWnM1eqY1FNL9PlStjaKWi\u003e\n\n[gundam-source]: \u003chttps://github.com/blurgyy/jaxngp/assets/44701880/022a7b3c-344d-418f-aba0-0ccb9bfeb374\u003e\n[metasun-source]: \u003chttps://twitter.com/filiptronicek/status/1654894133801103360\u003e\n\n[gundam-nvs]: \u003chttps://github.com/blurgyy/jaxngp/assets/44701880/d41c2039-4ef8-4acf-ac6e-419bbf20b9c8\u003e\n[gui-teaser]: \u003chttps://github.com/blurgyy/jaxngp/assets/44701880/b94dcd0f-a66d-404e-aee2-87f91ddf52fe\u003e\n\n[Demos]: \u003c#demos\u003e\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fblurgyy%2Fjaxngp","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fblurgyy%2Fjaxngp","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fblurgyy%2Fjaxngp/lists"}