Merge remote-tracking branch 'origin' into generative

nerfstudio/configs/base_config.py

@@ -23,6 +23,8 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, List, Optional, Tuple, Type
 
+from typing_extensions import Literal
+
 # model instances
 from nerfstudio.utils import writer
 
@@ -149,3 +151,5 @@ class ViewerConfig(PrintableConfig):
     skip_openrelay: bool = False
     """Avoid using openrelay to communicate with the viewer. Try disabling if you have trouble
     connecting to the viewer"""
+    codec: Literal["H264", "VP8"] = "VP8"
+    """Video codec that viewer will use."""

nerfstudio/process_data/metashape_utils.py

@@ -58,10 +58,14 @@ def metashape_to_json(  # pylint: disable=too-many-statements
     sensors = chunk.find("sensors")
 
     # TODO Add support for per-frame intrinsics
-    if sensors is None or len(sensors) != 1:
+    if sensors is None:
+        raise ValueError("No sensors found")
+
+    calibrated_sensors = [sensor for sensor in sensors if sensor.find("calibration")]
+    if len(calibrated_sensors) != 1:
         raise ValueError("Only one sensor is supported for now")
 
-    sensor = sensors.find("sensor")
+    sensor = calibrated_sensors[0]
 
     data = {}
 
@@ -98,6 +102,14 @@ def metashape_to_json(  # pylint: disable=too-many-statements
         if camera_label not in image_filename_map:
             continue
         frame["file_path"] = image_filename_map[camera_label].as_posix()
+
+        if camera.get("sensor_id") != sensor.get("id"):
+            # this should only happen when we have a sensor that doesn't have calibration
+            if verbose:
+                CONSOLE.print(f"Missing sensor calibration for {camera.get('label')}, Skipping")
+            num_skipped += 1
+            continue
+
         if camera.find("transform") is None:
             if verbose:
                 CONSOLE.print(f"Missing transforms data for {camera.get('label')}, Skipping")

nerfstudio/process_data/polycam_utils.py

@@ -72,7 +72,7 @@ def polycam_to_json(
         frame["h"] = frame_json["height"] - crop_border_pixels * 2
         frame["file_path"] = f"./images/frame_{i+1:05d}{image_filename.suffix}"
         if use_depth:
-            frame["depth_map_path"] = f"./depth/frame_{i+1:05d}{depth_filenames[i].suffix}"
+            frame["depth_file_path"] = f"./depth/frame_{i+1:05d}{depth_filenames[i].suffix}"
         # Transform matrix to nerfstudio format. Please refer to the documentation for coordinate system conventions.
         frame["transform_matrix"] = [
             [frame_json["t_20"], frame_json["t_21"], frame_json["t_22"], frame_json["t_23"]],

nerfstudio/process_data/process_data_utils.py

@@ -78,8 +78,8 @@ def get_num_frames_in_video(video: Path) -> int:
     Returns:
         The number of frames in a video.
     """
-    cmd = f"ffprobe -v error -select_streams v:0 -count_packets \
-            -show_entries stream=nb_read_packets -of csv=p=0 {video}"
+    cmd = f'ffprobe -v error -select_streams v:0 -count_packets \
+            -show_entries stream=nb_read_packets -of csv=p=0 "{video}"'
     output = run_command(cmd)
     assert output is not None
     output = output.strip(" ,\t\n\r")
@@ -114,7 +114,7 @@ def convert_video_to_images(
         print("Number of frames in video:", num_frames)
 
         out_filename = image_dir / "frame_%05d.png"
-        ffmpeg_cmd = f"ffmpeg -i {video_path}"
+        ffmpeg_cmd = f'ffmpeg -i "{video_path}"'
         spacing = num_frames // num_frames_target
 
         if spacing > 1:
@@ -168,7 +168,7 @@ def copy_images_list(
         file_type = image_paths[0].suffix
         filename = f"frame_%05d{file_type}"
         crop = f"crop=iw-{crop_border_pixels*2}:ih-{crop_border_pixels*2}"
-        ffmpeg_cmd = f"ffmpeg -y -noautorotate -i {image_dir / filename} -q:v 2 -vf {crop} {image_dir / filename}"
+        ffmpeg_cmd = f'ffmpeg -y -noautorotate -i "{image_dir / filename}" -q:v 2 -vf {crop} "{image_dir / filename}"'
         run_command(ffmpeg_cmd, verbose=verbose)
 
     num_frames = len(image_paths)
@@ -210,9 +210,9 @@ def copy_and_upscale_polycam_depth_maps_list(
         for idx, depth_map in enumerate(polycam_depth_image_filenames):
             destination = depth_dir / f"frame_{idx + 1:05d}{depth_map.suffix}"
             ffmpeg_cmd = [
-                f"ffmpeg -y -i {depth_map} ",
+                f'ffmpeg -y -i "{depth_map}" ',
                 f"-q:v 2 -vf scale=iw*{upscale_factor}:ih*{upscale_factor}:flags=neighbor ",
-                f"{destination}",
+                f'"{destination}"',
             ]
             ffmpeg_cmd = " ".join(ffmpeg_cmd)
             run_command(ffmpeg_cmd, verbose=verbose)
@@ -222,7 +222,7 @@ def copy_and_upscale_polycam_depth_maps_list(
         file_type = depth_dir.glob("frame_*").__next__().suffix
         filename = f"frame_%05d{file_type}"
         crop = f"crop=iw-{crop_border_pixels * 2}:ih-{crop_border_pixels * 2}"
-        ffmpeg_cmd = f"ffmpeg -y -i {depth_dir / filename} -q:v 2 -vf {crop} {depth_dir / filename}"
+        ffmpeg_cmd = f'ffmpeg -y -i "{depth_dir / filename}" -q:v 2 -vf {crop} "{depth_dir / filename}"'
         run_command(ffmpeg_cmd, verbose=verbose)
 
     CONSOLE.log("[bold green]:tada: Done upscaling depth maps.")
@@ -282,9 +282,9 @@ def downscale_images(image_dir: Path, num_downscales: int, folder_name: str = "i
             for f in files:
                 filename = f.name
                 ffmpeg_cmd = [
-                    f"ffmpeg -y -noautorotate -i {image_dir / filename} ",
+                    f'ffmpeg -y -noautorotate -i "{image_dir / filename}" ',
                     f"-q:v 2 -vf scale=iw/{downscale_factor}:ih/{downscale_factor} ",
-                    f"{downscale_dir / filename}",
+                    f'"{downscale_dir / filename}"',
                 ]
                 ffmpeg_cmd = " ".join(ffmpeg_cmd)
                 run_command(ffmpeg_cmd, verbose=verbose)

nerfstudio/viewer/app/package.json

@@ -1,7 +1,7 @@
 {
   "name": "viewer",
   "homepage": ".",
-  "version": "23-01-25-0",
+  "version": "23-02-3-0",
   "private": true,
   "dependencies": {
     "@emotion/react": "^11.10.4",

nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx

@@ -37,6 +37,12 @@ export function RenderControls() {
   const colormapChoice = useSelector(
     (state) => state.renderingState.colormap_choice,
   );
+  const colormapInvert = useSelector(
+    (state) => state.renderingState.colormap_invert,
+  );
+  const colormapNormalize = useSelector(
+    (state) => state.renderingState.colormap_normalize,
+  );
   const max_resolution = useSelector(
     (state) => state.renderingState.maxResolution,
   );
@@ -115,6 +121,51 @@ export function RenderControls() {
         },
         disabled: colormapOptions.length === 1,
       },
+      colormap_invert: {
+        label: '| Invert',
+        value: colormapInvert,
+        hint: 'Invert the colormap',
+        onChange: (v) => {
+          dispatch_and_send(
+            websocket,
+            dispatch,
+            'renderingState/colormap_invert',
+            v,
+          );
+        },
+        render: (get) => get('colormap_options') !== 'default',
+      },
+      colormap_normalize: {
+        label: '| Normalize',
+        value: colormapNormalize,
+        hint: 'Whether to normalize output between 0 and 1',
+        onChange: (v) => {
+          dispatch_and_send(
+            websocket,
+            dispatch,
+            'renderingState/colormap_normalize',
+            v,
+          );
+        },
+        render: (get) => get('colormap_options') !== 'default',
+      },
+      colormap_range: {
+        label: '| Range',
+        value: [0, 1],
+        step: 0.01,
+        min: -2,
+        max: 5,
+        hint: 'Min and max values of the colormap',
+        onChange: (v) => {
+          dispatch_and_send(
+            websocket,
+            dispatch,
+            'renderingState/colormap_range',
+            v,
+          );
+        },
+        render: (get) => get('colormap_options') !== 'default',
+      },
       // Dynamic Resolution
       target_train_util: {
         label: 'Train Util.',

nerfstudio/viewer/app/src/reducer.js

@@ -30,12 +30,17 @@ const initialState = {
 
     all_camera_paths: null, // object containing camera paths and names
 
-
     isTraining: true,
+
+    // colormap options
     output_options: ['rgb'], // populated by the possible Graph outputs
     output_choice: 'rgb', // the selected output
     colormap_options: ['default'], // populated by the output choice
     colormap_choice: 'default', // the selected colormap
+    colormap_invert: false, // whether to invert the colormap
+    colormap_normalize: false, // whether to normalize the colormap
+    colormap_range: [0.0, 1.0], // the range of the colormap
+
     maxResolution: 1024,
     targetTrainUtil: 0.9,
     eval_res: '?',
@@ -51,7 +56,7 @@ const initialState = {
 
     // Crop Box Options
     crop_enabled: false,
-    crop_bg_color: {r: 38, g:42, b:55},
+    crop_bg_color: { r: 38, g: 42, b: 55 },
     crop_scale: [2.0, 2.0, 2.0],
     crop_center: [0.0, 0.0, 0.0],
   },

nerfstudio/viewer/server/utils.py

@@ -63,7 +63,7 @@ def three_js_perspective_camera_focal_length(fov: float, image_height: int):
 
 
 def get_intrinsics_matrix_and_camera_to_world_h(
-    camera_object: Dict[str, Any], image_height: int
+    camera_object: Dict[str, Any], image_height: int, image_width: Optional[int] = None
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Returns the camera intrinsics matrix and the camera to world homogeneous matrix.
 
@@ -74,7 +74,8 @@ def get_intrinsics_matrix_and_camera_to_world_h(
     # intrinsics
     fov = camera_object["fov"]
     aspect = camera_object["aspect"]
-    image_width = aspect * image_height
+    if image_width is None:
+        image_width = aspect * image_height
     pp_w = image_width / 2.0
     pp_h = image_height / 2.0
     if (camera_object["camera_type"] == "perspective") | (camera_object["camera_type"] == "fisheye"):

nerfstudio/viewer/server/viewer_utils.py

@@ -82,7 +82,7 @@ def setup_viewer(config: cfg.ViewerConfig, log_filename: Path, datapath: str):
 
 
 class OutputTypes(str, enum.Enum):
-    """Noncomprehsnive list of output render types"""
+    """Noncomprehensive list of output render types"""
 
     INIT = "init"
     RGB = "rgb"
@@ -92,14 +92,14 @@ class OutputTypes(str, enum.Enum):
 
 
 class ColormapTypes(str, enum.Enum):
-    """Noncomprehsnive list of colormap render types"""
+    """List of colormap render types"""
 
-    INIT = "init"
     DEFAULT = "default"
     TURBO = "turbo"
-    DEPTH = "depth"
-    SEMANTIC = "semantic"
-    BOOLEAN = "boolean"
+    VIRIDIS = "viridis"
+    MAGMA = "magma"
+    INFERNO = "inferno"
+    CIVIDIS = "cividis"
 
 
 class IOChangeException(Exception):
@@ -220,12 +220,15 @@ class CheckThread(threading.Thread):
 
             # check colormap type
             colormap_type = self.state.vis["renderingState/colormap_choice"].read()
-            if colormap_type is None:
-                colormap_type = ColormapTypes.INIT
             if self.state.prev_colormap_type != colormap_type:
                 self.state.check_interrupt_vis = True
                 return
 
+            colormap_range = self.state.vis["renderingState/colormap_range"].read()
+            if self.state.prev_colormap_range != colormap_range:
+                self.state.check_interrupt_vis = True
+                return
+
             # check max render
             max_resolution = self.state.vis["renderingState/maxResolution"].read()
             if max_resolution is not None:
@@ -285,7 +288,10 @@ class ViewerState:
         self.prev_camera_matrix = None
         self.prev_render_time = 0
         self.prev_output_type = OutputTypes.INIT
-        self.prev_colormap_type = ColormapTypes.INIT
+        self.prev_colormap_type = None
+        self.prev_colormap_invert = False
+        self.prev_colormap_normalize = False
+        self.prev_colormap_range = [0, 1]
         self.prev_moving = False
         self.output_type_changed = True
         self.max_resolution = 1000
@@ -583,11 +589,21 @@ class ViewerState:
             self.camera_moving = True
 
         colormap_type = self.vis["renderingState/colormap_choice"].read()
-        if colormap_type is None:
-            colormap_type = ColormapTypes.INIT
         if self.prev_colormap_type != colormap_type:
             self.camera_moving = True
 
+        colormap_range = self.vis["renderingState/colormap_range"].read()
+        if self.prev_colormap_range != colormap_range:
+            self.camera_moving = True
+
+        colormap_invert = self.vis["renderingState/colormap_invert"].read()
+        if self.prev_colormap_invert != colormap_invert:
+            self.camera_moving = True
+
+        colormap_normalize = self.vis["renderingState/colormap_normalize"].read()
+        if self.prev_colormap_normalize != colormap_normalize:
+            self.camera_moving = True
+
         crop_bg_color = self.vis["renderingState/crop_bg_color"].read()
         if self.prev_crop_enabled:
             if self.prev_crop_bg_color != crop_bg_color:
@@ -621,37 +637,28 @@ class ViewerState:
             return outputs[reformatted_output]
 
         # rendering depth outputs
-        if self.prev_colormap_type == ColormapTypes.DEPTH or (
-            self.prev_colormap_type == ColormapTypes.DEFAULT
-            and outputs[reformatted_output].dtype == torch.float
-            and (torch.max(outputs[reformatted_output]) - 1.0) > eps  # handle floating point arithmetic
-        ):
-            accumulation_str = (
-                OutputTypes.ACCUMULATION
-                if OutputTypes.ACCUMULATION in self.output_list
-                else OutputTypes.ACCUMULATION_FINE
-            )
-            return colormaps.apply_depth_colormap(outputs[reformatted_output], accumulation=outputs[accumulation_str])
-
-        # rendering accumulation outputs
-        if self.prev_colormap_type == ColormapTypes.TURBO or (
-            self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[reformatted_output].dtype == torch.float
-        ):
-            return colormaps.apply_colormap(outputs[reformatted_output])
+        if outputs[reformatted_output].shape[-1] == 1 and outputs[reformatted_output].dtype == torch.float:
+            output = outputs[reformatted_output]
+            if self.prev_colormap_normalize:
+                output = output - torch.min(output)
+                output = output / (torch.max(output) + eps)
+            output = output * (self.prev_colormap_range[1] - self.prev_colormap_range[0]) + self.prev_colormap_range[0]
+            output = torch.clip(output, 0, 1)
+            if self.prev_colormap_invert:
+                output = 1 - output
+            if self.prev_colormap_type == ColormapTypes.DEFAULT:
+                return colormaps.apply_colormap(output, cmap=ColormapTypes.TURBO.value)
+            return colormaps.apply_colormap(output, cmap=self.prev_colormap_type)
 
         # rendering semantic outputs
-        if self.prev_colormap_type == ColormapTypes.SEMANTIC or (
-            self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[reformatted_output].dtype == torch.int
-        ):
+        if outputs[reformatted_output].dtype == torch.int:
             logits = outputs[reformatted_output]
             labels = torch.argmax(torch.nn.functional.softmax(logits, dim=-1), dim=-1)  # type: ignore
             assert colors is not None
             return colors[labels]
 
         # rendering boolean outputs
-        if self.prev_colormap_type == ColormapTypes.BOOLEAN or (
-            self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[reformatted_output].dtype == torch.bool
-        ):
+        if outputs[reformatted_output].dtype == torch.bool:
             return colormaps.apply_boolean_colormap(outputs[reformatted_output])
 
         raise NotImplementedError
@@ -689,7 +696,8 @@ class ViewerState:
         video = SingleFrameStreamTrack()
         self.video_tracks.add(video)
         video_sender = pc.addTrack(video)
-        force_codec(pc, video_sender, "video/VP8")
+        print(f"viewer using video/{self.config.codec}")
+        force_codec(pc, video_sender, f"video/{self.config.codec}")
 
         await pc.setRemoteDescription(offer)
         answer = await pc.createAnswer()
@@ -712,13 +720,12 @@ class ViewerState:
         for video_track in self.video_tracks:
             video_track.put_frame(image)
 
-    def _send_output_to_viewer(self, outputs: Dict[str, Any], colors: torch.Tensor = None, eps=1e-6):
+    def _send_output_to_viewer(self, outputs: Dict[str, Any], colors: torch.Tensor = None):
         """Chooses the correct output and sends it to the viewer
 
         Args:
             outputs: the dictionary of outputs to choose from, from the graph
             colors: is only set if colormap is for semantics. Defaults to None.
-            eps: epsilon to handle floating point comparisons
         """
         if self.output_list is None:
             self.output_list = list(outputs.keys())
@@ -731,16 +738,13 @@ class ViewerState:
 
         reformatted_output = self._process_invalid_output(self.prev_output_type)
         # re-register colormaps and send to viewer
-        if self.output_type_changed or self.prev_colormap_type == ColormapTypes.INIT:
+        if self.output_type_changed or self.prev_colormap_type is None:
             self.prev_colormap_type = ColormapTypes.DEFAULT
-            colormap_options = [ColormapTypes.DEFAULT]
-            if (
-                outputs[reformatted_output].shape[-1] != 3
-                and outputs[reformatted_output].dtype == torch.float
-                and (torch.max(outputs[reformatted_output]) - 1.0) <= eps  # handle floating point arithmetic
-            ):
-                # accumulation can also include depth
-                colormap_options.extend(["depth"])
+            colormap_options = []
+            if outputs[reformatted_output].shape[-1] == 3:
+                colormap_options = [ColormapTypes.DEFAULT]
+            if outputs[reformatted_output].shape[-1] == 1 and outputs[reformatted_output].dtype == torch.float:
+                colormap_options = list(ColormapTypes)
             self.output_type_changed = False
             self.vis["renderingState/colormap_choice"].write(self.prev_colormap_type)
             self.vis["renderingState/colormap_options"].write(colormap_options)
@@ -831,6 +835,11 @@ class ViewerState:
         if image_width > self.max_resolution:
             image_width = self.max_resolution
             image_height = int(image_width / aspect_ratio)
+        if self.config.codec != "VP8":
+            # force even values to allow hardware encoder usage
+            quantize = 2
+            image_width = int(image_width / quantize) * quantize
+            image_height = int(image_height / quantize) * quantize
         return image_height, image_width
 
     def _process_invalid_output(self, output_type: str) -> str:
@@ -879,9 +888,17 @@ class ViewerState:
 
         # check and perform colormap type updates
         colormap_type = self.vis["renderingState/colormap_choice"].read()
-        colormap_type = ColormapTypes.INIT if colormap_type is None else colormap_type
         self.prev_colormap_type = colormap_type
 
+        colormap_invert = self.vis["renderingState/colormap_invert"].read()
+        self.prev_colormap_invert = colormap_invert
+
+        colormap_normalize = self.vis["renderingState/colormap_normalize"].read()
+        self.prev_colormap_normalize = colormap_normalize
+
+        colormap_range = self.vis["renderingState/colormap_range"].read()
+        self.prev_colormap_range = colormap_range
+
         # update render aabb
         try:
             self._update_render_aabb(graph)
@@ -904,7 +921,7 @@ class ViewerState:
             return
 
         intrinsics_matrix, camera_to_world_h = get_intrinsics_matrix_and_camera_to_world_h(
-            camera_object, image_height=image_height
+            camera_object, image_height=image_height, image_width=image_width
         )
 
         camera_to_world = camera_to_world_h[:3, :]

scripts/process_data.py

@@ -350,7 +350,7 @@ class ProcessInsta360:
         if not filename_back.exists():
             raise FileNotFoundError(f"Could not find {filename_back}")
 
-        ffprobe_cmd = f"ffprobe -v quiet -print_format json -show_streams -select_streams v:0 {filename_back}"
+        ffprobe_cmd = f'ffprobe -v quiet -print_format json -show_streams -select_streams v:0 "{filename_back}"'
 
         ffprobe_output = process_data_utils.run_command(ffprobe_cmd)