Make HU to density optional

eigenvivek · Mar 18, 2024 · 3bb683c · 3bb683c
1 parent 5a19681
commit 3bb683c
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 6 deletions.
diff --git a/diffdrr/data.py b/diffdrr/data.py
@@ -27,14 +27,19 @@ def __init__(
         affine,
         origin,
         spacing,
-        bone_attenuation_multiplier,
+        bone_attenuation_multiplier=1.0,
+        convert_hu_to_density=True,  # Convert Hounsfield units to linear attenuation coefficient
     ):
         self.volume = torch.from_numpy(volume)
         self.affine = torch.from_numpy(affine)
         self.origin = torch.tensor(origin)
         self.spacing = torch.tensor(spacing)
-        self.density = self.parse_density(self.volume, bone_attenuation_multiplier)
-        self.bone_attenuation_multiplier = bone_attenuation_multiplier
+
+        # Set convert_hu_to_density=False if volume is not a CT (e.g., a segmentation or MRI)
+        if convert_hu_to_density:
+            self.density = self.parse_density(self.volume, bone_attenuation_multiplier)
+        else:
+            self.density = self.volume
 
     @staticmethod
     def parse_density(volume, bone_attenuation_multiplier):

diff --git a/notebooks/api/03_data.ipynb b/notebooks/api/03_data.ipynb
@@ -80,14 +80,19 @@
     "        affine,\n",
     "        origin,\n",
     "        spacing,\n",
-    "        bone_attenuation_multiplier,\n",
+    "        bone_attenuation_multiplier=1.0,\n",
+    "        convert_hu_to_density=True,  # Convert Hounsfield units to linear attenuation coefficient\n",
     "    ):\n",
     "        self.volume = torch.from_numpy(volume)\n",
     "        self.affine = torch.from_numpy(affine)\n",
     "        self.origin = torch.tensor(origin)\n",
     "        self.spacing = torch.tensor(spacing)\n",
-    "        self.density = self.parse_density(self.volume, bone_attenuation_multiplier)\n",
-    "        self.bone_attenuation_multiplier = bone_attenuation_multiplier\n",
+    "\n",
+    "        # Set convert_hu_to_density=False if volume is not a CT (e.g., a segmentation or MRI)\n",
+    "        if convert_hu_to_density:\n",
+    "            self.density = self.parse_density(self.volume, bone_attenuation_multiplier)\n",
+    "        else:\n",
+    "            self.density = self.volume\n",
     "\n",
     "    @staticmethod\n",
     "    def parse_density(volume, bone_attenuation_multiplier):\n",