change reconstruction function

mu_map/util.py

@@ -4,6 +4,8 @@ from typing import Optional
 import numpy as np
 from tomolab.Reconstruction.SPECT import SPECT_Static_Scan
 
+from mu_map.dataset.util import align_images
+
 COLOR_BLACK = (0, 0, 0)
 COLOR_WHITE = (255, 255, 255)
 
@@ -44,34 +46,33 @@ def grayscale_to_rgb(img: np.ndarray):
     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}")
-
+import random
+def reconstruct(recon, mu_map=None, use_gpu=True, seed=42):
+    random.seed(42)
 
     spect = SPECT_Static_Scan()
     spect.set_use_gpu(use_gpu)
 
+    recon_t = np.transpose(recon, (2, 1, 0)).astype(np.float32)
+    n_pixels = recon_t.shape[0]
+    print(f"NPixels {n_pixels}, {recon_t.shape}")
+    # spect.set_n_pixels(n_pixels, n_pixels)
     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)
+        mu_map_t = np.transpose(mu_map, (2, 1, 0)).astype(np.float32)
+
+        padding_lower = math.ceil(n_pixels - mu_map_t.shape[2])
+        padding_upper = math.floor(n_pixels - mu_map_t.shape[2])
+        mu_map_t = np.pad(mu_map_t, [(0, 0), (0, 0), (padding_lower, padding_upper)])
+
+        spect.set_attenuation(mu_map_t)
+
+    print("Forward projection .....", end="\r")
+    measurement = spect.project(recon_t)
+    print(f"Measurement: {measurement.data.shape}")
+    spect.set_measurement(measurement.data)
     print("Forward projection DONE!")
 
     spect.set_pixel_size(4.8, 4.8)
@@ -82,25 +83,19 @@ def reconstruct(recon, mu_map=None, use_gpu=True):
     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))
+    activity = activity.data
+    activity = np.transpose(activity, (2, 1, 0))
+    activity, _ = align_images(activity, recon)
+    return activity
 
 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
+    from mu_map.dataset.util import load_dcm_img
 
-    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
+    img = load_dcm_img("./data/second/images/0001-stress-recon_nac_nsc.dcm")
+    mu_map = load_dcm_img("./data/second/images/0001-stress-mu_map.dcm")
 
     reconstruct(img, mu_map)