Skip to content

Commit

Permalink
read write with yardl
Browse files Browse the repository at this point in the history
  • Loading branch information
Imraj-Singh committed Nov 14, 2023
1 parent ffc63f5 commit 9c7453c
Showing 3 changed files with 262 additions and 0 deletions.
52 changes: 52 additions & 0 deletions python/demo_read_reconstruct.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
# TODO: - absolute scale of recon (maybe with sinogram recon)
# - additive MLEM

from __future__ import annotations

import parallelproj
import array_api_compat.numpy as np
import matplotlib.pyplot as plt
from io_yardl import read_yardl


dev = "cpu"

# image properties
voxel_size = (2.66, 2.66, 2.66)
img_shape = (128, 128, 4)
img_origin = [-168.91, -168.91, -3.9900002]

event_det_id_1, event_det_id_2, scanner_lut = read_yardl("write_test.yardl")

# hack until we have the reader / writer implemented

xstart = scanner_lut[event_det_id_1, :]
xend = scanner_lut[event_det_id_2, :]

# ----------------------------------------------------------------------------
# ----------------------------------------------------------------------------
# ---- LM recon using the event detector IDs and the scanner LUT -------------
# ----------------------------------------------------------------------------
# ----------------------------------------------------------------------------

recon = np.ones(img_shape, dtype=np.float32, device=dev)
img = np.tile(np.load("../SL.npy")[..., None], (1, 1, 4))
num_iter = 2
num_subsets = 20
for it in range(num_iter):
for isub in range(num_subsets):
print(f"it {(it+1):03} / ss {(isub+1):03}", end="\r")
xs_sub = xstart[isub::num_subsets, :]
xe_sub = xend[isub::num_subsets, :]
exp = parallelproj.joseph3d_fwd(xs_sub, xe_sub, recon, img_origin, voxel_size)
tmp = parallelproj.joseph3d_back(
xs_sub, xe_sub, img_shape, img_origin, voxel_size, 1 / exp
)
recon *= tmp

fig, ax = plt.subplots(1, 3, figsize=(12, 4))
fig.colorbar(ax[0].imshow(img[:, :, 2]))
fig.colorbar(ax[1].imshow(recon[:, :, 2]))
# difference
fig.colorbar(ax[2].imshow(recon[:, :, 2] - img[:, :, 2]))
fig.savefig("read_reconstruct_test.png")
121 changes: 121 additions & 0 deletions python/demo_write.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,121 @@
# TODO: - absolute scale of recon (maybe with sinogram recon)
# - additive MLEM

from __future__ import annotations

import parallelproj
import utils
import array_api_compat.numpy as np
import matplotlib.pyplot as plt
from array_api_compat import to_device
from scipy.ndimage import gaussian_filter
from io_yardl import write_yardl, read_yardl

# device variable (cpu or cuda) that determines whether calculations
# are performed on the cpu or cuda gpu

dev = "cpu"
expected_num_trues = 1e6
num_iter = 2
num_subsets = 20
np.random.seed(1)

# ----------------------------------------------------------------------------
# ----------------------------------------------------------------------------
# --- setup the scanner / LOR geometry ---------------------------------------
# ----------------------------------------------------------------------------
# ----------------------------------------------------------------------------

# setup a line of response descriptor that describes the LOR start / endpoints of
# a "narrow" clinical PET scanner with 9 rings
lor_descriptor = utils.DemoPETScannerLORDescriptor(
np, dev, num_rings=2, radial_trim=141
)

# ----------------------------------------------------------------------------
# ----------------------------------------------------------------------------
# --- setup a simple 3D test image -------------------------------------------
# ----------------------------------------------------------------------------
# ----------------------------------------------------------------------------

# image properties
voxel_size = (2.66, 2.66, 2.66)
num_trans = 128
num_ax = 2 * lor_descriptor.scanner.num_modules

# setup a box like test image
img_shape = (num_trans, num_trans, num_ax)
n0, n1, n2 = img_shape

# setup an image containing a box
img = np.tile(np.load("../SL.npy")[..., None], (1, 1, 4))

# ----------------------------------------------------------------------------
# ----------------------------------------------------------------------------
# --- setup a non-TOF projector and project ----------------------------------
# ----------------------------------------------------------------------------
# ----------------------------------------------------------------------------

projector = utils.RegularPolygonPETProjector(lor_descriptor, img_shape, voxel_size)
projector.tof = False # set this to True to get a time of flight projector

# forward project the image
noise_free_sinogram = projector(img)

# rescale the forward projection and image such that we get the expected number of trues
scale = expected_num_trues / np.sum(noise_free_sinogram)
noise_free_sinogram *= scale
img *= scale

# calculate the sensitivity image
sens_img = projector.adjoint(
np.ones(noise_free_sinogram.shape, device=dev, dtype=np.float32)
)

# get the two dimensional indices of all sinogram bins
start_mods, end_mods, start_inds, end_inds = lor_descriptor.get_lor_indices()

# generate two sinograms that contain the linearized detector start and end indices
sino_det_start_index = (
lor_descriptor.scanner.num_lor_endpoints_per_module[0] * start_mods + start_inds
)
sino_det_end_index = (
lor_descriptor.scanner.num_lor_endpoints_per_module[0] * end_mods + end_inds
)

# add poisson noise to the noise free sinogram
noisy_sinogram = np.random.poisson(noise_free_sinogram)

# ravel the noisy sinogram and the detector start and end "index" sinograms
noisy_sinogram = noisy_sinogram.ravel()
sino_det_start_index = sino_det_start_index.ravel()
sino_det_end_index = sino_det_end_index.ravel()

# ----------------------------------------------------------------------------
# ----------------------------------------------------------------------------
# --- convert the index sinograms in to listmode data ------------------------
# ----------------------------------------------------------------------------
# ----------------------------------------------------------------------------

# generate listmode data from the noisy sinogram
event_sino_inds = np.repeat(np.arange(noisy_sinogram.shape[0]), noisy_sinogram)
# shuffle the event sinogram indices
np.random.shuffle(event_sino_inds)

event_det_id_1 = sino_det_start_index[event_sino_inds]
event_det_id_2 = sino_det_end_index[event_sino_inds]

print(f"number of events: {event_det_id_1.shape[0]}")

# ----------------------------------------------------------------------------
# ----------------------------------------------------------------------------
# ---- convert LM detector ID arrays into PRD here ---------------------------
# ----------------------------------------------------------------------------
# ----------------------------------------------------------------------------

# get a lookup table that contains the world coordinates of all scanner detectors
# this is a 2D array of shape (num_detectors, 3)
scanner_lut = lor_descriptor.scanner.all_lor_endpoints

# write and re-read
write_yardl(event_det_id_1, event_det_id_2, scanner_lut, output_file="write_test.yardl")
89 changes: 89 additions & 0 deletions python/io_yardl.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
from __future__ import annotations
import sys

sys.path.append("/workspaces/LMSimRecon/PETSIRD/python")
import prd
import numpy as np


def write_yardl(
det_1_array: np.array[int],
det_2_array: np.array[int],
scanner_lut: np.array[float],
output_file: str | None = None,
) -> None:
"""Writes a yardl list mode file from two detector arrays and a scanner lookup table.
Args:
output_file (str): File to write to.
det_1_array (np.array[int]): Indices corresponding to detector 1.
det_2_array (np.array[int]): Indices corresponding to detector 2.
scanner_lut (np.array[float]): Lookup table for detector world coordinates.
"""
num_counts = len(det_1_array)
# create list of detectors
detectors = []
for i in range(scanner_lut.shape[0]):
detectors.append(
prd.Detector(
id=i, x=scanner_lut[i, 0], y=scanner_lut[i, 1], z=scanner_lut[i, 2]
)
)
scanner_information = prd.ScannerInformation(detectors=detectors)
events = []
for i in range(num_counts):
events.append(
prd.CoincidenceEvent(
detector_1_id=int(det_1_array[i]), detector_2_id=int(det_2_array[i])
)
)
time_block = prd.TimeBlock(id=0, prompt_events=events)
if output_file is None:
with prd.BinaryPrdExperimentWriter(sys.stdout.buffer) as writer:
writer.write_header(prd.Header(scanner=scanner_information))
writer.write_time_blocks((time_block,))
else:
with prd.BinaryPrdExperimentWriter(output_file) as writer:
writer.write_header(prd.Header(scanner=scanner_information))
writer.write_time_blocks((time_block,))


def read_yardl(prd_file: str) -> tuple[np.array[int], np.array[int], np.array[float]]:
"""Reads a yardl list mode file and returns two detector arrays and a scanner lookup table.
Args:
prd_file (str): File to read from.
Returns:
tuple[np.array[int], np.array[int], np.array[float]]: Two detector arrays and a scanner lookup table.
"""
with prd.BinaryPrdExperimentReader(prd_file) as reader:
# Read header and build lookup table
header = reader.read_header()
scanner_information = header.scanner
detectors = scanner_information.detectors
scanner_lut = np.zeros((len(detectors), 3))
for i, detector in enumerate(detectors):
scanner_lut[i, 0] = detector.x
scanner_lut[i, 1] = detector.y
scanner_lut[i, 2] = detector.z
# Read events
detector_hits = []
for time_blocks in reader.read_time_blocks():
for event in time_blocks.prompt_events:
prompt = [event.detector_1_id, event.detector_2_id]
detector_hits.append(prompt)
det_1, det_2 = np.asarray(detector_hits).T
return det_1, det_2, scanner_lut


if __name__ == "__main__":
det_1 = np.asarray([0, 1, 2, 3, 4])
det_2 = np.asarray([5, 6, 7, 8, 9])
scanner_lut = np.random.rand(10, 3)
write_yardl(det_1, det_2, scanner_lut, output_file="test.yardl")
det_1_read, det_2_read, scanner_lut_read = read_yardl("test.yardl")
print(scanner_lut == scanner_lut_read)
print(np.isclose(scanner_lut, scanner_lut_read).all())
print(scanner_lut.dtype)
print(scanner_lut_read.dtype)

0 comments on commit 9c7453c

Please sign in to comment.