Skip to content

Commit

Permalink
Fix orientations
Browse files Browse the repository at this point in the history
  • Loading branch information
eigenvivek committed Dec 13, 2024
1 parent cbedf79 commit cbf9c46
Show file tree
Hide file tree
Showing 8 changed files with 74 additions and 55 deletions.
37 changes: 20 additions & 17 deletions diffdrr/data.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,35 +85,37 @@ def read(
# Frame-of-reference change
if orientation == "AP":
# Rotates the C-arm about the x-axis by 90 degrees
# Rotates the C-arm about the z-axis by -90 degrees
reorient = torch.tensor(
[
[1.0, 0.0, 0.0, 0.0],
[0.0, 0.0, -1.0, 0.0],
[0.0, -1.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 1.0],
]
[1, 0, 0, 0],
[0, 0, -1, 0],
[0, 1, 0, 0],
[0, 0, 0, 1],
],
dtype=torch.float32,
)
elif orientation == "PA":
# Rotates the C-arm about the x-axis by 90 degrees
# Rotates the C-arm about the z-axis by 90 degrees
# Reverses the direction of the y-axis
reorient = torch.tensor(
[
[1.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0],
[0.0, -1.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 1.0],
]
[1, 0, 0, 0],
[0, 0, 1, 0],
[0, 1, 0, 0],
[0, 0, 0, 1],
],
dtype=torch.float32,
)
elif orientation is None:
# Identity transform
reorient = torch.tensor(
[
[1.0, 0.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 1.0],
]
[1, 0, 0, 0],
[0, 1, 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 1],
],
dtype=torch.float32,
)
else:
raise ValueError(f"Unrecognized orientation {orientation}")
Expand All @@ -122,6 +124,7 @@ def read(
subject = Subject(
volume=volume,
mask=mask,
orientation=orientation,
reorient=reorient,
density=density,
fiducials=fiducials,
Expand Down
6 changes: 5 additions & 1 deletion diffdrr/detector.py
Original file line number Diff line number Diff line change
Expand Up @@ -117,8 +117,12 @@ def _initialize_carm(self: Detector):
# Construct equally spaced points along the basis vectors
t = torch.arange(-self.height // 2, self.height // 2, device=device) + h_off
s = torch.arange(-self.width // 2, self.width // 2, device=device) + w_off
if self.reverse_x_axis:

t = -t
s = -s
if not self.reverse_x_axis:
s = -s

coefs = torch.cartesian_prod(t, s).reshape(-1, 2)
target = torch.einsum("cd,nc->nd", basis, coefs)
target += center
Expand Down
2 changes: 1 addition & 1 deletion diffdrr/drr.py
Original file line number Diff line number Diff line change
Expand Up @@ -261,7 +261,7 @@ def perspective_projection(
x = torch.einsum("ij, bnj -> bni", self.detector.intrinsic, x)
z = x[..., -1].unsqueeze(-1).clone()
x = x / z
if self.detector.reverse_x_axis:
if not self.detector.reverse_x_axis:
x[..., 1] = self.detector.width - x[..., 1]
return x[..., :2] # .flip(-1)

Expand Down
18 changes: 10 additions & 8 deletions diffdrr/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ def resample(

# %% ../notebooks/api/07_utils.ipynb 6
from kornia.geometry.camera.pinhole import PinholeCamera as KorniaPinholeCamera
from torchio import Subject

from diffdrr.detector import Detector

Expand All @@ -66,9 +67,11 @@ def __init__(
height: torch.Tensor,
width: torch.Tensor,
detector: Detector,
subject: Subject,
):
super().__init__(intrinsics, extrinsics, height, width)
self.f = detector.sdd
multiplier = -1 if subject.orientation == "PA" else 1
self.f = multiplier * detector.sdd
self.delx = detector.delx
self.dely = detector.dely
self.x0 = detector.x0
Expand All @@ -94,9 +97,9 @@ def pose(self):

from kornia.geometry.calibration import solve_pnp_dlt

from .detector import make_intrinsic_matrix
from .drr import DRR
from .pose import RigidTransform
from .detector import make_intrinsic_matrix


def get_pinhole_camera(
Expand All @@ -107,14 +110,12 @@ def get_pinhole_camera(
pose = deepcopy(pose).to(device="cpu", dtype=dtype)

# Make the intrinsic matrix (in pixels)
fx = drr.detector.sdd / drr.detector.delx
fy = drr.detector.sdd / drr.detector.dely
multiplier = -1 if drr.subject.orientation == "PA" else 1
fx = multiplier * drr.detector.sdd / drr.detector.delx
fy = multiplier * drr.detector.sdd / drr.detector.dely
u0 = drr.detector.x0 / drr.detector.delx + drr.detector.width / 2
v0 = drr.detector.y0 / drr.detector.dely + drr.detector.height / 2
intrinsics = torch.eye(4)[None]
intrinsics[0, :3, :3] = make_intrinsic_matrix(drr.detector)

torch.tensor(
intrinsics = torch.tensor(
[
[
[fx, 0.0, u0, 0.0],
Expand Down Expand Up @@ -156,6 +157,7 @@ def get_pinhole_camera(
torch.tensor([drr.detector.height]),
torch.tensor([drr.detector.width]),
drr.detector,
drr.subject,
)

return camera
2 changes: 1 addition & 1 deletion notebooks/api/00_drr.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -404,7 +404,7 @@
" x = torch.einsum(\"ij, bnj -> bni\", self.detector.intrinsic, x)\n",
" z = x[..., -1].unsqueeze(-1).clone()\n",
" x = x / z\n",
" if self.detector.reverse_x_axis:\n",
" if not self.detector.reverse_x_axis:\n",
" x[..., 1] = self.detector.width - x[..., 1]\n",
" return x[..., :2]#.flip(-1)"
]
Expand Down
6 changes: 5 additions & 1 deletion notebooks/api/02_detector.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -180,8 +180,12 @@
" # Construct equally spaced points along the basis vectors\n",
" t = torch.arange(-self.height // 2, self.height // 2, device=device) + h_off\n",
" s = torch.arange(-self.width // 2, self.width // 2, device=device) + w_off\n",
" if self.reverse_x_axis:\n",
"\n",
" t = -t\n",
" s = -s\n",
" if not self.reverse_x_axis:\n",
" s = -s\n",
"\n",
" coefs = torch.cartesian_prod(t, s).reshape(-1, 2)\n",
" target = torch.einsum(\"cd,nc->nd\", basis, coefs)\n",
" target += center\n",
Expand Down
39 changes: 21 additions & 18 deletions notebooks/api/03_data.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -156,35 +156,37 @@
" # Frame-of-reference change\n",
" if orientation == \"AP\":\n",
" # Rotates the C-arm about the x-axis by 90 degrees\n",
" # Rotates the C-arm about the z-axis by -90 degrees\n",
" reorient = torch.tensor(\n",
" [\n",
" [1.0, 0.0, 0.0, 0.0],\n",
" [0.0, 0.0, -1.0, 0.0],\n",
" [0.0, -1.0, 0.0, 0.0],\n",
" [0.0, 0.0, 0.0, 1.0],\n",
" ]\n",
" [1, 0, 0, 0],\n",
" [0, 0, -1, 0],\n",
" [0, 1, 0, 0],\n",
" [0, 0, 0, 1],\n",
" ],\n",
" dtype=torch.float32,\n",
" )\n",
" elif orientation == \"PA\":\n",
" # Rotates the C-arm about the x-axis by 90 degrees\n",
" # Rotates the C-arm about the z-axis by 90 degrees\n",
" # Rotates the C-arm about the x-axis by 90 degrees \n",
" # Reverses the direction of the y-axis\n",
" reorient = torch.tensor(\n",
" [\n",
" [1.0, 0.0, 0.0, 0.0],\n",
" [0.0, 0.0, 1.0, 0.0],\n",
" [0.0, -1.0, 0.0, 0.0],\n",
" [0.0, 0.0, 0.0, 1.0],\n",
" ]\n",
" [1, 0, 0, 0],\n",
" [0, 0, 1, 0],\n",
" [0, 1, 0, 0],\n",
" [0, 0, 0, 1],\n",
" ],\n",
" dtype=torch.float32,\n",
" )\n",
" elif orientation is None:\n",
" # Identity transform\n",
" reorient = torch.tensor(\n",
" [\n",
" [1.0, 0.0, 0.0, 0.0],\n",
" [0.0, 1.0, 0.0, 0.0],\n",
" [0.0, 0.0, 1.0, 0.0],\n",
" [0.0, 0.0, 0.0, 1.0],\n",
" ]\n",
" [1, 0, 0, 0],\n",
" [0, 1, 0, 0],\n",
" [0, 0, 1, 0],\n",
" [0, 0, 0, 1],\n",
" ],\n",
" dtype=torch.float32,\n",
" )\n",
" else:\n",
" raise ValueError(f\"Unrecognized orientation {orientation}\")\n",
Expand All @@ -193,6 +195,7 @@
" subject = Subject(\n",
" volume=volume,\n",
" mask=mask,\n",
" orientation=orientation,\n",
" reorient=reorient,\n",
" density=density,\n",
" fiducials=fiducials,\n",
Expand Down
19 changes: 11 additions & 8 deletions notebooks/api/07_utils.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,7 @@
"source": [
"#| exporti\n",
"from kornia.geometry.camera.pinhole import PinholeCamera as KorniaPinholeCamera\n",
"from torchio import Subject\n",
"\n",
"from diffdrr.detector import Detector\n",
"\n",
Expand All @@ -131,9 +132,12 @@
" height: torch.Tensor,\n",
" width: torch.Tensor,\n",
" detector: Detector,\n",
" subject: Subject,\n",
" \n",
" ):\n",
" super().__init__(intrinsics, extrinsics, height, width)\n",
" self.f = detector.sdd\n",
" multiplier = -1 if subject.orientation == \"PA\" else 1\n",
" self.f = multiplier * detector.sdd\n",
" self.delx = detector.delx\n",
" self.dely = detector.dely\n",
" self.x0 = detector.x0\n",
Expand Down Expand Up @@ -167,9 +171,9 @@
"\n",
"from kornia.geometry.calibration import solve_pnp_dlt\n",
"\n",
"from diffdrr.detector import make_intrinsic_matrix\n",
"from diffdrr.drr import DRR\n",
"from diffdrr.pose import RigidTransform\n",
"from diffdrr.detector import make_intrinsic_matrix\n",
"\n",
"\n",
"def get_pinhole_camera(\n",
Expand All @@ -180,14 +184,12 @@
" pose = deepcopy(pose).to(device=\"cpu\", dtype=dtype)\n",
"\n",
" # Make the intrinsic matrix (in pixels)\n",
" fx = drr.detector.sdd / drr.detector.delx\n",
" fy = drr.detector.sdd / drr.detector.dely\n",
" multiplier = -1 if drr.subject.orientation == \"PA\" else 1\n",
" fx = multiplier * drr.detector.sdd / drr.detector.delx\n",
" fy = multiplier * drr.detector.sdd / drr.detector.dely\n",
" u0 = drr.detector.x0 / drr.detector.delx + drr.detector.width / 2\n",
" v0 = drr.detector.y0 / drr.detector.dely + drr.detector.height / 2\n",
" intrinsics = torch.eye(4)[None]\n",
" intrinsics[0, :3, :3] = make_intrinsic_matrix(drr.detector)\n",
" \n",
" torch.tensor(\n",
" intrinsics = torch.tensor(\n",
" [\n",
" [\n",
" [fx, 0.0, u0, 0.0],\n",
Expand Down Expand Up @@ -229,6 +231,7 @@
" torch.tensor([drr.detector.height]),\n",
" torch.tensor([drr.detector.width]),\n",
" drr.detector,\n",
" drr.subject,\n",
" )\n",
"\n",
" return camera"
Expand Down

0 comments on commit cbf9c46

Please # to comment.