Fixes.

app.py

@@ -58,6 +58,7 @@ def load_data():
 
 
 def render_pc(ncols, col, pc):
+    pc = pc[:2048]
     cols = st.columns(ncols)
     with cols[col]:
         rgb = (pc[:, 3:] * 255).astype(numpy.uint8)

openshape/caption.py

@@ -149,7 +149,7 @@ def pc_caption(pc_encoder: torch.nn.Module, pc, cond_scale):
     prefix = pc_encoder(torch.tensor(pc.T[None], device=ref_dev))
     prefix = prefix.float() * cond_scale
     prefix_embed = model.clip_project(prefix).reshape(1, prefix_length, -1)
-    text, _ = generate2(model, tokenizer, embed=prefix_embed)
+    text = generate2(model, tokenizer, embed=prefix_embed)
     return text
 
 

openshape/classification.py

@@ -7,7 +7,7 @@ from . import lvis
 @torch.no_grad()
 def pred_lvis_sims(pc_encoder: torch.nn.Module, pc):
     ref_dev = next(pc_encoder.parameters()).device
-    enc = pc_encoder(torch.tensor(pc[:, [0, 2, 1]].T[None], device=ref_dev)).cpu()
+    enc = pc_encoder(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)
     return OrderedDict((lvis.categories[i], sim[i]) for i in argsort)