[
{
"path": ".gitignore",
"content": "*~\n__pycache__/\n*.pyc\n.vscode\noutputs/\nrenders/\ndata/\nvenv/\nstatic/\nsource_videos\nsource_videos/\noverlay_images/\nindex.html\n"
},
{
"path": ".gitmodules",
"content": "[submodule \"nerfstudio\"]\n\tpath = nerfstudio\n\turl = https://github.com/SpectacularAI/nerfstudio\n[submodule \"gsplat\"]\n\tpath = gsplat\n\turl = https://github.com/SpectacularAI/gsplat\n"
},
{
"path": "CITATION.bib",
"content": "@misc{seiskari2024gaussian,\n title={Gaussian Splatting on the Move: Blur and Rolling Shutter Compensation for Natural Camera Motion}, \n author={Otto Seiskari and Jerry Ylilammi and Valtteri Kaatrasalo and Pekka Rantalankila and Matias Turkulainen and Juho Kannala and Arno Solin},\n year={2024},\n eprint={2403.13327},\n archivePrefix={arXiv},\n primaryClass={cs.CV}\n}\n"
},
{
"path": "LICENSE",
"content": " Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [yyyy] [name of copyright owner]\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n"
},
{
"path": "NOTICE",
"content": "The file process_synthetic_inputs.py contains source code from Deblur-NeRF\nhttps://github.com/limacv/Deblur-NeRF/blob/766ca3cfafa026ea45f75ee1d3186ec3d9e13d99/scripts/synthe2poses.py\nAnd is used under the following license:\n\n-----\n\nMIT License\n\nCopyright (c) 2020 bmild\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
},
{
"path": "README.md",
"content": "# Gaussian Splatting on the Move:
Blur and Rolling Shutter Compensation for Natural Camera Motion\n\n[](https://arxiv.org/abs/2403.13327)\n\n## Installation\n\nPrerequisites: run on a Linux system with a recent NVidia RTX GPU with at least 8 GB of VRAM.\nGit must be installed.\n\n 1. Activate a Conda environment with PyTorch that [supports Nerfstudio](https://github.com/nerfstudio-project/nerfstudio/?tab=readme-ov-file#dependencies)\n 2. Possibly required, depending on your environment: `conda install -c conda-forge gcc=12.1.0`\n 3. Run `./scripts/install.sh` (see steps within if something goes wrong)\n\n## Training with custom data\n\n**Custom video data** (_new in version 2_): The method can now be used for motion blur compensation with plain video data as follows \n\n ./scripts/process_and_train_video.sh /path/to/video.mp4\n\nor for rolling shutter compensation as\n\n ROLLING_SHUTTER=ON ./scripts/process_and_train_video.sh /path/to/video.mp4\n\nCurrently simultaneous motion blur and rolling-shutter compensation is only possible with known readout and exposure times. The easiest way to achieve this is using the Spectacular Rec application to record the data (see below).\n\n**Spectacular Rec app** ([v1.0.0+ for Android](https://play.google.com/store/apps/details?id=com.spectacularai.rec), [v1.2.0+ for iOS](https://apps.apple.com/us/app/spectacular-rec/id6473188128)) is needed for simultaneous rolling shutter and motion blur compensation. This approach is also expected to give the best results if the data collection app can be chosen, since it also allows automatic blurry frame filtering and VIO-based velocity initialization, both of which improve the final reconstruction quality. Instructions below.\n\nFirst, download and extract a recording created using the app, e.g., `/PATH/TO/spectacular-rec-MY_RECORDING`.\n\nThen process as\n\n ./scripts/process_and_train_sai_custom.sh /PATH/TO/spectacular-rec-MY_RECORDING\n\nor, for a faster version:\n\n SKIP_COLMAP=ON ./scripts/process_and_train_sai_custom.sh /PATH/TO/spectacular-rec-MY_RECORDING\n\nSee the contents of the script for more details.\n\n**Comparison videos** To train a custom recording with and without motion blur compensation and render a video comparing the two, use this script:\n\n * motion blur OR rolling shutter, COLMAP-based, from video:\n\n ./scripts/render_and_compile_comparison_video.sh /path/to/video.mp4\n ROLLING_SHUTTER=ON ./scripts/render_and_compile_comparison_video.sh /path/to/video.mp4\n\n * motion blur AND rolling shutter compensations (needs Spectacular Rec data)\n \n ./scripts/render_and_train_comparison_sai_custom.sh /PATH/TO/spectacular-rec-MY_RECORDING\n\n## Benchmark data\n\n[](https://doi.org/10.5281/zenodo.10848124)\n[](https://doi.org/10.5281/zenodo.10847884)\n\nThe inputs directly trainable with our fork of Nerfstudio are stored in `data/inputs-processed` folder.\nIts subfolders are called \"datasets\" in these scripts.\n\nThe data can be automatically downloaded by first installing: `pip install unzip` and then running\n\n python download_data.py --dataset synthetic\n # or 'sai' for processed real world smartphone data\n\n\n The data folder structure is as follows:
\n\n\n<3dgs-deblur>\n|---data\n |---inputs-processed\n |---colmap-sai-cli-vels-blur-scored/\n |---iphone-lego1\n |---images\n |---image 0\n |---image 1\n |---...\n |---sparse_pc.ply\n |---transforms.json\n |---...\n |---synthetic-mb\n |---cozyroom\n |---images\n |---image 0\n |---image 1\n |---...\n |---sparse_pc.ply\n |---transforms.json\n |---...\n |---...\n|---...\n\n
\n \n\n## Training\n\nExample: List trainable variants for the `synthetic-mb` dataset:\n\n python train.py --dataset=synthetic-mb\n\nTrain a single variant\n\n python train.py --dataset=synthetic-mb --case=2\n\nCommon useful options:\n\n * `--dry_run`\n * `--preview` (show Viser during training)\n\nAdditionally, any folder of the form `data/inputs-processed/CASE` can be trained directly with Nerfstudio\nusing the `ns-train splatfacto --data data/inputs-processed/CASE ...`. Use `--help` and see `train.py` for\nthe recommended parameters.\n\n## Viewing the results\n\nResults are written to `data/outputs/` by dataset. You can also run these on another machine\nand download these results on your machine. All of the below commands should then work for\nlocally examining the results.\n\n### Numeric\n\nList all numeric results\n\n python parse_outputs.py\n\n... or export to CSV\n\n python parse_outputs.py -f csv > data/results.csv\n\n### Visualizations\n\nOff-the-shelf:\n\n * Viser: `ns-viewer --load-config outputs/DATASET/VARIANT/splatfacto/TIMESTAMP/config.yml` (show actual results)\n * Tensorboard: `tensorboard --logdir outputs/DATASET/VARIANT/splatfacto/TIMESTAMP` (prerequisite `pip install tensorboard`)\n\nCustom:\n\n * Created by `train.py --render_images ...`: Renders of evaluation images and predictions are available in `outputs/DATASET/VARIANT/splatfacto/TIMESTAMP` (`/renders`, or `/demo_video*.mp4` if `render_video.py` has been run, see below)\n * Demo videos: see `render_video.py` and `scripts/render_and_combine_comparison_video.sh`\n\n## Processing the raw benchmark input data\n\nThis method also creates the extra variants discussed in the appendix/supplementary material of the paper,\nas well as all the relevant synthetic data variants.\n\n### Synthetic data\n\nFor synthetic data, we use different re-rendered versions of the [Deblur-NeRF](https://limacv.github.io/deblurnerf/) synthetic dataset.\nNote that there exists several, slightly different variation, which need to be trained with correct parameters for optimal results.\n\n**Our Deblur-NeRF re-render** (uses $\\gamma = 2.2$): Download and process as:\n\n python download_data.py --dataset synthetic-raw\n python process_synthetic_inputs.py\n\n**Other variants**\n\n 1. Download the data and extract as `inputs-raw/FOLDER_NAME` (see options below)\n 2. Run\n\n python process_deblur_nerf_inputs.py --dataset=FOLDER_NAME --manual_point_cloud all\n\nThis creates a dataset called `colmap-DATASET-synthetic-novel-view-manual-pc`\nNote that it may be necessary to run the last command multiple times until COLMAP succeeds\nin all cases (see also the `--case=N` argument in the script).\n\nSupported datasets (TODO: a bit messy):\n\n * Original Deblur-NeRF: `FOLDER_NAME` = `synthetic_camera_motion_blur`. Uses $\\gamma = 2.2$.\n * [BAD-NeRF](https://wangpeng000.github.io/BAD-NeRF/) re-render: `FOLDER_NAME` = `nerf_llff_data`. Uses $\\gamma = 1$.\n * [BAD-Gaussians](https://lingzhezhao.github.io/BAD-Gaussians/) re-render: `FOLDER_NAME` = `bad-nerf-gtK-colmap-nvs`\n \nThe last two are very similar except for the \"Tanabata\" scene, which is broken in the BAD-NeRF version:\nthe underlying 3D model is slightly different in the (sharp) and training (blurry) images (objects moved around).\n\n### Smartphone data\n\nDownload as:\n\n python download_data.py --dataset sai-raw\n\nand then process and convert using the following script:\n\n ./scripts/process_smartphone_dataset.sh\n # or \n # EXTRA_VARIANTS=ON ./scripts/process_smartphone_dataset.sh\n\nNote: all the components in this pipeline are not guaranteed to be deterministic, especially when executed on different machines.\nEspecially the COLMAP has a high level of randomness.\n\n## Changelog\n\n### Version 2 (2024-05)\n\n * Angular and linear velocities added as optimizable variables, which can be initialized to zero if VIO-estimated velocity data is not available (i.e., no IMU data available)\n * Added `--optimize-eval-cameras` mode, which allows optimizing evaluation camera poses and velocities (if `--optimize-eval-velocities=True`) without back-propagating information to the 3DSG reconstruction. This replaces the previous two-phase optimization mode (called \"rolling shutter pose optimization\" in the first paper revision)\n * Method can be run in motion blur OR rolling-shutter mode form plain video without a known exposure or readout times. Added a helper script `process_and_train_video.sh` for this.\n * Rebased on Nerfstudio version 1.1.0 and `gsplat` [409bcd3c](https://github.com/nerfstudio-project/gsplat/commit/409bcd3cf63491710444e60c29d3c44608d8eafd) (based on 0.1.11)\n * Fixed a bug in pixel velocity formulas\n * Tuned hyper-parameters (separate parameters for synthetic and real data)\n * Using [optimizable background color](https://github.com/nerfstudio-project/nerfstudio/pull/3100) by [KevinXu02](https://github.com/KevinXu02) for synthetic data\n * Using $\\gamma \\neq 1$ and `--min-rgb-level` only when motion blur compensation is enabled (for a more fair comparison to Splatfacto)\n * Added conversion scripts for other common Deblur-NeRF dataset variants\n\n### Version 1 (2024-03)\n\nInitial release where IMU data was mandatory to run the method, and the uncertainties in VIO-estimated velocities were addressed with a custom regularization scheme (see ยง3.6 in the [first revision of the paper](https://arxiv.org/pdf/2403.13327v1)).\nBased on Nerfstudio version 1.0.2 and `gsplat` 0.1.8.\n\n## License\n\nThe code in this repository (except the `gh-pages` website branch) is licensed under Apache 2.0.\nSee `LICENSE` and `NOTICE` files for more information.\n\nFor the source code of the website and its license, see the [`gh-pages` branch](https://github.com/SpectacularAI/3dgs-deblur/tree/gh-pages).\n\nThe licenses of the datasets (CC BY-SA 4.0 & CC BY 4.0) are detailed on the Zenodo pages.\n"
},
{
"path": "combine.py",
"content": "\"\"\"Combine COLMAP poses with sai-cli velocities\"\"\"\nimport os\nimport json\nimport shutil\n\ndef process(input_folder, args):\n if args.override_calibration is None:\n override_calibration = None\n else:\n with open(args.override_calibration, 'rt') as f:\n calib_json = json.load(f)\n calib_json_cam0, = calib_json['cameras']\n override_calibration = calib_json_cam0\n \n name = os.path.basename(os.path.normpath(input_folder))\n print('name', name)\n SAI_INPUT_ROOT = 'data/inputs-processed/' + args.dataset\n\n def read_json(path):\n with open(path) as f:\n return json.load(f)\n\n if args.sai_input_folder is None:\n sai_folder = os.path.join(SAI_INPUT_ROOT, name)\n else:\n sai_folder = args.sai_input_folder\n\n if args.pose_opt_pass_dir is None:\n src_poses = read_json(os.path.join(input_folder, 'transforms.json'))\n image_folder = os.path.join(input_folder, 'images')\n ply_pc = os.path.join(input_folder, 'sparse_pc.ply')\n else:\n model_f = os.path.join(input_folder, args.model_name)\n input_json_path = os.path.join(model_f, os.listdir(model_f)[0], 'transforms_train.json')\n src_poses = { 'frames': read_json(input_json_path) }\n image_folder = os.path.join(sai_folder, 'images')\n ply_pc = os.path.join(sai_folder, 'sparse_pc.ply')\n\n sai_poses = read_json(os.path.join(sai_folder, 'transforms.json'))\n\n src_poses_by_filename = { './images/' + os.path.basename(f['file_path']): f for f in src_poses['frames'] }\n if len(src_poses_by_filename) == 0:\n print('skipping: no source poses found')\n return\n\n # print([(k, src_poses_by_filename[k]['file_path']) for k in sorted(src_poses_by_filename.keys())])\n\n combined_frames = []\n\n import numpy as np\n frame_centers_sai = []\n frame_centers_src = []\n\n for sai_frame in sai_poses['frames']:\n id = sai_frame['file_path']\n if id.startswith('images'): id = './' + id\n frame = src_poses_by_filename.get(id, None)\n\n if frame is None:\n print('warning: could not find source pose for %s, skipping' % id)\n if not args.tolerate_missing: return\n continue\n # print('found frame', id)\n \n if 'transform' in frame:\n frame['transform_matrix'] = frame['transform']\n frame['transform_matrix'].append([0, 0, 0, 1])\n del frame['transform']\n\n frame['file_path'] = id\n\n frame_centers_sai.append(np.array(sai_frame['transform_matrix'])[:3, 3].tolist())\n frame_centers_src.append(np.array(frame['transform_matrix'])[:3, 3].tolist())\n\n for prop in ['camera_angular_velocity', 'camera_linear_velocity']:\n if prop in sai_frame:\n frame[prop] = sai_frame[prop]\n\n for prop in ['motion_blur_score']:\n if prop in sai_frame:\n frame[prop] = sai_frame[prop]\n\n for prop in ['colmap_im_id']:\n if prop in frame:\n del frame[prop]\n\n combined_frames.append(frame)\n\n # scale velocities to match COLMAP\n frame_centers_sai = np.array(frame_centers_sai)\n frame_centers_src = np.array(frame_centers_src)\n frame_centers_sai -= np.mean(frame_centers_sai, axis=0)\n frame_centers_src -= np.mean(frame_centers_src, axis=0)\n scale_factor = np.sqrt(np.sum(frame_centers_src**2)) / np.sqrt(np.sum(frame_centers_sai**2))\n print('scene scale factor %.12f' % scale_factor)\n if args.pose_opt_pass_dir is None: \n print('scaling linear velocities')\n\n for frame in combined_frames:\n # only linear velocity should be scaled\n frame['camera_linear_velocity'] = [v * scale_factor for v in frame['camera_linear_velocity']]\n \n processed_prefix = 'data/inputs-processed'\n \n if args.pose_opt_pass_dir is not None:\n output_prefix = os.path.join(processed_prefix, args.dataset + '-2nd-pass')\n combined_poses = sai_poses\n\n elif args.keep_intrinsics or override_calibration is not None:\n combined_poses = sai_poses\n\n if override_calibration is not None:\n assert(override_calibration['model'] == 'brown-conrady')\n def write_to_calib(names, values):\n for i, n in enumerate(names):\n combined_poses[n] = values[i]\n\n write_to_calib('k1 k2 p1 p2 k3'.split(), override_calibration['distortionCoefficients'][:5])\n write_to_calib('fl_x fl_y cx cy'.split(), [override_calibration[c] for c in 'focalLengthX focalLengthY principalPointX principalPointY'.split()])\n\n if override_calibration is None and args.set_rolling_shutter_to is None:\n intrinsics_postfix = 'orig'\n else:\n intrinsics_postfix = 'calib'\n\n output_prefix = os.path.join(processed_prefix, 'colmap-' + args.dataset + '-' + intrinsics_postfix + '-intrinsics')\n combined_poses['applied_transform'] = src_poses['applied_transform']\n\n for prop in ['orientation_override', 'auto_scale_poses_override', 'fx', 'fy']:\n if prop in combined_poses:\n del combined_poses[prop]\n else:\n output_prefix = os.path.join(processed_prefix, 'colmap-' + args.dataset + '-vels')\n combined_poses = src_poses\n for prop in ['exposure_time', 'rolling_shutter_time']:\n if prop in sai_poses:\n combined_poses[prop] = sai_poses[prop]\n\n combined_poses['frames'] = combined_frames\n if args.set_rolling_shutter_to is not None:\n combined_poses['rolling_shutter_time'] = args.set_rolling_shutter_to\n\n if args.output_folder is None:\n output_folder = os.path.join(output_prefix, name)\n else:\n output_folder = args.output_folder\n\n print('Output folder: ' + output_folder)\n if not args.dry_run:\n if os.path.exists(output_folder): shutil.rmtree(output_folder)\n shutil.copytree(image_folder, os.path.join(output_folder, 'images'))\n # shutil.copytree(colmap_folder, os.path.join(output_folder, 'colmap'))\n shutil.copyfile(ply_pc, os.path.join(output_folder, 'sparse_pc.ply'))\n with open (os.path.join(output_folder, 'transforms.json'), 'w') as f:\n json.dump(combined_poses, f, indent=4)\n\nif __name__ == '__main__':\n import argparse\n parser = argparse.ArgumentParser(description=__doc__)\n\n parser.add_argument(\"input_folder\", type=str, default=None, nargs='?')\n parser.add_argument('sai_input_folder', default=None, nargs='?')\n parser.add_argument('output_folder', default=None, nargs='?')\n parser.add_argument('--dataset', default='sai-cli')\n parser.add_argument('--set_rolling_shutter_to', default=None, type=float)\n parser.add_argument('--keep_intrinsics', action='store_true')\n parser.add_argument('--tolerate_missing', action='store_true')\n parser.add_argument('--override_calibration', type=str, default=None)\n parser.add_argument('--pose_opt_pass_dir', type=str, default=None)\n parser.add_argument('--model_name', default='splatfacto')\n parser.add_argument('--dry_run', action='store_true')\n parser.add_argument('--case_number', type=int, default=-1)\n args = parser.parse_args()\n\n if args.input_folder in ['all']:\n args.case_number = 0\n args.input_folder = None\n\n selected_cases = []\n\n if args.input_folder is None:\n if args.pose_opt_pass_dir is None:\n src_folder = 'data/inputs-processed/colmap-' + args.dataset + '-imgs'\n else:\n src_folder = args.pose_opt_pass_dir\n\n cases = [os.path.join(src_folder, f) for f in sorted(os.listdir(src_folder))]\n\n if args.case_number == -1:\n print('valid cases')\n for i, c in enumerate(cases): print(str(i+1) + ':\\t' + c)\n elif args.case_number == 0:\n selected_cases = cases\n else:\n selected_cases = [cases[args.case_number - 1]]\n else:\n selected_cases = [args.input_folder]\n\n for case in selected_cases:\n print('Processing ' + case)\n process(case, args)\n"
},
{
"path": "download_data.py",
"content": "\"\"\"Script to download processed datasets.\"\"\"\nimport os\nimport subprocess\nfrom dataclasses import dataclass\nfrom pathlib import Path\nfrom typing import Literal\n\nimport tyro\n\n\n@dataclass\nclass DownloadProcessedData:\n save_dir: Path = Path(os.getcwd() + \"/data\")\n \"\"\"Save directory. Default /data.\"\"\"\n dataset: Literal[\"synthetic\", \"sai\", \"synthetic-raw\", \"sai-raw\", \"all\"] = \"synthetic\"\n \"\"\"Dataset download name. Set to 'synthetic' to download all synthetic data. Set to 'spectacular' for real world smartphone captures.\"\"\"\n\n def main(self):\n self.save_dir.mkdir(parents=True, exist_ok=True)\n\n urls = {\n \"inputs-processed\": {\n \"synthetic-all\": \"https://zenodo.org/records/10847884/files/processed-nerfstudio.zip\",\n \"colmap-sai-cli-orig-intrinsics-blur-scored\": \"https://zenodo.org/records/10848124/files/colmap-sai-cli-orig-intrinsics-blur-scored.tar.xz\",\n \"colmap-sai-cli-calib-intrinsics-blur-scored\": \"https://zenodo.org/records/10848124/files/colmap-sai-cli-calib-intrinsics-blur-scored.tar.xz\",\n \"colmap-sai-cli-vels-blur-scored\": \"https://zenodo.org/records/10848124/files/colmap-sai-cli-vels-blur-scored.zip\",\n },\n \"inputs-raw\": {\n \"spectacular-rec\": \"https://zenodo.org/records/10848124/files/spectacular-rec.zip\",\n \"spectacular-rec-extras\": \"https://zenodo.org/records/10848124/files/spectacular-rec-extras.zip\",\n \"synthetic-raw\": \"https://zenodo.org/records/10847884/files/renders.zip\"\n }\n }\n\n def download_dataset(dataset):\n for subfolder, sub_urls in urls.items():\n if dataset not in sub_urls: continue\n \n save_dir = self.save_dir / subfolder\n save_dir.mkdir(parents=True, exist_ok=True)\n download_command = [\"wget\", \"-P\", str(self.save_dir), sub_urls[dataset]]\n\n # download\n try:\n subprocess.run(download_command, check=True)\n print(\"File file downloaded succesfully.\")\n except subprocess.CalledProcessError as e:\n print(f\"Error downloading file: {e}\")\n\n file_name = Path(sub_urls[dataset]).name\n\n # subsubfolder for sai data\n subsubfolder = dataset if \"sai\" in file_name or subfolder == \"inputs-raw\" else \"\"\n if subsubfolder:\n Path(self.save_dir / subfolder / subsubfolder).mkdir(\n parents=True, exist_ok=True\n )\n\n # deal with zip or tar formats\n if Path(sub_urls[dataset]).suffix == \".zip\":\n extract_command = [\n \"unzip\",\n self.save_dir / file_name,\n \"-d\",\n self.save_dir / Path(subfolder) / subsubfolder,\n ]\n else:\n extract_command = [\n \"tar\",\n \"-xvJf\",\n self.save_dir / file_name,\n \"-C\",\n self.save_dir / Path(subfolder) / subsubfolder,\n ]\n\n # extract\n try:\n subprocess.run(extract_command, check=True)\n os.remove(self.save_dir / file_name)\n print(\"Extraction complete.\")\n except subprocess.CalledProcessError as e:\n print(f\"Extraction failed: {e}\")\n\n def download_dataset_by_short_name(dataset):\n if dataset == \"synthetic\":\n for dataset in urls[\"inputs-processed\"].keys():\n if \"synthetic\" in dataset:\n download_dataset(dataset)\n elif dataset == \"sai\":\n for dataset in urls[\"inputs-processed\"].keys():\n if \"sai\" in dataset:\n download_dataset(dataset)\n elif dataset == \"synthetic-raw\":\n download_dataset(\"synthetic-raw\")\n\n elif dataset == \"sai-raw\":\n download_dataset(\"spectacular-rec\")\n download_dataset(\"spectacular-rec-extras\")\n\n else:\n raise NotImplementedError\n \n if self.dataset == \"all\":\n for ds in [\"synthetic\", \"sai\", \"synthetic-raw\", \"sai-raw\"]:\n download_dataset_by_short_name(ds)\n else:\n download_dataset_by_short_name(self.dataset)\n\nif __name__ == \"__main__\":\n tyro.cli(DownloadProcessedData).main()\n"
},
{
"path": "parse_outputs.py",
"content": "\"\"\"Parse output metrics from JSON files\"\"\"\nimport os\nimport json\n\ndef parse_metrics(metrics_path):\n with open(metrics_path) as f:\n return json.load(f)\n\ndef find_and_parse_directories_containing_splatting_metrics(root_dir):\n matching_dirs = []\n\n def parse_dir(dirpath, filename):\n run_name = dirpath[len(root_dir)+1:]\n dataset, _, rest = run_name.partition('/')\n\n rest_split = rest.split('/')\n if len(rest_split) != 4: return None\n variant, session, method, ts = rest_split\n if method != 'splatfacto': return None\n\n m = parse_metrics(os.path.join(dirpath, filename))\n\n d = {\n #'dataset': dataset[:1],\n 'dataset': dataset,\n 'variant': variant,\n 'session': session,\n 'path': dirpath,\n 'time': m.get('wall_clock_time_seconds', -1)\n }\n\n \n for k, v in m['results'].items(): d[k] = v\n # print(d)\n return d\n\n for dirpath, _, filenames in os.walk(root_dir):\n for filename in filenames:\n # print(dirpath, filename)\n if filename == 'metrics.json':\n parsed = parse_dir(dirpath, filename)\n if parsed is not None:\n matching_dirs.append(parsed)\n break\n\n return sorted(matching_dirs, key=lambda x: x['path'])\n\nif __name__ == '__main__':\n import argparse\n parser = argparse.ArgumentParser(description=__doc__)\n parser.add_argument('dataset', type=str, nargs='?', default=None)\n parser.add_argument('-f', '--output_format', choices=['csv', 'txt'], default='txt')\n args = parser.parse_args()\n\n import pandas as pd\n pd.set_option(\"display.max_rows\", None)\n df = pd.DataFrame(find_and_parse_directories_containing_splatting_metrics('data/outputs'))\n cols = 'dataset variant session psnr ssim lpips time'.split()\n df = df[cols]\n if args.dataset is not None:\n df = df[df['dataset'] == args.dataset].drop('dataset', axis=1)\n \n if args.output_format == 'csv':\n print(df.to_csv(index=False))\n elif args.output_format == 'txt':\n print(df)\n else:\n raise ValueError(f'Unknown format: {args.output_format}')"
},
{
"path": "process_deblur_nerf_inputs.py",
"content": "\"\"\"Run COLMAP on a single sequence through Nerfstudio scripts\"\"\"\nimport os\nimport subprocess\nimport shutil\nimport tempfile\nimport json\n\nfrom process_synthetic_inputs import generate_seed_points_match_and_triangulate\n\ndef process(input_folder, args, pass_no=1):\n\n name = os.path.basename(os.path.normpath(input_folder))\n\n # 'Wine' is 'Trolley' (see https://github.com/limacv/Deblur-NeRF/issues/39)\n out_name = name.replace('blur', '').replace('2', '').replace('wine', 'trolley')\n\n test_image_folder = None\n first_pass_folder = None\n input_image_folder = os.path.join(input_folder, 'images_1')\n\n if args.hloc:\n method = 'hloc'\n else:\n method = 'colmap'\n\n if args.dataset == 'synthetic_camera_motion_blur':\n paper = 'deblurnerf'\n if args.dataset == 'synthetic_release':\n paper = 'exblurf'\n elif args.dataset == 'nerf_llff_data':\n paper = 'bad-nerf'\n elif args.dataset == 'synthetic-mb':\n input_image_folder = os.path.join(input_folder, 'images')\n paper = 'sai-mb'\n elif args.dataset == 'synthetic-rs':\n input_image_folder = os.path.join(input_folder, 'images')\n paper = 'sai-rs'\n elif args.dataset == 'bad-nerf-gtK-colmap-nvs':\n # this data contains a fixed version of the Tanabata scene\n # where the wine trolley is in the same place in sharp and blurry images\n paper = 'bad-gaussians'\n input_image_folder = os.path.join(input_folder, 'images')\n elif args.dataset == 'colmap-bad-gaussians-synthetic-novel-view-deblurred-training':\n input_image_folder = os.path.join(input_folder, 'images')\n paper = 'mpr-deblurred'\n\n basename = method + '-' + paper + '-synthetic'\n\n if pass_no == 1:\n if args.use_all_images:\n dataset_name = basename + '-all'\n else:\n dataset_name = basename + '-novel-view-temp'\n elif pass_no == 2:\n first_pass_folder = os.path.join('data/inputs-processed/' + basename + '-novel-view-temp', out_name)\n dataset_name = basename + '-novel-view'\n elif pass_no == 3:\n dataset_name = basename + '-deblurring'\n input_image_folder = os.path.join(input_folder, 'images')\n test_image_folder = os.path.join(input_folder, 'images_test')\n else:\n assert False\n\n if pass_no != 1 or args.use_all_images:\n if args.exact_intrinsics:\n dataset_name += '-exact-intrinsics'\n if args.manual_point_cloud:\n dataset_name += '-manual-pc'\n\n output_folder = os.path.join('data/inputs-processed/' + dataset_name, out_name)\n\n temp_dir = tempfile.TemporaryDirectory()\n n = 0\n for index, f in enumerate(sorted(os.listdir(input_image_folder))):\n if 'depth' in f: continue\n if not args.dry_run:\n new_name = f\n if test_image_folder is not None:\n new_name = 'train_' + f\n if pass_no == 1 and index % 8 == 0 and not args.use_all_images:\n continue\n shutil.copyfile(os.path.join(input_image_folder, f), os.path.join(temp_dir.name, new_name))\n n += 1\n print('%d images (would be) copied in a temporary directory' % n)\n\n # Print the path to the temporary directory\n cmd = [\n 'ns-process-data',\n 'images',\n '--data', temp_dir.name,\n '--output-dir', output_folder,\n '--num-downscales', '1',\n '--matching-method', 'exhaustive',\n '--camera-type', 'simple_pinhole',\n ]\n\n if args.hloc:\n cmd.extend([\n '--feature-type', 'superpoint',\n '--matcher-type', 'superpoint+lightglue',\n ])\n\n if not args.post_process_only:\n print(cmd)\n if not args.dry_run:\n if os.path.exists(output_folder):\n shutil.rmtree(output_folder)\n subprocess.check_call(cmd)\n\n json_fn = os.path.join(output_folder, 'transforms.json')\n if os.path.exists(json_fn):\n with open(json_fn, 'r') as f:\n transforms = json.load(f)\n else:\n transforms = { 'frames': [] }\n assert args.dry_run\n\n if test_image_folder is not None:\n assert first_pass_folder is None\n\n test_images = sorted(os.listdir(test_image_folder))\n test_frames = []\n\n if not any('train_' in f['file_path'] for f in transforms['frames']):\n for index, frame in enumerate(sorted(transforms['frames'], key=lambda x: x['file_path'])):\n orig_fn = test_images[index]\n test_image_fn = 'eval_' + orig_fn\n test_image_path = 'images/' + test_image_fn\n\n if not args.dry_run:\n shutil.copyfile(os.path.join(test_image_folder, orig_fn), os.path.join(output_folder, test_image_path))\n\n if 'train_' not in frame['file_path']:\n train_path = 'images/train_' + orig_fn\n if not args.dry_run:\n shutil.move(os.path.join(output_folder, frame['file_path']), os.path.join(output_folder, train_path))\n frame['file_path'] = train_path\n\n test_frame = { k: v for k, v in frame.items() }\n test_frame['file_path'] = test_image_path\n test_frames.append(test_frame)\n\n transforms['frames'].extend(test_frames)\n\n elif first_pass_folder is not None:\n with open(os.path.join(first_pass_folder, 'transforms.json'), 'r') as f:\n first_pass_transforms = json.load(f)\n\n import numpy as np\n to_pose_mat = lambda f : np.array(f['transform_matrix'])\n get_frame_idx = lambda f: int(f['file_path'].split('_')[-1].split('.')[0], base=10) - 1\n\n train_frame_c2ws = { get_frame_idx(f): to_pose_mat(f) for f in first_pass_transforms['frames'] }\n all_frames_c2ws = { get_frame_idx(f): to_pose_mat(f) for f in transforms['frames'] }\n\n combined_transforms = { k: v for k, v in first_pass_transforms.items() }\n combined_transforms['frames'] = []\n\n orig_index = 0\n for index, frame in enumerate(sorted(transforms['frames'], key=lambda x: x['file_path'])):\n #print(frame['file_path'])\n if index % 8 == 0:\n ref_frame = index - 1\n ref_frame_orig_index = orig_index - 1\n if ref_frame < 0:\n ref_frame = index + 1\n ref_frame_orig_index = orig_index # the next frame\n\n # print(index, orig_index, ref_frame, ref_frame_orig_index)\n\n pose_cur_pred_c2w = train_frame_c2ws[ref_frame_orig_index] @ np.linalg.inv(all_frames_c2ws[ref_frame]) @ all_frames_c2ws[index]\n frame['transform_matrix'] = pose_cur_pred_c2w.tolist()\n else:\n frame['transform_matrix'] = train_frame_c2ws[orig_index].tolist()\n orig_index += 1\n\n combined_transforms['frames'].append(frame)\n\n transforms = combined_transforms\n if not args.dry_run:\n shutil.copyfile(os.path.join(first_pass_folder, 'sparse_pc.ply'), os.path.join(output_folder, 'sparse_pc.ply'))\n\n if args.exact_intrinsics:\n KNOWN_INTRINSICS = {\n \"w\": 600,\n \"h\": 400,\n \"cx\": 300.0,\n \"cy\": 200.0,\n \"fl_x\": 541.8502321581475,\n \"fl_y\": 541.8502321581475,\n \"k1\": 0,\n \"k2\": 0,\n \"p1\": 0,\n \"p2\": 0,\n }\n for k, v in KNOWN_INTRINSICS.items():\n transforms[k] = v\n\n print('writing %s' % json_fn)\n if not args.dry_run:\n with open(json_fn, 'wt') as f:\n json.dump(transforms, f, indent=4)\n\n if pass_no == 1 and args.manual_point_cloud:\n if os.path.exists(output_folder):\n if not args.dry_run:\n backup_ply = os.path.join(output_folder, 'sparse_pc_colmap.ply')\n backup_json = os.path.join(output_folder, 'transforms_colmap.json')\n if not os.path.exists(backup_ply):\n ply_fn = os.path.join(output_folder, 'sparse_pc.ply')\n assert os.path.exists(ply_fn) and os.path.exists(json_fn)\n shutil.copyfile(ply_fn, backup_ply)\n if not os.path.exists(backup_json):\n shutil.copyfile(json_fn, backup_json)\n generate_seed_points_match_and_triangulate(output_folder, dry_run=args.dry_run, visualize=args.dry_run)\n else:\n assert args.dry_run\n \n temp_dir.cleanup()\n\nif __name__ == '__main__':\n import argparse\n parser = argparse.ArgumentParser(description=__doc__)\n\n parser.add_argument(\"input_folder\", type=str, default=None, nargs='?')\n parser.add_argument('--dry_run', action='store_true')\n parser.add_argument('--dataset', default='synthetic_camera_motion_blur')\n parser.add_argument('--post_process_only', action='store_true')\n parser.add_argument('--manual_point_cloud', action='store_true')\n parser.add_argument('--deblurring_version', action='store_true')\n parser.add_argument('--exact_intrinsics', action='store_true')\n parser.add_argument('--hloc', action='store_true')\n parser.add_argument('--use_all_images', action='store_true',\n help='Use both blurry training and sharp test images for training pose registration')\n parser.add_argument('--case_number', type=int, default=-1)\n \n args = parser.parse_args()\n\n if args.input_folder in ['all']:\n args.case_number = 0\n args.input_folder = None\n \n selected_cases = []\n misc = False\n\n if args.dataset.endswith('/'): args.dataset = args.dataset[:-1]\n\n if args.input_folder is None:\n sai_dataset = args.dataset.startswith('synthetic-')\n\n if sai_dataset:\n input_root = os.path.join('data/inputs-processed/', args.dataset)\n else:\n input_root = os.path.join('data/inputs-raw/', args.dataset)\n cases = [os.path.join(input_root, f)\n for f in sorted(os.listdir(input_root))\n if f.startswith('blur') or sai_dataset or args.dataset == 'colmap-bad-gaussians-synthetic-novel-view-deblurred-training'\n ]\n\n if args.case_number == -1:\n print('valid cases')\n for i, c in enumerate(cases): print(str(i+1) + ':\\t' + c)\n elif args.case_number == 0:\n selected_cases = cases\n else:\n selected_cases = [cases[args.case_number - 1]]\n else:\n selected_cases = [args.input_folder]\n\n for case in selected_cases:\n print('Processing ' + case)\n process(case, args)\n if not args.use_all_images:\n if args.deblurring_version:\n process(case, args, pass_no=3)\n else:\n process(case, args, pass_no=2)\n"
},
{
"path": "process_sai_custom.py",
"content": "\"\"\"Process a single custom SAI input\"\"\"\nimport os\nimport subprocess\nimport shutil\nimport json\nimport tempfile\n\nfrom process_sai_inputs import SAI_CLI_PROCESS_PARAMS\n\nDEFAULT_OUT_FOLDER = 'data/inputs-processed/custom'\n\ndef ensure_exposure_time(target, input_folder):\n trans_fn = os.path.join(target, 'transforms.json')\n with open(trans_fn) as f:\n transforms = json.load(f)\n \n if 'exposure_time' in transforms: return\n\n with open(os.path.join(input_folder, 'data.jsonl')) as f:\n for line in f:\n d = json.loads(line)\n if 'frames' in d:\n e = d['frames'][0].get('exposureTimeSeconds', None)\n if e is not None:\n print('got exposure time %g from data.jsonl' % e)\n transforms['exposure_time'] = e\n with open(trans_fn, 'wt') as f:\n json.dump(transforms, f, indent=4)\n return\n \n raise RuntimeError(\"no exposure time available\")\n\ndef process(args):\n def maybe_run_cmd(cmd):\n print('COMMAND:', cmd)\n if not args.dry_run: subprocess.check_call(cmd)\n\n def maybe_unzip(fn):\n name = os.path.basename(fn)\n if name.endswith('.zip'):\n name = name[:-4]\n tempdir = tempfile.mkdtemp()\n input_folder = os.path.join(tempdir, 'recording')\n extract_command = [\n \"unzip\",\n fn,\n \"-d\",\n input_folder,\n ]\n maybe_run_cmd(extract_command)\n if not args.dry_run:\n # handle folder inside zip\n for f in os.listdir(input_folder):\n if f == name:\n input_folder = os.path.join(input_folder, f)\n break\n else:\n input_folder = fn\n \n return name, input_folder\n\n sai_params = json.loads(json.dumps(SAI_CLI_PROCESS_PARAMS))\n sai_params['key_frame_distance'] = args.key_frame_distance\n\n tempdir = None\n name, input_folder = maybe_unzip(args.spectacular_rec_input_folder_or_zip)\n\n sai_params_list = []\n for k, v in sai_params.items():\n if k == 'internal':\n for k2, v2 in v.items():\n sai_params_list.append(f'--{k}={k2}:{v2}')\n else:\n if v is None:\n sai_params_list.append(f'--{k}')\n else:\n sai_params_list.append(f'--{k}={v}')\n \n result_name = name\n\n if args.output_folder is None:\n final_target = os.path.join(DEFAULT_OUT_FOLDER, result_name)\n else:\n final_target = args.output_folder\n\n if not args.skip_colmap:\n if tempdir is None: tempdir = tempfile.mkdtemp()\n target = os.path.join(tempdir, 'sai-cli', result_name)\n else:\n target = final_target\n\n cmd = [\n 'sai-cli', 'process',\n input_folder,\n target\n ] + sai_params_list\n\n if args.preview:\n cmd.extend(['--preview', '--preview3d'])\n\n if os.path.exists(target): shutil.rmtree(target)\n maybe_run_cmd(cmd)\n if not args.dry_run: ensure_exposure_time(target, input_folder)\n\n if not args.skip_colmap:\n colmap_target = os.path.join(tempdir, 'colmap-sai-cli-imgs', result_name)\n colmap_cmd = [\n 'python', 'run_colmap.py',\n target,\n colmap_target\n ]\n maybe_run_cmd(colmap_cmd)\n \n combine_cmd = [\n 'python', 'combine.py',\n colmap_target,\n target,\n final_target,\n '--tolerate_missing'\n ]\n if args.keep_intrinsics:\n combine_cmd.append('--keep_intrinsics')\n \n if os.path.exists(final_target): shutil.rmtree(final_target)\n maybe_run_cmd(combine_cmd)\n\nif __name__ == '__main__':\n import argparse\n parser = argparse.ArgumentParser(description=__doc__)\n parser.add_argument(\"spectacular_rec_input_folder_or_zip\", type=str)\n parser.add_argument(\"output_folder\", type=str, default=None, nargs='?')\n parser.add_argument('--preview', action='store_true')\n parser.add_argument('--skip_colmap', action='store_true')\n parser.add_argument('--keep_intrinsics', action='store_true')\n parser.add_argument('--dry_run', action='store_true')\n parser.add_argument('--key_frame_distance', type=float, default=0.1,\n help=\"Minimum key frame distance in meters, default (0.1), increase for larger scenes\")\n args = parser.parse_args()\n\n process(args)\n"
},
{
"path": "process_sai_inputs.py",
"content": "\"\"\"Process raw input data to the main benchmark format\"\"\"\nimport os\nimport subprocess\nimport shutil\nimport json\n\nSAI_CLI_PROCESS_PARAMS = {\n 'image_format': 'png',\n 'no_undistort': None,\n 'key_frame_distance': 0.1,\n 'internal': {\n 'maxKeypoints': 2000,\n 'optimizerMaxIterations': 50,\n }\n}\n\nDATASET_SPECIFIC_PARAMETERS = {}\n\ndef process_subfolders(spec, output_folder, method='sai', only_this_case_number=None, dry_run=False, preview=False):\n def process(folder, counter, prefix, named):\n if named:\n name = os.path.basename(folder)\n else:\n name = \"%02d\" % counter\n \n if prefix is not None:\n name = prefix + '-' + name\n\n sai_params = json.loads(json.dumps(SAI_CLI_PROCESS_PARAMS)) # deep copy\n out_dataset_folder = output_folder\n if args.no_blur_score_filter:\n out_dataset_folder += '-no-blur-select'\n sai_params['blur_filter_range'] = 0\n sai_params['internal']['keyFrameCandidateSelectionBufferSize'] = 1\n\n for k, v in DATASET_SPECIFIC_PARAMETERS.get(prefix, {}).items():\n if k == 'internal':\n for k2, v2 in v.items():\n sai_params['internal'][k2] = v2\n else:\n sai_params[k] = v\n\n sai_params_list = []\n for k, v in sai_params.items():\n if k == 'internal':\n for k2, v2 in v.items():\n sai_params_list.append(f'--{k}={k2}:{v2}')\n else:\n if v is None:\n sai_params_list.append(f'--{k}')\n else:\n sai_params_list.append(f'--{k}={v}')\n \n target = os.path.join(out_dataset_folder, name.replace('_', '-').replace('-capture', ''))\n\n if method == 'sai':\n cmd = [\n 'sai-cli', 'process',\n folder,\n target\n ] + sai_params_list\n\n if preview:\n cmd.extend(['--preview', '--preview3d'])\n\n elif method == 'colmap-video':\n [\n 'ns-process-data',\n 'video',\n '--data', os.path.join(folder, 'data.mp4'),\n '--output-dir', target\n ]\n else:\n assert(False)\n\n if dry_run:\n print(cmd)\n return\n print(f\"Processing: {folder} -> {target}\")\n\n if os.path.exists(target): shutil.rmtree(target)\n subprocess.check_call(cmd)\n\n counter = 1\n for (base_folder, prefix, named) in spec:\n items = os.listdir(base_folder)\n directories = sorted([item for item in items if os.path.isdir(os.path.join(base_folder, item))])\n\n dir_counter = 1\n # Loop through each directory and run a command\n for directory in directories:\n full_path = os.path.join(base_folder, directory)\n if only_this_case_number is None or only_this_case_number == counter:\n print('case %d: %s' % (counter, full_path))\n process(full_path, dir_counter, prefix, named)\n counter += 1\n dir_counter += 1\n\nif __name__ == '__main__':\n import argparse\n parser = argparse.ArgumentParser(description=__doc__)\n parser.add_argument(\"--case_number\", type=int, default=None)\n parser.add_argument('--method', choices={'sai', 'colmap-video'}, default='sai')\n parser.add_argument('--no_blur_score_filter', action='store_true')\n parser.add_argument('--preview', action='store_true')\n parser.add_argument('--dry_run', action='store_true')\n args = parser.parse_args()\n\n if args.method == 'sai':\n out_folder ='data/inputs-processed/sai-cli'\n elif args.method == 'colmap-video':\n out_folder ='data/inputs-processed/colmap-video'\n else:\n assert(False)\n\n process_subfolders([\n ('data/inputs-raw/spectacular-rec', None, True),\n ],\n out_folder,\n method=args.method,\n only_this_case_number=args.case_number,\n dry_run=args.dry_run,\n preview=args.preview)\n"
},
{
"path": "process_synthetic_inputs.py",
"content": "\"\"\"Process raw synthetic input data to the main benchmark format\"\"\"\nimport os\nimport json\nimport shutil\nimport cv2\nimport numpy as np\n\nPOSE_POSITION_NOISE_REL = 0.05\nPOSE_ORIENTATION_NOISE_DEG = 1\n\nINTRINSIC_NOISE_REL = 0.01\n\ndef rotation_matrix_to_rotvec(R):\n # Using a proven/stable algorithm. Other options are sketchy for small rotation\n from scipy.spatial.transform import Rotation\n return Rotation.from_matrix(R).as_rotvec()\n\ndef quaternion_to_rotation_matrix(q_wxyz):\n q = q_wxyz\n return np.array([\n [q[0]*q[0]+q[1]*q[1]-q[2]*q[2]-q[3]*q[3], 2*q[1]*q[2] - 2*q[0]*q[3], 2*q[1]*q[3] + 2*q[0]*q[2]],\n [2*q[1]*q[2] + 2*q[0]*q[3], q[0]*q[0] - q[1]*q[1] + q[2]*q[2] - q[3]*q[3], 2*q[2]*q[3] - 2*q[0]*q[1]],\n [2*q[1]*q[3] - 2*q[0]*q[2], 2*q[2]*q[3] + 2*q[0]*q[1], q[0]*q[0] - q[1]*q[1] - q[2]*q[2] + q[3]*q[3]]\n ])\n\ndef deterministic_uniform_rand_generator(seed=1000):\n \"\"\"\n A simple pseudorandom number generator that returns the\n same random sequence on all machines. The quality of these\n random numbers is low but this is fine for this particular\n application.\n \"\"\"\n\n # see https://en.cppreference.com/w/cpp/numeric/random/linear_congruential_engine\n\n a, c, m = 48271, 0, 2147483647\n x = seed + 1\n uniform_steps = 999\n\n while True:\n x = (a * x + c) % m\n yield float(x % uniform_steps) / uniform_steps\n\ndef process(data_path, target, noisy_poses=False, noisy_intrinsics=False):\n \"\"\"\n # --- Based on\n # https://github.com/limacv/Deblur-NeRF/blob/766ca3cfafa026ea45f75ee1d3186ec3d9e13d99/scripts/synthe2poses.py\n # and used under the following license\n\n MIT License\n\n Copyright (c) 2020 bmild\n\n Permission is hereby granted, free of charge, to any person obtaining a copy\n of this software and associated documentation files (the \"Software\"), to deal\n in the Software without restriction, including without limitation the rights\n to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n copies of the Software, and to permit persons to whom the Software is\n furnished to do so, subject to the following conditions:\n\n The above copyright notice and this permission notice shall be included in all\n copies or substantial portions of the Software.\n\n THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n SOFTWARE.\n \"\"\"\n\n print(f\"Processing: {data_path} -> {target}\")\n if os.path.exists(target): shutil.rmtree(target)\n\n input_path = data_path\n json_path = os.path.join(input_path, \"transforms.json\")\n out_path = os.path.join(target, \"images\")\n converted_json_path = os.path.join(target, \"transforms.json\")\n os.makedirs(out_path, exist_ok=True)\n\n rand = deterministic_uniform_rand_generator()\n def rand3():\n nonlocal rand\n return np.array([next(rand) for _ in range(3)]) * 2 - 1\n\n def convert_pose_c2w(pose, scaling):\n pose = np.array(pose)\n pose[:3, :] *= scaling\n return pose\n\n def get_scaling(m):\n return 1.0 / np.sqrt((m[:3,:3].transpose() @ m[:3,:3])[0,0])\n\n with open(json_path, 'r') as metaf:\n meta = json.load(metaf)\n frames_data = meta[\"frames\"]\n fov = meta[\"fov\"]\n h, w = meta['h'], meta['w']\n exposure_time = meta[\"exposure_time\"]\n rolling_shutter_time = meta[\"rolling_shutter_time\"]\n\n focal_length = w / 2 / np.tan(fov / 2)\n\n if noisy_intrinsics:\n # slight (fixed) error in intrinsics\n intrinsic_noisy_scaling_x = 1 + INTRINSIC_NOISE_REL\n intrinsic_noisy_scaling_y = 1 - INTRINSIC_NOISE_REL\n else:\n intrinsic_noisy_scaling_x = 1\n intrinsic_noisy_scaling_y = 1\n\n converted_meta = {\n \"aabb_scale\": 16,\n \"w\": w,\n \"h\": h,\n \"cx\": w/2,\n \"cy\": h/2,\n \"orientation_override\": \"none\",\n \"exposure_time\": exposure_time,\n \"rolling_shutter_time\": rolling_shutter_time,\n \"fl_x\": focal_length * intrinsic_noisy_scaling_x,\n \"fl_y\": focal_length * intrinsic_noisy_scaling_y,\n \"k1\": 0,\n \"k2\": 0,\n \"p1\": 0,\n \"p2\": 0,\n \"frames\": []\n }\n\n scaling = None\n\n cam_positions = []\n\n for frame_data in frames_data:\n pose = np.array(frame_data[\"transform_matrix\"])\n if scaling is None:\n scaling = get_scaling(pose)\n pose = convert_pose_c2w(pose, scaling)\n cam_positions.append(pose[:3, 3])\n img_path = os.path.join(data_path, frame_data[\"filename\"])\n img_name = os.path.basename(img_path)\n img_out = os.path.join(out_path, img_name)\n\n if frame_data[\"blurcount\"] == 0:\n img = cv2.imread(img_path)\n cv2.imwrite(img_out, img)\n\n velocity_cam = np.array([0, 0, 0])\n ang_vel_cam = np.array([0, 0, 0])\n else:\n img = cv2.imread(img_path)\n blur_poses = []\n for bluri in range(frame_data[\"blurcount\"]):\n blur_poses.append(convert_pose_c2w(frame_data['blur_matrices'][bluri], scaling))\n\n velocity_w = (blur_poses[-1][:3, 3] - blur_poses[0][:3, 3]) / (exposure_time + rolling_shutter_time)\n rot = blur_poses[-1][:3, :3] @ blur_poses[0][:3, :3].transpose()\n rot_vec = rotation_matrix_to_rotvec(rot)\n # print(rot, rot_vec, np.linalg.norm(rot_vec))\n ang_vel_w = rot_vec / (exposure_time + rolling_shutter_time)\n\n R_w2c = pose[:3, :3].transpose()\n velocity_cam = R_w2c @ velocity_w\n ang_vel_cam = R_w2c @ ang_vel_w\n # print(velocity_cam, ang_vel_cam)\n cv2.imwrite(img_out, img)\n\n print(f\"frame {img_name} saved!\")\n\n converted_meta[\"frames\"].append({\n \"camera_linear_velocity\": velocity_cam.tolist(),\n \"camera_angular_velocity\": ang_vel_cam.tolist(),\n \"file_path\": f\"./images/{img_name}\",\n \"transform_matrix\": pose.tolist()\n })\n \n if noisy_poses:\n center = np.mean(cam_positions, axis=0)\n scene_motion_scale = np.max(np.linalg.norm(cam_positions - center, axis=1))\n pos_noise_scale = POSE_POSITION_NOISE_REL * scene_motion_scale\n print('center point of scene cameras %s scale %g, pose noise scale +-%g' % (\n str(center.tolist()),\n scene_motion_scale,\n pos_noise_scale))\n for f in converted_meta['frames']:\n pose = np.array(f['transform_matrix'])\n pose[:3, 3] + rand3() * pos_noise_scale\n noise_ang = 0\n while noise_ang < 1e-6:\n noise_rot_vec = rand3() * POSE_ORIENTATION_NOISE_DEG / 180.0 * np.pi\n noise_ang = np.linalg.norm(noise_rot_vec)\n\n noise_rot_dir = noise_rot_vec / noise_ang\n noise_quat = [np.cos(noise_ang*0.5)] + (np.sin(noise_ang*0.5) * noise_rot_dir).tolist()\n noise_R = quaternion_to_rotation_matrix(noise_quat)\n pose[:3, :3] = pose[:3, :3] @ noise_R\n f['transform_matrix'] = pose.tolist()\n\n with open(converted_json_path, 'wt') as f:\n json.dump(converted_meta, f, indent=4)\n\ndef point_cloud_to_ply(xyzrgbs, out_fn):\n with open(out_fn, 'wt') as f:\n f.write('\\n'.join([\n 'ply',\n 'format ascii 1.0',\n 'element vertex %d' % len(xyzrgbs),\n 'property float x',\n 'property float y',\n 'property float z',\n 'property uint8 red',\n 'property uint8 green',\n 'property uint8 blue',\n 'end_header'\n ]) + '\\n')\n for r in xyzrgbs:\n for i in range(3): r[i+3] = int(r[i+3])\n f.write(' '.join([str(v) for v in r]) + '\\n')\n\ndef triangulate_point(o1, d1, o2, d2):\n A = np.stack([d1, -d2]).T\n b = o2 - o1\n x, _, _, _ = np.linalg.lstsq(A, b, rcond=None)\n P1 = o1 + x[0] * d1\n P2 = o2 + x[1] * d2\n P = (P1 + P2) / 2\n return P\n\ndef reproject_point(p, c2w, intrinsics):\n p_cam = c2w[:3, :3].transpose() @ (p - c2w[:3, 3])\n MIN_D = 1e-6\n\n if -p_cam[2] <= MIN_D: return None\n\n p_img = p_cam[:2] / -p_cam[2]\n p_px = [p_img[0] * intrinsics['fl_x'] + intrinsics['cx'], -p_img[1] * intrinsics['fl_y'] + intrinsics['cy']]\n return p_px\n\ndef reprojection_error(p_reproj, p_orig):\n if p_reproj is None: return 1e6\n return np.linalg.norm(p_reproj - np.array(p_orig))\n\ndef triangulate(points1, points2, c2w_i, c2w_j, matches, intrinsics, reprojection_error_pixels):\n filtered_matches = []\n points3d = []\n rejected_matches = []\n\n for match in matches:\n i, j = match.queryIdx, match.trainIdx\n\n def to_dir(p):\n px = (p[0] - intrinsics['cx']) / intrinsics['fl_x']\n py = -(p[1] - intrinsics['cy']) / intrinsics['fl_y']\n h = [px, py, -1]\n return np.array(h) / np.linalg.norm(h)\n\n p1 = points1[i].pt\n p2 = points2[j].pt\n\n dir_i_cam = to_dir(p1)\n dir_j_cam = to_dir(p2)\n\n dir_i = c2w_i[:3, :3] @ dir_i_cam\n dir_j = c2w_j[:3, :3] @ dir_j_cam\n\n P = triangulate_point(c2w_i[:3, 3], dir_i, c2w_j[:3, 3], dir_j)\n\n rp1 = reproject_point(P, c2w_i, intrinsics)\n rp2 = reproject_point(P, c2w_j, intrinsics)\n\n err = max(\n reprojection_error(rp1, p1),\n reprojection_error(rp2, p2))\n\n if err > reprojection_error_pixels:\n rejected_matches.append((match, rp1, rp2))\n continue\n\n filtered_matches.append(match)\n points3d.append(P)\n\n return filtered_matches, points3d, rejected_matches\n\ndef generate_seed_points_match_and_triangulate(target, visualize=False, dry_run=False, reprojection_error_pixels=10):\n json_path = os.path.join(target, \"transforms.json\")\n def is_eval_frame(i, frame):\n if i % 8 == 0:\n if 'camera_linear_velocity' in frame:\n vel = np.linalg.norm(frame['camera_linear_velocity']) + np.linalg.norm(frame['camera_angular_velocity'])\n assert(vel == 0)\n return True\n return False\n\n with open(json_path, 'rt') as f: transforms = json.load(f)\n training_frames = [f for i, f in enumerate(sorted(transforms['frames'], key=lambda fr: fr['file_path'])) if not is_eval_frame(i, f)]\n\n transforms['ply_file_path'] = './sparse_pc.ply'\n converted_json = transforms\n\n images = [cv2.imread(os.path.join(target, frame['file_path'])) for frame in training_frames]\n\n # --- By ChatGPT\n def find_keypoints_and_descriptors(images, detector):\n \"\"\"Find keypoints and descriptors for each image using the given detector.\"\"\"\n keypoints_and_descriptors = []\n for image in images:\n keypoints, descriptors = detector.detectAndCompute(image, None)\n keypoints_and_descriptors.append((keypoints, descriptors))\n return keypoints_and_descriptors\n\n def match_descriptors_and_triangulate(descriptor_pairs, matcher, frames, intrinsics):\n \"\"\"Match descriptors between all pairs of images.\"\"\"\n matches = {}\n n = len(descriptor_pairs)\n for i in range(n):\n for j in range(i+1, n):\n matches_ij = matcher.match(descriptor_pairs[i][1], descriptor_pairs[j][1])\n matches_ij = sorted(matches_ij, key=lambda x: x.distance)\n\n c2w_i = np.array(frames[i]['transform_matrix'])\n c2w_j = np.array(frames[j]['transform_matrix'])\n\n matches_ij, points3d, rejected_matches = triangulate(\n descriptor_pairs[i][0],\n descriptor_pairs[j][0],\n c2w_i, c2w_j,\n matches_ij, intrinsics, reprojection_error_pixels)\n matches[(i, j)] = (matches_ij, points3d, rejected_matches)\n\n return matches\n\n def visualize_matches(images, keypoints_and_descriptors, matches, pair):\n \"\"\"Visualize the matches for a specific pair of images.\"\"\"\n img1, img2 = images[pair[0]], images[pair[1]]\n kp1, kp2 = keypoints_and_descriptors[pair[0]][0], keypoints_and_descriptors[pair[1]][0]\n matches_ij, points, rejected_matches = matches[pair]\n img_matches = cv2.drawMatches(img1, kp1, img2, kp2, matches_ij, None, flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS)\n for rm in rejected_matches:\n match, rp1, rp2 = rm\n p1_orig = tuple(map(int, kp1[match.queryIdx].pt))\n p2_orig_x, p2_orig_y = tuple(map(int, kp2[match.trainIdx].pt))\n p2_orig_x += img1.shape[1]\n p2_orig = (p2_orig_x, p2_orig_y)\n cv2.circle(img_matches, p1_orig, 3, (0, 0, 255), 1)\n cv2.circle(img_matches, p2_orig, 3, (0, 0, 255), 1)\n if rp1 is not None:\n cv2.line(img_matches, p1_orig, tuple(map(int, rp1)), (0, 0, 255), 1)\n if rp2 is not None:\n rp2_x, rp2_y = tuple(map(int, rp2))\n cv2.line(img_matches, p2_orig, (rp2_x + img1.shape[1], rp2_y), (0, 0, 255), 1)\n cv2.imshow(f\"Matches between image {pair[0]} and {pair[1]}\", img_matches)\n cv2.waitKey(0)\n cv2.destroyAllWindows()\n\n detector = cv2.SIFT_create()\n print('finding keypoints and descriptors...')\n keypoints_and_descriptors = find_keypoints_and_descriptors(images, detector)\n print('matching descriptors...')\n #bf_matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)\n bf_matcher = cv2.BFMatcher(cv2.NORM_L2, crossCheck=True)\n matches = match_descriptors_and_triangulate(keypoints_and_descriptors, bf_matcher, training_frames, transforms)\n if visualize:\n visualize_matches(images, keypoints_and_descriptors, matches, (0, 1))\n\n xyzrgbs = []\n for i in range(len(images)):\n for j in range(i+1, len(images)):\n matches_ij, points, rejected_matches = matches[(i, j)]\n for (k, match) in enumerate(matches_ij):\n p = points[k]\n kp1 = keypoints_and_descriptors[i][0][match.queryIdx].pt\n color = images[i][int(kp1[1]), int(kp1[0]), [2, 1, 0]]\n xyzrgbs.append(p.tolist() + color.tolist())\n print('Triangulated %d points' % len(xyzrgbs))\n\n if not dry_run:\n with open(json_path, 'wt') as f:\n json.dump(converted_json, f, indent=4)\n\n seed_ply_path = os.path.join(target, \"sparse_pc.ply\")\n point_cloud_to_ply(xyzrgbs, seed_ply_path)\n\ndef process_dataset_folder(\n base_folder, \n output_folder,\n subfolder,\n points_only=False,\n noisy_poses=False,\n noisy_intrinsics=False,\n dry_run=False,\n visualize=False):\n items = os.listdir(base_folder)\n directories = sorted([item for item in items if os.path.isdir(os.path.join(base_folder, item))])\n\n for directory in directories:\n print(directory)\n full_path = os.path.join(base_folder, directory, subfolder)\n if not os.path.exists(full_path): continue\n out_path = os.path.join(output_folder, directory)\n if not points_only and not dry_run:\n process(full_path, out_path, noisy_poses=noisy_poses, noisy_intrinsics=noisy_intrinsics)\n if os.path.exists(out_path):\n generate_seed_points_match_and_triangulate(out_path, visualize=visualize)\n\nif __name__ == '__main__':\n import argparse\n parser = argparse.ArgumentParser(__doc__)\n parser.add_argument('--dry_run', action='store_true')\n parser.add_argument('--points_only', action='store_true')\n parser.add_argument('--visualize', action='store_true')\n args = parser.parse_args()\n\n process_dataset_folder(\n 'data/inputs-raw/synthetic-raw',\n 'data/inputs-processed/synthetic-posenoise',\n subfolder='raw_clear',\n noisy_poses=True,\n **vars(args))\n\n process_dataset_folder(\n 'data/inputs-raw/synthetic-raw',\n 'data/inputs-processed/synthetic-rs',\n subfolder='raw_rs',\n **vars(args))\n\n process_dataset_folder(\n 'data/inputs-raw/synthetic-raw',\n 'data/inputs-processed/synthetic-mb',\n subfolder='raw_mb',\n **vars(args))\n\n process_dataset_folder(\n 'data/inputs-raw/synthetic-raw',\n 'data/inputs-processed/synthetic-mb-posenoise',\n subfolder='raw_mb',\n noisy_poses=True,\n **vars(args))\n\n process_dataset_folder(\n 'data/inputs-raw/synthetic-raw',\n 'data/inputs-processed/synthetic-clear',\n subfolder='raw_clear',\n **vars(args))\n\n process_dataset_folder(\n 'data/inputs-raw/synthetic-raw',\n 'data/inputs-processed/synthetic-mbrs',\n subfolder='raw_mbrs',\n **vars(args))\n\n process_dataset_folder(\n 'data/inputs-raw/synthetic-raw',\n 'data/inputs-processed/synthetic-mbrs-posenoise',\n subfolder='raw_mbrs',\n noisy_poses=True,\n **vars(args))\n\n process_dataset_folder(\n 'data/inputs-raw/synthetic-raw',\n 'data/inputs-processed/synthetic-mbrs-pose-calib-noise',\n subfolder='raw_mbrs',\n noisy_poses=True,\n noisy_intrinsics=True,\n **vars(args))"
},
{
"path": "render_model.py",
"content": "\"\"\"Load g model and render all outputs to disc\"\"\"\n\nfrom dataclasses import dataclass\nfrom pathlib import Path\nimport torch\nimport tyro\nimport os\nimport numpy as np\nimport shutil\n\nfrom nerfstudio.cameras.cameras import Cameras\nfrom nerfstudio.models.splatfacto import SplatfactoModel\nfrom nerfstudio.utils.eval_utils import eval_setup\nfrom nerfstudio.utils import colormaps\nfrom nerfstudio.data.datasets.base_dataset import InputDataset\nfrom PIL import Image\nfrom torch import Tensor\n\nfrom typing import List, Literal, Optional, Union\n\ndef save_img(image, image_path, verbose=True) -> None:\n \"\"\"helper to save images\n\n Args:\n image: image to save (numpy, Tensor)\n image_path: path to save\n verbose: whether to print save path\n\n Returns:\n None\n \"\"\"\n if image.shape[-1] == 1 and torch.is_tensor(image):\n image = image.repeat(1, 1, 3)\n if torch.is_tensor(image):\n image = image.detach().cpu().numpy() * 255\n image = image.astype(np.uint8)\n if not Path(os.path.dirname(image_path)).exists():\n Path(os.path.dirname(image_path)).mkdir(parents=True)\n im = Image.fromarray(image)\n if verbose:\n print(\"saving to: \", image_path)\n im.save(image_path)\n\n# Depth Scale Factor m to mm\nSCALE_FACTOR = 0.001\nSAVE_RAW_DEPTH = False\n\ndef save_depth(depth, depth_path, verbose=True, scale_factor=SCALE_FACTOR) -> None:\n \"\"\"helper to save metric depths\n\n Args:\n depth: image to save (numpy, Tensor)\n depth_path: path to save\n verbose: whether to print save path\n scale_factor: depth metric scaling factor\n\n Returns:\n None\n \"\"\"\n if torch.is_tensor(depth):\n depth = depth.float() / scale_factor\n depth = depth.detach().cpu().numpy()\n else:\n depth = depth / scale_factor\n if not Path(os.path.dirname(depth_path)).exists():\n Path(os.path.dirname(depth_path)).mkdir(parents=True)\n if verbose:\n print(\"saving to: \", depth_path)\n np.save(depth_path, depth)\n\ndef save_outputs_helper(\n rgb_out: Optional[Tensor],\n gt_img: Optional[Tensor],\n depth_color: Optional[Tensor],\n depth_gt_color: Optional[Tensor],\n depth_gt: Optional[Tensor],\n depth: Optional[Tensor],\n normal_gt: Optional[Tensor],\n normal: Optional[Tensor],\n render_output_path: Path,\n image_name: Optional[str],\n) -> None:\n \"\"\"Helper to save model rgb/depth/gt outputs to disk\n\n Args:\n rgb_out: rgb image\n gt_img: gt rgb image\n depth_color: colored depth image\n depth_gt_color: gt colored depth image\n depth_gt: gt depth map\n depth: depth map\n render_output_path: save directory path\n image_name: stem of save name\n\n Returns:\n None\n \"\"\"\n if image_name is None:\n image_name = \"\"\n\n if rgb_out is not None and gt_img is not None:\n # easier consecutive compare\n save_img(rgb_out, os.getcwd() + f\"/{render_output_path}/{image_name}_pred.png\", False)\n save_img(gt_img, os.getcwd() + f\"/{render_output_path}/{image_name}_gt.png\", False)\n\n if depth_color is not None:\n save_img(\n depth_color,\n os.getcwd()\n + f\"/{render_output_path}/pred/depth/colorised/{image_name}.png\",\n False,\n )\n if depth_gt_color is not None:\n save_img(\n depth_gt_color,\n os.getcwd() + f\"/{render_output_path}/gt/depth/colorised/{image_name}.png\",\n False,\n )\n if depth_gt is not None:\n # save metric depths\n save_depth(\n depth_gt,\n os.getcwd() + f\"/{render_output_path}/gt/depth/raw/{image_name}.npy\",\n False,\n )\n\n if SAVE_RAW_DEPTH:\n if depth is not None:\n save_depth(\n depth,\n os.getcwd() + f\"/{render_output_path}/pred/depth/raw/{image_name}.npy\",\n False,\n )\n\n if normal is not None:\n save_normal(\n normal,\n os.getcwd() + f\"/{render_output_path}/pred/normal/{image_name}.png\",\n verbose=False,\n )\n\n if normal_gt is not None:\n save_normal(\n normal_gt,\n os.getcwd() + f\"/{render_output_path}/gt/normal/{image_name}.png\",\n verbose=False,\n )\n\n@dataclass\nclass RenderModel:\n \"\"\"Render outputs of a GS model.\"\"\"\n\n load_config: Path = Path(\"outputs/\")\n \"\"\"Path to the config YAML file.\"\"\"\n output_dir: Path = Path(\"./data/renders/\")\n \"\"\"Path to the output directory.\"\"\"\n set: Literal[\"train\", \"eval\"] = \"eval\"\n \"\"\"Dataset to test with (train or eval)\"\"\"\n output_same_dir: bool = True\n \"\"\"Output to the subdirectory of the load_config path\"\"\"\n\n def main(self):\n if self.output_same_dir:\n self.output_dir = os.path.join(os.path.dirname(self.load_config), 'renders')\n\n if os.path.exists(self.output_dir):\n shutil.rmtree(self.output_dir)\n os.makedirs(self.output_dir)\n print('writing %s' % str(self.output_dir))\n\n _, pipeline, _, _ = eval_setup(self.load_config)\n\n assert isinstance(pipeline.model, SplatfactoModel)\n\n model: SplatfactoModel = pipeline.model\n dataset: InputDataset\n\n with torch.no_grad():\n if self.set == \"train\":\n dataset = pipeline.datamanager.train_dataset\n images = pipeline.datamanager.cached_train\n elif self.set == \"eval\":\n dataset = pipeline.datamanager.eval_dataset\n images = pipeline.datamanager.cached_eval\n else:\n raise RuntimeError(\"Invalid set\")\n \n cameras: Cameras = dataset.cameras # type: ignore\n for image_idx in range(len(dataset)): # type: ignore\n data = images[image_idx]\n\n # process batch gt data\n mask = None\n if \"mask\" in data:\n mask = data[\"mask\"]\n\n gt_img = 256 - data[\"image\"] # not sure why negative\n if \"sensor_depth\" in data:\n depth_gt = data[\"sensor_depth\"]\n depth_gt_color = colormaps.apply_depth_colormap(\n data[\"sensor_depth\"]\n )\n else:\n depth_gt = None\n depth_gt_color = None\n if \"normal\" in data:\n normal_gt = data[\"normal\"]\n else:\n normal_gt = None\n\n # process pred outputs\n camera = cameras[image_idx : image_idx + 1].to(\"cpu\")\n #if self.set == \"train\":\n # camera idx is used to fetch camera optimizer adjustments\n # and should not be used for 'eval' data\n camera.metadata['cam_idx'] = image_idx\n outputs = model.get_outputs_for_camera(camera=camera)\n\n rgb_out, depth_out = outputs[\"rgb\"], outputs[\"depth\"]\n\n normal = None\n if \"normal\" in outputs:\n normal = outputs[\"normal\"]\n\n seq_name = Path(dataset.image_filenames[image_idx])\n image_name = f\"{seq_name.stem}\"\n\n depth_color = colormaps.apply_depth_colormap(depth_out)\n depth = depth_out.detach().cpu().numpy()\n\n if mask is not None:\n rgb_out = rgb_out * mask\n gt_img = gt_img * mask\n if depth_color is not None:\n depth_color = depth_color * mask\n if depth_gt_color is not None:\n depth_gt_color = depth_gt_color * mask\n if depth_gt is not None:\n depth_gt = depth_gt * mask\n if depth is not None:\n depth = depth * mask\n if normal_gt is not None:\n normal_gt = normal_gt * mask\n if normal is not None:\n normal = normal * mask\n\n # save all outputs\n save_outputs_helper(\n rgb_out,\n gt_img,\n depth_color,\n depth_gt_color,\n depth_gt,\n depth,\n normal_gt,\n normal,\n self.output_dir,\n image_name,\n )\n\n\nif __name__ == \"__main__\":\n tyro.cli(RenderModel).main()\n"
},
{
"path": "render_video.py",
"content": "\"\"\"Generate demo video camera trajectory\"\"\"\nimport os\nimport json\nimport subprocess\nimport numpy as np\n\nclass SplineInterpolator: \n def __init__(self, target, frames_per_transition):\n self.target = target\n self.positions = []\n self.orientations = []\n self.loop = False\n self.tension = 0.0\n self.model_frame = None\n self.frames_per_transition = frames_per_transition\n\n def push(self, frame):\n from scipy.spatial.transform import Rotation\n m = np.array(frame['camera_to_world'])\n self.positions.append(m[:3, 3].tolist())\n q_xyzw = Rotation.from_matrix(m[:3, :3]).as_quat().tolist()\n self.orientations.append(q_xyzw)\n if self.model_frame is None:\n self.model_frame = frame\n\n def finish(self):\n import splines\n import splines.quaternion\n from scipy.spatial.transform import Rotation\n\n # as in Nerfstudio\n end_cond = \"closed\" if self.loop else \"natural\"\n\n orientation_spline = splines.quaternion.KochanekBartels(\n [\n splines.quaternion.UnitQuaternion.from_unit_xyzw(q)\n for q in self.orientations\n ],\n tcb=(self.tension, 0.0, 0.0),\n endconditions=end_cond,\n )\n\n position_spline = splines.KochanekBartels(\n self.positions,\n tcb=(self.tension, 0.0, 0.0),\n endconditions=end_cond,\n )\n\n n = len(self.positions)\n for t in np.linspace(0, n-1, num=(n-1)*self.frames_per_transition, endpoint=True):\n f = { k: v for k, v in self.model_frame.items() }\n\n q = orientation_spline.evaluate(t)\n p = position_spline.evaluate(t)\n m = np.eye(4)\n m[:3, 3] = p\n m[:3, :3] = Rotation.from_quat([*q.vector, q.scalar]).as_matrix()\n\n f['camera_to_world'] = m.tolist()\n self.target.append(f)\n\ndef look_at(cam_pos, cam_target, up_dir=np.array([0, 0, 1])):\n z = cam_target - cam_pos\n z = z / np.linalg.norm(z)\n x = np.cross(z, up_dir)\n x = x / np.linalg.norm(x)\n y = np.cross(z, x)\n y = y / np.linalg.norm(y)\n m = np.eye(4)\n m[:3, 3] = cam_pos\n m[:3, :3] = np.column_stack((x, -y, -z))\n return m\n\ndef get_original_length_seconds(raw_input_data_jsonl):\n with open(raw_input_data_jsonl, 'rt') as f:\n first_ts = None\n for line in f:\n d = json.loads(line)\n if 'time' in d:\n last_ts = d['time']\n if first_ts is None:\n first_ts = last_ts\n return last_ts - first_ts\n\ndef add_velocities(camera_path, loop=False):\n from scipy.spatial.transform import Rotation\n\n path = camera_path['camera_path']\n for i in range(len(path)):\n if loop:\n i_prev = (i - 1) % len(path)\n i_next = (i + 1) % len(path)\n else:\n i_prev = max(0, i - 1)\n i_next = min(len(path) - 1, i + 1)\n \n delta_t = i_next - i_prev\n\n prev_pose = np.array(path[i_prev]['camera_to_world'])\n next_pose = np.array(path[i_next]['camera_to_world'])\n\n velocity_w = (next_pose[:3, 3] - prev_pose[:3, 3]) / delta_t\n\n cur_pose = np.array(path[i]['camera_to_world'])\n\n rot = next_pose[:3, :3] @ prev_pose[:3, :3].transpose()\n rot_vec = Rotation.from_matrix(rot).as_rotvec()\n ang_vel_w = rot_vec / delta_t\n\n R_w2c = cur_pose[:3, :3].transpose()\n velocity_cam = R_w2c @ velocity_w\n ang_vel_cam = R_w2c @ ang_vel_w\n\n path[i]['camera_linear_velocity'] = velocity_cam.tolist()\n path[i]['camera_angular_velocity'] = ang_vel_cam.tolist()\n\ndef process(out_folder, args):\n import numpy as np\n\n path = os.path.normpath(out_folder)\n name = os.path.basename(path)\n variant_folder = os.path.split(path)[0]\n # variant = os.path.basename(variant_folder)\n dataset_folder = os.path.split(variant_folder)[0]\n dataset = os.path.basename(dataset_folder)\n result_folder = os.path.join(out_folder, 'splatfacto', os.listdir(os.path.join(out_folder, 'splatfacto'))[0])\n config_file = os.path.join(result_folder, 'config.yml')\n\n input_folder = os.path.join('data/inputs-processed', dataset, name)\n\n with open(os.path.join(input_folder, 'transforms.json'), 'rt') as f:\n transforms = json.load(f)\n\n with open(os.path.join(result_folder, 'dataparser_transforms.json'), 'rt') as f:\n parser_transforms = json.load(f)\n\n def transform_func(m):\n if 'applied_transform' in transforms:\n M1 = np.array(transforms['applied_transform'] + [[0,0,0,1]])\n else:\n M1 = np.eye(4)\n M = np.array(parser_transforms['transform'] + [[0,0,0,1]])\n\n m = np.array(m)\n M = M @ np.linalg.inv(M1)\n m = M @ m\n m[:3, 3] *= parser_transforms['scale']\n return m\n\n if args.original_trajectory: \n raw_input_data_jsonl = os.path.join('data', 'inputs-raw', 'spectacular-rec', name, 'data.jsonl')\n \n if os.path.exists(raw_input_data_jsonl):\n length_seconds = get_original_length_seconds(raw_input_data_jsonl)\n print('original length %g' % length_seconds)\n else:\n length_seconds = len(transforms['frames']) * 0.3\n print('approx. length %g' % length_seconds)\n\n length_seconds /= args.playback_speed\n \n def get_frame_number(frame):\n return int(frame['file_path'].rpartition('_')[-1].split('.')[0])\n \n frames = sorted(transforms['frames'], key=get_frame_number)\n frames = frames[::args.key_frame_stride]\n\n if args.max_duration is not None:\n max_frames = round(args.max_duration / length_seconds * len(frames))\n if max_frames < len(frames):\n length_seconds = length_seconds * max_frames / len(frames)\n print('keeping %d/%d key frames to cut duration to %g' % (max_frames, len(frames), length_seconds))\n frames = frames[:max_frames]\n\n frame_poses = [transform_func(frame['transform_matrix']) for frame in frames]\n loop = False\n else:\n length_seconds = args.artificial_length_seconds\n loop = True\n\n rough_up_dir = np.array([0, 0, 1])\n\n frame_poses_np = [transform_func(frame['transform_matrix']) for frame in transforms['frames']]\n scene_cam_center = np.mean([m[:3, 3] for m in frame_poses_np], axis=0)\n scene_cam_mean_dir = np.mean([-m[:3, 2] for m in frame_poses_np], axis=0)\n scene_cam_mean_dir = scene_cam_mean_dir / np.linalg.norm(scene_cam_mean_dir)\n\n scene_scale = np.max([np.linalg.norm(m[:3, 3] - scene_cam_center) for m in frame_poses_np])\n cam_target = scene_cam_center + scene_cam_mean_dir * scene_scale * args.artificial_relative_look_at_distance\n left = np.cross(rough_up_dir, scene_cam_mean_dir)\n left = left / np.linalg.norm(left)\n up = np.cross(scene_cam_mean_dir, left)\n\n up_dim = np.max(np.abs(np.dot([m[:3, 3] - scene_cam_center for m in frame_poses_np], up)))\n left_dim = np.max(np.abs(np.dot([m[:3, 3] - scene_cam_center for m in frame_poses_np], left)))\n \n frame_poses = []\n for t in np.linspace(0, 2*np.pi, endpoint=False, num=100):\n frame_poses.append(look_at(\n scene_cam_center + args.artificial_relative_motion_scale * (\n up_dim * up * np.sin(t * args.artificial_y_rounds) +\n left_dim * left * np.cos(t)\n ),\n cam_target,\n rough_up_dir\n ))\n\n center_cam_to_world = look_at(scene_cam_center, cam_target, rough_up_dir)\n\n fov = 2.0 * np.arctan(0.5 * transforms['h'] / transforms['fl_y']) / np.pi * 180.0 / args.zoom\n frames_per_transition = round((length_seconds * args.fps) / (len(frame_poses) - 1))\n\n width = transforms['w']\n height = transforms['h']\n if args.resolution is not None:\n width, height = [int(x) for x in args.resolution.split('x')]\n\n aspect = width / float(height)\n\n cam_path = {\n 'render_width': width,\n 'render_height': height,\n 'fps': args.fps,\n 'seconds': length_seconds,\n 'camera_path': []\n }\n \n interpolator = SplineInterpolator(cam_path['camera_path'], frames_per_transition=frames_per_transition)\n interpolator.loop = loop\n\n for pose in frame_poses:\n # print(frame['file_path'])\n interpolator.push({\n 'aspect': aspect,\n 'fov': fov,\n 'camera_to_world': pose\n })\n\n interpolator.finish()\n\n add_velocities(cam_path)\n cam_path['rolling_shutter_time'] = args.rolling_shutter_time\n cam_path['exposure_time'] = args.exposure_time\n\n if args.artificial_keep_center_pose:\n for c in cam_path['camera_path']: c['camera_to_world'] = center_cam_to_world.tolist()\n\n trajectory_file = os.path.join(result_folder, 'demo_video_camera_path.json')\n\n if args.output_video_file is None:\n video_fn = ['demo_video']\n if args.rolling_shutter_time > 0:\n video_fn.append('rs')\n if args.exposure_time > 0:\n video_fn.append('mb')\n \n video_file = os.path.join(result_folder, '-'.join(video_fn) + '.mp4')\n else:\n video_file = args.output_video_file\n\n render_cmd = [\n 'ns-render',\n 'camera-path',\n '--load-config', config_file,\n '--camera-path-filename', trajectory_file,\n '--output-path', video_file\n ]\n\n if args.video_crf is not None:\n render_cmd.extend(['--crf', str(args.video_crf)])\n\n if not args.dry_run:\n with open(trajectory_file, 'wt') as f:\n json.dump(cam_path, f, indent=4)\n\n subprocess.check_call(render_cmd)\n\nif __name__ == '__main__':\n import argparse\n parser = argparse.ArgumentParser(description=__doc__)\n\n parser.add_argument(\"input_folder\", type=str, default=None, nargs='?')\n parser.add_argument('--output_variant_folder', default='data/outputs/colmap-sai-cli-imgs/baseline', type=str)\n parser.add_argument('-o', '--output_video_file', default=None, type=str)\n parser.add_argument('--key_frame_stride', default=3, type=int)\n parser.add_argument('--dry_run', action='store_true')\n parser.add_argument('--original_trajectory', action='store_true')\n parser.add_argument('--fps', default=30, type=int)\n parser.add_argument('--playback_speed', default=0.5, type=float)\n parser.add_argument('--artificial_relative_motion_scale', default=0.6, type=float)\n parser.add_argument('--artificial_relative_look_at_distance', default=3, type=float)\n parser.add_argument('--artificial_y_rounds', default=1, type=int)\n parser.add_argument('--artificial_length_seconds', default=8, type=float)\n parser.add_argument('--artificial_keep_center_pose', action='store_true')\n parser.add_argument('--rolling_shutter_time', default=0.0, type=float)\n parser.add_argument('--max_duration', default=None, type=float)\n parser.add_argument('--resolution', type=str, default=None)\n parser.add_argument('--exposure_time', default=0.0, type=float)\n parser.add_argument('--zoom', default=1.0, type=float)\n parser.add_argument('--video_crf', default=None, type=int)\n parser.add_argument('--case_number', type=int, default=-1)\n args = parser.parse_args()\n\n if args.input_folder in ['all']:\n args.case_number = 0\n args.input_folder = None\n\n selected_cases = []\n\n if args.input_folder is None:\n src_folder = args.output_variant_folder\n cases = [os.path.join(src_folder, f) for f in sorted(os.listdir(src_folder))]\n\n if args.case_number == -1:\n print('valid cases')\n for i, c in enumerate(cases): print(str(i+1) + ':\\t' + c)\n elif args.case_number == 0:\n selected_cases = cases\n else:\n selected_cases = [cases[args.case_number - 1]]\n else:\n selected_cases = [args.input_folder]\n\n for case in selected_cases:\n print('Processing ' + case)\n process(case, args)\n"
},
{
"path": "run_colmap.py",
"content": "\"\"\"Run COLMAP on a single sequence through Nerfstudio scripts\"\"\"\nimport os\nimport subprocess\nimport shutil\nimport sys\nimport tempfile\n\ndef process(input_folder, args):\n name = os.path.basename(os.path.normpath(input_folder))\n postf = 'colmap-' + args.dataset + '-imgs'\n if args.output_folder is None:\n output_folder = os.path.join('data/inputs-processed/' + postf, name)\n else:\n output_folder = args.output_folder\n\n input_image_folder = os.path.join(input_folder, 'images')\n\n temp_dir = tempfile.TemporaryDirectory()\n n = 0\n for f in os.listdir(input_image_folder):\n if 'depth' in f: continue\n if not args.dry_run:\n shutil.copyfile(os.path.join(input_image_folder, f), os.path.join(temp_dir.name, f))\n n += 1\n print('%d images (would be) copied in a temporary directory' % n)\n\n # Print the path to the temporary directory\n cmd = [\n 'ns-process-data',\n 'images',\n '--data', temp_dir.name,\n '--output-dir', output_folder\n ]\n\n print(cmd)\n success = False\n if not args.dry_run:\n for itr in range(args.max_retries):\n if os.path.exists(output_folder):\n shutil.rmtree(output_folder)\n\n ret = subprocess.run(cmd, check=True, capture_output=True)\n success = any([b'CONGRATS' in s for s in [ret.stdout, ret.stderr]]) # hacky\n if success:\n break\n else:\n print('COLMAP failed')\n print('--- stdout ---')\n sys.stdout.buffer.write(ret.stdout)\n print('--- stderr ---')\n sys.stderr.buffer.write(ret.stderr)\n if itr != args.max_retries - 1:\n print('Retrying...')\n\n if not success:\n raise RuntimeError('Could not get COLMAP to succeed')\n \n temp_dir.cleanup()\n\nif __name__ == '__main__':\n import argparse\n parser = argparse.ArgumentParser(description=__doc__)\n\n parser.add_argument(\"input_folder\", type=str, default=None, nargs='?')\n parser.add_argument(\"output_folder\", type=str, default=None, nargs='?')\n parser.add_argument('--dry_run', action='store_true')\n parser.add_argument('--dataset', default='sai-cli')\n parser.add_argument('--case_number', type=int, default=-1)\n parser.add_argument('--max_retries', type=int, default=1)\n \n args = parser.parse_args()\n\n if args.input_folder in ['all']:\n args.case_number = 0\n args.input_folder = None\n \n selected_cases = []\n misc = False\n\n PROCESSED_PREFIX = 'data/inputs-processed/'\n if args.dataset.startswith(PROCESSED_PREFIX):\n args.dataset = args.dataset[len(PROCESSED_PREFIX):]\n if args.dataset.endswith('/'): args.dataset = args.dataset[:-1]\n\n if args.input_folder is None:\n input_root = os.path.join(PROCESSED_PREFIX, args.dataset)\n cases = [os.path.join(input_root, f) for f in sorted(os.listdir(input_root))]\n\n if args.case_number == -1:\n print('valid cases')\n for i, c in enumerate(cases): print(str(i+1) + ':\\t' + c)\n elif args.case_number == 0:\n selected_cases = cases\n else:\n selected_cases = [cases[args.case_number - 1]]\n else:\n selected_cases = [args.input_folder]\n\n for case in selected_cases:\n print('Processing ' + case)\n process(case, args)\n"
},
{
"path": "scripts/compile_comparison_video.sh",
"content": "#!/bin/bash\nset -eux\n\nINPUT_BASELINE=\"$1\"\nINPUT_OURS=\"$2\"\nOUTPUT=\"$3\"\n: \"${OURS_NAME:=Deblurred}\"\n\n#: \"${VIDEO_MODE:=HALF}\"\n: \"${VIDEO_MODE:=SWEEP}\"\n\n: \"${DRAW_TEXT:=ON}\"\n: \"${DRAW_BAR:=ON}\"\n: \"${CROP_TO_HD_ASPECT:=ON}\"\n\nif [ $CROP_TO_HD_ASPECT == \"ON\" ]; then\n BASE_FILTER=\"\n [0:v]crop=iw:'min(ih,iw/16*9)'[base];\\\n [1:v]crop=iw:'min(ih,iw/16*9)'[ours]\"\nelse\n BASE_FILTER=\"[0:v]copy[base];[1:v]copy[ours]\"\nfi\nif [ $DRAW_TEXT == \"ON\" ]; then\n BASE_FILTER=\"\n $BASE_FILTER;\\\n [base]drawtext=text='Baseline':fontcolor=white:fontsize=h/50:x=w/50:y=h/50[base];\\\n [ours]drawtext=text='$OURS_NAME':fontcolor=white:fontsize=h/50:x=w-tw-w/50:y=h/50[ours]\"\nfi\nif [ $DRAW_BAR == \"ON\" ]; then\n BASE_FILTER=\"\n $BASE_FILTER;\\\n color=0x80ff80,format=rgba[bar];\\\n [bar][base]scale2ref[bar][base];\\\n [bar]crop=iw:ih/200:0:0[bar];\\\n [ours][bar]overlay=x=0:y=0[ours]\"\nfi\n\ncase $VIDEO_MODE in\n HALF)\n VIDEO_FILTER=\"\n $BASE_FILTER;\\\n [base]crop=iw/2:ih:0:0[left_crop];\\\n [ours]crop=iw/2:ih:iw/2:0[right_crop];\\\n [left_crop][right_crop]hstack\"\n ;;\n\n SWEEP)\n LEN=8\n VIDEO_FILTER=\"\n $BASE_FILTER;\\\n color=0x00000000,format=rgba,scale=[black];\\\n color=0xffffffff,format=rgba[white];\\\n [black][base]scale2ref[black][base];\\\n [white][base]scale2ref[white][base];\\\n [white][black]blend=all_expr='if(lte(X,W*abs(1-mod(T,$LEN)/$LEN*2)),B,A)'[mask];\\\n [ours][mask]alphamerge[overlayalpha]; \\\n [base][overlayalpha]overlay=shortest=1\"\n ;;\n\n *)\n echo -n \"unknown video mode $VIDEO_MODE\"\n exit 1\n ;;\nesac\n\nffmpeg -i \"$INPUT_BASELINE\" -i \"$INPUT_OURS\" -filter_complex \"$VIDEO_FILTER\" -hide_banner -y \"$OUTPUT\""
},
{
"path": "scripts/install.sh",
"content": "#!/bin/bash\nset -eux\n\n: \"${BUILD_NERFSTUDIO:=ON}\"\n: \"${INSTALL_SAI:=ON}\"\n\n# You may also need to run this\n# pip install --upgrade pip setuptools\n\nif [ $BUILD_NERFSTUDIO == \"ON\" ]; then\n # Install the custom fork of Nerfstudio\n cd nerfstudio\n pip install -e .\n cd ..\nfi\n\n# ... then install the custom gsplat (order may matter here!)\nif [ $BUILD_NERFSTUDIO == \"ON\" ]; then\n cd gsplat\n pip install -e .\n cd ..\nfi\n\nif [ $INSTALL_SAI == \"ON\" ]; then\n pip install spectacularAI[full]==1.31.0\nfi"
},
{
"path": "scripts/process_and_train_sai_custom.sh",
"content": "#!/bin/bash\n\n# Process and train a custom recording created with Spectacular Rec.\n#\n# This version uses both motion blur compensation only and should work\n# well with iPhone data and other devices with short rolling shutter\n# readout times (or global shutter cameras)\n#\n# Run as\n#\n# ./scripts/process_and_train_sai_custom.sh /PATH/TO/RECORDING.zip\n#\n# or, in headless mode\n#\n# SAI_PREVIEW=OFF ./scripts/process_and_train_sai_custom.sh \\\n# /PATH/TO/RECORDING.zip\n\nset -eux\n\nNAME_W_EXT=`basename \"$1\"`\nNAME=${NAME_W_EXT%.zip}\n\n: \"${SAI_PREVIEW:=ON}\"\n: \"${SKIP_COLMAP:=OFF}\"\nif [ $SAI_PREVIEW == \"ON\" ]; then\n PREVIEW_FLAG=\"--preview\"\nelse\n PREVIEW_FLAG=\"\"\nfi\nif [ $SKIP_COLMAP == \"ON\" ]; then\n COLMAP_FLAG=\"--skip_colmap\"\nelse\n COLMAP_FLAG=\"\"\nfi\n\npython process_sai_custom.py \"$1\" $COLMAP_FLAG $PREVIEW_FLAG\npython train.py data/inputs-processed/custom/$NAME --no_eval --train_all $PREVIEW_FLAG"
},
{
"path": "scripts/process_and_train_video.sh",
"content": "#!/bin/bash\n\n# Process and train deblurred 3DGS from a video.\n# set ROLLING_SHUTTER=ON to train a rolling shutter compensated model instead\n# of a deblurred one. For simultaneous MB and RS compensation, see\n# process_and_train_sai_custom.sh\n\nset -eux\n\nNAME_W_EXT=`basename \"$1\"`\nNAME=\"${NAME_W_EXT%.*}\"\n\n: \"${ROLLING_SHUTTER:=OFF}\"\n\n: \"${PREVIEW:=ON}\"\nif [ $PREVIEW == \"ON\" ]; then\n PREVIEW_FLAG=\"--preview\"\nelse\n PREVIEW_FLAG=\"\"\nfi\n\nif [ $ROLLING_SHUTTER == \"ON\" ]; then\n MODE_FLAGS=\"--no_motion_blur\"\nelse\n MODE_FLAGS=\"--no_rolling_shutter\"\nfi\n\nmkdir -p \"data/inputs-processed/custom\"\nTARGET_DIR=\"data/inputs-processed/custom/$NAME\"\n\nns-process-data video --num-frames-target 100 --data \"$1\" --output-dir \"$TARGET_DIR\"\npython train.py \"$TARGET_DIR\" $MODE_FLAGS --velocity_opt_zero_init --train_all --no_eval $PREVIEW_FLAG"
},
{
"path": "scripts/process_smartphone_dataset.sh",
"content": "#!/bin/bash\nset -eux\n\n# Process raw input data. If set to OFF, then sai-cli and\n# colmap-sai-cli-imgs intermediary datasets must have been\n# fully generated or downloaded\n: \"${PROCESS_RAW:=ON}\"\n\n# Extra variants in the supplementary\n: \"${EXTRA_VARIANTS:=OFF}\"\n\n# Show preview in sai-cli\n: \"${PREVIEW:=ON}\"\n\nif [ $PREVIEW == \"ON\" ]; then\n PREVIEW_FLAG=\"--preview\"\nelse\n PREVIEW_FLAG=\"\"\nfi\n\nif [ $PROCESS_RAW == \"ON\" ]; then\n\t# Process and convert using the Spectacular AI SDK to get VIO velocity and pose estimates\n\tpython process_sai_inputs.py $PREVIEW_FLAG\n\t# you can also run individual failing cases with: python run_colmap.py all --case=N\n\tpython run_colmap.py all --max_retries=10\nfi\n\nrm -rf data/inputs-processed/colmap-sai-cli-vels*\nrm -rf data/inputs-processed/colmap-sai-cli-orig-intrinsics*\nrm -rf data/inputs-processed/sai-cli-blur-scored\n\n# --- real data, COLMAP intrinsics\npython combine.py all\npython train_eval_split_by_blur_score.py colmap-sai-cli-vels all\n\n# --- real data, factory intrinsics\npython combine.py --keep_intrinsics all\npython train_eval_split_by_blur_score.py colmap-sai-cli-orig-intrinsics all\n\n# --- real data, calibrated intrinsics\nrm -rf data/inputs-processed/colmap-sai-cli-calib-intrinsics*\n\nfor i in 1 2 3 4 5; do\n\tpython combine.py --case=$i --keep_intrinsics --set_rolling_shutter_to=0.005\ndone\n\nfor i in 6 7 8; do\n\tpython combine.py --case=$i --override_calibration=data/inputs-raw/spectacular-rec-extras/calibration/manual-calibration-result-pixel5.json\ndone\n\nfor i in 9 10 11; do\n\tpython combine.py --case=$i --override_calibration=data/inputs-raw/spectacular-rec-extras/calibration/manual-calibration-result-s20.json\ndone\n\npython train_eval_split_by_blur_score.py colmap-sai-cli-calib-intrinsics all\n\nif [ $EXTRA_VARIANTS == \"ON\" ]; then\n\trm -rf data/inputs-processed/colmap-sai-cli-no-blur-select-imgs*\n\n\tif [ $PROCESS_RAW == \"ON\" ]; then\n\t\t# --- real data, no blur score filter\n\t\tpython process_sai_inputs.py --no_blur_score_filter $PREVIEW_FLAG\n\tfi\n\n\t# NOTE: run this until success\n\tpython run_colmap.py --dataset=sai-cli-no-blur-select all --max_retries=10\n\n\t# --- real data, no blur score filter, COLMAP intrinsics\n\tpython combine.py --dataset=sai-cli-no-blur-select all\n\n\t# --- real data, no blur score filter, factory intrinsics\n\tpython combine.py --keep_intrinsics --dataset=sai-cli-no-blur-select all\nfi\n\n"
},
{
"path": "scripts/render_and_compile_comparison_video.sh",
"content": "#!/bin/bash\nset -eux\n\nINPUT_BASE=\"$1\"\nINPUT_OURS=\"$2\"\n\n# zoom 2x original focal length to highlight details, slow speed (approx.)\nRENDER_ARGS=\"--zoom=1.5 --original_trajectory --playback_speed=0.25\"\n\nNAME=`basename \"$INPUT_BASE\"`\n\nmkdir -p data/renders\n\nBASE_VID=\"data/renders/$NAME-baseline.mp4\"\nOURS_VID=\"data/renders/$NAME-deblurred.mp4\"\nCOMP_VID=\"data/renders/$NAME-comparison.mp4\"\n\npython render_video.py $RENDER_ARGS \"$INPUT_BASE\" -o \"$BASE_VID\"\npython render_video.py $RENDER_ARGS \"$INPUT_OURS\" -o \"$OURS_VID\"\n\n./scripts/compile_comparison_video.sh \"$BASE_VID\" \"$OURS_VID\" \"$COMP_VID\""
},
{
"path": "scripts/render_and_train_comparison_sai_custom.sh",
"content": "#!/bin/bash\n\n# Process and train a custom recording created with Spectacular Rec.\n# Trains two versions: baseline and deblurred and renders a video that\n# shows their differences. With normal, not-very-blurry recordings, the\n# expected improvement is subtle but noticeable.\n\nset -eu\n\nNAME_W_EXT=`basename \"$1\"`\nNAME=${NAME_W_EXT%.zip}\n\necho \"============= Training motion-blur compensated model ==========\"\n# Note: do not set SKIP_COLMAP here: the 3DGS reconstruction may work\n# fine but the comparison video will often be misaligned\n./scripts/process_and_train_sai_custom.sh \"$1\"\n\necho \"============= Training baseline model ==========\"\npython train.py data/inputs-processed/custom/$NAME \\\n --no_eval --train_all --no_rolling_shutter --no_pose_opt --no_motion_blur --no_velocity_opt --preview\n\necho \"============= Rendering comparison video ==========\"\n./scripts/render_and_compile_comparison_video.sh \\\n \"data/outputs/custom/baseline/$NAME\" \\\n \"data/outputs/custom/pose_opt-motion_blur-rolling_shutter-velocity_opt/$NAME\"\n\necho \"Success: see data/renders/$NAME-comparison.mp4\""
},
{
"path": "scripts/render_and_train_comparison_video.sh",
"content": "#!/bin/bash\n\n# Process and train 3DGS from a video with and without deblurring\n# (or rolling shutter compensation if ROLLINGS_SHUTTER=ON) and\n# render a comparison video\n\nset -eux\n\nNAME_W_EXT=`basename \"$1\"`\nNAME=\"${NAME_W_EXT%.*}\"\n\n: \"${ROLLING_SHUTTER:=OFF}\"\n\nif [ $ROLLING_SHUTTER == \"ON\" ]; then\n MODE_NAME=\"rolling_shutter\"\n export OURS_NAME=\"Compensated\"\nelse\n MODE_NAME=\"motion_blur\"\n export OURS_NAME=\"Deblurred\"\nfi\nexport ROLLING_SHUTTER\n\necho \"============= Training $MODE_NAME compensated model ==========\"\n./scripts/process_and_train_video.sh \"$1\"\n\necho \"============= Training baseline model ==========\"\nTARGET_DIR=\"data/inputs-processed/custom/$NAME\"\npython train.py \"$TARGET_DIR\" --no_rolling_shutter --no_pose_opt \\\n --no_motion_blur --no_velocity_opt --train_all --no_eval\n\necho \"============= Rendering comparison video ==========\"\n./scripts/render_and_compile_comparison_video.sh \\\n \"data/outputs/custom/baseline/$NAME\" \\\n \"data/outputs/custom/pose_opt-${MODE_NAME}-velocity_opt-zero_init/$NAME\"\n\necho \"Success: see data/renders/$NAME-comparison.mp4\""
},
{
"path": "train.py",
"content": "\"\"\"Train a single instance\"\"\"\nimport os\nimport subprocess\nimport shutil\nimport sys\nimport time\nimport datetime\nimport json\nimport re\n\nDATASET_SPECIFIC_PARAMETERS = {\n r\".*synthetic.*\": [\n # '--max-num-iterations', '20000', # this would be enough, usually\n '--pipeline.model.num-downscales', '0', # low resolution -> no downscaling\n # These help reconstructing large areas with very smooth color,\n # i.e., the synthetic sky. With defaults, large holes can easily appear\n '--pipeline.model.background-color', 'auto',\n '--pipeline.model.cull-scale-thresh', '2.0',\n # Evaluation data is known to be static. Don't try to optimize camera velocities\n '--pipeline.model.optimize-eval-velocities=False',\n # Hight motion blur, needs more samples\n '--pipeline.model.blur-samples=10',\n ]\n}\n\ndef print_cmd(cmd):\n print('RUNNING COMMAND: ' + ' '.join(cmd))\n\ndef flags_to_variant_name_and_cmd(args): \n cmd = []\n variant = []\n\n use_gamma_correction = False\n optimize_eval_cameras = False\n\n if not args.get('no_pose_opt', False):\n optimize_eval_cameras = True\n variant.append('pose_opt')\n cmd.extend([\n '--pipeline.model.camera-optimizer.mode=SO3xR3',\n ## '--pipeline.model.sh-degree=0'\n ])\n\n if not args.get('no_motion_blur', False):\n variant.append('motion_blur')\n # default blur samples: 5\n use_gamma_correction = not args.get('no_gamma', False)\n if not use_gamma_correction:\n variant.append('no_gamma')\n else:\n cmd.append('--pipeline.model.blur-samples=0')\n\n if not args.get('no_rolling_shutter', False):\n variant.append('rolling_shutter')\n else:\n cmd.append('--pipeline.model.rolling-shutter-compensation=False')\n\n if use_gamma_correction:\n # min RGB level only seems necessary with gamma correction\n cmd.append('--pipeline.model.min-rgb-level=10')\n else:\n cmd.append('--pipeline.model.gamma=1')\n\n if not args.get('no_velocity_opt', False):\n optimize_eval_cameras = True\n cmd.append('--pipeline.model.camera-velocity-optimizer.enabled=True')\n variant.append('velocity_opt')\n\n if args.get('velocity_opt_zero_init', False):\n cmd.append('--pipeline.model.camera-velocity-optimizer.zero-initial-velocities=True')\n variant.append('zero_init')\n\n if len(variant) == 0:\n variant.append('baseline')\n\n return '-'.join(variant), cmd, optimize_eval_cameras\n\ndef evaluate(output_folder, elapsed_time, dry_run=False, render_images=True):\n result_paths = find_config_path(output_folder)\n if result_paths is None:\n if dry_run: return\n assert(False)\n\n out_path, config_path = result_paths\n metrics_path = os.path.join(out_path, 'metrics.json')\n elapsed_time\n eval_cmd = [\n 'ns-eval',\n '--load-config', config_path,\n '--output-path', metrics_path\n ]\n\n print_cmd(eval_cmd)\n if not dry_run:\n subprocess.check_call(eval_cmd)\n with open(metrics_path) as f:\n metrics = json.load(f)\n metrics['wall_clock_time_seconds'] = elapsed_time\n with open(metrics_path, 'w') as f:\n json.dump(metrics, f, indent=4)\n \n if render_images:\n render_cmd = [\n 'python', 'render_model.py',\n '--load-config', config_path\n ]\n print_cmd(render_cmd)\n if not dry_run:\n subprocess.check_call(render_cmd)\n\ndef process(input_folder, args):\n name = os.path.split(input_folder)[-1]\n\n cmd = [\n 'ns-train',\n 'splatfacto',\n '--data', input_folder,\n '--viewer.quit-on-train-completion', 'True',\n '--pipeline.model.rasterize-mode', 'antialiased',\n '--pipeline.model.use-scale-regularization', 'True',\n # '--logging.local-writer.max-log-size=0'\n ]\n\n for pattern, values in DATASET_SPECIFIC_PARAMETERS.items():\n if re.match(pattern, args.dataset):\n cmd.extend(values)\n\n if '--max-num-iterations' not in cmd:\n if args.draft:\n cmd.extend(['--max-num-iterations', '3000'])\n else:\n cmd.extend(['--max-num-iterations', '20000'])\n\n if args.preview:\n cmd.extend([\n '--vis=viewer+tensorboard',\n '--viewer.websocket-host=127.0.0.1'\n ])\n else:\n cmd.append('--vis=tensorboard')\n\n variant, variant_cmd, optimize_eval_cameras = flags_to_variant_name_and_cmd(vars(args))\n cmd.extend(variant_cmd)\n\n if args.case_number is None:\n dataset_folder = 'custom'\n else:\n dataset_folder = args.dataset\n \n variant_folder = os.path.join(dataset_folder, variant)\n\n output_prefix = 'data/outputs'\n\n # note: 'name' is automatically added by Nerfstudio\n output_root = os.path.join(output_prefix, variant_folder)\n\n cmd.extend(['--output-dir', output_root])\n\n cmd.extend([\n 'nerfstudio-data',\n '--orientation-method', 'none',\n ])\n\n if args.train_all:\n cmd.extend([\n '--eval-mode', 'all'\n ])\n optimize_eval_cameras = False\n elif '-scored' in args.input_folder or args.dataset == 'colmap-bad-nerf-synthetic-deblurring':\n cmd.extend([\n '--eval-mode', 'filename'\n ])\n else:\n cmd.extend([\n '--eval-mode', 'interval',\n '--eval-interval', '8'\n ])\n #cmd.extend(['--eval-mode', 'all'])\n\n if optimize_eval_cameras:\n cmd.extend([\n '--optimize-eval-cameras', 'True',\n ])\n\n print_cmd(cmd)\n output_folder = os.path.join(output_root, name)\n elapsed_time = 0\n if not args.dry_run and not args.eval_only:\n if os.path.exists(output_folder):\n shutil.rmtree(output_folder)\n\n start_time = time.time()\n subprocess.check_call(cmd)\n end_time = time.time()\n elapsed_time = end_time - start_time\n print('Training time: %s' % str(datetime.timedelta(seconds=elapsed_time)))\n \n if not args.no_eval:\n evaluate(output_folder, elapsed_time,\n dry_run=args.dry_run,\n render_images=args.render_images)\n\ndef find_config_path(output_folder):\n model_folder = os.path.join(output_folder, 'splatfacto')\n paths = []\n if os.path.exists(model_folder):\n for subdir in os.listdir(model_folder):\n out_path = os.path.join(model_folder, subdir)\n config_path = os.path.join(out_path, 'config.yml')\n if os.path.exists(config_path):\n paths.append((out_path, config_path))\n if len(paths) == 0: return None\n assert(len(paths) == 1)\n return paths[0]\n\ndef add_velocity_opt_variants(variants, dataset):\n has_velocity_info = ('sai-' in dataset\n or 'spectacular-rec' in dataset\n or ('synthetic-' in dataset and 'colmap' not in dataset and 'hloc' not in dataset)\n )\n\n new_variants = []\n for v in variants:\n v1 = v.copy()\n no_velocity_to_optimize = 'no_rolling_shutter' in v and 'no_motion_blur' in v\n if has_velocity_info or no_velocity_to_optimize:\n v1.add('no_velocity_opt')\n new_variants.append(v1)\n\n if no_velocity_to_optimize: continue\n\n if has_velocity_info:\n new_variants.append(v)\n\n v2 = v.copy()\n v2.add('velocity_opt_zero_init')\n new_variants.append(v2)\n\n return new_variants\n\nif __name__ == '__main__':\n import argparse\n parser = argparse.ArgumentParser(description=__doc__)\n\n # note: velocity optimization arguments are auto-added to all of these\n baseline = {\n 'no_pose_opt',\n 'no_motion_blur',\n 'no_rolling_shutter'\n }\n\n no_rolling_shutter_variants = [\n baseline,\n { 'no_rolling_shutter', 'no_pose_opt' },\n { 'no_rolling_shutter', 'no_motion_blur' },\n { 'no_rolling_shutter' }\n ]\n \n full_variants = no_rolling_shutter_variants + [\n { 'no_pose_opt', 'no_motion_blur' },\n { 'no_pose_opt' },\n { 'no_motion_blur' },\n set([])\n ]\n\n default_variants = full_variants\n bad_nerf_variants = [\n baseline,\n { 'no_rolling_shutter', 'no_pose_opt' },\n { 'no_rolling_shutter' }\n ]\n\n add_popt = lambda a: a + [o - {'no_pose_opt'} for o in a if 'no_pose_opt' in o]\n\n variants_by_dataset = {\n 'synthetic-clear': [\n baseline\n ],\n 'synthetic-mb': add_popt([\n baseline,\n { 'no_pose_opt', 'no_rolling_shutter' }\n ]),\n 'synthetic-rs': add_popt([\n baseline,\n { 'no_pose_opt', 'no_motion_blur' }\n ]),\n 'synthetic-posenoise': add_popt([\n baseline,\n { 'no_rolling_shutter', 'no_motion_blur' }\n ]),\n 'synthetic-mbrs': add_popt([\n baseline,\n { 'no_pose_opt' },\n { 'no_pose_opt', 'no_motion_blur' },\n { 'no_pose_opt', 'no_rolling_shutter' }\n ]),\n 'synthetic-posenoise-2nd-pass': [\n baseline\n ],\n 'colmap-bad-nerf-synthetic-deblurring': bad_nerf_variants,\n 'colmap-bad-nerf-synthetic-novel-view': bad_nerf_variants,\n 'colmap-bad-nerf-synthetic-novel-view-manual-pc': add_popt(bad_nerf_variants),\n 'colmap-exblurf-synthetic-novel-view-manual-pc': bad_nerf_variants,\n 'hloc-exblurf-synthetic-novel-view-manual-pc': bad_nerf_variants,\n 'hloc-bad-nerf-synthetic-novel-view-manual-pc': bad_nerf_variants,\n 'hloc-bad-nerf-synthetic-novel-view-exact-intrinsics-manual-pc': bad_nerf_variants,\n 'hloc-bad-gaussians-synthetic-novel-view-manual-pc': bad_nerf_variants,\n 'colmap-bad-gaussians-synthetic-novel-view-manual-pc': bad_nerf_variants,\n 'colmap-mpr-deblurred-synthetic-all-manual-pc': bad_nerf_variants,\n 'colmap-mpr-deblurred-synthetic-novel-view-manual-pc': bad_nerf_variants + [{ 'no_rolling_shutter', 'no_motion_blur' }],\n }\n\n parser.add_argument(\"input_folder\", type=str, default=None, nargs='?')\n parser.add_argument(\"--preview\", action='store_true', help='show Viser preview')\n parser.add_argument(\"--no_pose_opt\", action='store_true')\n parser.add_argument(\"--no_motion_blur\", action='store_true')\n parser.add_argument('--no_rolling_shutter', action='store_true')\n parser.add_argument('--no_velocity_opt', action='store_true')\n parser.add_argument('--velocity_opt_zero_init', action='store_true')\n parser.add_argument('--dataset', type=str, default='colmap-sai-cli-vels-blur-scored')\n parser.add_argument('--draft', action='store_true')\n parser.add_argument('--no_gamma', action='store_true')\n parser.add_argument('--dry_run', action='store_true')\n parser.add_argument('--render_images', action='store_true')\n parser.add_argument('--eval_only', action='store_true')\n parser.add_argument('--no_eval', action='store_true')\n parser.add_argument('--train_all', action='store_true')\n\n parser.add_argument('--case_number', type=int, default=None)\n args = parser.parse_args()\n\n if args.input_folder is None and args.case_number is None:\n args.case_number = -1\n\n if args.case_number is not None:\n INPUT_ROOT = 'data/inputs-processed/' + args.dataset\n sessions = [os.path.join(INPUT_ROOT, f) for f in sorted(os.listdir(INPUT_ROOT))]\n variants = add_velocity_opt_variants(variants_by_dataset.get(args.dataset, default_variants), args.dataset)\n cases = [(s, v) for v in variants for s in sessions]\n\n if args.case_number <= 0:\n print('valid cases')\n for i, (c, v) in enumerate(cases):\n variant = flags_to_variant_name_and_cmd({k: True for k in v})[0]\n print(str(i+1) + ':\\t' + variant + '\\t' + c)\n sys.exit(0)\n else:\n args.input_folder, variant = cases[args.case_number - 1]\n for p in variant: setattr(args, p, True)\n print('Running %s %s' % (args.input_folder, str(variant)))\n\n process(args.input_folder, args)\n"
},
{
"path": "train_eval_split_by_blur_score.py",
"content": "\"\"\"Combine COLMAP poses with sai-cli velocities\"\"\"\nimport os\nimport json\nimport shutil\n\ndef process(input_folder, output_prefix, args):\n name = os.path.basename(os.path.normpath(input_folder))\n print('name', name)\n\n def read_json(folder):\n with open(os.path.join(folder, 'transforms.json')) as f:\n return json.load(f)\n\n print(input_folder)\n output_folder = os.path.join(output_prefix, name)\n\n input_image_folder = os.path.join(input_folder, 'images')\n output_image_folder = os.path.join(output_folder, 'images')\n \n poses = read_json(input_folder)\n poses['frames'].sort(key=lambda x: x['file_path'])\n\n if not args.dry_run:\n if os.path.exists(output_folder): shutil.rmtree(output_folder)\n os.makedirs(output_image_folder)\n \n ival_start = 0\n while ival_start < len(poses['frames']):\n ival_end = ival_start + args.interval\n least_blur = sorted(poses['frames'][ival_start:ival_end], key=lambda x: x['motion_blur_score'])[0]['file_path']\n\n for frame in poses['frames'][ival_start:ival_end]:\n id = frame['file_path']\n if id == least_blur:\n new_name = f'eval_' + os.path.basename(id)\n else:\n new_name = f'train_' + os.path.basename(id)\n\n old_file_name = os.path.join(input_image_folder, os.path.basename(id))\n new_file_name = os.path.join(output_image_folder, new_name)\n\n frame['file_path'] = os.path.join('images', new_name)\n print(\"%s -> %s (%g)\" % (old_file_name, new_file_name, frame['motion_blur_score']))\n if not args.dry_run:\n shutil.copyfile(old_file_name, new_file_name)\n\n ival_start = ival_end\n\n # colmap_folder = os.path.join(args.input_folder, 'colmap')\n ply_pc = os.path.join(input_folder, 'sparse_pc.ply')\n\n print('Output folder: ' + output_folder)\n if not args.dry_run:\n # shutil.copytree(colmap_folder, os.path.join(output_folder, 'colmap'))\n shutil.copyfile(ply_pc, os.path.join(output_folder, 'sparse_pc.ply'))\n with open (os.path.join(output_folder, 'transforms.json'), 'w') as f:\n json.dump(poses, f, indent=4)\n\nif __name__ == '__main__':\n import argparse\n parser = argparse.ArgumentParser(description=__doc__)\n\n parser.add_argument('dataset')\n parser.add_argument(\"input_folder\", type=str, default=None, nargs='?')\n parser.add_argument('--interval', type=int, default=8)\n parser.add_argument('--dry_run', action='store_true')\n parser.add_argument('--case_number', type=int, default=-1)\n args = parser.parse_args()\n\n if args.input_folder in ['all']:\n args.case_number = 0\n args.input_folder = None\n\n selected_cases = []\n\n PROCESSED_PREFIX = 'data/inputs-processed/'\n if args.dataset.startswith(PROCESSED_PREFIX):\n args.dataset = args.dataset[len(PROCESSED_PREFIX):]\n\n out_folder = os.path.join(PROCESSED_PREFIX, args.dataset + '-blur-scored')\n\n if args.input_folder is None:\n processed_prefix = os.path.join(PROCESSED_PREFIX, args.dataset)\n cases = [os.path.join(processed_prefix, f) for f in sorted(os.listdir(processed_prefix))]\n\n if args.case_number == -1:\n print('valid cases')\n for i, c in enumerate(cases): print(str(i+1) + ':\\t' + c)\n elif args.case_number == 0:\n selected_cases = cases\n else:\n selected_cases = [cases[args.case_number - 1]]\n else:\n selected_cases = [args.input_folder]\n\n for case in selected_cases:\n print('Processing ' + case)\n process(case, out_folder, args)\n"
}
]