Update app.py

app.py

@@ -33,11 +33,10 @@ def load_openshape(name, to_cpu=False):
 
 
 def load_tripletmix(name, to_cpu=False):
-    pce, pca = openshape.load_pc_encoder_mix(name)
+    pce = openshape.load_pc_encoder_mix(name)
     if to_cpu:
         pce = pce.cpu()
-        pca = pca.cpu()
-    return pce, pca
+    return pce
 
 
 
@@ -81,10 +80,8 @@ def classification_lvis(load_data):
     pc = load_data(prog)
     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))
-    if model_name == "pb-sn-M":
-        enc = pc_adapter(enc)
+    ref_dev = next(model_classification.parameters()).device
+    enc = model_classification(torch.tensor(pc[:, [0, 2, 1, 3, 4, 5]].T[None], device=ref_dev))
     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))
@@ -104,10 +101,8 @@ def classification_custom(load_data, cats):
     feats = clip_model.get_text_features(**tn).float().cpu()
     
     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))
-    if model_name == "pb-sn-M":
-        enc = pc_adapter(enc)  
+    ref_dev = next(model_classification.parameters()).device
+    enc = model_classification(torch.tensor(pc[:, [0, 2, 1, 3, 4, 5]].T[None], device=ref_dev)) 
     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))
@@ -207,14 +202,7 @@ try:
 
     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-shapenet')
+
     task = st.sidebar.selectbox(
         'Task Selection', 
         ("3D Classification", "Cross-modal retrieval", "Cross-modal generation")
@@ -225,6 +213,14 @@ try:
             'Choose the source of categories', 
             ("LVIS Categories", "Custom Categories")
         )
+        model_name = st.sidebar.selectbox(
+            'Model Selection',
+            ("pb-Mix", "pb")
+        )
+        if model_name == "pb-Mix":
+            model_classification = load_tripletmix('tripletmix-pointbert-all-modelnet40')
+        elif model_name == "pb":
+            model_classification = load_openshape('openshape-pointbert-vitg14-rgb')
         load_data = utils.input_3d_shape('rpcinput')
         if cls_mode == "LVIS Categories":
             st.title("Classification with LVIS Categories")