argus.py: correspond() returns dense {matches, scores, grid} dict; lazy-load class names

argus.py

@@ -1640,12 +1640,46 @@ class Argus(PreTrainedModel):
                     if torch.isnan(buf).any() or torch.isinf(buf).any():
                         buf.data.zero_()
 
+    def _load_imagenet_classes(self):
+        if getattr(self, "_imagenet_classes_loaded", False):
+            return
+        self._imagenet_classes_loaded = True
+        import json
+        import os as _os
+        candidates = []
+        here = _os.path.dirname(_os.path.abspath(__file__))
+        candidates.append(_os.path.join(here, "imagenet_classes.json"))
+        name_or_path = getattr(self.config, "_name_or_path", None)
+        if name_or_path and _os.path.isdir(name_or_path):
+            candidates.append(_os.path.join(name_or_path, "imagenet_classes.json"))
+        for path in candidates:
+            if _os.path.isfile(path):
+                with open(path) as f:
+                    data = json.load(f)
+                self.config.class_ids = data.get("class_ids", [])
+                self.config.class_names = data.get("class_names", [])
+                return
+        if name_or_path and not _os.path.isdir(name_or_path):
+            try:
+                from huggingface_hub import hf_hub_download
+                path = hf_hub_download(name_or_path, "imagenet_classes.json")
+                with open(path) as f:
+                    data = json.load(f)
+                self.config.class_ids = data.get("class_ids", [])
+                self.config.class_names = data.get("class_names", [])
+            except Exception:
+                pass
+
     @property
     def class_ids(self):
+        if not self.config.class_ids:
+            self._load_imagenet_classes()
         return self.config.class_ids
 
     @property
     def class_names(self):
+        if not self.config.class_names:
+            self._load_imagenet_classes()
         return self.config.class_names
 
     def quantize_int8(self):
@@ -1782,34 +1816,32 @@ class Argus(PreTrainedModel):
         self,
         src_image: Image.Image,
         tgt_image: Image.Image,
-        src_keypoints: list,
         resolution: int = 512,
     ):
-        sw, sh = src_image.size
-        tw, th = tgt_image.size
+        """Dense patch correspondence between two images.
+
+        Returns a dict with keys `matches` (numpy array of length grid*grid mapping
+        each source patch to its argmax target patch), `scores` (cosine similarity
+        at the match), and `grid` (the patch-grid side length).
+        """
         transform = make_eupe_transform(resolution)
         src_t = transform(src_image).unsqueeze(0).to(self.device)
         tgt_t = transform(tgt_image).unsqueeze(0).to(self.device)
 
-        _, src_feats = self._extract(src_t)
-        _, tgt_feats = self._extract(tgt_t)
-
-        src_feats = F.interpolate(src_feats, size=(resolution, resolution), mode="bilinear", align_corners=False)
-        tgt_feats = F.interpolate(tgt_feats, size=(resolution, resolution), mode="bilinear", align_corners=False)
-
-        src_feats = F.normalize(src_feats[0].permute(1, 2, 0), dim=-1)
-        tgt_feats = F.normalize(tgt_feats[0].permute(1, 2, 0), dim=-1)
-
-        preds = []
-        for kp in src_keypoints:
-            sx = min(max(int(kp[0] / sw * resolution), 0), resolution - 1)
-            sy = min(max(int(kp[1] / sh * resolution), 0), resolution - 1)
-            src_vec = src_feats[sy, sx]
-            sim_map = torch.einsum("d,hwd->hw", src_vec, tgt_feats)
-            flat = sim_map.argmax().item()
-            py, px = flat // resolution, flat % resolution
-            preds.append([px / resolution * tw, py / resolution * th])
-        return preds
+        with torch.autocast(self.device.type, dtype=torch.bfloat16, enabled=self.device.type == "cuda"):
+            oa = self.backbone.forward_features(src_t)
+            ob = self.backbone.forward_features(tgt_t)
+        pa = F.normalize(oa['x_norm_patchtokens'].float().squeeze(0), dim=-1)
+        pb = F.normalize(ob['x_norm_patchtokens'].float().squeeze(0), dim=-1)
+        sim = pa @ pb.t()
+        m = sim.argmax(dim=-1)
+        s = sim.max(dim=-1).values
+        grid = int(np.sqrt(pa.shape[0]))
+        return {
+            "matches": m.cpu().numpy(),
+            "scores": s.cpu().numpy(),
+            "grid": grid,
+        }
 
     @torch.inference_mode()
     def detect(