Add multiscale variant of NCC

Author: eigenvivek

diffdrr/_modidx.py

@@ -32,6 +32,12 @@
                                                                                                     'diffdrr/metrics.py'),
                                  'diffdrr.metrics.GradientNormalizedCrossCorrelation2d.forward': ( 'api/metrics.html#gradientnormalizedcrosscorrelation2d.forward',
                                                                                                    'diffdrr/metrics.py'),
+                                 'diffdrr.metrics.MultiscaleNormalizedCrossCorrelation2d': ( 'api/metrics.html#multiscalenormalizedcrosscorrelation2d',
+                                                                                             'diffdrr/metrics.py'),
+                                 'diffdrr.metrics.MultiscaleNormalizedCrossCorrelation2d.__init__': ( 'api/metrics.html#multiscalenormalizedcrosscorrelation2d.__init__',
+                                                                                                      'diffdrr/metrics.py'),
+                                 'diffdrr.metrics.MultiscaleNormalizedCrossCorrelation2d.forward': ( 'api/metrics.html#multiscalenormalizedcrosscorrelation2d.forward',
+                                                                                                     'diffdrr/metrics.py'),
                                  'diffdrr.metrics.NormalizedCrossCorrelation2d': ( 'api/metrics.html#normalizedcrosscorrelation2d',
                                                                                    'diffdrr/metrics.py'),
                                  'diffdrr.metrics.NormalizedCrossCorrelation2d.__init__': ( 'api/metrics.html#normalizedcrosscorrelation2d.__init__',

diffdrr/metrics.py

@@ -7,7 +7,7 @@
 import torch.nn as nn
 
 # %% auto 0
-__all__ = ['NormalizedCrossCorrelation2d', 'GradientNormalizedCrossCorrelation2d']
+__all__ = ['NormalizedCrossCorrelation2d', 'MultiscaleNormalizedCrossCorrelation2d', 'GradientNormalizedCrossCorrelation2d']
 
 # %% ../notebooks/api/05_metrics.ipynb 4
 class NormalizedCrossCorrelation2d(torch.nn.Module):
@@ -29,6 +29,26 @@ def forward(self, x1, x2):
         score /= c * h * w
         return score
 
+
+class MultiscaleNormalizedCrossCorrelation2d(torch.nn.Module):
+    """Compute Normalized Cross Correlation between two batches of images at multiple scales."""
+
+    def __init__(self, patch_sizes=[None], patch_weights=[1.0]):
+        super().__init__()
+        self.norm = torch.nn.InstanceNorm2d(num_features=1)
+
+        assert len(patch_sizes) == len(patch_weights), "Each scale must have a weight"
+        self.nccs = [
+            NormalizedCrossCorrelation2d(patch_size) for patch_size in patch_sizes
+        ]
+        self.patch_weights = patch_weights
+
+    def forward(self, x1, x2):
+        scores = []
+        for weight, ncc in zip(self.patch_weights, self.nccs):
+            scores.append(weight * ncc(x1, x2))
+        return torch.stack(scores, dim=0).sum(dim=0)
+
 # %% ../notebooks/api/05_metrics.ipynb 5
 from einops import rearrange
 

notebooks/api/05_metrics.ipynb

@@ -72,7 +72,25 @@
     "        x1, x2 = self.norm(x1), self.norm(x2)\n",
     "        score = torch.einsum(\"b...,b...->b\", x1, x2)\n",
     "        score /= c * h * w\n",
-    "        return score"
+    "        return score\n",
+    "\n",
+    "\n",
+    "class MultiscaleNormalizedCrossCorrelation2d(torch.nn.Module):\n",
+    "    \"\"\"Compute Normalized Cross Correlation between two batches of images at multiple scales.\"\"\"\n",
+    "\n",
+    "    def __init__(self, patch_sizes=[None], patch_weights=[1.0]):\n",
+    "        super().__init__()\n",
+    "        self.norm = torch.nn.InstanceNorm2d(num_features=1)\n",
+    "        \n",
+    "        assert len(patch_sizes) == len(patch_weights), \"Each scale must have a weight\"\n",
+    "        self.nccs = [NormalizedCrossCorrelation2d(patch_size) for patch_size in patch_sizes]\n",
+    "        self.patch_weights = patch_weights\n",
+    "\n",
+    "    def forward(self, x1, x2):\n",
+    "        scores = []\n",
+    "        for weight, ncc in zip(self.patch_weights, self.nccs):\n",
+    "            scores.append(weight * ncc(x1, x2))\n",
+    "        return torch.stack(scores, dim=0).sum(dim=0)"
    ]
   },
   {
@@ -154,7 +172,7 @@
     {
      "data": {
       "text/plain": [
-       "tensor([ 0.0039, -0.0092,  0.0008,  0.0022, -0.0049, -0.0204, -0.0088,  0.0056])"
+       "tensor([-0.0190,  0.0077,  0.0057,  0.0140, -0.0191,  0.0089, -0.0021,  0.0083])"
       ]
      },
      "execution_count": null,
@@ -172,6 +190,9 @@
     "ncc = NormalizedCrossCorrelation2d(patch_size=9)\n",
     "ncc(x1, x2)\n",
     "\n",
+    "msncc = MultiscaleNormalizedCrossCorrelation2d(patch_sizes=[9, None], patch_weights=[0.5, 0.5])\n",
+    "msncc(x1, x2)\n",
+    "\n",
     "gncc = GradientNormalizedCrossCorrelation2d()\n",
     "gncc(x1, x2)\n",
     "\n",