Create a base_encoder object, + refactor

projects-lerf/lerf/encoders/clip_encoder.py

@@ -0,0 +1,87 @@
+from dataclasses import dataclass, field
+from typing import Tuple, Type
+
+import torch
+import torchvision
+
+try:
+    import clip
+except ImportError:
+    assert False, "clip is not installed, install it with `pip install clip`"
+
+from lerf.encoders.image_encoder import BaseImageEncoder, BaseImageEncoderConfig
+
+@dataclass
+class CLIPNetworkConfig(BaseImageEncoderConfig):
+    _target: Type = field(default_factory=lambda: CLIPNetwork)
+    name: str = "clip"
+    description: str = "OpenAI CLIP network, images and text"
+    """target class to instantiate"""
+    clip_model_type: str = "ViT-B-16"
+    clip_n_dims: int = 512
+    negatives: Tuple[str] = ("object", "things", "stuff", "texture")
+
+class CLIPNetwork(BaseImageEncoder):
+    def __init__(self, config: CLIPNetworkConfig):
+        self.config = config
+        self.process = torchvision.transforms.Compose(
+            [
+                torchvision.transforms.Resize((224, 224)),
+                torchvision.transforms.Normalize(
+                    mean=[0.48145466, 0.4578275, 0.40821073],
+                    std=[0.26862954, 0.26130258, 0.27577711],
+                ),
+            ]
+        )
+        model, _ = clip.load(self.config.clip_model_type)
+        model.eval()
+        self.tokenizer = clip.tokenize
+        self.model = model.to("cuda")
+        self.clip_n_dims = self.config.clip_n_dims
+
+        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")
+            self.pos_embeds = model.encode_text(tok_phrases)
+            tok_phrases = torch.cat([self.tokenizer(phrase) for phrase in self.negatives]).to("cuda")
+            self.neg_embeds = model.encode_text(tok_phrases)
+        self.pos_embeds /= self.pos_embeds.norm(dim=-1, keepdim=True)
+        self.neg_embeds /= self.neg_embeds.norm(dim=-1, keepdim=True)
+
+        assert (
+            self.pos_embeds.shape[1] == self.neg_embeds.shape[1]
+        ), "Positive and negative embeddings must have the same dimensionality"
+        assert (
+            self.pos_embeds.shape[1] == self.clip_n_dims
+        ), "Embedding dimensionality must match the model dimensionality"
+
+    @property
+    def embedding_dim(self) -> int:
+        return self.config.clip_n_dims
+
+    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.positives]).to("cuda")
+            self.pos_embeds = self.model.encode_text(tok_phrases)
+        self.pos_embeds /= self.pos_embeds.norm(dim=-1, keepdim=True)
+
+    def get_relevancy(self, embed: torch.Tensor, positive_id: int) -> torch.Tensor:
+        phrases_embeds = torch.cat([self.pos_embeds, self.neg_embeds], dim=0)
+        p = phrases_embeds.to(embed.dtype)  # phrases x 512
+        output = torch.mm(embed, p.T)  # rays x phrases
+        positive_vals = output[..., positive_id : positive_id + 1]  # rays x 1
+        negative_vals = output[..., len(self.positives) :]  # rays x N_phrase
+        repeated_pos = positive_vals.repeat(1, len(self.negatives))  # rays x N_phrase
+
+        sims = torch.stack((repeated_pos, negative_vals), dim=-1)  # rays x N-phrase x 2
+        softmax = torch.softmax(10 * sims, dim=-1)  # rays x n-phrase x 2
+        best_id = softmax[..., 0].argmin(dim=1)  # rays x 2
+        return torch.gather(softmax, 1, best_id[..., None, None].expand(best_id.shape[0], len(self.negatives), 2))[
+            :, 0, :
+        ]
+
+    def encode_image(self, input):
+        processed_input = self.process(input).half()
+        return self.model.encode_image(processed_input)

projects-lerf/lerf/encoders/image_encoder.py

@@ -0,0 +1,35 @@
+from abc import abstractmethod, abstractproperty
+from dataclasses import dataclass, field
+from typing import Type
+
+import torch
+
+from nerfstudio.configs import base_config as cfg
+
+# TODO(cmk) want to rename ImageEncoder to be more generic
+@dataclass
+class BaseImageEncoderConfig(cfg.InstantiateConfig):
+    _target: Type = field(default_factory=lambda: BaseImageEncoder)
+    name: str = ""
+    """name of the encoder"""
+    description: str = ""
+    """description of the encoder"""
+
+class BaseImageEncoder:
+    @abstractproperty
+    def embedding_dim(self) -> int:
+        """
+        returns the dimension of the embeddings
+        """
+
+    @abstractmethod
+    def encode_image(self, input: torch.Tensor) -> torch.Tensor:
+        """
+        Given a batch of input images, return their encodings
+        """
+
+    @abstractmethod
+    def get_relevancy(self, embed: torch.Tensor, positive_id: int) -> torch.Tensor:
+        """
+        Given a batch of embeddings, return the relevancy to the given positive id
+        """

projects-lerf/lerf/encoders/openclip_encoder.py

@@ -1,25 +1,21 @@
-import typing
-from abc import abstractmethod, abstractproperty
 from dataclasses import dataclass, field
-from typing import List, Literal, Tuple, Type
+from typing import Tuple, Type
 
+import torch
 import torchvision
 
-from nerfstudio.configs import base_config as cfg
-
 try:
     import open_clip
 except ImportError:
     assert False, "open_clip is not installed, install it with `pip install open-clip-torch`"
 
-import torch
-
+from lerf.encoders.image_encoder import BaseImageEncoder, BaseImageEncoderConfig
 
 @dataclass
-class OpenCLIPNetworkConfig(cfg.InstantiateConfig):
-    """Configuration for network instantiation"""
-
+class OpenCLIPNetworkConfig(BaseImageEncoderConfig):
     _target: Type = field(default_factory=lambda: OpenCLIPNetwork)
+    name: str = "openclip"
+    description: str = "OpenCLIP network, images and text"
     """target class to instantiate"""
     clip_model_type: str = "ViT-B-16"
     clip_model_pretrained: str = "laion2b_s34b_b88k"
@@ -29,29 +25,7 @@ class OpenCLIPNetworkConfig(cfg.InstantiateConfig):
     # clip_n_dims: int = 768
     negatives: Tuple[str] = ("object", "things", "stuff", "texture")
 
-
-class ImageEncoder:
-    @abstractproperty
-    def embedding_dim(self) -> int:
-        """
-        returns the dimension of the embeddings
-        """
-
-    @abstractmethod
-    def encode_image(self, input: torch.Tensor) -> torch.Tensor:
-        """
-        Given a batch of input images, return their encodings
-        """
-
-    @abstractmethod
-    def get_relevancy(self, embed: torch.Tensor, positive_id: int) -> torch.Tensor:
-        """
-        Given a batch of embeddings, return the relevancy to the given positive id
-        """
-
-
-# ambivalent about this being here, it's still technically rendering but I want to keep it with the OpenCLIPNetworkConfig
-class OpenCLIPNetwork(ImageEncoder):
+class OpenCLIPNetwork(BaseImageEncoder):
     def __init__(self, config: OpenCLIPNetworkConfig):
         self.config = config
         self.process = torchvision.transforms.Compose(
@@ -118,4 +92,4 @@ class OpenCLIPNetwork(ImageEncoder):
 
     def encode_image(self, input):
         processed_input = self.process(input).half()
-        return self.model.encode_image(processed_input)
+        return self.model.encode_image(processed_input)

projects-lerf/lerf/lerf.py

@@ -18,7 +18,7 @@ from nerfstudio.model_components.renderers import (
     MeanRenderer,
 )
 from nerfstudio.models.nerfacto import NerfactoModel, NerfactoModelConfig
-from lerf.lerf_encoders import ImageEncoder
+from lerf.encoders.image_encoder import BaseImageEncoder
 from lerf.lerf_field import LERFField
 
 
@@ -42,7 +42,7 @@ class LERFModel(NerfactoModel):
         self.renderer_clip = CLIPRenderer()
         self.renderer_mean = MeanRenderer()
 
-        self.image_encoder: ImageEncoder = self.kwargs["image_encoder"]
+        self.image_encoder: BaseImageEncoder = self.kwargs["image_encoder"]
         self.lerf_field = LERFField(clip_n_dims=self.image_encoder.embedding_dim)
 
     def get_max_across(self, ray_samples, weights, hashgrid_field, scales_shape, preset_scales=None):

projects-lerf/lerf/lerf_config.py

@@ -13,7 +13,12 @@ from nerfstudio.plugins.types import MethodSpecification
 from lerf.data.lerf_datamanager import LERFDataManagerConfig
 from lerf.lerf import LERFModelConfig
 from lerf.lerf_pipeline import LERFPipelineConfig
-from lerf.lerf_encoders import OpenCLIPNetworkConfig
+
+"""
+Swap out the network config to use OpenCLIP or CLIP here.
+"""
+from lerf.encoders.openclip_encoder import OpenCLIPNetworkConfig
+from lerf.encoders.clip_encoder import CLIPNetworkConfig
 
 lerf_method = MethodSpecification(
   config=TrainerConfig(

projects-lerf/lerf/lerf_pipeline.py

@@ -17,7 +17,7 @@ from lerf.data.lerf_datamanager import (
     LERFDataManagerConfig,
 )
 from lerf.lerf import LERFModel, LERFModelConfig
-from lerf.lerf_encoders import OpenCLIPNetwork, OpenCLIPNetworkConfig
+from lerf.encoders.image_encoder import BaseImageEncoderConfig, BaseImageEncoder
 
 
 @dataclass
@@ -30,7 +30,7 @@ class LERFPipelineConfig(VanillaPipelineConfig):
     """specifies the datamanager config"""
     model: ModelConfig = LERFModelConfig()
     """specifies the model config"""
-    network: OpenCLIPNetworkConfig = OpenCLIPNetworkConfig()
+    network: BaseImageEncoderConfig = BaseImageEncoderConfig()
     """specifies the vision-language network config"""
 
 
@@ -47,7 +47,7 @@ class LERFPipeline(VanillaPipeline):
         self.config = config
         self.test_mode = test_mode
 
-        self.image_encoder: OpenCLIPNetwork = config.network.setup()
+        self.image_encoder: BaseImageEncoder = config.network.setup()
 
         self.datamanager: LERFDataManager = config.datamanager.setup(
             device=device,