From 67943291129464802294820a01303c93e61286f8 Mon Sep 17 00:00:00 2001 From: Andrew Herzing Date: Fri, 31 May 2024 17:01:39 -0400 Subject: [PATCH] Added docstrings for DART functions --- tomotools/recon.py | 80 ++++++++++++++++++++++++++++++++++++---------- 1 file changed, 63 insertions(+), 17 deletions(-) diff --git a/tomotools/recon.py b/tomotools/recon.py index 57fd377f..7eaa1531 100644 --- a/tomotools/recon.py +++ b/tomotools/recon.py @@ -22,11 +22,11 @@ def run_alg(sino, iters, sino_id, alg_id, rec_id): """ - Run reconstruction algorithm. + Run FBP or SIRT reconstruction algorithm. Args ---------- - sion : NumPy array + sino : NumPy array Sinogram of shape (nangles, ny) iters : int Number of iterations for the SIRT reconstruction @@ -48,16 +48,52 @@ def run_alg(sino, iters, sino_id, alg_id, rec_id): return astra.data2d.get(rec_id) -def run_dart(sinogram, iters, dart_iterations, p, - alg, proj_id, mask_id, rec_id, sino_id, - thresholds, gray_levels, thickness, ny): - astra.data2d.store(sino_id, sinogram) +def run_dart(sino, iters, dart_iters, p, + alg_id, proj_id, mask_id, rec_id, sino_id, + thresholds, gray_levels): + """ + Run FBP or SIRT reconstruction algorithm. + + Args + ---------- + sino : NumPy array + Sinogram of shape (nangles, ny) + iters : int + Number of iterations for the SART reconstruction + dart_iters : int + Number of iterations for the DART reconstruction + p : float + Probability for free pixel determination + alg_id : int + ASTRA algorithm identity + proj_id : int + ASTRA projector identity + mask_id : boolean + ASTRA mask identity + rec_id : boolean + ASTRA reconstruction identity + sino_id : int + ASTRA sinogram identity + thresholds : list or NumPy array + Thresholds for DART reconstruction + gray_levels : list or NumPy array + Gray levels for DART reconstruction + + + Returns + ---------- + Numpy array + Reconstruction of input sinogram + + """ + thickness, ny = astra.data2d.get(rec_id).shape + astra.data2d.store(sino_id, sino) astra.data2d.store(rec_id, np.zeros([thickness, ny])) astra.data2d.store(mask_id, np.ones([thickness, ny])) - astra.algorithm.run(alg, iters) + astra.algorithm.run(alg_id, iters) curr_rec = astra.data2d.get(rec_id) dart_rec = copy.deepcopy(curr_rec) - for j in range(dart_iterations): + for j in range(dart_iters): segmented = dart_segment(dart_rec, thresholds, gray_levels) boundary = get_dart_boundaries(segmented) @@ -76,17 +112,17 @@ def run_dart(sinogram, iters, dart_iterations, p, fixed_rec = copy.deepcopy(dart_rec) fixed_rec[free_idx[0], free_idx[1]] = 0 _, fixed_sino = astra.creators.create_sino(fixed_rec, proj_id) - free_sino = sinogram - fixed_sino + free_sino = sino - fixed_sino # Run SART reconstruction on free sinogram with free pixel mask astra.data2d.store(rec_id, dart_rec) astra.data2d.store(mask_id, free) astra.data2d.store(sino_id, free_sino) - astra.algorithm.run(alg, iters) + astra.algorithm.run(alg_id, iters) dart_rec = astra.data2d.get(rec_id) # Smooth reconstruction - if j < dart_iterations - 1: + if j < dart_iters - 1: smooth = gaussian_filter(dart_rec, sigma=1) curr_rec[free_idx[0], free_idx[1]] = smooth[free_idx[0], free_idx[1]] else: @@ -117,6 +153,18 @@ def run(stack, method, niterations=20, constrain=None, thresh=0, cuda=None, thic cuda : boolean If True, use the CUDA-accelerated Astra algorithms. Otherwise, use the CPU-based algorithms + thickness : int + Limit for the height of the reconstruction + ncores : int + Number of cores to use for multithreaded CPU-based reconstructions + filter : str + Filter to use for filtered backprojection + gray_levels : list or NumPy array + Gray levels for DART reconstruction + dart_iterations : int + Number of DART iterations + p : float + Probability for setting free pixels in DART reconstruction Returns ---------- @@ -200,8 +248,7 @@ def run(stack, method, niterations=20, constrain=None, thresh=0, cuda=None, thic astra.data2d.store(rec_id, np.zeros([thickness, ny])) astra.data2d.store(mask_id, np.ones([thickness, ny])) rec[i, :, :] = run_dart(sinogram, niterations, dart_iterations, p, - alg, proj_id, mask_id, rec_id, sino_id, thresholds, - gray_levels, thickness, ny) + alg, proj_id, mask_id, rec_id, sino_id, thresholds, gray_levels) else: if ncores is None: ncores = min(nx, int(0.9 * mp.cpu_count())) @@ -256,16 +303,15 @@ def run(stack, method, niterations=20, constrain=None, thresh=0, cuda=None, thic if ncores == 1: for i in tqdm.tqdm(range(0, nx)): rec[i] = run_dart(stack.data[:, :, i], niterations, dart_iterations, p, - alg, proj_id, mask_id, rec_id, sino_id, thresholds, - gray_levels, thickness, ny) + alg, proj_id, mask_id, rec_id, sino_id, thresholds, gray_levels) else: logger.info("Using %s CPU cores to reconstruct %s slices" % (ncores, nx)) with mp.Pool(ncores) as pool: for i, result in enumerate( pool.starmap(run_dart, [(stack.data[:, :, i], niterations, dart_iterations, p, - alg, proj_id, mask_id, rec_id, sino_id, thresholds, - gray_levels, thickness, ny) for i in range(0, nx)],)): + alg, proj_id, mask_id, rec_id, sino_id, thresholds, gray_levels) + for i in range(0, nx)],)): rec[i] = result astra.clear() return rec