All lerf-changes moved to projects-lerf, except viewer changes

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):

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

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

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
 

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"

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