diff --git a/nerfstudio/field_components/field_heads.py b/nerfstudio/field_components/field_heads.py index 89e2dd381..6fffb5111 100644 --- a/nerfstudio/field_components/field_heads.py +++ b/nerfstudio/field_components/field_heads.py @@ -40,9 +40,6 @@ class FieldHeadNames(Enum): SDF = "sdf" ALPHA = "alpha" GRADIENT = "gradient" - HASHGRID = "hashgrid" - CLIP = "clip" - DINO = "dino" class FieldHead(FieldComponent): diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 74310570c..5e488bc35 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -342,31 +342,3 @@ class NormalsRenderer(nn.Module): if normalize: n = safe_normalize(n) return n - -class CLIPRenderer(nn.Module): - """Calculate CLIP embeddings along ray.""" - - @classmethod - def forward( - cls, - embeds: TensorType["bs":..., "num_samples", "num_classes"], - weights: TensorType["bs":..., "num_samples", 1], - ) -> TensorType["bs":..., "num_classes"]: - """Calculate semantics along the ray.""" - output = torch.sum(weights * embeds, dim=-2) - output = output / torch.linalg.norm(output, dim=-1, keepdim=True) - return output - - -class MeanRenderer(nn.Module): - """Calculate average of embeddings along ray.""" - - @classmethod - def forward( - cls, - embeds: TensorType["bs":..., "num_samples", "num_classes"], - weights: TensorType["bs":..., "num_samples", 1], - ) -> TensorType["bs":..., "num_classes"]: - """Calculate semantics along the ray.""" - output = torch.sum(weights * embeds, dim=-2) - return output diff --git a/projects-lerf/lerf/lerf.py b/projects-lerf/lerf/lerf.py index 42cb17bed..836e2a85b 100644 --- a/projects-lerf/lerf/lerf.py +++ b/projects-lerf/lerf/lerf.py @@ -12,14 +12,13 @@ from nerfstudio.data.scene_box import SceneBox from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.field_components.spatial_distortions import SceneContraction from nerfstudio.model_components.ray_samplers import PDFSampler -from nerfstudio.model_components.renderers import ( - CLIPRenderer, - DepthRenderer, - MeanRenderer, -) +from nerfstudio.model_components.renderers import DepthRenderer from nerfstudio.models.nerfacto import NerfactoModel, NerfactoModelConfig + from lerf.encoders.image_encoder import BaseImageEncoder from lerf.lerf_field import LERFField +from lerf.lerf_renderers import CLIPRenderer, MeanRenderer +from lerf.lerf_fieldheadnames import LERFFieldHeadNames @dataclass @@ -108,15 +107,15 @@ class LERFModel(NerfactoModel): if self.training: outputs["clip"] = self.renderer_clip( - embeds=lerf_field_outputs[FieldHeadNames.CLIP], weights=lerf_weights.detach() + embeds=lerf_field_outputs[LERFFieldHeadNames.CLIP], weights=lerf_weights.detach() ) outputs["dino"] = self.renderer_mean( - embeds=lerf_field_outputs[FieldHeadNames.DINO], weights=lerf_weights.detach() + embeds=lerf_field_outputs[LERFFieldHeadNames.DINO], weights=lerf_weights.detach() ) if not self.training: max_across, best_scales = self.get_max_across( - lerf_samples, lerf_weights, lerf_field_outputs[FieldHeadNames.HASHGRID], clip_scales.shape + lerf_samples, lerf_weights, lerf_field_outputs[LERFFieldHeadNames.HASHGRID], clip_scales.shape ) multiphrase = max_across[0] # normalization for sanity check TODO(cmk) remove diff --git a/projects-lerf/lerf/lerf_field.py b/projects-lerf/lerf/lerf_field.py index 2b9063c83..91193de36 100644 --- a/projects-lerf/lerf/lerf_field.py +++ b/projects-lerf/lerf/lerf_field.py @@ -16,6 +16,8 @@ from nerfstudio.field_components.spatial_distortions import ( ) from nerfstudio.fields.base_field import Field +from lerf.lerf_fieldheadnames import LERFFieldHeadNames + try: import tinycudann as tcnn except ImportError: @@ -76,7 +78,7 @@ class LERFField(Field): ) return enc - def get_outputs(self, ray_samples: RaySamples, clip_scales): + def get_outputs(self, ray_samples: RaySamples, clip_scales) -> Dict[LERFFieldHeadNames, TensorType]: # random scales, one scale outputs = {} @@ -87,13 +89,13 @@ class LERFField(Field): xs = [e(positions.view(-1, 3)) for e in self.clip_encs] x = torch.concat(xs, dim=-1) - outputs[FieldHeadNames.HASHGRID] = x.view(*ray_samples.frustums.shape, -1) + outputs[LERFFieldHeadNames.HASHGRID] = x.view(*ray_samples.frustums.shape, -1) clip_pass = self.clip_net(torch.cat([x, clip_scales.view(-1, 1)], dim=-1)).view(*ray_samples.frustums.shape, -1) - outputs[FieldHeadNames.CLIP] = clip_pass / clip_pass.norm(dim=-1, keepdim=True) + outputs[LERFFieldHeadNames.CLIP] = clip_pass / clip_pass.norm(dim=-1, keepdim=True) dino_pass = self.dino_net(x).view(*ray_samples.frustums.shape, -1) - outputs[FieldHeadNames.DINO] = dino_pass + outputs[LERFFieldHeadNames.DINO] = dino_pass return outputs diff --git a/projects-lerf/lerf/lerf_fieldheadnames.py b/projects-lerf/lerf/lerf_fieldheadnames.py new file mode 100644 index 000000000..738fb22c3 --- /dev/null +++ b/projects-lerf/lerf/lerf_fieldheadnames.py @@ -0,0 +1,8 @@ +from enum import Enum +# from nerfstudio.field_components.field_heads import FieldHeadNames + +class LERFFieldHeadNames(Enum): + """Possible field outputs""" + HASHGRID = "hashgrid" + CLIP = "clip" + DINO = "dino" \ No newline at end of file diff --git a/projects-lerf/lerf/lerf_renderers.py b/projects-lerf/lerf/lerf_renderers.py new file mode 100644 index 000000000..92f6fc47a --- /dev/null +++ b/projects-lerf/lerf/lerf_renderers.py @@ -0,0 +1,31 @@ +import torch +from torch import nn +from torchtyping import TensorType + +class CLIPRenderer(nn.Module): + """Calculate CLIP embeddings along ray.""" + + @classmethod + def forward( + cls, + embeds: TensorType["bs":..., "num_samples", "num_classes"], + weights: TensorType["bs":..., "num_samples", 1], + ) -> TensorType["bs":..., "num_classes"]: + """Calculate semantics along the ray.""" + output = torch.sum(weights * embeds, dim=-2) + output = output / torch.linalg.norm(output, dim=-1, keepdim=True) + return output + + +class MeanRenderer(nn.Module): + """Calculate average of embeddings along ray.""" + + @classmethod + def forward( + cls, + embeds: TensorType["bs":..., "num_samples", "num_classes"], + weights: TensorType["bs":..., "num_samples", 1], + ) -> TensorType["bs":..., "num_classes"]: + """Calculate semantics along the ray.""" + output = torch.sum(weights * embeds, dim=-2) + return output \ No newline at end of file