Update app.py

app.py

@@ -7,7 +7,7 @@ import openshape
 import transformers
 from PIL import Image
 from huggingface_hub import HfFolder, snapshot_download
-from demo_support import retrieval
+from demo_support import retrieval, utils
 
 @st.cache_resource
 def load_openclip():
@@ -22,6 +22,12 @@ def load_openclip():
             clip_model.cuda()
     return clip_model, clip_prep
 
+@st.cache_resource
+def load_openshape(name, to_cpu=False):
+    pce = openshape.load_pc_encoder(name)
+    if to_cpu:
+        pce = pce.cpu()
+    return pce
 
 def retrieval_filter_expand(key):
     with st.expander("Filters"):
@@ -85,11 +91,19 @@ def demo_retrieval():
     with tab_pc:
         with st.form("rpcform"):
             k = st.slider("Number of items to retrieve", 1, 100, 16, key='rpc')
-            #pc = utils.load_3D_shape('rpcinput')
+            load_data = utils.load_3D_shape('rpcinput')
             sim_th, filter_fn = retrieval_filter_expand('pc')
             if st.form_submit_button("Retrieve with Point Cloud"):
                 prog.progress(0.49, "Computing Embeddings")
-
+                pc = load_data(prog)
+                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()
+                
+                prog.progress(0.7, "Running Retrieval")
+                retrieval_results(retrieval.retrieve(enc, k, sim_th, filter_fn))
+                
+                prog.progress(1.0, "Idle")
 
     with tab_img:
         with st.form("rimgform"):
@@ -141,6 +155,7 @@ tab_cls, tab_pc, tab_img, tab_text, tab_sd, tab_cap = st.tabs([
 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')
 
 try:
     with tab_cls: