-
Tamino Huxohl authoredTamino Huxohl authored
compute_recons.py 4.30 KiB
"""
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}")