Examples.

.gitignore

@@ -0,0 +1 @@
+__pycache__

app.py

@@ -42,7 +42,12 @@ model_b32 = load_openshape('openshape-pointbert-vitb32-rgb').cpu()
 model_l14 = load_openshape('openshape-pointbert-vitl14-rgb')
 model_g14 = load_openshape('openshape-pointbert-vitg14-rgb')
 torch.set_grad_enabled(False)
+for kc, vc in st.session_state.get('state_queue', []):
+    st.session_state[kc] = vc
+st.session_state.state_queue = []
 
+
+import samples_index
 from openshape.demo import misc_utils, classification, caption, sd_pc2img, retrieval
 
 
@@ -59,6 +64,67 @@ tab_cls, tab_img, tab_text, tab_pc, tab_sd, tab_cap = st.tabs([
 ])
 
 
+def sq(kc, vc):
+    st.session_state.state_queue.append((kc, vc))
+
+
+def reset_3d_shape_input(key):
+    objaid_key = key + "_objaid"
+    model_key = key + "_model"
+    npy_key = key + "_npy"
+    swap_key = key + "_swap"
+    sq(objaid_key, "")
+    sq(model_key, None)
+    sq(npy_key, None)
+    sq(swap_key, "Y is up (for most Objaverse shapes)")
+
+
+def auto_submit(key):
+    if st.session_state.get(key):
+        st.session_state[key] = False
+        return True
+    return False
+
+
+def queue_auto_submit(key):
+    st.session_state[key] = True
+    st.experimental_rerun()
+
+
+img_example_counter = 0
+
+
+def image_examples(samples, ncols, return_key=None):
+    global img_example_counter
+    trigger = False
+    with st.expander("Examples", True):
+        for i in range(len(samples) // ncols):
+            cols = st.columns(ncols)
+            for j in range(ncols):
+                idx = i * ncols + j
+                if idx >= len(samples):
+                    continue
+                entry = samples[idx]
+                with cols[j]:
+                    st.image(entry['dispi'])
+                    img_example_counter += 1
+                    with st.columns(5)[2]:
+                        this_trigger = st.button('\+', key='imgexuse%d' % img_example_counter)
+                    trigger = trigger or this_trigger
+                    if this_trigger:
+                        if return_key is None:
+                            for k, v in entry.items():
+                                if not k.startswith('disp'):
+                                    sq(k, v)
+                        else:
+                            trigger = entry[return_key]
+    return trigger
+
+
+def text_examples(samples):
+    return st.selectbox("Or pick an example", samples)
+
+
 def demo_classification():
     load_data = misc_utils.input_3d_shape('cls')
     cats = st.text_input("Custom Categories (64 max, separated with comma)")
@@ -68,7 +134,7 @@ def demo_classification():
         return
     lvis_run = st.button("Run Classification on LVIS Categories")
     custom_run = st.button("Run Classification on Custom Categories")
-    if lvis_run:
+    if lvis_run or auto_submit("clsauto"):
         pc = load_data(prog)
         col2 = misc_utils.render_pc(pc)
         prog.progress(0.5, "Running Classification")
@@ -92,31 +158,35 @@ def demo_classification():
                 st.text(cat)
                 st.caption("Similarity %.4f" % sim)
         prog.progress(1.0, "Idle")
+    if image_examples(samples_index.classification, 3):
+        queue_auto_submit("clsauto")
 
 
 def demo_captioning():
     with st.form("capform"):
         load_data = misc_utils.input_3d_shape('cap')
-        cond_scale = st.slider('Conditioning Scale', 0.0, 4.0, 2.0)
-        if st.form_submit_button("Generate a Caption"):
+        cond_scale = st.slider('Conditioning Scale', 0.0, 4.0, 2.0, 0.1, key='capcondscl')
+        if st.form_submit_button("Generate a Caption") or auto_submit("capauto"):
             pc = load_data(prog)
             col2 = misc_utils.render_pc(pc)
             prog.progress(0.5, "Running Generation")
             cap = caption.pc_caption(model_b32, pc, cond_scale)
             st.text(cap)
             prog.progress(1.0, "Idle")
+    if image_examples(samples_index.cap, 3):
+        queue_auto_submit("capauto")
 
 
 def demo_pc2img():
     with st.form("sdform"):
         load_data = misc_utils.input_3d_shape('sd')
-        prompt = st.text_input("Prompt (Optional)")
+        prompt = st.text_input("Prompt (Optional)", key='sdtprompt')
         noise_scale = st.slider('Variation Level', 0, 5, 1)
         cfg_scale = st.slider('Guidance Scale', 0.0, 30.0, 10.0)
         steps = st.slider('Diffusion Steps', 8, 50, 25)
         width = 640  # st.slider('Width', 480, 640, step=32)
         height = 640  # st.slider('Height', 480, 640, step=32)
-        if st.form_submit_button("Generate"):
+        if st.form_submit_button("Generate") or auto_submit("sdauto"):
             pc = load_data(prog)
             col2 = misc_utils.render_pc(pc)
             prog.progress(0.49, "Running Generation")
@@ -131,6 +201,8 @@ def demo_pc2img():
             with col2:
                 st.image(img)
             prog.progress(1.0, "Idle")
+    if image_examples(samples_index.sd, 3):
+        queue_auto_submit("sdauto")
 
 
 def retrieval_results(results):
@@ -155,35 +227,44 @@ def demo_retrieval():
         with st.form("rtextform"):
             k = st.slider("# Shapes to Retrieve", 1, 100, 16, key='rtext')
             text = st.text_input("Input Text")
+            picked_sample = text_examples(samples_index.retrieval_texts)
             if st.form_submit_button("Run with Text"):
                 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)
+                tn = clip_prep(
+                    text=[text or picked_sample], return_tensors='pt', truncation=True, max_length=76
+                ).to(device)
                 enc = clip_model.get_text_features(**tn).float().cpu()
                 prog.progress(0.7, "Running Retrieval")
                 retrieval_results(retrieval.retrieve(enc, k))
                 prog.progress(1.0, "Idle")
 
     with tab_img:
+        submit = False
         with st.form("rimgform"):
             k = st.slider("# Shapes to Retrieve", 1, 100, 16, key='rimage')
-            pic = st.file_uploader("Upload an Image")
+            pic = st.file_uploader("Upload an Image", key='rimageinput')
             if st.form_submit_button("Run with Image"):
-                img = Image.open(pic)
-                st.image(img)
-                prog.progress(0.49, "Computing Embeddings")
-                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))
-                prog.progress(1.0, "Idle")
+                submit = True
+        sample_got = image_examples(samples_index.iret, 4, 'rimageinput')
+        if sample_got:
+            pic = sample_got
+        if sample_got or submit:
+            img = Image.open(pic)
+            st.image(img)
+            prog.progress(0.49, "Computing Embeddings")
+            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))
+            prog.progress(1.0, "Idle")
 
     with tab_pc:
         with st.form("rpcform"):
             k = st.slider("# Shapes to Retrieve", 1, 100, 16, key='rpc')
             load_data = misc_utils.input_3d_shape('retpc')
-            if st.form_submit_button("Run with Shape"):
+            if st.form_submit_button("Run with Shape") or auto_submit('rpcauto'):
                 pc = load_data(prog)
                 col2 = misc_utils.render_pc(pc)
                 prog.progress(0.49, "Computing Embeddings")
@@ -192,6 +273,8 @@ def demo_retrieval():
                 prog.progress(0.7, "Running Retrieval")
                 retrieval_results(retrieval.retrieve(enc, k))
                 prog.progress(1.0, "Idle")
+        if image_examples(samples_index.pret, 3):
+            queue_auto_submit("rpcauto")
 
 
 try:

samples/retrieval-text.txt

@@ -1,16 +0,0 @@
-shark
-swordfish
-dolphin
-goldfish
-high heels
-boots
-slippers
-sneakers
-tiki mug
-viking mug
-animal-shaped mug
-travel mug
-white conical mug
-green cubic mug
-blue spherical mug
-orange cylinder mug

samples/sd-text.txt

@@ -1,3 +0,0 @@
-b8db8dc5caad4fa5842a9ed6dbd2e9d6,falcon
-ff2875fb1a5b4771805a5fd35c8fe7bb,in the woods
-tpvzmLUXAURQ7ZxccJIBZvcIDlr,above the fields

samples_index.py

@@ -0,0 +1,61 @@
+import os
+
+
+cap_base = 'samples/caption'
+cap = [
+    dict(cap_objaid=os.path.splitext(x)[0], dispi=os.path.join(cap_base, x))
+    for x in sorted(os.listdir(cap_base))
+]
+
+cls_base = 'samples/classification'
+classification = [
+    dict(cls_objaid=os.path.splitext(x)[0], dispi=os.path.join(cls_base, x))
+    for x in sorted(os.listdir(cls_base))
+]
+
+sd_base = 'samples/sd'
+sd_texts = {
+    'b8db8dc5caad4fa5842a9ed6dbd2e9d6': 'falcon',
+    'ff2875fb1a5b4771805a5fd35c8fe7bb': 'in the woods',
+    'tpvzmLUXAURQ7ZxccJIBZvcIDlr': 'above the fields'
+}
+sd = [
+    dict(
+        sd_objaid=os.path.splitext(x)[0],
+        dispi=os.path.join(sd_base, x),
+        sdtprompt=sd_texts.get(os.path.splitext(x)[0], '')
+    )
+    for x in sorted(os.listdir(sd_base))
+]
+
+retrieval_texts = """
+shark
+swordfish
+dolphin
+goldfish
+high heels
+boots
+slippers
+sneakers
+tiki mug
+viking mug
+animal-shaped mug
+travel mug
+white conical mug
+green cubic mug
+blue spherical mug
+orange cylinder mug
+""".splitlines()
+retrieval_texts = [x.strip() for x in retrieval_texts if x.strip()]
+
+pret_base = 'samples/retrieval-pc'
+pret = [
+    dict(retpc_objaid=os.path.splitext(x)[0], dispi=os.path.join(pret_base, x))
+    for x in sorted(os.listdir(pret_base))
+]
+
+iret_base = 'samples/retrieval-img'
+iret = [
+    dict(rimageinput=os.path.join(iret_base, x), dispi=os.path.join(iret_base, x))
+    for x in sorted(os.listdir(iret_base))
+]