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 .Minkowski import MinkResNet34
 
 
 def module(state_dict: dict, name):
@@ -32,10 +33,16 @@ def B32(s):
     return model
 
 
+def Mk34(s):
+    model = MinkResNet34()
+    model.load_state_dict(module(s, 'pc_encoder'))
+    return model
+
 model_list = {
     "openshape-pointbert-vitb32-rgb": B32,
     "openshape-pointbert-vitl14-rgb": L14,
     "openshape-pointbert-vitg14-rgb": G14,
+    "tripletmix-spconv-all": Mk34,
 }
 
 
@@ -45,3 +52,11 @@ def load_pc_encoder(name):
     if torch.cuda.is_available():
         model.cuda()
     return model
+
+
+def load_pc_encoder_mix(name):
+    s = torch.load(hf_hub_download("TripletMix/" + name, "model.pt"), map_location='cpu')
+    model = model_list[name](s).eval()
+    if torch.cuda.is_available():
+        model.cuda()
+    return model