implement image pre-loading in dataset

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