Update openshape/__init__.py

openshape/__init__.py

@@ -2,6 +2,7 @@ import torch
 import torch.nn as nn
 from huggingface_hub import hf_hub_download
 from .ppat_rgb import Projected, PointPatchTransformer
+from .mlp import MLP
 
 
 def module(state_dict: dict, name):
@@ -32,7 +33,7 @@ def B32(s):
     return model
 
 
-def G14_shapenet_mix(s):
+def G14_M(s):
     model = Projected(
         PointPatchTransformer(512, 12, 8, 512*3, 256, 384, 0.2, 64, 6),
         nn.Linear(512, 1280)
@@ -40,11 +41,17 @@ def G14_shapenet_mix(s):
     model.load_state_dict(module(s['state_dict'], 'module'))
     return model
 
+def PCA(s):
+    model = MLP(in_out_features=1280)
+    model.load_state_dict(module(s['pc_augment_adapter'], 'module'))
+    return model
+
 model_list = {
     "openshape-pointbert-vitb32-rgb": B32,
     "openshape-pointbert-vitl14-rgb": L14,
     "openshape-pointbert-vitg14-rgb": G14,
-    "tripletmix-pointbert-shapenet": G14_shapenet_mix,
+    "tripletmix-pointbert-shapenet": G14_M,
+    "pc_adapter": PCA,
 }
 
 
@@ -57,8 +64,13 @@ def load_pc_encoder(name):
 
 
 def load_pc_encoder_mix(name):
+    pc_adapter = None
     s = torch.load(hf_hub_download("TripletMix/" + name, "model.pt"), map_location='cpu')
     model = model_list[name](s).eval()
+    if name ==  "tripletmix-pointbert-shapenet":
+        pc_adapter = model_list["pc_adapter"](s).eval()
+        if torch.cuda.is_available():
+            pc_adapter.cuda()
     if torch.cuda.is_available():
         model.cuda()
-    return model
+    return model, pc_adapter