""" Compute reconstructions with CTAC and DLAC (attenuation maps created by `compute_mu_maps`). This script implements the (post-reconstructions attenuation correction) PRAC algorithm by first projection a non-attenuation-corrected reconstruction and then reconstructing this projection with an attenuation map. """ import argparse import os import pydicom import stir stir.Verbosity_set(0) from mu_map.file.dicom import change_uid, load_dcm, update_dcm from mu_map.file.util import load_as_interfile from mu_map.logging import add_logging_args, get_logger_by_args from mu_map.recon.filter import GaussianFilter from mu_map.recon.osem import reconstruct from mu_map.recon.project import forward_project parser = argparse.ArgumentParser() parser.add_argument( "--mu_maps", type=str, required=True, help="directory of mu maps computed with `mu_map.scripts.compute_mu_maps`", ) parser.add_argument( "--out_dir", type=str, required=True, help="directory where resulting reconstructions are written to", ) parser.add_argument( "--description_postfix", type=str, help="postfix for the DICOM SeriesDescription", ) parser.add_argument( "--dataset_dir", type=str, default="data/second/", help="directory of the dataset with reconstructions and real mu maps", ) parser.add_argument( "--images_dir", type=str, default="images", help="directory under `--dataset_dir` actually containing the images", ) parser.add_argument( "--n_subsets", type=int, default=3, help="number of subsets used for the OSEM reconstruction", ) parser.add_argument( "--n_iterations", type=int, default=10, help="number of iterations performed by the OSEM reconstruction", ) add_logging_args(parser, defaults={"--logfile": "compute_recons.log"}) args = parser.parse_args() args.images_dir = os.path.join(args.dataset_dir, args.images_dir) args.description_postfix = ( "" if args.description_postfix is None else f" - {args.description_postfix}" ) args.logfile = os.path.join(args.out_dir, args.logfile) logger = get_logger_by_args(args) logger.info(args) files_mu_maps_dl = os.listdir(args.mu_maps) files_mu_maps_dl = filter(lambda f: f.split(".")[-1] == "dcm", files_mu_maps_dl) files_mu_maps_dl = list(files_mu_maps_dl) files_mu_maps_dl.sort() for file_mu_map_dl in files_mu_maps_dl: logger.info(f"Process {file_mu_map_dl}") file_recon_ac = file_mu_map_dl.replace("mu_map_dl", "recon_ac_nsc") file_recon_ac = os.path.join(args.images_dir, file_recon_ac) recon_ac_dcm, _ = load_dcm(file_recon_ac, direction=1) logger.info(f"Found AC recon: {file_recon_ac}") file_recon_nac = file_mu_map_dl.replace("mu_map_dl", "recon_nac_nsc") file_recon_nac = os.path.join(args.images_dir, file_recon_nac) recon_nac = load_as_interfile(file_recon_nac, direction=1) logger.info(f"Found nAC recon: {file_recon_nac}") file_mu_map_ct = file_mu_map_dl.replace("mu_map_dl", "mu_map") file_mu_map_ct = os.path.join(args.images_dir, file_mu_map_ct) mu_map_ct = load_as_interfile(file_mu_map_ct, direction=1) logger.info(f"Found CT-based mu map: {file_mu_map_ct}") mu_map_dl = load_as_interfile( os.path.join(args.mu_maps, file_mu_map_dl), direction=1 ) for mu_map, label in zip([mu_map_ct, mu_map_dl], ["CTAC", "DLAC"]): logger.info(f"Compute {label} reconstruction") projection = forward_project(recon_nac, n_slices=mu_map.image.shape[0]) _recon_ac = reconstruct( projection, mu_map=mu_map, n_subsets=args.n_subsets, n_iterations=args.n_iterations, postfilter=GaussianFilter(projection), ) # for some reason STIR reconstruction are inverted in the x-direction # this is reverted here _recon_ac.image = _recon_ac.image[:, :, ::-1] _recon_ac_dcm = change_uid(update_dcm(recon_ac_dcm.copy(), _recon_ac.image)) _recon_ac_dcm.SeriesDescription = _recon_ac_dcm.SeriesDescription.replace( "NoSC - AC ", f"{label}{args.description_postfix}" ) _file_recon_ac = file_mu_map_dl.replace("mu_map_dl", f"recon_{label.lower()}") _file_recon_ac = os.path.join(args.out_dir, _file_recon_ac) pydicom.dcmwrite(_file_recon_ac, _recon_ac_dcm) logger.info(f"Store {label} recon at {_file_recon_ac}")