From 18208f7aaa4f8602a8835220def2f0a3ea47d5b5 Mon Sep 17 00:00:00 2001 From: Chung Min Kim Date: Fri, 31 Mar 2023 22:38:23 +0000 Subject: [PATCH] Create a base_encoder object, + refactor --- projects-lerf/lerf/encoders/clip_encoder.py | 87 +++++++++++++++++++ projects-lerf/lerf/encoders/image_encoder.py | 35 ++++++++ .../openclip_encoder.py} | 42 ++------- projects-lerf/lerf/lerf.py | 4 +- projects-lerf/lerf/lerf_config.py | 7 +- projects-lerf/lerf/lerf_pipeline.py | 6 +- 6 files changed, 141 insertions(+), 40 deletions(-) create mode 100644 projects-lerf/lerf/encoders/clip_encoder.py create mode 100644 projects-lerf/lerf/encoders/image_encoder.py rename projects-lerf/lerf/{lerf_encoders.py => encoders/openclip_encoder.py} (78%) diff --git a/projects-lerf/lerf/encoders/clip_encoder.py b/projects-lerf/lerf/encoders/clip_encoder.py new file mode 100644 index 000000000..097e08dd2 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/projects-lerf/lerf/encoders/image_encoder.py b/projects-lerf/lerf/encoders/image_encoder.py new file mode 100644 index 000000000..ffbf75aa3 --- /dev/null +++ b/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 + """ \ No newline at end of file diff --git a/projects-lerf/lerf/lerf_encoders.py b/projects-lerf/lerf/encoders/openclip_encoder.py similarity index 78% rename from projects-lerf/lerf/lerf_encoders.py rename to projects-lerf/lerf/encoders/openclip_encoder.py index 5a7ccc7ca..ad68a3897 100644 --- a/projects-lerf/lerf/lerf_encoders.py +++ b/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) \ No newline at end of file diff --git a/projects-lerf/lerf/lerf.py b/projects-lerf/lerf/lerf.py index f9d82d77d..42cb17bed 100644 --- a/projects-lerf/lerf/lerf.py +++ b/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): diff --git a/projects-lerf/lerf/lerf_config.py b/projects-lerf/lerf/lerf_config.py index f8f3c4914..d8df60219 100644 --- a/projects-lerf/lerf/lerf_config.py +++ b/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( diff --git a/projects-lerf/lerf/lerf_pipeline.py b/projects-lerf/lerf/lerf_pipeline.py index 8ab166536..ea50118f8 100644 --- a/projects-lerf/lerf/lerf_pipeline.py +++ b/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,