Merge branch 'main' of https://gitlab.ub.uni-bielefeld.de/thuxohl/mu-map

mu_map/dataset/default.py

@@ -62,6 +62,7 @@ class MuMapDataset(Dataset):
         bed_contours_file: Optional[str] = DEFAULT_BED_CONTOURS_FILENAME,
         discard_mu_map_slices: bool = True,
         align: bool = True,
+        scatter_correction: bool = False,
         transform_normalization: Transform = Transform(),
         transform_augmentation: Transform = Transform(),
         logger=None,
@@ -92,6 +93,8 @@ class MuMapDataset(Dataset):
 
         self.discard_mu_map_slices = discard_mu_map_slices
         self.align = align
+        self.scatter_correction = scatter_correction
+        self.header_recon = headers.file_recon_nac_sc if self.scatter_correction else headers.file_recon_nac_nsc
 
         self.reconstructions = {}
         self.mu_maps = {}
@@ -113,7 +116,7 @@ class MuMapDataset(Dataset):
                 for i in range(mu_map.shape[0]):
                     mu_map[i] = cv.drawContours(mu_map[i], [bed_contour], -1, 0.0, -1)
 
-            recon_file = os.path.join(self.dir_images, row[headers.file_recon_nac_nsc])
+            recon_file = os.path.join(self.dir_images, row[self.header_recon])
             recon = pydicom.dcmread(recon_file)
             recon = recon.pixel_array / recon[DCM_TAG_PIXEL_SCALE_FACTOR].value
             if self.align: