Update openshape/__init__.py

openshape/__init__.py

@@ -2,7 +2,6 @@ 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):
@@ -33,16 +32,19 @@ def B32(s):
     return model
 
 
-def Mk34(s):
-    model = MinkResNet34()
-    model.load_state_dict(module(s, 'pc_encoder'))
+def G14_shapenet_mix(s):
+    model = Projected(
+        PointPatchTransformer(512, 12, 8, 512*3, 256, 384, 0.2, 64, 6),
+        nn.Linear(512, 1280)
+    )
+    model.load_state_dict(module(s['state_dict'], 'module'))
     return model
 
 model_list = {
     "openshape-pointbert-vitb32-rgb": B32,
     "openshape-pointbert-vitl14-rgb": L14,
     "openshape-pointbert-vitg14-rgb": G14,
-    "tripletmix-spconv-all": Mk34,
+    "tripletmix-pointbert-shapenet": G14_shapenet_mix,
 }