init

.gitignore

@@ -0,0 +1,2 @@
+.cache
+__pycache__/

README.md

@@ -1,10 +1,10 @@
 ---
 title: LD T3D
-emoji: 🚀
+emoji: 🐳
 colorFrom: indigo
 colorTo: yellow
-sdk: gradio
-sdk_version: 4.22.0
+sdk: docker
+app_port: 7860
 app_file: app.py
 pinned: false
 ---

app.py

@@ -0,0 +1,129 @@
+import os
+import gradio as gr
+import numpy as np
+import torch
+import functools
+from datasets import load_dataset
+from feature_extractors.uni3d_embedding_encoder import Uni3dEmbeddingEncoder
+
+# os.environ['HTTP_PROXY'] = 'http://192.168.48.17:18000'
+# os.environ['HTTPS_PROXY'] = 'http://192.168.48.17:18000'
+
+MAX_BATCH_SIZE = 16
+MAX_QUEUE_SIZE = 10
+MAX_K_RETRIEVAL = 20
+cache_dir = "./.cache"
+
+encoder = Uni3dEmbeddingEncoder(cache_dir)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+source_id_list = torch.load("data/source_id_list.pt")
+source_to_id = {source_id: i for i, source_id in enumerate(source_id_list)}
+dataset = load_dataset("VAST-AI/LD-T3D", name=f"rendered_imgs_diag_above", split="base", cache_dir=cache_dir)
+
+@functools.lru_cache()
+def get_embedding(option, modality, angle=None):
+    save_path = f'data/objaverse_{option}_{modality + (("_" + str(angle)) if angle is not None else "")}_embeddings.pt'
+    if os.path.exists(save_path):
+        return torch.load(save_path)
+    else:
+        return gr.Error(f"Embedding file not found: {save_path}")
+
+def predict(xb, xq, top_k):
+    xb = xb.to(xq.device)
+    sim = xq @ xb.T # (nq, nb)
+    _, indices = sim.topk(k=top_k, largest=True)
+    return indices
+
+def get_image(index):
+    return dataset[index]["image"]
+
+def retrieve_3D_models(textual_query, top_k, modality_list):
+    if textual_query == "":
+        raise gr.Error("Please enter a textual query")
+    if len(textual_query.split()) > 20:
+        gr.Warning("Retrieval result may be inaccurate due to long textual query")
+    if len(modality_list) == 0:
+        raise gr.Error("Please select at least one modality")
+    
+    def _retrieve_3D_models(query, top_k, modals:list):
+        option = "uni3d"
+        op = "add"
+        is_text = True if "text" in modals else False
+        is_3D = True if "3D" in modals else False
+        if is_text:
+            modals.remove("text")
+        if is_3D:
+            modals.remove("3D")
+        angles = modals
+
+        # get base embeddings
+        embeddings = []
+        if is_text:
+            embeddings.append(get_embedding(option, "text"))
+        if len(angles) > 0:
+            for angle in angles:
+                embeddings.append(get_embedding(option, "image", angle=angle))
+        if is_3D:
+            embeddings.append(get_embedding(option, "3D"))
+            
+        ## fuse base embeddings
+        if len(embeddings) > 1:
+            if op == "concat":
+                embeddings = torch.cat(embeddings, dim=-1)
+            elif op == "add":
+                embeddings = sum(embeddings)
+            else:
+                raise ValueError(f"Unsupported operation: {op}")
+            embeddings /= embeddings.norm(dim=-1, keepdim=True)
+        else:
+            embeddings = embeddings[0]
+
+        # encode query embeddings
+        xq = encoder.encode_query(query)
+        if op == "concat":
+            xq = xq.repeat(1, embeddings.shape[-1] // xq.shape[-1]) # repeat to be aligned with the xb
+            xq /= xq.norm(dim=-1, keepdim=True)
+        
+        pred_ind_list = predict(embeddings, xq, top_k)
+        return pred_ind_list[0].cpu().tolist() # we have only one query
+
+    indices = _retrieve_3D_models(textual_query, top_k, modality_list)
+    return [get_image(index) for index in indices]
+
+def launch():
+    with gr.Blocks() as demo:
+        with gr.Row():
+            textual_query = gr.Textbox(label="Textual Query", autofocus=True,
+                                       placeholder="A chair with a wooden frame and a cushioned seat")
+            modality_list = gr.CheckboxGroup(label="Modality List", value=[],
+                                             choices=["text", "front", "back", "left", "right", "above", 
+                                                      "below", "diag_above", "diag_below", "3D"])
+        with gr.Row():
+            top_k = gr.Slider(minimum=1, maximum=MAX_K_RETRIEVAL, step=1, label="Top K Retrieval Result", 
+                              value=5, scale=2)
+            run = gr.Button("Search", scale=1)
+            clear_button = gr.ClearButton(scale=1)
+        with gr.Row():
+            output = gr.Gallery(format="webp", label="Retrieval Result", columns=5, type="pil")
+        run.click(retrieve_3D_models, [textual_query, top_k, modality_list], output, 
+                #   batch=True, max_batch_size=MAX_BATCH_SIZE
+                  )
+        clear_button.click(lambda: ["", 5, [], []], outputs=[textual_query, top_k, modality_list, output])
+        examples = gr.Examples(examples=[["An ice cream with a cherry on top", 10, ["text", "front", "back", "left", "right", "above", "below", "diag_above", "diag_below", "3D"]],
+                                         ["A mid-age castle", 10, ["text", "front", "back", "left", "right", "above", "below", "diag_above", "diag_below", "3D"]],
+                                         ["A coke", 10, ["text", "front", "back", "left", "right", "above", "below", "diag_above", "diag_below", "3D"]]],
+                            inputs=[textual_query, top_k, modality_list],
+                            # cache_examples=True,
+                            outputs=output,
+                            fn=retrieve_3D_models)
+        
+    demo.queue(max_size=10)
+
+    # os.environ.pop('HTTP_PROXY')
+    # os.environ.pop('HTTPS_PROXY')
+
+    demo.launch(server_name='0.0.0.0')
+
+if __name__ == "__main__":
+    launch()
+    # print(len(retrieve_3D_models("A chair with a wooden frame and a cushioned seat", 5, ["3D", "diag_above", "diag_below"])))

change_setup.txt

@@ -0,0 +1,38 @@
+import glob
+import os
+import os.path as osp
+
+from setuptools import find_packages, setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+this_dir = osp.dirname(osp.abspath(__file__))
+_ext_src_root = osp.join("pointnet2_ops", "_ext-src")
+_ext_sources = glob.glob(osp.join(_ext_src_root, "src", "*.cpp")) + glob.glob(
+    osp.join(_ext_src_root, "src", "*.cu")
+)
+_ext_headers = glob.glob(osp.join(_ext_src_root, "include", "*"))
+
+requirements = ["torch>=1.4"]
+
+exec(open(osp.join("pointnet2_ops", "_version.py")).read())
+
+setup(
+    name="pointnet2_ops",
+    version=__version__,
+    author="Erik Wijmans",
+    packages=find_packages(),
+    install_requires=requirements,
+    ext_modules=[
+        CUDAExtension(
+            name="pointnet2_ops._ext",
+            sources=_ext_sources,
+            extra_compile_args={
+                "cxx": ["-O3"],
+                "nvcc": ["-O3", "-Xfatbin", "-compress-all"],
+            },
+            include_dirs=[osp.join(this_dir, _ext_src_root, "include")],
+        )
+    ],
+    cmdclass={"build_ext": BuildExtension},
+    include_package_data=True,
+)

data/objaverse_uni3d_3D_embeddings.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b05400ab75009785535bd78d859db0a902176fbeb5df2ef73e55a95990ded1b8
+size 365511995

data/objaverse_uni3d_image_above_embeddings.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c0708d9bfb4df4e6f86a21bd5a1096401c8c037e84575e6d0397efdb1b138289
+size 365512104

data/objaverse_uni3d_image_back_embeddings.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5667981bc1215e1f60c034ff8e2d214da6186a2f3212061b8ed3e1c32073ad6e
+size 365512104

data/objaverse_uni3d_image_below_embeddings.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f91df0329424657666dd9a5b3181d52f9155ad545dc22a2f725f24f9b854abbd
+size 365512104

data/objaverse_uni3d_image_diag_above_embeddings.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b44e2ee38885128e9080c75ee1d311fee8f718375e867c2209273649455c89a7
+size 365512035

data/objaverse_uni3d_image_diag_below_embeddings.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:79eb0da600d75874e22bbfcca6001669eb14f06ec37326bf5148521db82f3e34
+size 365512035

data/objaverse_uni3d_image_front_embeddings.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:016208fa7a76e959840c128c30e178a0b43a570cf7a8e6cfd6fcdb442f6b72db
+size 365512104

data/objaverse_uni3d_image_left_embeddings.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5db0c17a56ebbb0fa1323b105dfe04386f8d7f88c876bc24b943e8713a01076
+size 365512035

data/objaverse_uni3d_image_right_embeddings.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f5fb149475c79b465157d5b2cfe2af4ad8947ff23f99577da264c2632bc9d770
+size 365512035

data/objaverse_uni3d_text_embeddings.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a2d908630bcc8a5a231e8b5d11714c63a3e8b6d78427a82a833da9219b2a7263
+size 365512020

data/source_id_list.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c218ccb58d0045b0b6671c1378ee43362054b890f9895d7cac3de727683a9a76
+size 3747900

dockerfile

@@ -0,0 +1,19 @@
+
+FROM nvcr.io/nvidia/pytorch:23.08
+
+LABEL maintainer="yuanze"
+LABEL email="yuanze1024@gmail.com"
+
+# Install webp support
+RUN apt update && apt install libwebp-dev -y
+
+RUN pip install -r requirements.txt
+
+# note that you may need to modify the TORCH_CUDA_ARCH_LIST in the setup.py file
+ENV TORCH_CUDA_ARCH_LIST="8.6"
+
+# Install Pointnet2_PyTorch
+RUN git clone https://github.com/erikwijmans/Pointnet2_PyTorch.git \
+    && mv -f backup_install.txt Pointnet2_PyTorch/pointnet2_ops_lib/setup.py \
+    && cd Pointnet2_PyTorch/pointnet2_ops_lib \
+    && python install .

feature_extractors/__init__.py

@@ -0,0 +1,56 @@
+from collections.abc import Sequence
+from abc import ABC, abstractmethod
+import torch
+from PIL.Image import Image
+
+class FeatureExtractor(ABC):
+    @abstractmethod
+    def encode_image(self, img_list: Sequence[Image]) -> torch.Tensor:
+        """
+        Encode the input images and return the corresponding embeddings.
+
+        Args:
+            img_list: A list of PIL.Image.Image objects.
+
+        Returns:
+            The embeddings of the input images. The shape should be (len(img_list), embedding_dim).
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def encode_text(self, text_list: Sequence[str]) -> torch.Tensor:
+        """
+        Encode the input text data and return the corresponding embeddings.
+
+        Args:
+            text_list: A list of strings.
+
+        Returns:
+            The embeddings of the input text data. The shape should be (len(text_list), embedding_dim).
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def encode_3D(self, pc_tensor: torch.Tensor) -> torch.Tensor:
+        """
+        Encode the input 3D point cloud and return the corresponding embeddings.
+
+        Args:
+            pc_tensor: A tensor of shape (B, N, 3 + 3).
+        
+        Returns:
+            The embeddings of the input 3D point cloud. The shape should be (B, embedding_dim).
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def encode_query(self, queries: Sequence[str]) -> torch.Tensor:
+        """Encode the queries and return the corresponding embeddings.
+
+        Args:
+            queries: A list of strings.
+
+        Returns:
+            The embeddings of the input text data. The shape should be (len(input_text), embedding_dim).
+        """
+        raise NotImplementedError

feature_extractors/uni3d_embedding_encoder.py

@@ -0,0 +1,337 @@
+"""
+See https://github.com/baaivision/Uni3D for source code
+"""
+import os
+import torch
+import torch.nn as nn
+import timm
+import numpy as np
+from pointnet2_ops import pointnet2_utils
+import open_clip
+from huggingface_hub import hf_hub_download
+import sys
+sys.path.append('')
+from feature_extractors import FeatureExtractor
+from utils.tokenizer import SimpleTokenizer
+
+import logging
+
+def fps(data, number):
+    '''
+        data B N 3
+        number int
+    '''
+    fps_idx = pointnet2_utils.furthest_point_sample(data, number) 
+    fps_data = pointnet2_utils.gather_operation(data.transpose(1, 2).contiguous(), fps_idx).transpose(1,2).contiguous()
+    return fps_data
+
+# https://github.com/Strawberry-Eat-Mango/PCT_Pytorch/blob/main/util.py 
+def knn_point(nsample, xyz, new_xyz):
+    """
+    Input:
+        nsample: max sample number in local region
+        xyz: all points, [B, N, C]
+        new_xyz: query points, [B, S, C]
+    Return:
+        group_idx: grouped points index, [B, S, nsample]
+    """
+    sqrdists = square_distance(new_xyz, xyz)
+    _, group_idx = torch.topk(sqrdists, nsample, dim = -1, largest=False, sorted=False)
+    return group_idx
+
+def square_distance(src, dst):
+    """
+    Calculate Euclid distance between each two points.
+    src^T * dst = xn * xm + yn * ym + zn * zm;
+    sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
+    sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
+    dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
+         = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst
+    Input:
+        src: source points, [B, N, C]
+        dst: target points, [B, M, C]
+    Output:
+        dist: per-point square distance, [B, N, M]
+    """
+    B, N, _ = src.shape
+    _, M, _ = dst.shape
+    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
+    dist += torch.sum(src ** 2, -1).view(B, N, 1)
+    dist += torch.sum(dst ** 2, -1).view(B, 1, M)
+    return dist    
+
+
+class PatchDropout(nn.Module):
+    """
+    https://arxiv.org/abs/2212.00794
+    """
+
+    def __init__(self, prob, exclude_first_token=True):
+        super().__init__()
+        assert 0 <= prob < 1.
+        self.prob = prob
+        self.exclude_first_token = exclude_first_token  # exclude CLS token
+        logging.info("patch dropout prob is {}".format(prob))
+
+    def forward(self, x):
+        # if not self.training or self.prob == 0.:
+        #     return x
+
+        if self.exclude_first_token:
+            cls_tokens, x = x[:, :1], x[:, 1:]
+        else:
+            cls_tokens = torch.jit.annotate(torch.Tensor, x[:, :1])
+
+        batch = x.size()[0]
+        num_tokens = x.size()[1]
+
+        batch_indices = torch.arange(batch)
+        batch_indices = batch_indices[..., None]
+
+        keep_prob = 1 - self.prob
+        num_patches_keep = max(1, int(num_tokens * keep_prob))
+
+        rand = torch.randn(batch, num_tokens)
+        patch_indices_keep = rand.topk(num_patches_keep, dim=-1).indices
+
+        x = x[batch_indices, patch_indices_keep]
+
+        if self.exclude_first_token:
+            x = torch.cat((cls_tokens, x), dim=1)
+
+        return x
+
+
+class Group(nn.Module):
+    def __init__(self, num_group, group_size):
+        super().__init__()
+        self.num_group = num_group
+        self.group_size = group_size
+
+    def forward(self, xyz, color):
+        '''
+            input: B N 3
+            ---------------------------
+            output: B G M 3
+            center : B G 3
+        '''
+        batch_size, num_points, _ = xyz.shape
+        # fps the centers out
+        center = fps(xyz, self.num_group) # B G 3
+        # knn to get the neighborhood
+        # _, idx = self.knn(xyz, center) # B G M
+        idx = knn_point(self.group_size, xyz, center) # B G M
+        assert idx.size(1) == self.num_group
+        assert idx.size(2) == self.group_size
+        idx_base = torch.arange(0, batch_size, device=xyz.device).view(-1, 1, 1) * num_points
+        idx = idx + idx_base
+        idx = idx.view(-1)
+        neighborhood = xyz.view(batch_size * num_points, -1)[idx, :]
+        neighborhood = neighborhood.view(batch_size, self.num_group, self.group_size, 3).contiguous()
+
+        neighborhood_color = color.view(batch_size * num_points, -1)[idx, :]
+        neighborhood_color = neighborhood_color.view(batch_size, self.num_group, self.group_size, 3).contiguous()
+
+        # normalize
+        neighborhood = neighborhood - center.unsqueeze(2)
+
+        features = torch.cat((neighborhood, neighborhood_color), dim=-1)
+        return neighborhood, center, features
+
+class Encoder(nn.Module):
+    def __init__(self, encoder_channel):
+        super().__init__()
+        self.encoder_channel = encoder_channel
+        self.first_conv = nn.Sequential(
+            nn.Conv1d(6, 128, 1),
+            nn.BatchNorm1d(128),
+            nn.ReLU(inplace=True),
+            nn.Conv1d(128, 256, 1)
+        )
+        self.second_conv = nn.Sequential(
+            nn.Conv1d(512, 512, 1),
+            nn.BatchNorm1d(512),
+            nn.ReLU(inplace=True),
+            nn.Conv1d(512, self.encoder_channel, 1)
+        )
+    def forward(self, point_groups):
+        '''
+            point_groups : B G N 3
+            -----------------
+            feature_global : B G C
+        '''
+        bs, g, n , _ = point_groups.shape
+        point_groups = point_groups.reshape(bs * g, n, 6)
+        # encoder
+        feature = self.first_conv(point_groups.transpose(2,1))  # BG 256 n
+        feature_global = torch.max(feature,dim=2,keepdim=True)[0]  # BG 256 1
+        feature = torch.cat([feature_global.expand(-1,-1,n), feature], dim=1)# BG 512 n
+        feature = self.second_conv(feature) # BG 1024 n
+        feature_global = torch.max(feature, dim=2, keepdim=False)[0] # BG 1024
+        return feature_global.reshape(bs, g, self.encoder_channel)
+
+class PointcloudEncoder(nn.Module):
+    def __init__(self, point_transformer):
+        # use the giant branch of uni3d
+        super().__init__()
+        from easydict import EasyDict
+        self.trans_dim = 1408
+        self.embed_dim = 1024
+        self.group_size = 64
+        self.num_group = 512
+        # grouper
+        self.group_divider = Group(num_group = self.num_group, group_size = self.group_size)
+        # define the encoder
+        self.encoder_dim = 512
+        self.encoder = Encoder(encoder_channel = self.encoder_dim)
+       
+        # bridge encoder and transformer
+        self.encoder2trans = nn.Linear(self.encoder_dim,  self.trans_dim)
+        
+        # bridge transformer and clip embedding
+        self.trans2embed = nn.Linear(self.trans_dim,  self.embed_dim)
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.trans_dim))
+        self.cls_pos = nn.Parameter(torch.randn(1, 1, self.trans_dim))
+
+        self.pos_embed = nn.Sequential(
+            nn.Linear(3, 128),
+            nn.GELU(),
+            nn.Linear(128, self.trans_dim)
+        )  
+        # setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn
+        self.patch_dropout = PatchDropout(0.) if 0. > 0. else nn.Identity()
+        self.visual = point_transformer
+
+
+    def forward(self, pts, colors):
+        # divide the point cloud in the same form. This is important
+        _, center, features = self.group_divider(pts, colors)
+
+        # encoder the input cloud patches
+        group_input_tokens = self.encoder(features)  #  B G N
+        group_input_tokens = self.encoder2trans(group_input_tokens)
+        # prepare cls
+        cls_tokens = self.cls_token.expand(group_input_tokens.size(0), -1, -1)  
+        cls_pos = self.cls_pos.expand(group_input_tokens.size(0), -1, -1)  
+        # add pos embedding
+        pos = self.pos_embed(center)
+        # final input
+        x = torch.cat((cls_tokens, group_input_tokens), dim=1)
+        pos = torch.cat((cls_pos, pos), dim=1)
+        # transformer
+        x = x + pos
+        # x = x.half()
+        
+        # a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in
+        x = self.patch_dropout(x)
+
+        x = self.visual.pos_drop(x)
+
+        # ModuleList not support forward
+        for i, blk in enumerate(self.visual.blocks):
+            x = blk(x)
+        x = self.visual.norm(x[:, 0, :])
+        x = self.visual.fc_norm(x)
+
+        x = self.trans2embed(x)
+        return x
+
+class Uni3D(nn.Module):
+    def __init__(self, point_encoder):
+        super().__init__()
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+        self.point_encoder = point_encoder
+
+    def encode_pc(self, pc):
+        xyz = pc[:,:,:3].contiguous()
+        color = pc[:,:,3:].contiguous()
+        pc_feat = self.point_encoder(xyz, color)
+        return pc_feat
+
+    def forward(self, pc, text, image):
+        text_embed_all = text
+        image_embed = image   
+        pc_embed = self.encode_pc(pc)
+        return {'text_embed': text_embed_all,
+                'pc_embed': pc_embed,
+                'image_embed': image_embed,
+                'logit_scale': self.logit_scale.exp()}
+
+def get_metric_names(model):
+    return ['loss', 'uni3d_loss', 'pc_image_acc', 'pc_text_acc']
+
+def create_uni3d(uni3d_path):  
+    # create transformer blocks for point cloud via timm
+    point_transformer = timm.create_model("eva_giant_patch14_560")
+
+    # create whole point cloud encoder
+    point_encoder = PointcloudEncoder(point_transformer)
+
+    # uni3d model
+    model = Uni3D(point_encoder=point_encoder,)
+
+    checkpoint = torch.load(uni3d_path, map_location='cpu')
+    logging.info('loaded checkpoint {}'.format(uni3d_path))
+    sd = checkpoint['module']
+    if next(iter(sd.items()))[0].startswith('module'):
+        sd = {k[len('module.'):]: v for k, v in sd.items()}
+    model.load_state_dict(sd)
+    return model
+
+class Uni3dEmbeddingEncoder(FeatureExtractor):
+    def __init__(self, cache_dir, **kwargs) -> None:
+        bpe_path = "utils/bpe_simple_vocab_16e6.txt.gz"
+        uni3d_path = os.path.join(cache_dir, "Uni3D", "modelzoo", "uni3d-g", "model.pt") # concat the subfolder as hf_hub_download will put it here
+        clip_path = os.path.join(cache_dir, "Uni3D", "open_clip_pytorch_model.bin")
+
+        if not os.path.exists(uni3d_path):
+            hf_hub_download("BAAI/Uni3D", "model.pt", subfolder="modelzoo/uni3d-g", cache_dir=cache_dir, 
+                            local_dir=cache_dir + os.sep + "Uni3D")
+        if not os.path.exists(clip_path):
+            hf_hub_download("timm/eva02_enormous_patch14_plus_clip_224.laion2b_s9b_b144k", "open_clip_pytorch_model.bin", 
+                            cache_dir=cache_dir, local_dir=cache_dir + os.sep + "Uni3D")
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.tokenizer = SimpleTokenizer(bpe_path)
+        self.model = create_uni3d(uni3d_path)
+        self.model.eval()
+        self.model.to(self.device)
+        self.clip_model, _, self.preprocess = open_clip.create_model_and_transforms(model_name="EVA02-E-14-plus", pretrained=clip_path)
+        self.clip_model.to(self.device)
+        
+    def pc_norm(self, pc):
+        """ pc: NxC, return NxC """
+        centroid = np.mean(pc, axis=0)
+        pc = pc - centroid
+        m = np.max(np.sqrt(np.sum(pc ** 2, axis=1)))
+        pc = pc / m
+        return pc
+
+    @torch.no_grad()
+    def encode_3D(self, data):
+        pc = data.to(device=self.device, non_blocking=True)
+        pc_features = self.model.encode_pc(pc)
+        pc_features = pc_features / pc_features.norm(dim=-1, keepdim=True)
+        return pc_features.float()
+
+    @torch.no_grad()
+    def encode_text(self, input_text):
+        texts = self.tokenizer(input_text).to(device=self.device, non_blocking=True)
+        if len(texts.shape) < 2:
+            texts = texts[None, ...]
+        class_embeddings = self.clip_model.encode_text(texts)
+        class_embeddings = class_embeddings / class_embeddings.norm(dim=-1, keepdim=True)
+        return class_embeddings.float()
+
+    @torch.no_grad()
+    def encode_image(self, img_tensor_list):
+        image = img_tensor_list.to(device=self.device, non_blocking=True)
+        image_features = self.clip_model.encode_image(image)
+        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
+        return image_features.float()
+
+    def encode_query(self, query_list):
+        return self.encode_text(query_list)
+    
+    def get_img_transform(self):
+        return self.preprocess

packages

@@ -0,0 +1 @@
+libwebp-dev

requirements.txt

@@ -0,0 +1,9 @@
+gradio
+datasets
+timm
+pillow
+open-clip-torch
+huggingface_hub
+ftfy
+regex
+easydict

utils/bpe_simple_vocab_16e6.txt.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a
+size 1356917

utils/tokenizer.py

@@ -0,0 +1,147 @@
+# copied from github.com/baaivision/Uni3D
+# # Modified from github.com/openai/CLIP
+import gzip
+import html
+import os
+from functools import lru_cache
+
+import ftfy
+import regex as re
+import torch
+
+
+@lru_cache()
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a corresponding list of unicode strings.
+    The reversible bpe codes work on unicode strings.
+    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
+    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
+    This is a signficant percentage of your normal, say, 32K bpe vocab.
+    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
+    And avoids mapping to whitespace/control characters the bpe code barfs on.
+    """
+    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8+n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+    """Return set of symbol pairs in a word.
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    return text
+
+
+class SimpleTokenizer(object):
+    def __init__(self, bpe_path):
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        merges = gzip.open(bpe_path).read().decode("utf-8").split('\n')
+        merges = merges[1:49152-256-2+1]
+        merges = [tuple(merge.split()) for merge in merges]
+        vocab = list(bytes_to_unicode().values())
+        vocab = vocab + [v+'</w>' for v in vocab]
+        for merge in merges:
+            vocab.append(''.join(merge))
+        vocab.extend(['<|startoftext|>', '<|endoftext|>'])
+        self.encoder = dict(zip(vocab, range(len(vocab))))
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.bpe_ranks = dict(zip(merges, range(len(merges))))
+        self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
+        self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE)
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token[:-1]) + ( token[-1] + '</w>',)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token+'</w>'
+
+        while True:
+            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except:
+                    new_word.extend(word[i:])
+                    break
+
+                if word[i] == first and i < len(word)-1 and word[i+1] == second:
+                    new_word.append(first+second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = ' '.join(word)
+        self.cache[token] = word
+        return word
+
+    def encode(self, text):
+        bpe_tokens = []
+        text = whitespace_clean(basic_clean(text)).lower()
+        for token in re.findall(self.pat, text):
+            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
+            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
+        return bpe_tokens
+
+    def decode(self, tokens):
+        text = ''.join([self.decoder[token] for token in tokens])
+        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
+        return text
+
+    def __call__(self, texts, context_length=77):
+        if isinstance(texts, str):
+            texts = [texts]
+
+        sot_token = self.encoder["<|startoftext|>"]
+        eot_token = self.encoder["<|endoftext|>"]
+        all_tokens = [[sot_token] + self.encode(text) + [eot_token] for text in texts]
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+
+        for i, tokens in enumerate(all_tokens):
+            tokens = tokens[:context_length]
+            result[i, :len(tokens)] = torch.tensor(tokens)
+
+        if len(result) == 1:
+            return result[0]
+        return result