{"id":20569924,"url":"https://github.com/zubair-irshad/nerf-mae","last_synced_at":"2025-04-07T13:06:58.549Z","repository":{"id":230076741,"uuid":"778034087","full_name":"zubair-irshad/NeRF-MAE","owner":"zubair-irshad","description":"[ECCV 2024] Pytorch code for our ECCV'24 paper NeRF-MAE: Masked AutoEncoders for Self-Supervised 3D Representation Learning for Neural Radiance Fields","archived":false,"fork":false,"pushed_at":"2025-03-20T14:29:57.000Z","size":4682,"stargazers_count":96,"open_issues_count":0,"forks_count":5,"subscribers_count":7,"default_branch":"main","last_synced_at":"2025-03-28T05:51:06.320Z","etag":null,"topics":["3d","3d-deep-learning","3d-detection","3d-unet","differentiable-rendering","feature-pyramid-network","instant-ngp","masked-autoencoder","multi-view","nerf","neural-radiance-fields","neural-rendering","region-proposal-network","representation-learning","self-supervised-learning","semantic-segmantation","super-resoluion","transformers","vision-transformers","vit"],"latest_commit_sha":null,"homepage":"https://nerf-mae.github.io/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/zubair-irshad.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-03-27T00:29:34.000Z","updated_at":"2025-03-24T10:44:55.000Z","dependencies_parsed_at":"2024-07-31T06:04:34.365Z","dependency_job_id":"37c245c4-3d88-4a92-b339-86dfbe2e5e5c","html_url":"https://github.com/zubair-irshad/NeRF-MAE","commit_stats":null,"previous_names":["zubair-irshad/nerf-mae"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/zubair-irshad%2FNeRF-MAE","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/zubair-irshad%2FNeRF-MAE/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/zubair-irshad%2FNeRF-MAE/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/zubair-irshad%2FNeRF-MAE/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/zubair-irshad","download_url":"https://codeload.github.com/zubair-irshad/NeRF-MAE/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247657281,"owners_count":20974345,"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":["3d","3d-deep-learning","3d-detection","3d-unet","differentiable-rendering","feature-pyramid-network","instant-ngp","masked-autoencoder","multi-view","nerf","neural-radiance-fields","neural-rendering","region-proposal-network","representation-learning","self-supervised-learning","semantic-segmantation","super-resoluion","transformers","vision-transformers","vit"],"created_at":"2024-11-16T05:09:43.781Z","updated_at":"2025-04-07T13:06:58.543Z","avatar_url":"https://github.com/zubair-irshad.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003c!-- # NeRF-MAE: Masked AutoEncoders for Self-Supervised 3D Representation Learning for Neural Radiance Fields --\u003e\n\n\u003cdiv align=\"center\"\u003e\n    \u003cimg src=\"demo/nerf-mae_teaser.png\" width=\"85%\"\u003e\n    \u003cimg src=\"demo/nerf-mae_teaser.jpeg\" width=\"85%\"\u003e\n\u003c/div\u003e\n\u003c!-- \u003cp align=\"center\"\u003e\n\u003cimg src=\"demo/nerf-mae_teaser.jpeg\" width=\"100%\"\u003e\n\u003c/p\u003e --\u003e\n\n\u003cbr\u003e\n\u003cdiv align=\"center\"\u003e\n\n[![arXiv](https://img.shields.io/badge/arXiv-2404.01300-gray?style=for-the-badge\u0026logo=arxiv\u0026logoColor=white\u0026color=B31B1B)](https://arxiv.org/abs/2404.01300)\n[![Project Page](https://img.shields.io/badge/Project-Page-orange?style=for-the-badge\u0026logoColor=white\u0026labelColor=gray\u0026link=https%3A%2F%2Fnerf-mae.github.io%2F)](https://nerf-mae.github.io)\n[![Pytorch](https://img.shields.io/badge/Pytorch-%3E1.12-gray?style=for-the-badge\u0026logo=pytorch\u0026logoColor=white\u0026labelColor=gray\u0026color=ee4c2c\u0026link=https%3A%2F%2Fnerf-mae.github.io%2F)](https://pytorch.org/)\n[![Cite](https://img.shields.io/badge/Cite-Bibtex-gray?style=for-the-badge\u0026logoColor=white\u0026color=F7A41D\n)](https://github.com/zubair-irshad/NeRF-MAE?tab=readme-ov-file#citation)\n[![Video](https://img.shields.io/badge/youtube-video-CD201F?style=for-the-badge\u0026logo=youtube\u0026labelColor=grey\n)](https://youtu.be/D60hlhmeuJI?si=d4RfHAwBJgLJXdKj)\n\n\n\n\u003c/div\u003e\n\n---\n\n\u003ca href=\"https://www.tri.global/\" target=\"_blank\"\u003e\n \u003cimg align=\"right\" src=\"demo/GeorgiaTech_RGB.png\" width=\"18%\"/\u003e\n\u003c/a\u003e\n\n\u003ca href=\"https://www.tri.global/\" target=\"_blank\"\u003e\n \u003cimg align=\"right\" src=\"demo/tri-logo.png\" width=\"17%\"/\u003e\n\u003c/a\u003e\n\n### [Project Page](https://nerf-mae.github.io/) | [arXiv](https://arxiv.org/abs/2308.12967) | [PDF](https://arxiv.org/pdf/2308.12967.pdf)\n\n\n\n**NeRF-MAE : Masked AutoEncoders for Self-Supervised 3D Representation Learning for Neural Radiance Fields**\n\n\u003ca href=\"https://zubairirshad.com\"\u003e\u003cstrong\u003eMuhammad Zubair Irshad\u003c/strong\u003e\u003c/a\u003e\n·\n\u003ca href=\"https://zakharos.github.io/\"\u003e\u003cstrong\u003eSergey Zakharov\u003c/strong\u003e\u003c/a\u003e\n·\n\u003ca href=\"https://www.linkedin.com/in/vitorguizilini\"\u003e\u003cstrong\u003eVitor Guizilini\u003c/strong\u003e\u003c/a\u003e\n·\n\u003ca href=\"https://adriengaidon.com/\"\u003e\u003cstrong\u003eAdrien Gaidon\u003c/strong\u003e\u003c/a\u003e\n·\n\u003ca href=\"https://faculty.cc.gatech.edu/~zk15/\"\u003e\u003cstrong\u003eZsolt Kira\u003c/strong\u003e\u003c/a\u003e\n·\n\u003ca href=\"https://www.tri.global/about-us/dr-rares-ambrus\"\u003e\u003cstrong\u003eRares Ambrus\u003c/strong\u003e\u003c/a\u003e\n\u003cbr\u003e **European Conference on Computer Vision, ECCV 2024**\u003cbr\u003e\n\n\u003cb\u003e Toyota Research Institute \u0026nbsp; | \u0026nbsp; Georgia Institute of Technology\u003c/b\u003e\n\n## 💡 Highlights\n- **NeRF-MAE**: The first large-scale pretraining utilizing Neural Radiance Fields (NeRF) as an input modality. We pretrain a single Transformer model on thousands of NeRFs for 3D representation learning.\n- **NeRF-MAE Dataset**: A large-scale NeRF pretraining and downstream task finetuning dataset.\n\n## 🏷️ TODO 🚀\n\n- [x] Release large-scale pretraining code 🚀\n- [x] Release NeRF-MAE dataset comprising radiance and density grids 🚀\n- [x] Release 3D object detection finetuning and eval code 🚀\n- [x] Pretrained NeRF-MAE checkpoints and out-of-the-box model usage 🚀\n\n## NeRF-MAE Model Architecture\n\u003cp align=\"center\"\u003e\n\u003cimg src=\"demo/nerf-mae_architecture.jpg\" width=\"90%\"\u003e\n\u003c/p\u003e\n\n\u003c!-- _________________ \n\n\n\u003cp align=\"center\"\u003e\n\u003cimg src=\"demo/comparison_mae.jpeg\" width=\"100%\"\u003e\n\u003c/p\u003e --\u003e\n\n## Citation\n\nIf you find this repository or our dataset useful, please star ⭐ this repository and consider citing 📝:\n\n```\n@inproceedings{irshad2024nerfmae,\n    title={NeRF-MAE: Masked AutoEncoders for Self-Supervised 3D Representation Learning for Neural Radiance Fields},\n    author={Muhammad Zubair Irshad and Sergey Zakharov and Vitor Guizilini and Adrien Gaidon and Zsolt Kira and Rares Ambrus},\n    booktitle={European Conference on Computer Vision (ECCV)},\n    year={2024}\n    }\n```\n\n### Contents\n - [🌇  Environment](#-environment)\n - [⛳ Model Usage and Checkpoints](#-model-usage-and-checkpoints)\n - [🗂️ Dataset](#-dataset)\n - [📉 Pretraining](#-pretraining)\n - [📊 Finetuning](#-finetuning)\n - [📌 FAQ](#-faq)\n\n ## 🌇  Environment\n\nCreate a python 3.7 virtual environment and install requirements:\n\n```bash\ncd $NeRF-MAE repo\nconda create -n nerf-mae python=3.9\nconda activate nerf-mae\npip install --upgrade pip\npip install -r requirements.txt\npip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html\n```\nThe code was built and tested on **cuda 11.3**\n\nCompile CUDA extension, to run downstream task finetuning, as described in [NeRF-RPN](https://github.com/lyclyc52/NeRF_RPN):\n\n```bash\ncd $NeRF-MAE repo\ncd nerf_rpn/model/rotated_iou/cuda_op\npython setup.py install\ncd ../../../..\n\n```\n\n## ⛳ Model Usage and Checkpoints\n\n- [Hugginface repo to download pretrained and finetuned checkpoints](https://huggingface.co/mirshad7/NeRF-MAE)\n\nNeRF-MAE is structured to provide easy access to pretrained NeRF-MAE models (and reproductions), to facilitate use for various downstream tasks. This is for extracting good visual features from NeRFs if you don't have resources for large-scale pretraining. Our pretraining provides an easy-to-access embedding of any NeRF scene, which can be used for a variety of downstream tasks in a straightforwaed way. \n\nWe have released pretrained and finetuned checkpoints to start using our codebase out-of-the-box. There are two types of usages. 1. Most common one is using the features directly in a downstream task such as an FPN head for 3D Object Detection and 2. Reconstruct the original grid for enforcing losses such as masked reconstruction loss. Below is a sample useage of our model with spelled out comments.\n\n\n1. Get the features to be used in a downstream task\n\n```\nimport torch\n\n# Define Swin Transformer configurations\nswin_config = {\n    \"swin_t\": {\"embed_dim\": 96, \"depths\": [2, 2, 6, 2], \"num_heads\": [3, 6, 12, 24]},\n    \"swin_s\": {\"embed_dim\": 96, \"depths\": [2, 2, 18, 2], \"num_heads\": [3, 6, 12, 24]},\n    \"swin_b\": {\"embed_dim\": 128, \"depths\": [2, 2, 18, 2], \"num_heads\": [3, 6, 12, 24]},\n    \"swin_l\": {\"embed_dim\": 192, \"depths\": [2, 2, 18, 2], \"num_heads\": [6, 12, 24, 48]},\n}\n\n# Set the desired backbone type\nbackbone_type = \"swin_s\"\nconfig = swin_config[backbone_type]\n\n# Initialize Swin Transformer model\nmodel = SwinTransformer_MAE3D_New(\n    patch_size=[4, 4, 4],\n    embed_dim=config[\"embed_dim\"],\n    depths=config[\"depths\"],\n    num_heads=config[\"num_heads\"],\n    window_size=[4, 4, 4],\n    stochastic_depth_prob=0.1,\n    expand_dim=True,\n    resolution=resolution,\n)\n\n# Load checkpoint and remove unused layers\ncheckpoint = torch.load(checkpoint_path, map_location=\"cpu\")\nmodel.load_state_dict(checkpoint[\"state_dict\"])\nfor attr in [\"decoder4\", \"decoder3\", \"decoder2\", \"decoder1\", \"out\", \"mask_token\"]:\n    delattr(model, attr)\n\n# Extract features using Swin Transformer backbone. input_grid has sample shape torch.randn((1, 4, 160, 160, 160))\nfeatures = []\ninput_grid = model.patch_partition(input_grid) + model.pos_embed.type_as(input_grid).to(input_grid.device).clone().detach()\nfor stage in model.stages:\n    input_grid = stage(input_grid)\n    features.append(torch.permute(input_grid, [0, 4, 1, 2, 3]).contiguous())  # Format: [N, C, H, W, D]\n\n#Multi-scale features have shape:  [torch.Size([1, 96, 40, 40, 40]), torch.Size([1, 192, 20, 20, 20]), torch.Size([1, 384, 10, 10, 10]), torch.Size([1, 768, 5, 5, 5])] \n\n# Process features through FPN\n```\n\n2. Get the Original Grid Output \n```\nimport torch\n# Load data from the specified folder and filename with the given resolution.\nres, rgbsigma = load_data(folder_name, filename, resolution=args.resolution)\n\n# rgbsigma has sample shape torch.randn((1, 4, 160, 160, 160))\n\n# Build the model using provided arguments.\nmodel = build_model(args)\n\n# Load checkpoint if provided.\nif args.checkpoint:\n    model.load_state_dict(torch.load(args.checkpoint, map_location=\"cpu\")[\"state_dict\"])\n    model.eval()  # Set model to evaluation mode.\n\n# Run inference getting the features out for downsteam usage\nwith torch.no_grad():\n    pred = model([rgbsigma], is_eval=True)[3]  # Extract only predictions.\n\n```\n\n### 1. How to plug these features for downstream 3D bounding detection from NeRFs (i.e. plug-and-play with a [NeRF-RPN](https://github.com/lyclyc52/NeRF_RPN) OBB prediction head)\n\nPlease also see the section on [Finetuning](#-finetuning). Our released finetuned checkpoint achieves state-of-the-art on 3D object detection in NeRFs. To run evaluation using our finetuned checkpoint on the dataset provided by NeRF-RPN, please run the below script, after updating the paths to the pretrained checkpoint i.e. --checkpoint and  DATA_ROOT depending on evaluation done for ```Front3D``` or ```Scannet```:\n\n```\nbash test_fcos_pretrained.sh\n```\n\nAlso see the cooresponding run file i.e. ```run_fcos_pretrained.py``` and our model adaptation i.e. ```SwinTransformer_FPN_Pretrained_Skip```. This is a minimal adaptation to plug and play our weights with a NeRF-RPN architecture and achieve significant boost in performance. \n\n\n## 🗂️ Dataset\n\nDownload the preprocessed datasets here. \n\n- Pretraining dataset (comprising NeRF radiance and density grids). [Download link](https://s3.amazonaws.com/tri-ml-public.s3.amazonaws.com/github/nerfmae/NeRF-MAE_pretrain.tar.gz)\n- Finetuning dataset (comprising NeRF radiance and density grids and bounding box/semantic labelling annotations). [3D Object Detection (Provided by NeRF-RPN)](https://drive.google.com/drive/folders/1q2wwLi6tSXu1hbEkMyfAKKdEEGQKT6pj), [3D Semantic Segmentation (Coming Soon)](), [Voxel-Super Resolution (Coming Soon)]()\n\n\nExtract pretraining and finetuning dataset under ```NeRF-MAE/datasets```. The directory structure should look like this:\n\n```\nNeRF-MAE\n├── pretrain\n│   ├── features\n│   └── nerfmae_split.npz\n└── finetune\n    └── front3d_rpn_data\n        ├── features\n        ├── aabb\n        └── obb\n```\n\n\nNote: The above datasets are all you need to train and evaluate our method. Bonus: we will be releasing our multi-view rendered posed RGB images from FRONT3D, HM3D and Hypersim as well as Instant-NGP trained checkpoints soon (these comprise over 1M+ images and 3k+ NeRF checkpoints)\n\nPlease note that our dataset was generated using the instruction from [NeRF-RPN](https://github.com/lyclyc52/NeRF_RPN) and [3D-CLR](https://vis-www.cs.umass.edu/3d-clr/). Please consider citing our work, NeRF-RPN and 3D-CLR if you find this dataset useful in your research. \n\nPlease also note that our dataset uses [Front3D](https://arxiv.org/abs/2011.09127), [Habitat-Matterport3D](https://arxiv.org/abs/2109.08238), [HyperSim](https://github.com/apple/ml-hypersim) and [ScanNet](https://www.scan-net.org/) as the base version of the dataset i.e. we train a NeRF per scene and extract radiance and desnity grid as well as aligned NeRF-grid 3D annotations. Please read the term of use for each dataset if you want to utilize the posed multi-view images for each of these datasets. \n\n## 📉 Pretraining\n\nOfcourse, you can also pretrain your own NeRF-MAE models. Navigate to **nerf-mae** folder and run pretraining script. \n\n```\ncd nerf-mae\nbash train_mae3d.sh\n```\n\nCheckout **train_mae3d.sh** file for a complete list of all hyperparameters such as ```num_epochs```, ```lr```, ```masking_prob``` etc. \n\nCheckpoints will be saved at a regular interval of 200 epochs. For reproducing the paper results, we utilize the checkpoints at 1200 epochs.\n\n\n**Notes**: \n1. with default settings i.e. ```batch_size 32``` and gpus ```0,1,2,3,4,5,6,7``` on ```A100``` GPU, the expected time it takes to pretrain is around 2 days. Please set these accoringly based on your machine's capacity.  \n\n2. The dataset_name is set to default as ```dataset_name=\"nerfmae\"```. This is for convenince for the dataloader as it describes the format. Our pretraining data comprises of scenes from Front3D, Habitat Matterport3D and Hypersim. \n\n\n## 📊 Finetuning\n\nOur finetuning code is largely based on [NeRF-RPN](https://github.com/lyclyc52/NeRF_RPN). Infact, we use the same dataset as NeRF-RPN (unseen during pretraining), for finetuning. This makes sure our comparison with NeRF-RPN is based on the same architecture, the only difference is the network weights are started from scratch for NeRF-RPN, whereas in our case, we start with our pretrained network weights. Please see our paper for more details.\n\n**Note**: We do not see ScanNet dataset during our pretraining. ScanNet 3D OBB prediction finetuning is a challenging case of cross-dataset transfer. \n\n\n### 3D Object Detection\nNavigate to **nerf-rpn** folder and run finetuning script. \n\nTo run 3D Swin Transformer + FPN model finetuning with our pretrained weights:\n\n```\ncd nerf-rpn\nbash train_fcos_pretrained.sh\n```\n\nTo train the 3D Swin Transformer + FPN model model with weights started from scratch:\n\n```\ncd nerf-rpn\nbash train_fcos.sh\n```\n\n**Note**: only 3D Swin Transformer weights are started from our pretraining. FPN weights for both cases are started from scratch. For evaluating our pretrained weights or finetuning from scratch, use ```bash test_fcos_pretrained.sh``` or ```bash test_fcos.sh```\n\nCheckout **train_fcos_pretraining.sh** and ***test_fcos_pretrained.sh*** file for a complete list of all hyperparameters such as ```mae_checkpoint```, ```num_epochs```, ```lr```, ```masking_prob``` etc. Code for finetuning and eval for our downstream tasks are based on [NeRF-RPN's](https://github.com/lyclyc52/NeRF_RPN) implementation.\n\n## FAQ\n\n1. How do I generate bounding-box overlays from the predicted proposals as shown in Figure 9 in our paper? i.e. after running ```test_fcos_pretrained.sh```?\n\nTo visualize proposals, you first need to convert them to ngb boxes, we followed the code from NeRF-RPN. You can use the [code here](https://github.com/lyclyc52/NeRF_RPN/blob/59e90de86e66458aaf852ff94802ae9dc8576306/nerf_rpn/scripts/proposals2ngp.py#L105), it will take in features, proposal generated by our method and input json transforms path and output a new json file with saved boxes in instant-ngp format. \n\nYou will then need to clone [this version](https://github.com/zymk9/instant-ngp/tree/10f337f3467b3992e1ad48a0851aeb029d6642a3) of instant-ngp (also provided in the NeRF-RPN repo) and if you start the instant-ngp visualizer with the new transforms json file, you will see boxes displayed in the native visualizer. If you'd like to render them, the same instant-ngp fork also provides a run file (see [this line](https://github.com/zymk9/instant-ngp/blob/10f337f3467b3992e1ad48a0851aeb029d6642a3/scripts/run.py#L385)), you can repurpose that to render those bounding boxes from the new transforms.json file. Thanks to the authors of NeRF-RPN for open-sourcing the visualization code. \n\n## Acknowledgments\nThis code is built upon the implementation from [NeRF-RPN](https://github.com/lyclyc52/NeRF_RPN). We appreciate the authors for releasing their open-source implementation. \n\n## Licenses\nThis repository and dataset is released under the [CC BY-NC 4.0](https://github.com/zubair-irshad/NeO-360/blob/master/LICENSE.md) license.","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fzubair-irshad%2Fnerf-mae","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fzubair-irshad%2Fnerf-mae","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fzubair-irshad%2Fnerf-mae/lists"}