Implementation of Thermal Restoration NeRF (TRNeRF)
[project page] [paper (IEEE)] [paper (CVF, open access)]
The dataset is available through the project page. See the Data Format documentation page for a detailed description of the files. We also provide scripts for converting the data from the H5 format to ROS1 rosbags or common formats (CSV files and folders of PNG images). See the Conversion documentation page for more details.
To run training and evaluation (except the detection metric) with the recommended settings on all sequences, the minimally required subset of the dataset and the expected file structure is as follows:
/path/to/trnerf_dataset_folder/
├── cam_calibration
│ └── camchain.yaml
├── slow_indoor
│ ├── adk_right.h5
│ └── pseudo_ground_truth_poses.h5
├── slow_outdoor
│ └── ...
├── medium_indoor
│ ├── adk_right.h5
│ ├── pseudo_ground_truth_adk_right.h5
│ └── pseudo_ground_truth_poses.h5
├── medium_outdoor
│ └── ...
├── fast_indoor
│ └── ...
└── fast_outdoor
└── ...
Additionally:
- Computing the detection metric requires the
cam_calibration/aprilgrid.yamlfile and themono_left.h5file for each non-calibration sequence. - Running the full ablation study requires the
adk_left.h5file for each non-calibration sequence. - Running pseudo ground truth pose generation and pseudo ground truth image generation requires the
mono_left.h5,mono_right.h5, andadk_left.h5files for each non-calibration sequence.
In this section we describe how to perform training with TRNeRF.
We recommend running the code with Docker, as explained below. However, if you already have nerfstudio installed, it may be possible to install TRNeRF as follows:
cd /path/to/trnerf/trnerf_utils/
pip3 install -e .
cd ../trnerf_main/
pip3 install -e .
ns-install-cli
You can check the install by running ns-train -h and confirming that trnerf appears in the list of subcommands.
To run TRNeRF with Docker, first ensure Docker is installed and has access to your NVIDIA GPU(s). On Ubuntu, this can be done by installing Docker Engine, performing this post-installation step, updating your NVIDIA drivers, and finally installing the NVIDIA Container Toolkit.
After installing Docker, a Docker image for running TRNeRF can be built as follows:
cd /path/to/trnerf/
docker build \
--build-arg CUDA_ARCHITECTURES=<compute capability> \
--build-arg USER_ID=$(id -u) \
--tag trnerf_training \
--file docker/Dockerfile_nerfstudio \
.
where <compute capability> should be replaced with the compute capability of your GPU(s) with the decimal point removed (e.g., if the compute capability is 8.6, then put 86 here). The result will be a new Docker image named trnerf_training.
Run the Docker image created in the previous section with the following command:
docker run -it --gpus '"device=0"' -p 7007:7007 --shm-size=12gb -v ~/.cache/:/home/user/.cache -v /path/to/trnerf_dataset_folder/:/data -v /path/to/output_folder/:/output trnerf_training
where /path/to/trnerf_dataset_folder/ should be replaced with the path to the downloaded TRNeRF dataset (see Dataset Download, Format, and Conversion) and /path/to/output_folder/ should be replaced with the folder you would like training output to be stored in. The folders will appear within the container as /data and /output. The command -p 7007:7007 binds port 7007 of the container to that of the host as the nerfstudio viewer uses port 7007 by default. You can also change --gpus '"device=0"' to another GPU (e.g., --gpus '"device=2"'), a list of GPUs (e.g., --gpus '"device=0, 2"'), or all GPUs (--gpus all).
The command above will start an interactive bash session in the running container. From here, you can view the TRNeRF and TRNeRF dataparser arguments with ns-train trnerf -h and ns-train trnerf trnerf_dataparser -h, respectively. To run TRNeRF training with recommended settings, we've provided the run_training_rec.sh bash script that can be executed in the container as follows:
/home/user/trnerf/scripts/training/run_training_rec.sh \
/output/ \
/data/ \
<speed> \
<scene>
where <speed> can be slow, medium, or fast and <scene> can be indoor or outdoor. The result will be a trial folder /output/right_<speed>_<scene>_recommended/trnerf/<date and time>/ that contains configuration information (config.yml and dataparser_transforms.json) and the trained model (nerfstudio_models/step-*.ckpt). If training crashes due to memory limitations, see Reducing CPU & GPU Memory Consumption below.
We've also provided the run_training_rec_all.sh and run_training_ablations.sh bash scripts in the folder scripts/training to perform training on all scenes with recommended settings and to perform training with all of the settings and data tested in the ablation study, respectively. The run_training_rec_all.sh bash script can be executed in the container as follows:
/home/user/trnerf/scripts/training/run_training_rec_all.sh \
/output/ \
/data/
and run_training_ablations.sh can be run with the same arguments.
If you are training locally, you can use the nerfstudio viewer by opening a browser and going to localhost:7007. If you are training on a remote machine, run the following in a terminal window on your local machine:
ssh -L 7007:localhost:7007 <username>@<training-host-ip>
and you should now be able to use the viewer by navigating to localhost:7007 as described above.
If you cannot see the camera poses, uncheck Composite depth. You can click on a camera pose to jump to its view. It may take a while for training to converge to something visually recognizable.
By default, the visualization will be scaled between the minimum and maximum intensities observed in the scene. To switch to a higher contrast visualization select viewer_output in the Output type dropdown.
To use the viewer after training, run the following from within the Docker container:
ns-viewer --load-config /path/to/config.yml
where /path/to/config.yml should be replaced with the path to the config.yml file in the trial folder that was generated during training.
The CPU and GPU memory consumption can be reduced by editing the run_training.sh bash script.
To reduce the CPU memory usage you can set --pipeline.datamanager.train-num-images-to-sample-from to a lower value. You may also want to increase --pipeline.datamanager.train-num-times-to-repeat-images accordingly.
To reduce the GPU memory consumption, we suggest:
- setting
--mixed-precisiontoFalse - multiplying the
--gradient-accumulation-stepsfor both fields by some factor - dividing the
--pipeline.datamanager.train-num-rays-per-batchby the same factor - multiplying the
--max-num-iterationsby the same factor
In this section we describe how to render restored images, compute metrics from the restored images, and compile and visualize the results.
Performing the evaluation requires two Docker images. The first is the trnerf_training Docker image described above for training. The second Docker image has Kalibr installed for Aprilgrid detection and can be built as follows:
cd /path/to/trnerf/
docker build \
--build-arg USER_ID=$(id -u) \
--tag trnerf_detection_metric \
--file docker/Dockerfile_kalibr \
.
The result will be a new Docker image named trnerf_detection_metric.
Run the trnerf_training Docker image as described above. To render a restored image for each training image in a sequence, run the following in the container:
python3 /home/user/trnerf/trnerf_main/trnerf/render_restored.py --path-config /path/to/config.yml
where /path/to/config.yml should be replaced with the path to the config.yml file in the trial folder that was generated during training. The result will be a file, restored.h5, that contains the restored images, in the same folder as config.yml.
We've provided the run_rendering_all.sh bash script to perform rendering for all trials. It can be run as follows:
/home/user/trnerf/scripts/evaluation/rendering/run_rendering_all.sh /output/
Run the trnerf_training Docker image as described above. You can view the arguments for computing the LPIPS metric with python3 /home/user/trnerf/evaluation/compute_lpips_metric.py -h. To compute the LPIPS metric for a single trial, we've provided the run_lpips_metric_trial.sh bash script that can be executed in the container as follows:
/home/user/trnerf/scripts/evaluation/computing_metrics/run_lpips_metric_trial.sh /path/to/trial_folder/ /data/
where /path/to/trial_folder/ should be replaced with the path to the trial folder that was generated during training. The result will be a file results_lpips_metric.pickle in the trial folder. Note that the LPIPS metric is only computed for the medium and fast sequences as pseudo ground truth images are not generated for the slow sequences.
We've also provided the run_lpips_metric_all.sh bash script to compute the LPIPS metric for all medium/fast trials. It can be run in the container as follows:
/home/user/trnerf/scripts/evaluation/computing_metrics/run_lpips_metric_all.sh /output/ /data/
Additionally, the run_lpips_metric_raw.sh script can be run with the same arguments as run_lpips_metric_all.sh to compute the LPIPS metric with the raw images from all medium/fast sequences. If needed, this will create new trial folders with the pattern /output/<camera>_<speed>_<scene>_raw/ to store the results.
Note that run_lpips_metric_trial.sh and run_lpips_metric_all.sh assume that the speed (slow/medium/fast), scene (indoor/outdoor), and camera (left/right) can be inferred from the trial folder paths (which is true for any trial folder generated from the provided training bash scripts). Alternatively, run_lpips_metric.sh can be run directly with additional arguments.
Run the trnerf_detection_metric Docker image created in the section above with the following command:
docker run -it -v /path/to/trnerf_dataset_folder/:/data -v /path/to/output_folder/:/output trnerf_detection_metric
where /path/to/trnerf_dataset_folder/ should be replaced with the path to the downloaded TRNeRF dataset (see Dataset Download, Format, and Conversion) and /path/to/output_folder/ should be replaced with the output folder used during training. The folders will appear within the container as /data and /output.
The command above will start an interactive bash session in the running container. Analogous to the previous section, you can view the arguments for computing the detection metric with python3 /home/user/trnerf/evaluation/compute_detection_metric.py -h. To compute the detection metric for a single trial, we've provided the run_detection_metric_trial.sh bash script that can be executed in the container as follows:
/home/user/trnerf/scripts/evaluation/computing_metrics/run_detection_metric_trial.sh /path/to/trial_folder/ /data/
where /path/to/trial_folder/ should be replaced with the path to the trial folder that was generated during training. The result will be a file results_detection_metric.pickle in the trial folder.
We've also provided the run_detection_metric_all.sh bash script to compute the detection metric for all trials. It can be run in the container as follows:
/home/user/trnerf/scripts/evaluation/computing_metrics/run_detection_metric_all.sh /output/ /data/
Additionally, the run_detection_metric_raw.sh script can be run with the same arguments as run_detection_metric_all.sh to compute the detection metric with the raw images from all sequences. If needed, this will create new trial folders with the pattern /output/<camera>_<speed>_<scene>_raw/ to store the results.
Note that run_detection_metric_trial.sh and run_detection_metric_all.sh assume that the speed (slow/medium/fast), scene (indoor/outdoor), and camera (left/right) can be inferred from the trial folder paths (which is true for any trial folder generated from the provided training bash scripts). Alternatively, run_detection_metric.sh can be run directly with additional arguments.
In either Docker container, the detection and LPIPS results across all trials can be compiled into spreadsheet by running:
python3 /home/user/trnerf/evaluation/create_spreadsheet.py /output/
The result will be a file results.xlsx in the /output/ folder.
In either Docker container, you can view the arguments for creating a comparison video with python3 /home/user/trnerf/evaluation/create_comparison_video.py -h. As an example, the following command:
python3 /home/user/trnerf/evaluation/create_comparison_video.py \
--paths \
/output/right_medium_indoor_raw/ \
/output/right_medium_indoor_recommended/trnerf/<date and time>/ \
--method-list \
Raw \
TRNeRF \
--include-pseudo-gt \
--annotate-results \
--grid-shape 1 3 \
--path-output-video /output/example_video.mp4
will create a video /output/example_video.mp4 that includes pseudo ground truth, raw, and restored thermal images side-by-side, labeled by type, and annotated with detected corners and per-frame metrics.
In this section, we describe the steps that we took to prepare our dataset for use with TRNeRF. If you want to run TRNeRF with a custom dataset, you may need to adapt these steps. Note that these steps do not include calibration and timestamp synchronization. It is assumed that the data is already precisely time synchronized and that accurate calibration results are available.
Run the trnerf_training Docker image as described above. To compute the pseudo ground truth poses for all sequences we've provided the run_generate_poses.sh bash script that can be executed in the container as follows:
/home/user/trnerf/scripts/pseudo_gt_generation/run_generate_poses.sh /output/ /data/
The result will be one file for each sequence following the pattern /output/<speed>_<scene>/pseudo_ground_truth_poses.h5 that contains the pseudo ground truth poses. There will also be a folder, /output/pose_estimation, that contains intermediate outputs that are used below but can otherwise be deleted.
Run the trnerf_training Docker image as described above. To compute the pseudo ground truth images for both cameras in the medium and fast sequences we've provided the run_generate_images.sh bash script that can be executed in the container as follows:
/home/user/trnerf/scripts/pseudo_gt_generation/run_generate_images.sh /output/ /data/
The result will be two files (one for each camera) for each medium and fast sequence following the pattern /output/<speed>_<scene>/pseudo_ground_truth_adk_<camera>.h5 and for each slow sequence following the pattern /output/<speed>_<scene>/two_point_nuc_results_adk_<camera>.h5. The pseudo_ground_truth_adk_<camera>.h5 files contain the pseudo ground truth images, and the two_point_nuc_results_adk_<camera>.h5 contain the two point NUC results (pixelwise gains and offsets) that were used in generating the pseudo ground truth images. There will also be a folder, /output/training_w_two_point_output, that contains intermediate outputs and can be deleted.
There are scene specific arguments which are hardcoded in run_training.sh and other bash scripts. This section details how they were determined.
The --rotation, --translation, and --scale-factor arguments are used by the TRNeRF data parser to transform the camera poses such that the poses (and the scene) fit (almost entirely) within the scene box (a cube with side length 2 centered about the origin). The script visualize_scene_box.py can be used to determine these arguments by visualizing how well they transform the COLMAP estimated poses and point cloud into the scene box. For example, the following produces this visualization for the medium indoor sequence:
python3 /home/user/trnerf/dataset_prep/visualize_scene_box.py \
--path-poses /data/medium_indoor/pseudo_ground_truth_poses.h5 \
--path-colmap-model /output/pose_estimation/colmap/indoor/sparse/0/ \
--name-camera adk_right \
--path-calibration-results /data/cam_calibration/camchain.yaml \
--rotation 0.0 0.0 0.0 \
--translation -0.37 -1.875 -0.95 \
--scale-factor 0.35
Note that this script requires the COLMAP model produced during the pseudo ground truth pose generation to visualize the point cloud estimate of the scene.
As described in the paper, the thermal image thresholds used during training are set to the minimum and maximum values observed in the each scene. Similarly, the thermal image thresholds used for visualization and evaluation are set to the 0.1 and 99.9 percentiles. The script compute_thresholds.py can be used to compute these thresholds across multiple sequences. For example, the following computes the thresholds used for the indoor scene:
python3 /home/user/trnerf/dataset_prep/compute_thresholds.py \
--paths-image-file \
/data/slow_indoor/adk_right.h5 \
/data/medium_indoor/adk_right.h5 \
/data/fast_indoor/adk_right.h5 \
--paths-poses \
/data/slow_indoor/pseudo_ground_truth_poses.h5 \
/data/medium_indoor/pseudo_ground_truth_poses.h5 \
/data/fast_indoor/pseudo_ground_truth_poses.h5
Note that the pose files are required only to determine the subset of images that are to be used for training.
@InProceedings{trnerf_2025,
author = {Carmichael, Spencer and Bhat, Manohar and Ramanagopal, Mani and Buchan, Austin and Vasudevan, Ram and Skinner, Katherine A.},
title = {{TRNeRF}: Restoring Blurry, Rolling Shutter, and Noisy Thermal Images with Neural Radiance Fields},
booktitle = {Proceedings of the Winter Conference on Applications of Computer Vision (WACV)},
month = {February},
year = {2025},
pages = {7969-7979}
}
TRNeRF, by the UM & Ford Center for Autonomous Vehicles (FCAV), is licensed under CC BY-NC-SA 4.0. For inquiries about commercial licensing, please reach out to the authors.