Compute raylength in world coordinates (#328)

Author: eigenvivek

diffdrr/drr.py

@@ -173,6 +173,10 @@ def forward(
             )
             img.append(partial)
         img = torch.cat(img, dim=-1)
+
+    # Multiply by the raylength (in world coordinate units)
+    img *= self.affine(target - source).norm(dim=-1).unsqueeze(1)
+
     return self.reshape_transform(img, batch_size=len(pose))
 
 # %% ../notebooks/api/00_drr.ipynb 11

diffdrr/renderers.py

@@ -82,9 +82,6 @@ def forward(
                 .scatter_add_(1, channels.transpose(-1, -2), img.transpose(-1, -2))
             )
 
-        # Multiply by ray length such that the proportion of attenuated energy is unitless
-        raylength = (target - source + self.eps).norm(dim=-1)
-        img *= raylength.unsqueeze(1)
         return img
 
 # %% ../notebooks/api/01_renderers.ipynb 8
@@ -202,6 +199,10 @@ def forward(
         # Sample the volume with trilinear interpolation
         img = _get_voxel(volume, xyzs, self.mode, align_corners=align_corners)
 
+        # Multiply by the step size to compute the rectangular rule for integration
+        step_size = (alphamax - alphamin) / (n_points - 1)
+        img = img * step_size
+
         # Handle optional masking
         if mask is None:
             img = img.sum(dim=-1).unsqueeze(1)
@@ -217,8 +218,4 @@ def forward(
                 .scatter_add_(1, channels.transpose(-1, -2), img.transpose(-1, -2))
             )
 
-        # Compute the line integral with the rectangular rule and return the DRR
-        raylength = (target - source + self.eps).norm(dim=-1).unsqueeze(1)
-        step_size = (alphamax - alphamin) / (n_points - 1)
-        img *= raylength * step_size
         return img

notebooks/api/00_drr.ipynb

@@ -195,7 +195,7 @@
     "        self.patch_size = patch_size\n",
     "        if self.patch_size is not None:\n",
     "            self.n_patches = (height * width) // (self.patch_size**2)\n",
-    "\n",
+    "        \n",
     "    def reshape_transform(self, img, batch_size):\n",
     "        if self.reshape:\n",
     "            if self.detector.n_subsample is None:\n",
@@ -292,6 +292,10 @@
     "            )\n",
     "            img.append(partial)\n",
     "        img = torch.cat(img, dim=-1)\n",
+    "    \n",
+    "    # Multiply by the raylength (in world coordinate units)\n",
+    "    img *= self.affine(target - source).norm(dim=-1).unsqueeze(1)\n",
+    "\n",
     "    return self.reshape_transform(img, batch_size=len(pose))"
    ]
   },

notebooks/api/01_renderers.ipynb

@@ -184,9 +184,6 @@
     "                .scatter_add_(1, channels.transpose(-1, -2), img.transpose(-1, -2))\n",
     "            )\n",
     "\n",
-    "        # Multiply by ray length such that the proportion of attenuated energy is unitless\n",
-    "        raylength = (target - source + self.eps).norm(dim=-1)\n",
-    "        img *= raylength.unsqueeze(1)\n",
     "        return img"
    ]
   },
@@ -332,6 +329,10 @@
     "\n",
     "        # Sample the volume with trilinear interpolation\n",
     "        img = _get_voxel(volume, xyzs, self.mode, align_corners=align_corners)\n",
+    "        \n",
+    "        # Multiply by the step size to compute the rectangular rule for integration\n",
+    "        step_size = (alphamax - alphamin) / (n_points - 1)\n",
+    "        img = img * step_size\n",
     "\n",
     "        # Handle optional masking\n",
     "        if mask is None:\n",
@@ -348,10 +349,6 @@
     "                .scatter_add_(1, channels.transpose(-1, -2), img.transpose(-1, -2))\n",
     "            )\n",
     "\n",
-    "        # Compute the line integral with the rectangular rule and return the DRR\n",
-    "        raylength = (target - source + self.eps).norm(dim=-1).unsqueeze(1)\n",
-    "        step_size = (alphamax - alphamin) / (n_points - 1)\n",
-    "        img *= raylength * step_size\n",
     "        return img"
    ]
   },