From 81dde9cff10b4b2d835e6fc3bcab809844024be0 Mon Sep 17 00:00:00 2001
From: Tamino Huxohl <thuxohl@techfak.uni-bielefeld.de>
Date: Thu, 1 Sep 2022 16:04:26 +0200
Subject: [PATCH] implement image pre-loading in dataset
---
mu_map/data/datasets.py | 54 ++++++++++++++++++++++++++++++++---------
1 file changed, 42 insertions(+), 12 deletions(-)
diff --git a/mu_map/data/datasets.py b/mu_map/data/datasets.py
index 7d0d6c7..deadaed 100644
--- a/mu_map/data/datasets.py
+++ b/mu_map/data/datasets.py
@@ -46,6 +46,7 @@ class MuMapDataset(Dataset):
bed_contours_file: Optional[str] = DEFAULT_BED_CONTOURS_FILENAME,
discard_mu_map_slices: bool = True,
align: bool = True,
+ pathes_per_image
):
super().__init__()
@@ -67,25 +68,54 @@ class MuMapDataset(Dataset):
self.discard_mu_map_slices = discard_mu_map_slices
self.align = align
+ self.reconstructions = {}
+ self.mu_maps = {}
+ self.pre_load_images()
+
+ def pre_load_images(self):
+ for i in range(len(self.table)):
+ row = self.table.iloc[index]
+ _id = row["id"]
+
+ mu_map_file = os.path.join(self.dir_images, row[headers.file_mu_map])
+ mu_map = pydicom.dcmread(mu_map_file)
+ if self.discard_mu_map_slices:
+ mu_map = discard_slices(row, mu_map)
+ if self.bed_contours:
+ bed_contour = self.bed_contours[row["id"]]
+ for i in range(mu_map.shape[0]):
+ mu_map[i] = cv.drawContours(mu_map[i], [bed_contour], -1, 0.0, -1)
+ self.mu_maps[_id] = mu_map
+
+ recon_file = os.path.join(self.dir_images, row[headers.file_recon_nac_nsc])
+ recon = pydicom.dcmread(recon_file)
+ if self.align:
+ recon = align_images(recon, mu_map)
+ self.reconstructions[_id] = recon
+
def __getitem__(self, index: int):
row = self.table.iloc[index]
+ _id = row["id"]
+
+ recon = self.reconstructions[_id]
+ mu_map = self.mu_maps[_id]
- 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_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
- mu_map = pydicom.dcmread(mu_map_file).pixel_array
+ # recon = pydicom.dcmread(recon_file).pixel_array
+ # mu_map = pydicom.dcmread(mu_map_file).pixel_array
- if self.discard_mu_map_slices:
- mu_map = discard_slices(row, mu_map)
+ # if self.discard_mu_map_slices:
+ # mu_map = discard_slices(row, mu_map)
- if self.bed_contours:
- bed_contour = self.bed_contours[row["id"]]
- for i in range(mu_map.shape[0]):
- mu_map[i] = cv.drawContours(mu_map[i], [bed_contour], -1, 0.0, -1)
+ # if self.bed_contours:
+ # bed_contour = self.bed_contours[row["id"]]
+ # for i in range(mu_map.shape[0]):
+ # mu_map[i] = cv.drawContours(mu_map[i], [bed_contour], -1, 0.0, -1)
- if self.align:
- recon = align_images(recon, mu_map)
+ # if self.align:
+ # recon = align_images(recon, mu_map)
return recon, mu_map
--
GitLab