Implement mutual information image similarity (#322)

Author: eigenvivek

diffdrr/_modidx.py

@@ -56,6 +56,11 @@
                                                                                                       'diffdrr/metrics.py'),
                                  'diffdrr.metrics.MultiscaleNormalizedCrossCorrelation2d.forward': ( 'api/metrics.html#multiscalenormalizedcrosscorrelation2d.forward',
                                                                                                      'diffdrr/metrics.py'),
+                                 'diffdrr.metrics.MutualInformation': ('api/metrics.html#mutualinformation', 'diffdrr/metrics.py'),
+                                 'diffdrr.metrics.MutualInformation.__init__': ( 'api/metrics.html#mutualinformation.__init__',
+                                                                                 'diffdrr/metrics.py'),
+                                 'diffdrr.metrics.MutualInformation.forward': ( 'api/metrics.html#mutualinformation.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 @@
 
 # %% auto 0
 __all__ = ['NormalizedCrossCorrelation2d', 'MultiscaleNormalizedCrossCorrelation2d', 'GradientNormalizedCrossCorrelation2d',
-           'LogGeodesicSE3', 'DoubleGeodesicSE3']
+           'MutualInformation', 'LogGeodesicSE3', 'DoubleGeodesicSE3']
 
 # %% ../notebooks/api/05_metrics.ipynb 6
 from einops import rearrange
@@ -100,7 +100,42 @@ def __init__(self, patch_size=None, sigma=1.0, **kwargs):
     def forward(self, x1, x2):
         return super().forward(self.sobel(x1), self.sobel(x2))
 
-# %% ../notebooks/api/05_metrics.ipynb 14
+# %% ../notebooks/api/05_metrics.ipynb 11
+from kornia.enhance.histogram import marginal_pdf, joint_pdf
+
+
+class MutualInformation(torch.nn.Module):
+    """Mutual Information."""
+
+    def __init__(self, sigma=0.1, num_bins=256, epsilon=1e-10, normalize=True):
+        super().__init__()
+        self.register_buffer("sigma", torch.tensor(sigma))
+        self.register_buffer("bins", torch.linspace(0.0, 1.0, num_bins))
+        self.epsilon = epsilon
+        self.normalize = normalize
+
+    def forward(self, x1, x2):
+        assert x1.shape == x2.shape
+        B, C, H, W = x1.shape
+
+        x1 = x1.view(B, H * W, C)
+        x2 = x2.view(B, H * W, C)
+
+        pdf_x1, kernel_values1 = marginal_pdf(x1, self.bins, self.sigma, self.epsilon)
+        pdf_x2, kernel_values2 = marginal_pdf(x2, self.bins, self.sigma, self.epsilon)
+        pdf_x1x2 = joint_pdf(kernel_values1, kernel_values2)
+
+        H_x1 = -(pdf_x1 * (pdf_x1 + self.epsilon).log2()).sum(dim=1)
+        H_x2 = -(pdf_x2 * (pdf_x2 + self.epsilon).log2()).sum(dim=1)
+        H_x1x2 = -(pdf_x1x2 * (pdf_x1x2 + self.epsilon).log2()).sum(dim=(1, 2))
+
+        mutual_information = H_x1 + H_x2 - H_x1x2
+        if self.normalize:
+            mutual_information = 2 * mutual_information / (H_x1 + H_x2)
+
+        return mutual_information
+
+# %% ../notebooks/api/05_metrics.ipynb 15
 from .pose import RigidTransform, convert
 
 
@@ -119,7 +154,7 @@ def forward(
     ) -> Float[torch.Tensor, "b"]:
         return pose_2.compose(pose_1.inverse()).get_se3_log().norm(dim=1)
 
-# %% ../notebooks/api/05_metrics.ipynb 17
+# %% ../notebooks/api/05_metrics.ipynb 18
 from .pose import so3_log_map
 
 

notebooks/api/05_metrics.ipynb

@@ -203,6 +203,49 @@
     "        return super().forward(self.sobel(x1), self.sobel(x2))"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6e422fc1-8100-4120-b226-f6f54602fe3c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "from kornia.enhance.histogram import marginal_pdf, joint_pdf\n",
+    "\n",
+    "\n",
+    "class MutualInformation(torch.nn.Module):\n",
+    "    \"\"\"Mutual Information.\"\"\"\n",
+    "\n",
+    "    def __init__(self, sigma=0.1, num_bins=256, epsilon=1e-10, normalize=True):\n",
+    "        super().__init__()\n",
+    "        self.register_buffer(\"sigma\", torch.tensor(sigma))\n",
+    "        self.register_buffer(\"bins\", torch.linspace(0.0, 1.0, num_bins))\n",
+    "        self.epsilon = epsilon\n",
+    "        self.normalize = normalize\n",
+    "\n",
+    "    def forward(self, x1, x2):\n",
+    "        assert(x1.shape == x2.shape)\n",
+    "        B, C, H, W = x1.shape\n",
+    "\n",
+    "        x1 = x1.view(B, H * W, C)\n",
+    "        x2 = x2.view(B, H * W, C)\n",
+    "\n",
+    "        pdf_x1, kernel_values1 = marginal_pdf(x1, self.bins, self.sigma, self.epsilon)\n",
+    "        pdf_x2, kernel_values2 = marginal_pdf(x2, self.bins, self.sigma, self.epsilon)\n",
+    "        pdf_x1x2 = joint_pdf(kernel_values1, kernel_values2)\n",
+    "\n",
+    "        H_x1 = -(pdf_x1 * (pdf_x1 + self.epsilon).log2()).sum(dim=1)\n",
+    "        H_x2 = -(pdf_x2 * (pdf_x2 + self.epsilon).log2()).sum(dim=1)\n",
+    "        H_x1x2 = -(pdf_x1x2 * (pdf_x1x2 + self.epsilon).log2()).sum(dim=(1, 2))\n",
+    "\n",
+    "        mutual_information = H_x1 + H_x2 - H_x1x2\n",
+    "        if self.normalize:\n",
+    "            mutual_information = 2 * mutual_information / (H_x1 + H_x2)\n",
+    "\n",
+    "        return mutual_information"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -212,7 +255,7 @@
     {
      "data": {
       "text/plain": [
-       "tensor([-0.0019, -0.0004,  0.0035, -0.0198, -0.0078, -0.0175,  0.0171,  0.0019])"
+       "tensor([ 0.0002,  0.0006,  0.0003,  0.0005,  0.0003,  0.0005,  0.0012, -0.0001])"
       ]
      },
      "execution_count": null,
@@ -243,7 +286,10 @@
     "gncc(x1, x2)\n",
     "\n",
     "gncc = GradientNormalizedCrossCorrelation2d(patch_size=9)\n",
-    "gncc(x1, x2)"
+    "gncc(x1, x2)\n",
+    "\n",
+    "mi = MutualInformation()\n",
+    "mi(x1, x2)"
    ]
   },
   {