add: rdd sparse and dense match

.gitignore

@@ -7,7 +7,7 @@ cmake-build-debug/
 *.pyc
 flagged
 .ipynb_checkpoints
-__pycache__
+**__pycache__**
 Untitled*
 experiments
 third_party/REKD

README.md

@@ -44,8 +44,9 @@ The tool currently supports various popular image matching algorithms, namely:
 
 | Algorithm        | Supported | Conference/Journal | Year | GitHub Link |
 |------------------|-----------|--------------------|------|-------------|
-| DaD            | ✅ | ARXIV   | 2025 | [Link](https://github.com/Parskatt/dad) |
 | LiftFeat       | ✅ | ICRA    | 2025 | [Link](https://github.com/lyp-deeplearning/LiftFeat) |
+| RDD            | ✅ | CVPR    | 2025 | [Link](https://github.com/xtcpete/rdd)  |
+| DaD            | ✅ | ARXIV   | 2025 | [Link](https://github.com/Parskatt/dad) |
 | MINIMA         | ✅ | ARXIV   | 2024 | [Link](https://github.com/LSXI7/MINIMA) |
 | XoFTR          | ✅ | CVPR    | 2024 | [Link](https://github.com/OnderT/XoFTR) |
 | EfficientLoFTR | ✅ | CVPR    | 2024 | [Link](https://github.com/zju3dv/EfficientLoFTR) |

config/config.yaml

@@ -267,6 +267,27 @@ matcher_zoo:
       paper: https://arxiv.org/abs/2505.0342
       project: null
       display: true
+  rdd(sparse):
+    matcher: NN-mutual
+    feature: rdd
+    dense: false
+    info:
+      name: RDD(sparse) #dispaly name
+      source: "CVPR 2025"
+      github: hhttps://github.com/xtcpete/rdd
+      paper: https://arxiv.org/abs/2505.08013
+      project: https://xtcpete.github.io/rdd
+      display: true
+  rdd(dense):
+    matcher: rdd_dense
+    dense: true
+    info:
+      name: RDD(dense) #dispaly name
+      source: "CVPR 2025"
+      github: hhttps://github.com/xtcpete/rdd
+      paper: https://arxiv.org/abs/2505.08013
+      project: https://xtcpete.github.io/rdd
+      display: true
   dedode:
     matcher: Dual-Softmax
     feature: dedode

imcui/hloc/extract_features.py

@@ -225,6 +225,17 @@ confs = {
             "resize_max": 1600,
         },
     },
+    "rdd": {
+        "output": "feats-rdd-n5000-r1600",
+        "model": {
+            "name": "rdd",
+            "max_keypoints": 5000,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "resize_max": 1600,
+        },
+    },
     "aliked-n16-rot": {
         "output": "feats-aliked-n16-rot",
         "model": {

imcui/hloc/extractors/liftfeat.py

@@ -1,13 +1,10 @@
-import logging
 import sys
 from pathlib import Path
-import torch
-import random
 from ..utils.base_model import BaseModel
 from .. import logger, MODEL_REPO_ID
 
-fire_path = Path(__file__).parent / "../../third_party/LiftFeat"
-sys.path.append(str(fire_path))
+liftfeat_path = Path(__file__).parent / "../../third_party/LiftFeat"
+sys.path.append(str(liftfeat_path))
 
 from models.liftfeat_wrapper import LiftFeat
 
@@ -25,9 +22,7 @@ class Liftfeat(BaseModel):
         logger.info("Loading LiftFeat model...")
         model_path = self._download_model(
             repo_id=MODEL_REPO_ID,
-            filename="{}/{}".format(
-                Path(__file__).stem, self.conf["model_name"]
-            ),
+            filename="{}/{}".format(Path(__file__).stem, self.conf["model_name"]),
         )
         self.net = LiftFeat(
             weight=model_path,

imcui/hloc/extractors/rdd.py

@@ -0,0 +1,56 @@
+import sys
+import yaml
+from pathlib import Path
+from ..utils.base_model import BaseModel
+from .. import logger, MODEL_REPO_ID, DEVICE
+
+rdd_path = Path(__file__).parent / "../../third_party/rdd"
+sys.path.append(str(rdd_path))
+
+from RDD.RDD import build as build_rdd
+
+class Rdd(BaseModel):
+    default_conf = {
+        "keypoint_threshold": 0.1,
+        "max_keypoints": 4096,
+        "model_name": "RDD-v2.pth",
+    }
+
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        logger.info("Loading RDD model...")
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                Path(__file__).stem, self.conf["model_name"]
+            ),
+        )
+        config_path = rdd_path / "configs/default.yaml"
+        with open(config_path, "r") as file:
+            config = yaml.safe_load(file)
+        config["top_k"] = conf["max_keypoints"]
+        config["detection_threshold"] = conf["keypoint_threshold"]
+        config["device"] = DEVICE
+        self.net = build_rdd(config=config, weights=model_path)
+        self.net.eval()
+        logger.info("Loading RDD model done!")
+
+    def _forward(self, data):
+        image = data["image"]
+        pred = self.net.extract(image)[0]
+        keypoints = pred["keypoints"]
+        descriptors = pred["descriptors"]
+        scores = pred["scores"]
+        if self.conf["max_keypoints"] < len(keypoints):
+            idxs = scores.argsort()[-self.conf["max_keypoints"] or None :]
+            keypoints = keypoints[idxs, :2]
+            descriptors = descriptors[idxs]
+            scores = scores[idxs]
+
+        pred = {
+            "keypoints": keypoints[None],
+            "descriptors": descriptors[None].permute(0, 2, 1),
+            "scores": scores[None],
+        }
+        return pred

imcui/hloc/match_dense.py

@@ -337,6 +337,23 @@ confs = {
             "dfactor": 8,
         },
     },
+    "rdd_dense": {
+        "output": "matches-rdd_dense",
+        "model": {
+            "name": "rdd_dense",
+            "model_name": "RDD-v2.pth",
+            "max_keypoints": 2000,
+            "match_threshold": 0.2,
+        },
+        "preprocessing": {
+            "grayscale": False,
+            "force_resize": True,
+            "resize_max": 1024,
+            "width": 320,
+            "height": 240,
+            "dfactor": 8,
+        },
+    },
     "minima_roma": {
         "output": "matches-minima_roma",
         "model": {

imcui/hloc/matchers/rdd_dense.py

@@ -0,0 +1,52 @@
+import sys
+import yaml
+import torch
+from pathlib import Path
+from ..utils.base_model import BaseModel
+from .. import logger, MODEL_REPO_ID, DEVICE
+
+rdd_path = Path(__file__).parent / "../../third_party/rdd"
+sys.path.append(str(rdd_path))
+
+from RDD.RDD import build as build_rdd
+from RDD.RDD_helper import RDD_helper
+
+class RddDense(BaseModel):
+    default_conf = {
+        "keypoint_threshold": 0.1,
+        "max_keypoints": 4096,
+        "model_name": "RDD-v2.pth",
+        "match_threshold": 0.1,
+    }
+
+    required_inputs = ["image0", "image1"]
+
+    def _init(self, conf):
+        logger.info("Loading RDD model...")
+        model_path = self._download_model(
+            repo_id=MODEL_REPO_ID,
+            filename="{}/{}".format(
+                "rdd", self.conf["model_name"]
+            ),
+        )
+        config_path = rdd_path / "configs/default.yaml"
+        with open(config_path, "r") as file:
+            config = yaml.safe_load(file)
+        config["top_k"] = conf["max_keypoints"]
+        config["detection_threshold"] = conf["keypoint_threshold"]
+        config["device"] = DEVICE
+        rdd_net = build_rdd(config=config, weights=model_path)
+        rdd_net.eval()
+        self.net = RDD_helper(rdd_net)
+        logger.info("Loading RDD model done!")
+
+    def _forward(self, data):
+        img0 = data["image0"]
+        img1 = data["image1"]
+        mkpts_0, mkpts_1, conf = self.net.match_dense(img0, img1, thr=self.conf["match_threshold"])
+        pred = {
+            "keypoints0": torch.from_numpy(mkpts_0),
+            "keypoints1": torch.from_numpy(mkpts_1),
+            "mconf": torch.from_numpy(conf),
+        }
+        return pred

imcui/third_party/rdd/.gitignore

@@ -0,0 +1,8 @@
+.venv
+/build/
+**.egg-info
+**.pyc
+/.idea/
+**/__pycache__/
+weights/
+outputs

imcui/third_party/rdd/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

imcui/third_party/rdd/RDD/RDD.py

@@ -0,0 +1,260 @@
+# Description: RDD model
+import torch
+import torch.nn.functional as F
+from torch import nn
+import numpy as np
+from .utils import NestedTensor, nested_tensor_from_tensor_list, to_pixel_coords, read_config
+from .models.detector import build_detector
+from .models.descriptor import build_descriptor
+from .models.soft_detect import SoftDetect
+from .models.interpolator import InterpolateSparse2d
+
+class RDD(nn.Module):
+
+    def __init__(self, detector, descriptor, detection_threshold=0.5, top_k=4096, train_detector=False, device='cuda'):
+        super().__init__()
+        self.detector = detector
+        self.descriptor = descriptor
+        self.interpolator = InterpolateSparse2d('bicubic')
+        self.detection_threshold = detection_threshold
+        self.top_k = top_k
+        self.device = device
+        if train_detector:
+            for p in self.detector.parameters():
+                p.requires_grad = True
+            for p in self.descriptor.parameters():
+                p.requires_grad = False
+        else:
+            for p in self.detector.parameters():
+                p.requires_grad = False
+            for p in self.descriptor.parameters():
+                p.requires_grad = True
+        
+        self.softdetect = None
+        self.stride = descriptor.stride
+
+    def train(self, mode=True):
+        super().train(mode)
+        self.set_softdetect(top_k=500, scores_th=0.2)
+        
+    def eval(self):
+        super().eval()
+        self.set_softdetect(top_k=self.top_k, scores_th=0.01)
+        
+    def forward(self, samples: NestedTensor):
+        
+        if not isinstance(samples, NestedTensor):
+            samples = nested_tensor_from_tensor_list(samples)
+        
+        scoremap = self.detector(samples)
+        
+        feats, matchibility = self.descriptor(samples)
+        
+        return feats, scoremap, matchibility
+    
+    def set_softdetect(self, top_k=4096, scores_th=0.01):
+        self.softdetect = SoftDetect(radius=2, top_k=top_k, scores_th=scores_th)
+    
+    @torch.inference_mode()
+    def filter(self, matchibility):
+        # Filter out keypoints on the border
+        B, _, H, W = matchibility.shape
+        frame = torch.zeros(B, H, W, device=matchibility.device)
+        frame[:, self.stride:-self.stride, self.stride:-self.stride] = 1
+        matchibility = matchibility * frame
+        return matchibility
+    
+    @torch.inference_mode()
+    def extract(self, x):
+        if self.softdetect is None:
+            self.eval()
+        
+        x, rh1, rw1 = self.preprocess_tensor(x)
+        x = x.to(self.device).float()
+        B, _, _H1, _W1 = x.shape
+        M1, K1, H1 = self.forward(x)
+        M1 = F.normalize(M1, dim=1)
+        
+        keypoints, kptscores, scoredispersitys = self.softdetect(K1)
+        
+        keypoints = torch.vstack([keypoints[b].unsqueeze(0) for b in range(B)])
+        kptscores = torch.vstack([kptscores[b].unsqueeze(0) for b in range(B)])
+        
+        keypoints = to_pixel_coords(keypoints, _H1, _W1)
+        
+        feats = self.interpolator(M1, keypoints, H = _H1, W = _W1)
+        
+        feats = F.normalize(feats, dim=-1)
+		
+        # Correct kpt scale
+        keypoints = keypoints * torch.tensor([rw1,rh1], device=keypoints.device).view(1, -1)
+        valid = kptscores > self.detection_threshold
+
+        return [  
+                    {'keypoints': keypoints[b][valid[b]],
+                    'scores': kptscores[b][valid[b]],
+                    'descriptors': feats[b][valid[b]]} for b in range(B) 
+                ]	
+        
+    @torch.inference_mode()
+    def extract_3rd_party(self, x, model='aliked'):
+        """
+        one image per batch
+        """
+        x, rh1, rw1 = self.preprocess_tensor(x)
+        B, _, _H1, _W1 = x.shape
+        if model == 'aliked':
+            from third_party import extract_aliked_kpts
+            img = x
+            mkpts, scores = extract_aliked_kpts(img, self.device)
+        else:
+            raise ValueError('Unknown model')
+    
+        M1, _ = self.descriptor(x)
+        M1 = F.normalize(M1, dim=1)
+        
+        if mkpts.shape[1] > self.top_k:
+            idx = torch.argsort(scores, descending=True)[0][:self.top_k]
+            mkpts = mkpts[:,idx]
+            scores = scores[:,idx]
+
+        feats = self.interpolator(M1, mkpts, H = _H1, W = _W1)
+        feats = F.normalize(feats, dim=-1)
+        mkpts = mkpts * torch.tensor([rw1,rh1], device=mkpts.device).view(1, 1, -1)
+        
+        return [  
+				   {'keypoints': mkpts[b],
+                    'scores': scores[b],
+					'descriptors': feats[b]} for b in range(B) 
+			   ]
+        
+    @torch.inference_mode()
+    def extract_dense(self, x, n_limit=30000, thr=0.01):
+        self.set_softdetect(top_k=n_limit, scores_th=-1)
+            
+        x, rh1, rw1 = self.preprocess_tensor(x)
+
+        B, _, _H1, _W1 = x.shape
+
+        M1, K1, H1 = self.forward(x)
+        M1 = F.normalize(M1, dim=1)
+        
+        keypoints, kptscores, scoredispersitys = self.softdetect(K1)
+        
+        keypoints = torch.vstack([keypoints[b].unsqueeze(0) for b in range(B)])
+        kptscores = torch.vstack([kptscores[b].unsqueeze(0) for b in range(B)])
+        
+        keypoints = to_pixel_coords(keypoints, _H1, _W1)
+        
+        feats = self.interpolator(M1, keypoints, H = _H1, W = _W1)
+        
+        feats = F.normalize(feats, dim=-1)
+        
+        H1 = self.filter(H1)
+        
+        dense_kpts, dense_scores, inds = self.sample_dense_kpts(H1, n_limit=n_limit)
+ 
+        dense_keypoints = to_pixel_coords(dense_kpts, _H1, _W1)
+
+        dense_feats = self.interpolator(M1, dense_keypoints, H = _H1, W = _W1)
+        
+        dense_feats = F.normalize(dense_feats, dim=-1)
+        
+        keypoints = keypoints * torch.tensor([rw1,rh1], device=keypoints.device).view(1, -1)
+        dense_keypoints = dense_keypoints * torch.tensor([rw1,rh1], device=dense_keypoints.device).view(1, -1)	
+
+        valid = kptscores > self.detection_threshold
+        valid_dense = dense_scores > thr		
+
+        return [  
+                    {'keypoints': keypoints[b][valid[b]],
+                    'scores': kptscores[b][valid[b]],
+                    'descriptors': feats[b][valid[b]], 
+                    'keypoints_dense': dense_keypoints[b][valid_dense[b]],
+                    'scores_dense': dense_scores[b][valid_dense[b]],
+                    'descriptors_dense': dense_feats[b][valid_dense[b]]} for b in range(B)
+                ]
+        
+    @torch.inference_mode()
+    def sample_dense_kpts(self, keypoint_logits, threshold=0.01, n_limit=30000, force_kpts = True):
+        
+        B, K, H, W = keypoint_logits.shape
+
+        if n_limit < 0 or n_limit > H*W:
+            n_limit = min(H*W - 1, n_limit)
+
+        scoremap = keypoint_logits.permute(0,2,3,1)
+
+        scoremap = scoremap.reshape(B, H, W)
+
+        frame = torch.zeros(B, H, W, device=keypoint_logits.device)
+
+        frame[:, 1:-1, 1:-1] = 1
+
+        scoremap = scoremap * frame
+
+        scoremap = scoremap.reshape(B, H*W)
+
+        grid = self.get_grid(B, H, W, device = keypoint_logits.device)
+
+        inds = torch.topk(scoremap, n_limit, dim=1).indices
+
+        # inds = torch.multinomial(scoremap, top_k, replacement=False)
+        kpts = torch.gather(grid, 1, inds[..., None].expand(B, n_limit, 2))
+        scoremap = torch.gather(scoremap, 1, inds)
+        if force_kpts:
+            valid = scoremap > threshold
+            kpts = kpts[valid][None]
+            scoremap = scoremap[valid][None]
+
+        return kpts, scoremap, inds
+
+    def preprocess_tensor(self, x):
+        """ Guarantee that image is divisible by 32 to avoid aliasing artifacts. """
+        if isinstance(x, np.ndarray) and len(x.shape) == 3:
+            x = torch.tensor(x).permute(2,0,1)[None]
+        x = x.to(self.device).float()
+
+        H, W = x.shape[-2:]
+
+        _H, _W = (H//32) * 32, (W//32) * 32
+
+        rh, rw = H/_H, W/_W
+
+        x = F.interpolate(x, (_H, _W), mode='bilinear', align_corners=False)
+        return x, rh, rw
+
+    @torch.inference_mode()
+    def get_grid(self, B, H, W, device = None):
+        x1_n = torch.meshgrid(
+        *[
+            torch.linspace(
+                -1 + 1 / n, 1 - 1 / n, n, device=device
+            )
+            for n in (B, H, W)
+        ]
+        )
+        x1_n = torch.stack((x1_n[2], x1_n[1]), dim=-1).reshape(B, H * W, 2)
+        return x1_n
+
+def build(config=None, weights=None):
+    if config is None:
+        config = read_config('./configs/default.yaml')
+    if weights is not None:
+        config['weights'] = weights
+    device = torch.device(config['device'])
+    print('config', config)
+    detector = build_detector(config)
+    descriptor = build_descriptor(config)
+    model = RDD(
+        detector, 
+        descriptor, 
+        detection_threshold=config['detection_threshold'], 
+        top_k=config['top_k'], 
+        train_detector=config['train_detector'],
+        device=device
+    )
+    if 'weights' in config and config['weights'] is not None:
+        model.load_state_dict(torch.load(config['weights'], map_location='cpu'))
+    model.to(device)
+    return model

imcui/third_party/rdd/RDD/RDD_helper.py

@@ -0,0 +1,179 @@
+from .matchers import DualSoftmaxMatcher, DenseMatcher, LightGlue
+import torch
+import torch.nn as nn
+import numpy as np
+import torch.nn.functional as F
+import kornia
+
+class RDD_helper(nn.Module):
+    def __init__(self, RDD):
+        super().__init__()
+        self.matcher = DualSoftmaxMatcher(inv_temperature = 20, thr = 0.01)
+        self.dense_matcher = DenseMatcher(inv_temperature=20, thr=0.01)
+        self.RDD = RDD
+        self.lg_matcher = None
+        
+    @torch.inference_mode()
+    def match(self, img0, img1, thr=0.01, resize=None, top_k=4096):
+        if top_k is not None and top_k != self.RDD.top_k:
+            self.RDD.top_k = top_k
+            self.RDD.set_softdetect(top_k=top_k)
+        
+        img0, scale0 = self.parse_input(img0, resize)
+        img1, scale1 = self.parse_input(img1, resize)
+
+        out0 = self.RDD.extract(img0)[0]
+        out1 = self.RDD.extract(img1)[0]
+        
+        # get top_k confident matches
+        mkpts0, mkpts1, conf = self.matcher(out0, out1, thr)
+        
+        scale0 = 1.0 / scale0
+        scale1 = 1.0 / scale1
+        
+        mkpts0 = mkpts0 * scale0
+        mkpts1 = mkpts1 * scale1
+        
+        return mkpts0.cpu().numpy(), mkpts1.cpu().numpy(), conf.cpu().numpy()
+    
+    @torch.inference_mode()
+    def match_lg(self, img0, img1, thr=0.01, resize=None, top_k=4096):
+        if self.lg_matcher is None:
+            lg_conf = {
+                "name": "lightglue",  # just for interfacing
+                "input_dim": 256,  # input descriptor dimension (autoselected from weights)
+                "descriptor_dim": 256,
+                "add_scale_ori": False,
+                "n_layers": 9,
+                "num_heads": 4,
+                "flash": True,  # enable FlashAttention if available.
+                "mp": False,  # enable mixed precision
+                "filter_threshold": 0.01,  # match threshold
+                "depth_confidence": -1,  # depth confidence threshold
+                "width_confidence": -1,  # width confidence threshold
+                "weights": './weights/RDD_lg-v2.pth',  # path to the weights
+            }
+            self.lg_matcher = LightGlue(features='rdd', conf=lg_conf).to(self.RDD.device)
+            
+        if top_k is not None and top_k != self.RDD.top_k:
+            self.RDD.top_k = top_k
+            self.RDD.set_softdetect(top_k=top_k)
+            
+        img0, scale0 = self.parse_input(img0, resize=resize)
+        img1, scale1 = self.parse_input(img1, resize=resize)
+        
+        size0 = torch.tensor(img0.shape[-2:])[None]
+        size1 = torch.tensor(img1.shape[-2:])[None]
+        
+        out0 = self.RDD.extract(img0)[0]
+        out1 = self.RDD.extract(img1)[0]
+        
+        # get top_k confident matches
+        image0_data = {
+            'keypoints': out0['keypoints'][None],
+            'descriptors': out0['descriptors'][None],
+            'image_size': size0,
+        }
+
+        image1_data = {
+            'keypoints': out1['keypoints'][None],
+            'descriptors': out1['descriptors'][None],
+            'image_size': size1,
+        }
+        
+        pred = {}
+        
+        with torch.no_grad():
+            pred.update({'image0': image0_data, 'image1': image1_data})
+            pred.update(self.lg_matcher({**pred}))
+        
+        kpts0 = pred['image0']['keypoints'][0]
+        kpts1 = pred['image1']['keypoints'][0]
+        
+        matches = pred['matches'][0]
+
+        mkpts0 = kpts0[matches[... , 0]]
+        mkpts1 = kpts1[matches[... , 1]]
+        conf = pred['scores'][0]
+        
+        valid_mask = conf > thr
+        mkpts0 = mkpts0[valid_mask]
+        mkpts1 = mkpts1[valid_mask]
+        conf = conf[valid_mask]
+        
+        scale0 = 1.0 / scale0
+        scale1 = 1.0 / scale1
+        mkpts0 = mkpts0 * scale0
+        mkpts1 = mkpts1 * scale1
+        
+        return mkpts0.cpu().numpy(), mkpts1.cpu().numpy(), conf.cpu().numpy()
+    
+    @torch.inference_mode()
+    def match_dense(self, img0, img1, thr=0.01, resize=None):
+        
+        img0, scale0 = self.parse_input(img0, resize=resize)
+        img1, scale1 = self.parse_input(img1, resize=resize)
+
+        out0 = self.RDD.extract_dense(img0)[0]
+        out1 = self.RDD.extract_dense(img1)[0]
+        
+        # get top_k confident matches
+        mkpts0, mkpts1, conf = self.dense_matcher(out0, out1, thr, err_thr=self.RDD.stride)
+        
+        scale0 = 1.0 / scale0
+        scale1 = 1.0 / scale1
+        
+        mkpts0 = mkpts0 * scale0
+        mkpts1 = mkpts1 * scale1
+        
+        return mkpts0.cpu().numpy(), mkpts1.cpu().numpy(), conf.cpu().numpy()
+        
+    @torch.inference_mode()
+    def match_3rd_party(self, img0, img1, model='aliked', resize=None, thr=0.01):
+        img0, scale0 = self.parse_input(img0, resize=resize)
+        img1, scale1 = self.parse_input(img1, resize=resize)
+
+        out0 = self.RDD.extract_3rd_party(img0, model=model)[0]
+        out1 = self.RDD.extract_3rd_party(img1, model=model)[0]
+        
+        mkpts0, mkpts1, conf = self.matcher(out0, out1, thr)
+        
+        scale0 = 1.0 / scale0
+        scale1 = 1.0 / scale1
+        
+        mkpts0 = mkpts0 * scale0
+        mkpts1 = mkpts1 * scale1
+        
+        return mkpts0.cpu().numpy(), mkpts1.cpu().numpy(), conf.cpu().numpy()
+    
+    def parse_input(self, x, resize=None):
+        if len(x.shape) == 3:
+            x = x[None, ...]
+
+        if isinstance(x, np.ndarray):
+            x = torch.tensor(x).permute(0,3,1,2)/255
+        
+        h, w = x.shape[-2:]
+        size = h, w
+        
+        if resize is not None:
+            size = self.get_new_image_size(h, w, resize)
+            x = kornia.geometry.transform.resize(
+                x,
+                size,
+                side='long',
+                antialias=True,
+                align_corners=None,
+                interpolation='bilinear',
+            )
+        scale = torch.Tensor([x.shape[-1] / w, x.shape[-2] / h]).to(self.RDD.device)
+        
+        return x, scale
+    
+    def get_new_image_size(self, h, w, resize=1600):
+        aspect_ratio = w / h
+        size = int(resize / aspect_ratio), resize
+
+        size = list(map(lambda x: int(x // 32 * 32), size)) # make sure size is divisible by 32
+        
+        return size

imcui/third_party/rdd/RDD/dataset/megadepth/__init__.py

@@ -0,0 +1,2 @@
+from .megadepth import *
+from .megadepth_warper import *

imcui/third_party/rdd/RDD/dataset/megadepth/megadepth.py

@@ -0,0 +1,313 @@
+import os
+import copy
+import h5py
+import torch
+import pickle
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+from pathlib import Path
+from torchvision import transforms
+from torch.utils.data import Dataset
+
+import cv2
+from .utils import scale_intrinsics, warp_depth, warp_points2d
+
+class MegaDepthDataset(Dataset):
+    def __init__(
+            self,
+            root,
+            npz_path,
+            num_per_scene=100,
+            image_size=256,
+            min_overlap_score=0.1,
+            max_overlap_score=0.9,
+            gray=False,
+            crop_or_scale='scale',  # crop, scale, crop_scale
+            train=True,
+    ):
+        self.data_path = Path(root)
+        self.num_per_scene = num_per_scene
+        self.train = train
+        self.image_size = image_size
+        self.gray = gray
+        self.crop_or_scale = crop_or_scale
+
+        self.scene_info = dict(np.load(npz_path, allow_pickle=True))
+        self.pair_infos = self.scene_info['pair_infos'].copy()
+        del self.scene_info['pair_infos']
+        self.pair_infos = [pair_info for pair_info in self.pair_infos if pair_info[1] > min_overlap_score and pair_info[1] < max_overlap_score]
+        if len(self.pair_infos) > num_per_scene:
+            indices = np.random.choice(len(self.pair_infos), num_per_scene, replace=False)
+            self.pair_infos = [self.pair_infos[idx] for idx in indices]
+        self.transforms = transforms.Compose([transforms.ToPILImage(),
+                                                  transforms.ToTensor()])
+            
+    def __len__(self):
+        return len(self.pair_infos)
+
+    def recover_pair(self, idx):
+        (idx0, idx1), overlap_score, central_matches = self.pair_infos[idx % len(self)]
+        
+        img_name1 = self.scene_info['image_paths'][idx0]
+        img_name2 = self.scene_info['image_paths'][idx1]
+        
+        depth1 = '/'.join([self.scene_info['depth_paths'][idx0].replace('phoenix/S6/zl548/MegaDepth_v1', 'depth_undistorted').split('/')[i] for i in [0, 1, -1]])
+        depth2 = '/'.join([self.scene_info['depth_paths'][idx1].replace('phoenix/S6/zl548/MegaDepth_v1', 'depth_undistorted').split('/')[i] for i in [0, 1, -1]])
+        
+        depth_path1 = self.data_path / depth1
+        with h5py.File(depth_path1, 'r') as hdf5_file:
+            depth1 = np.array(hdf5_file['/depth'])
+        assert (np.min(depth1) >= 0)
+        image_path1 = self.data_path / img_name1
+        image1 = Image.open(image_path1)
+        if image1.mode != 'RGB':
+            image1 = image1.convert('RGB')
+        image1 = np.array(image1)
+        assert (image1.shape[0] == depth1.shape[0] and image1.shape[1] == depth1.shape[1])
+        intrinsics1 = self.scene_info['intrinsics'][idx0].copy()
+        pose1 = self.scene_info['poses'][idx0]
+        
+        depth_path2 = self.data_path / depth2
+        with h5py.File(depth_path2, 'r') as hdf5_file:
+            depth2 = np.array(hdf5_file['/depth'])
+        assert (np.min(depth2) >= 0)
+        image_path2 = self.data_path / img_name2
+        image2 = Image.open(image_path2)
+        if image2.mode != 'RGB':
+            image2 = image2.convert('RGB')
+        image2 = np.array(image2)
+        assert (image2.shape[0] == depth2.shape[0] and image2.shape[1] == depth2.shape[1])
+        intrinsics2 = self.scene_info['intrinsics'][idx1].copy()
+        pose2 = self.scene_info['poses'][idx1]
+
+        pose12 = pose2 @ np.linalg.inv(pose1)
+        pose21 = np.linalg.inv(pose12)
+
+        if self.train:
+            if "crop" in self.crop_or_scale:
+                # ================================================= compute central_match
+                DOWNSAMPLE = 10
+                # resize to speed up
+                depth1s = cv2.resize(depth1, (depth1.shape[1] // DOWNSAMPLE, depth1.shape[0] // DOWNSAMPLE))
+                depth2s = cv2.resize(depth2, (depth2.shape[1] // DOWNSAMPLE, depth2.shape[0] // DOWNSAMPLE))
+                intrinsic1s = scale_intrinsics(intrinsics1, (DOWNSAMPLE, DOWNSAMPLE))
+                intrinsic2s = scale_intrinsics(intrinsics2, (DOWNSAMPLE, DOWNSAMPLE))
+
+                # warp
+                depth12s = warp_depth(depth1s, intrinsic1s, intrinsic2s, pose12, depth2s.shape)
+                depth21s = warp_depth(depth2s, intrinsic2s, intrinsic1s, pose21, depth1s.shape)
+
+                depth12s[depth12s < 0] = 0
+                depth21s[depth21s < 0] = 0
+
+                valid12s = np.logical_and(depth12s > 0, depth2s > 0)
+                valid21s = np.logical_and(depth21s > 0, depth1s > 0)
+
+                pos1 = np.array(valid21s.nonzero())
+                try:
+                    idx1_random = np.random.choice(np.arange(pos1.shape[1]), 1)
+                    uv1s = pos1[:, idx1_random][[1, 0]].reshape(1, 2)
+                    d1s = np.array(depth1s[uv1s[0, 1], uv1s[0, 0]]).reshape(1, 1)
+
+                    uv12s, z12s = warp_points2d(uv1s, d1s, intrinsic1s, intrinsic2s, pose12)
+
+                    uv1 = uv1s[0] * DOWNSAMPLE
+                    uv2 = uv12s[0] * DOWNSAMPLE
+                except ValueError:
+                    uv1 = [depth1.shape[1] / 2, depth1.shape[0] / 2]
+                    uv2 = [depth2.shape[1] / 2, depth2.shape[0] / 2]
+
+                central_match = [uv1[1], uv1[0], uv2[1], uv2[0]]
+                # ================================================= compute central_match
+
+            if self.crop_or_scale == 'crop':
+                # =============== padding
+                h1, w1, _ = image1.shape
+                h2, w2, _ = image2.shape
+                if h1 < self.image_size:
+                    padding = np.zeros((self.image_size - h1, w1, 3))
+                    image1 = np.concatenate([image1, padding], axis=0).astype(np.uint8)
+                    depth1 = np.concatenate([depth1, padding[:, :, 0]], axis=0).astype(np.float32)
+                    h1, w1, _ = image1.shape
+                if w1 < self.image_size:
+                    padding = np.zeros((h1, self.image_size - w1, 3))
+                    image1 = np.concatenate([image1, padding], axis=1).astype(np.uint8)
+                    depth1 = np.concatenate([depth1, padding[:, :, 0]], axis=1).astype(np.float32)
+                if h2 < self.image_size:
+                    padding = np.zeros((self.image_size - h2, w2, 3))
+                    image2 = np.concatenate([image2, padding], axis=0).astype(np.uint8)
+                    depth2 = np.concatenate([depth2, padding[:, :, 0]], axis=0).astype(np.float32)
+                    h2, w2, _ = image2.shape
+                if w2 < self.image_size:
+                    padding = np.zeros((h2, self.image_size - w2, 3))
+                    image2 = np.concatenate([image2, padding], axis=1).astype(np.uint8)
+                    depth2 = np.concatenate([depth2, padding[:, :, 0]], axis=1).astype(np.float32)
+                # =============== padding
+                image1, bbox1, image2, bbox2 = self.crop(image1, image2, central_match)
+
+                depth1 = depth1[bbox1[0]: bbox1[0] + self.image_size, bbox1[1]: bbox1[1] + self.image_size]
+                depth2 = depth2[bbox2[0]: bbox2[0] + self.image_size, bbox2[1]: bbox2[1] + self.image_size]
+            elif self.crop_or_scale == 'scale':
+                image1, depth1, intrinsics1 = self.scale(image1, depth1, intrinsics1)
+                image2, depth2, intrinsics2 = self.scale(image2, depth2, intrinsics2)
+                bbox1 = bbox2 = np.array([0., 0.])
+            elif self.crop_or_scale == 'crop_scale':
+                bbox1 = bbox2 = np.array([0., 0.])
+                image1, depth1, intrinsics1 = self.crop_scale(image1, depth1, intrinsics1, central_match[:2])
+                image2, depth2, intrinsics2 = self.crop_scale(image2, depth2, intrinsics2, central_match[2:])
+            else:
+                raise RuntimeError(f"Unkown type {self.crop_or_scale}")
+        else:
+            bbox1 = bbox2 = np.array([0., 0.])
+
+        return (image1, depth1, intrinsics1, pose12, bbox1,
+                image2, depth2, intrinsics2, pose21, bbox2)
+
+    def scale(self, image, depth, intrinsic):
+        img_size_org = image.shape
+        image = cv2.resize(image, (self.image_size, self.image_size))
+        depth = cv2.resize(depth, (self.image_size, self.image_size))
+        intrinsic = scale_intrinsics(intrinsic, (img_size_org[1] / self.image_size, img_size_org[0] / self.image_size))
+        return image, depth, intrinsic
+
+    def crop_scale(self, image, depth, intrinsic, centeral):
+        h_org, w_org, three = image.shape
+        image_size = min(h_org, w_org)
+        if h_org > w_org:
+            if centeral[1] - image_size // 2 < 0:
+                h_start = 0
+            elif centeral[1] + image_size // 2 > h_org:
+                h_start = h_org - image_size
+            else:
+                h_start = int(centeral[1]) - image_size // 2
+            w_start = 0
+        else:
+            if centeral[0] - image_size // 2 < 0:
+                w_start = 0
+            elif centeral[0] + image_size // 2 > w_org:
+                w_start = w_org - image_size
+            else:
+                w_start = int(centeral[0]) - image_size // 2
+            h_start = 0
+
+        croped_image = image[h_start: h_start + image_size, w_start: w_start + image_size]
+        croped_depth = depth[h_start: h_start + image_size, w_start: w_start + image_size]
+        intrinsic[0, 2] = intrinsic[0, 2] - w_start
+        intrinsic[1, 2] = intrinsic[1, 2] - h_start
+
+        image = cv2.resize(croped_image, (self.image_size, self.image_size))
+        depth = cv2.resize(croped_depth, (self.image_size, self.image_size))
+        intrinsic = scale_intrinsics(intrinsic, (image_size / self.image_size, image_size / self.image_size))
+
+        return image, depth, intrinsic
+
+    def crop(self, image1, image2, central_match):
+        bbox1_i = max(int(central_match[0]) - self.image_size // 2, 0)
+        if bbox1_i + self.image_size >= image1.shape[0]:
+            bbox1_i = image1.shape[0] - self.image_size
+        bbox1_j = max(int(central_match[1]) - self.image_size // 2, 0)
+        if bbox1_j + self.image_size >= image1.shape[1]:
+            bbox1_j = image1.shape[1] - self.image_size
+
+        bbox2_i = max(int(central_match[2]) - self.image_size // 2, 0)
+        if bbox2_i + self.image_size >= image2.shape[0]:
+            bbox2_i = image2.shape[0] - self.image_size
+        bbox2_j = max(int(central_match[3]) - self.image_size // 2, 0)
+        if bbox2_j + self.image_size >= image2.shape[1]:
+            bbox2_j = image2.shape[1] - self.image_size
+
+        return (image1[bbox1_i: bbox1_i + self.image_size, bbox1_j: bbox1_j + self.image_size],
+                np.array([bbox1_i, bbox1_j]),
+                image2[bbox2_i: bbox2_i + self.image_size, bbox2_j: bbox2_j + self.image_size],
+                np.array([bbox2_i, bbox2_j])
+                )
+
+    def __getitem__(self, idx):
+        (image1, depth1, intrinsics1, pose12, bbox1,
+         image2, depth2, intrinsics2, pose21, bbox2) \
+            = self.recover_pair(idx)
+
+        if self.gray:
+            gray1 = cv2.cvtColor(image1, cv2.COLOR_RGB2GRAY)
+            gray2 = cv2.cvtColor(image2, cv2.COLOR_RGB2GRAY)
+            gray1 = transforms.ToTensor()(gray1)
+            gray2 = transforms.ToTensor()(gray2)
+        if self.transforms is not None:
+            image1, image2 = self.transforms(image1), self.transforms(image2)  # [C,H,W]
+        ret = {'image0': image1,
+               'image1': image2,
+               'angle': 0,
+               'overlap': self.pair_infos[idx][1],
+               'warp01_params': {'mode': 'se3',
+                                 'width': self.image_size if self.train else image1.shape[2],
+                                 'height': self.image_size if self.train else image1.shape[1],
+                                 'pose01': torch.from_numpy(pose12.astype(np.float32)),
+                                 'bbox0': torch.from_numpy(bbox1.astype(np.float32)),
+                                 'bbox1': torch.from_numpy(bbox2.astype(np.float32)),
+                                 'depth0': torch.from_numpy(depth1.astype(np.float32)),
+                                 'depth1': torch.from_numpy(depth2.astype(np.float32)),
+                                 'intrinsics0': torch.from_numpy(intrinsics1.astype(np.float32)),
+                                 'intrinsics1': torch.from_numpy(intrinsics2.astype(np.float32))},
+               'warp10_params': {'mode': 'se3',
+                                 'width': self.image_size if self.train else image2.shape[2],
+                                 'height': self.image_size if self.train else image2.shape[2],
+                                 'pose01': torch.from_numpy(pose21.astype(np.float32)),
+                                 'bbox0': torch.from_numpy(bbox2.astype(np.float32)),
+                                 'bbox1': torch.from_numpy(bbox1.astype(np.float32)),
+                                 'depth0': torch.from_numpy(depth2.astype(np.float32)),
+                                 'depth1': torch.from_numpy(depth1.astype(np.float32)),
+                                 'intrinsics0': torch.from_numpy(intrinsics2.astype(np.float32)),
+                                 'intrinsics1': torch.from_numpy(intrinsics1.astype(np.float32))},
+               }
+        if self.gray:
+            ret['gray0'] = gray1
+            ret['gray1'] = gray2
+        return ret
+
+
+if __name__ == '__main__':
+    from torch.utils.data import DataLoader
+    import matplotlib.pyplot as plt
+
+
+    def visualize(image0, image1, depth0, depth1):
+        # visualize image and depth
+        plt.figure(figsize=(9, 9))
+        plt.subplot(2, 2, 1)
+        plt.imshow(image0, cmap='gray')
+        plt.subplot(2, 2, 2)
+        plt.imshow(depth0)
+        plt.subplot(2, 2, 3)
+        plt.imshow(image1, cmap='gray')
+        plt.subplot(2, 2, 4)
+        plt.imshow(depth1)
+        plt.show()
+
+
+    dataset = MegaDepthDataset(  # root='../data/megadepth',
+        root='../data/imw2020val',
+        train=False,
+        using_cache=True,
+        pairs_per_scene=100,
+        image_size=256,
+        colorjit=True,
+        gray=False,
+        crop_or_scale='scale',
+    )
+    dataset.build_dataset()
+
+    batch_size = 2
+
+    loader = DataLoader(dataset, batch_size=1, shuffle=True, num_workers=0)
+
+    for idx, batch in enumerate(tqdm(loader)):
+        image0, image1 = batch['image0'], batch['image1']  # [B,3,H,W]
+        depth0, depth1 = batch['warp01_params']['depth0'], batch['warp01_params']['depth1']  # [B,H,W]
+        intrinsics0, intrinsics1 = batch['warp01_params']['intrinsics0'], batch['warp01_params'][
+            'intrinsics1']  # [B,3,3]
+
+        batch_size, channels, h, w = image0.shape
+
+        for b_idx in range(batch_size):
+            visualize(image0[b_idx].permute(1, 2, 0), image1[b_idx].permute(1, 2, 0), depth0[b_idx], depth1[b_idx])

imcui/third_party/rdd/RDD/dataset/megadepth/megadepth_warper.py

@@ -0,0 +1,75 @@
+import torch
+from kornia.utils import create_meshgrid
+import matplotlib.pyplot as plt
+import pdb
+from .utils import warp
+
+@torch.no_grad()
+def spvs_coarse(data, scale = 8):
+    N, _, H0, W0 = data['image0'].shape
+    _, _, H1, W1 = data['image1'].shape
+    device = data['image0'].device
+    corrs = []
+    for idx in range(N):
+        warp01_params = {}
+        for k, v in data['warp01_params'].items():
+            if isinstance(v[idx], torch.Tensor):
+                warp01_params[k] = v[idx].to(device)
+            else:
+                warp01_params[k] = v[idx]
+        warp10_params = {}
+        for k, v in data['warp10_params'].items():
+            if isinstance(v[idx], torch.Tensor):
+                warp10_params[k] = v[idx].to(device)
+            else:
+                warp10_params[k] = v[idx]
+            
+        # create kpts
+        h0, w0, h1, w1 = map(lambda x: x // scale, [H0, W0, H1, W1])
+        grid_pt1_c = create_meshgrid(h1, w1, False, device).reshape(h1*w1, 2)    # [N, hw, 2]
+        
+        # normalize kpts
+        grid_pt1_c = grid_pt1_c * scale
+
+        
+
+        try:
+            grid_pt1_c_valid, grid_pt10_c, ids1, ids1_out = warp(grid_pt1_c, warp10_params)
+            grid_pt10_c_valid, grid_pt01_c, ids0, ids0_out = warp(grid_pt10_c, warp01_params)
+            
+            # check reproj error
+            grid_pt1_c_valid = grid_pt1_c_valid[ids0]
+            dist = torch.linalg.norm(grid_pt1_c_valid - grid_pt01_c, dim=-1)
+            
+            mask_mutual = (dist < 1.5) 
+            
+            #get correspondences
+            pts = torch.cat([grid_pt10_c_valid[mask_mutual] / scale,
+                                grid_pt01_c[mask_mutual] / scale], dim=-1)
+            #remove repeated correspondences
+            lut_mat12 = torch.ones((h1, w1, 4), device = device, dtype = torch.float32) * -1
+            lut_mat21 = torch.clone(lut_mat12)
+            src_pts = pts[:, :2]
+            tgt_pts = pts[:, 2:]
+        
+            lut_mat12[src_pts[:,1].long(), src_pts[:,0].long()] = torch.cat([src_pts, tgt_pts], dim=1)
+            mask_valid12 = torch.all(lut_mat12 >= 0, dim=-1)
+            points = lut_mat12[mask_valid12]
+
+            #Target-src check
+            src_pts, tgt_pts = points[:, :2], points[:, 2:]
+            lut_mat21[tgt_pts[:,1].long(), tgt_pts[:,0].long()] = torch.cat([src_pts, tgt_pts], dim=1)
+            mask_valid21 = torch.all(lut_mat21 >= 0, dim=-1)
+            points = lut_mat21[mask_valid21]
+
+            corrs.append(points)
+        except:
+            corrs.append(torch.zeros((0, 4), device = device))
+            #pdb.set_trace()
+            #print('..')
+
+    #Plot for debug purposes    
+    # for i in range(len(corrs)):
+    #     plot_corrs(data['image0'][i], data['image1'][i], corrs[i][:, :2]*8, corrs[i][:, 2:]*8)
+
+    return corrs

imcui/third_party/rdd/RDD/dataset/megadepth/utils.py

@@ -0,0 +1,848 @@
+"""
+	"XFeat: Accelerated Features for Lightweight Image Matching, CVPR 2024."
+	https://www.verlab.dcc.ufmg.br/descriptors/xfeat_cvpr24/
+
+    MegaDepth data handling was adapted from 
+    LoFTR official code: https://github.com/zju3dv/LoFTR/blob/master/src/datasets/megadepth.py
+"""
+
+import io
+import cv2
+import numpy as np
+import h5py
+import torch
+from numpy.linalg import inv
+from kornia.geometry.epipolar import essential_from_Rt
+from kornia.geometry.epipolar import fundamental_from_essential
+
+import cv2
+import torch
+import numpy as np
+from numba import jit
+from copy import deepcopy
+
+try:
+    from utils.project_depth_nn_cython_pkg import project_depth_nn_cython
+
+    nn_cython = True
+except:
+    print('\033[1;41;37mWarning: using python to project depth!!!\033[0m')
+
+    nn_cython = False
+
+
+class EmptyTensorError(Exception):
+    pass
+
+
+def mutual_NN(cross_matrix, mode: str = 'min'):
+    """
+    compute mutual nearest neighbor from a cross_matrix, non-differentiable function
+    :param cross_matrix: N0xN1
+    :param mode: 'min': mutual minimum; 'max':mutual maximum
+    :return: index0,index1, Mx2
+    """
+    if mode == 'min':
+        nn0 = cross_matrix == cross_matrix.min(dim=1, keepdim=True)[0]
+        nn1 = cross_matrix == cross_matrix.min(dim=0, keepdim=True)[0]
+    elif mode == 'max':
+        nn0 = cross_matrix == cross_matrix.max(dim=1, keepdim=True)[0]
+        nn1 = cross_matrix == cross_matrix.max(dim=0, keepdim=True)[0]
+    else:
+        raise TypeError("error mode, must be 'min' or 'max'.")
+
+    mutual_nn = nn0 * nn1
+
+    return torch.nonzero(mutual_nn, as_tuple=False)
+
+
+def mutual_argmax(value, mask=None, as_tuple=True):
+    """
+    Args:
+        value: MxN
+        mask:  MxN
+
+    Returns:
+
+    """
+    value = value - value.min()  # convert to non-negative tensor
+    if mask is not None:
+        value = value * mask
+
+    max0 = value.max(dim=1, keepdim=True)  # the col index the max value in each row
+    max1 = value.max(dim=0, keepdim=True)
+
+    valid_max0 = value == max0[0]
+    valid_max1 = value == max1[0]
+
+    mutual = valid_max0 * valid_max1
+    if mask is not None:
+        mutual = mutual * mask
+
+    return mutual.nonzero(as_tuple=as_tuple)
+
+
+def mutual_argmin(value, mask=None):
+    return mutual_argmax(-value, mask)
+
+
+def compute_keypoints_distance(kpts0, kpts1, p=2):
+    """
+    Args:
+        kpts0: torch.tensor [M,2]
+        kpts1: torch.tensor [N,2]
+        p: (int, float, inf, -inf, 'fro', 'nuc', optional): the order of norm
+
+    Returns:
+        dist, torch.tensor [N,M]
+    """
+    dist = kpts0[:, None, :] - kpts1[None, :, :]  # [M,N,2]
+    dist = torch.norm(dist, p=p, dim=2)  # [M,N]
+    return dist
+
+
+def keypoints_normal2pixel(kpts_normal, w, h):
+    wh = kpts_normal[0].new_tensor([[w - 1, h - 1]])
+    kpts_pixel = [(kpts + 1) / 2 * wh for kpts in kpts_normal]
+    return kpts_pixel
+
+
+def plot_keypoints(image, kpts, radius=2, color=(255, 0, 0)):
+    image = image.cpu().detach().numpy() if isinstance(image, torch.Tensor) else image
+    kpts = kpts.cpu().detach().numpy() if isinstance(kpts, torch.Tensor) else kpts
+
+    if image.dtype is not np.dtype('uint8'):
+        image = image * 255
+        image = image.astype(np.uint8)
+
+    if len(image.shape) == 2 or image.shape[2] == 1:
+        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+
+    out = np.ascontiguousarray(deepcopy(image))
+    kpts = np.round(kpts).astype(int)
+
+    for kpt in kpts:
+        y0, x0 = kpt
+        cv2.drawMarker(out, (x0, y0), color, cv2.MARKER_CROSS, radius)
+
+        # cv2.circle(out, (x0, y0), radius, color, -1, lineType=cv2.LINE_4)
+    return out
+
+
+def plot_matches(image0, image1, kpts0, kpts1, radius=2, color=(255, 0, 0), mcolor=(0, 255, 0), layout='lr'):
+    image0 = image0.cpu().detach().numpy() if isinstance(image0, torch.Tensor) else image0
+    image1 = image1.cpu().detach().numpy() if isinstance(image1, torch.Tensor) else image1
+    kpts0 = kpts0.cpu().detach().numpy() if isinstance(kpts0, torch.Tensor) else kpts0
+    kpts1 = kpts1.cpu().detach().numpy() if isinstance(kpts1, torch.Tensor) else kpts1
+
+    out0 = plot_keypoints(image0, kpts0, radius, color)
+    out1 = plot_keypoints(image1, kpts1, radius, color)
+
+    H0, W0 = image0.shape[0], image0.shape[1]
+    H1, W1 = image1.shape[0], image1.shape[1]
+
+    if layout == "lr":
+        H, W = max(H0, H1), W0 + W1
+        out = 255 * np.ones((H, W, 3), np.uint8)
+        out[:H0, :W0, :] = out0
+        out[:H1, W0:, :] = out1
+    elif layout == "ud":
+        H, W = H0 + H1, max(W0, W1)
+        out = 255 * np.ones((H, W, 3), np.uint8)
+        out[:H0, :W0, :] = out0
+        out[H0:, :W1, :] = out1
+    else:
+        raise ValueError("The layout must be 'lr' or 'ud'!")
+
+    kpts0 = np.round(kpts0).astype(int)
+    kpts1 = np.round(kpts1).astype(int)
+
+    for kpt0, kpt1 in zip(kpts0, kpts1):
+        (y0, x0), (y1, x1) = kpt0, kpt1
+
+        if layout == "lr":
+            cv2.line(out, (x0, y0), (x1 + W0, y1), color=mcolor, thickness=1, lineType=cv2.LINE_AA)
+        elif layout == "ud":
+            cv2.line(out, (x0, y0), (x1, y1 + H0), color=mcolor, thickness=1, lineType=cv2.LINE_AA)
+
+    return out
+
+
+def interpolate_depth(pos, depth):
+    pos = pos.t()[[1, 0]]  # Nx2 -> 2xN; w,h -> h,w(i,j)
+
+    # =============================================== from d2-net
+    device = pos.device
+
+    ids = torch.arange(0, pos.size(1), device=device)
+
+    h, w = depth.size()
+
+    i = pos[0, :].detach()  # TODO: changed here
+    j = pos[1, :].detach()  # TODO: changed here
+
+    # Valid corners
+    i_top_left = torch.floor(i).long()
+    j_top_left = torch.floor(j).long()
+    valid_top_left = torch.min(i_top_left >= 0, j_top_left >= 0)
+
+    i_top_right = torch.floor(i).long()
+    j_top_right = torch.ceil(j).long()
+    valid_top_right = torch.min(i_top_right >= 0, j_top_right < w)
+
+    i_bottom_left = torch.ceil(i).long()
+    j_bottom_left = torch.floor(j).long()
+    valid_bottom_left = torch.min(i_bottom_left < h, j_bottom_left >= 0)
+
+    i_bottom_right = torch.ceil(i).long()
+    j_bottom_right = torch.ceil(j).long()
+    valid_bottom_right = torch.min(i_bottom_right < h, j_bottom_right < w)
+
+    valid_corners = torch.min(torch.min(valid_top_left, valid_top_right),
+                              torch.min(valid_bottom_left, valid_bottom_right))
+
+    i_top_left = i_top_left[valid_corners]
+    j_top_left = j_top_left[valid_corners]
+
+    i_top_right = i_top_right[valid_corners]
+    j_top_right = j_top_right[valid_corners]
+
+    i_bottom_left = i_bottom_left[valid_corners]
+    j_bottom_left = j_bottom_left[valid_corners]
+
+    i_bottom_right = i_bottom_right[valid_corners]
+    j_bottom_right = j_bottom_right[valid_corners]
+
+    ids = ids[valid_corners]
+    ids_valid_corners = deepcopy(ids)
+    if ids.size(0) == 0:
+        # raise ValueError('empty tensor: ids')
+        raise EmptyTensorError
+
+    # Valid depth
+    valid_depth = torch.min(torch.min(depth[i_top_left, j_top_left] > 0,
+                                      depth[i_top_right, j_top_right] > 0),
+                            torch.min(depth[i_bottom_left, j_bottom_left] > 0,
+                                      depth[i_bottom_right, j_bottom_right] > 0))
+
+    i_top_left = i_top_left[valid_depth]
+    j_top_left = j_top_left[valid_depth]
+
+    i_top_right = i_top_right[valid_depth]
+    j_top_right = j_top_right[valid_depth]
+
+    i_bottom_left = i_bottom_left[valid_depth]
+    j_bottom_left = j_bottom_left[valid_depth]
+
+    i_bottom_right = i_bottom_right[valid_depth]
+    j_bottom_right = j_bottom_right[valid_depth]
+
+    ids = ids[valid_depth]
+    ids_valid_depth = deepcopy(ids)
+    if ids.size(0) == 0:
+        # raise ValueError('empty tensor: ids')
+        raise EmptyTensorError
+
+    # Interpolation
+    i = i[ids]
+    j = j[ids]
+    dist_i_top_left = i - i_top_left.float()
+    dist_j_top_left = j - j_top_left.float()
+    w_top_left = (1 - dist_i_top_left) * (1 - dist_j_top_left)
+    w_top_right = (1 - dist_i_top_left) * dist_j_top_left
+    w_bottom_left = dist_i_top_left * (1 - dist_j_top_left)
+    w_bottom_right = dist_i_top_left * dist_j_top_left
+
+    interpolated_depth = (w_top_left * depth[i_top_left, j_top_left] +
+                          w_top_right * depth[i_top_right, j_top_right] +
+                          w_bottom_left * depth[i_bottom_left, j_bottom_left] +
+                          w_bottom_right * depth[i_bottom_right, j_bottom_right])
+
+    # pos = torch.cat([i.view(1, -1), j.view(1, -1)], dim=0)
+
+    pos = pos[:, ids]
+
+    # =============================================== from d2-net
+    pos = pos[[1, 0]].t()  # 2xN -> Nx2;  h,w(i,j) -> w,h
+
+    # interpolated_depth: valid interpolated depth
+    # pos: valid position (keypoint)
+    # ids: indices of valid position (keypoint)
+
+    return [interpolated_depth, pos, ids, ids_valid_corners, ids_valid_depth]
+
+
+def to_homogeneous(kpts):
+    '''
+    :param kpts: Nx2
+    :return: Nx3
+    '''
+    ones = kpts.new_ones([kpts.shape[0], 1])
+    return torch.cat((kpts, ones), dim=1)
+
+
+def warp_homography(kpts0, params):
+    '''
+    :param kpts: Nx2
+    :param homography_matrix: 3x3
+    :return:
+    '''
+    homography_matrix = params['homography_matrix']
+    w, h = params['width'], params['height']
+    kpts0_homogeneous = to_homogeneous(kpts0)
+    kpts01_homogeneous = torch.einsum('ij,kj->ki', homography_matrix, kpts0_homogeneous)
+    kpts01 = kpts01_homogeneous[:, :2] / kpts01_homogeneous[:, 2:]
+
+    kpts01_ = kpts01.detach()
+    # due to float coordinates, the upper boundary should be (w-1) and (h-1).
+    # For example, if the image size is 480, then the coordinates should in [0~470].
+    # 470.5 is not acceptable.
+    valid01 = (kpts01_[:, 0] >= 0) * (kpts01_[:, 0] <= w - 1) * (kpts01_[:, 1] >= 0) * (kpts01_[:, 1] <= h - 1)
+    kpts0_valid = kpts0[valid01]
+    kpts01_valid = kpts01[valid01]
+    ids = torch.nonzero(valid01, as_tuple=False)[:, 0]
+    ids_out = torch.nonzero(~valid01, as_tuple=False)[:, 0]
+
+    # kpts0_valid: valid keypoints0, the invalid and inconsistance keypoints are removed
+    # kpts01_valid: the warped valid keypoints0
+    # ids: the valid indices
+    return kpts0_valid, kpts01_valid, ids, ids_out
+
+
+def project(points3d, K):
+    """
+    project 3D points to image plane
+
+    Args:
+        points3d: [N,3]
+        K: [3,3]
+
+    Returns:
+        uv, (u,v), [N,2]
+    """
+    if type(K) == torch.Tensor:
+        zuv1 = torch.einsum('jk,nk->nj', K, points3d)  # z*(u,v,1) = K*points3d -> [N,3]
+    elif type(K) == np.ndarray:
+        zuv1 = np.einsum('jk,nk->nj', K, points3d)
+    else:
+        raise TypeError("Input type should be 'torch.tensor' or 'numpy.ndarray'")
+    uv1 = zuv1 / zuv1[:, -1][:, None]  # (u,v,1) -> [N,3]
+    uv = uv1[:, 0:2]  # (u,v) -> [N,2]
+    return uv, zuv1[:, -1]
+
+
+def unproject(uv, d, K):
+    """
+    unproject pixels uv to 3D points
+
+    Args:
+        uv: [N,2]
+        d: depth, [N,1]
+        K: [3,3]
+
+    Returns:
+        3D points, [N,3]
+    """
+    duv = uv * d  # (u,v) [N,2]
+    if type(K) == torch.Tensor:
+        duv1 = torch.cat([duv, d], dim=1)  # z*(u,v,1) [N,3]
+        K_inv = torch.inverse(K)  # [3,3]
+        points3d = torch.einsum('jk,nk->nj', K_inv, duv1)  # [N,3]
+    elif type(K) == np.ndarray:
+        duv1 = np.concatenate((duv, d), axis=1)  # z*(u,v,1) [N,3]
+        K_inv = np.linalg.inv(K)  # [3,3]
+        points3d = np.einsum('jk,nk->nj', K_inv, duv1)  # [N,3]
+    else:
+        raise TypeError("Input type should be 'torch.tensor' or 'numpy.ndarray'")
+    return points3d
+
+
+def warp_se3(kpts0, params):
+    pose01 = params['pose01']  # relative motion
+    bbox0 = params['bbox0']  # row, col
+    bbox1 = params['bbox1']
+    depth0 = params['depth0']
+    depth1 = params['depth1']
+    intrinsics0 = params['intrinsics0']
+    intrinsics1 = params['intrinsics1']
+
+    # kpts0_valid: valid kpts0
+    # z0_valid: depth of valid kpts0
+    # ids0: the indices of valid kpts0 ( valid corners and valid depth)
+    # ids0_valid_corners: the valid indices of kpts0 in image ( 0<=x<w, 0<=y<h )
+    # ids0_valid_depth: the valid indices of kpts0 with valid depth ( depth > 0 )
+    z0_valid, kpts0_valid, ids0, ids0_valid_corners, ids0_valid_depth = interpolate_depth(kpts0, depth0)
+
+    # COLMAP convention
+    bkpts0_valid = kpts0_valid + bbox0[[1, 0]][None, :] + 0.5
+
+    # unproject pixel coordinate to 3D points (camera coordinate system)
+    bpoints3d0 = unproject(bkpts0_valid, z0_valid.unsqueeze(1), intrinsics0)  # [:,3]
+    bpoints3d0_homo = to_homogeneous(bpoints3d0)  # [:,4]
+
+    # warp 3D point (camera 0 coordinate system) to 3D point (camera 1 coordinate system)
+    bpoints3d01_homo = torch.einsum('jk,nk->nj', pose01, bpoints3d0_homo)  # [:,4]
+    bpoints3d01 = bpoints3d01_homo[:, 0:3]  # [:,3]
+
+    # project 3D point (camera coordinate system) to pixel coordinate
+    buv01, z01 = project(bpoints3d01, intrinsics1)  # uv: [:,2], (h,w); z1: [N]
+
+    uv01 = buv01 - bbox1[None, [1, 0]] - .5
+
+    # kpts01_valid: valid kpts01
+    # z01_valid: depth of valid kpts01
+    # ids01: the indices of valid kpts01 ( valid corners and valid depth)
+    # ids01_valid_corners: the valid indices of kpts01 in image ( 0<=x<w, 0<=y<h )
+    # ids01_valid_depth: the valid indices of kpts01 with valid depth ( depth > 0 )
+    z01_interpolate, kpts01_valid, ids01, ids01_valid_corners, ids01_valid_depth = interpolate_depth(uv01, depth1)
+
+    outimage_mask = torch.ones(ids0.shape[0], device=ids0.device).bool()
+    outimage_mask[ids01_valid_corners] = 0
+    ids01_invalid_corners = torch.arange(0, ids0.shape[0], device=ids0.device)[outimage_mask]
+    ids_outside = ids0[ids01_invalid_corners]
+
+    # ids_valid: matched kpts01 without occlusion
+    ids_valid = ids0[ids01]
+    kpts0_valid = kpts0_valid[ids01]
+    z01_proj = z01[ids01]
+
+    inlier_mask = torch.abs(z01_proj - z01_interpolate) < 0.05
+
+    # indices of kpts01 with occlusion
+    ids_occlude = ids_valid[~inlier_mask]
+
+    ids_valid = ids_valid[inlier_mask]
+    if ids_valid.size(0) == 0:
+        # raise ValueError('empty tensor: ids')
+        raise EmptyTensorError
+
+    kpts01_valid = kpts01_valid[inlier_mask]
+    kpts0_valid = kpts0_valid[inlier_mask]
+
+    # indices of kpts01 which are no matches in image1 for sure,
+    # other projected kpts01 are not sure because of no depth in image0 or imgae1
+    ids_out = torch.cat([ids_outside, ids_occlude])
+
+    # kpts0_valid: valid keypoints0, the invalid and inconsistance keypoints are removed
+    # kpts01_valid: the warped valid keypoints0
+    # ids: the valid indices
+    return kpts0_valid, kpts01_valid, ids_valid, ids_out
+
+
+def warp(kpts0, params: dict):
+    mode = params['mode']
+    if mode == 'homo':
+        return warp_homography(kpts0, params)
+    elif mode == 'se3':
+        return warp_se3(kpts0, params)
+    else:
+        raise ValueError('unknown mode!')
+
+
+def warp_xy(kpts0_xy, params: dict):
+    w, h = params['width'], params['height']
+    kpts0 = (kpts0_xy / 2 + 0.5) * kpts0_xy.new_tensor([[w - 1, h - 1]])
+    kpts0, kpts01, ids = warp(kpts0, params)
+    kpts01_xy = 2 * kpts01 / kpts01.new_tensor([[w - 1, h - 1]]) - 1
+    kpts0_xy = 2 * kpts0 / kpts0.new_tensor([[w - 1, h - 1]]) - 1
+    return kpts0_xy, kpts01_xy, ids
+
+
+def scale_intrinsics(K, scales):
+    scales = np.diag([1. / scales[0], 1. / scales[1], 1.])
+    return np.dot(scales, K)
+
+
+def warp_points3d(points3d0, pose01):
+    points3d0_homo = np.concatenate((points3d0, np.ones(points3d0.shape[0])[:, np.newaxis]), axis=1)  # [:,4]
+
+    points3d01_homo = np.einsum('jk,nk->nj', pose01, points3d0_homo)  # [N,4]
+    points3d01 = points3d01_homo[:, 0:3]  # [N,3]
+
+    return points3d01
+
+
+def unproject_depth(depth, K):
+    h, w = depth.shape
+
+    wh_range = np.mgrid[0:w, 0:h].transpose(2, 1, 0)  # [H,W,2]
+
+    uv = wh_range.reshape(-1, 2)
+    d = depth.reshape(-1, 1)
+    points3d = unproject(uv, d, K)
+
+    valid = np.logical_and((d[:, 0] > 0), (points3d[:, 2] > 0))
+
+    return points3d, valid
+
+
+@jit(nopython=True)
+def project_depth_nn_python(uv, z, depth):
+    h, w = depth.shape
+    # TODO: speed up the for loop
+    for idx in range(len(uv)):
+        uvi = uv[idx]
+        x = int(round(uvi[0]))
+        y = int(round(uvi[1]))
+
+        if x < 0 or y < 0 or x >= w or y >= h:
+            continue
+
+        if depth[y, x] == 0. or depth[y, x] > z[idx]:
+            depth[y, x] = z[idx]
+    return depth
+
+
+def project_nn(uv, z, depth):
+    """
+    uv: pixel coordinates [N,2]
+    z: projected depth (xyz -> z) [N]
+    depth: output depth array: [h,w]
+    """
+    if nn_cython:
+        return project_depth_nn_cython(uv.astype(np.float64),
+                                       z.astype(np.float64),
+                                       depth.astype(np.float64))
+    else:
+        return project_depth_nn_python(uv, z, depth)
+
+
+def warp_depth(depth0, intrinsics0, intrinsics1, pose01, shape1):
+    points3d0, valid0 = unproject_depth(depth0, intrinsics0)  # [:,3]
+    points3d0 = points3d0[valid0]
+
+    points3d01 = warp_points3d(points3d0, pose01)
+
+    uv01, z01 = project(points3d01, intrinsics1)  # uv: [N,2], (h,w); z1: [N]
+
+    depth01 = project_nn(uv01, z01, depth=np.zeros(shape=shape1))
+
+    return depth01
+
+
+def warp_points2d(uv0, d0, intrinsics0, intrinsics1, pose01):
+    points3d0 = unproject(uv0, d0, intrinsics0)
+    points3d01 = warp_points3d(points3d0, pose01)
+    uv01, z01 = project(points3d01, intrinsics1)
+    return uv01, z01
+
+
+def display_image_in_actual_size(image):
+    import matplotlib.pyplot as plt
+
+    dpi = 100
+    height, width = image.shape[:2]
+
+    # What size does the figure need to be in inches to fit the image?
+    figsize = width / float(dpi), height / float(dpi)
+
+    # Create a figure of the right size with one axes that takes up the full figure
+    fig = plt.figure(figsize=figsize)
+    ax = fig.add_axes([0, 0, 1, 1])
+
+    # Hide spines, ticks, etc.
+    ax.axis('off')
+
+    # Display the image.
+    if len(image.shape) == 3:
+        ax.imshow(image, cmap='gray')
+    elif len(image.shape) == 2:
+        if image.dtype == np.uint8:
+            ax.imshow(image, cmap='gray')
+        else:
+            ax.imshow(image)
+            ax.text(20, 20, f"Range: {image.min():g}~{image.max():g}", color='red')
+    plt.show()
+
+
+# ====================================== copied from ASLFeat
+from datetime import datetime
+
+
+class ClassProperty(property):
+    """For dynamically obtaining system time"""
+
+    def __get__(self, cls, owner):
+        return classmethod(self.fget).__get__(None, owner)()
+
+
+class Notify(object):
+    """Colorful printing prefix.
+    A quick example:
+    print(Notify.INFO, YOUR TEXT, Notify.ENDC)
+    """
+
+    def __init__(self):
+        pass
+
+    @ClassProperty
+    def HEADER(cls):
+        return str(datetime.now()) + ': \033[95m'
+
+    @ClassProperty
+    def INFO(cls):
+        return str(datetime.now()) + ': \033[92mI'
+
+    @ClassProperty
+    def OKBLUE(cls):
+        return str(datetime.now()) + ': \033[94m'
+
+    @ClassProperty
+    def WARNING(cls):
+        return str(datetime.now()) + ': \033[93mW'
+
+    @ClassProperty
+    def FAIL(cls):
+        return str(datetime.now()) + ': \033[91mF'
+
+    @ClassProperty
+    def BOLD(cls):
+        return str(datetime.now()) + ': \033[1mB'
+
+    @ClassProperty
+    def UNDERLINE(cls):
+        return str(datetime.now()) + ': \033[4mU'
+
+    ENDC = '\033[0m'
+
+def get_essential(T0, T1):
+    R0 = T0[:3, :3]
+    R1 = T1[:3, :3]
+    
+    t0 = T0[:3, 3].reshape(3, 1)
+    t1 = T1[:3, 3].reshape(3, 1)
+    
+    R0 = torch.tensor(R0, dtype=torch.float32)
+    R1 = torch.tensor(R1, dtype=torch.float32)
+    t0 = torch.tensor(t0, dtype=torch.float32)
+    t1 = torch.tensor(t1, dtype=torch.float32)
+    
+    E = essential_from_Rt(R0, t0, R1, t1)
+    
+    return E
+
+def get_fundamental(E, K0, K1):
+    F = fundamental_from_essential(E, K0, K1)
+    
+    return F
+try:
+    # for internel use only
+    from .client import MEGADEPTH_CLIENT, SCANNET_CLIENT
+except Exception:
+    MEGADEPTH_CLIENT = SCANNET_CLIENT = None
+
+# --- DATA IO ---
+
+def load_array_from_s3(
+    path, client, cv_type,
+    use_h5py=False,
+):
+    byte_str = client.Get(path)
+    try:
+        if not use_h5py:
+            raw_array = np.fromstring(byte_str, np.uint8)
+            data = cv2.imdecode(raw_array, cv_type)
+        else:
+            f = io.BytesIO(byte_str)
+            data = np.array(h5py.File(f, 'r')['/depth'])
+    except Exception as ex:
+        print(f"==> Data loading failure: {path}")
+        raise ex
+
+    assert data is not None
+    return data
+
+
+def imread_gray(path, augment_fn=None, client=SCANNET_CLIENT):
+    cv_type = cv2.IMREAD_GRAYSCALE if augment_fn is None \
+                else cv2.IMREAD_COLOR
+    if str(path).startswith('s3://'):
+        image = load_array_from_s3(str(path), client, cv_type)
+    else:
+        image = cv2.imread(str(path), 1)
+
+    if augment_fn is not None:
+        image = cv2.imread(str(path), cv2.IMREAD_COLOR)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        image = augment_fn(image)
+        image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+    return image  # (h, w)
+
+
+def get_resized_wh(w, h, resize=None):
+    if resize is not None:  # resize the longer edge
+        scale = resize / max(h, w)
+        w_new, h_new = int(round(w*scale)), int(round(h*scale))
+    else:
+        w_new, h_new = w, h
+    return w_new, h_new
+
+
+def get_divisible_wh(w, h, df=None):
+    if df is not None:
+        w_new, h_new = map(lambda x: int(x // df * df), [w, h])
+    else:
+        w_new, h_new = w, h
+    return w_new, h_new
+
+
+def pad_bottom_right(inp, pad_size, ret_mask=False):
+    assert isinstance(pad_size, int) and pad_size >= max(inp.shape[-2:]), f"{pad_size} < {max(inp.shape[-2:])}"
+    mask = None
+    if inp.ndim == 2:
+        padded = np.zeros((pad_size, pad_size), dtype=inp.dtype)
+        padded[:inp.shape[0], :inp.shape[1]] = inp
+        if ret_mask:
+            mask = np.zeros((pad_size, pad_size), dtype=bool)
+            mask[:inp.shape[0], :inp.shape[1]] = True
+    elif inp.ndim == 3:
+        padded = np.zeros((inp.shape[0], pad_size, pad_size), dtype=inp.dtype)
+        padded[:, :inp.shape[1], :inp.shape[2]] = inp
+        if ret_mask:
+            mask = np.zeros((inp.shape[0], pad_size, pad_size), dtype=bool)
+            mask[:, :inp.shape[1], :inp.shape[2]] = True
+    else:
+        raise NotImplementedError()
+    return padded, mask
+
+
+# --- MEGADEPTH ---
+
+def fix_path_from_d2net(path):
+    if not path:
+        return None
+
+    path = path.replace('Undistorted_SfM/', '')
+    path = path.replace('images', 'dense0/imgs')
+    path = path.replace('phoenix/S6/zl548/MegaDepth_v1/', '')
+
+    return path
+
+def read_megadepth_gray(path, resize=None, df=None, padding=False, augment_fn=None):
+    """
+    Args:
+        resize (int, optional): the longer edge of resized images. None for no resize.
+        padding (bool): If set to 'True', zero-pad resized images to squared size.
+        augment_fn (callable, optional): augments images with pre-defined visual effects
+    Returns:
+        image (torch.tensor): (1, h, w)
+        mask (torch.tensor): (h, w)
+        scale (torch.tensor): [w/w_new, h/h_new]        
+    """
+    # read image
+    image = imread_gray(path, augment_fn, client=MEGADEPTH_CLIENT)
+
+    # resize image
+    w, h = image.shape[1], image.shape[0]
+
+    if len(resize) == 2:
+        w_new, h_new = resize
+    else:
+        resize = resize[0]
+        w_new, h_new = get_resized_wh(w, h, resize)
+        w_new, h_new = get_divisible_wh(w_new, h_new, df)
+
+
+    image = cv2.resize(image, (w_new, h_new))
+    scale = torch.tensor([w/w_new, h/h_new], dtype=torch.float)
+
+    if padding:  # padding
+        pad_to = max(h_new, w_new)
+        image, mask = pad_bottom_right(image, pad_to, ret_mask=True)
+    else:
+        mask = None
+
+    #image = torch.from_numpy(image).float()[None] / 255  # (h, w) -> (1, h, w) and normalized
+    image = torch.from_numpy(image).float().permute(2,0,1) / 255  # (h, w) -> (1, h, w) and normalized
+    mask = torch.from_numpy(mask) if mask is not None else None
+
+    return image, mask, scale
+
+def imread_color(path, augment_fn=None, client=SCANNET_CLIENT):
+    cv_type = cv2.IMREAD_COLOR
+    # if str(path).startswith('s3://'):
+    #     image = load_array_from_s3(str(path), client, cv_type)
+    # else:
+    #     image = cv2.imread(str(path), cv_type)
+
+    image = cv2.imread(str(path), cv2.IMREAD_COLOR)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    if augment_fn is not None:
+        image = augment_fn(image)
+        # image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+    return image  # (3, h, w)
+
+
+def read_megadepth_color(path,
+                         resize=None,
+                         df=None,
+                         padding=False,
+                         augment_fn=None,
+                         rotation=0):
+    """
+    Args:
+        resize (int, optional): the longer edge of resized images. None for no resize.
+        padding (bool): If set to 'True', zero-pad resized images to squared size.
+        augment_fn (callable, optional): augments images with pre-defined visual effects
+    Returns:
+        image (torch.tensor): (3, h, w)
+        mask (torch.tensor): (h, w)
+        scale (torch.tensor): [w/w_new, h/h_new]
+    """
+    # read image
+    image = imread_color(path, augment_fn, client=MEGADEPTH_CLIENT)
+    
+    if rotation != 0:
+        image = np.rot90(image, k=rotation).copy()
+
+    # resize image
+    if resize is not None:
+        w, h = image.shape[1], image.shape[0]
+        if len(resize) == 2:
+            w_new, h_new = resize
+        else:
+            resize = resize[0]
+            w_new, h_new = get_resized_wh(w, h, resize)
+            w_new, h_new = get_divisible_wh(w_new, h_new, df)
+
+
+        image = cv2.resize(image, (w_new, h_new))
+        scale = torch.tensor([w/w_new, h/h_new], dtype=torch.float)
+        scale_wh = torch.tensor([w_new, h_new], dtype=torch.float)
+    else:
+        scale = torch.tensor([1., 1.], dtype=torch.float)
+        scale_wh = torch.tensor([image.shape[1], image.shape[0]], dtype=torch.float)
+        
+    image = image.transpose(2, 0, 1)
+    
+    if padding:  # padding
+        if resize is not None:
+            pad_to = max(h_new, w_new)
+        else:
+            pad_to = 2000
+        image, mask = pad_bottom_right(image, pad_to, ret_mask=True)
+    else:
+        mask = None
+
+    image = torch.from_numpy(image).float() / 255  # (h, w) -> (1, h, w) and normalized
+    mask = torch.from_numpy(mask) if mask is not None else None
+
+    return image, mask, scale
+
+def read_megadepth_depth(path, pad_to=None):
+
+    if str(path).startswith('s3://'):
+        depth = load_array_from_s3(path, MEGADEPTH_CLIENT, None, use_h5py=True)
+    else:
+        depth = np.array(h5py.File(path, 'r')['depth'])
+    if pad_to is not None:
+        depth, _ = pad_bottom_right(depth, pad_to, ret_mask=False)
+    depth = torch.from_numpy(depth).float()  # (h, w)
+    return depth
+
+def get_image_name(path):
+    return path.split('/')[-1]
+
+def scale_intrinsics(K, scales):
+    scales = np.diag([1. / scales[0], 1. / scales[1], 1.])
+    return np.dot(scales, K)

imcui/third_party/rdd/RDD/matchers/__init__.py

@@ -0,0 +1,3 @@
+from .dual_softmax_matcher import DualSoftmaxMatcher
+from .dense_matcher import DenseMatcher
+from .lightglue import LightGlue

imcui/third_party/rdd/RDD/matchers/dense_matcher.py

@@ -0,0 +1,88 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import poselib
+
+class DenseMatcher(nn.Module):
+    def __init__(self, inv_temperature = 20, thr = 0.01):
+        super().__init__()
+        self.inv_temperature = inv_temperature
+        self.thr = thr
+
+    def forward(self, info0, info1, thr = None, err_thr=4, min_num_inliers=30):
+        
+        desc0 = info0['descriptors']
+        desc1 = info1['descriptors']
+        
+        inds, P = self.dual_softmax(desc0, desc1, thr=thr)
+        
+        mkpts_0 = info0['keypoints'][inds[:,0]]
+        mkpts_1 = info1['keypoints'][inds[:,1]]
+        mconf = P[inds[:,0], inds[:,1]]
+        Fm, inliers = self.get_fundamental_matrix(mkpts_0, mkpts_1)
+        
+        if inliers.sum() >= min_num_inliers:
+            desc1_dense = info0['descriptors_dense']
+            desc2_dense = info1['descriptors_dense']
+
+            inds_dense, P_dense = self.dual_softmax(desc1_dense, desc2_dense, thr=thr)
+            
+            mkpts_0_dense = info0['keypoints_dense'][inds_dense[:,0]]
+            mkpts_1_dense = info1['keypoints_dense'][inds_dense[:,1]]
+            mconf_dense = P_dense[inds_dense[:,0], inds_dense[:,1]]
+            
+            mkpts_0_dense, mkpts_1_dense, mconf_dense = self.refine_matches(mkpts_0_dense, mkpts_1_dense, mconf_dense, Fm, err_thr=err_thr)
+            mkpts_0 = mkpts_0[inliers]
+            mkpts_1 = mkpts_1[inliers]
+            mconf = mconf[inliers]
+            # concatenate the matches
+            mkpts_0 = torch.cat([mkpts_0, mkpts_0_dense], dim=0)
+            mkpts_1 = torch.cat([mkpts_1, mkpts_1_dense], dim=0)
+            mconf = torch.cat([mconf, mconf_dense], dim=0)
+
+        return mkpts_0, mkpts_1, mconf
+
+    def get_fundamental_matrix(self, kpts_0, kpts_1):
+        Fm, info = poselib.estimate_fundamental(kpts_0.cpu().numpy(), kpts_1.cpu().numpy(), {'max_epipolar_error': 1, 'progressive_sampling': True}, {})
+        inliers = info['inliers']
+        Fm = torch.tensor(Fm, device=kpts_0.device, dtype=kpts_0.dtype)
+        inliers = torch.tensor(inliers, device=kpts_0.device, dtype=torch.bool)
+        return Fm, inliers    
+    
+    def dual_softmax(self, desc0, desc1, thr = None):
+        if thr is None:
+            thr = self.thr
+        dist_mat = (desc0 @ desc1.t()) * self.inv_temperature
+        P = dist_mat.softmax(dim = -2) * dist_mat.softmax(dim= -1)
+        inds = torch.nonzero((P == P.max(dim=-1, keepdim = True).values) 
+                        * (P == P.max(dim=-2, keepdim = True).values) * (P >= thr))
+        
+        return inds, P
+    
+    @torch.inference_mode()
+    def refine_matches(self, mkpts_0, mkpts_1, mconf, Fm, err_thr=4):    
+        mkpts_0_h = torch.cat([mkpts_0, torch.ones(mkpts_0.shape[0], 1, device=mkpts_0.device)], dim=1)  # (N, 3)
+        mkpts_1_h = torch.cat([mkpts_1, torch.ones(mkpts_1.shape[0], 1, device=mkpts_1.device)], dim=1)  # (N, 3)
+        
+        lines_1 = torch.matmul(Fm, mkpts_0_h.T).T
+        
+        a, b, c = lines_1[:, 0], lines_1[:, 1], lines_1[:, 2]  
+
+        x1, y1 = mkpts_1[:, 0], mkpts_1[:, 1]
+        
+        denom = a**2 + b**2 + 1e-8  
+        
+        x_offset = (b * (b * x1 - a * y1) - a * c) / denom - x1
+        y_offset = (a * (a * y1 - b * x1) - b * c) / denom - y1
+
+        inds = (x_offset.abs() < err_thr) | (y_offset.abs() < err_thr)
+
+        x_offset = x_offset[inds]
+        y_offset = y_offset[inds]
+
+        mkpts_0 = mkpts_0[inds]
+        mkpts_1 = mkpts_1[inds]
+        
+        refined_mkpts_1 = mkpts_1 + torch.stack([x_offset, y_offset], dim=1)
+  
+        return mkpts_0, refined_mkpts_1, mconf[inds]

imcui/third_party/rdd/RDD/matchers/dual_softmax_matcher.py

@@ -0,0 +1,31 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class DualSoftmaxMatcher(nn.Module):
+    def __init__(self, inv_temperature = 20, thr = 0.01):
+        super().__init__()
+        self.inv_temperature = inv_temperature
+        self.thr = thr
+
+    def forward(self, info0, info1, thr = None):
+        desc0 = info0['descriptors']
+        desc1 = info1['descriptors']
+        
+        inds, P = self.dual_softmax(desc0, desc1, thr)
+        mkpts0 = info0['keypoints'][inds[:,0]]
+        mkpts1 = info1['keypoints'][inds[:,1]]
+        mconf = P[inds[:,0], inds[:,1]]
+        
+        return mkpts0, mkpts1, mconf
+
+    def dual_softmax(self, desc0, desc1, thr = None):
+        if thr is None:
+            thr = self.thr
+        dist_mat = (desc0 @ desc1.t()) * self.inv_temperature
+        P = dist_mat.softmax(dim = -2) * dist_mat.softmax(dim= -1)
+        
+        inds = torch.nonzero((P == P.max(dim=-1, keepdim = True).values) 
+                        * (P == P.max(dim=-2, keepdim = True).values) * (P >= thr))
+        
+        return inds, P

imcui/third_party/rdd/RDD/matchers/lightglue.py

@@ -0,0 +1,667 @@
+"""
+Modified from 
+https://github.com/cvg/LightGlue
+"""
+
+import warnings
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Callable, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+try:
+    from flash_attn.modules.mha import FlashCrossAttention
+except ModuleNotFoundError:
+    FlashCrossAttention = None
+
+if FlashCrossAttention or hasattr(F, "scaled_dot_product_attention"):
+    FLASH_AVAILABLE = True
+else:
+    FLASH_AVAILABLE = False
+
+torch.backends.cudnn.deterministic = True
+
+
+@torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+def normalize_keypoints(
+    kpts: torch.Tensor, size: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    if size is None:
+        size = 1 + kpts.max(-2).values - kpts.min(-2).values
+    elif not isinstance(size, torch.Tensor):
+        size = torch.tensor(size, device=kpts.device, dtype=kpts.dtype)
+    size = size.to(kpts)
+    shift = size / 2
+    scale = size.max(-1).values / 2
+    kpts = (kpts - shift[..., None, :]) / scale[..., None, None]
+    return kpts
+
+
+def pad_to_length(x: torch.Tensor, length: int) -> Tuple[torch.Tensor]:
+    if length <= x.shape[-2]:
+        return x, torch.ones_like(x[..., :1], dtype=torch.bool)
+    pad = torch.ones(
+        *x.shape[:-2], length - x.shape[-2], x.shape[-1], device=x.device, dtype=x.dtype
+    )
+    y = torch.cat([x, pad], dim=-2)
+    mask = torch.zeros(*y.shape[:-1], 1, dtype=torch.bool, device=x.device)
+    mask[..., : x.shape[-2], :] = True
+    return y, mask
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x = x.unflatten(-1, (-1, 2))
+    x1, x2 = x.unbind(dim=-1)
+    return torch.stack((-x2, x1), dim=-1).flatten(start_dim=-2)
+
+
+def apply_cached_rotary_emb(freqs: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+    return (t * freqs[0]) + (rotate_half(t) * freqs[1])
+
+
+class LearnableFourierPositionalEncoding(nn.Module):
+    def __init__(self, M: int, dim: int, F_dim: int = None, gamma: float = 1.0) -> None:
+        super().__init__()
+        F_dim = F_dim if F_dim is not None else dim
+        self.gamma = gamma
+        self.Wr = nn.Linear(M, F_dim // 2, bias=False)
+        nn.init.normal_(self.Wr.weight.data, mean=0, std=self.gamma**-2)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """encode position vector"""
+        projected = self.Wr(x)
+        cosines, sines = torch.cos(projected), torch.sin(projected)
+        emb = torch.stack([cosines, sines], 0).unsqueeze(-3)
+        return emb.repeat_interleave(2, dim=-1)
+
+
+class TokenConfidence(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.token = nn.Sequential(nn.Linear(dim, 1), nn.Sigmoid())
+
+    def forward(self, desc0: torch.Tensor, desc1: torch.Tensor):
+        """get confidence tokens"""
+        return (
+            self.token(desc0.detach()).squeeze(-1),
+            self.token(desc1.detach()).squeeze(-1),
+        )
+
+
+class Attention(nn.Module):
+    def __init__(self, allow_flash: bool) -> None:
+        super().__init__()
+        if allow_flash and not FLASH_AVAILABLE:
+            warnings.warn(
+                "FlashAttention is not available. For optimal speed, "
+                "consider installing torch >= 2.0 or flash-attn.",
+                stacklevel=2,
+            )
+        self.enable_flash = allow_flash and FLASH_AVAILABLE
+        self.has_sdp = hasattr(F, "scaled_dot_product_attention")
+        if allow_flash and FlashCrossAttention:
+            self.flash_ = FlashCrossAttention()
+        if self.has_sdp:
+            torch.backends.cuda.enable_flash_sdp(allow_flash)
+
+    def forward(self, q, k, v, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if q.shape[-2] == 0 or k.shape[-2] == 0:
+            return q.new_zeros((*q.shape[:-1], v.shape[-1]))
+        if self.enable_flash and q.device.type == "cuda":
+            # use torch 2.0 scaled_dot_product_attention with flash
+            if self.has_sdp:
+                args = [x.half().contiguous() for x in [q, k, v]]
+                v = F.scaled_dot_product_attention(*args, attn_mask=mask).to(q.dtype)
+                return v if mask is None else v.nan_to_num()
+            else:
+                assert mask is None
+                q, k, v = [x.transpose(-2, -3).contiguous() for x in [q, k, v]]
+                m = self.flash_(q.half(), torch.stack([k, v], 2).half())
+                return m.transpose(-2, -3).to(q.dtype).clone()
+        elif self.has_sdp:
+            args = [x.contiguous() for x in [q, k, v]]
+            v = F.scaled_dot_product_attention(*args, attn_mask=mask)
+            return v if mask is None else v.nan_to_num()
+        else:
+            s = q.shape[-1] ** -0.5
+            sim = torch.einsum("...id,...jd->...ij", q, k) * s
+            if mask is not None:
+                sim.masked_fill(~mask, -float("inf"))
+            attn = F.softmax(sim, -1)
+            return torch.einsum("...ij,...jd->...id", attn, v)
+
+
+class SelfBlock(nn.Module):
+    def __init__(
+        self, embed_dim: int, num_heads: int, flash: bool = False, bias: bool = True
+    ) -> None:
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        assert self.embed_dim % num_heads == 0
+        self.head_dim = self.embed_dim // num_heads
+        self.Wqkv = nn.Linear(embed_dim, 3 * embed_dim, bias=bias)
+        self.inner_attn = Attention(flash)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.ffn = nn.Sequential(
+            nn.Linear(2 * embed_dim, 2 * embed_dim),
+            nn.LayerNorm(2 * embed_dim, elementwise_affine=True),
+            nn.GELU(),
+            nn.Linear(2 * embed_dim, embed_dim),
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        encoding: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        qkv = self.Wqkv(x)
+        qkv = qkv.unflatten(-1, (self.num_heads, -1, 3)).transpose(1, 2)
+        q, k, v = qkv[..., 0], qkv[..., 1], qkv[..., 2]
+        q = apply_cached_rotary_emb(encoding, q)
+        k = apply_cached_rotary_emb(encoding, k)
+        context = self.inner_attn(q, k, v, mask=mask)
+        message = self.out_proj(context.transpose(1, 2).flatten(start_dim=-2))
+        return x + self.ffn(torch.cat([x, message], -1))
+
+
+class CrossBlock(nn.Module):
+    def __init__(
+        self, embed_dim: int, num_heads: int, flash: bool = False, bias: bool = True
+    ) -> None:
+        super().__init__()
+        self.heads = num_heads
+        dim_head = embed_dim // num_heads
+        self.scale = dim_head**-0.5
+        inner_dim = dim_head * num_heads
+        self.to_qk = nn.Linear(embed_dim, inner_dim, bias=bias)
+        self.to_v = nn.Linear(embed_dim, inner_dim, bias=bias)
+        self.to_out = nn.Linear(inner_dim, embed_dim, bias=bias)
+        self.ffn = nn.Sequential(
+            nn.Linear(2 * embed_dim, 2 * embed_dim),
+            nn.LayerNorm(2 * embed_dim, elementwise_affine=True),
+            nn.GELU(),
+            nn.Linear(2 * embed_dim, embed_dim),
+        )
+        if flash and FLASH_AVAILABLE:
+            self.flash = Attention(True)
+        else:
+            self.flash = None
+
+    def map_(self, func: Callable, x0: torch.Tensor, x1: torch.Tensor):
+        return func(x0), func(x1)
+
+    def forward(
+        self, x0: torch.Tensor, x1: torch.Tensor, mask: Optional[torch.Tensor] = None
+    ) -> List[torch.Tensor]:
+        qk0, qk1 = self.map_(self.to_qk, x0, x1)
+        v0, v1 = self.map_(self.to_v, x0, x1)
+        qk0, qk1, v0, v1 = map(
+            lambda t: t.unflatten(-1, (self.heads, -1)).transpose(1, 2),
+            (qk0, qk1, v0, v1),
+        )
+        if self.flash is not None and qk0.device.type == "cuda":
+            m0 = self.flash(qk0, qk1, v1, mask)
+            m1 = self.flash(
+                qk1, qk0, v0, mask.transpose(-1, -2) if mask is not None else None
+            )
+        else:
+            qk0, qk1 = qk0 * self.scale**0.5, qk1 * self.scale**0.5
+            sim = torch.einsum("bhid, bhjd -> bhij", qk0, qk1)
+            if mask is not None:
+                sim = sim.masked_fill(~mask, -float("inf"))
+            attn01 = F.softmax(sim, dim=-1)
+            attn10 = F.softmax(sim.transpose(-2, -1).contiguous(), dim=-1)
+            m0 = torch.einsum("bhij, bhjd -> bhid", attn01, v1)
+            m1 = torch.einsum("bhji, bhjd -> bhid", attn10.transpose(-2, -1), v0)
+            if mask is not None:
+                m0, m1 = m0.nan_to_num(), m1.nan_to_num()
+        m0, m1 = self.map_(lambda t: t.transpose(1, 2).flatten(start_dim=-2), m0, m1)
+        m0, m1 = self.map_(self.to_out, m0, m1)
+        x0 = x0 + self.ffn(torch.cat([x0, m0], -1))
+        x1 = x1 + self.ffn(torch.cat([x1, m1], -1))
+        return x0, x1
+
+
+class TransformerLayer(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.self_attn = SelfBlock(*args, **kwargs)
+        self.cross_attn = CrossBlock(*args, **kwargs)
+
+    def forward(
+        self,
+        desc0,
+        desc1,
+        encoding0,
+        encoding1,
+        mask0: Optional[torch.Tensor] = None,
+        mask1: Optional[torch.Tensor] = None,
+    ):
+        if mask0 is not None and mask1 is not None:
+            return self.masked_forward(desc0, desc1, encoding0, encoding1, mask0, mask1)
+        else:
+            desc0 = self.self_attn(desc0, encoding0)
+            desc1 = self.self_attn(desc1, encoding1)
+            return self.cross_attn(desc0, desc1)
+
+    # This part is compiled and allows padding inputs
+    def masked_forward(self, desc0, desc1, encoding0, encoding1, mask0, mask1):
+        mask = mask0 & mask1.transpose(-1, -2)
+        mask0 = mask0 & mask0.transpose(-1, -2)
+        mask1 = mask1 & mask1.transpose(-1, -2)
+        desc0 = self.self_attn(desc0, encoding0, mask0)
+        desc1 = self.self_attn(desc1, encoding1, mask1)
+        return self.cross_attn(desc0, desc1, mask)
+
+
+def sigmoid_log_double_softmax(
+    sim: torch.Tensor, z0: torch.Tensor, z1: torch.Tensor
+) -> torch.Tensor:
+    """create the log assignment matrix from logits and similarity"""
+    b, m, n = sim.shape
+    certainties = F.logsigmoid(z0) + F.logsigmoid(z1).transpose(1, 2)
+    scores0 = F.log_softmax(sim, 2)
+    scores1 = F.log_softmax(sim.transpose(-1, -2).contiguous(), 2).transpose(-1, -2)
+    scores = sim.new_full((b, m + 1, n + 1), 0)
+    scores[:, :m, :n] = scores0 + scores1 + certainties
+    scores[:, :-1, -1] = F.logsigmoid(-z0.squeeze(-1))
+    scores[:, -1, :-1] = F.logsigmoid(-z1.squeeze(-1))
+    return scores
+
+
+class MatchAssignment(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.dim = dim
+        self.matchability = nn.Linear(dim, 1, bias=True)
+        self.final_proj = nn.Linear(dim, dim, bias=True)
+
+    def forward(self, desc0: torch.Tensor, desc1: torch.Tensor):
+        """build assignment matrix from descriptors"""
+        mdesc0, mdesc1 = self.final_proj(desc0), self.final_proj(desc1)
+        _, _, d = mdesc0.shape
+        mdesc0, mdesc1 = mdesc0 / d**0.25, mdesc1 / d**0.25
+        sim = torch.einsum("bmd,bnd->bmn", mdesc0, mdesc1)
+        z0 = self.matchability(desc0)
+        z1 = self.matchability(desc1)
+        scores = sigmoid_log_double_softmax(sim, z0, z1)
+        return scores, sim
+
+    def get_matchability(self, desc: torch.Tensor):
+        return torch.sigmoid(self.matchability(desc)).squeeze(-1)
+
+
+def filter_matches(scores: torch.Tensor, th: float):
+    """obtain matches from a log assignment matrix [Bx M+1 x N+1]"""
+    max0, max1 = scores[:, :-1, :-1].max(2), scores[:, :-1, :-1].max(1)
+    m0, m1 = max0.indices, max1.indices
+    indices0 = torch.arange(m0.shape[1], device=m0.device)[None]
+    indices1 = torch.arange(m1.shape[1], device=m1.device)[None]
+    mutual0 = indices0 == m1.gather(1, m0)
+    mutual1 = indices1 == m0.gather(1, m1)
+    max0_exp = max0.values.exp()
+    zero = max0_exp.new_tensor(0)
+    mscores0 = torch.where(mutual0, max0_exp, zero)
+    mscores1 = torch.where(mutual1, mscores0.gather(1, m1), zero)
+    valid0 = mutual0 & (mscores0 > th)
+    valid1 = mutual1 & valid0.gather(1, m1)
+    m0 = torch.where(valid0, m0, -1)
+    m1 = torch.where(valid1, m1, -1)
+    return m0, m1, mscores0, mscores1
+
+
+class LightGlue(nn.Module):
+    default_conf = {
+        "name": "lightglue",  # just for interfacing
+        "input_dim": 256,  # input descriptor dimension (autoselected from weights)
+        "descriptor_dim": 256,
+        "add_scale_ori": False,
+        "n_layers": 9,
+        "num_heads": 4,
+        "flash": True,  # enable FlashAttention if available.
+        "mp": False,  # enable mixed precision
+        "depth_confidence": -1,  # early stopping, disable with -1
+        "width_confidence": -1,  # point pruning, disable with -1
+        "filter_threshold": 0.01,  # match threshold
+        "weights": None,
+    }
+
+    # Point pruning involves an overhead (gather).
+    # Therefore, we only activate it if there are enough keypoints.
+    pruning_keypoint_thresholds = {
+        "cpu": -1,
+        "mps": -1,
+        "cuda": 1024,
+        "flash": 1536,
+    }
+
+    required_data_keys = ["image0", "image1"]
+
+    version = "v0.1_arxiv"
+    url = "https://github.com/cvg/LightGlue/releases/download/{}/{}_lightglue.pth"
+
+    features = {
+        "superpoint": {
+            "weights": "superpoint_lightglue",
+            "input_dim": 256,
+        },
+        "disk": {
+            "weights": "disk_lightglue",
+            "input_dim": 128,
+        },
+        "aliked": {
+            "weights": "aliked_lightglue",
+            "input_dim": 128,
+        },
+        "sift": {
+            "weights": "sift_lightglue",
+            "input_dim": 128,
+            "add_scale_ori": True,
+        },
+        "doghardnet": {
+            "weights": "doghardnet_lightglue",
+            "input_dim": 128,
+            "add_scale_ori": True,
+        },
+        "rdd": {
+            "weights": './weights/RDD_lg-v2.pth',
+            "input_dim": 256,
+        },
+    }
+
+    def __init__(self, features="rdd", **conf) -> None:
+        super().__init__()
+        self.conf = conf = SimpleNamespace(**{**self.default_conf, **conf})
+        if features is not None:
+            if features not in self.features:
+                raise ValueError(
+                    f"Unsupported features: {features} not in "
+                    f"{{{','.join(self.features)}}}"
+                )
+            for k, v in self.features[features].items():
+                setattr(conf, k, v)
+
+        if conf.input_dim != conf.descriptor_dim:
+            self.input_proj = nn.Linear(conf.input_dim, conf.descriptor_dim, bias=True)
+        else:
+            self.input_proj = nn.Identity()
+
+        head_dim = conf.descriptor_dim // conf.num_heads
+        self.posenc = LearnableFourierPositionalEncoding(
+            2 + 2 * self.conf.add_scale_ori, head_dim, head_dim
+        )
+
+        h, n, d = conf.num_heads, conf.n_layers, conf.descriptor_dim
+
+        self.transformers = nn.ModuleList(
+            [TransformerLayer(d, h, conf.flash) for _ in range(n)]
+        )
+
+        self.log_assignment = nn.ModuleList([MatchAssignment(d) for _ in range(n)])
+        self.token_confidence = nn.ModuleList(
+            [TokenConfidence(d) for _ in range(n - 1)]
+        )
+        self.register_buffer(
+            "confidence_thresholds",
+            torch.Tensor(
+                [self.confidence_threshold(i) for i in range(self.conf.n_layers)]
+            ),
+        )
+
+        state_dict = None
+        if features is not None and features != 'rdd':
+            fname = f"{conf.weights}_{self.version.replace('.', '-')}.pth"
+            state_dict = torch.hub.load_state_dict_from_url(
+                self.url.format(self.version, features), file_name=fname
+            )
+            self.load_state_dict(state_dict, strict=False)
+        elif conf.weights is not None:
+            if features == 'rdd':
+                path = Path(conf.weights)
+            else:
+                path = Path(__file__).parent
+                path = path / "weights/{}.pth".format(self.conf.weights)
+            state_dict = torch.load(str(path), map_location="cpu")
+
+        if state_dict:
+            # rename old state dict entries
+            for i in range(self.conf.n_layers):
+                pattern = f"self_attn.{i}", f"transformers.{i}.self_attn"
+                state_dict = {k.replace(*pattern): v for k, v in state_dict.items()}
+                pattern = f"cross_attn.{i}", f"transformers.{i}.cross_attn"
+                state_dict = {k.replace(*pattern): v for k, v in state_dict.items()}
+            self.load_state_dict(state_dict, strict=False)
+
+        # static lengths LightGlue is compiled for (only used with torch.compile)
+        self.static_lengths = None
+
+    def compile(
+        self, mode="reduce-overhead", static_lengths=[256, 512, 768, 1024, 1280, 1536]
+    ):
+        if self.conf.width_confidence != -1:
+            warnings.warn(
+                "Point pruning is partially disabled for compiled forward.",
+                stacklevel=2,
+            )
+
+        torch._inductor.cudagraph_mark_step_begin()
+        for i in range(self.conf.n_layers):
+            self.transformers[i].masked_forward = torch.compile(
+                self.transformers[i].masked_forward, mode=mode, fullgraph=True
+            )
+
+        self.static_lengths = static_lengths
+
+    def forward(self, data: dict) -> dict:
+        """
+        Match keypoints and descriptors between two images
+
+        Input (dict):
+            image0: dict
+                keypoints: [B x M x 2]
+                descriptors: [B x M x D]
+                image: [B x C x H x W] or image_size: [B x 2]
+            image1: dict
+                keypoints: [B x N x 2]
+                descriptors: [B x N x D]
+                image: [B x C x H x W] or image_size: [B x 2]
+        Output (dict):
+            matches0: [B x M]
+            matching_scores0: [B x M]
+            matches1: [B x N]
+            matching_scores1: [B x N]
+            matches: List[[Si x 2]]
+            scores: List[[Si]]
+            stop: int
+            prune0: [B x M]
+            prune1: [B x N]
+        """
+        with torch.autocast(enabled=self.conf.mp, device_type="cuda"):
+            return self._forward(data)
+
+    def _forward(self, data: dict) -> dict:
+        for key in self.required_data_keys:
+            assert key in data, f"Missing key {key} in data"
+        data0, data1 = data["image0"], data["image1"]
+        kpts0, kpts1 = data0["keypoints"], data1["keypoints"]
+        b, m, _ = kpts0.shape
+        b, n, _ = kpts1.shape
+        device = kpts0.device
+        size0, size1 = data0.get("image_size"), data1.get("image_size")
+        kpts0 = normalize_keypoints(kpts0, size0).clone()
+        kpts1 = normalize_keypoints(kpts1, size1).clone()
+
+        if self.conf.add_scale_ori:
+            kpts0 = torch.cat(
+                [kpts0] + [data0[k].unsqueeze(-1) for k in ("scales", "oris")], -1
+            )
+            kpts1 = torch.cat(
+                [kpts1] + [data1[k].unsqueeze(-1) for k in ("scales", "oris")], -1
+            )
+        desc0 = data0["descriptors"].detach().contiguous()
+        desc1 = data1["descriptors"].detach().contiguous()
+
+        assert desc0.shape[-1] == self.conf.input_dim
+        assert desc1.shape[-1] == self.conf.input_dim
+
+        if torch.is_autocast_enabled():
+            desc0 = desc0.half()
+            desc1 = desc1.half()
+
+        mask0, mask1 = None, None
+        c = max(m, n)
+        do_compile = self.static_lengths and c <= max(self.static_lengths)
+        if do_compile:
+            kn = min([k for k in self.static_lengths if k >= c])
+            desc0, mask0 = pad_to_length(desc0, kn)
+            desc1, mask1 = pad_to_length(desc1, kn)
+            kpts0, _ = pad_to_length(kpts0, kn)
+            kpts1, _ = pad_to_length(kpts1, kn)
+        desc0 = self.input_proj(desc0)
+        desc1 = self.input_proj(desc1)
+        # cache positional embeddings
+        encoding0 = self.posenc(kpts0)
+        encoding1 = self.posenc(kpts1)
+
+        # GNN + final_proj + assignment
+        do_early_stop = self.conf.depth_confidence > 0
+        do_point_pruning = self.conf.width_confidence > 0 and not do_compile
+        pruning_th = self.pruning_min_kpts(device)
+        if do_point_pruning:
+            ind0 = torch.arange(0, m, device=device)[None]
+            ind1 = torch.arange(0, n, device=device)[None]
+            # We store the index of the layer at which pruning is detected.
+            prune0 = torch.ones_like(ind0)
+            prune1 = torch.ones_like(ind1)
+        token0, token1 = None, None
+        for i in range(self.conf.n_layers):
+            if desc0.shape[1] == 0 or desc1.shape[1] == 0:  # no keypoints
+                break
+            desc0, desc1 = self.transformers[i](
+                desc0, desc1, encoding0, encoding1, mask0=mask0, mask1=mask1
+            )
+            if i == self.conf.n_layers - 1:
+                continue  # no early stopping or adaptive width at last layer
+
+            if do_early_stop:
+                token0, token1 = self.token_confidence[i](desc0, desc1)
+                if self.check_if_stop(token0[..., :m], token1[..., :n], i, m + n):
+                    break
+            if do_point_pruning and desc0.shape[-2] > pruning_th:
+                scores0 = self.log_assignment[i].get_matchability(desc0)
+                prunemask0 = self.get_pruning_mask(token0, scores0, i)
+                keep0 = torch.where(prunemask0)[1]
+                ind0 = ind0.index_select(1, keep0)
+                desc0 = desc0.index_select(1, keep0)
+                encoding0 = encoding0.index_select(-2, keep0)
+                prune0[:, ind0] += 1
+            if do_point_pruning and desc1.shape[-2] > pruning_th:
+                scores1 = self.log_assignment[i].get_matchability(desc1)
+                prunemask1 = self.get_pruning_mask(token1, scores1, i)
+                keep1 = torch.where(prunemask1)[1]
+                ind1 = ind1.index_select(1, keep1)
+                desc1 = desc1.index_select(1, keep1)
+                encoding1 = encoding1.index_select(-2, keep1)
+                prune1[:, ind1] += 1
+
+        if desc0.shape[1] == 0 or desc1.shape[1] == 0:  # no keypoints
+            m0 = desc0.new_full((b, m), -1, dtype=torch.long)
+            m1 = desc1.new_full((b, n), -1, dtype=torch.long)
+            mscores0 = desc0.new_zeros((b, m))
+            mscores1 = desc1.new_zeros((b, n))
+            matches = desc0.new_empty((b, 0, 2), dtype=torch.long)
+            mscores = desc0.new_empty((b, 0))
+            if not do_point_pruning:
+                prune0 = torch.ones_like(mscores0) * self.conf.n_layers
+                prune1 = torch.ones_like(mscores1) * self.conf.n_layers
+            return {
+                "matches0": m0,
+                "matches1": m1,
+                "matching_scores0": mscores0,
+                "matching_scores1": mscores1,
+                "stop": i + 1,
+                "matches": matches,
+                "scores": mscores,
+                "prune0": prune0,
+                "prune1": prune1,
+            }
+
+        desc0, desc1 = desc0[..., :m, :], desc1[..., :n, :]  # remove padding
+        scores, _ = self.log_assignment[i](desc0, desc1)
+        m0, m1, mscores0, mscores1 = filter_matches(scores, self.conf.filter_threshold)
+        matches, mscores = [], []
+        for k in range(b):
+            valid = m0[k] > -1
+            m_indices_0 = torch.where(valid)[0]
+            m_indices_1 = m0[k][valid]
+            if do_point_pruning:
+                m_indices_0 = ind0[k, m_indices_0]
+                m_indices_1 = ind1[k, m_indices_1]
+            matches.append(torch.stack([m_indices_0, m_indices_1], -1))
+            mscores.append(mscores0[k][valid])
+
+        # TODO: Remove when hloc switches to the compact format.
+        if do_point_pruning:
+            m0_ = torch.full((b, m), -1, device=m0.device, dtype=m0.dtype)
+            m1_ = torch.full((b, n), -1, device=m1.device, dtype=m1.dtype)
+            m0_[:, ind0] = torch.where(m0 == -1, -1, ind1.gather(1, m0.clamp(min=0)))
+            m1_[:, ind1] = torch.where(m1 == -1, -1, ind0.gather(1, m1.clamp(min=0)))
+            mscores0_ = torch.zeros((b, m), device=mscores0.device)
+            mscores1_ = torch.zeros((b, n), device=mscores1.device)
+            mscores0_[:, ind0] = mscores0
+            mscores1_[:, ind1] = mscores1
+            m0, m1, mscores0, mscores1 = m0_, m1_, mscores0_, mscores1_
+        else:
+            prune0 = torch.ones_like(mscores0) * self.conf.n_layers
+            prune1 = torch.ones_like(mscores1) * self.conf.n_layers
+
+        return {
+            "matches0": m0,
+            "matches1": m1,
+            "matching_scores0": mscores0,
+            "matching_scores1": mscores1,
+            "stop": i + 1,
+            "matches": matches,
+            "scores": mscores,
+            "prune0": prune0,
+            "prune1": prune1,
+        }
+
+    def confidence_threshold(self, layer_index: int) -> float:
+        """scaled confidence threshold"""
+        threshold = 0.8 + 0.1 * np.exp(-4.0 * layer_index / self.conf.n_layers)
+        return np.clip(threshold, 0, 1)
+
+    def get_pruning_mask(
+        self, confidences: torch.Tensor, scores: torch.Tensor, layer_index: int
+    ) -> torch.Tensor:
+        """mask points which should be removed"""
+        keep = scores > (1 - self.conf.width_confidence)
+        if confidences is not None:  # Low-confidence points are never pruned.
+            keep |= confidences <= self.confidence_thresholds[layer_index]
+        return keep
+
+    def check_if_stop(
+        self,
+        confidences0: torch.Tensor,
+        confidences1: torch.Tensor,
+        layer_index: int,
+        num_points: int,
+    ) -> torch.Tensor:
+        """evaluate stopping condition"""
+        confidences = torch.cat([confidences0, confidences1], -1)
+        threshold = self.confidence_thresholds[layer_index]
+        ratio_confident = 1.0 - (confidences < threshold).float().sum() / num_points
+        return ratio_confident > self.conf.depth_confidence
+
+    def pruning_min_kpts(self, device: torch.device):
+        if self.conf.flash and FLASH_AVAILABLE and device.type == "cuda":
+            return self.pruning_keypoint_thresholds["flash"]
+        else:
+            return self.pruning_keypoint_thresholds[device.type]

imcui/third_party/rdd/RDD/models/backbone.py

@@ -0,0 +1,147 @@
+# Modified from Deformable DETR
+# https://github.com/fundamentalvision/Deformable-DETR
+# ------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Modified from DETR (https://github.com/facebookresearch/detr)
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# ------------------------------------------------------------------------
+
+"""
+Backbone modules.
+"""
+from collections import OrderedDict
+
+import torch
+import torch.nn.functional as F
+import torchvision
+from torch import nn
+from torchvision.models._utils import IntermediateLayerGetter
+from typing import Dict, List
+import torch.distributed as dist
+from .position_encoding import build_position_encoding
+
+from ..utils.misc import NestedTensor, is_main_process
+
+class FrozenBatchNorm2d(torch.nn.Module):
+    """
+    BatchNorm2d where the batch statistics and the affine parameters are fixed.
+
+    Copy-paste from torchvision.misc.ops with added eps before rqsrt,
+    without which any other models than torchvision.models.resnet[18,34,50,101]
+    produce nans.
+    """
+
+    def __init__(self, n, eps=1e-5):
+        super(FrozenBatchNorm2d, self).__init__()
+        self.register_buffer("weight", torch.ones(n))
+        self.register_buffer("bias", torch.zeros(n))
+        self.register_buffer("running_mean", torch.zeros(n))
+        self.register_buffer("running_var", torch.ones(n))
+        self.eps = eps
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
+                              missing_keys, unexpected_keys, error_msgs):
+        num_batches_tracked_key = prefix + 'num_batches_tracked'
+        if num_batches_tracked_key in state_dict:
+            del state_dict[num_batches_tracked_key]
+
+        super(FrozenBatchNorm2d, self)._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict,
+            missing_keys, unexpected_keys, error_msgs)
+
+    def forward(self, x):
+        # move reshapes to the beginning
+        # to make it fuser-friendly
+        w = self.weight.reshape(1, -1, 1, 1)
+        b = self.bias.reshape(1, -1, 1, 1)
+        rv = self.running_var.reshape(1, -1, 1, 1)
+        rm = self.running_mean.reshape(1, -1, 1, 1)
+        eps = self.eps
+        scale = w * (rv + eps).rsqrt()
+        bias = b - rm * scale
+        return x * scale + bias
+
+
+class BackboneBase(nn.Module):
+
+    def __init__(self, backbone: nn.Module, train_backbone: bool, return_interm_layers: bool, n_layers = 4):
+        super().__init__()
+        for name, parameter in backbone.named_parameters():
+            if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:
+                parameter.requires_grad_(False)
+        if return_interm_layers:
+            if n_layers == 4:
+                return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"}
+                self.strides = [4, 8, 16, 32]
+                self.num_channels = [256, 512, 1024, 2048]
+            elif n_layers == 3:
+                return_layers = {"layer2": "0", "layer3": "1", "layer4": "2"}
+                self.strides = [8, 16, 32]
+                self.num_channels = [512, 1024, 2048]
+            else:
+                raise ValueError("n_layers should be 3 or 4")
+            
+        else:
+            return_layers = {'layer4': "0"}
+            self.strides = [32]
+            self.num_channels = [2048]
+        self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)
+
+    def forward(self, tensor_list: NestedTensor):
+        xs = self.body(tensor_list.tensors)
+        out: Dict[str, NestedTensor] = {}
+        for name, x in xs.items(): 
+            m = tensor_list.mask
+            assert m is not None
+            mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0]
+            out[name] = NestedTensor(x, mask)
+        return out
+
+
+class Backbone(BackboneBase):
+    """ResNet backbone with frozen BatchNorm."""
+    def __init__(self, name: str,
+                 train_backbone: bool,
+                 return_interm_layers: bool,
+                 dilation: bool,
+                 n_layers = 4):
+        norm_layer = FrozenBatchNorm2d
+        backbone = getattr(torchvision.models, name)(
+            replace_stride_with_dilation=[False, False, dilation],
+            weights='ResNet50_Weights.IMAGENET1K_V1', norm_layer=norm_layer)
+        assert name not in ('resnet18', 'resnet34'), "number of channels are hard coded"
+        super().__init__(backbone, train_backbone, return_interm_layers, n_layers)
+        if dilation:
+            self.strides[-1] = self.strides[-1] // 2
+
+class Joiner(nn.Sequential):
+    def __init__(self, backbone, position_embedding):
+        super().__init__(backbone, position_embedding)
+        self.strides = backbone.strides
+        self.num_channels = backbone.num_channels
+
+    def forward(self, tensor_list: NestedTensor):
+        xs = self[0](tensor_list)
+        out: List[NestedTensor] = []
+        pos = []
+        for name, x in sorted(xs.items()):
+            out.append(x)
+
+        # position encoding
+        for x in out:
+            pos.append(self[1](x).to(x.tensors.dtype))
+
+        return out, pos
+
+
+def build_backbone(config):
+    position_embedding = build_position_encoding(config)
+    train_backbone = config['lr_backbone'] > 0
+    return_interm_layers = True
+    n_layers = config['num_feature_levels'] - 1
+    backbone = Backbone('resnet50', train_backbone, return_interm_layers, False, n_layers=n_layers)
+    model = Joiner(backbone, position_embedding)
+    return model

imcui/third_party/rdd/RDD/models/deformable_transformer.py

@@ -0,0 +1,270 @@
+# Modified from Deformable DETR 
+# https://github.com/fundamentalvision/Deformable-DETR
+# ------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Modified from DETR (https://github.com/facebookresearch/detr)
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# ------------------------------------------------------------------------
+
+import copy
+from typing import Optional, List
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+from torch.nn.init import xavier_uniform_, constant_, uniform_, normal_
+
+from ..utils.misc import inverse_sigmoid
+from .ops.modules import MSDeformAttn
+
+class MLP(nn.Module):
+    """ Very simple multi-layer perceptron (also called FFN)"""
+
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        return x
+    
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+class DeformableTransformerEncoderLayer(nn.Module):
+    def __init__(self,
+                 d_model=256, d_ffn=1024,
+                 dropout=0.1, activation="relu",
+                 n_levels=4, n_heads=8, n_points=4):
+        super().__init__()
+
+        # self attention
+        self.self_attn = MSDeformAttn(d_model, n_levels, n_heads, n_points)
+        self.dropout1 = nn.Dropout(dropout)
+        self.norm1 = nn.LayerNorm(d_model)
+
+        # ffn
+        self.linear1 = nn.Linear(d_model, d_ffn)
+        self.activation = _get_activation_fn(activation)
+        self.dropout2 = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(d_ffn, d_model)
+        self.dropout3 = nn.Dropout(dropout)
+        self.norm2 = nn.LayerNorm(d_model)
+
+    @staticmethod
+    def with_pos_embed(tensor, pos):
+        return tensor if pos is None else tensor + pos
+
+    def forward_ffn(self, src):
+        src2 = self.linear2(self.dropout2(self.activation(self.linear1(src))))
+        src = src + self.dropout3(src2)
+        src = self.norm2(src)
+        return src
+
+    def forward(self, src, pos, reference_points, spatial_shapes, level_start_index, padding_mask=None, debug=False):
+        # self attention
+        if debug:
+            src2, sampled_points = self.self_attn(self.with_pos_embed(src, pos), reference_points, src, spatial_shapes, level_start_index, padding_mask)
+        else:
+            src2 = self.self_attn(self.with_pos_embed(src, pos), reference_points, src, spatial_shapes, level_start_index, padding_mask)
+        src = src + self.dropout1(src2)
+        src = self.norm1(src)
+
+        # ffn
+        src = self.forward_ffn(src)
+        if debug:
+            return src, sampled_points
+        return src
+
+class DeformableTransformerEncoder(nn.Module):
+    def __init__(self, encoder_layer, num_layers):
+        super().__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+
+    @staticmethod
+    def get_reference_points(spatial_shapes, valid_ratios, device):
+        reference_points_list = []
+        for lvl, (H_, W_) in enumerate(spatial_shapes):
+
+            ref_y, ref_x = torch.meshgrid(torch.linspace(0.5, H_ - 0.5, H_, dtype=torch.float32, device=device),
+                                          torch.linspace(0.5, W_ - 0.5, W_, dtype=torch.float32, device=device))
+            ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, lvl, 1] * H_)
+            ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, lvl, 0] * W_)
+            ref = torch.stack((ref_x, ref_y), -1)
+            reference_points_list.append(ref)
+        reference_points = torch.cat(reference_points_list, 1)
+        reference_points = reference_points[:, :, None] * valid_ratios[:, None]
+        return reference_points
+
+    def forward(self, src, spatial_shapes, level_start_index, valid_ratios, pos=None, padding_mask=None, debug=False):
+        output = src
+        reference_points = self.get_reference_points(spatial_shapes, valid_ratios, device=src.device)
+        for _, layer in enumerate(self.layers):
+            if debug:
+                output, sampled_points = layer(output, pos, reference_points, spatial_shapes, level_start_index, padding_mask, debug=debug)
+            else:
+                output = layer(output, pos, reference_points, spatial_shapes, level_start_index, padding_mask)
+        if debug:
+            return output, reference_points, sampled_points
+        return output
+
+class DecoderLayer(nn.Module):
+    def __init__(self, d_model=256, n_head=8, dropout=0.1):
+        super().__init__()
+        self.nhead = n_head
+        self.dim = d_model // n_head
+        self.attention = LinearAttention()
+        self.dropout1 = nn.Dropout(dropout)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.q_proj = nn.Linear(d_model, d_model, bias=False)
+        self.k_proj = nn.Linear(d_model, d_model, bias=False)
+        self.v_proj = nn.Linear(d_model, d_model, bias=False)
+        
+        self.mlp = nn.Sequential(
+            nn.Linear(d_model*2, d_model*2, bias=False),
+            nn.ReLU(True),
+            nn.Linear(d_model*2, d_model, bias=False),
+        )
+        
+        self.norm2 = nn.LayerNorm(d_model)
+        
+    def forward(self, tgt, src, tgt_mask=None, src_mask=None):
+        
+        bs = tgt.size(0)
+        query, key, value = tgt, src, src
+        
+        query = self.q_proj(query).view(bs, -1, self.nhead, self.dim)  # [N, L, (H, D)]
+        key = self.k_proj(key).view(bs, -1, self.nhead, self.dim)  # [N, S, (H, D)]
+        value = self.v_proj(value).view(bs, -1, self.nhead, self.dim)
+        
+        tgt2 = self.attention(query, key, value, q_mask=tgt_mask, kv_mask=src_mask)
+        tgt2 = tgt2.view(bs, -1, self.nhead*self.dim)
+        tgt2 = self.norm1(self.dropout1(tgt2))
+        tgt2 = self.mlp(torch.cat([tgt, tgt2], dim=2))
+        
+        tgt2 = self.norm2(tgt2)
+
+        return tgt + tgt2
+
+class MLP(nn.Module):
+    """ Very simple multi-layer perceptron (also called FFN)"""
+
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        return x
+    
+
+class Decoder(nn.Module):
+    def __init__(self, decoder_layer, num_layers):
+        super().__init__()
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+
+    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None):
+        for layer in self.layers:
+            tgt = layer(tgt, memory, tgt_mask=tgt_mask, src_mask=memory_mask)
+        
+        return tgt
+
+import math
+
+class DeformableTransformer(nn.Module):
+    def __init__(self, d_model=256, nhead=8,
+                 num_encoder_layers=4, dim_feedforward=1024, dropout=0.1,
+                 activation="relu",
+                 num_feature_levels=5,  enc_n_points=8):
+        super().__init__()
+
+        self.d_model = d_model
+        self.nhead = nhead
+
+        # Encoder and Decoder
+        encoder_layer = DeformableTransformerEncoderLayer(d_model, dim_feedforward,
+                                                          dropout, activation,
+                                                          num_feature_levels, nhead, enc_n_points)
+        self.encoder = DeformableTransformerEncoder(encoder_layer, num_encoder_layers)
+
+        # Embedding for feature levels (multi-scale)
+        self.level_embed = nn.Parameter(torch.Tensor(num_feature_levels, d_model))
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+        normal_(self.level_embed)
+
+    def get_valid_ratio(self, mask):
+        _, H, W = mask.shape
+        valid_H = torch.sum(~mask[:, :, 0], 1)
+        valid_W = torch.sum(~mask[:, 0, :], 1)
+        valid_ratio_h = valid_H.float() / H
+        valid_ratio_w = valid_W.float() / W
+        valid_ratio = torch.stack([valid_ratio_w, valid_ratio_h], -1)
+        return valid_ratio
+    
+    def forward(self, srcs, masks, pos_embeds):
+
+        # Prepare inputs for encoder
+        src_flatten = []
+        mask_flatten = []
+        lvl_pos_embed_flatten = []
+        spatial_shapes = []
+        for lvl, (src, mask, pos_embed) in enumerate(zip(srcs, masks, pos_embeds)):
+            bs, c, h, w = src.shape
+            spatial_shape = (h, w)
+            spatial_shapes.append(spatial_shape)
+            src = src.flatten(2).transpose(1, 2)
+            mask = mask.flatten(1)
+            pos_embed = pos_embed.flatten(2).transpose(1, 2).to(src.device)
+            lvl_pos_embed = pos_embed + self.level_embed[lvl].view(1, 1, -1)
+            lvl_pos_embed_flatten.append(lvl_pos_embed)
+            src_flatten.append(src)
+            mask_flatten.append(mask)
+        src_flatten = torch.cat(src_flatten, 1)
+        mask_flatten = torch.cat(mask_flatten, 1)
+        lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1)
+        spatial_shapes = torch.as_tensor(spatial_shapes, dtype=torch.long, device=src_flatten.device)
+        level_start_index = torch.cat((spatial_shapes.new_zeros((1, )), spatial_shapes.prod(1).cumsum(0)[:-1]))
+        valid_ratios = torch.stack([self.get_valid_ratio(m) for m in masks], 1)
+    
+        memory = self.encoder(src_flatten, spatial_shapes, level_start_index, valid_ratios, lvl_pos_embed_flatten, mask_flatten, debug=False)
+        
+        return memory, spatial_shapes, level_start_index
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def _get_activation_fn(activation):
+    """Return an activation function given a string"""
+    if activation == "relu":
+        return F.relu
+    if activation == "gelu":
+        return F.gelu
+    if activation == "glu":
+        return F.glu
+    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")
+
+
+def build_deforamble_transformer(config):
+    return DeformableTransformer(d_model=config['d_model'], nhead=config['nhead'],
+        num_encoder_layers=config['num_encoder_layers'], dim_feedforward=config['dim_feedforward'], dropout=config['dropout'],
+        activation=config['activation'],
+        num_feature_levels=config['num_feature_levels'], enc_n_points=config['enc_n_points'])

imcui/third_party/rdd/RDD/models/descriptor.py

@@ -0,0 +1,116 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from ..utils.misc import NestedTensor, nested_tensor_from_tensor_list
+import torchvision.transforms as transforms
+from .backbone import build_backbone
+from .deformable_transformer import build_deforamble_transformer
+
+class BasicLayer(nn.Module):
+	"""
+	  Basic Convolutional Layer: Conv2d -> BatchNorm -> ReLU
+	"""
+	def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, bias=False):
+		super().__init__()
+		self.layer = nn.Sequential(
+									  nn.Conv2d( in_channels, out_channels, kernel_size, padding = padding, stride=stride, dilation=dilation, bias = bias),
+									  nn.BatchNorm2d(out_channels, affine=False),
+									  nn.ReLU(inplace = False),
+									)
+
+	def forward(self, x):
+	  return self.layer(x)
+
+class RDD_Descriptor(nn.Module):
+    def __init__(self, backbone, transformer, num_feature_levels):
+        super().__init__()
+        self.transformer = transformer
+        self.hidden_dim = transformer.d_model
+        self.num_feature_levels = num_feature_levels
+        
+        self.matchibility_head = nn.Sequential(
+										BasicLayer(256, 128, 1, padding=0),
+										BasicLayer(128, 64, 1, padding=0),
+										nn.Conv2d (64, 1, 1),
+										nn.Sigmoid()
+									)
+        
+        if num_feature_levels > 1:
+            num_backbone_outs = len(backbone.strides)
+            input_proj_list = []
+            for _ in range(num_backbone_outs):
+                in_channels = backbone.num_channels[_]
+                input_proj_list.append(nn.Sequential(
+                    nn.Conv2d(in_channels, self.hidden_dim, kernel_size=1),
+                    nn.GroupNorm(32, self.hidden_dim),
+                ))
+            for _ in range(num_feature_levels - num_backbone_outs):
+                input_proj_list.append(nn.Sequential(
+                    nn.Conv2d(in_channels, self.hidden_dim, kernel_size=3, stride=2, padding=1),
+                    nn.GroupNorm(32, self.hidden_dim),
+                ))
+                in_channels = self.hidden_dim
+            self.input_proj = nn.ModuleList(input_proj_list)
+        else:
+            self.input_proj = nn.ModuleList([
+                nn.Sequential(
+                    nn.Conv2d(backbone.num_channels[0], self.hidden_dim, kernel_size=1),
+                    nn.GroupNorm(32, self.hidden_dim),
+                )])
+        self.backbone = backbone
+        self.stride = backbone.strides[0]
+        for proj in self.input_proj:
+            nn.init.xavier_uniform_(proj[0].weight, gain=1)
+            nn.init.constant_(proj[0].bias, 0)
+            
+    def forward(self, samples: NestedTensor):
+        
+        if not isinstance(samples, NestedTensor):
+            samples = nested_tensor_from_tensor_list(samples)
+        
+        features, pos = self.backbone(samples)
+
+        srcs = []
+        masks = []
+        for l, feat in enumerate(features):
+            src, mask = feat.decompose()
+            srcs.append(self.input_proj[l](src))
+            masks.append(mask)
+            assert mask is not None
+        if self.num_feature_levels > len(srcs):
+            _len_srcs = len(srcs)
+            for l in range(_len_srcs, self.num_feature_levels):
+                if l == _len_srcs:
+                    src = self.input_proj[l](features[-1].tensors)
+                else:
+                    src = self.input_proj[l](srcs[-1])
+                m = samples.mask
+                mask = F.interpolate(m[None].float(), size=src.shape[-2:]).to(torch.bool)[0]
+                pos_l = self.backbone[1](NestedTensor(src, mask)).to(src.dtype)
+                srcs.append(src)
+                masks.append(mask)
+                pos.append(pos_l)
+        
+        flatten_feats, spatial_shapes, level_start_index = self.transformer(srcs, masks, pos)
+        # Reshape the flattened features back to the original spatial shapes
+        feats = []
+        level_start_index = torch.cat((level_start_index, torch.tensor([flatten_feats.shape[1]+1]).to(level_start_index.device)))
+        for i, shape in enumerate(spatial_shapes):
+            assert len(shape) == 2
+            temp = flatten_feats[:, level_start_index[i] : level_start_index[i+1], :]
+            feats.append(temp.transpose(1, 2).view(-1, self.hidden_dim, *shape))
+        
+        # Sum up the features from different levels
+        final_feature = feats[0]
+        for feat in feats[1:]:
+            final_feature = final_feature + F.interpolate(feat, size=final_feature.shape[-2:], mode='bilinear', align_corners=True)
+        
+        matchibility = self.matchibility_head(final_feature)
+        
+        return final_feature, matchibility
+    
+    
+def build_descriptor(config):
+    backbone = build_backbone(config)
+    transformer = build_deforamble_transformer(config)
+    return RDD_Descriptor(backbone, transformer, config['num_feature_levels'])

imcui/third_party/rdd/RDD/models/detector.py

@@ -0,0 +1,141 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models import resnet
+from typing import Optional, Callable
+from ..utils.misc import NestedTensor
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels,
+                 gate: Optional[Callable[..., nn.Module]] = None,
+                 norm_layer: Optional[Callable[..., nn.Module]] = None):
+        super().__init__()
+        if gate is None:
+            self.gate = nn.ReLU(inplace=False)
+        else:
+            self.gate = gate
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self.conv1 = resnet.conv3x3(in_channels, out_channels)
+        self.bn1 = norm_layer(out_channels)
+        self.conv2 = resnet.conv3x3(out_channels, out_channels)
+        self.bn2 = norm_layer(out_channels)
+
+    def forward(self, x):
+        x = self.gate(self.bn1(self.conv1(x)))  # B x in_channels x H x W
+        x = self.gate(self.bn2(self.conv2(x)))  # B x out_channels x H x W
+        return x
+
+class ResBlock(nn.Module):
+    expansion: int = 1
+
+    def __init__(
+            self,
+            inplanes: int,
+            planes: int,
+            stride: int = 1,
+            downsample: Optional[nn.Module] = None,
+            groups: int = 1,
+            base_width: int = 64,
+            dilation: int = 1,
+            gate: Optional[Callable[..., nn.Module]] = None,
+            norm_layer: Optional[Callable[..., nn.Module]] = None
+    ) -> None:
+        super(ResBlock, self).__init__()
+        if gate is None:
+            self.gate = nn.ReLU(inplace=False)
+        else:
+            self.gate = gate
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError('ResBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in ResBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = resnet.conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.conv2 = resnet.conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.gate(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out = out + identity
+        out = self.gate(out)
+
+        return out
+    
+class RDD_detector(nn.Module):
+    def __init__(self, block_dims, hidden_dim=128):
+        super().__init__()
+        self.gate = nn.ReLU(inplace=False)
+        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.pool4 = nn.MaxPool2d(kernel_size=4, stride=4)
+        self.block1 = ConvBlock(3, block_dims[0], self.gate, nn.BatchNorm2d)
+        self.block2 = ResBlock(inplanes=block_dims[0], planes=block_dims[1], stride=1,
+                            downsample=nn.Conv2d(block_dims[0], block_dims[1], 1),
+                            gate=self.gate,
+                            norm_layer=nn.BatchNorm2d)
+        self.block3 = ResBlock(inplanes=block_dims[1], planes=block_dims[2], stride=1,
+                            downsample=nn.Conv2d(block_dims[1], block_dims[2], 1),
+                            gate=self.gate,
+                            norm_layer=nn.BatchNorm2d)
+        self.block4 = ResBlock(inplanes=block_dims[2], planes=block_dims[3], stride=1,
+                            downsample=nn.Conv2d(block_dims[2], block_dims[3], 1),
+                            gate=self.gate,
+                            norm_layer=nn.BatchNorm2d)
+
+        self.conv1 = resnet.conv1x1(block_dims[0], hidden_dim // 4)
+        self.conv2 = resnet.conv1x1(block_dims[1], hidden_dim // 4)
+        self.conv3 = resnet.conv1x1(block_dims[2], hidden_dim // 4)
+        self.conv4 = resnet.conv1x1(block_dims[3], hidden_dim // 4)
+        
+        self.upsample2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        self.upsample4 = nn.Upsample(scale_factor=4, mode='bilinear', align_corners=True)
+        self.upsample8 = nn.Upsample(scale_factor=8, mode='bilinear', align_corners=True)
+        self.upsample32 = nn.Upsample(scale_factor=32, mode='bilinear', align_corners=True)
+
+        self.convhead2 = nn.Sequential(
+                                resnet.conv1x1(hidden_dim, 1),
+                                nn.Sigmoid()
+        )
+
+    def forward(self, samples: NestedTensor):
+        x1 = self.block1(samples.tensors)
+        x2 = self.pool2(x1)
+        x2 = self.block2(x2)  # B x c2 x H/2 x W/2
+        x3 = self.pool4(x2)
+        x3 = self.block3(x3)  # B x c3 x H/8 x W/8
+        x4 = self.pool4(x3)
+        x4 = self.block4(x4)
+        
+        x1 = self.gate(self.conv1(x1))  # B x dim//4 x H x W
+        x2 = self.gate(self.conv2(x2))  # B x dim//4 x H//2 x W//2
+        x3 = self.gate(self.conv3(x3))  # B x dim//4 x H//8 x W//8
+        x4 = self.gate(self.conv4(x4))  # B x dim//4 x H//32 x W//32
+        
+        x2_up = self.upsample2(x2)
+        x3_up = self.upsample8(x3)
+        x4_up = self.upsample32(x4)
+        
+        x1234 = torch.cat([x1, x2_up, x3_up, x4_up], dim=1)
+        scoremap = self.convhead2(x1234)
+        
+        return scoremap
+    
+def build_detector(config):
+    block_dims = config['block_dims']
+    return RDD_detector(block_dims, block_dims[-1])

imcui/third_party/rdd/RDD/models/interpolator.py

@@ -0,0 +1,33 @@
+"""
+    "XFeat: Accelerated Features for Lightweight Image Matching, CVPR 2024."
+	https://www.verlab.dcc.ufmg.br/descriptors/xfeat_cvpr24/
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class InterpolateSparse2d(nn.Module):
+    """ Efficiently interpolate tensor at given sparse 2D positions. """ 
+    def __init__(self, mode = 'bilinear', align_corners = True): 
+        super().__init__()
+        self.mode = mode
+        self.align_corners = align_corners
+
+    def normgrid(self, x, H, W):
+        """ Normalize coords to [-1,1]. """
+        return 2. * (x/(torch.tensor([W-1, H-1], device = x.device, dtype = x.dtype))) - 1.
+
+    def forward(self, x, pos, H, W):
+        """
+        Input
+            x: [B, C, H, W] feature tensor
+            pos: [B, N, 2] tensor of positions
+            H, W: int, original resolution of input 2d positions -- used in normalization [-1,1]
+
+        Returns
+            [B, N, C] sampled channels at 2d positions
+        """
+        grid = self.normgrid(pos, H, W).unsqueeze(-2).to(x.dtype)
+        x = F.grid_sample(x, grid, mode = self.mode , align_corners = self.align_corners)
+        return x.permute(0,2,3,1).squeeze(-2)

imcui/third_party/rdd/RDD/models/ops/functions/__init__.py

@@ -0,0 +1,13 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+
+from .ms_deform_attn_func import MSDeformAttnFunction
+

imcui/third_party/rdd/RDD/models/ops/functions/ms_deform_attn_func.py

@@ -0,0 +1,72 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import torch
+import torch.nn.functional as F
+from torch.autograd import Function
+from torch.autograd.function import once_differentiable
+
+try:
+    import MultiScaleDeformableAttention as MSDA
+except ModuleNotFoundError as e:
+    info_string = (
+        "\n\nPlease compile MultiScaleDeformableAttention CUDA op with the following commands:\n"
+        "\t`cd mask2former/modeling/pixel_decoder/ops`\n"
+        "\t`sh make.sh`\n"
+    )
+    print(info_string)
+
+
+class MSDeformAttnFunction(Function):
+    @staticmethod
+    def forward(ctx, value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights, im2col_step):
+        ctx.im2col_step = im2col_step
+        output = MSDA.ms_deform_attn_forward(
+            value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights, ctx.im2col_step)
+        ctx.save_for_backward(value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights)
+        return output
+
+    @staticmethod
+    @once_differentiable
+    def backward(ctx, grad_output):
+        value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights = ctx.saved_tensors
+        grad_value, grad_sampling_loc, grad_attn_weight = \
+            MSDA.ms_deform_attn_backward(
+                value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights, grad_output, ctx.im2col_step)
+
+        return grad_value, None, None, grad_sampling_loc, grad_attn_weight, None
+
+
+def ms_deform_attn_core_pytorch(value, value_spatial_shapes, sampling_locations, attention_weights):
+    # for debug and test only,
+    # need to use cuda version instead
+    N_, S_, M_, D_ = value.shape
+    _, Lq_, M_, L_, P_, _ = sampling_locations.shape
+    value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1)
+    sampling_grids = 2 * sampling_locations - 1
+    sampling_value_list = []
+    for lid_, (H_, W_) in enumerate(value_spatial_shapes):
+        # N_, H_*W_, M_, D_ -> N_, H_*W_, M_*D_ -> N_, M_*D_, H_*W_ -> N_*M_, D_, H_, W_
+        value_l_ = value_list[lid_].flatten(2).transpose(1, 2).reshape(N_*M_, D_, H_, W_)
+        # N_, Lq_, M_, P_, 2 -> N_, M_, Lq_, P_, 2 -> N_*M_, Lq_, P_, 2
+        sampling_grid_l_ = sampling_grids[:, :, :, lid_].transpose(1, 2).flatten(0, 1)
+        # N_*M_, D_, Lq_, P_
+        sampling_value_l_ = F.grid_sample(value_l_, sampling_grid_l_,
+                                          mode='bilinear', padding_mode='zeros', align_corners=False)
+        sampling_value_list.append(sampling_value_l_)
+    # (N_, Lq_, M_, L_, P_) -> (N_, M_, Lq_, L_, P_) -> (N_, M_, 1, Lq_, L_*P_)
+    attention_weights = attention_weights.transpose(1, 2).reshape(N_*M_, 1, Lq_, L_*P_)
+    output = (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights).sum(-1).view(N_, M_*D_, Lq_)
+    return output.transpose(1, 2).contiguous()

imcui/third_party/rdd/RDD/models/ops/make.sh

@@ -0,0 +1,13 @@
+#!/usr/bin/env bash
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+
+python setup.py build install

imcui/third_party/rdd/RDD/models/ops/modules/__init__.py

@@ -0,0 +1,12 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+
+from .ms_deform_attn import MSDeformAttn

imcui/third_party/rdd/RDD/models/ops/modules/ms_deform_attn.py

@@ -0,0 +1,125 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import warnings
+import math
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn.init import xavier_uniform_, constant_
+
+from ..functions import MSDeformAttnFunction
+from ..functions.ms_deform_attn_func import ms_deform_attn_core_pytorch
+
+
+def _is_power_of_2(n):
+    if (not isinstance(n, int)) or (n < 0):
+        raise ValueError("invalid input for _is_power_of_2: {} (type: {})".format(n, type(n)))
+    return (n & (n-1) == 0) and n != 0
+
+
+class MSDeformAttn(nn.Module):
+    def __init__(self, d_model=256, n_levels=4, n_heads=8, n_points=4):
+        """
+        Multi-Scale Deformable Attention Module
+        :param d_model      hidden dimension
+        :param n_levels     number of feature levels
+        :param n_heads      number of attention heads
+        :param n_points     number of sampling points per attention head per feature level
+        """
+        super().__init__()
+        if d_model % n_heads != 0:
+            raise ValueError('d_model must be divisible by n_heads, but got {} and {}'.format(d_model, n_heads))
+        _d_per_head = d_model // n_heads
+        # you'd better set _d_per_head to a power of 2 which is more efficient in our CUDA implementation
+        if not _is_power_of_2(_d_per_head):
+            warnings.warn("You'd better set d_model in MSDeformAttn to make the dimension of each attention head a power of 2 "
+                          "which is more efficient in our CUDA implementation.")
+
+        self.im2col_step = 128
+
+        self.d_model = d_model
+        self.n_levels = n_levels
+        self.n_heads = n_heads
+        self.n_points = n_points
+
+        self.sampling_offsets = nn.Linear(d_model, n_heads * n_levels * n_points * 2)
+        self.attention_weights = nn.Linear(d_model, n_heads * n_levels * n_points)
+        self.value_proj = nn.Linear(d_model, d_model)
+        self.output_proj = nn.Linear(d_model, d_model)
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        constant_(self.sampling_offsets.weight.data, 0.)
+        thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads)
+        grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
+        grid_init = (grid_init / grid_init.abs().max(-1, keepdim=True)[0]).view(self.n_heads, 1, 1, 2).repeat(1, self.n_levels, self.n_points, 1)
+        for i in range(self.n_points):
+            grid_init[:, :, i, :] *= i + 1
+        with torch.no_grad():
+            self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))
+        constant_(self.attention_weights.weight.data, 0.)
+        constant_(self.attention_weights.bias.data, 0.)
+        xavier_uniform_(self.value_proj.weight.data)
+        constant_(self.value_proj.bias.data, 0.)
+        xavier_uniform_(self.output_proj.weight.data)
+        constant_(self.output_proj.bias.data, 0.)
+
+    def forward(self, query, reference_points, input_flatten, input_spatial_shapes, input_level_start_index, input_padding_mask=None):
+        """
+        :param query                       (N, Length_{query}, C)
+        :param reference_points            (N, Length_{query}, n_levels, 2), range in [0, 1], top-left (0,0), bottom-right (1, 1), including padding area
+                                        or (N, Length_{query}, n_levels, 4), add additional (w, h) to form reference boxes
+        :param input_flatten               (N, \sum_{l=0}^{L-1} H_l \cdot W_l, C)
+        :param input_spatial_shapes        (n_levels, 2), [(H_0, W_0), (H_1, W_1), ..., (H_{L-1}, W_{L-1})]
+        :param input_level_start_index     (n_levels, ), [0, H_0*W_0, H_0*W_0+H_1*W_1, H_0*W_0+H_1*W_1+H_2*W_2, ..., H_0*W_0+H_1*W_1+...+H_{L-1}*W_{L-1}]
+        :param input_padding_mask          (N, \sum_{l=0}^{L-1} H_l \cdot W_l), True for padding elements, False for non-padding elements
+
+        :return output                     (N, Length_{query}, C)
+        """
+        N, Len_q, _ = query.shape
+        N, Len_in, _ = input_flatten.shape
+        assert (input_spatial_shapes[:, 0] * input_spatial_shapes[:, 1]).sum() == Len_in
+
+        value = self.value_proj(input_flatten)
+        if input_padding_mask is not None:
+            value = value.masked_fill(input_padding_mask[..., None], float(0))
+        value = value.view(N, Len_in, self.n_heads, self.d_model // self.n_heads)
+        sampling_offsets = self.sampling_offsets(query).view(N, Len_q, self.n_heads, self.n_levels, self.n_points, 2)
+        attention_weights = self.attention_weights(query).view(N, Len_q, self.n_heads, self.n_levels * self.n_points)
+        attention_weights = F.softmax(attention_weights, -1).view(N, Len_q, self.n_heads, self.n_levels, self.n_points)
+        # N, Len_q, n_heads, n_levels, n_points, 2
+        if reference_points.shape[-1] == 2:
+            offset_normalizer = torch.stack([input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1)
+            sampling_locations = reference_points[:, :, None, :, None, :] \
+                                 + sampling_offsets / offset_normalizer[None, None, None, :, None, :]
+        elif reference_points.shape[-1] == 4:
+            sampling_locations = reference_points[:, :, None, :, None, :2] \
+                                 + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5
+        else:
+            raise ValueError(
+                'Last dim of reference_points must be 2 or 4, but get {} instead.'.format(reference_points.shape[-1]))
+        try:
+            output = MSDeformAttnFunction.apply(
+                value, input_spatial_shapes, input_level_start_index, sampling_locations, attention_weights, self.im2col_step)
+        except:
+            # CPU
+            output = ms_deform_attn_core_pytorch(value, input_spatial_shapes, sampling_locations, attention_weights)
+        # # For FLOPs calculation only
+        # output = ms_deform_attn_core_pytorch(value, input_spatial_shapes, sampling_locations, attention_weights)
+        output = self.output_proj(output)
+        return output

imcui/third_party/rdd/RDD/models/ops/setup.py

@@ -0,0 +1,78 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+
+import os
+import glob
+
+import torch
+
+from torch.utils.cpp_extension import CUDA_HOME
+from torch.utils.cpp_extension import CppExtension
+from torch.utils.cpp_extension import CUDAExtension
+
+from setuptools import find_packages
+from setuptools import setup
+
+requirements = ["torch", "torchvision"]
+
+def get_extensions():
+    this_dir = os.path.dirname(os.path.abspath(__file__))
+    extensions_dir = os.path.join(this_dir, "src")
+
+    main_file = glob.glob(os.path.join(extensions_dir, "*.cpp"))
+    source_cpu = glob.glob(os.path.join(extensions_dir, "cpu", "*.cpp"))
+    source_cuda = glob.glob(os.path.join(extensions_dir, "cuda", "*.cu"))
+
+    sources = main_file + source_cpu
+    extension = CppExtension
+    extra_compile_args = {"cxx": []}
+    define_macros = []
+
+    # Force cuda since torch ask for a device, not if cuda is in fact available.
+    if (os.environ.get('FORCE_CUDA') or torch.cuda.is_available()) and CUDA_HOME is not None:
+        extension = CUDAExtension
+        sources += source_cuda
+        define_macros += [("WITH_CUDA", None)]
+        extra_compile_args["nvcc"] = [
+            "-DCUDA_HAS_FP16=1",
+            "-D__CUDA_NO_HALF_OPERATORS__",
+            "-D__CUDA_NO_HALF_CONVERSIONS__",
+            "-D__CUDA_NO_HALF2_OPERATORS__",
+        ]
+    else:
+        if CUDA_HOME is None:
+            raise NotImplementedError('CUDA_HOME is None. Please set environment variable CUDA_HOME.')
+        else:
+            raise NotImplementedError('No CUDA runtime is found. Please set FORCE_CUDA=1 or test it by running torch.cuda.is_available().')
+
+    sources = [os.path.join(extensions_dir, s) for s in sources]
+    include_dirs = [extensions_dir]
+    ext_modules = [
+        extension(
+            "MultiScaleDeformableAttention",
+            sources,
+            include_dirs=include_dirs,
+            define_macros=define_macros,
+            extra_compile_args=extra_compile_args,
+        )
+    ]
+    return ext_modules
+
+setup(
+    name="MultiScaleDeformableAttention",
+    version="1.0",
+    author="Weijie Su",
+    url="https://github.com/fundamentalvision/Deformable-DETR",
+    description="PyTorch Wrapper for CUDA Functions of Multi-Scale Deformable Attention",
+    packages=find_packages(exclude=("configs", "tests",)),
+    ext_modules=get_extensions(),
+    cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension},
+)

imcui/third_party/rdd/RDD/models/ops/src/cpu/ms_deform_attn_cpu.cpp

@@ -0,0 +1,46 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+/*!
+* Copyright (c) Facebook, Inc. and its affiliates.
+* Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+*/
+
+#include <vector>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+
+at::Tensor
+ms_deform_attn_cpu_forward(
+    const at::Tensor &value, 
+    const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index,
+    const at::Tensor &sampling_loc,
+    const at::Tensor &attn_weight,
+    const int im2col_step)
+{
+    AT_ERROR("Not implement on cpu");
+}
+
+std::vector<at::Tensor>
+ms_deform_attn_cpu_backward(
+    const at::Tensor &value, 
+    const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index,
+    const at::Tensor &sampling_loc,
+    const at::Tensor &attn_weight,
+    const at::Tensor &grad_output,
+    const int im2col_step)
+{
+    AT_ERROR("Not implement on cpu");
+}
+

imcui/third_party/rdd/RDD/models/ops/src/cpu/ms_deform_attn_cpu.h

@@ -0,0 +1,38 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+/*!
+* Copyright (c) Facebook, Inc. and its affiliates.
+* Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+*/
+
+#pragma once
+#include <torch/extension.h>
+
+at::Tensor
+ms_deform_attn_cpu_forward(
+    const at::Tensor &value, 
+    const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index,
+    const at::Tensor &sampling_loc,
+    const at::Tensor &attn_weight,
+    const int im2col_step);
+
+std::vector<at::Tensor>
+ms_deform_attn_cpu_backward(
+    const at::Tensor &value, 
+    const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index,
+    const at::Tensor &sampling_loc,
+    const at::Tensor &attn_weight,
+    const at::Tensor &grad_output,
+    const int im2col_step);
+
+

imcui/third_party/rdd/RDD/models/ops/src/cuda/ms_deform_attn_cuda.cu

@@ -0,0 +1,158 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+/*!
+* Copyright (c) Facebook, Inc. and its affiliates.
+* Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+*/
+
+#include <vector>
+#include "cuda/ms_deform_im2col_cuda.cuh"
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+
+at::Tensor ms_deform_attn_cuda_forward(
+    const at::Tensor &value, 
+    const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index,
+    const at::Tensor &sampling_loc,
+    const at::Tensor &attn_weight,
+    const int im2col_step)
+{
+    AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous");
+    AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous");
+    AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous");
+    AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous");
+    AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous");
+
+    AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor");
+    AT_ASSERTM(spatial_shapes.type().is_cuda(), "spatial_shapes must be a CUDA tensor");
+    AT_ASSERTM(level_start_index.type().is_cuda(), "level_start_index must be a CUDA tensor");
+    AT_ASSERTM(sampling_loc.type().is_cuda(), "sampling_loc must be a CUDA tensor");
+    AT_ASSERTM(attn_weight.type().is_cuda(), "attn_weight must be a CUDA tensor");
+
+    const int batch = value.size(0);
+    const int spatial_size = value.size(1);
+    const int num_heads = value.size(2);
+    const int channels = value.size(3);
+
+    const int num_levels = spatial_shapes.size(0);
+
+    const int num_query = sampling_loc.size(1);
+    const int num_point = sampling_loc.size(4);
+
+    const int im2col_step_ = std::min(batch, im2col_step);
+
+    AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_);
+    
+    auto output = at::zeros({batch, num_query, num_heads, channels}, value.options());
+
+    const int batch_n = im2col_step_;
+    auto output_n = output.view({batch/im2col_step_, batch_n, num_query, num_heads, channels});
+    auto per_value_size = spatial_size * num_heads * channels;
+    auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2;
+    auto per_attn_weight_size = num_query * num_heads * num_levels * num_point;
+    for (int n = 0; n < batch/im2col_step_; ++n)
+    {
+        auto columns = output_n.select(0, n);
+        AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_forward_cuda", ([&] {
+            ms_deformable_im2col_cuda(at::cuda::getCurrentCUDAStream(),
+                value.data<scalar_t>() + n * im2col_step_ * per_value_size,
+                spatial_shapes.data<int64_t>(),
+                level_start_index.data<int64_t>(),
+                sampling_loc.data<scalar_t>() + n * im2col_step_ * per_sample_loc_size,
+                attn_weight.data<scalar_t>() + n * im2col_step_ * per_attn_weight_size,
+                batch_n, spatial_size, num_heads, channels, num_levels, num_query, num_point,
+                columns.data<scalar_t>());
+
+        }));
+    }
+
+    output = output.view({batch, num_query, num_heads*channels});
+
+    return output;
+}
+
+
+std::vector<at::Tensor> ms_deform_attn_cuda_backward(
+    const at::Tensor &value, 
+    const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index,
+    const at::Tensor &sampling_loc,
+    const at::Tensor &attn_weight,
+    const at::Tensor &grad_output,
+    const int im2col_step)
+{
+
+    AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous");
+    AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous");
+    AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous");
+    AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous");
+    AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous");
+    AT_ASSERTM(grad_output.is_contiguous(), "grad_output tensor has to be contiguous");
+
+    AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor");
+    AT_ASSERTM(spatial_shapes.type().is_cuda(), "spatial_shapes must be a CUDA tensor");
+    AT_ASSERTM(level_start_index.type().is_cuda(), "level_start_index must be a CUDA tensor");
+    AT_ASSERTM(sampling_loc.type().is_cuda(), "sampling_loc must be a CUDA tensor");
+    AT_ASSERTM(attn_weight.type().is_cuda(), "attn_weight must be a CUDA tensor");
+    AT_ASSERTM(grad_output.type().is_cuda(), "grad_output must be a CUDA tensor");
+
+    const int batch = value.size(0);
+    const int spatial_size = value.size(1);
+    const int num_heads = value.size(2);
+    const int channels = value.size(3);
+
+    const int num_levels = spatial_shapes.size(0);
+
+    const int num_query = sampling_loc.size(1);
+    const int num_point = sampling_loc.size(4);
+
+    const int im2col_step_ = std::min(batch, im2col_step);
+
+    AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_);
+
+    auto grad_value = at::zeros_like(value);
+    auto grad_sampling_loc = at::zeros_like(sampling_loc);
+    auto grad_attn_weight = at::zeros_like(attn_weight);
+
+    const int batch_n = im2col_step_;
+    auto per_value_size = spatial_size * num_heads * channels;
+    auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2;
+    auto per_attn_weight_size = num_query * num_heads * num_levels * num_point;
+    auto grad_output_n = grad_output.view({batch/im2col_step_, batch_n, num_query, num_heads, channels});
+    
+    for (int n = 0; n < batch/im2col_step_; ++n)
+    {
+        auto grad_output_g = grad_output_n.select(0, n);
+        AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_backward_cuda", ([&] {
+            ms_deformable_col2im_cuda(at::cuda::getCurrentCUDAStream(),
+                                    grad_output_g.data<scalar_t>(),
+                                    value.data<scalar_t>() + n * im2col_step_ * per_value_size,
+                                    spatial_shapes.data<int64_t>(),
+                                    level_start_index.data<int64_t>(),
+                                    sampling_loc.data<scalar_t>() + n * im2col_step_ * per_sample_loc_size,
+                                    attn_weight.data<scalar_t>() + n * im2col_step_ * per_attn_weight_size,
+                                    batch_n, spatial_size, num_heads, channels, num_levels, num_query, num_point,
+                                    grad_value.data<scalar_t>() +  n * im2col_step_ * per_value_size,
+                                    grad_sampling_loc.data<scalar_t>() + n * im2col_step_ * per_sample_loc_size,
+                                    grad_attn_weight.data<scalar_t>() + n * im2col_step_ * per_attn_weight_size);
+
+        }));
+    }
+
+    return {
+        grad_value, grad_sampling_loc, grad_attn_weight
+    };
+}

imcui/third_party/rdd/RDD/models/ops/src/cuda/ms_deform_attn_cuda.h

@@ -0,0 +1,35 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+/*!
+* Copyright (c) Facebook, Inc. and its affiliates.
+* Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+*/
+
+#pragma once
+#include <torch/extension.h>
+
+at::Tensor ms_deform_attn_cuda_forward(
+    const at::Tensor &value, 
+    const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index,
+    const at::Tensor &sampling_loc,
+    const at::Tensor &attn_weight,
+    const int im2col_step);
+
+std::vector<at::Tensor> ms_deform_attn_cuda_backward(
+    const at::Tensor &value, 
+    const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index,
+    const at::Tensor &sampling_loc,
+    const at::Tensor &attn_weight,
+    const at::Tensor &grad_output,
+    const int im2col_step);
+

imcui/third_party/rdd/RDD/models/ops/src/cuda/ms_deform_im2col_cuda.cuh

@@ -0,0 +1,1332 @@
+/*!
+**************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************
+* Modified from DCN (https://github.com/msracver/Deformable-ConvNets)
+* Copyright (c) 2018 Microsoft
+**************************************************************************
+*/
+
+/*!
+* Copyright (c) Facebook, Inc. and its affiliates.
+* Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+*/
+
+#include <cstdio>
+#include <algorithm>
+#include <cstring>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <THC/THCAtomics.cuh>
+
+#define CUDA_KERNEL_LOOP(i, n)                          \
+  for (int i = blockIdx.x * blockDim.x + threadIdx.x;   \
+      i < (n);                                          \
+      i += blockDim.x * gridDim.x)
+
+const int CUDA_NUM_THREADS = 1024;
+inline int GET_BLOCKS(const int N, const int num_threads)
+{
+  return (N + num_threads - 1) / num_threads;
+}
+
+
+template <typename scalar_t>
+__device__ scalar_t ms_deform_attn_im2col_bilinear(const scalar_t* &bottom_data, 
+                                                   const int &height, const int &width, const int &nheads, const int &channels,
+                                                   const scalar_t &h, const scalar_t &w, const int &m, const int &c)
+{
+  const int h_low = floor(h);
+  const int w_low = floor(w);
+  const int h_high = h_low + 1;
+  const int w_high = w_low + 1;
+
+  const scalar_t lh = h - h_low;
+  const scalar_t lw = w - w_low;
+  const scalar_t hh = 1 - lh, hw = 1 - lw;
+
+  const int w_stride = nheads * channels;
+  const int h_stride = width * w_stride;
+  const int h_low_ptr_offset = h_low * h_stride;
+  const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
+  const int w_low_ptr_offset = w_low * w_stride;
+  const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
+  const int base_ptr = m * channels + c;
+
+  scalar_t v1 = 0;
+  if (h_low >= 0 && w_low >= 0)
+  {
+    const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
+    v1 = bottom_data[ptr1];
+  }
+  scalar_t v2 = 0;
+  if (h_low >= 0 && w_high <= width - 1)
+  {
+    const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
+    v2 = bottom_data[ptr2];
+  }
+  scalar_t v3 = 0;
+  if (h_high <= height - 1 && w_low >= 0)
+  {
+    const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
+    v3 = bottom_data[ptr3];
+  }
+  scalar_t v4 = 0;
+  if (h_high <= height - 1 && w_high <= width - 1)
+  {
+    const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
+    v4 = bottom_data[ptr4];
+  }
+
+  const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+
+  const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+  return val;
+}
+
+
+template <typename scalar_t>
+__device__ void ms_deform_attn_col2im_bilinear(const scalar_t* &bottom_data, 
+                                                   const int &height, const int &width, const int &nheads, const int &channels,
+                                                   const scalar_t &h, const scalar_t &w, const int &m, const int &c,
+                                                   const scalar_t &top_grad,
+                                                   const scalar_t &attn_weight,
+                                                   scalar_t* &grad_value, 
+                                                   scalar_t* grad_sampling_loc,
+                                                   scalar_t* grad_attn_weight)
+{
+  const int h_low = floor(h);
+  const int w_low = floor(w);
+  const int h_high = h_low + 1;
+  const int w_high = w_low + 1;
+
+  const scalar_t lh = h - h_low;
+  const scalar_t lw = w - w_low;
+  const scalar_t hh = 1 - lh, hw = 1 - lw;
+
+  const int w_stride = nheads * channels;
+  const int h_stride = width * w_stride;
+  const int h_low_ptr_offset = h_low * h_stride;
+  const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
+  const int w_low_ptr_offset = w_low * w_stride;
+  const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
+  const int base_ptr = m * channels + c;
+
+  const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+  const scalar_t top_grad_value = top_grad * attn_weight;
+  scalar_t grad_h_weight = 0, grad_w_weight = 0;
+
+  scalar_t v1 = 0;
+  if (h_low >= 0 && w_low >= 0)
+  {
+    const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
+    v1 = bottom_data[ptr1];
+    grad_h_weight -= hw * v1;
+    grad_w_weight -= hh * v1;
+    atomicAdd(grad_value+ptr1, w1*top_grad_value);
+  }
+  scalar_t v2 = 0;
+  if (h_low >= 0 && w_high <= width - 1)
+  {
+    const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
+    v2 = bottom_data[ptr2];
+    grad_h_weight -= lw * v2;
+    grad_w_weight += hh * v2;
+    atomicAdd(grad_value+ptr2, w2*top_grad_value);
+  }
+  scalar_t v3 = 0;
+  if (h_high <= height - 1 && w_low >= 0)
+  {
+    const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
+    v3 = bottom_data[ptr3];
+    grad_h_weight += hw * v3;
+    grad_w_weight -= lh * v3;
+    atomicAdd(grad_value+ptr3, w3*top_grad_value); 
+  }
+  scalar_t v4 = 0;
+  if (h_high <= height - 1 && w_high <= width - 1)
+  {
+    const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
+    v4 = bottom_data[ptr4];
+    grad_h_weight += lw * v4;
+    grad_w_weight += lh * v4;
+    atomicAdd(grad_value+ptr4, w4*top_grad_value);
+  }
+
+  const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+  *grad_attn_weight = top_grad * val;
+  *grad_sampling_loc = width * grad_w_weight * top_grad_value;
+  *(grad_sampling_loc + 1) = height * grad_h_weight * top_grad_value;
+}
+
+
+template <typename scalar_t>
+__device__ void ms_deform_attn_col2im_bilinear_gm(const scalar_t* &bottom_data, 
+                                                   const int &height, const int &width, const int &nheads, const int &channels,
+                                                   const scalar_t &h, const scalar_t &w, const int &m, const int &c,
+                                                   const scalar_t &top_grad,
+                                                   const scalar_t &attn_weight,
+                                                   scalar_t* &grad_value, 
+                                                   scalar_t* grad_sampling_loc,
+                                                   scalar_t* grad_attn_weight)
+{
+  const int h_low = floor(h);
+  const int w_low = floor(w);
+  const int h_high = h_low + 1;
+  const int w_high = w_low + 1;
+
+  const scalar_t lh = h - h_low;
+  const scalar_t lw = w - w_low;
+  const scalar_t hh = 1 - lh, hw = 1 - lw;
+
+  const int w_stride = nheads * channels;
+  const int h_stride = width * w_stride;
+  const int h_low_ptr_offset = h_low * h_stride;
+  const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
+  const int w_low_ptr_offset = w_low * w_stride;
+  const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
+  const int base_ptr = m * channels + c;
+
+  const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+  const scalar_t top_grad_value = top_grad * attn_weight;
+  scalar_t grad_h_weight = 0, grad_w_weight = 0;
+
+  scalar_t v1 = 0;
+  if (h_low >= 0 && w_low >= 0)
+  {
+    const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
+    v1 = bottom_data[ptr1];
+    grad_h_weight -= hw * v1;
+    grad_w_weight -= hh * v1;
+    atomicAdd(grad_value+ptr1, w1*top_grad_value);
+  }
+  scalar_t v2 = 0;
+  if (h_low >= 0 && w_high <= width - 1)
+  {
+    const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
+    v2 = bottom_data[ptr2];
+    grad_h_weight -= lw * v2;
+    grad_w_weight += hh * v2;
+    atomicAdd(grad_value+ptr2, w2*top_grad_value);
+  }
+  scalar_t v3 = 0;
+  if (h_high <= height - 1 && w_low >= 0)
+  {
+    const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
+    v3 = bottom_data[ptr3];
+    grad_h_weight += hw * v3;
+    grad_w_weight -= lh * v3;
+    atomicAdd(grad_value+ptr3, w3*top_grad_value); 
+  }
+  scalar_t v4 = 0;
+  if (h_high <= height - 1 && w_high <= width - 1)
+  {
+    const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
+    v4 = bottom_data[ptr4];
+    grad_h_weight += lw * v4;
+    grad_w_weight += lh * v4;
+    atomicAdd(grad_value+ptr4, w4*top_grad_value);
+  }
+
+  const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+  atomicAdd(grad_attn_weight, top_grad * val); 
+  atomicAdd(grad_sampling_loc, width * grad_w_weight * top_grad_value);
+  atomicAdd(grad_sampling_loc + 1, height * grad_h_weight * top_grad_value);
+}
+
+
+template <typename scalar_t>
+__global__ void ms_deformable_im2col_gpu_kernel(const int n,
+                                                const scalar_t *data_value, 
+                                                const int64_t *data_spatial_shapes,
+                                                const int64_t *data_level_start_index, 
+                                                const scalar_t *data_sampling_loc,
+                                                const scalar_t *data_attn_weight,
+                                                const int batch_size, 
+                                                const int spatial_size, 
+                                                const int num_heads,
+                                                const int channels, 
+                                                const int num_levels,
+                                                const int num_query,
+                                                const int num_point,
+                                                scalar_t *data_col)
+{
+  CUDA_KERNEL_LOOP(index, n)
+  {
+    int _temp = index;
+    const int c_col = _temp % channels;
+    _temp /= channels;
+    const int sampling_index = _temp; 
+    const int m_col = _temp % num_heads;
+    _temp /= num_heads;
+    const int q_col = _temp % num_query;
+    _temp /= num_query;
+    const int b_col = _temp;
+
+    scalar_t *data_col_ptr = data_col + index;
+    int data_weight_ptr = sampling_index * num_levels * num_point;
+    int data_loc_w_ptr = data_weight_ptr << 1;
+    const int qid_stride = num_heads * channels;
+    const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+    scalar_t col = 0;
+    
+    for (int l_col=0; l_col < num_levels; ++l_col)
+    {
+      const int level_start_id = data_level_start_index[l_col];
+      const int spatial_h_ptr = l_col << 1;
+      const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+      const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+      const scalar_t *data_value_ptr = data_value + (data_value_ptr_init_offset + level_start_id * qid_stride);
+      for (int p_col=0; p_col < num_point; ++p_col)
+      {
+        const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+        const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+        const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+        const scalar_t h_im = loc_h * spatial_h - 0.5;
+        const scalar_t w_im = loc_w * spatial_w - 0.5;
+
+        if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)
+        {
+          col += ms_deform_attn_im2col_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col) * weight;
+        }
+
+        data_weight_ptr += 1;
+        data_loc_w_ptr += 2;
+      }
+    }
+    *data_col_ptr = col;
+  }
+}
+
+template <typename scalar_t, unsigned int blockSize>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1(const int n,
+                                                const scalar_t *grad_col,
+                                                const scalar_t *data_value,
+                                                const int64_t *data_spatial_shapes,
+                                                const int64_t *data_level_start_index, 
+                                                const scalar_t *data_sampling_loc,
+                                                const scalar_t *data_attn_weight,
+                                                const int batch_size, 
+                                                const int spatial_size, 
+                                                const int num_heads,
+                                                const int channels, 
+                                                const int num_levels,
+                                                const int num_query,
+                                                const int num_point,
+                                                scalar_t *grad_value,
+                                                scalar_t *grad_sampling_loc,
+                                                scalar_t *grad_attn_weight)
+{
+  CUDA_KERNEL_LOOP(index, n)
+  {
+    __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2];
+    __shared__ scalar_t cache_grad_attn_weight[blockSize];
+    unsigned int tid = threadIdx.x;
+    int _temp = index;
+    const int c_col = _temp % channels;
+    _temp /= channels;
+    const int sampling_index = _temp; 
+    const int m_col = _temp % num_heads;
+    _temp /= num_heads;
+    const int q_col = _temp % num_query;
+    _temp /= num_query;
+    const int b_col = _temp;
+
+    const scalar_t top_grad = grad_col[index];
+
+    int data_weight_ptr = sampling_index * num_levels * num_point;
+    int data_loc_w_ptr = data_weight_ptr << 1;
+    const int grad_sampling_ptr = data_weight_ptr;
+    grad_sampling_loc += grad_sampling_ptr << 1;
+    grad_attn_weight += grad_sampling_ptr;
+    const int grad_weight_stride = 1;
+    const int grad_loc_stride = 2;
+    const int qid_stride = num_heads * channels;
+    const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+    for (int l_col=0; l_col < num_levels; ++l_col)
+    {
+      const int level_start_id = data_level_start_index[l_col];
+      const int spatial_h_ptr = l_col << 1;
+      const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+      const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+      const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+      const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+      scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+      for (int p_col=0; p_col < num_point; ++p_col)
+      {
+        const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+        const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+        const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+        const scalar_t h_im = loc_h * spatial_h - 0.5;
+        const scalar_t w_im = loc_w * spatial_w - 0.5;
+        *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0;
+        *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0;
+        *(cache_grad_attn_weight+threadIdx.x)=0;
+        if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)
+        {
+          ms_deform_attn_col2im_bilinear(
+            data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,
+            top_grad, weight, grad_value_ptr, 
+            cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x);
+        }
+        
+        __syncthreads();
+        if (tid == 0)
+        {
+          scalar_t _grad_w=cache_grad_sampling_loc[0], _grad_h=cache_grad_sampling_loc[1], _grad_a=cache_grad_attn_weight[0];
+          int sid=2;
+          for (unsigned int tid = 1; tid < blockSize; ++tid)
+          {
+            _grad_w += cache_grad_sampling_loc[sid];
+            _grad_h += cache_grad_sampling_loc[sid + 1];
+            _grad_a += cache_grad_attn_weight[tid];
+            sid += 2;
+          }
+          
+          
+          *grad_sampling_loc = _grad_w;
+          *(grad_sampling_loc + 1) = _grad_h;
+          *grad_attn_weight = _grad_a;
+        }
+        __syncthreads();
+
+        data_weight_ptr += 1;
+        data_loc_w_ptr += 2;
+        grad_attn_weight += grad_weight_stride;
+        grad_sampling_loc += grad_loc_stride;
+      }
+    }
+  }
+}
+
+
+template <typename scalar_t, unsigned int blockSize>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2(const int n,
+                                                const scalar_t *grad_col,
+                                                const scalar_t *data_value,
+                                                const int64_t *data_spatial_shapes,
+                                                const int64_t *data_level_start_index, 
+                                                const scalar_t *data_sampling_loc,
+                                                const scalar_t *data_attn_weight,
+                                                const int batch_size, 
+                                                const int spatial_size, 
+                                                const int num_heads,
+                                                const int channels, 
+                                                const int num_levels,
+                                                const int num_query,
+                                                const int num_point,
+                                                scalar_t *grad_value,
+                                                scalar_t *grad_sampling_loc,
+                                                scalar_t *grad_attn_weight)
+{
+  CUDA_KERNEL_LOOP(index, n)
+  {
+    __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2];
+    __shared__ scalar_t cache_grad_attn_weight[blockSize];
+    unsigned int tid = threadIdx.x;
+    int _temp = index;
+    const int c_col = _temp % channels;
+    _temp /= channels;
+    const int sampling_index = _temp; 
+    const int m_col = _temp % num_heads;
+    _temp /= num_heads;
+    const int q_col = _temp % num_query;
+    _temp /= num_query;
+    const int b_col = _temp;
+
+    const scalar_t top_grad = grad_col[index];
+
+    int data_weight_ptr = sampling_index * num_levels * num_point;
+    int data_loc_w_ptr = data_weight_ptr << 1;
+    const int grad_sampling_ptr = data_weight_ptr;
+    grad_sampling_loc += grad_sampling_ptr << 1;
+    grad_attn_weight += grad_sampling_ptr;
+    const int grad_weight_stride = 1;
+    const int grad_loc_stride = 2;
+    const int qid_stride = num_heads * channels;
+    const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+    for (int l_col=0; l_col < num_levels; ++l_col)
+    {
+      const int level_start_id = data_level_start_index[l_col];
+      const int spatial_h_ptr = l_col << 1;
+      const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+      const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+      const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+      const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+      scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+      for (int p_col=0; p_col < num_point; ++p_col)
+      {
+        const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+        const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+        const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+        const scalar_t h_im = loc_h * spatial_h - 0.5;
+        const scalar_t w_im = loc_w * spatial_w - 0.5;
+        *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0;
+        *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0;
+        *(cache_grad_attn_weight+threadIdx.x)=0;
+        if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)
+        {
+          ms_deform_attn_col2im_bilinear(
+            data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,
+            top_grad, weight, grad_value_ptr, 
+            cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x);
+        }
+        
+        __syncthreads();
+
+        for (unsigned int s=blockSize/2; s>0; s>>=1)
+        {
+          if (tid < s) {
+            const unsigned int xid1 = tid << 1;
+            const unsigned int xid2 = (tid + s) << 1;
+            cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];
+            cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];
+            cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];
+          }
+          __syncthreads();
+        }
+
+        if (tid == 0)
+        { 
+          *grad_sampling_loc = cache_grad_sampling_loc[0];
+          *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1];
+          *grad_attn_weight = cache_grad_attn_weight[0];
+        }
+        __syncthreads();
+
+        data_weight_ptr += 1;
+        data_loc_w_ptr += 2;
+        grad_attn_weight += grad_weight_stride;
+        grad_sampling_loc += grad_loc_stride;
+      }
+    }
+  }
+}
+
+
+template <typename scalar_t>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v1(const int n,
+                                                const scalar_t *grad_col,
+                                                const scalar_t *data_value,
+                                                const int64_t *data_spatial_shapes,
+                                                const int64_t *data_level_start_index, 
+                                                const scalar_t *data_sampling_loc,
+                                                const scalar_t *data_attn_weight,
+                                                const int batch_size, 
+                                                const int spatial_size, 
+                                                const int num_heads,
+                                                const int channels, 
+                                                const int num_levels,
+                                                const int num_query,
+                                                const int num_point,
+                                                scalar_t *grad_value,
+                                                scalar_t *grad_sampling_loc,
+                                                scalar_t *grad_attn_weight)
+{
+  CUDA_KERNEL_LOOP(index, n)
+  {
+    extern __shared__ int _s[];
+    scalar_t* cache_grad_sampling_loc = (scalar_t*)_s;
+    scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;
+    unsigned int tid = threadIdx.x;
+    int _temp = index;
+    const int c_col = _temp % channels;
+    _temp /= channels;
+    const int sampling_index = _temp; 
+    const int m_col = _temp % num_heads;
+    _temp /= num_heads;
+    const int q_col = _temp % num_query;
+    _temp /= num_query;
+    const int b_col = _temp;
+
+    const scalar_t top_grad = grad_col[index];
+
+    int data_weight_ptr = sampling_index * num_levels * num_point;
+    int data_loc_w_ptr = data_weight_ptr << 1;
+    const int grad_sampling_ptr = data_weight_ptr;
+    grad_sampling_loc += grad_sampling_ptr << 1;
+    grad_attn_weight += grad_sampling_ptr;
+    const int grad_weight_stride = 1;
+    const int grad_loc_stride = 2;
+    const int qid_stride = num_heads * channels;
+    const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+    for (int l_col=0; l_col < num_levels; ++l_col)
+    {
+      const int level_start_id = data_level_start_index[l_col];
+      const int spatial_h_ptr = l_col << 1;
+      const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+      const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+      const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+      const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+      scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+      for (int p_col=0; p_col < num_point; ++p_col)
+      {
+        const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+        const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+        const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+        const scalar_t h_im = loc_h * spatial_h - 0.5;
+        const scalar_t w_im = loc_w * spatial_w - 0.5;
+        *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0;
+        *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0;
+        *(cache_grad_attn_weight+threadIdx.x)=0;
+        if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)
+        {
+          ms_deform_attn_col2im_bilinear(
+            data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,
+            top_grad, weight, grad_value_ptr, 
+            cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x);
+        }
+        
+        __syncthreads();
+        if (tid == 0)
+        {
+          scalar_t _grad_w=cache_grad_sampling_loc[0], _grad_h=cache_grad_sampling_loc[1], _grad_a=cache_grad_attn_weight[0];
+          int sid=2;
+          for (unsigned int tid = 1; tid < blockDim.x; ++tid)
+          {
+            _grad_w += cache_grad_sampling_loc[sid];
+            _grad_h += cache_grad_sampling_loc[sid + 1];
+            _grad_a += cache_grad_attn_weight[tid];
+            sid += 2;
+          }
+          
+          
+          *grad_sampling_loc = _grad_w;
+          *(grad_sampling_loc + 1) = _grad_h;
+          *grad_attn_weight = _grad_a;
+        }
+        __syncthreads();
+
+        data_weight_ptr += 1;
+        data_loc_w_ptr += 2;
+        grad_attn_weight += grad_weight_stride;
+        grad_sampling_loc += grad_loc_stride;
+      }
+    }
+  }
+}
+
+template <typename scalar_t>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2(const int n,
+                                                const scalar_t *grad_col,
+                                                const scalar_t *data_value,
+                                                const int64_t *data_spatial_shapes,
+                                                const int64_t *data_level_start_index, 
+                                                const scalar_t *data_sampling_loc,
+                                                const scalar_t *data_attn_weight,
+                                                const int batch_size, 
+                                                const int spatial_size, 
+                                                const int num_heads,
+                                                const int channels, 
+                                                const int num_levels,
+                                                const int num_query,
+                                                const int num_point,
+                                                scalar_t *grad_value,
+                                                scalar_t *grad_sampling_loc,
+                                                scalar_t *grad_attn_weight)
+{
+  CUDA_KERNEL_LOOP(index, n)
+  {
+    extern __shared__ int _s[];
+    scalar_t* cache_grad_sampling_loc = (scalar_t*)_s;
+    scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;
+    unsigned int tid = threadIdx.x;
+    int _temp = index;
+    const int c_col = _temp % channels;
+    _temp /= channels;
+    const int sampling_index = _temp; 
+    const int m_col = _temp % num_heads;
+    _temp /= num_heads;
+    const int q_col = _temp % num_query;
+    _temp /= num_query;
+    const int b_col = _temp;
+
+    const scalar_t top_grad = grad_col[index];
+
+    int data_weight_ptr = sampling_index * num_levels * num_point;
+    int data_loc_w_ptr = data_weight_ptr << 1;
+    const int grad_sampling_ptr = data_weight_ptr;
+    grad_sampling_loc += grad_sampling_ptr << 1;
+    grad_attn_weight += grad_sampling_ptr;
+    const int grad_weight_stride = 1;
+    const int grad_loc_stride = 2;
+    const int qid_stride = num_heads * channels;
+    const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+    for (int l_col=0; l_col < num_levels; ++l_col)
+    {
+      const int level_start_id = data_level_start_index[l_col];
+      const int spatial_h_ptr = l_col << 1;
+      const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+      const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+      const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+      const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+      scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+      for (int p_col=0; p_col < num_point; ++p_col)
+      {
+        const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+        const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+        const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+        const scalar_t h_im = loc_h * spatial_h - 0.5;
+        const scalar_t w_im = loc_w * spatial_w - 0.5;
+        *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0;
+        *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0;
+        *(cache_grad_attn_weight+threadIdx.x)=0;
+        if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)
+        {
+          ms_deform_attn_col2im_bilinear(
+            data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,
+            top_grad, weight, grad_value_ptr, 
+            cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x);
+        }
+        
+        __syncthreads();
+
+        for (unsigned int s=blockDim.x/2, spre=blockDim.x; s>0; s>>=1, spre>>=1)
+        {
+          if (tid < s) {
+            const unsigned int xid1 = tid << 1;
+            const unsigned int xid2 = (tid + s) << 1;
+            cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];
+            cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];
+            cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];
+            if (tid + (s << 1) < spre)
+            {
+              cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)];
+              cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)];
+              cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)];
+            } 
+          }
+          __syncthreads();
+        }
+
+        if (tid == 0)
+        {
+          *grad_sampling_loc = cache_grad_sampling_loc[0];
+          *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1];
+          *grad_attn_weight = cache_grad_attn_weight[0];
+        }
+        __syncthreads();
+
+        data_weight_ptr += 1;
+        data_loc_w_ptr += 2;
+        grad_attn_weight += grad_weight_stride;
+        grad_sampling_loc += grad_loc_stride;
+      }
+    }
+  }
+}
+
+template <typename scalar_t>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks(const int n,
+                                                const scalar_t *grad_col,
+                                                const scalar_t *data_value,
+                                                const int64_t *data_spatial_shapes,
+                                                const int64_t *data_level_start_index, 
+                                                const scalar_t *data_sampling_loc,
+                                                const scalar_t *data_attn_weight,
+                                                const int batch_size, 
+                                                const int spatial_size, 
+                                                const int num_heads,
+                                                const int channels, 
+                                                const int num_levels,
+                                                const int num_query,
+                                                const int num_point,
+                                                scalar_t *grad_value,
+                                                scalar_t *grad_sampling_loc,
+                                                scalar_t *grad_attn_weight)
+{
+  CUDA_KERNEL_LOOP(index, n)
+  {
+    extern __shared__ int _s[];
+    scalar_t* cache_grad_sampling_loc = (scalar_t*)_s;
+    scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;
+    unsigned int tid = threadIdx.x;
+    int _temp = index;
+    const int c_col = _temp % channels;
+    _temp /= channels;
+    const int sampling_index = _temp; 
+    const int m_col = _temp % num_heads;
+    _temp /= num_heads;
+    const int q_col = _temp % num_query;
+    _temp /= num_query;
+    const int b_col = _temp;
+
+    const scalar_t top_grad = grad_col[index];
+
+    int data_weight_ptr = sampling_index * num_levels * num_point;
+    int data_loc_w_ptr = data_weight_ptr << 1;
+    const int grad_sampling_ptr = data_weight_ptr;
+    grad_sampling_loc += grad_sampling_ptr << 1;
+    grad_attn_weight += grad_sampling_ptr;
+    const int grad_weight_stride = 1;
+    const int grad_loc_stride = 2;
+    const int qid_stride = num_heads * channels;
+    const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+    for (int l_col=0; l_col < num_levels; ++l_col)
+    {
+      const int level_start_id = data_level_start_index[l_col];
+      const int spatial_h_ptr = l_col << 1;
+      const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+      const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+      const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+      const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+      scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+      for (int p_col=0; p_col < num_point; ++p_col)
+      {
+        const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+        const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+        const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+        const scalar_t h_im = loc_h * spatial_h - 0.5;
+        const scalar_t w_im = loc_w * spatial_w - 0.5;
+        *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0;
+        *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0;
+        *(cache_grad_attn_weight+threadIdx.x)=0;
+        if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)
+        {
+          ms_deform_attn_col2im_bilinear(
+            data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,
+            top_grad, weight, grad_value_ptr, 
+            cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x);
+        }
+        
+        __syncthreads();
+
+        for (unsigned int s=blockDim.x/2, spre=blockDim.x; s>0; s>>=1, spre>>=1)
+        {
+          if (tid < s) {
+            const unsigned int xid1 = tid << 1;
+            const unsigned int xid2 = (tid + s) << 1;
+            cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];
+            cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];
+            cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];
+            if (tid + (s << 1) < spre)
+            {
+              cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)];
+              cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)];
+              cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)];
+            }
+          }
+          __syncthreads();
+        }
+
+        if (tid == 0)
+        {
+          atomicAdd(grad_sampling_loc, cache_grad_sampling_loc[0]);
+          atomicAdd(grad_sampling_loc + 1, cache_grad_sampling_loc[1]);
+          atomicAdd(grad_attn_weight, cache_grad_attn_weight[0]);
+        }
+        __syncthreads();
+
+        data_weight_ptr += 1;
+        data_loc_w_ptr += 2;
+        grad_attn_weight += grad_weight_stride;
+        grad_sampling_loc += grad_loc_stride;
+      }
+    }
+  }
+}
+
+
+template <typename scalar_t>
+__global__ void ms_deformable_col2im_gpu_kernel_gm(const int n,
+                                                const scalar_t *grad_col,
+                                                const scalar_t *data_value,
+                                                const int64_t *data_spatial_shapes,
+                                                const int64_t *data_level_start_index, 
+                                                const scalar_t *data_sampling_loc,
+                                                const scalar_t *data_attn_weight,
+                                                const int batch_size, 
+                                                const int spatial_size, 
+                                                const int num_heads,
+                                                const int channels, 
+                                                const int num_levels,
+                                                const int num_query,
+                                                const int num_point,
+                                                scalar_t *grad_value,
+                                                scalar_t *grad_sampling_loc,
+                                                scalar_t *grad_attn_weight)
+{
+  CUDA_KERNEL_LOOP(index, n)
+  {
+    int _temp = index;
+    const int c_col = _temp % channels;
+    _temp /= channels;
+    const int sampling_index = _temp; 
+    const int m_col = _temp % num_heads;
+    _temp /= num_heads;
+    const int q_col = _temp % num_query;
+    _temp /= num_query;
+    const int b_col = _temp;
+
+    const scalar_t top_grad = grad_col[index];
+
+    int data_weight_ptr = sampling_index * num_levels * num_point;
+    int data_loc_w_ptr = data_weight_ptr << 1;
+    const int grad_sampling_ptr = data_weight_ptr;
+    grad_sampling_loc += grad_sampling_ptr << 1;
+    grad_attn_weight += grad_sampling_ptr;
+    const int grad_weight_stride = 1;
+    const int grad_loc_stride = 2;
+    const int qid_stride = num_heads * channels;
+    const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+    for (int l_col=0; l_col < num_levels; ++l_col)
+    {
+      const int level_start_id = data_level_start_index[l_col];
+      const int spatial_h_ptr = l_col << 1;
+      const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+      const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+      const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+      const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+      scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+      for (int p_col=0; p_col < num_point; ++p_col)
+      {
+        const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+        const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+        const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+        const scalar_t h_im = loc_h * spatial_h - 0.5;
+        const scalar_t w_im = loc_w * spatial_w - 0.5;
+        if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w)
+        {
+          ms_deform_attn_col2im_bilinear_gm(
+            data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col,
+            top_grad, weight, grad_value_ptr, 
+            grad_sampling_loc, grad_attn_weight);
+        }
+        data_weight_ptr += 1;
+        data_loc_w_ptr += 2;
+        grad_attn_weight += grad_weight_stride;
+        grad_sampling_loc += grad_loc_stride;
+      }
+    }
+  }
+}
+
+
+template <typename scalar_t>
+void ms_deformable_im2col_cuda(cudaStream_t stream,
+                              const scalar_t* data_value,
+                              const int64_t* data_spatial_shapes, 
+                              const int64_t* data_level_start_index, 
+                              const scalar_t* data_sampling_loc,
+                              const scalar_t* data_attn_weight,
+                              const int batch_size,
+                              const int spatial_size, 
+                              const int num_heads, 
+                              const int channels, 
+                              const int num_levels, 
+                              const int num_query,
+                              const int num_point,
+                              scalar_t* data_col)
+{
+  const int num_kernels = batch_size * num_query * num_heads * channels;
+  const int num_actual_kernels = batch_size * num_query * num_heads * channels;
+  const int num_threads = CUDA_NUM_THREADS;
+  ms_deformable_im2col_gpu_kernel<scalar_t>
+      <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+          0, stream>>>(
+      num_kernels, data_value, data_spatial_shapes, data_level_start_index, data_sampling_loc, data_attn_weight, 
+      batch_size, spatial_size, num_heads, channels, num_levels, num_query, num_point, data_col);
+  
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess)
+  {
+    printf("error in ms_deformable_im2col_cuda: %s\n", cudaGetErrorString(err));
+  }
+
+}
+
+template <typename scalar_t>
+void ms_deformable_col2im_cuda(cudaStream_t stream,
+                              const scalar_t* grad_col,
+                              const scalar_t* data_value,
+                              const int64_t * data_spatial_shapes,
+                              const int64_t * data_level_start_index,
+                              const scalar_t * data_sampling_loc,
+                              const scalar_t * data_attn_weight,
+                              const int batch_size, 
+                              const int spatial_size, 
+                              const int num_heads,
+                              const int channels, 
+                              const int num_levels,
+                              const int num_query,
+                              const int num_point, 
+                              scalar_t* grad_value,
+                              scalar_t* grad_sampling_loc,
+                              scalar_t* grad_attn_weight)
+{
+  const int num_threads = (channels > CUDA_NUM_THREADS)?CUDA_NUM_THREADS:channels;
+  const int num_kernels = batch_size * num_query * num_heads * channels;
+  const int num_actual_kernels = batch_size * num_query * num_heads * channels;
+  if (channels > 1024)
+  {
+    if ((channels & 1023) == 0)
+    {
+      ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks<scalar_t>
+          <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+              num_threads*3*sizeof(scalar_t), stream>>>(
+                        num_kernels, 
+                        grad_col,
+                        data_value,
+                        data_spatial_shapes,
+                        data_level_start_index, 
+                        data_sampling_loc,
+                        data_attn_weight,
+                        batch_size, 
+                        spatial_size, 
+                        num_heads,
+                        channels, 
+                        num_levels,
+                        num_query,
+                        num_point,
+                        grad_value,
+                        grad_sampling_loc,
+                        grad_attn_weight);
+    }
+    else
+    {
+      ms_deformable_col2im_gpu_kernel_gm<scalar_t>
+        <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+            0, stream>>>(
+                      num_kernels, 
+                      grad_col,
+                      data_value,
+                      data_spatial_shapes,
+                      data_level_start_index, 
+                      data_sampling_loc,
+                      data_attn_weight,
+                      batch_size, 
+                      spatial_size, 
+                      num_heads,
+                      channels, 
+                      num_levels,
+                      num_query,
+                      num_point,
+                      grad_value,
+                      grad_sampling_loc,
+                      grad_attn_weight);
+    }
+  }
+  else{
+    switch(channels)
+    {
+      case 1:
+        ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 1>
+        <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+            0, stream>>>(
+                      num_kernels, 
+                      grad_col,
+                      data_value,
+                      data_spatial_shapes,
+                      data_level_start_index, 
+                      data_sampling_loc,
+                      data_attn_weight,
+                      batch_size, 
+                      spatial_size, 
+                      num_heads,
+                      channels, 
+                      num_levels,
+                      num_query,
+                      num_point,
+                      grad_value,
+                      grad_sampling_loc,
+                      grad_attn_weight);
+        break;
+      case 2:
+        ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 2>
+        <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+            0, stream>>>(
+                      num_kernels, 
+                      grad_col,
+                      data_value,
+                      data_spatial_shapes,
+                      data_level_start_index, 
+                      data_sampling_loc,
+                      data_attn_weight,
+                      batch_size, 
+                      spatial_size, 
+                      num_heads,
+                      channels, 
+                      num_levels,
+                      num_query,
+                      num_point,
+                      grad_value,
+                      grad_sampling_loc,
+                      grad_attn_weight);
+        break;
+      case 4:
+        ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 4>
+        <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+            0, stream>>>(
+                      num_kernels, 
+                      grad_col,
+                      data_value,
+                      data_spatial_shapes,
+                      data_level_start_index, 
+                      data_sampling_loc,
+                      data_attn_weight,
+                      batch_size, 
+                      spatial_size, 
+                      num_heads,
+                      channels, 
+                      num_levels,
+                      num_query,
+                      num_point,
+                      grad_value,
+                      grad_sampling_loc,
+                      grad_attn_weight);
+        break;
+      case 8:
+        ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 8>
+        <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+            0, stream>>>(
+                      num_kernels, 
+                      grad_col,
+                      data_value,
+                      data_spatial_shapes,
+                      data_level_start_index, 
+                      data_sampling_loc,
+                      data_attn_weight,
+                      batch_size, 
+                      spatial_size, 
+                      num_heads,
+                      channels, 
+                      num_levels,
+                      num_query,
+                      num_point,
+                      grad_value,
+                      grad_sampling_loc,
+                      grad_attn_weight);
+        break;
+      case 16:
+        ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 16>
+        <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+            0, stream>>>(
+                      num_kernels, 
+                      grad_col,
+                      data_value,
+                      data_spatial_shapes,
+                      data_level_start_index, 
+                      data_sampling_loc,
+                      data_attn_weight,
+                      batch_size, 
+                      spatial_size, 
+                      num_heads,
+                      channels, 
+                      num_levels,
+                      num_query,
+                      num_point,
+                      grad_value,
+                      grad_sampling_loc,
+                      grad_attn_weight);
+        break;
+      case 32:
+        ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 32>
+        <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+            0, stream>>>(
+                      num_kernels, 
+                      grad_col,
+                      data_value,
+                      data_spatial_shapes,
+                      data_level_start_index, 
+                      data_sampling_loc,
+                      data_attn_weight,
+                      batch_size, 
+                      spatial_size, 
+                      num_heads,
+                      channels, 
+                      num_levels,
+                      num_query,
+                      num_point,
+                      grad_value,
+                      grad_sampling_loc,
+                      grad_attn_weight);
+        break;
+      case 64:
+        ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 64>
+        <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+            0, stream>>>(
+                      num_kernels, 
+                      grad_col,
+                      data_value,
+                      data_spatial_shapes,
+                      data_level_start_index, 
+                      data_sampling_loc,
+                      data_attn_weight,
+                      batch_size, 
+                      spatial_size, 
+                      num_heads,
+                      channels, 
+                      num_levels,
+                      num_query,
+                      num_point,
+                      grad_value,
+                      grad_sampling_loc,
+                      grad_attn_weight);
+        break;
+      case 128:
+        ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 128>
+        <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+            0, stream>>>(
+                      num_kernels, 
+                      grad_col,
+                      data_value,
+                      data_spatial_shapes,
+                      data_level_start_index, 
+                      data_sampling_loc,
+                      data_attn_weight,
+                      batch_size, 
+                      spatial_size, 
+                      num_heads,
+                      channels, 
+                      num_levels,
+                      num_query,
+                      num_point,
+                      grad_value,
+                      grad_sampling_loc,
+                      grad_attn_weight);
+        break;
+      case 256:
+        ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 256>
+        <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+            0, stream>>>(
+                      num_kernels, 
+                      grad_col,
+                      data_value,
+                      data_spatial_shapes,
+                      data_level_start_index, 
+                      data_sampling_loc,
+                      data_attn_weight,
+                      batch_size, 
+                      spatial_size, 
+                      num_heads,
+                      channels, 
+                      num_levels,
+                      num_query,
+                      num_point,
+                      grad_value,
+                      grad_sampling_loc,
+                      grad_attn_weight);
+        break;
+      case 512:
+        ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 512>
+        <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+            0, stream>>>(
+                      num_kernels, 
+                      grad_col,
+                      data_value,
+                      data_spatial_shapes,
+                      data_level_start_index, 
+                      data_sampling_loc,
+                      data_attn_weight,
+                      batch_size, 
+                      spatial_size, 
+                      num_heads,
+                      channels, 
+                      num_levels,
+                      num_query,
+                      num_point,
+                      grad_value,
+                      grad_sampling_loc,
+                      grad_attn_weight);
+        break;
+      case 1024:
+        ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 1024>
+        <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+            0, stream>>>(
+                      num_kernels, 
+                      grad_col,
+                      data_value,
+                      data_spatial_shapes,
+                      data_level_start_index, 
+                      data_sampling_loc,
+                      data_attn_weight,
+                      batch_size, 
+                      spatial_size, 
+                      num_heads,
+                      channels, 
+                      num_levels,
+                      num_query,
+                      num_point,
+                      grad_value,
+                      grad_sampling_loc,
+                      grad_attn_weight);
+        break;
+      default:
+        if (channels < 64)
+        {
+          ms_deformable_col2im_gpu_kernel_shm_reduce_v1<scalar_t>
+          <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+              num_threads*3*sizeof(scalar_t), stream>>>(
+                        num_kernels, 
+                        grad_col,
+                        data_value,
+                        data_spatial_shapes,
+                        data_level_start_index, 
+                        data_sampling_loc,
+                        data_attn_weight,
+                        batch_size, 
+                        spatial_size, 
+                        num_heads,
+                        channels, 
+                        num_levels,
+                        num_query,
+                        num_point,
+                        grad_value,
+                        grad_sampling_loc,
+                        grad_attn_weight);
+        }
+        else
+        {
+          ms_deformable_col2im_gpu_kernel_shm_reduce_v2<scalar_t>
+          <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads,
+              num_threads*3*sizeof(scalar_t), stream>>>(
+                        num_kernels, 
+                        grad_col,
+                        data_value,
+                        data_spatial_shapes,
+                        data_level_start_index, 
+                        data_sampling_loc,
+                        data_attn_weight,
+                        batch_size, 
+                        spatial_size, 
+                        num_heads,
+                        channels, 
+                        num_levels,
+                        num_query,
+                        num_point,
+                        grad_value,
+                        grad_sampling_loc,
+                        grad_attn_weight);
+        }
+    }
+  }
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess)
+  {
+    printf("error in ms_deformable_col2im_cuda: %s\n", cudaGetErrorString(err));
+  }
+
+}

imcui/third_party/rdd/RDD/models/ops/src/ms_deform_attn.h

@@ -0,0 +1,67 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+/*!
+* Copyright (c) Facebook, Inc. and its affiliates.
+* Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+*/
+
+#pragma once
+
+#include "cpu/ms_deform_attn_cpu.h"
+
+#ifdef WITH_CUDA
+#include "cuda/ms_deform_attn_cuda.h"
+#endif
+
+
+at::Tensor
+ms_deform_attn_forward(
+    const at::Tensor &value, 
+    const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index,
+    const at::Tensor &sampling_loc,
+    const at::Tensor &attn_weight,
+    const int im2col_step)
+{
+    if (value.type().is_cuda())
+    {
+#ifdef WITH_CUDA
+        return ms_deform_attn_cuda_forward(
+            value, spatial_shapes, level_start_index, sampling_loc, attn_weight, im2col_step);
+#else
+        AT_ERROR("Not compiled with GPU support");
+#endif
+    }
+    AT_ERROR("Not implemented on the CPU");
+}
+
+std::vector<at::Tensor>
+ms_deform_attn_backward(
+    const at::Tensor &value, 
+    const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index,
+    const at::Tensor &sampling_loc,
+    const at::Tensor &attn_weight,
+    const at::Tensor &grad_output,
+    const int im2col_step)
+{
+    if (value.type().is_cuda())
+    {
+#ifdef WITH_CUDA
+        return ms_deform_attn_cuda_backward(
+            value, spatial_shapes, level_start_index, sampling_loc, attn_weight, grad_output, im2col_step);
+#else
+        AT_ERROR("Not compiled with GPU support");
+#endif
+    }
+    AT_ERROR("Not implemented on the CPU");
+}
+

imcui/third_party/rdd/RDD/models/ops/src/vision.cpp

@@ -0,0 +1,21 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+/*!
+* Copyright (c) Facebook, Inc. and its affiliates.
+* Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+*/
+
+#include "ms_deform_attn.h"
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("ms_deform_attn_forward", &ms_deform_attn_forward, "ms_deform_attn_forward");
+  m.def("ms_deform_attn_backward", &ms_deform_attn_backward, "ms_deform_attn_backward");
+}

imcui/third_party/rdd/RDD/models/ops/test.py

@@ -0,0 +1,92 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Modified by Bowen Cheng from https://github.com/fundamentalvision/Deformable-DETR
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import time
+import torch
+import torch.nn as nn
+from torch.autograd import gradcheck
+
+from functions.ms_deform_attn_func import MSDeformAttnFunction, ms_deform_attn_core_pytorch
+
+
+N, M, D = 1, 2, 2
+Lq, L, P = 2, 2, 2
+shapes = torch.as_tensor([(6, 4), (3, 2)], dtype=torch.long).cuda()
+level_start_index = torch.cat((shapes.new_zeros((1, )), shapes.prod(1).cumsum(0)[:-1]))
+S = sum([(H*W).item() for H, W in shapes])
+
+
+torch.manual_seed(3)
+
+
+@torch.no_grad()
+def check_forward_equal_with_pytorch_double():
+    value = torch.rand(N, S, M, D).cuda() * 0.01
+    sampling_locations = torch.rand(N, Lq, M, L, P, 2).cuda()
+    attention_weights = torch.rand(N, Lq, M, L, P).cuda() + 1e-5
+    attention_weights /= attention_weights.sum(-1, keepdim=True).sum(-2, keepdim=True)
+    im2col_step = 2
+    output_pytorch = ms_deform_attn_core_pytorch(value.double(), shapes, sampling_locations.double(), attention_weights.double()).detach().cpu()
+    output_cuda = MSDeformAttnFunction.apply(value.double(), shapes, level_start_index, sampling_locations.double(), attention_weights.double(), im2col_step).detach().cpu()
+    fwdok = torch.allclose(output_cuda, output_pytorch)
+    max_abs_err = (output_cuda - output_pytorch).abs().max()
+    max_rel_err = ((output_cuda - output_pytorch).abs() / output_pytorch.abs()).max()
+
+    print(f'* {fwdok} check_forward_equal_with_pytorch_double: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+
+
+@torch.no_grad()
+def check_forward_equal_with_pytorch_float():
+    value = torch.rand(N, S, M, D).cuda() * 0.01
+    sampling_locations = torch.rand(N, Lq, M, L, P, 2).cuda()
+    attention_weights = torch.rand(N, Lq, M, L, P).cuda() + 1e-5
+    attention_weights /= attention_weights.sum(-1, keepdim=True).sum(-2, keepdim=True)
+    im2col_step = 2
+    output_pytorch = ms_deform_attn_core_pytorch(value, shapes, sampling_locations, attention_weights).detach().cpu()
+    output_cuda = MSDeformAttnFunction.apply(value, shapes, level_start_index, sampling_locations, attention_weights, im2col_step).detach().cpu()
+    fwdok = torch.allclose(output_cuda, output_pytorch, rtol=1e-2, atol=1e-3)
+    max_abs_err = (output_cuda - output_pytorch).abs().max()
+    max_rel_err = ((output_cuda - output_pytorch).abs() / output_pytorch.abs()).max()
+
+    print(f'* {fwdok} check_forward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+
+
+def check_gradient_numerical(channels=4, grad_value=True, grad_sampling_loc=True, grad_attn_weight=True):
+
+    value = torch.rand(N, S, M, channels).cuda() * 0.01
+    sampling_locations = torch.rand(N, Lq, M, L, P, 2).cuda()
+    attention_weights = torch.rand(N, Lq, M, L, P).cuda() + 1e-5
+    attention_weights /= attention_weights.sum(-1, keepdim=True).sum(-2, keepdim=True)
+    im2col_step = 2
+    func = MSDeformAttnFunction.apply
+
+    value.requires_grad = grad_value
+    sampling_locations.requires_grad = grad_sampling_loc
+    attention_weights.requires_grad = grad_attn_weight
+
+    gradok = gradcheck(func, (value.double(), shapes, level_start_index, sampling_locations.double(), attention_weights.double(), im2col_step))
+
+    print(f'* {gradok} check_gradient_numerical(D={channels})')
+
+
+if __name__ == '__main__':
+    check_forward_equal_with_pytorch_double()
+    check_forward_equal_with_pytorch_float()
+
+    for channels in [30, 32, 64, 71, 1025, 2048, 3096]:
+        check_gradient_numerical(channels, True, True, True)
+
+
+

imcui/third_party/rdd/RDD/models/position_encoding.py

@@ -0,0 +1,48 @@
+import math
+import torch
+from torch import nn
+from ..utils.misc import NestedTensor
+
+class PositionEmbeddingSine(nn.Module):
+    """
+    This is a more standard version of the position embedding, very similar to the one
+    used by the Attention is all you need paper, generalized to work on images.
+    """
+    def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
+        super().__init__()
+        self.num_pos_feats = num_pos_feats
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError("normalize should be True if scale is passed")
+        if scale is None:
+            scale = 2 * math.pi
+        self.scale = scale
+
+    def forward(self, tensor_list: NestedTensor):
+        x = tensor_list.tensors
+        mask = tensor_list.mask
+        assert mask is not None
+        not_mask = ~mask
+        y_embed = not_mask.cumsum(1, dtype=torch.float32)
+        x_embed = not_mask.cumsum(2, dtype=torch.float32)
+        if self.normalize:
+            eps = 1e-6
+            y_embed = (y_embed - 0.5) / (y_embed[:, -1:, :] + eps) * self.scale
+            x_embed = (x_embed - 0.5) / (x_embed[:, :, -1:] + eps) * self.scale
+
+        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
+        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
+
+        pos_x = x_embed[:, :, :, None] / dim_t
+        pos_y = y_embed[:, :, :, None] / dim_t
+        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
+        return pos
+    
+def build_position_encoding(config):
+    N_steps = config['hidden_dim'] // 2
+    # TODO find a better way of exposing other arguments
+    position_embedding = PositionEmbeddingSine(N_steps, normalize=True)
+    return position_embedding

imcui/third_party/rdd/RDD/models/soft_detect.py

@@ -0,0 +1,250 @@
+# ALIKE: https://github.com/Shiaoming/ALIKE
+import torch
+from torch import nn
+import numpy as np
+import torch.nn.functional as F
+
+
+# coordinates system
+#  ------------------------------>  [ x: range=-1.0~1.0; w: range=0~W ]
+#  | -----------------------------
+#  | |                           |
+#  | |                           |
+#  | |                           |
+#  | |         image             |
+#  | |                           |
+#  | |                           |
+#  | |                           |
+#  | |---------------------------|
+#  v
+# [ y: range=-1.0~1.0; h: range=0~H ]
+
+def simple_nms(scores, nms_radius: int):
+    """ Fast Non-maximum suppression to remove nearby points """
+    assert (nms_radius >= 0)
+
+    def max_pool(x):
+        return torch.nn.functional.max_pool2d(
+            x, kernel_size=nms_radius * 2 + 1, stride=1, padding=nms_radius)
+
+    zeros = torch.zeros_like(scores)
+    max_mask = scores == max_pool(scores)
+
+    for _ in range(2):
+        supp_mask = max_pool(max_mask.float()) > 0
+        supp_scores = torch.where(supp_mask, zeros, scores)
+        new_max_mask = supp_scores == max_pool(supp_scores)
+        max_mask = max_mask | (new_max_mask & (~supp_mask))
+    return torch.where(max_mask, scores, zeros)
+
+
+"""
+	"XFeat: Accelerated Features for Lightweight Image Matching, CVPR 2024."
+	https://www.verlab.dcc.ufmg.br/descriptors/xfeat_cvpr24/
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class InterpolateSparse2d(nn.Module):
+    """ Efficiently interpolate tensor at given sparse 2D positions. """ 
+    def __init__(self, mode = 'bicubic', align_corners = False): 
+        super().__init__()
+        self.mode = mode
+        self.align_corners = align_corners
+
+    def normgrid(self, x, H, W):
+        """ Normalize coords to [-1,1]. """
+        return 2. * (x/(torch.tensor([W-1, H-1], device = x.device, dtype = x.dtype))) - 1.
+
+    def forward(self, x, pos, H, W):
+        """
+        Input
+            x: [B, C, H, W] feature tensor
+            pos: [B, N, 2] tensor of positions
+            H, W: int, original resolution of input 2d positions -- used in normalization [-1,1]
+
+        Returns
+            [B, N, C] sampled channels at 2d positions
+        """
+        grid = self.normgrid(pos, H, W).unsqueeze(-2).to(x.dtype)
+        x = F.grid_sample(x, grid, mode = self.mode , align_corners = False)
+        return x.permute(0,2,3,1).squeeze(-2)
+
+
+class SoftDetect(nn.Module):
+    def __init__(self, radius=2, top_k=0, scores_th=0.2, n_limit=20000):
+        """
+        Args:
+            radius: soft detection radius, kernel size is (2 * radius + 1)
+            top_k: top_k > 0: return top k keypoints
+            scores_th: top_k <= 0 threshold mode:  scores_th > 0: return keypoints with scores>scores_th
+                                                   else: return keypoints with scores > scores.mean()
+            n_limit: max number of keypoint in threshold mode
+        """
+        super().__init__()
+        self.radius = radius
+        self.top_k = top_k
+        self.scores_th = scores_th
+        self.n_limit = n_limit
+        self.kernel_size = 2 * self.radius + 1
+        self.temperature = 0.1  # tuned temperature
+        self.unfold = nn.Unfold(kernel_size=self.kernel_size, padding=self.radius)
+        self.sample_descriptor = InterpolateSparse2d('bicubic')
+        # local xy grid
+        x = torch.linspace(-self.radius, self.radius, self.kernel_size)
+        # (kernel_size*kernel_size) x 2 : (w,h)
+        self.hw_grid = torch.stack(torch.meshgrid([x, x])).view(2, -1).t()[:, [1, 0]]
+
+    def detect_keypoints(self, scores_map, normalized_coordinates=True):
+        b, c, h, w = scores_map.shape
+        scores_nograd = scores_map.detach()
+        
+        # nms_scores = simple_nms(scores_nograd, self.radius)
+        nms_scores = simple_nms(scores_nograd, 2)
+
+        # remove border
+        nms_scores[:, :, :self.radius + 1, :] = 0
+        nms_scores[:, :, :, :self.radius + 1] = 0
+        nms_scores[:, :, h - self.radius:, :] = 0
+        nms_scores[:, :, :, w - self.radius:] = 0
+
+        # detect keypoints without grad
+        if self.top_k > 0:
+            topk = torch.topk(nms_scores.view(b, -1), self.top_k)
+            indices_keypoints = topk.indices  # B x top_k
+        else:
+            if self.scores_th > 0:
+                masks = nms_scores > self.scores_th
+                if masks.sum() == 0:
+                    th = scores_nograd.reshape(b, -1).mean(dim=1)  # th = self.scores_th
+                    masks = nms_scores > th.reshape(b, 1, 1, 1)
+            else:
+                th = scores_nograd.reshape(b, -1).mean(dim=1)  # th = self.scores_th
+                masks = nms_scores > th.reshape(b, 1, 1, 1)
+            masks = masks.reshape(b, -1)
+
+            indices_keypoints = []  # list, B x (any size)
+            scores_view = scores_nograd.reshape(b, -1)
+            for mask, scores in zip(masks, scores_view):
+                indices = mask.nonzero(as_tuple=False)[:, 0]
+                if len(indices) > self.n_limit:
+                    kpts_sc = scores[indices]
+                    sort_idx = kpts_sc.sort(descending=True)[1]
+                    sel_idx = sort_idx[:self.n_limit]
+                    indices = indices[sel_idx]
+                indices_keypoints.append(indices)
+
+        # detect soft keypoints with grad backpropagation
+        patches = self.unfold(scores_map)  # B x (kernel**2) x (H*W)
+        self.hw_grid = self.hw_grid.to(patches)  # to device
+        keypoints = []
+        scoredispersitys = []
+        kptscores = []
+        for b_idx in range(b):
+            patch = patches[b_idx].t()  # (H*W) x (kernel**2)
+            indices_kpt = indices_keypoints[b_idx]  # one dimension vector, say its size is M
+            patch_scores = patch[indices_kpt]  # M x (kernel**2)
+
+            # max is detached to prevent undesired backprop loops in the graph
+            max_v = patch_scores.max(dim=1).values.detach()[:, None]
+            x_exp = ((patch_scores - max_v) / self.temperature).exp()  # M * (kernel**2), in [0, 1]
+
+            # \frac{ \sum{(i,j) \times \exp(x/T)} }{ \sum{\exp(x/T)} }
+            xy_residual = x_exp @ self.hw_grid / x_exp.sum(dim=1)[:, None]  # Soft-argmax, Mx2
+
+            hw_grid_dist2 = torch.norm((self.hw_grid[None, :, :] - xy_residual[:, None, :]) / self.radius,
+                                       dim=-1) ** 2
+            scoredispersity = (x_exp * hw_grid_dist2).sum(dim=1) / x_exp.sum(dim=1)
+
+            # compute result keypoints
+            keypoints_xy_nms = torch.stack([indices_kpt % w, indices_kpt // w], dim=1)  # Mx2
+            keypoints_xy = keypoints_xy_nms + xy_residual
+            if normalized_coordinates:
+                keypoints_xy = keypoints_xy / keypoints_xy.new_tensor([w - 1, h - 1]) * 2 - 1  # (w,h) -> (-1~1,-1~1)
+
+            kptscore = torch.nn.functional.grid_sample(scores_map[b_idx].unsqueeze(0), keypoints_xy.view(1, 1, -1, 2),
+                                                       mode='bilinear', align_corners=True)[0, 0, 0, :]  # CxN
+
+            keypoints.append(keypoints_xy)
+            scoredispersitys.append(scoredispersity)
+            kptscores.append(kptscore)
+
+        return keypoints, scoredispersitys, kptscores
+
+    def forward(self, scores_map, normalized_coordinates=True):
+        """
+        :param scores_map:  Bx1xHxW
+        :return: kpts: list[Nx2,...]; kptscores: list[N,....] normalised position: -1.0 ~ 1.0
+        """
+        B, _, H, W = scores_map.shape
+
+        keypoints, scoredispersitys, kptscores = self.detect_keypoints(scores_map,
+                                                                       normalized_coordinates)
+        
+        # keypoints: B M 2
+        # scoredispersitys:
+        return keypoints, kptscores, scoredispersitys
+    
+import torch
+import torch.nn as nn
+
+class Detect(nn.Module):
+    def __init__(self, stride=4, top_k=0, scores_th=0, n_limit=20000):
+        super().__init__()
+        self.stride = stride
+        self.top_k = top_k
+        self.scores_th = scores_th
+        self.n_limit = n_limit
+
+    def forward(self, scores, coords, w, h):
+        """
+        scores: B x N x 1 (keypoint confidence scores)
+        coords: B x N x 2 (offsets within stride x stride window)
+        w, h: Image dimensions
+        """
+        b, n, _ = scores.shape
+        kpts_list = []
+        scores_list = []
+
+        for b_idx in range(b):
+            score = scores[b_idx].squeeze(-1)  # Shape: (N,)
+            coord = coords[b_idx]  # Shape: (N, 2)
+
+            # Apply score thresholding
+            if self.scores_th >= 0:
+                valid = score > self.scores_th
+            else:
+                valid = score > score.mean()
+
+            valid_indices = valid.nonzero(as_tuple=True)[0]  # Get valid indices
+            if valid_indices.numel() == 0:
+                kpts_list.append(torch.empty((0, 2), device=scores.device))
+                scores_list.append(torch.empty((0,), device=scores.device))
+                continue
+
+            # Compute keypoint locations in original image space
+            i_ids = valid_indices  # Indices where keypoints exist
+            kpts = torch.stack([i_ids % w, i_ids // w], dim=1).to(torch.float) * self.stride  # Grid position
+            kpts += coord[i_ids] * self.stride  # Apply offset
+
+            # Normalize keypoints to [-1, 1] range
+            kpts = (kpts / torch.tensor([w - 1, h - 1], device=kpts.device, dtype=kpts.dtype)) * 2 - 1
+
+            # Filter top-k keypoints if needed
+            scores_valid = score[valid_indices]
+            if self.top_k > 0 and len(kpts) > self.top_k:
+                topk = torch.topk(scores_valid, self.top_k, dim=0)
+                kpts = kpts[topk.indices]
+                scores_valid = topk.values
+            elif self.top_k < 0:
+                if len(kpts) > self.n_limit:
+                    sorted_idx = scores_valid.argsort(descending=True)[:self.n_limit]
+                    kpts = kpts[sorted_idx]
+                    scores_valid = scores_valid[sorted_idx]
+
+            kpts_list.append(kpts)
+            scores_list.append(scores_valid)
+
+        return kpts_list, scores_list

imcui/third_party/rdd/RDD/utils/__init__.py

@@ -0,0 +1 @@
+from .misc import *

imcui/third_party/rdd/RDD/utils/misc.py

@@ -0,0 +1,531 @@
+# Modified from Deformable DETR 
+# https://github.com/fundamentalvision/Deformable-DETR
+# ------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Modified from DETR (https://github.com/facebookresearch/detr)
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# ------------------------------------------------------------------------
+import os
+import subprocess
+import time
+from collections import defaultdict, deque
+import datetime
+import pickle
+from typing import Optional, List
+import yaml
+import torch
+import torch.nn as nn
+import torch.distributed as dist
+from torch import Tensor
+
+# needed due to empty tensor bug in pytorch and torchvision 0.5
+import torchvision
+if float(torchvision.__version__.split('.')[1]) < 0.5:
+    import math
+    from torchvision.ops.misc import _NewEmptyTensorOp
+    def _check_size_scale_factor(dim, size, scale_factor):
+        # type: (int, Optional[List[int]], Optional[float]) -> None
+        if size is None and scale_factor is None:
+            raise ValueError("either size or scale_factor should be defined")
+        if size is not None and scale_factor is not None:
+            raise ValueError("only one of size or scale_factor should be defined")
+        if not (scale_factor is not None and len(scale_factor) != dim):
+            raise ValueError(
+                "scale_factor shape must match input shape. "
+                "Input is {}D, scale_factor size is {}".format(dim, len(scale_factor))
+            )
+    def _output_size(dim, input, size, scale_factor):
+        # type: (int, Tensor, Optional[List[int]], Optional[float]) -> List[int]
+        assert dim == 2
+        _check_size_scale_factor(dim, size, scale_factor)
+        if size is not None:
+            return size
+        # if dim is not 2 or scale_factor is iterable use _ntuple instead of concat
+        assert scale_factor is not None and isinstance(scale_factor, (int, float))
+        scale_factors = [scale_factor, scale_factor]
+        # math.floor might return float in py2.7
+        return [
+            int(math.floor(input.size(i + 2) * scale_factors[i])) for i in range(dim)
+        ]
+elif float(torchvision.__version__.split('.')[1]) < 7:
+    from torchvision.ops import _new_empty_tensor
+    from torchvision.ops.misc import _output_size
+
+
+class SmoothedValue(object):
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=self.max,
+            value=self.value)
+
+
+def all_gather(data):
+    """
+    Run all_gather on arbitrary picklable data (not necessarily tensors)
+    Args:
+        data: any picklable object
+    Returns:
+        list[data]: list of data gathered from each rank
+    """
+    world_size = get_world_size()
+    if world_size == 1:
+        return [data]
+
+    # serialized to a Tensor
+    buffer = pickle.dumps(data)
+    storage = torch.ByteStorage.from_buffer(buffer)
+    tensor = torch.ByteTensor(storage).to("cuda")
+
+    # obtain Tensor size of each rank
+    local_size = torch.tensor([tensor.numel()], device="cuda")
+    size_list = [torch.tensor([0], device="cuda") for _ in range(world_size)]
+    dist.all_gather(size_list, local_size)
+    size_list = [int(size.item()) for size in size_list]
+    max_size = max(size_list)
+
+    # receiving Tensor from all ranks
+    # we pad the tensor because torch all_gather does not support
+    # gathering tensors of different shapes
+    tensor_list = []
+    for _ in size_list:
+        tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device="cuda"))
+    if local_size != max_size:
+        padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device="cuda")
+        tensor = torch.cat((tensor, padding), dim=0)
+    dist.all_gather(tensor_list, tensor)
+
+    data_list = []
+    for size, tensor in zip(size_list, tensor_list):
+        buffer = tensor.cpu().numpy().tobytes()[:size]
+        data_list.append(pickle.loads(buffer))
+
+    return data_list
+
+
+def reduce_dict(input_dict, average=True):
+    """
+    Args:
+        input_dict (dict): all the values will be reduced
+        average (bool): whether to do average or sum
+    Reduce the values in the dictionary from all processes so that all processes
+    have the averaged results. Returns a dict with the same fields as
+    input_dict, after reduction.
+    """
+    world_size = get_world_size()
+    if world_size < 2:
+        return input_dict
+    with torch.no_grad():
+        names = []
+        values = []
+        # sort the keys so that they are consistent across processes
+        for k in sorted(input_dict.keys()):
+            names.append(k)
+            values.append(input_dict[k])
+        values = torch.stack(values, dim=0)
+        dist.all_reduce(values)
+        if average:
+            values /= world_size
+        reduced_dict = {k: v for k, v in zip(names, values)}
+    return reduced_dict
+
+
+class MetricLogger(object):
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError("'{}' object has no attribute '{}'".format(
+            type(self).__name__, attr))
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(
+                "{}: {}".format(name, str(meter))
+            )
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ''
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt='{avg:.4f}')
+        data_time = SmoothedValue(fmt='{avg:.4f}')
+        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
+        if torch.cuda.is_available():
+            log_msg = self.delimiter.join([
+                header,
+                '[{0' + space_fmt + '}/{1}]',
+                'eta: {eta}',
+                '{meters}',
+                'time: {time}',
+                'data: {data}',
+                'max mem: {memory:.0f}'
+            ])
+        else:
+            log_msg = self.delimiter.join([
+                header,
+                '[{0' + space_fmt + '}/{1}]',
+                'eta: {eta}',
+                '{meters}',
+                'time: {time}',
+                'data: {data}'
+            ])
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(iterable) - 1:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time),
+                        memory=torch.cuda.max_memory_allocated() / MB))
+                else:
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time)))
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print('{} Total time: {} ({:.4f} s / it)'.format(
+            header, total_time_str, total_time / len(iterable)))
+
+
+def get_sha():
+    cwd = os.path.dirname(os.path.abspath(__file__))
+
+    def _run(command):
+        return subprocess.check_output(command, cwd=cwd).decode('ascii').strip()
+    sha = 'N/A'
+    diff = "clean"
+    branch = 'N/A'
+    try:
+        sha = _run(['git', 'rev-parse', 'HEAD'])
+        subprocess.check_output(['git', 'diff'], cwd=cwd)
+        diff = _run(['git', 'diff-index', 'HEAD'])
+        diff = "has uncommited changes" if diff else "clean"
+        branch = _run(['git', 'rev-parse', '--abbrev-ref', 'HEAD'])
+    except Exception:
+        pass
+    message = f"sha: {sha}, status: {diff}, branch: {branch}"
+    return message
+
+
+def collate_fn(batch):
+    batch = list(zip(*batch))
+    batch[0] = nested_tensor_from_tensor_list(batch[0])
+    return tuple(batch)
+
+
+def _max_by_axis(the_list):
+    # type: (List[List[int]]) -> List[int]
+    maxes = the_list[0]
+    for sublist in the_list[1:]:
+        for index, item in enumerate(sublist):
+            maxes[index] = max(maxes[index], item)
+    return maxes
+
+
+def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):
+    # TODO make this more general
+    if tensor_list[0].ndim == 3:
+        # TODO make it support different-sized images
+        max_size = _max_by_axis([list(img.shape) for img in tensor_list])
+        # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))
+        batch_shape = [len(tensor_list)] + max_size
+        b, c, h, w = batch_shape
+        dtype = tensor_list[0].dtype
+        device = tensor_list[0].device
+        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
+        mask = torch.ones((b, h, w), dtype=torch.bool, device=device)
+        for img, pad_img, m in zip(tensor_list, tensor, mask):
+            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
+            m[: img.shape[1], :img.shape[2]] = False
+    else:
+        raise ValueError('not supported')
+    return NestedTensor(tensor, mask)
+
+
+class NestedTensor(object):
+    def __init__(self, tensors, mask: Optional[Tensor]):
+        self.tensors = tensors
+        self.mask = mask
+
+    def to(self, device, non_blocking=False):
+        # type: (Device) -> NestedTensor # noqa
+        cast_tensor = self.tensors.to(device, non_blocking=non_blocking)
+        mask = self.mask
+        if mask is not None:
+            assert mask is not None
+            cast_mask = mask.to(device, non_blocking=non_blocking)
+        else:
+            cast_mask = None
+        return NestedTensor(cast_tensor, cast_mask)
+
+    def record_stream(self, *args, **kwargs):
+        self.tensors.record_stream(*args, **kwargs)
+        if self.mask is not None:
+            self.mask.record_stream(*args, **kwargs)
+
+    def decompose(self):
+        return self.tensors, self.mask
+
+    def __repr__(self):
+        return str(self.tensors)
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    import builtins as __builtin__
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop('force', False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def get_local_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return int(os.environ['LOCAL_SIZE'])
+
+
+def get_local_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return int(os.environ['LOCAL_RANK'])
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def init_distributed_mode(args):
+    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ['WORLD_SIZE'])
+        args.gpu = int(os.environ['LOCAL_RANK'])
+        args.dist_url = 'env://'
+        os.environ['LOCAL_SIZE'] = str(torch.cuda.device_count())
+    elif 'SLURM_PROCID' in os.environ:
+        proc_id = int(os.environ['SLURM_PROCID'])
+        ntasks = int(os.environ['SLURM_NTASKS'])
+        node_list = os.environ['SLURM_NODELIST']
+        num_gpus = torch.cuda.device_count()
+        addr = subprocess.getoutput(
+            'scontrol show hostname {} | head -n1'.format(node_list))
+        os.environ['MASTER_PORT'] = os.environ.get('MASTER_PORT', '29500')
+        os.environ['MASTER_ADDR'] = addr
+        os.environ['WORLD_SIZE'] = str(ntasks)
+        os.environ['RANK'] = str(proc_id)
+        os.environ['LOCAL_RANK'] = str(proc_id % num_gpus)
+        os.environ['LOCAL_SIZE'] = str(num_gpus)
+        args.dist_url = 'env://'
+        args.world_size = ntasks
+        args.rank = proc_id
+        args.gpu = proc_id % num_gpus
+    else:
+        print('Not using distributed mode')
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = 'nccl'
+    print('| distributed init (rank {}): {}'.format(
+        args.rank, args.dist_url), flush=True)
+    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
+                                         world_size=args.world_size, rank=args.rank)
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)
+
+
+@torch.no_grad()
+def accuracy(output, target, topk=(1,)):
+    """Computes the precision@k for the specified values of k"""
+    if target.numel() == 0:
+        return [torch.zeros([], device=output.device)]
+    maxk = max(topk)
+    batch_size = target.size(0)
+
+    _, pred = output.topk(maxk, 1, True, True)
+    pred = pred.t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+    res = []
+    for k in topk:
+        correct_k = correct[:k].view(-1).float().sum(0)
+        res.append(correct_k.mul_(100.0 / batch_size))
+    return res
+
+
+def interpolate(input, size=None, scale_factor=None, mode="nearest", align_corners=None):
+    # type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor
+    """
+    Equivalent to nn.functional.interpolate, but with support for empty batch sizes.
+    This will eventually be supported natively by PyTorch, and this
+    class can go away.
+    """
+    if float(torchvision.__version__[:3]) < 0.7:
+        if input.numel() > 0:
+            return torch.nn.functional.interpolate(
+                input, size, scale_factor, mode, align_corners
+            )
+
+        output_shape = _output_size(2, input, size, scale_factor)
+        output_shape = list(input.shape[:-2]) + list(output_shape)
+        if float(torchvision.__version__[:3]) < 0.5:
+            return _NewEmptyTensorOp.apply(input, output_shape)
+        return _new_empty_tensor(input, output_shape)
+    else:
+        return torchvision.ops.misc.interpolate(input, size, scale_factor, mode, align_corners)
+
+
+def get_total_grad_norm(parameters, norm_type=2):
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+    norm_type = float(norm_type)
+    device = parameters[0].grad.device
+    total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]),
+                            norm_type)
+    return total_norm
+
+def inverse_sigmoid(x, eps=1e-5):
+    x = x.clamp(min=0, max=1)
+    x1 = x.clamp(min=eps)
+    x2 = (1 - x).clamp(min=eps)
+    return torch.log(x1/x2)
+
+
+def to_pixel_coords(flow, h1, w1):
+    flow = (
+        torch.stack(
+            (
+                w1 * (flow[..., 0] + 1) / 2,
+                h1 * (flow[..., 1] + 1) / 2,
+            ),
+            axis=-1,
+        )
+    )
+    return flow
+
+def read_config(file_path):
+    with open(file_path, 'r') as file:
+        config = yaml.safe_load(file)
+    return config

imcui/third_party/rdd/README.md

@@ -0,0 +1,197 @@
+## RDD: Robust Feature Detector and Descriptor using Deformable Transformer (CVPR 2025)
+[Gonglin Chen](https://xtcpete.com/) · [Tianwen Fu](https://twfu.me/) · [Haiwei Chen](https://scholar.google.com/citations?user=LVWRssoAAAAJ&hl=en) · [Wenbin Teng](https://wbteng9526.github.io/) · [Hanyuan Xiao](https://corneliushsiao.github.io/index.html) · [Yajie Zhao](https://ict.usc.edu/about-us/leadership/research-leadership/yajie-zhao/)
+
+[Project Page](https://xtcpete.github.io/rdd/) 
+
+## Table of Contents
+- [Updates](#updates)
+- [Installation](#installation)
+- [Usage](#usage)
+  - [Inference](#inference)
+  - [Evaluation](#evaluation)
+  - [Training](#training)
+- [Citation](#citation)
+- [License](#license)
+- [Acknowledgements](#acknowledgements)
+
+## Updates
+
+- SfM reconstruction through [COLMAP](https://github.com/colmap/colmap.git) added. We provide a ready-to-use [notebook](./demo_sfm.ipynb) for a simple example. Code adopted from [hloc](https://github.com/cvg/Hierarchical-Localization.git).
+
+- Training code and new weights released.
+
+- We have updated the training code compared to what was described in the paper. In the original setup, the RDD was trained on the MegaDepth and Air-to-Ground datasets by resizing all images to the training resolution. In this release, we retrained RDD on MegaDepth only, using a combination of resizing and cropping, a strategy used by [ALIKE](https://github.com/Shiaoming/ALIKE). This change significantly improves robustness.
+
+<table>
+  <tr>
+    <th></th>
+    <th colspan="3">MegaDepth-1500</th>
+    <th colspan="3">MegaDepth-View</th>
+    <th colspan="3">Air-to-Ground</th>
+  </tr>
+  <tr>
+    <td></td>
+    <td>AUC 5&deg</td><td>AUC 10&deg</td><td>AUC 20&deg</td>
+    <td>AUC 5&deg</td><td>AUC 10&deg</td><td>AUC 20&deg</td>
+    <td>AUC 5&deg</td><td>AUC 10&deg</td><td>AUC 20&deg</td>
+  </tr>
+  <tr>
+    <td>RDD-v2</td>
+    <td>52.4</td><td>68.5</td><td>80.1</td>
+    <td>52.0</td><td>67.1</td><td>78.2</td>
+    <td>45.8</td><td>58.6</td><td>71.0</td>
+  </tr>
+  <tr>
+    <td>RDD-v1</td>
+    <td>48.2</td><td>65.2</td><td>78.3</td>
+    <td>38.3</td><td>53.1</td><td>65.6</td>
+    <td>41.4</td><td>56.0</td><td>67.8</td>
+  </tr>
+  <tr>
+    <td>RDD-v2+LG</td>
+    <td>53.3</td><td>69.8</td><td>82.0</td>
+    <td>59.0</td><td>74.2</td><td>84.0</td>
+    <td>54.8</td><td>69.0</td><td>79.1</td>
+  </tr>
+  <tr>
+    <td>RDD-v1+LG</td>
+    <td>52.3</td><td>68.9</td><td>81.8</td>
+    <td>54.2</td><td>69.3</td><td>80.3</td>
+    <td>55.1</td><td>68.9</td><td>78.9</td>
+  </tr>
+</table>
+
+## Installation
+
+```bash
+git clone --recursive https://github.com/xtcpete/rdd
+cd RDD
+
+# Create conda env
+conda create -n rdd python=3.10 pip
+conda activate rdd
+
+# Install CUDA 
+conda install -c nvidia/label/cuda-11.8.0 cuda-toolkit
+# Install torch
+pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu118
+# Install all dependencies
+pip install -r requirements.txt
+# Compile custom operations
+cd ./RDD/models/ops
+pip install -e .
+```
+
+We provide the [download link](https://drive.google.com/drive/folders/1QgVaqm4iTUCqbWb7_Fi6mX09EHTId0oA?usp=sharing) to:
+  - the MegaDepth-1500 test set
+  - the MegaDepth-View test set
+  - the Air-to-Ground test set
+  - 2 pretrained models, RDD and LightGlue for matching RDD
+
+Create and unzip downloaded test data to the `data` folder.
+
+Create and add weights to the `weights` folder and you are ready to go.
+
+## Usage
+For your convenience, we provide a ready-to-use [notebook](./demo_matching.ipynb) for some examples.
+
+### Inference
+
+```python
+from RDD.RDD import build
+
+RDD_model = build()
+
+output = RDD_model.extract(torch.randn(1, 3, 480, 640))
+```
+
+### Evaluation
+
+Please note that due to the different GPU architectures and the stochastic nature of RANSAC, you may observe slightly different results; however, they should be very close to those reported in the paper. To reproduce the number in paper, use v1 weights instead.
+
+Results can be visualized by passing argument --plot
+
+**MegaDepth-1500**
+
+```bash
+# Sparse matching
+python ./benchmarks/mega_1500.py
+
+# Dense matching
+python ./benchmarks/mega_1500.py --method dense
+
+# LightGlue
+python ./benchmarks/mega_1500.py --method lightglue
+```
+
+**MegaDepth-View**
+
+```bash
+# Sparse matching
+python ./benchmarks/mega_view.py
+
+# Dense matching
+python ./benchmarks/mega_view.py --method dense
+
+# LightGlue
+python ./benchmarks/mega_view.py --method lightglue
+```
+
+**Air-to-Ground**
+
+```bash
+# Sparse matching
+python ./benchmarks/air_ground.py
+
+# Dense matching
+python ./benchmarks/air_ground.py --method dense
+
+# LightGlue
+python ./benchmarks/air_ground.py --method lightglue
+```
+
+### Training
+
+1. Download MegaDepth dataset using [download.sh](./data/megadepth/download.sh) and megadepth_indices from [LoFTR](https://github.com/zju3dv/LoFTR/blob/master/docs/TRAINING.md#download-datasets). Then the MegaDepth root folder should look like the following:
+```bash
+./data/megadepth/megadepth_indices # indices
+./data/megadepth/depth_undistorted # depth maps
+./data/megadepth/Undistorted_SfM # images and poses
+./data/megadepth/scene_info # indices for training LightGlue
+```
+2. Then you can train RDD in two steps; Descriptor first
+```bash
+# distributed training with 8 gpus
+python -m training.train --ckpt_save_path ./ckpt_descriptor --distributed --batch_size 32
+
+# single gpu 
+python -m training.train --ckpt_save_path ./ckpt_descriptor
+```
+and then Detector
+```bash
+python -m training.train --ckpt_save_path ./ckpt_detector --weights ./ckpt_descriptor/RDD_best.pth --train_detector --training_res 480
+```
+
+I am working on recollecting the Air-to-Ground dataset because of licensing issues.
+
+## Citation
+```
+@inproceedings{gonglin2025rdd,
+    title     = {RDD: Robust Feature Detector and Descriptor using Deformable Transformer},
+    author    = {Chen, Gonglin and Fu, Tianwen and Chen, Haiwei and Teng, Wenbin and Xiao, Hanyuan and Zhao, Yajie},
+    booktitle = {IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+    year      = {2025}
+}
+```
+
+
+## License
+[![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](LICENSE)
+
+## Acknowledgements
+
+We thank these great repositories: [ALIKE](https://github.com/Shiaoming/ALIKE), [LoFTR](https://github.com/zju3dv/LoFTR), [DeDoDe](https://github.com/Parskatt/DeDoDe), [XFeat](https://github.com/verlab/accelerated_features), [LightGlue](https://github.com/cvg/LightGlue), [Kornia](https://github.com/kornia/kornia), and [Deformable DETR](https://github.com/fundamentalvision/Deformable-DETR), and many other inspiring works in the community.
+
+LightGlue is trained with [Glue Factory](https://github.com/cvg/glue-factory).
+
+Supported by the Intelligence Advanced Research Projects Activity (IARPA) via Department of Interior/Interior Business Center (DOI/IBC) contract number 140D0423C0075. The U.S. Government is authorized to reproduce and distribute reprints for governmental purposes notwithstanding any copyright annotation thereon. Disclaimer: The views and conclusions contained herein are those of the authors and should not be interpreted as necessarily representing the official policies or endorsements, either expressed or implied, of IARPA, DOI/IBC, or the U.S. Government. We would like to thank Yayue Chen for her help with visualization.

imcui/third_party/rdd/benchmarks/air_ground.py

@@ -0,0 +1,247 @@
+import sys
+sys.path.append(".")
+import numpy as np
+import torch
+from PIL import Image
+import tqdm
+import cv2
+import argparse
+import matplotlib.pyplot as plt
+import matplotlib
+from RDD.RDD_helper import RDD_helper
+from RDD.RDD import build
+import os
+from benchmarks.utils import pose_auc, angle_error_vec, angle_error_mat, symmetric_epipolar_distance, compute_symmetrical_epipolar_errors, compute_pose_error, compute_relative_pose, estimate_pose, dynamic_alpha
+
+def make_matching_figure(
+        img0, img1, mkpts0, mkpts1, color,
+        kpts0=None, kpts1=None, text=[], dpi=75, path=None):
+    # draw image pair
+    assert mkpts0.shape[0] == mkpts1.shape[0], f'mkpts0: {mkpts0.shape[0]} v.s. mkpts1: {mkpts1.shape[0]}'
+    fig, axes = plt.subplots(1, 2, figsize=(10, 6), dpi=dpi)
+    axes[0].imshow(img0, cmap='gray')
+    axes[1].imshow(img1, cmap='gray')
+    for i in range(2):   # clear all frames
+        axes[i].get_yaxis().set_ticks([])
+        axes[i].get_xaxis().set_ticks([])
+        for spine in axes[i].spines.values():
+            spine.set_visible(False)
+    plt.tight_layout(pad=1)
+    
+    if kpts0 is not None:
+        assert kpts1 is not None
+        axes[0].scatter(kpts0[:, 0], kpts0[:, 1], c='w', s=2)
+        axes[1].scatter(kpts1[:, 0], kpts1[:, 1], c='w', s=2)
+
+    # draw matches
+    if mkpts0.shape[0] != 0 and mkpts1.shape[0] != 0:
+        fig.canvas.draw()
+        transFigure = fig.transFigure.inverted()
+        fkpts0 = transFigure.transform(axes[0].transData.transform(mkpts0))
+        fkpts1 = transFigure.transform(axes[1].transData.transform(mkpts1))
+        fig.lines = [matplotlib.lines.Line2D((fkpts0[i, 0], fkpts1[i, 0]),
+                                            (fkpts0[i, 1], fkpts1[i, 1]),
+                                            transform=fig.transFigure, c=color[i], linewidth=1)
+                                        for i in range(len(mkpts0))]
+        
+        axes[0].scatter(mkpts0[:, 0], mkpts0[:, 1], c=color, s=4)
+        axes[1].scatter(mkpts1[:, 0], mkpts1[:, 1], c=color, s=4)
+
+    # put txts
+    txt_color = 'k' if img0[:100, :200].mean() > 200 else 'w'
+    fig.text(
+        0.01, 0.99, '\n'.join(text), transform=fig.axes[0].transAxes,
+        fontsize=15, va='top', ha='left', color=txt_color)
+
+    # save or return figure
+    if path:
+        plt.savefig(str(path), bbox_inches='tight', pad_inches=0)
+        plt.close()
+    else:
+        return fig
+
+def error_colormap(err, thr, alpha=1.0):
+    assert alpha <= 1.0 and alpha > 0, f"Invaid alpha value: {alpha}"
+    x = 1 - np.clip(err / (thr * 2), 0, 1)
+    return np.clip(
+        np.stack([2-x*2, x*2, np.zeros_like(x), np.ones_like(x)*alpha], -1), 0, 1)
+
+def _make_evaluation_figure(img0, img1, kpts0, kpts1, epi_errs, e_t, e_R, alpha='dynamic', path=None):
+    conf_thr = 1e-4
+    
+    img0 = np.array(img0)
+    img1 = np.array(img1)
+    
+    kpts0 = kpts0
+    kpts1 = kpts1
+    
+    epi_errs = epi_errs.cpu().numpy()
+    correct_mask = epi_errs < conf_thr
+    precision = np.mean(correct_mask) if len(correct_mask) > 0 else 0
+    n_correct = np.sum(correct_mask)
+    
+    # recall might be larger than 1, since the calculation of conf_matrix_gt
+    # uses groundtruth depths and camera poses, but epipolar distance is used here.
+
+    # matching info
+    if alpha == 'dynamic':
+        alpha = dynamic_alpha(len(correct_mask))
+    color = error_colormap(epi_errs, conf_thr, alpha=alpha)
+    
+    text = [
+        f'#Matches {len(kpts0)}',
+        f'Precision({conf_thr:.2e}) ({100 * precision:.1f}%): {n_correct}/{len(kpts0)}',
+        f'e_t: {e_t:.2f} | e_R: {e_R:.2f}',
+    ]
+    
+    # make the figure
+    figure = make_matching_figure(img0, img1, kpts0, kpts1,
+                                  color, text=text, path=path)
+    return figure
+
+class AirGroundPoseMNNBenchmark:
+    def __init__(self, data_root="./data/air_ground", scene_names = None) -> None:
+        if scene_names is None:
+            self.scene_names = [
+                "indices.npz",
+            ]
+            # self.scene_names = ["0022_0.5_0.7.npz",]
+        else:
+            self.scene_names = scene_names
+        self.scenes = [
+            np.load(f"{data_root}/{scene}", allow_pickle=True)
+            for scene in self.scene_names
+        ]
+        self.data_root = data_root
+
+    def benchmark(self, model_helper, model_name = None, scale_intrinsics = False, calibrated = True, plot_every_iter=10, plot=False, method='sparse'):
+        with torch.no_grad():
+            data_root = self.data_root
+            tot_e_t, tot_e_R, tot_e_pose = [], [], []
+            thresholds = [5, 10, 20]
+            for scene_ind in range(len(self.scenes)):
+                import os
+                scene_name = os.path.splitext(self.scene_names[scene_ind])[0]
+                scene = self.scenes[scene_ind]
+                indices = scene['pair_info']
+                idx = 0
+                for pair in tqdm.tqdm(indices):
+                    
+                    pairs = pair['pair_names']
+                    K0 = pair['intrinsic'][0].copy().astype(np.float32)
+                    T0 = pair['pose'][0].copy().astype(np.float32)
+                    R0, t0 = T0[:3, :3], T0[:3, 3]
+                    K1 = pair['intrinsic'][1].copy().astype(np.float32)    
+                    T1 = pair['pose'][1].copy().astype(np.float32)
+                    R1, t1 = T1[:3, :3], T1[:3, 3]
+                    R, t = compute_relative_pose(R0, t0, R1, t1)
+                    T0_to_1 = np.concatenate((R,t[:,None]), axis=-1)
+                    im_A_path = f"{data_root}/images/{pairs[0]}"
+                    im_B_path = f"{data_root}/images/{pairs[1]}"
+                    
+                    im_A = cv2.imread(im_A_path)
+                    im_B = cv2.imread(im_B_path)
+                    
+                    if method == 'dense':
+                        kpts0, kpts1, conf = model_helper.match_dense(im_A, im_B, thr=0.01, resize=1600)
+                    elif method == 'lightglue':
+                        kpts0, kpts1, conf = model_helper.match_lg(im_A, im_B, thr=0.01, resize=1600)
+                    elif method == 'sparse':
+                        kpts0, kpts1, conf = model_helper.match(im_A, im_B, thr=0.01, resize=1600)
+                    else:
+                        raise ValueError(f"Invalid method {method}")
+
+                    im_A = Image.open(im_A_path)
+                    w0, h0 = im_A.size
+                    im_B = Image.open(im_B_path)
+                    w1, h1 = im_B.size
+                    if scale_intrinsics:
+                        scale0 = 840 / max(w0, h0)
+                        scale1 = 840 / max(w1, h1)
+                        w0, h0 = scale0 * w0, scale0 * h0
+                        w1, h1 = scale1 * w1, scale1 * h1
+                        K0, K1 = K0.copy(), K1.copy()
+                        K0[:2] = K0[:2] * scale0
+                        K1[:2] = K1[:2] * scale1
+                    
+                    threshold = 0.5 
+                    if calibrated:
+                        norm_threshold = threshold / (np.mean(np.abs(K0[:2, :2])) + np.mean(np.abs(K1[:2, :2])))
+                        ret = estimate_pose(
+                            kpts0,
+                            kpts1,
+                            K0,
+                            K1,
+                            norm_threshold,
+                            conf=0.99999,
+                        )
+                    if ret is not None:
+                        R_est, t_est, mask = ret
+                        T0_to_1_est = np.concatenate((R_est, t_est), axis=-1)  #
+                        T0_to_1 = np.concatenate((R, t[:,None]), axis=-1)
+                        e_t, e_R = compute_pose_error(T0_to_1_est, R, t)
+                        
+                        epi_errs = compute_symmetrical_epipolar_errors(T0_to_1, kpts0, kpts1, K0, K1)
+                        if scene_ind % plot_every_iter == 0 and plot:
+
+                            if not os.path.exists(f'outputs/air_ground/{model_name}_{method}'):
+                                os.mkdir(f'outputs/air_ground/{model_name}_{method}')
+                            name = f'outputs/air_ground/{model_name}_{method}/{scene_name}_{idx}.png'
+                            _make_evaluation_figure(im_A, im_B, kpts0, kpts1, epi_errs, e_t, e_R, path=name)
+                        e_pose = max(e_t, e_R)
+                        
+                        tot_e_t.append(e_t)
+                        tot_e_R.append(e_R)
+                        tot_e_pose.append(e_pose)
+                    idx += 1
+                    
+            tot_e_pose = np.array(tot_e_pose)
+            auc = pose_auc(tot_e_pose, thresholds)
+            acc_5 = (tot_e_pose < 5).mean()
+            acc_10 = (tot_e_pose < 10).mean()
+            acc_15 = (tot_e_pose < 15).mean()
+            acc_20 = (tot_e_pose < 20).mean()
+            map_5 = acc_5
+            map_10 = np.mean([acc_5, acc_10])
+            map_20 = np.mean([acc_5, acc_10, acc_15, acc_20])
+            print(f"{model_name} auc: {auc}")
+            return {
+                "auc_5": auc[0],
+                "auc_10": auc[1],
+                "auc_20": auc[2],
+                "map_5": map_5,
+                "map_10": map_10,
+                "map_20": map_20,
+            }
+            
+    
+            
+def parse_arguments():
+    parser = argparse.ArgumentParser(description="Testing script.")
+    
+    parser.add_argument("--data_root", type=str, default="./data/air_ground", help="Path to the Air-to-Ground test dataset.")
+
+    parser.add_argument("--weights", type=str, default="./weights/RDD-v2.pth", help="Path to the model checkpoint.")
+
+    parser.add_argument("--plot", action="store_true", help="Whether to plot the results.")
+
+    parser.add_argument("--method", type=str, default="sparse", help="Method for matching.")
+    
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_arguments()    
+    
+    if not os.path.exists('outputs'):
+        os.mkdir('outputs')
+    if not os.path.exists(f'outputs/air_ground'):
+        os.mkdir(f'outputs/air_ground')
+    model = build(weights=args.weights)
+    benchmark = AirGroundPoseMNNBenchmark(data_root=args.data_root)
+    model.eval()
+    model_helper = RDD_helper(model)
+    with torch.no_grad():
+        method = args.method
+        out = benchmark.benchmark(model_helper, model_name='RDD', plot_every_iter=1, plot=args.plot, method=method)
+        with open(f'outputs/air_ground/RDD_{method}.txt', 'w') as f:
+            f.write(str(out))

imcui/third_party/rdd/benchmarks/mega_1500.py

@@ -0,0 +1,255 @@
+import sys
+sys.path.append(".")
+import numpy as np
+import torch
+from PIL import Image
+import tqdm
+import cv2
+import argparse
+from RDD.RDD_helper import RDD_helper
+from RDD.RDD import build
+import matplotlib.pyplot as plt
+import matplotlib
+import os
+from benchmarks.utils import pose_auc, angle_error_vec, angle_error_mat, symmetric_epipolar_distance, compute_symmetrical_epipolar_errors, compute_pose_error, compute_relative_pose, estimate_pose, dynamic_alpha
+
+def make_matching_figure(
+        img0, img1, mkpts0, mkpts1, color,
+        kpts0=None, kpts1=None, text=[], dpi=75, path=None):
+    # draw image pair
+    assert mkpts0.shape[0] == mkpts1.shape[0], f'mkpts0: {mkpts0.shape[0]} v.s. mkpts1: {mkpts1.shape[0]}'
+    fig, axes = plt.subplots(1, 2, figsize=(10, 6), dpi=dpi)
+    axes[0].imshow(img0, cmap='gray')
+    axes[1].imshow(img1, cmap='gray')
+    for i in range(2):   # clear all frames
+        axes[i].get_yaxis().set_ticks([])
+        axes[i].get_xaxis().set_ticks([])
+        for spine in axes[i].spines.values():
+            spine.set_visible(False)
+    plt.tight_layout(pad=1)
+    
+    if kpts0 is not None:
+        assert kpts1 is not None
+        axes[0].scatter(kpts0[:, 0], kpts0[:, 1], c='w', s=2)
+        axes[1].scatter(kpts1[:, 0], kpts1[:, 1], c='w', s=2)
+
+    # draw matches
+    if mkpts0.shape[0] != 0 and mkpts1.shape[0] != 0:
+        fig.canvas.draw()
+        transFigure = fig.transFigure.inverted()
+        fkpts0 = transFigure.transform(axes[0].transData.transform(mkpts0))
+        fkpts1 = transFigure.transform(axes[1].transData.transform(mkpts1))
+        fig.lines = [matplotlib.lines.Line2D((fkpts0[i, 0], fkpts1[i, 0]),
+                                            (fkpts0[i, 1], fkpts1[i, 1]),
+                                            transform=fig.transFigure, c=color[i], linewidth=1)
+                                        for i in range(len(mkpts0))]
+        
+        axes[0].scatter(mkpts0[:, 0], mkpts0[:, 1], c=color, s=4)
+        axes[1].scatter(mkpts1[:, 0], mkpts1[:, 1], c=color, s=4)
+
+    # put txts
+    txt_color = 'k' if img0[:100, :200].mean() > 200 else 'w'
+    fig.text(
+        0.01, 0.99, '\n'.join(text), transform=fig.axes[0].transAxes,
+        fontsize=15, va='top', ha='left', color=txt_color)
+
+    # save or return figure
+    if path:
+        plt.savefig(str(path), bbox_inches='tight', pad_inches=0)
+        plt.close()
+    else:
+        return fig
+
+def error_colormap(err, thr, alpha=1.0):
+    assert alpha <= 1.0 and alpha > 0, f"Invaid alpha value: {alpha}"
+    x = 1 - np.clip(err / (thr * 2), 0, 1)
+    return np.clip(
+        np.stack([2-x*2, x*2, np.zeros_like(x), np.ones_like(x)*alpha], -1), 0, 1)
+
+def _make_evaluation_figure(img0, img1, kpts0, kpts1, epi_errs, e_t, e_R, alpha='dynamic', path=None):
+    conf_thr = 1e-4
+    
+    img0 = np.array(img0)
+    img1 = np.array(img1)
+    
+    kpts0 = kpts0
+    kpts1 = kpts1
+    
+    epi_errs = epi_errs.cpu().numpy()
+    correct_mask = epi_errs < conf_thr
+    precision = np.mean(correct_mask) if len(correct_mask) > 0 else 0
+    n_correct = np.sum(correct_mask)
+    
+    # recall might be larger than 1, since the calculation of conf_matrix_gt
+    # uses groundtruth depths and camera poses, but epipolar distance is used here.
+
+    # matching info
+    if alpha == 'dynamic':
+        alpha = dynamic_alpha(len(correct_mask))
+    color = error_colormap(epi_errs, conf_thr, alpha=alpha)
+    
+    text = [
+        f'#Matches {len(kpts0)}',
+        f'Precision({conf_thr:.2e}) ({100 * precision:.1f}%): {n_correct}/{len(kpts0)}',
+        f'e_t: {e_t:.2f} | e_R: {e_R:.2f}',
+    ]
+    
+    # make the figure
+    figure = make_matching_figure(img0, img1, kpts0, kpts1,
+                                  color, text=text, path=path)
+    return figure
+
+class MegaDepthPoseMNNBenchmark:
+    def __init__(self, data_root="./megadepth_test_1500", scene_names = None) -> None:
+        if scene_names is None:
+            self.scene_names = [
+                "0015_0.1_0.3.npz",
+                "0015_0.3_0.5.npz",
+                "0022_0.1_0.3.npz",
+                "0022_0.3_0.5.npz",
+                "0022_0.5_0.7.npz",
+            ]
+
+        else:
+            self.scene_names = scene_names
+        self.scenes = [
+            np.load(f"{data_root}/{scene}", allow_pickle=True)
+            for scene in self.scene_names
+        ]
+        self.data_root = data_root
+
+    def benchmark(self, model_helper, model_name = None, scale_intrinsics = False, calibrated = True, plot_every_iter=1, plot=False, method='sparse'):
+        
+        with torch.no_grad():
+            data_root = self.data_root
+            tot_e_t, tot_e_R, tot_e_pose = [], [], []
+            thresholds = [5, 10, 20]
+            for scene_ind in range(len(self.scenes)):
+                import os
+                scene_name = os.path.splitext(self.scene_names[scene_ind])[0]
+                print(f"Processing {scene_name}")
+                scene = self.scenes[scene_ind]
+                pairs = scene["pair_infos"]
+                intrinsics = scene["intrinsics"]
+                poses = scene["poses"]
+                im_paths = scene["image_paths"]
+                pair_inds = range(len(pairs))
+                for pairind in tqdm.tqdm(pair_inds):
+                    idx0, idx1 = pairs[pairind][0]
+                    K0 = intrinsics[idx0].copy()
+                    T0 = poses[idx0].copy()
+                    R0, t0 = T0[:3, :3], T0[:3, 3]
+                    K1 = intrinsics[idx1].copy()
+                    T1 = poses[idx1].copy()
+                    R1, t1 = T1[:3, :3], T1[:3, 3]
+                    R, t = compute_relative_pose(R0, t0, R1, t1)
+                    T0_to_1 = np.concatenate((R,t[:,None]), axis=-1)
+                    im_A_path = f"{data_root}/{im_paths[idx0]}"
+                    im_B_path = f"{data_root}/{im_paths[idx1]}"
+                    
+                    im_A = cv2.imread(im_A_path)
+                    im_B = cv2.imread(im_B_path)
+                
+                    if method == 'dense':
+                        kpts0, kpts1, conf = model_helper.match_dense(im_A, im_B, thr=0.01, resize=1600)
+                    elif method == 'lightglue':
+                        kpts0, kpts1, conf = model_helper.match_lg(im_A, im_B, thr=0.01, resize=1600)
+                    elif method == 'sparse':
+                        kpts0, kpts1, conf = model_helper.match(im_A, im_B, thr=0.01, resize=1600)
+                    else:
+                        kpts0, kpts1, conf = model_helper.match_3rd_party(im_A, im_B, thr=0.01, resize=1600, model=method)
+
+                    im_A = Image.open(im_A_path)
+                    w0, h0 = im_A.size
+                    im_B = Image.open(im_B_path)
+                    w1, h1 = im_B.size
+                    if scale_intrinsics:
+                        scale0 = 840 / max(w0, h0)
+                        scale1 = 840 / max(w1, h1)
+                        w0, h0 = scale0 * w0, scale0 * h0
+                        w1, h1 = scale1 * w1, scale1 * h1
+                        K0, K1 = K0.copy(), K1.copy()
+                        K0[:2] = K0[:2] * scale0
+                        K1[:2] = K1[:2] * scale1
+                    
+                        
+                    threshold = 0.5 
+                    if calibrated:
+                        norm_threshold = threshold / (np.mean(np.abs(K0[:2, :2])) + np.mean(np.abs(K1[:2, :2])))
+                        ret = estimate_pose(
+                            kpts0,
+                            kpts1,
+                            K0,
+                            K1,
+                            norm_threshold,
+                            conf=0.99999,
+                        )
+                    if ret is not None:
+                        R_est, t_est, mask = ret
+                        T0_to_1_est = np.concatenate((R_est, t_est), axis=-1)  #
+                        T0_to_1 = np.concatenate((R, t[:,None]), axis=-1)
+                        e_t, e_R = compute_pose_error(T0_to_1_est, R, t)
+                        
+                        epi_errs = compute_symmetrical_epipolar_errors(T0_to_1, kpts0, kpts1, K0, K1)
+                        if scene_ind % plot_every_iter == 0 and plot:
+
+                            if not os.path.exists(f'outputs/mega_1500/{model_name}_{method}'):
+                                os.mkdir(f'outputs/mega_1500/{model_name}_{method}')
+                            name = f'outputs/mega_1500/{model_name}_{method}/{scene_name}_{pairind}.png'
+                            _make_evaluation_figure(im_A, im_B, kpts0, kpts1, epi_errs, e_t, e_R, path=name)
+                        e_pose = max(e_t, e_R)
+                        
+                        tot_e_t.append(e_t)
+                        tot_e_R.append(e_R)
+                        tot_e_pose.append(e_pose)
+                            
+            tot_e_pose = np.array(tot_e_pose)
+            auc = pose_auc(tot_e_pose, thresholds)
+            acc_5 = (tot_e_pose < 5).mean()
+            acc_10 = (tot_e_pose < 10).mean()
+            acc_15 = (tot_e_pose < 15).mean()
+            acc_20 = (tot_e_pose < 20).mean()
+            map_5 = acc_5
+            map_10 = np.mean([acc_5, acc_10])
+            map_20 = np.mean([acc_5, acc_10, acc_15, acc_20])
+            print(f"{model_name} auc: {auc}")
+            return {
+                "auc_5": auc[0],
+                "auc_10": auc[1],
+                "auc_20": auc[2],
+                "map_5": map_5,
+                "map_10": map_10,
+                "map_20": map_20,
+            }
+            
+            
+def parse_arguments():
+    parser = argparse.ArgumentParser(description="Testing script.")
+    
+    parser.add_argument("--data_root", type=str, default="./data/megadepth_test_1500", help="Path to the MegaDepth dataset.")
+
+    parser.add_argument("--weights", type=str, default="./weights/RDD-v2.pth", help="Path to the model checkpoint.")
+
+    parser.add_argument("--plot", action="store_true", help="Whether to plot the results.")
+
+    parser.add_argument("--method", type=str, default="sparse", help="Method for matching.")
+    
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_arguments()    
+    if not os.path.exists('outputs'):
+        os.mkdir('outputs')
+
+    if not os.path.exists(f'outputs/mega_1500'):
+        os.mkdir(f'outputs/mega_1500')
+        
+    model = build(weights=args.weights)
+    benchmark = MegaDepthPoseMNNBenchmark(data_root=args.data_root)
+    model.eval()
+    model_helper = RDD_helper(model)
+    with torch.no_grad():
+        method = args.method
+        out = benchmark.benchmark(model_helper, model_name='RDD', plot_every_iter=1, plot=args.plot, method=method)
+        with open(f'outputs/mega_1500/RDD_{method}.txt', 'w') as f:
+            f.write(str(out))

imcui/third_party/rdd/benchmarks/mega_view.py

@@ -0,0 +1,250 @@
+import sys
+sys.path.append(".")
+import numpy as np
+import torch
+from PIL import Image
+import tqdm
+import cv2
+import argparse
+import matplotlib.pyplot as plt
+import matplotlib
+from RDD.RDD import build
+from RDD.RDD_helper import RDD_helper
+import os
+from benchmarks.utils import pose_auc, angle_error_vec, angle_error_mat, symmetric_epipolar_distance, compute_symmetrical_epipolar_errors, compute_pose_error, compute_relative_pose, estimate_pose, dynamic_alpha
+
+def make_matching_figure(
+        img0, img1, mkpts0, mkpts1, color,
+        kpts0=None, kpts1=None, text=[], dpi=75, path=None):
+    # draw image pair
+    assert mkpts0.shape[0] == mkpts1.shape[0], f'mkpts0: {mkpts0.shape[0]} v.s. mkpts1: {mkpts1.shape[0]}'
+    fig, axes = plt.subplots(1, 2, figsize=(10, 6), dpi=dpi)
+    axes[0].imshow(img0, cmap='gray')
+    axes[1].imshow(img1, cmap='gray')
+    for i in range(2):   # clear all frames
+        axes[i].get_yaxis().set_ticks([])
+        axes[i].get_xaxis().set_ticks([])
+        for spine in axes[i].spines.values():
+            spine.set_visible(False)
+    plt.tight_layout(pad=1)
+    
+    if kpts0 is not None:
+        assert kpts1 is not None
+        axes[0].scatter(kpts0[:, 0], kpts0[:, 1], c='w', s=2)
+        axes[1].scatter(kpts1[:, 0], kpts1[:, 1], c='w', s=2)
+
+    # draw matches
+    if mkpts0.shape[0] != 0 and mkpts1.shape[0] != 0:
+        fig.canvas.draw()
+        transFigure = fig.transFigure.inverted()
+        fkpts0 = transFigure.transform(axes[0].transData.transform(mkpts0))
+        fkpts1 = transFigure.transform(axes[1].transData.transform(mkpts1))
+        fig.lines = [matplotlib.lines.Line2D((fkpts0[i, 0], fkpts1[i, 0]),
+                                            (fkpts0[i, 1], fkpts1[i, 1]),
+                                            transform=fig.transFigure, c=color[i], linewidth=1)
+                                        for i in range(len(mkpts0))]
+        
+        axes[0].scatter(mkpts0[:, 0], mkpts0[:, 1], c=color, s=4)
+        axes[1].scatter(mkpts1[:, 0], mkpts1[:, 1], c=color, s=4)
+
+    # put txts
+    txt_color = 'k' if img0[:100, :200].mean() > 200 else 'w'
+    fig.text(
+        0.01, 0.99, '\n'.join(text), transform=fig.axes[0].transAxes,
+        fontsize=15, va='top', ha='left', color=txt_color)
+
+    # save or return figure
+    if path:
+        plt.savefig(str(path), bbox_inches='tight', pad_inches=0)
+        plt.close()
+    else:
+        return fig
+
+def error_colormap(err, thr, alpha=1.0):
+    assert alpha <= 1.0 and alpha > 0, f"Invaid alpha value: {alpha}"
+    x = 1 - np.clip(err / (thr * 2), 0, 1)
+    return np.clip(
+        np.stack([2-x*2, x*2, np.zeros_like(x), np.ones_like(x)*alpha], -1), 0, 1)
+
+def _make_evaluation_figure(img0, img1, kpts0, kpts1, epi_errs, e_t, e_R, alpha='dynamic', path=None):
+    conf_thr = 1e-4
+    
+    img0 = np.array(img0)
+    img1 = np.array(img1)
+    
+    kpts0 = kpts0
+    kpts1 = kpts1
+    
+    epi_errs = epi_errs.cpu().numpy()
+    correct_mask = epi_errs < conf_thr
+    precision = np.mean(correct_mask) if len(correct_mask) > 0 else 0
+    n_correct = np.sum(correct_mask)
+    
+    # recall might be larger than 1, since the calculation of conf_matrix_gt
+    # uses groundtruth depths and camera poses, but epipolar distance is used here.
+
+    # matching info
+    if alpha == 'dynamic':
+        alpha = dynamic_alpha(len(correct_mask))
+    color = error_colormap(epi_errs, conf_thr, alpha=alpha)
+    
+    text = [
+        f'#Matches {len(kpts0)}',
+        f'Precision({conf_thr:.2e}) ({100 * precision:.1f}%): {n_correct}/{len(kpts0)}',
+        f'e_t: {e_t:.2f} | e_R: {e_R:.2f}',
+    ]
+    
+    # make the figure
+    figure = make_matching_figure(img0, img1, kpts0, kpts1,
+                                  color, text=text, path=path)
+    return figure
+
+class MegaDepthPoseMNNBenchmark:
+    def __init__(self, data_root="./megadepth_test_1500", scene_names = None) -> None:
+        if scene_names is None:
+            self.scene_names = [
+                "hard_indices.npz",
+            ]
+            # self.scene_names = ["0022_0.5_0.7.npz",]
+        else:
+            self.scene_names = scene_names
+        self.scenes = [
+            np.load(f"{data_root}/{scene}", allow_pickle=True)
+            for scene in self.scene_names
+        ]
+        self.data_root = data_root
+
+    def benchmark(self, model_helper, model_name = None, scale_intrinsics = False, calibrated = True, plot_every_iter=1, plot=False, method='sparse'):
+        with torch.no_grad():
+            data_root = self.data_root
+            tot_e_t, tot_e_R, tot_e_pose = [], [], []
+            thresholds = [5, 10, 20]
+            for scene_ind in range(len(self.scenes)):
+                scene_name = os.path.splitext(self.scene_names[scene_ind])[0]
+                scene = self.scenes[scene_ind]
+                indices = scene['indices']
+                idx = 0
+    
+                for pair in tqdm.tqdm(indices):
+                    
+                    pairs = pair['pair_names']
+                    K0 = pair['intrisinic'][0].copy().astype(np.float32)
+                    T0 = pair['pose'][0].copy().astype(np.float32)
+                    R0, t0 = T0[:3, :3], T0[:3, 3]
+                    K1 = pair['intrisinic'][1].copy().astype(np.float32)    
+                    T1 = pair['pose'][1].copy().astype(np.float32)
+                    R1, t1 = T1[:3, :3], T1[:3, 3]
+                    R, t = compute_relative_pose(R0, t0, R1, t1)
+                    T0_to_1 = np.concatenate((R,t[:,None]), axis=-1)
+                    im_A_path = f"{data_root}/images/{pairs[0]}"
+                    im_B_path = f"{data_root}/images/{pairs[1]}"
+                    
+                    im_A = cv2.imread(im_A_path)
+                    im_B = cv2.imread(im_B_path)
+
+                    if method == 'dense':
+                        kpts0, kpts1, conf = model_helper.match_dense(im_A, im_B, thr=0.01, resize=1600)
+                    elif method == 'lightglue':
+                        kpts0, kpts1, conf = model_helper.match_lg(im_A, im_B, thr=0.01,  resize=1600)
+                    elif method == 'sparse':
+                        kpts0, kpts1, conf = model_helper.match(im_A, im_B, thr=0.01,  resize=1600)
+                    else:
+                        raise ValueError(f"Invalid method {method}")
+            
+                    im_A = Image.open(im_A_path)
+                    w0, h0 = im_A.size
+                    im_B = Image.open(im_B_path)
+                    w1, h1 = im_B.size
+                    
+                    if scale_intrinsics:
+                        scale0 = 840 / max(w0, h0)
+                        scale1 = 840 / max(w1, h1)
+                        w0, h0 = scale0 * w0, scale0 * h0
+                        w1, h1 = scale1 * w1, scale1 * h1
+                        K0, K1 = K0.copy(), K1.copy()
+                        K0[:2] = K0[:2] * scale0
+                        K1[:2] = K1[:2] * scale1
+                    
+                    threshold = 0.5 
+                    if calibrated:
+                        norm_threshold = threshold / (np.mean(np.abs(K0[:2, :2])) + np.mean(np.abs(K1[:2, :2])))
+                        ret = estimate_pose(
+                            kpts0,
+                            kpts1,
+                            K0,
+                            K1,
+                            norm_threshold,
+                            conf=0.99999,
+                        )
+                    if ret is not None:
+                        R_est, t_est, mask = ret
+                        T0_to_1_est = np.concatenate((R_est, t_est), axis=-1)  #
+                        T0_to_1 = np.concatenate((R, t[:,None]), axis=-1)
+                        e_t, e_R = compute_pose_error(T0_to_1_est, R, t)
+                        
+                        epi_errs = compute_symmetrical_epipolar_errors(T0_to_1, kpts0, kpts1, K0, K1)
+                        if scene_ind % plot_every_iter == 0 and plot:
+
+                            if not os.path.exists(f'outputs/mega_view/{model_name}_{method}'):
+                                os.mkdir(f'outputs/mega_view/{model_name}_{method}')
+                            name = f'outputs/mega_view/{model_name}_{method}/{scene_name}_{idx}.png'
+                            _make_evaluation_figure(im_A, im_B, kpts0, kpts1, epi_errs, e_t, e_R, path=name)
+                        e_pose = max(e_t, e_R)
+                        
+                        tot_e_t.append(e_t)
+                        tot_e_R.append(e_R)
+                        tot_e_pose.append(e_pose)
+                        idx += 1
+                        
+            tot_e_pose = np.array(tot_e_pose)
+            auc = pose_auc(tot_e_pose, thresholds)
+            acc_5 = (tot_e_pose < 5).mean()
+            acc_10 = (tot_e_pose < 10).mean()
+            acc_15 = (tot_e_pose < 15).mean()
+            acc_20 = (tot_e_pose < 20).mean()
+            map_5 = acc_5
+            map_10 = np.mean([acc_5, acc_10])
+            map_20 = np.mean([acc_5, acc_10, acc_15, acc_20])
+            print(f"{model_name} auc: {auc}")
+            return {
+                "auc_5": auc[0],
+                "auc_10": auc[1],
+                "auc_20": auc[2],
+                "map_5": map_5,
+                "map_10": map_10,
+                "map_20": map_20,
+            }
+            
+    
+            
+def parse_arguments():
+    parser = argparse.ArgumentParser(description="Testing script.")
+    
+    parser.add_argument("--data_root", type=str, default="./data/megadepth_view", help="Path to the MegaDepth dataset.")
+
+    parser.add_argument("--weights", type=str, default="./weights/RDD-v2.pth", help="Path to the model checkpoint.")
+
+    parser.add_argument("--plot", action="store_true", help="Whether to plot the results.")
+
+    parser.add_argument("--method", type=str, default="sparse", help="Method for matching.")
+    
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_arguments()    
+    if not os.path.exists('outputs'):
+        os.mkdir('outputs')
+    
+    if not os.path.exists(f'outputs/mega_view'):
+        os.mkdir(f'outputs/mega_view')
+    model = build(weights=args.weights)
+    benchmark = MegaDepthPoseMNNBenchmark(data_root=args.data_root)
+    model.eval()
+    model_helper = RDD_helper(model)
+    with torch.no_grad():
+        method = args.method
+        out = benchmark.benchmark(model_helper, model_name='RDD', plot_every_iter=1, plot=args.plot, method=method)
+        with open(f'outputs/mega_view/RDD_{method}.txt', 'w') as f:
+            f.write(str(out))
+
+

imcui/third_party/rdd/benchmarks/utils.py

@@ -0,0 +1,112 @@
+import numpy as np
+import torch
+from kornia.geometry.epipolar import numeric
+from kornia.geometry.conversions import convert_points_to_homogeneous
+import cv2
+
+def pose_auc(errors, thresholds):
+    sort_idx = np.argsort(errors)
+    errors = np.array(errors.copy())[sort_idx]
+    recall = (np.arange(len(errors)) + 1) / len(errors)
+    errors = np.r_[0.0, errors]
+    recall = np.r_[0.0, recall]
+    aucs = []
+    for t in thresholds:
+        last_index = np.searchsorted(errors, t)
+        r = np.r_[recall[:last_index], recall[last_index - 1]]
+        e = np.r_[errors[:last_index], t]
+        aucs.append(np.trapz(r, x=e) / t)
+    return aucs
+
+def angle_error_vec(v1, v2):
+    n = np.linalg.norm(v1) * np.linalg.norm(v2)
+    return np.rad2deg(np.arccos(np.clip(np.dot(v1, v2) / n, -1.0, 1.0)))
+
+def angle_error_mat(R1, R2):
+    cos = (np.trace(np.dot(R1.T, R2)) - 1) / 2
+    cos = np.clip(cos, -1.0, 1.0)  # numercial errors can make it out of bounds
+    return np.rad2deg(np.abs(np.arccos(cos)))
+
+def symmetric_epipolar_distance(pts0, pts1, E, K0, K1):
+    """Squared symmetric epipolar distance.
+    This can be seen as a biased estimation of the reprojection error.
+    Args:
+        pts0 (torch.Tensor): [N, 2]
+        E (torch.Tensor): [3, 3]
+    """
+    pts0 = (pts0 - K0[[0, 1], [2, 2]][None]) / K0[[0, 1], [0, 1]][None]
+    pts1 = (pts1 - K1[[0, 1], [2, 2]][None]) / K1[[0, 1], [0, 1]][None]
+    pts0 = convert_points_to_homogeneous(pts0)
+    pts1 = convert_points_to_homogeneous(pts1)
+
+    Ep0 = pts0 @ E.T  # [N, 3]
+    p1Ep0 = torch.sum(pts1 * Ep0, -1)  # [N,]
+    Etp1 = pts1 @ E  # [N, 3]
+
+    d = p1Ep0**2 * (1.0 / (Ep0[:, 0]**2 + Ep0[:, 1]**2) + 1.0 / (Etp1[:, 0]**2 + Etp1[:, 1]**2))  # N
+    return d
+
+def compute_symmetrical_epipolar_errors(T_0to1, pts0, pts1, K0, K1, device='cuda'):
+    """ 
+    Update:
+        data (dict):{"epi_errs": [M]}
+    """
+    pts0 = torch.tensor(pts0, device=device)
+    pts1 = torch.tensor(pts1, device=device)
+    K0 = torch.tensor(K0, device=device)
+    K1 = torch.tensor(K1, device=device)
+    T_0to1 = torch.tensor(T_0to1, device=device)
+    Tx = numeric.cross_product_matrix(T_0to1[:3, 3])
+    E_mat = Tx @ T_0to1[:3, :3]
+    
+    epi_err = symmetric_epipolar_distance(pts0, pts1, E_mat, K0, K1)
+    return epi_err
+
+def compute_pose_error(T_0to1, R, t):
+    R_gt = T_0to1[:3, :3]
+    t_gt = T_0to1[:3, 3]
+    error_t = angle_error_vec(t.squeeze(), t_gt)
+    error_t = np.minimum(error_t, 180 - error_t)  # ambiguity of E estimation
+    error_R = angle_error_mat(R, R_gt)
+    return error_t, error_R
+
+def compute_relative_pose(R1, t1, R2, t2):
+    rots = R2 @ (R1.T)
+    trans = -rots @ t1 + t2
+    return rots, trans
+
+def estimate_pose(kpts0, kpts1, K0, K1, norm_thresh, conf=0.99999):
+    if len(kpts0) < 5:
+        return None
+    K0inv = np.linalg.inv(K0[:2,:2])
+    K1inv = np.linalg.inv(K1[:2,:2])
+
+    kpts0 = (K0inv @ (kpts0-K0[None,:2,2]).T).T 
+    kpts1 = (K1inv @ (kpts1-K1[None,:2,2]).T).T
+    E, mask = cv2.findEssentialMat(
+        kpts0, kpts1, np.eye(3), threshold=norm_thresh, prob=conf
+    )
+
+    ret = None
+    if E is not None:
+        best_num_inliers = 0
+
+        for _E in np.split(E, len(E) / 3):
+            n, R, t, _ = cv2.recoverPose(_E, kpts0, kpts1, np.eye(3), 1e9, mask=mask)
+            if n > best_num_inliers:
+                best_num_inliers = n
+                ret = (R, t, mask.ravel() > 0)
+    return ret
+
+def dynamic_alpha(n_matches,
+                  milestones=[0, 300, 1000, 2000],
+                  alphas=[1.0, 0.8, 0.4, 0.2]):
+    if n_matches == 0:
+        return 1.0
+    ranges = list(zip(alphas, alphas[1:] + [None]))
+    loc = bisect.bisect_right(milestones, n_matches) - 1
+    _range = ranges[loc]
+    if _range[1] is None:
+        return _range[0]
+    return _range[1] + (milestones[loc + 1] - n_matches) / (
+        milestones[loc + 1] - milestones[loc]) * (_range[0] - _range[1])

imcui/third_party/rdd/configs/default.yaml

@@ -0,0 +1,19 @@
+activation: relu
+block_dims:
+- 8
+- 16
+- 32
+- 64
+d_model: 256
+detection_threshold: 0.1
+device: cuda
+dim_feedforward: 1024
+dropout: 0.1
+enc_n_points: 8
+hidden_dim: 256
+lr_backbone: 2.0e-05
+nhead: 8
+num_encoder_layers: 4
+num_feature_levels: 5
+top_k: 4096
+train_detector: False