|
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)) |
|
|
