Improve documentation

diffdrrdata/deepfluoro.py

@@ -1,7 +1,7 @@
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/00_deepfluoro.ipynb.
 
 # %% auto 0
-__all__ = ['DeepFluoroDataset']
+__all__ = ['DeepFluoroDataset', 'preprocess', 'Transforms']
 
 # %% ../notebooks/00_deepfluoro.ipynb 3
 from pathlib import Path
@@ -18,7 +18,7 @@
 
 from .utils import load_file
 
-# %% ../notebooks/00_deepfluoro.ipynb 5
+# %% ../notebooks/00_deepfluoro.ipynb 6
 class DeepFluoroDataset(torch.utils.data.Dataset):
     """
     A `torch.utils.data.Dataset` that stores the imaging data for subjects
@@ -31,9 +31,9 @@ class DeepFluoroDataset(torch.utils.data.Dataset):
     def __init__(
         self,
         id_number: int,  # Subject ID in {1, ..., 6}
-        preprocess: bool = True,  # Preprocess raw X-rays
+        preprocess: bool = True,  # Convert X-rays from exponentiated to linear form
+        bone_attenuation_multiplier: float = 1.0,  # Scalar multiplier on density of high attenuation voxels (from `DiffDRR`, see [here](https://vivekg.dev/DiffDRR/tutorials/introduction.html#changing-the-appearance-of-the-rendered-drrs))
         batchless: bool = False,  # Return unbatched images and poses (e.g., to interface with a `torch.utils.data.DataLoader`)
-        bone_attenuation_multiplier: float = 1.0,  # Scalar multiplier on density of high attenuation voxels
     ):
         super().__init__()
 
@@ -111,7 +111,7 @@ def __getitem__(self, idx):
     def rot_180_for_up(self, idx):
         return self.projections[f"{idx:03d}/rot-180-for-up"][()]
 
-# %% ../notebooks/00_deepfluoro.ipynb 6
+# %% ../notebooks/00_deepfluoro.ipynb 8
 def parse_volume(subject, bone_attenuation_multiplier):
     # Get all parts of the volume
     volume = subject["vol/pixels"][:]
@@ -234,11 +234,28 @@ def load(id_number, bone_attenuation_multiplier):
         y0,
     )
 
-
-def preprocess(img, size=None, initial_energy=torch.tensor(65487.0)):
+# %% ../notebooks/00_deepfluoro.ipynb 9
+def preprocess(img, initial_energy=torch.tensor(65487.0)):
     """Convert X-ray fluoroscopy from the exponentiated form to the linear form."""
     img = center_crop(img, (1436, 1436))
     img = gaussian_blur(img, (5, 5), sigma=1.0)
     img = initial_energy.log() - img.log()
-    img = (img - img.min()) / (img.max() - img.min())
     return img
+
+# %% ../notebooks/00_deepfluoro.ipynb 11
+from torchvision.transforms import Compose, Lambda, Normalize, Resize
+
+
+class Transforms:
+    def __init__(self, height: int, eps: float = 1e-6):
+        """Standardize, resize, and normalize X-rays and DRRs before inputting to a deep learning model."""
+        self.transforms = Compose(
+            [
+                Lambda(lambda x: (x - x.min()) / (x.max() - x.min() + eps)),
+                Resize((height, height), antialias=True),
+                Normalize(mean=0.3080, std=0.1494),
+            ]
+        )
+
+    def __call__(self, x):
+        return self.transforms(x)