Update app.py

app.py

@@ -66,17 +66,6 @@ def retrieval_results(results):
 
 
 
-def demo_classification():
-    with st.form("clsform"):
-        #load_data = misc_utils.input_3d_shape('cls')
-        cats = st.text_input("Custom Categories (64 max, separated with comma)")
-        cats = [a.strip() for a in cats.split(',')]
-        if len(cats) > 64:
-            st.error('Maximum 64 custom categories supported in the demo')
-            return
-        lvis_run = st.form_submit_button("Run Classification on LVIS Categories")
-        custom_run = st.form_submit_button("Run Classification on Custom Categories")
-
 def demo_captioning():
     with st.form("capform"):
         cond_scale = st.slider('Conditioning Scale', 0.0, 4.0, 2.0, 0.1, key='capcondscl')
@@ -146,6 +135,43 @@ def demo_retrieval():
                 prog.progress(1.0, "Idle")
 
 
+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)).cpu()
+
+    sim = torch.matmul(torch.nn.functional.normalize(lvis.feats, dim=-1), torch.nn.functional.normalize(enc, 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:
+        for i, (cat, sim) in zip(range(5), pred.items()):
+            st.text(cat)
+            st.caption("Similarity %.4f" % sim)
+    prog.progress(1.0, "Idle")
+
+def classification_custom(load_data, cats):
+    pc = load_data(prog)
+    col2 = utils.render_pc(pc)
+    prog.progress(0.5, "Computing Category Embeddings")
+    device = clip_model.device
+    tn = clip_prep(text=cats, return_tensors='pt', truncation=True, max_length=76, padding=True).to(device)
+    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)).cpu()
+    sim = torch.matmul(torch.nn.functional.normalize(feats, dim=-1), torch.nn.functional.normalize(enc, 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:
+        for i, (cat, sim) in zip(range(5), pred.items()):
+            st.text(cat)
+            st.caption("Similarity %.4f" % sim)
+    prog.progress(1.0, "Idle")
+
+
 def retrieval_pc(load_data, k, sim_th, filter_fn):
     pc = load_data(prog)
     prog.progress(0.49, "Computing Embeddings")
@@ -208,21 +234,21 @@ try:
             'Choose the source of categories', 
             ("LVIS Categories", "Custom Categories")
         )
-        pc = st.sidebar.text_input("Input pc", key='rtextinput')
+        load_data = utils.input_3d_shape('rpcinput')
         if cls_mode == "LVIS Categories":
+            st.title("Classification with LVIS Categories")
+            prog = st.progress(0.0, "Idle")
             if st.sidebar.button("submit"):
-                st.title("Classification with LVIS Categories")
-                prog = st.progress(0.0, "Idle")
-
+                classification_lvis(load_data)
         elif cls_mode == "Custom Categories":
+            st.title("Classification with Custom Categories")
+            prog = st.progress(0.0, "Idle")
             cats = st.sidebar.text_input("Custom Categories (64 max, separated with comma)")
             cats = [a.strip() for a in cats.split(',')]
             if len(cats) > 64:
                 st.error('Maximum 64 custom categories supported in the demo')
             if st.sidebar.button("submit"):
-                st.title("Classification with Custom Categories")
-                prog = st.progress(0.0, "Idle")
-
+                classification_custom(load_data, cats)
     elif task == "Cross-modal retrieval":
         input_mode = st.sidebar.selectbox(
             'Choose an input modality', 
@@ -231,22 +257,22 @@ try:
         k = st.sidebar.slider("Number of items to retrieve", 1, 100, 16, key='rnum')
         sim_th, filter_fn = retrieval_filter_expand()
         if input_mode == "Point Cloud":
+            st.title("Retrieval with Point Cloud")
+            prog = st.progress(0.0, "Idle")
             load_data = utils.input_3d_shape('rpcinput')
             if st.sidebar.button("submit"):
-                st.title("Retrieval with Point Cloud")
-                prog = st.progress(0.0, "Idle")
                 retrieval_pc(load_data, k, sim_th, filter_fn)
         elif input_mode == "Image":
+            st.title("Retrieval with Image")
+            prog = st.progress(0.0, "Idle") 
             pic = st.sidebar.file_uploader("Upload an Image", key='rimageinput')
             if st.sidebar.button("submit"):
-                st.title("Retrieval with Image")
-                prog = st.progress(0.0, "Idle") 
                 retrieval_img(pic, k, sim_th, filter_fn)
         elif input_mode == "Text":
+            st.title("Retrieval with Text")
+            prog = st.progress(0.0, "Idle")
             text = st.sidebar.text_input("Input Text", key='rtextinput')
             if st.sidebar.button("submit"):
-                st.title("Retrieval with Text")
-                prog = st.progress(0.0, "Idle")
                 retrieval_text(text, k, sim_th, filter_fn)
     elif task == "Cross-modal generation":
         generation_mode = st.sidebar.selectbox(
@@ -255,14 +281,15 @@ try:
         )
         pc = st.sidebar.text_input("Input pc", key='rtextinput')
         if generation_mode == "PointCloud-to-Image":
+            st.title("Image Generation")
+            prog = st.progress(0.0, "Idle")
             if st.sidebar.button("submit"):
-                st.title("Image Generation")
-                prog = st.progress(0.0, "Idle")
-
+                pc = st.text_input("Input pc", key='rtextinput')
         elif generation_mode == "PointCloud-to-Text":
+            st.title("Text Generation")
+            prog = st.progress(0.0, "Idle")
             if st.sidebar.button("submit"):
-                st.title("Text Generation")
-                prog = st.progress(0.0, "Idle")
+                pc = st.text_input("Input pc", key='rtextinput')
 
 except Exception:
     import traceback