added scripts for evaluation

README.md

@@ -1,19 +1,33 @@
-# LAENeRF
-We present LAENeRF, a method for **L**ocal **A**ppearance **E**diting of **Ne**ural **R**adiance **F**ields. \
-\
-<img src="assets/teaser.png" width=1200px> \
-\
-<img src="assets/flower.gif" width="280"/>
-<img src="assets/horns.gif" width="280"/>
-<img src="assets/trex.gif" width="280"/>
-<img src="assets/bonsai_2.gif" width="280"/>
+# LAENeRF: Local Appearance Editing for Neural Radiance Fields
+### [Lukas Radl](https://r4dl.github.io/), [Michael Steiner](https://online.tugraz.at/tug_online/visitenkarte.show_vcard?pPersonenId=2999CD2743C7E977&pPersonenGruppe=3), [Andreas Kurz](https://scholar.google.com/citations?user=3yD4NZgAAAAJ&hl=de&oi=ao), [Markus Steinberger](https://www.markussteinberger.net/)
+##### Institute of Computer Graphics and Vision, Graz University of Technology
+#### [Project Page](https://r4dl.github.io/nerfinternals/) | [Preprint](https://arxiv.org/pdf/2306.00696.pdf)
+
+<p align="middle">
+    <img src="assets/teaser.png" width=1200px>
+</p>
+<p align="middle">
+    <img src="assets/flower.gif" width="280"/>
+    <img src="assets/horns.gif" width="280"/>
+    <img src="assets/trex.gif" width="280"/>
+    <img src="assets/bonsai_2.gif" width="280"/>
+</p>
+
+We present LAENeRF, a method for **L**ocal **A**ppearance **E**diting of **Ne**ural **R**adiance **F**ields.
+This repository contains the official implementation, based on [torch-ngp](https://github.com/ashawkey/torch-ngp).
+
 ## Installation
-First, setup the conda environment using:
+First, clone our repository
+```bash
+git clone https://github.com/r4dl/LAENeRF.git
+cd LAENeRF
+```
+Next, setup the conda environment using:
 ```bash
 conda env create -f environment_laenerf.yml
 conda activate LAENeRF
 ```
-Next, install tiny-cuda-nn.
+Finally, install [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn).
 ```bash
 pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
 ```
@@ -33,7 +47,7 @@ We use the same data format torch-ngp and use the scripts from PaletteNeRF.
 Currently, three datasets are supported: [NeRF Synthetic](https://drive.google.com/drive/folders/128yBriW1IG_3NJ5Rp7APSTZsJqdJdfc1), [LLFF](https://drive.google.com/drive/folders/14boI-o5hGO9srnWaaogTU5_ji7wkX2S7) and [mip-NeRF 360](http://storage.googleapis.com/gresearch/refraw360/360_v2.zip). 
 The folder structure should look like the following:
 ```
-code
+LAENeRF
 ├── ...
 ├── data
 │   ├── nerf_synthetic
@@ -52,7 +66,7 @@ For LLFF and mip-NeRF 360, data conversion is required:
 # mip-NeRF 360
 python scripts/llff2nerf.py /path/to/the/scene/directory --images images_4 --downscale 4 --datatype mip360
 # LLFF
-python scripts/llff2nerf.py /path/to/the/scene/directory --images images_4 --downscale 4 --datatype llff
+python scripts/llff2nerf.py /path/to/the/scene/directory --images images_4 --downscale 4 --datatype masks
 ```
 
 ## Training a NeRF
@@ -73,16 +87,104 @@ bash scripts/run_{dataset}.sh scripts/configs_{dataset}/{scene}.sh -m style
 ```
 ### Results
 By default, we write the results to `test/{scene}_{style, recolor}`.
-In this directory, we additionally save test/val set images, a video, LAENeRF (which may be loaded for recolorable stylization),
+In this directory, we additionally save test/val set images, a video, the trained LAENeRF (which may be loaded for recolorable stylization),
 the fine-tuned NeRF and the grids.
 
 ### Style Images
 We provide a subset of style images in the `style_images` directory. 
 Per-default, *the great wave* is loaded for stylization.
-\
-\
-<img src="style_images/ju-jitsu-style.png" style="width:200px; height:200px">
-<img src="style_images/style_image_paint.png" style="width:200px; height:200px">
-<img src="style_images/wave_style.png" style="width:200px; height:200px">
-<img src="style_images/style_image_.png" style="width:200px; height:200px">
-<img src="style_images/picasso2.jpeg" style="width:200px; height:200px">
+
+<p align="middle">
+    <img src="style_images/ju-jitsu-style.png" style="width:200px; height:200px">
+    <img src="style_images/style_image_paint.png" style="width:200px; height:200px">
+    <img src="style_images/wave_style.png" style="width:200px; height:200px">
+    <img src="style_images/style_image_.png" style="width:200px; height:200px">
+    <img src="style_images/picasso2.jpeg" style="width:200px; height:200px">
+</p>
+
+## Evaluation
+We provide our evaluation scripts for our quantitative results to enable more future research in this direction.
+All necessary scripts are located in the `scripts/eval` directory, with the required masks contained in the `scripts/eval/masks` directory.
+
+### MSE in the background
+To evaluate the MSE in the background with respect to the ground truth test set images, run
+```bash
+# mip-NeRF 360
+python scripts/eval/mse_background_mip360.py --scene {scene} --results_dir {results_dir}
+# LLFF
+python scripts/eval/mse_background_llff.py --scene {scene} --results_dir {results_dir}
+```
+Further, to compare to images from a pre-trained NeRF, provide a `--comparison_dir`. 
+To save the difference images as well as the results as a json file, provide a `--save_dir`.
+
+### View Consistency
+To evaluate view-consistency, clone [RAFT](https://github.com/princeton-vl/RAFT), 
+install the conda environment according to their README.md and download their models.
+Copy `scripts/eval/consistency_metrics.py` in the `RAFT` directory.
+Finally, to obtain view-consistency metrics:
+```bash
+# Short-Range Consistency
+python consistency_metrics.py --model=models/raft-sintel.pth --directory_gt={path_to_nerf_frames} --directory={path_to_stylized_frames} --step 1
+# Long-Range Consistency
+python consistency_metrics.py --model=models/raft-sintel.pth --directory_gt={path_to_nerf_frames} --directory={path_to_stylized_frames} --step 7
+```
+Here, `{path_to_nerf_frames}` is a directory which contains frames from a video rendered by a pre-trained NeRF. 
+`{path_to_stylized_frames}` is a directory which contains the same frames from a recoloring/stylization method. 
+Importantly, we compute the optical flow from the former, and measure MSE with the latter.
+
+## Citation
+If you find our work useful, consider citing us with:
+```bibtex
+@article(Radl2023laenerf,
+  title     = {{LAENeRF: Local Appearance Editing of Neural Radiance Fields}},
+  author    = {Radl, Lukas and Steiner, Michael and Kurz, Andreas and Steinberger, Markus},
+  journal   = {arXiv preprint arXiv:},
+  year      = {2023},
+}
+```
+
+## Acknowledgements
+The implementation of Instant-NGP is adapted from [torch-ngp](https://github.com/ashawkey/torch-ngp):
+```bibtex
+@misc{torch-ngp,
+  Author    = {Jiaxiang Tang},
+  Year      = {2022},
+  Note      = {https://github.com/ashawkey/torch-ngp},
+  Title     = {Torch-ngp: a PyTorch implementation of instant-ngp}
+}
+
+@article{tang2022compressible,
+  title     = {Compressible-composable NeRF via Rank-residual Decomposition},
+  author    = {Tang, Jiaxiang and Chen, Xiaokang and Wang, Jingbo and Zeng, Gang},
+  journal   = {arXiv preprint arXiv:2205.14870},
+  year      = {2022}
+}
+```
+Some elements of the GUI as well as scripts for data conversion and running our method are adapted from [PaletteNeRF](https://github.com/zfkuang/PaletteNeRF/). 
+```bibtex
+@article{kuang2022palettenerf,
+  title     = {PaletteNeRF: Palette-based Appearance Editing of Neural Radiance Fields},
+  author    = {Kuang, Zhengfei and Luan, Fujun and Bi, Sai and Shu, Zhixin and Wetzstein, Gordon and Sunkavalli, Kalyan},
+  journal   = {arXiv preprint arXiv:2212.10699},
+  year      = {2022}
+}
+```
+The segmentation masks for Mip-NeRF 360 were extracted with [Segment Anything](https://github.com/facebookresearch/segment-anything),
+and the masks for LLFF were given to use by the authors of [ICE-NeRF](https://heuyklee.github.io/ice-nerf/).
+```bibtex
+@article{kirillov2023segany,
+  title     = {Segment Anything},
+  author    = {Kirillov, Alexander and Mintun, Eric and Ravi, Nikhila and Mao, Hanzi and Rolland, Chloe and Gustafson, Laura and Xiao, Tete and Whitehead, Spencer and Berg, Alexander C. and Lo, Wan-Yen and Doll{\'a}r, Piotr and Girshick, Ross},
+  journal   = {arXiv:2304.02643},
+  year      = {2023}
+}
+
+@InProceedings{Lee_2023_ICCV,
+  author    = {Lee, Jae-Hyeok and Kim, Dae-Shik},
+  title     = {ICE-NeRF: Interactive Color Editing of NeRFs via Decomposition-Aware Weight Optimization},
+  booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
+  month     = {October},
+  year      = {2023},
+  pages     = {3491-3501}
+}
+```

nerf/gui.py

@@ -1905,7 +1905,7 @@ def callback_palet_close(sender, app_data):
             dpg.bind_font(default_font)
 
 
-        dpg.create_viewport(title='torch-ngp', width=self.W, height=self.H, resizable=False)
+        dpg.create_viewport(title='LAENeRF', width=self.W, height=self.H, resizable=False)
 
         ### global theme
         with dpg.theme() as theme_no_padding:

scripts/eval/consistency_metrics.py

@@ -0,0 +1,176 @@
+import sys
+sys.path.append('core')
+
+from PIL import Image
+import argparse
+import os
+import time
+import numpy as np
+import torch
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+import cv2, json
+import lpips
+
+import datasets
+from utils import flow_viz
+from utils import frame_utils
+
+from raft import RAFT
+from utils.utils import InputPadder, forward_interpolate
+
+def compute_flow_magnitude(flow):
+
+    flow_mag = flow[:, :, 0] ** 2 + flow[:, :, 1] ** 2
+
+    return flow_mag
+
+def compute_flow_gradients(flow):
+    H = flow.shape[0]
+    W = flow.shape[1]
+
+    flow_x_du = np.zeros((H, W))
+    flow_x_dv = np.zeros((H, W))
+    flow_y_du = np.zeros((H, W))
+    flow_y_dv = np.zeros((H, W))
+
+    flow_x = flow[:, :, 0]
+    flow_y = flow[:, :, 1]
+
+    flow_x_du[:, :-1] = flow_x[:, :-1] - flow_x[:, 1:]
+    flow_x_dv[:-1, :] = flow_x[:-1, :] - flow_x[1:, :]
+    flow_y_du[:, :-1] = flow_y[:, :-1] - flow_y[:, 1:]
+    flow_y_dv[:-1, :] = flow_y[:-1, :] - flow_y[1:, :]
+
+    return flow_x_du, flow_x_dv, flow_y_du, flow_y_dv
+
+@torch.no_grad()
+def detect_occlusion(fw_flow, bw_flow):
+    ## fw-flow: img1 => img2
+    ## bw-flow: img2 => img1
+
+    fw_flow_t = torch.from_numpy(fw_flow).cuda()
+    bw_flow_t = torch.from_numpy(bw_flow).cuda()
+
+    ## warp fw-flow to img2
+    flow = torch.from_numpy(fw_flow).permute(2, 0, 1).float()
+    flow[0, :, :] += torch.arange(flow.shape[2])
+    flow[1, :, :] += torch.arange(flow.shape[1])[:, None]
+    fw_flow_w = cv2.remap((bw_flow_t).cpu().numpy(), flow.permute(1, 2, 0).cpu().numpy(), None, cv2.INTER_LINEAR)
+
+    ## convert to numpy array
+    fw_flow_w = (fw_flow_w)
+
+    ## occlusion
+    fb_flow_sum = fw_flow_w + bw_flow
+    fb_flow_mag = compute_flow_magnitude(fb_flow_sum)
+    fw_flow_w_mag = compute_flow_magnitude(fw_flow_w)
+    bw_flow_mag = compute_flow_magnitude(bw_flow)
+
+    mask1 = fb_flow_mag > 0.01 * (fw_flow_w_mag + bw_flow_mag) + 0.5
+
+    ## motion boundary
+    fx_du, fx_dv, fy_du, fy_dv = compute_flow_gradients(bw_flow)
+    fx_mag = fx_du ** 2 + fx_dv ** 2
+    fy_mag = fy_du ** 2 + fy_dv ** 2
+
+    mask2 = (fx_mag + fy_mag) > 0.01 * bw_flow_mag + 0.002
+
+    ## combine mask
+    mask = np.logical_or(mask1, mask2)
+    occlusion = np.zeros((fw_flow.shape[0], fw_flow.shape[1]))
+    occlusion[mask == 1] = 1
+
+    return occlusion
+
+@torch.no_grad()
+def validate_longrange(model, frames_directory, frames_directory_gt, step=7):
+    """ Peform validation using the Sintel (train) split """
+    model.eval()
+    results = {}
+    iters=32
+
+    lpipsmeter = lpips.LPIPS(net='alex').cuda()
+
+    frames_gt = sorted([i for i in os.listdir(frames_directory_gt) if 'png' in i])
+    frames_ours = sorted([i for i in os.listdir(frames_directory) if 'png' in i])
+    flow_prev, sequence_prev = None, None
+
+    errors_ours = []
+    lpips_ours = []
+
+    for test_id in range(len(frames_directory) - step):
+        image1 = frame_utils.read_gen(os.path.join(frames_directory_gt, frames_gt[test_id]))
+        image2 = frame_utils.read_gen(os.path.join(frames_directory_gt, frames_gt[test_id + step]))
+        image1 = np.array(image1).astype(np.uint8)[..., :3]
+        image2 = np.array(image2).astype(np.uint8)[..., :3]
+        image1 = torch.from_numpy(image1).permute(2, 0, 1).float()
+        image2 = torch.from_numpy(image2).permute(2, 0, 1).float()
+
+        padder = InputPadder(image1.shape)
+        image1, image2 = padder.pad(image1[None].cuda(), image2[None].cuda())
+
+        flow_low, flow_pr = model(image2, image1, iters=iters, flow_init=flow_prev, test_mode=True)
+        flow = padder.unpad(flow_pr[0]).permute(1, 2, 0).cpu().numpy()
+
+        flow_low, flow_pr = model(image1, image2, iters=iters, flow_init=flow_prev, test_mode=True)
+        flow_bw = padder.unpad(flow_pr[0]).permute(1, 2, 0).cpu().numpy()
+
+        flow_clone = np.copy(flow)
+        occ = detect_occlusion(fw_flow=flow_clone, bw_flow=flow_bw)
+
+        flow = torch.from_numpy(flow).permute(2, 0, 1).float()
+        flow[0, :, :] += torch.arange(flow.shape[2])
+        flow[1, :, :] += torch.arange(flow.shape[1])[:, None]
+
+        image1 = frame_utils.read_gen(os.path.join(frames_directory, frames_ours[test_id]))
+        image2 = frame_utils.read_gen(os.path.join(frames_directory, frames_ours[test_id + step]))
+        image1 = np.array(image1).astype(np.uint8)[..., :3]
+        image2 = np.array(image2).astype(np.uint8)[..., :3]
+        image1 = torch.from_numpy(image1).permute(2, 0, 1).float() / 255.
+        image2 = torch.from_numpy(image2).permute(2, 0, 1).float() / 255.
+
+        warped_result = cv2.remap((image1).permute(1,2,0).cpu().numpy(), flow.permute(1,2,0).cpu().numpy(), None, cv2.INTER_LINEAR)
+        lpi = lpipsmeter(torch.from_numpy(warped_result).permute(-1,0,1).cuda(), image2.cuda()).item()
+
+        mse = (((torch.from_numpy(warped_result).permute(2,0,1) - image2)**2).sum(0) * occ)
+        error = (mse.sum() / occ.sum())
+
+        errors_ours.append(error.item())
+        lpips_ours.append(lpi)
+
+    print(f'RSME of {np.mean(np.array(errors_ours)):.4f}')
+    print(f'LPIPS of {np.mean(np.array(lpips_ours))}')
+
+    dict_ = {
+        'rsme': np.mean(np.array(errors_ours)),
+        'lpips': np.mean(np.array(lpips_ours))
+    }
+
+    with open('results.json', 'w') as outfile:
+        json.dump(dict_, outfile, indent=2)
+
+    return results
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model', help="restore checkpoint")
+    parser.add_argument('--step', type=int, help='number of steps to compute the optical flow')
+    parser.add_argument('--directory', type=str, help='directory in which the images to be evaluated lie')
+    parser.add_argument('--directory_gt', type=str, help='directory in which the gt images lie')
+    args = parser.parse_args()
+
+    # custom args
+    args.small = False
+    args.mixed_precision = False
+
+    model = torch.nn.DataParallel(RAFT(args))
+    model.load_state_dict(torch.load(args.model))
+
+    model.cuda()
+    model.eval()
+
+    with torch.no_grad():
+        validate_longrange(model.module, frames_directory=args.directory, frames_directory_gt=args.directory_gt,
+                           step=args.step)