diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index 548324069..a85b5ddaa 100644 --- a/nerfstudio/configs/base_config.py +++ b/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.""" diff --git a/nerfstudio/process_data/metashape_utils.py b/nerfstudio/process_data/metashape_utils.py index 491c06a2c..17d7771e8 100644 --- a/nerfstudio/process_data/metashape_utils.py +++ b/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") diff --git a/nerfstudio/process_data/polycam_utils.py b/nerfstudio/process_data/polycam_utils.py index d1887833b..079ff16f7 100644 --- a/nerfstudio/process_data/polycam_utils.py +++ b/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"]], diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index d8d9206e9..cedfc2b27 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/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) diff --git a/nerfstudio/viewer/app/package.json b/nerfstudio/viewer/app/package.json index 37f2b7bbb..171ec8dd1 100644 --- a/nerfstudio/viewer/app/package.json +++ b/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", diff --git a/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx b/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx index a1c331fad..66b92458e 100644 --- a/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx +++ b/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.', diff --git a/nerfstudio/viewer/app/src/reducer.js b/nerfstudio/viewer/app/src/reducer.js index 6de1ea7c2..d8335c80d 100644 --- a/nerfstudio/viewer/app/src/reducer.js +++ b/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], }, diff --git a/nerfstudio/viewer/server/utils.py b/nerfstudio/viewer/server/utils.py index b0763b6d1..a21b02d82 100644 --- a/nerfstudio/viewer/server/utils.py +++ b/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"): diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 93dc908b6..ce0858bb4 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/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, :] diff --git a/scripts/process_data.py b/scripts/process_data.py index 4727e2b2c..4d8e9acf7 100755 --- a/scripts/process_data.py +++ b/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)