add script to visualize scanner and events

python/petsird_plot_scanner_and_events.py

@@ -0,0 +1,167 @@
+#  Copyright (C) 2024 University College London
+#
+#  SPDX-License-Identifier: Apache-2.0
+
+# basic plotting of the scanner geometry
+# preliminary code!
+
+import sys
+import numpy
+import numpy.typing as npt
+import petsird
+
+from petsird_helpers import (
+    get_num_det_els,
+    get_module_and_element,
+    get_detection_efficiency,
+)
+
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+
+
+def transform_to_mat44(
+    transform: petsird.RigidTransformation,
+) -> npt.NDArray[numpy.float32]:
+    return numpy.vstack([transform.matrix, [0, 0, 0, 1]])
+
+
+def mat44_to_transform(mat: npt.NDArray[numpy.float32]) -> petsird.RigidTransformation:
+    return petsird.RigidTransformation(matrix=mat[0:3, :])
+
+
+def coordinate_to_homogeneous(coord: petsird.Coordinate) -> npt.NDArray[numpy.float32]:
+    return numpy.hstack([coord.c, 1])
+
+
+def homogeneous_to_coordinate(
+    hom_coord: npt.NDArray[numpy.float32],
+) -> petsird.Coordinate:
+    return petsird.Coordinate(c=hom_coord[0:3])
+
+
+def mult_transforms(
+    transforms: list[petsird.RigidTransformation],
+) -> petsird.RigidTransformation:
+    """multiply rigid transformations"""
+    mat = numpy.array(
+        ((1, 0, 0, 0), (0, 1, 0, 0), (0, 0, 1, 0), (0, 0, 0, 1)),
+        dtype="float32",
+    )
+
+    for t in reversed(transforms):
+        mat = numpy.matmul(transform_to_mat44(t), mat)
+    return mat44_to_transform(mat)
+
+
+def mult_transforms_coord(
+    transforms: list[petsird.RigidTransformation], coord: petsird.Coordinate
+) -> petsird.Coordinate:
+    """apply list of transformations to coordinate"""
+    # TODO better to multiply with coordinates in sequence, as first multiplying the matrices
+    hom = numpy.matmul(
+        transform_to_mat44(mult_transforms(transforms)),
+        coordinate_to_homogeneous(coord),
+    )
+    return homogeneous_to_coordinate(hom)
+
+
+def transform_BoxShape(
+    transform: petsird.RigidTransformation, box_shape: petsird.BoxShape
+) -> petsird.BoxShape:
+    return petsird.BoxShape(
+        corners=[mult_transforms_coord([transform], c) for c in box_shape.corners]
+    )
+
+
+def draw_BoxShape(ax, box: petsird.BoxShape) -> None:
+    vertices = numpy.array([c.c for c in box.corners])
+    edges = [
+        [vertices[j] for j in [0, 1, 2, 3]],
+        [vertices[j] for j in [4, 5, 6, 7]],
+        [vertices[j] for j in [0, 1, 5, 4]],
+        [vertices[j] for j in [2, 3, 7, 6]],
+        [vertices[j] for j in [1, 2, 6, 5]],
+        [vertices[j] for j in [4, 7, 3, 0]],
+    ]
+    box = Poly3DCollection(edges, alpha=0.25, linewidths=1, edgecolors="r")
+    ax.add_collection3d(box)
+
+
+if __name__ == "__main__":
+
+    # Create a new figure
+    fig = plt.figure()
+    ax = fig.add_subplot(111, projection="3d")
+
+    # dictionary to store the transformations and centers of the detecting elements
+    # here we assume that we only have BoxShapes
+    element_transforms = dict()
+    element_centers = dict()
+
+    # with petsird.BinaryPETSIRDReader(sys.stdin.buffer) as reader:
+    with petsird.BinaryPETSIRDReader("test.bin") as reader:
+        header = reader.read_header()
+
+        # draw all crystals
+        for rep_module in header.scanner.scanner_geometry.replicated_modules:
+            det_el = rep_module.object.detecting_elements
+            for i_mod, mod_transform in enumerate(rep_module.transforms):
+                for rep_volume in det_el:
+                    for i_el, transform in enumerate(rep_volume.transforms):
+
+                        combined_transform = mult_transforms([mod_transform, transform])
+                        transformed_boxshape = transform_BoxShape(
+                            combined_transform, rep_volume.object.shape
+                        )
+
+                        transformed_boxshape_vertices = numpy.array(
+                            [c.c for c in transformed_boxshape.corners]
+                        )
+
+                        draw_BoxShape(ax, transformed_boxshape)
+
+                        element_transforms[(i_mod, i_el)] = combined_transform
+                        element_centers[(i_mod, i_el)] = (
+                            transformed_boxshape_vertices.mean(axis=0)
+                        )
+
+        # ----
+        # read and draw events
+
+        num_prompts = 0
+        event_counter = 0
+
+        for i_time_block, time_block in enumerate(reader.read_time_blocks()):
+            if isinstance(time_block, petsird.TimeBlock.EventTimeBlock):
+                num_prompts += len(time_block.value.prompt_events)
+
+                for i_event, event in enumerate(time_block.value.prompt_events):
+                    event_mods_and_els = get_module_and_element(
+                        header.scanner.scanner_geometry, event.detector_ids
+                    )
+
+                    event_start_coord = element_centers[
+                        event_mods_and_els[0].module, event_mods_and_els[0].el
+                    ]
+                    event_end_coord = element_centers[
+                        event_mods_and_els[1].module, event_mods_and_els[1].el
+                    ]
+
+                    # draw line between the two 3D points (event_start_coord, event_end_coord)
+                    # for the first 2 events in the time block
+                    if i_event < 2:
+                        ax.plot(
+                            [event_start_coord[0], event_end_coord[0]],
+                            [event_start_coord[1], event_end_coord[1]],
+                            [event_start_coord[2], event_end_coord[2]],
+                        )
+
+                    event_eff = get_detection_efficiency(header.scanner, event)
+
+                    event_counter += 1
+
+    ax.set_xlim(-400, 400)
+    ax.set_ylim(-400, 400)
+    ax.set_zlim(-25, 50)
+    fig.show()