fix reconstruction impl

mu_map/util.py

@@ -47,52 +47,47 @@ def grayscale_to_rgb(img: np.ndarray):
     return img.repeat(3).reshape((*img.shape, 3))
 
 
-def reconstruct(recon, mu_map=None, use_gpu=True, seed=42):
+def reconstruct(recon, mu_map=None, iterations=10, use_gpu=True, seed=42, test_flag=True):
     random.seed(seed)
 
-    print("Init scan")
     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"Set to NPixels {n_pixels}, {recon_t.shape}")
-    # spect.set_n_pixels(n_pixels, n_pixels)
-    spect.set_n_pixels(128, 128)
+
+    padding = (n_pixels - recon_t.shape[2]) / 2
+    padding_lower = math.ceil(padding)
+    padding_upper = math.floor(padding)
+    recon_t = np.pad(recon_t, [(0, 0), (0, 0), (padding_lower, padding_upper)])
+
+    spect.set_n_pixels(n_pixels, n_pixels)
     spect.set_gantry_angular_positions(0.0, 360.0, 59)
 
+    # 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!")
+
     if mu_map is not None:
         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)])
-
-        print("Set attenuation map...")
         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!")
-
-    print("init more parameters...")
     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")
+    # print("Reconstruction .....", end="\r")
     activity = spect.estimate_activity(
-        iterations=10, subset_size=16, subset_mode="random", method="EM"
+        iterations=iterations, subset_size=16, subset_mode="random", method="EM"
     )
-    print("Reconstruction DONE!")
+    # print("Reconstruction DONE!")
 
     activity = activity.data
     activity = np.transpose(activity, (2, 1, 0))
-    _shape = activity.shape
     activity, _ = align_images(activity, recon)
-    print(f"Realign activity from {_shape} to {activity.shape}")
     return activity
 
 if __name__ == "__main__":
@@ -125,6 +120,7 @@ if __name__ == "__main__":
 
         s = np.full((1024, 10), 239, np.uint8)
         return np.hstack((img1, s, img2))
+    exit(0)
 
     i = 0
     wname = "Test"