Fix Ljubljana preprocessing (#19)

* Fix Ljubljana preprocessing

* Remove unused inputs

* Update README

* Switch DeepFluoro orientation to PA

* Rotate pose

* Rerun notebooks

* Move dim flip

Author: eigenvivek

README.md

@@ -1,5 +1,5 @@
-# DiffDRR Datasets
-
+DiffDRR Datasets
+================
 
 <!-- WARNING: THIS FILE WAS AUTOGENERATED! DO NOT EDIT! -->
 
@@ -90,15 +90,14 @@ drr = DRR(
     deepfluoro.delx * subsample,
     x0=deepfluoro.x0,
     y0=deepfluoro.y0,
-    reverse_x_axis=False,
 )
 transform = Transforms(deepfluoro.height // subsample)
 
 # Render a DRR from the ground truth camera pose
 gt, pose = deepfluoro[0]
 img = drr(pose)
 gt, img = transform(gt), transform(img)
-plot_drr(torch.concat([gt, img]), title=["Downsampled X-ray", "DRR"])
+plot_drr(torch.concat([gt, img, gt - img]), title=["Downsampled X-ray", "DRR", "Difference"])
 plt.show()
 ```
 
@@ -141,7 +140,7 @@ drr = DRR(
     delx * subsample,
     width // subsample,
     dely * subsample,
-    x0=-x0,
+    x0=x0,
     y0=y0,
     reverse_x_axis=False,
 )
@@ -150,7 +149,7 @@ transform = Transforms(height // subsample, width // subsample)
 # Render a DRR from the ground truth camera pose
 img = drr(pose)
 gt, img = transform(gt), transform(img)
-plot_drr(torch.concat([gt, img]), title=["Downsampled X-ray", "DRR"])
+plot_drr(torch.concat([gt, img, gt - img]), title=["Downsampled X-ray", "DRR", "Difference"])
 plt.show()
 ```
 

diffdrrdata/deepfluoro.py

@@ -9,13 +9,14 @@
 import h5py
 import numpy as np
 import torch
-from diffdrr.data import read
-from diffdrr.pose import RigidTransform
-from diffdrr.utils import parse_intrinsic_matrix
 from torchio import LabelMap, ScalarImage, Subject
 from torchio.transforms.preprocessing import ToCanonical
 from torchvision.transforms.functional import center_crop, gaussian_blur
 
+from diffdrr.data import read
+from diffdrr.pose import RigidTransform
+from diffdrr.utils import parse_intrinsic_matrix
+
 from .utils import load_file
 
 # %% ../notebooks/00_deepfluoro.ipynb 6
@@ -64,7 +65,7 @@ def __init__(
             torch.tensor(
                 [
                     [0, 1, 0, 0],
-                    [1, 0, 0, 0],
+                    [-1, 0, 0, 0],
                     [0, 0, -1, 0],
                     [0, 0, 0, 1],
                 ]
@@ -93,10 +94,12 @@ def __getitem__(self, idx):
         pose = self.projections[f"{idx:03d}/gt-poses/cam-to-pelvis-vol"][:]
         pose = RigidTransform(torch.from_numpy(pose))
         pose = (
-            self.flip_z
+            self.rot_180
+            .compose(self.flip_z)
             .compose(self.world2camera.inverse())
             .compose(pose)
             .compose(self.anatomical2world)
+            .compose(self.rot_180)
         )
         if self.rot_180_for_up(idx):
             img = torch.rot90(img, k=2)
@@ -118,11 +121,11 @@ def parse_volume(subject, bone_attenuation_multiplier, labels):
     # Get all parts of the volume
     volume = subject["vol/pixels"][:]
     volume = np.swapaxes(volume, 0, 2).copy()
-    volume = torch.from_numpy(volume).unsqueeze(0)
+    volume = torch.from_numpy(volume).unsqueeze(0).flip(1).flip(2)
 
     mask = subject["vol-seg/image/pixels"][:]
     mask = np.swapaxes(mask, 0, 2).copy()
-    mask = torch.from_numpy(mask).unsqueeze(0)
+    mask = torch.from_numpy(mask).unsqueeze(0).flip(1).flip(2)
 
     affine = np.eye(4)
     affine[:3, :3] = subject["vol/dir-mat"][:]
@@ -161,10 +164,13 @@ def parse_volume(subject, bone_attenuation_multiplier, labels):
         volume=volume,
         labelmap=labelmap,
         labels=labels,
+        orientation="PA",
         bone_attenuation_multiplier=bone_attenuation_multiplier,
         label_def=defns,
         fiducials=fiducials,
     )
+    reorient = RigidTransform(torch.diag(torch.tensor([-1.0, -1.0, 1.0, 1.0])))
+    subject.fiducials = reorient(subject.fiducials)
 
     return subject, anatomical2world
 

diffdrrdata/ljubljana.py

@@ -57,7 +57,9 @@ def __getitem__(self, idx):
         )
 
         if self.preprocess:
-            img += 1
+            img = img - img.mode().values.mode().values.item()  # Subtract background color
+            img = torch.clamp(img, -1, 0) + 1  # Restrict to [0, 1]
+            img += 1  # Convert to log-scale
             img = img.max().log() - img.log()
 
         pose = (
@@ -75,9 +77,9 @@ def parse_volume(f, subject_id):
 
     # Get the volume
     volume = subject["volume/pixels"][:]
-    volume = volume[::-1].copy()
-    volume = torch.from_numpy(volume).unsqueeze(0)
-    volume[volume < 250] = -1000.0
+    volume = volume.copy()
+    volume = torch.from_numpy(volume).unsqueeze(0).flip(1)
+    volume[volume < 1000] = 0.0  # Discard a lot of the background from the 3D DSA
 
     affine = np.eye((4))
     spacing = subject["volume/spacing"][:]
@@ -162,7 +164,7 @@ def __init__(self, height: int, width: int, eps: float = 1e-6):
             [
                 Lambda(lambda x: (x - x.min()) / (x.max() - x.min() + eps)),
                 Resize((height, width), antialias=True),
-                Normalize(mean=0.0774, std=0.0569),
+                Normalize(mean=0.0306, std=0.0564),
             ]
         )