openshape/classification.py

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