diff --git a/mu_map/dataset/default.py b/mu_map/dataset/default.py index 385ca93470888ba413da1701b5680dafd82d0dc2..bb5723571e6804765b3385fe3c10350260276cf0 100644 --- a/mu_map/dataset/default.py +++ b/mu_map/dataset/default.py @@ -17,7 +17,8 @@ from mu_map.data.remove_bed import ( from mu_map.data.review_mu_map import discard_slices from mu_map.data.split import split_csv from mu_map.dataset.transform import Transform -from mu_map.dataset.util import align_images, load_dcm_img +from mu_map.dataset.util import align_images +from mu_map.file.dicom import load_dcm_img from mu_map.logging import get_logger @@ -97,9 +98,13 @@ class MuMapDataset(Dataset): _id = row[headers.id] mu_map_file = os.path.join(self.dir_images, row[headers.file_mu_map]) - mu_map = load_dcm_img(mu_map_file) + mu_map = load_dcm_img(mu_map_file, direction=1) + + recon_file = os.path.join(self.dir_images, row[self.header_recon]) + recon = load_dcm_img(recon_file, direction=1) + if self.discard_mu_map_slices: - mu_map = discard_slices(row, mu_map) + mu_map, recon = discard_slices(row, mu_map, recon) if self.bed_contours: if _id in self.bed_contours: bed_contour = self.bed_contours[_id] @@ -107,8 +112,6 @@ class MuMapDataset(Dataset): else: logger.warning(f"Could not find bed contour for id {_id}") - recon_file = os.path.join(self.dir_images, row[self.header_recon]) - recon = load_dcm_img(recon_file) if self.align: recon, mu_map = align_images(recon, mu_map)