add reconstruction function to util

mu_map/util.py

+import math
 from typing import Optional
 
 import numpy as np
+from tomolab.Reconstruction.SPECT import SPECT_Static_Scan
 
 COLOR_BLACK = (0, 0, 0)
 COLOR_WHITE = (255, 255, 255)
@@ -40,3 +42,65 @@ def grayscale_to_rgb(img: np.ndarray):
     """
     assert img.ndim == 2, f"grascale image has more than 2 dimensions {img.shape}"
     return img.repeat(3).reshape((*img.shape, 3))
+
+
+def reconstruct(recon, mu_map=None, use_gpu=True):
+    recon = np.transpose(recon, (2, 1, 0)).astype(np.float32)
+    print(f"Transpose recon to {recon.shape}")
+    if mu_map is not None:
+        mu_map = np.transpose(mu_map, (2, 1, 0)).astype(np.float32)
+
+        s = recon.shape[0]
+        padding_lower = math.ceil(s - mu_map.shape[2])
+        padding_upper = math.floor(s - mu_map.shape[2])
+        _shape = mu_map.shape
+        mu_map = np.pad(mu_map, [(0, 0), (0, 0), (padding_lower, padding_upper)])
+        print(f"Change mu_map shape from {_shape} to {mu_map.shape}")
+
+
+    spect = SPECT_Static_Scan()
+    spect.set_use_gpu(use_gpu)
+
+    spect.set_n_pixels(128, 128)
+    spect.set_gantry_angular_positions(0.0, 360.0, 59)
+
+    print("Forward projection .....", end="\r")
+    if mu_map is not None:
+        measurement = spect.project(recon, attenuation=mu_map)
+        spect.set_measurement(measurement.data)
+        spect.set_attenuation(mu_map)
+    else:
+        measurement = spect.project(recon)
+        spect.set_measurement(measurement.data)
+    print("Forward projection DONE!")
+
+    spect.set_pixel_size(4.8, 4.8)
+    spect.set_radius(200.0)
+    spect.set_psf(fwhm0_mm=5.0, depth_dependence=0.0001)
+
+    print("Reconstruction .....", end="\r")
+    activity = spect.estimate_activity(
+        iterations=10, subset_size=16, subset_mode="random", method="EM"
+    )
+    activity = activity.data
+    print("Reconstruction DONE!")
+
+    s = recon.shape[2] / 2
+    c = activity.shape[2] // 2
+    activity = activity[:, :, c-math.ceil(s):c+math.floor(s)]
+    print(f"Reshape activity to {activity.shape}")
+
+    return np.transpose(activity, (2, 1, 0))
+
+if __name__ == "__main__":
+    import pydicom
+    
+    from mu_map.dataset.default import DCM_TAG_PIXEL_SCALE_FACTOR
+
+    img = pydicom.dcmread("./data/second/images/0001-stress-recon_nac_nsc.dcm")
+    img = img.pixel_array / img[DCM_TAG_PIXEL_SCALE_FACTOR].value
+
+    mu_map = pydicom.dcmread("./data/second/images/0001-stress-mu_map.dcm")
+    mu_map = mu_map.pixel_array / mu_map[DCM_TAG_PIXEL_SCALE_FACTOR].value
+
+    reconstruct(img, mu_map)