Update app.py

app.py

@@ -7,7 +7,8 @@ import openshape
 import transformers
 from PIL import Image
 from huggingface_hub import HfFolder, snapshot_download
-from demo_support import retrieval, utils
+from demo_support import retrieval, utils, lvis
+from collections import OrderedDict
 
 @st.cache_resource
 def load_openclip():
@@ -99,6 +100,14 @@ def demo_retrieval():
                 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(F.normalize(lvis.feats, dim=-1), F.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))