fix image pre-loading in dataset

mu_map/data/datasets.py

@@ -46,7 +46,6 @@ 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__()
 
@@ -73,12 +72,13 @@ class MuMapDataset(Dataset):
         self.pre_load_images()
 
     def pre_load_images(self):
+        print("Pre-loading images ...", end="\r")
         for i in range(len(self.table)):
-            row = self.table.iloc[index]
+            row = self.table.iloc[i]
             _id = row["id"]
 
             mu_map_file = os.path.join(self.dir_images, row[headers.file_mu_map])
-            mu_map = pydicom.dcmread(mu_map_file)
+            mu_map = pydicom.dcmread(mu_map_file).pixel_array
             if self.discard_mu_map_slices:
                 mu_map = discard_slices(row, mu_map)
             if self.bed_contours:
@@ -88,10 +88,11 @@ class MuMapDataset(Dataset):
             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)
+            recon = pydicom.dcmread(recon_file).pixel_array
             if self.align:
                 recon = align_images(recon, mu_map)
             self.reconstructions[_id] = recon
+        print("Pre-loading images done!")
 
     def __getitem__(self, index: int):
         row = self.table.iloc[index]
@@ -107,15 +108,15 @@ class MuMapDataset(Dataset):
         # mu_map = pydicom.dcmread(mu_map_file).pixel_array
 
         # if self.discard_mu_map_slices:
-            # mu_map = discard_slices(row, mu_map)
+        # 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)
+        # 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)
+        # recon = align_images(recon, mu_map)
 
         return recon, mu_map
 
@@ -129,21 +130,43 @@ if __name__ == "__main__":
     import argparse
 
     import cv2 as cv
-    
+
     from mu_map.util import to_grayscale, COLOR_WHITE
 
-    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")
+    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()
 
     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)
+    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)
@@ -156,7 +179,9 @@ if __name__ == "__main__":
         _image = to_grayscale(image[_slice], min_val=image.min(), max_val=image.max())
         _image = cv.resize(_image, (1024, 1024), cv.INTER_AREA)
         _text = f"{str(_slice):>{len(str(image.shape[0]))}}/{str(image.shape[0])}"
-        _image = cv.putText(_image, _text, (0, 30), cv.FONT_HERSHEY_SIMPLEX, 1, COLOR_WHITE, 3)
+        _image = cv.putText(
+            _image, _text, (0, 30), cv.FONT_HERSHEY_SIMPLEX, 1, COLOR_WHITE, 3
+        )
         return _image
 
     def combine_images(images, slices):
@@ -189,8 +214,8 @@ if __name__ == "__main__":
                 exit(0)
             elif key == ord("p"):
                 timeout = 0 if timeout > 0 else 100
-            elif key == 83: # right arrow key
+            elif key == 83:  # right arrow key
                 continue
-            elif key == 81: # left arrow key
+            elif key == 81:  # left arrow key
                 ir = max(ir - 2, 0)
                 im = max(im - 2, 0)