Spaces:
Running
Running
| import torch | |
| from ..utils.base_model import BaseModel | |
| import numpy as np | |
| # borrow from dedode | |
| def dual_softmax_matcher( | |
| desc_A: tuple["B", "C", "N"], | |
| desc_B: tuple["B", "C", "M"], | |
| threshold=0.1, | |
| inv_temperature=20, | |
| normalize=True, | |
| ): | |
| B, C, N = desc_A.shape | |
| if len(desc_A.shape) < 3: | |
| desc_A, desc_B = desc_A[None], desc_B[None] | |
| if normalize: | |
| desc_A = desc_A / desc_A.norm(dim=1, keepdim=True) | |
| desc_B = desc_B / desc_B.norm(dim=1, keepdim=True) | |
| sim = torch.einsum("b c n, b c m -> b n m", desc_A, desc_B) * inv_temperature | |
| P = sim.softmax(dim=-2) * sim.softmax(dim=-1) | |
| mask = torch.nonzero( | |
| (P == P.max(dim=-1, keepdim=True).values) | |
| * (P == P.max(dim=-2, keepdim=True).values) | |
| * (P > threshold) | |
| ) | |
| mask = mask.cpu().numpy() | |
| matches0 = np.ones((B, P.shape[-2]), dtype=int) * (-1) | |
| scores0 = np.zeros((B, P.shape[-2]), dtype=float) | |
| matches0[:, mask[:, 1]] = mask[:, 2] | |
| tmp_P = P.cpu().numpy() | |
| scores0[:, mask[:, 1]] = tmp_P[mask[:, 0], mask[:, 1], mask[:, 2]] | |
| matches0 = torch.from_numpy(matches0).to(P.device) | |
| scores0 = torch.from_numpy(scores0).to(P.device) | |
| return matches0, scores0 | |
| class DualSoftMax(BaseModel): | |
| default_conf = { | |
| "match_threshold": 0.2, | |
| "inv_temperature": 20, | |
| } | |
| # shape: B x DIM x M | |
| required_inputs = ["descriptors0", "descriptors1"] | |
| def _init(self, conf): | |
| pass | |
| def _forward(self, data): | |
| if data["descriptors0"].size(-1) == 0 or data["descriptors1"].size(-1) == 0: | |
| matches0 = torch.full( | |
| data["descriptors0"].shape[:2], -1, device=data["descriptors0"].device | |
| ) | |
| return { | |
| "matches0": matches0, | |
| "matching_scores0": torch.zeros_like(matches0), | |
| } | |
| matches0, scores0 = dual_softmax_matcher( | |
| data["descriptors0"], | |
| data["descriptors1"], | |
| threshold=self.conf["match_threshold"], | |
| inv_temperature=self.conf["inv_temperature"], | |
| ) | |
| return { | |
| "matches0": matches0, # 1 x M | |
| "matching_scores0": scores0, | |
| } | |