import torch import torch.nn.functional as F from collections import OrderedDict 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, 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 if i < len(lvis.categories)) @torch.no_grad() def pred_custom_sims(pc_encoder: torch.nn.Module, pc, cats, feats): ref_dev = next(pc_encoder.parameters()).device enc = pc_encoder(torch.tensor(pc[:, [0, 2, 1, 3, 4, 5]].T[None], device=ref_dev)).cpu() sim = torch.matmul(F.normalize(feats, dim=-1), F.normalize(enc, dim=-1).squeeze()) argsort = torch.argsort(sim, descending=True) return OrderedDict((cats[i], sim[i]) for i in argsort if i < len(cats))