diff --git a/mu_map/data/datasets.py b/mu_map/data/datasets.py
index 64fd191d1421b303ea451d2bba85ebbb97e9b466..7d0d6c746bf2d3fb6bd3ca1ef8aa99683790c98a 100644
--- a/mu_map/data/datasets.py
+++ b/mu_map/data/datasets.py
@@ -70,7 +70,7 @@ class MuMapDataset(Dataset):
     def __getitem__(self, index: int):
         row = self.table.iloc[index]
 
-        recon_file = os.path.join(self.dir_images, row[header.file_recon_nac_nsc])
+        recon_file = os.path.join(self.dir_images, row[headers.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
@@ -105,9 +105,15 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Visualize the images of a MuMapDataset", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument("dataset_dir", type=str, help="the directory from which the dataset is loaded")
     parser.add_argument("--unaligned", action="store_true", help="do not perform center alignment of reconstruction an mu-map slices")
+    parser.add_argument("--show_bed", action="store_true", help="do not remove the bed contour from the mu map")
+    parser.add_argument("--full_mu_map", action="store_true", help="do not remove broken slices of the mu map")
     args = parser.parse_args()
 
-    dataset = MuMapDataset(args.dataset_dir, align=not args.unaligned)
+    align = not args.unaligned
+    discard_mu_map_slices = not args.full_mu_map
+    bed_contours_file = None if args.show_bed else DEFAULT_BED_CONTOURS_FILENAME
+
+    dataset = MuMapDataset(args.dataset_dir, align=align, discard_mu_map_slices=discard_mu_map_slices, bed_contours_file=bed_contours_file)
 
     wname = "Dataset"
     cv.namedWindow(wname, cv.WINDOW_NORMAL)