From 60da97263ca53790dac9f703a08fa408a8abf018 Mon Sep 17 00:00:00 2001
From: JC <29726242+jc211@users.noreply.github.com>
Date: Fri, 3 Feb 2023 16:07:32 +1000
Subject: [PATCH] Clean up imports and docs

---
 docs/nerf_dataset_tips.md   | 14 +++++------
 scripts/nerfcapture2nerf.py | 48 +++++++++++++++----------------------
 2 files changed, 25 insertions(+), 37 deletions(-)

diff --git a/docs/nerf_dataset_tips.md b/docs/nerf_dataset_tips.md
index 735d1d2..5535e4a 100644
--- a/docs/nerf_dataset_tips.md
+++ b/docs/nerf_dataset_tips.md
@@ -53,12 +53,12 @@ See [nerf_loader.cu](src/nerf_loader.cu) for implementation details and addition
 
 ## Preparing new NeRF datasets
 
-To train on self-captured data, one has to process the data into an existing format supported by Instant-NGP. We provide scripts to support three approaches:
+To train on self-captured data, one has to process the data into an existing format supported by `instant-ngp`. We provide scripts to support three approaches:
 - [COLMAP](#COLMAP) to create a dataset from a set of photos or a video you took
 
 - [Record3D](#Record3D) to create a dataset with an iPhone 12 Pro or newer (based on ARKit)
 
-- [NeRFCapture](https://github.com/jc211/NeRFCapture) to create a dataset or stream posed images directly to InsantNGP with an iOS device.
+- [NeRFCapture](#NeRFCapture) to create a dataset or stream posed images directly to `instant-ngp` with an iOS device.
 
 Both require [Python](https://www.python.org/) 3.7 or higher to be installed and available in your PATH.
 
@@ -116,7 +116,7 @@ instant-ngp$ ./instant-ngp [path to training data folder containing transforms.j
 
 ### Record3D
 
-With an >=iPhone 12 Pro, one can use [Record3D](https://record3d.app/) to collect data and avoid COLMAP. [Record3D](https://record3d.app/) is an iOS app that relies on ARKit to estimate each image's camera pose. It is more robust than COLMAP for scenes that lack textures or contain repetitive patterns. To train Instant-NGPs with Record3D data, follow these steps: 
+With an >=iPhone 12 Pro, one can use [Record3D](https://record3d.app/) to collect data and avoid COLMAP. [Record3D](https://record3d.app/) is an iOS app that relies on ARKit to estimate each image's camera pose. It is more robust than COLMAP for scenes that lack textures or contain repetitive patterns. To train `instant-ngp` with Record3D data, follow these steps: 
 
 1. Record a video and export with the "Shareable/Internal format (.r3d)".
 2. Send the exported data to your computer.
@@ -127,19 +127,17 @@ With an >=iPhone 12 Pro, one can use [Record3D](https://record3d.app/) to collec
 	```
 	If you capture the scene in the landscape orientation, add `--rotate`.
 
-5. Launch Instant-NGP training:
+5. Launch `instant-ngp` training:
 	```
 	instant-ngp$ ./instant-ngp path/to/data
 	```
 	
 ### NeRFCapture
-[NeRFCapture](https://github.com/jc211/NeRFCapture) is an iOS app that runs on any ARKit device. It allows you to stream images directly from your phone to InstantNGP thus enabling a more interactive experience. It can also collect an offline dataset for later use. 
+[NeRFCapture](ttps://github.com/jc211/NeRFCapture) is an iOS app that runs on any ARKit device. It allows you to stream images directly from your phone to `instant-ngp` thus enabling a more interactive experience. It can also collect an offline dataset for later use. 
 
 The following dependencies are needed to run the NeRFCapture script:
 ```
 pip install cyclonedds
-pip install opencv-python
-pip install Pillow
 ```
 
 To stream:
@@ -149,7 +147,7 @@ To stream:
 	instant-ngp$ python scripts/nerfcapture2nerf.py --stream
 	```
 3. Wait for the connection between the app and the script to occur. This will be indicated on the app.
-4. Click the send button on the app. The frame captured will be sent to InstantNGP. 
+4. Click the send button on the app. The frame captured will be sent to `instant-ngp`. 
 5. Toggle training
 
 To save a dataset:
diff --git a/scripts/nerfcapture2nerf.py b/scripts/nerfcapture2nerf.py
index 55ddf85..89f387b 100644
--- a/scripts/nerfcapture2nerf.py
+++ b/scripts/nerfcapture2nerf.py
@@ -5,9 +5,7 @@ import argparse
 import cv2
 from pathlib import Path
 import json
-import time
 import shutil
-from PIL import Image
 
 import cyclonedds.idl as idl
 import cyclonedds.idl.annotations as annotate
@@ -27,7 +25,7 @@ def parse_args():
     parser.add_argument("--stream", action="store_true", help="Stream images directly to InstantNGP.")
     parser.add_argument("--n_frames", default=10, type=int, help="Number of frames before saving the dataset. Also used as the number of cameras to remember when streaming.")
     parser.add_argument("--save_path", required='--stream' not in sys.argv, type=str, help="Path to save the dataset.")
-    parser.add_argument("--depth_scale", default=4.0, type=float, help="Depth scale used when saving depth. Only used when saving dataset.")
+    parser.add_argument("--depth_scale", default=10.0, type=float, help="Depth scale used when saving depth. Only used when saving dataset.")
     parser.add_argument("--overwrite", action="store_true", help="Rewrite over dataset if it exists.")
     return parser.parse_args()
 
@@ -80,8 +78,8 @@ def set_frame(testbed, frame_idx: int, rgb: np.ndarray, depth: np.ndarray, depth
 def live_streaming_loop(reader: DataReader, max_cameras: int):
     # Start InstantNGP
     testbed = ngp.Testbed(ngp.TestbedMode.Nerf)
-    testbed.init_window(640, 480)
-    testbed.reload_network_from_file(f"configs/nerf/base.json")
+    testbed.init_window(1920, 1080)
+    testbed.reload_network_from_file()
     testbed.visualize_unit_cube = True
     testbed.nerf.visualize_cameras = True
 
@@ -97,7 +95,7 @@ def live_streaming_loop(reader: DataReader, max_cameras: int):
     while testbed.frame():
         sample = reader.read_next() # Get frame from NeRFCapture
         if sample:
-            print(f"Frame received")
+            print(f"Frame {total_frames + 1} received")
 
             # RGB
             image = np.asarray(sample.image, dtype=np.uint8).reshape(
@@ -105,7 +103,7 @@ def live_streaming_loop(reader: DataReader, max_cameras: int):
             image = np.concatenate(
                 [image, np.zeros((sample.height, sample.width, 1), dtype=np.float32)], axis=-1)
 
-            # Depth if avaiable
+            # Depth if available
             depth = None
             if sample.has_depth:
                 depth = np.asarray(sample.depth_image, dtype=np.uint8).view(
@@ -113,7 +111,6 @@ def live_streaming_loop(reader: DataReader, max_cameras: int):
                 depth = cv2.resize(depth, dsize=(
                     sample.width, sample.height), interpolation=cv2.INTER_NEAREST)
 
-
             # Transform
             X_WV = np.asarray(sample.transform_matrix,
                             dtype=np.float32).reshape((4, 4)).T[:3, :]
@@ -140,19 +137,17 @@ def live_streaming_loop(reader: DataReader, max_cameras: int):
                 testbed.render_groundtruth = True
 
 def dataset_capture_loop(reader: DataReader, save_path: Path, overwrite: bool, n_frames: int):
-
     if save_path.exists():
         if overwrite:
             # Prompt user to confirm deletion
-            res = input("Warning, directory exists already. Press Y to delete anyway: ")
-            if res == 'Y':
-                shutil.rmtree(save_path)
-            else:
-                exit()
+            if (input(f"warning! folder '{save_path}' will be deleted/replaced. continue? (Y/n)").lower().strip()+"y")[:1] != "y":
+                sys.exit(1)
+            shutil.rmtree(save_path)
         else:
-            print("save_path already exists")
-            exit()
+            print(f"save_path {save_path} already exists")
+            sys.exit(1)
     
+    print("Waiting for frames...")
     # Make directory
     images_dir = save_path.joinpath("images")
 
@@ -170,10 +165,9 @@ def dataset_capture_loop(reader: DataReader, save_path: Path, overwrite: bool, n
 
     # Start DDS Loop
     while True:
-        time.sleep(0.001)
         sample = reader.read_next() # Get frame from NeRFCapture
         if sample:
-            print(f"Frame received")
+            print(f"{total_frames + 1}/{n_frames} frames received")
 
             if total_frames == 0:
                 save_path.mkdir(parents=True)
@@ -187,11 +181,8 @@ def dataset_capture_loop(reader: DataReader, save_path: Path, overwrite: bool, n
                 manifest["integer_depth_scale"] = float(args.depth_scale)/65535.0
 
             # RGB
-            image = np.asarray(sample.image, dtype=np.uint8).reshape(
-                (sample.height, sample.width, 3))
-            image = np.concatenate(
-                [image, 255*np.ones((sample.height, sample.width, 1), dtype=np.uint8)], axis=-1)
-            Image.fromarray(image).save(images_dir.joinpath(f"{total_frames}.png"))
+            image = np.asarray(sample.image, dtype=np.uint8).reshape((sample.height, sample.width, 3))
+            cv2.imwrite(str(images_dir.joinpath(f"{total_frames}.png")), cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
 
             # Depth if avaiable
             depth = None
@@ -201,13 +192,11 @@ def dataset_capture_loop(reader: DataReader, save_path: Path, overwrite: bool, n
                 depth = (depth*65535/float(args.depth_scale)).astype(np.uint16)
                 depth = cv2.resize(depth, dsize=(
                     sample.width, sample.height), interpolation=cv2.INTER_NEAREST)
-                Image.fromarray(depth).save(images_dir.joinpath(f"{total_frames}.depth.png"))
-
+                cv2.imwrite(str(images_dir.joinpath(f"{total_frames}.depth.png")), depth) 
 
             # Transform
             X_WV = np.asarray(sample.transform_matrix,
                               dtype=np.float32).reshape((4, 4)).T
-
             
             frame = {
                 "transform_matrix": X_WV.tolist(),
@@ -226,14 +215,15 @@ def dataset_capture_loop(reader: DataReader, save_path: Path, overwrite: bool, n
             manifest["frames"].append(frame)
 
             # Update index
-            total_frames += 1
-            if total_frames == n_frames:
+            if total_frames == n_frames - 1:
+                print("Saving manifest...")
                 # Write manifest as json
                 manifest_json = json.dumps(manifest, indent=4)
                 with open(save_path.joinpath("transforms.json"), "w") as f:
                     f.write(manifest_json)
                 print("Done")
-                exit()
+                sys.exit(0)
+            total_frames += 1
 
 
 if __name__ == "__main__":
-- 
GitLab