Merge pull request #160 from eigenvivek/eps

Expose InstanceNorm eps argument

Author: eigenvivek

diffdrr/metrics.py

@@ -13,9 +13,9 @@
 class NormalizedCrossCorrelation2d(torch.nn.Module):
     """Compute Normalized Cross Correlation between two batches of images."""
 
-    def __init__(self, patch_size=None):
+    def __init__(self, patch_size=None, eps=1e-5):
         super().__init__()
-        self.norm = torch.nn.InstanceNorm2d(num_features=1)
+        self.norm = torch.nn.InstanceNorm2d(num_features=1, eps=eps)
         self.patch_size = patch_size
 
     def forward(self, x1, x2):
@@ -33,9 +33,9 @@ def forward(self, x1, x2):
 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]):
+    def __init__(self, patch_sizes=[None], patch_weights=[1.0], eps=1e-5):
         super().__init__()
-        self.norm = torch.nn.InstanceNorm2d(num_features=1)
+        self.norm = torch.nn.InstanceNorm2d(num_features=1, eps=eps)
 
         assert len(patch_sizes) == len(patch_weights), "Each scale must have a weight"
         self.nccs = [
@@ -61,8 +61,8 @@ def to_patches(x, patch_size):
 class GradientNormalizedCrossCorrelation2d(NormalizedCrossCorrelation2d):
     """Compute Normalized Cross Correlation between the image gradients of two batches of images."""
 
-    def __init__(self, patch_size=None, sigma=1.0):
-        super().__init__(patch_size)
+    def __init__(self, patch_size=None, sigma=1.0, **kwargs):
+        super().__init__(patch_size, **kwargs)
         self.sobel = Sobel(sigma)
 
     def forward(self, x1, x2):

notebooks/api/05_metrics.ipynb

@@ -58,9 +58,9 @@
     "class NormalizedCrossCorrelation2d(torch.nn.Module):\n",
     "    \"\"\"Compute Normalized Cross Correlation between two batches of images.\"\"\"\n",
     "\n",
-    "    def __init__(self, patch_size=None):\n",
+    "    def __init__(self, patch_size=None, eps=1e-5):\n",
     "        super().__init__()\n",
-    "        self.norm = torch.nn.InstanceNorm2d(num_features=1)\n",
+    "        self.norm = torch.nn.InstanceNorm2d(num_features=1, eps=eps)\n",
     "        self.patch_size = patch_size\n",
     "\n",
     "    def forward(self, x1, x2):\n",
@@ -78,9 +78,9 @@
     "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",
+    "    def __init__(self, patch_sizes=[None], patch_weights=[1.0], eps=1e-5):\n",
     "        super().__init__()\n",
-    "        self.norm = torch.nn.InstanceNorm2d(num_features=1)\n",
+    "        self.norm = torch.nn.InstanceNorm2d(num_features=1, eps=eps)\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",
@@ -120,8 +120,8 @@
     "class GradientNormalizedCrossCorrelation2d(NormalizedCrossCorrelation2d):\n",
     "    \"\"\"Compute Normalized Cross Correlation between the image gradients of two batches of images.\"\"\"\n",
     "\n",
-    "    def __init__(self, patch_size=None, sigma=1.0):\n",
-    "        super().__init__(patch_size)\n",
+    "    def __init__(self, patch_size=None, sigma=1.0, **kwargs):\n",
+    "        super().__init__(patch_size, **kwargs)\n",
     "        self.sobel = Sobel(sigma)\n",
     "\n",
     "    def forward(self, x1, x2):\n",
@@ -172,7 +172,7 @@
     {
      "data": {
       "text/plain": [
-       "tensor([-0.0190,  0.0077,  0.0057,  0.0140, -0.0191,  0.0089, -0.0021,  0.0083])"
+       "tensor([ 0.0220, -0.0021,  0.0235,  0.0114, -0.0079, -0.0089,  0.0051, -0.0099])"
       ]
      },
      "execution_count": null,
@@ -187,6 +187,9 @@
     "ncc = NormalizedCrossCorrelation2d()\n",
     "ncc(x1, x2)\n",
     "\n",
+    "ncc = NormalizedCrossCorrelation2d(eps=1e-1)\n",
+    "ncc(x1, x2)\n",
+    "\n",
     "ncc = NormalizedCrossCorrelation2d(patch_size=9)\n",
     "ncc(x1, x2)\n",
     "\n",