diff --git a/mu_map/data/datasets.py b/mu_map/data/datasets.py
index 752ac6a9c6d82f5cf6a0390325b490e00ccd9b9a..64fd191d1421b303ea451d2bba85ebbb97e9b466 100644
--- a/mu_map/data/datasets.py
+++ b/mu_map/data/datasets.py
@@ -6,6 +6,7 @@ import pydicom
 import numpy as np
 from torch.utils.data import Dataset
 
+from mu_map.data.prepare import headers
 from mu_map.data.remove_bed import DEFAULT_BED_CONTOURS_FILENAME, load_contours
 from mu_map.data.review_mu_map import discard_slices
 
@@ -69,8 +70,8 @@ class MuMapDataset(Dataset):
     def __getitem__(self, index: int):
         row = self.table.iloc[index]
 
-        recon_file = os.path.join(self.dir_images, row["file_recon_no_ac"])
-        mu_map_file = os.path.join(self.dir_images, row["file_mu_map"])
+        recon_file = os.path.join(self.dir_images, row[header.file_recon_nac_nsc])
+        mu_map_file = os.path.join(self.dir_images, row[headers.file_mu_map])
 
         recon = pydicom.dcmread(recon_file).pixel_array
         mu_map = pydicom.dcmread(mu_map_file).pixel_array