From 6ec2920da3ffb2d96b2eab42f65342d880c7310b Mon Sep 17 00:00:00 2001 From: Chung Min Kim Date: Fri, 24 Mar 2023 23:54:01 +0000 Subject: [PATCH] Added textbox + minor bugfixes --- nerfstudio/fields/lerf_field.py | 8 ++++---- nerfstudio/models/lerf.py | 7 +------ nerfstudio/pipelines/lerf_pipeline.py | 5 ++--- .../app/src/modules/ConfigPanel/ConfigPanel.jsx | 17 +++++++++++++++++ .../viewer/app/src/modules/Scene/Scene.jsx | 2 +- nerfstudio/viewer/app/src/reducer.js | 2 ++ nerfstudio/viewer/server/viewer_utils.py | 8 ++++++++ 7 files changed, 35 insertions(+), 14 deletions(-) diff --git a/nerfstudio/fields/lerf_field.py b/nerfstudio/fields/lerf_field.py index 242e472b9..705344574 100644 --- a/nerfstudio/fields/lerf_field.py +++ b/nerfstudio/fields/lerf_field.py @@ -27,10 +27,10 @@ class LERFField(Field): super().__init__() self.spatial_distortion = spatial_distortion - self.clip_encs = torch.nn.ModuleList( - # [LERFField._get_encoding(16, 512, 16, indim=3)] - [LERFField._get_encoding(16, 128, 12, indim=3), LERFField._get_encoding(128, 512, 12, indim=3)] - ) + self.clip_encs = torch.nn.ModuleList([ + LERFField._get_encoding( 16, 128, 12, indim=3), + LERFField._get_encoding(128, 512, 12, indim=3) + ]) tot_out_dims = sum([e.n_output_dims for e in self.clip_encs]) self.clip_net = tcnn.Network( diff --git a/nerfstudio/models/lerf.py b/nerfstudio/models/lerf.py index 825e577c7..dd7a17ac6 100644 --- a/nerfstudio/models/lerf.py +++ b/nerfstudio/models/lerf.py @@ -88,11 +88,6 @@ class LERFModel(NerfactoModel): dataclass_fn = lambda dc: dc._apply_fn_to_fields(gather_fn, dataclass_fn) lerf_samples = ray_samples._apply_fn_to_fields(gather_fn, dataclass_fn) - # lerf_samples = self.lerf_sampler(ray_bundle, ray_samples, weights) - # with torch.no_grad(): - # dens, _ = self.field.get_density(lerf_samples) - # lerf_weights = lerf_samples.get_weights(dens) - if self.training: clip_scales = ray_bundle.metadata["clip_scales"] clip_scales = clip_scales[..., None] @@ -108,7 +103,7 @@ class LERFModel(NerfactoModel): for i in range(self.config.num_proposal_iterations): outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i]) - lerf_field_outputs = self.lerf_field.get_outputs(ray_samples, clip_scales) + lerf_field_outputs = self.lerf_field.get_outputs(lerf_samples, clip_scales) if self.training: outputs['clip'] = self.renderer_clip(embeds=lerf_field_outputs[FieldHeadNames.CLIP], weights=lerf_weights.detach()) diff --git a/nerfstudio/pipelines/lerf_pipeline.py b/nerfstudio/pipelines/lerf_pipeline.py index 46b087667..6f8fa7578 100644 --- a/nerfstudio/pipelines/lerf_pipeline.py +++ b/nerfstudio/pipelines/lerf_pipeline.py @@ -29,7 +29,6 @@ except ImportError: import torch - @dataclass class OpenCLIPNetworkConfig(cfg.InstantiateConfig): """Configuration for network instantiation""" @@ -55,7 +54,7 @@ class OpenCLIPNetwork: self.model = model.to("cuda") self.clip_n_dims = self.config.clip_n_dims - self.positives = ["'wes' text"] + self.positives = ["hand sanitizer"] self.negatives = self.config.negatives with torch.no_grad(): tok_phrases = torch.cat([self.tokenizer(phrase) for phrase in self.positives]).to("cuda") @@ -75,7 +74,7 @@ class OpenCLIPNetwork: def set_positives(self, text_list): self.positives = text_list with torch.no_grad(): - tok_phrases = torch.cat([self.tokenizer(phrase) for phrase in self.config.positives]).to("cuda") + tok_phrases = torch.cat([self.tokenizer(phrase) for phrase in self.positives]).to("cuda") self.pos_embeds = self.model.encode_text(tok_phrases) self.pos_embeds /= self.pos_embeds.norm(dim=-1, keepdim=True) diff --git a/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx b/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx index 66b92458e..70e6c2eaf 100644 --- a/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx +++ b/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx @@ -53,6 +53,9 @@ export function RenderControls() { const crop_enabled = useSelector( (state) => state.renderingState.crop_enabled, ); + const clip_pos_word = useSelector( + (state) => state.renderingState.clip_pos_word, + ) const crop_bg_color = useSelector( (state) => state.renderingState.crop_bg_color, @@ -206,6 +209,18 @@ export function RenderControls() { '1024px': () => setControls({ max_resolution: 1024 }), '2048px': () => setControls({ max_resolution: 2048 }), }), + clip_pos_word: { + label: 'clip_pos_word', + value: clip_pos_word, + onChange: (v) => { + dispatch_and_send( + websocket, + dispatch, + 'renderingState/clip_pos_word', + v, + ); + }, + }, // Enable Crop crop_enabled: { label: 'Crop Viewport', @@ -304,6 +319,7 @@ export function RenderControls() { setControls({ max_resolution }); setControls({ output_options: outputChoice }); setControls({ colormap_options: colormapChoice }); + setControls({ clip_pos_word }); setControls({ crop_enabled }); setControls({ crop_bg_color }); setControls({ crop_scale }); @@ -317,6 +333,7 @@ export function RenderControls() { max_resolution, target_train_util, render_time, + clip_pos_word, crop_enabled, crop_bg_color, crop_scale, diff --git a/nerfstudio/viewer/app/src/modules/Scene/Scene.jsx b/nerfstudio/viewer/app/src/modules/Scene/Scene.jsx index 8388f7735..a3257914b 100644 --- a/nerfstudio/viewer/app/src/modules/Scene/Scene.jsx +++ b/nerfstudio/viewer/app/src/modules/Scene/Scene.jsx @@ -190,7 +190,7 @@ export function get_scene_tree() { window.addEventListener('keyup', onKeyUp, true); window.addEventListener('keydown', (event) => { - if (event.code === 'Space') { + if (event.code === 'Space' && scene_state.value.get('mouse_in_scene')) { camera_controls.setLookAt(0.7, -0.7, 0.3, 0, 0, 0); } }); diff --git a/nerfstudio/viewer/app/src/reducer.js b/nerfstudio/viewer/app/src/reducer.js index 5aefe9c8a..af719731e 100644 --- a/nerfstudio/viewer/app/src/reducer.js +++ b/nerfstudio/viewer/app/src/reducer.js @@ -54,6 +54,8 @@ const initialState = { clipping_center: [0.0, 0.0, 0.0], clipping_box_scale: [2.0, 2.0, 2.0], + clip_pos_word: "", + // Crop Box Options crop_enabled: false, crop_bg_color: { r: 38, g: 42, b: 55 }, diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 53efc0a3b..9d245fe47 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -303,6 +303,7 @@ class ViewerState: self.camera_moving = False self.prev_camera_timestamp = 0 self.prev_crop_enabled = False + self.prev_clip_pos_word = "" self.prev_crop_bg_color = None self.prev_crop_scale = None self.prev_crop_center = None @@ -816,6 +817,13 @@ class ViewerState: colormap_range = self.vis["renderingState/colormap_range"].read() self.prev_colormap_range = colormap_range + clip_pos_word = self.vis["renderingState/clip_pos_word"].read() + if self.prev_clip_pos_word != clip_pos_word and clip_pos_word is not None: + print("setting target word to {}".format(clip_pos_word)) + graph.network.set_positives([clip_pos_word]) + print("done") + self.prev_clip_pos_word = clip_pos_word + # update render aabb try: self._update_render_aabb(graph)