import argparse
import os
from typing import List
from mu_map.file.dicom import DICOM, DICOMTime
def get_files_recursive(_dir: str) -> List[str]:
"""
Recursively get all files in a directory. This means that all
sub-directories are recursively searched until a file is reached.
:param _dir: the directory of which files are listed
:return: a list of files in the directory and its sub-directories
"""
files: List[str] = []
for _file in os.listdir(_dir):
_file = os.path.join(_dir, _file)
if os.path.isdir(_file):
files.extend(get_files_recursive(_file))
else:
files.append(_file)
return files
def is_scatter_corrected(dcm: DICOM) -> bool:
return "SC " in dcm.SeriesDescription
def is_attenuation_corrected(dcm: DICOM) -> bool:
return not ("NoAC" in dcm.SeriesDescription)
def shows_segments(dcm: DICOM) -> bool:
return "segment" in dcm.SeriesDescription.lower()
def get_type(dcm: DICOM) -> str:
description = dcm.SeriesDescription.lower()
if "DLAC" in description:
return "dl"
elif "CTAC" in description:
return "ct"
else:
return "symbia"
headers = argparse.Namespace()
headers.id = "id"
headers.scatter_correction = "scatter_correction"
headers.sc = headers.scatter_correction
headers.attenuation_correction = "attenuation_correction"
headers.ac = headers.attenuation_correction
headers.segments = "segments"
headers.mu_map_type = "mu_map_type"
headers.type = headers.mu_map_type
headers.file = "file"
if __name__ == "__main__":
import cv2 as cv
import numpy as np
import pandas as pd
from mu_map.data.prepare import get_protocol
from mu_map.data.prepare import headers as meta_headers
from mu_map.file.dicom import load_dcm, DCM_TAG_PIXEL_SCALE_FACTOR, DICOM
from mu_map.logging import add_logging_args, get_logger_by_args
parser = argparse.ArgumentParser(
description="Store/Sort in polar maps for further processing",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument(
"--polar_map_dir",
type=str,
required=True,
help="directory where the raw DICOM files of polar maps are stored",
)
parser.add_argument(
"--out_dir", type=str, required=True, help="directory where files are output to"
)
parser.add_argument(
"--images_dir",
type=str,
default="images",
help="directory under <out_dir> where images of polar maps are stored",
)
parser.add_argument(
"--csv",
type=str,
default="polar_maps.csv",
help="file unter <out_dir> where meta information is stored",
)
parser.add_argument(
"--meta_csv",
type=str,
required=True,
help="the csv file containing meta information of the dataset of which the generated polar maps are",
)
parser.add_argument(
"--rect_upper",
type=int,
nargs=4,
default=[22, 1499, 22 + 588, 1500 + 588],
help="rectangle as [top, left, bottom, right] coordinates where the polar map is found in the DICOM images",
)
parser.add_argument(
"--rect_lower",
type=int,
nargs=4,
default=[634, 1499, 634 + 588, 1500 + 588],
help="rectangle as [top, left, bottom, right] coordinates where the polar map is found in the DICOM images",
)
parser.add_argument(
"-i",
"--interactive",
action="store_true",
help="start interactive mode where every polar map is displayed and confirmation is queried",
)
add_logging_args(parser, defaults={"--logfile": "prepare.log"})
args = parser.parse_args()
args.images_dir = os.path.join(args.out_dir, args.images_dir)
args.csv = os.path.join(args.out_dir, args.csv)
args.logfile = os.path.join(args.out_dir, args.logfile)
args.loglevel = "DEBUG" if args.interactive else args.loglevel
if args.interactive:
print(
"""
Start in interactive mode.
This mode automatically sets the loglevel to debug so that you can see what is going on.
Use the following keys on the displayed polar map:
n(ext): go to the next image without saving the polar map
s(ave): save the polar map and go to the next
q(uit): quit the application
"""
)
logger = get_logger_by_args(args)
logger.info(args)
if not os.path.exists(args.out_dir):
os.mkdir(args.out_dir)
if not os.path.exists(args.images_dir):
os.mkdir(args.images_dir)
meta = pd.read_csv(args.meta_csv)
meta[meta_headers.datetime_acquisition] = pd.to_datetime(meta[meta_headers.datetime_acquisition])
dcm_files = sorted(get_files_recursive(args.polar_map_dir))
data = pd.DataFrame(
{
headers.id: [],
headers.sc: [],
headers.ac: [],
headers.segments: [],
headers.type: [],
headers.file: [],
}
)
data[headers.id] = data[headers.id].astype(int)
data[headers.sc] = data[headers.sc].astype(bool)
data[headers.ac] = data[headers.ac].astype(bool)
data[headers.segments] = data[headers.segments].astype(bool)
data[headers.type] = data[headers.type].astype(str)
data[headers.file] = data[headers.file].astype(str)
for dcm_file in dcm_files:
logger.debug(f"Process file {dcm_file}")
dcm, img = load_dcm(dcm_file)
protocol = get_protocol(dcm)
patient_id = int(dcm.PatientID)
meta_rows = meta[
(meta[meta_headers.patient_id] == patient_id)
& (meta[meta_headers.protocol] == protocol)
& (meta[meta_headers.datetime_acquisition] == DICOMTime.Acquisition.to_datetime(dcm))
]
assert len(meta_rows) == 1
meta_row = meta_rows.iloc[0]
row = {
headers.id: meta_row[meta_headers.id],
headers.sc: is_scatter_corrected(dcm),
headers.ac: is_attenuation_corrected(dcm),
headers.segments: shows_segments(dcm),
}
types = dcm.SeriesDescription.split("-")[-1]
types = types[:-1].strip()
types = types.split("/")
for _type, rect in zip(types, [args.rect_upper, args.rect_lower]):
top, left, bottom, right = rect
polar_map = img[top:bottom, left:right]
polar_map = cv.cvtColor(polar_map, cv.COLOR_RGB2BGR)
row[headers.type] = _type[:2].lower()
_seg = "segments" if row[headers.segments] else "gray"
_file = f"{row[headers.id]:04d}-{_seg}_ac_nsc_{row[headers.type]}.png"
row[headers.file] = _file
_file = os.path.join(args.images_dir, _file)
logger.debug(f"For series {dcm.SeriesDescription} store row {row}")
store = not args.interactive
if args.interactive:
while True:
cv.imshow("Polar Map", cv.resize(polar_map, (1024, 1024)))
key = cv.waitKey(0)
if key == ord("q"):
exit(0)
elif key == ord("n"):
break
elif key == ord("s"):
store = True
break
if store:
cv.imwrite(_file, polar_map)
data = pd.concat((data, pd.DataFrame(row, index=[0])), ignore_index=True)
data = data.sort_values(by=[headers.id, headers.file])
data.to_csv(args.csv, index=False)