Update app.py

app.py

@@ -30,25 +30,22 @@ def load_openshape(name, to_cpu=False):
         pce = pce.cpu()
     return pce
 
-def retrieval_filter_expand(key):
-    with st.expander("Filters"):
-        sim_th = st.slider("Similarity Threshold", 0.05, 0.5, 0.1, key=key + 'rtsimth')
-        tag = st.text_input("Has Tag", "", key=key + 'rthastag')
-        col1, col2 = st.columns(2)
-        face_min = int(col1.text_input("Face Count Min", "0", key=key + 'rtfcmin'))
-        face_max = int(col2.text_input("Face Count Max", "34985808", key=key + 'rtfcmax'))
-        col1, col2 = st.columns(2)
-        anim_min = int(col1.text_input("Animation Count Min", "0", key=key + 'rtacmin'))
-        anim_max = int(col2.text_input("Animation Count Max", "563", key=key + 'rtacmax'))
-        tag_n = not bool(tag.strip())
-        anim_n = not (anim_min > 0 or anim_max < 563)
-        face_n = not (face_min > 0 or face_max < 34985808)
-        filter_fn = lambda x: (
-            (anim_n or anim_min <= x['anims'] <= anim_max)
-            and (face_n or face_min <= x['faces'] <= face_max)
-            and (tag_n or tag in x['tags'])
-        )
-        return sim_th, filter_fn
+def retrieval_filter_expand():
+    sim_th = st.sidebar.slider("Similarity Threshold", 0.05, 0.5, 0.1, key='rsimth')
+    tag = ""
+    face_min = 0
+    face_max = 34985808
+    anim_min = 0
+    anim_max = 563
+    tag_n = not bool(tag.strip())
+    anim_n = not (anim_min > 0 or anim_max < 563)
+    face_n = not (face_min > 0 or face_max < 34985808)
+    filter_fn = lambda x: (
+        (anim_n or anim_min <= x['anims'] <= anim_max)
+        and (face_n or face_min <= x['faces'] <= face_max)
+        and (tag_n or tag in x['tags'])
+    )
+    return sim_th, filter_fn
 
 def retrieval_results(results):
     st.caption("Click the link to view the 3D shape")
@@ -148,32 +145,125 @@ def demo_retrieval():
                 
                 prog.progress(1.0, "Idle")
 
-st.title("TripletMix Demo")
-st.caption("For faster inference without waiting in queue, you may clone the space and run it yourself.")
-prog = st.progress(0.0, "Idle")
-tab_cls, tab_pc, tab_img, tab_text, tab_sd, tab_cap = st.tabs([
-    "Classification",
-    "Retrieval w/ 3D",
-    "Retrieval w/ Image",
-    "Retrieval w/ Text",
-    "Image Generation",
-    "Captioning",
-])
 
+def retrieval_pc(load_data, k, sim_th, filter_fn):
+    pc = load_data(prog)
+    prog.progress(0.49, "Computing Embeddings")
+    col2 = utils.render_pc(pc)
+    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(0.7, "Running Retrieval")
+    retrieval_results(retrieval.retrieve(enc, k, sim_th, filter_fn))
+    
+    prog.progress(1.0, "Idle")
+
+def retrieval_img(pic, k, sim_th, filter_fn):
+    img = Image.open(pic)
+    prog.progress(0.49, "Computing Embeddings")
+    st.image(img)
+    device = clip_model.device
+    tn = clip_prep(images=[img], return_tensors="pt").to(device)
+    enc = clip_model.get_image_features(pixel_values=tn['pixel_values'].type(half)).float().cpu()
+    
+    prog.progress(0.7, "Running Retrieval")
+    retrieval_results(retrieval.retrieve(enc, k, sim_th, filter_fn))
+    
+    prog.progress(1.0, "Idle")
+
+def retrieval_text(text, k, sim_th, filter_fn):
+    prog.progress(0.49, "Computing Embeddings")
+    device = clip_model.device
+    tn = clip_prep(text=[text], return_tensors='pt', truncation=True, max_length=76).to(device)
+    enc = clip_model.get_text_features(**tn).float().cpu()
 
-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')
+    prog.progress(0.7, "Running Retrieval")
+    retrieval_results(retrieval.retrieve(enc, k, sim_th, filter_fn))
+    
+    prog.progress(1.0, "Idle")
 
 try:
-    with tab_cls:
-        demo_classification()
-    with tab_cap:
-        demo_captioning()
-    with tab_sd:
-        demo_pc2img()
-    demo_retrieval()
+    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')
+
+    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")
+    task = st.sidebar.selectbox(
+        'Task Selection', 
+        ("3D Classification", "Cross-modal retrieval", "Cross-modal generation")
+    )
+
+    if task == "3D Classification":
+        cls_mode = st.sidebar.selectbox(
+            'Choose the source of categories', 
+            ("LVIS Categories", "Custom Categories")
+        )
+        pc = st.sidebar.text_input("Input pc", key='rtextinput')
+        if cls_mode == "LVIS Categories":
+            if st.sidebar.button("submit"):
+                st.title("Classification with LVIS Categories")
+                prog = st.progress(0.0, "Idle")
+
+        elif cls_mode == "Custom Categories":
+            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")
+
+    elif task == "Cross-modal retrieval":
+        input_mode = st.sidebar.selectbox(
+            'Choose an input modality', 
+            ("Point Cloud", "Image", "Text")
+        )
+        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":
+            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":
+            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":
+            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(
+            'Choose the mode of generation', 
+            ("PointCloud-to-Image", "PointCloud-to-Text")
+        )
+        pc = st.sidebar.text_input("Input pc", key='rtextinput')
+        if generation_mode == "PointCloud-to-Image":
+            if st.sidebar.button("submit"):
+                st.title("Image Generation")
+                prog = st.progress(0.0, "Idle")
+
+        elif generation_mode == "PointCloud-to-Text":
+            if st.sidebar.button("submit"):
+                st.title("Text Generation")
+                prog = st.progress(0.0, "Idle")
+
 except Exception:
     import traceback
     st.error(traceback.format_exc().replace("\n", "  \n"))