Update app.py

app.py

@@ -33,10 +33,11 @@ def load_openshape(name, to_cpu=False):
 
 
 def load_tripletmix(name, to_cpu=False):
-    pce = openshape.load_pc_encoder_mix(name)
+    pce, pca = openshape.load_pc_encoder_mix(name)
     if to_cpu:
         pce = pce.cpu()
-    return pce
+        pca = pca.cpu()
+    return pce, pca
 
 
 
@@ -81,9 +82,10 @@ def classification_lvis(load_data):
     col2 = utils.render_pc(pc)
     prog.progress(0.5, "Running Classification")
     ref_dev = next(model_g14.parameters()).device
-    enc = model_g14(torch.tensor(pc[:, [0, 2, 1, 3, 4, 5]].T[None], device=ref_dev)).cpu()
-
-    sim = torch.matmul(torch.nn.functional.normalize(lvis.feats, dim=-1), torch.nn.functional.normalize(enc, dim=-1).squeeze())
+    enc = model_g14(torch.tensor(pc[:, [0, 2, 1, 3, 4, 5]].T[None], device=ref_dev))
+    if model_name == "pb-sn-M":
+        enc = pc_adapter(enc)
+    sim = torch.matmul(torch.nn.functional.normalize(lvis.feats, dim=-1), torch.nn.functional.normalize(enc.cpu(), dim=-1).squeeze())
     argsort = torch.argsort(sim, descending=True)
     pred = OrderedDict((lvis.categories[i], sim[i]) for i in argsort if i < len(lvis.categories))
     with col2:
@@ -103,8 +105,10 @@ def classification_custom(load_data, cats):
     
     prog.progress(0.5, "Running Classification")
     ref_dev = next(model_g14.parameters()).device
-    enc = model_g14(torch.tensor(pc[:, [0, 2, 1, 3, 4, 5]].T[None], device=ref_dev)).cpu()
-    sim = torch.matmul(torch.nn.functional.normalize(feats, dim=-1), torch.nn.functional.normalize(enc, dim=-1).squeeze())
+    enc = model_g14(torch.tensor(pc[:, [0, 2, 1, 3, 4, 5]].T[None], device=ref_dev))
+    if model_name == "pb-sn-M":
+        enc = pc_adapter(enc)    
+    sim = torch.matmul(torch.nn.functional.normalize(feats, dim=-1), torch.nn.functional.normalize(enc.cpu(), dim=-1).squeeze())
     argsort = torch.argsort(sim, descending=True)
     pred = OrderedDict((cats[i], sim[i]) for i in argsort if i < len(cats))
     with col2:
@@ -197,11 +201,19 @@ try:
     f32 = numpy.float32
     half = torch.float16 if torch.cuda.is_available() else torch.bfloat16
     clip_model, clip_prep = load_openclip()
-    model_g14 = load_openshape('openshape-pointbert-vitg14-rgb')
+    #model_g14 = load_openshape('openshape-pointbert-vitg14-rgb')
     #model_g14 = load_tripletmix('tripletmix-spconv-all')
 
     st.caption("This demo presents three tasks: 3D classification, cross-modal retrieval, and cross-modal generation. Examples are provided for demonstration purposes. You're encouraged to fine-tune task parameters and upload files for customized testing as required.")
     st.sidebar.title("TripletMix Demo Configuration Panel")
+    model_name = st.sidebar.selectbox(
+        'Model Selection',
+        ("pb-sn-M", "pb-sn")
+    )
+    if model_name == "pb-sn-M":
+        model_g14, pc_adapter = load_tripletmix('tripletmix-pointbert-shapenet')
+    elif model_name == "pb-sn":
+        model_g14 = load_openshape('openshape-pointbert-vitg14-rgb')
     task = st.sidebar.selectbox(
         'Task Selection', 
         ("3D Classification", "Cross-modal retrieval", "Cross-modal generation")