tripletmix-demo / app.py
winfred2027's picture
Update app.py
64fa430 verified
raw
history blame
2.61 kB
import streamlit as st
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_img, tab_text, tab_pc, tab_sd, tab_cap = st.tabs([
"Classification",
"Retrieval w/ Image",
"Retrieval w/ Text",
"Retrieval w/ 3D",
"Image Generation",
"Captioning",
])
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")
"""
if lvis_run or auto_submit("clsauto"):
pc = load_data(prog)
col2 = misc_utils.render_pc(pc)
prog.progress(0.5, "Running Classification")
pred = classification.pred_lvis_sims(model_g14, pc)
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")
if custom_run:
pc = load_data(prog)
col2 = misc_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")
pred = classification.pred_custom_sims(model_g14, pc, cats, feats)
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")
"""
"""
if image_examples(samples_index.classification, 3, example_text="Examples (Choose one of the following 3D shapes)"):
queue_auto_submit("clsauto")
"""
try:
with tab_cls:
demo_classification()
"""
with tab_cap:
demo_captioning()
with tab_sd:
demo_pc2img()
demo_retrieval()
"""
except Exception:
import traceback
st.error(traceback.format_exc().replace("\n", " \n"))