Make volume loading more flexible (#320)

* Makes centering optional
* Adds the option to pass a pretransform to the volume's affine matrix

Author: eigenvivek

diffdrr/data.py

@@ -37,13 +37,18 @@ def load_example_ct(
     )
 
 # %% ../notebooks/api/03_data.ipynb 6
+from .pose import RigidTransform
+
+
 def read(
     volume: str | Path | ScalarImage,  # CT volume
     labelmap: str | Path | LabelMap = None,  # Labelmap for the CT volume
     labels: int | list = None,  # Labels from the mask of structures to render
     orientation: str | None = "AP",  # Frame-of-reference change
     bone_attenuation_multiplier: float = 1.0,  # Scalar multiplier on density of high attenuation voxels
     fiducials: torch.Tensor = None,  # 3D fiducials in world coordinates
+    transform: RigidTransform = None,  # RigidTransform to apply to the volume's affine
+    center_volume: bool = True,  # Move the volume's isocenter to the world origin
     **kwargs,  # Any additional information to be stored in the torchio.Subject
 ) -> Subject:
     """
@@ -57,10 +62,6 @@ def read(
     else:
         volume = ScalarImage(volume)
 
-    # Convert the volume to density
-    density = transform_hu_to_density(volume.data, bone_attenuation_multiplier)
-    density = ScalarImage(tensor=density, affine=volume.affine)
-
     # Read the mask if passed
     if labelmap is not None:
         if isinstance(labelmap, LabelMap):
@@ -71,6 +72,15 @@ def read(
     else:
         mask = None
 
+    # Optionally apply transform
+    if transform is not None:
+        T = transform.matrix[0].numpy()
+        volume = ScalarImage(tensor=volume.data, affine=T.dot(volume.affine))
+
+    # Convert the volume to density
+    density = transform_hu_to_density(volume.data, bone_attenuation_multiplier)
+    density = ScalarImage(tensor=density, affine=volume.affine)
+
     # Frame-of-reference change
     if orientation == "AP":
         # Rotates the C-arm about the x-axis by 90 degrees
@@ -117,9 +127,9 @@ def read(
         **kwargs,
     )
 
-    # Canonicalize the images by converting to RAS+ and moving the
-    # Subject's isocenter to the origin in world coordinates
-    subject = canonicalize(subject)
+    # Move the subject's isocenter to the origin in world coordinates
+    if center_volume:
+        subject = canonicalize(subject)
 
     # Apply mask
     if labels is not None:

notebooks/api/03_data.ipynb

@@ -108,13 +108,18 @@
    "outputs": [],
    "source": [
     "#| export\n",
+    "from diffdrr.pose import RigidTransform\n",
+    "\n",
+    "\n",
     "def read(\n",
     "    volume: str | Path | ScalarImage,  # CT volume\n",
     "    labelmap: str | Path | LabelMap = None,  # Labelmap for the CT volume\n",
     "    labels: int | list = None,  # Labels from the mask of structures to render\n",
     "    orientation: str | None = \"AP\",  # Frame-of-reference change\n",
     "    bone_attenuation_multiplier: float = 1.0,  # Scalar multiplier on density of high attenuation voxels\n",
     "    fiducials: torch.Tensor = None,  # 3D fiducials in world coordinates\n",
+    "    transform: RigidTransform = None,  # RigidTransform to apply to the volume's affine\n",
+    "    center_volume: bool = True,  # Move the volume's isocenter to the world origin\n",
     "    **kwargs,  # Any additional information to be stored in the torchio.Subject\n",
     ") -> Subject:\n",
     "    \"\"\"\n",
@@ -128,10 +133,6 @@
     "    else:\n",
     "        volume = ScalarImage(volume)\n",
     "\n",
-    "    # Convert the volume to density\n",
-    "    density = transform_hu_to_density(volume.data, bone_attenuation_multiplier)\n",
-    "    density = ScalarImage(tensor=density, affine=volume.affine)\n",
-    "\n",
     "    # Read the mask if passed\n",
     "    if labelmap is not None:\n",
     "        if isinstance(labelmap, LabelMap):\n",
@@ -142,6 +143,15 @@
     "    else:\n",
     "        mask = None\n",
     "\n",
+    "    # Optionally apply transform\n",
+    "    if transform is not None:\n",
+    "        T = transform.matrix[0].numpy()\n",
+    "        volume = ScalarImage(tensor=volume.data, affine=T.dot(volume.affine))\n",
+    "    \n",
+    "    # Convert the volume to density\n",
+    "    density = transform_hu_to_density(volume.data, bone_attenuation_multiplier)\n",
+    "    density = ScalarImage(tensor=density, affine=volume.affine)\n",
+    "\n",
     "    # Frame-of-reference change\n",
     "    if orientation == \"AP\":\n",
     "        # Rotates the C-arm about the x-axis by 90 degrees\n",
@@ -188,9 +198,9 @@
     "        **kwargs,\n",
     "    )\n",
     "\n",
-    "    # Canonicalize the images by converting to RAS+ and moving the\n",
-    "    # Subject's isocenter to the origin in world coordinates\n",
-    "    subject = canonicalize(subject)\n",
+    "    # Move the subject's isocenter to the origin in world coordinates\n",
+    "    if center_volume:\n",
+    "        subject = canonicalize(subject)\n",
     "\n",
     "    # Apply mask\n",
     "    if labels is not None:\n",