update pydicom reader

Signed-off-by: Yiheng Wang <[email protected]>

Author: yiheng-wang-nv

monai/data/image_reader.py

@@ -418,6 +418,10 @@ class PydicomReader(ImageReader):
             If provided, only the matched files will be included. For example, to include the file name
             "image_0001.dcm", the regular expression could be `".*image_(\\d+).dcm"`. Default to `""`.
             Set it to `None` to use `pydicom.misc.is_dicom` to match valid files.
+        to_gpu: If True, load the image into GPU memory using CuPy and Kvikio. This can accelerate data loading.
+            Default is False. CuPy and Kvikio are required for this option.
+            In practical use, it's recommended to add a warm up call before the actual loading.
+            A related tutorial will be prepared in the future, and the document will be updated accordingly.
         kwargs: additional args for `pydicom.dcmread` API. more details about available args:
             https://pydicom.github.io/pydicom/stable/reference/generated/pydicom.filereader.dcmread.html
             If the `get_data` function will be called
@@ -434,6 +438,7 @@ def __init__(
         prune_metadata: bool = True,
         label_dict: dict | None = None,
         fname_regex: str = "",
+        to_gpu: bool = False,
         **kwargs,
     ):
         super().__init__()
@@ -444,6 +449,33 @@ def __init__(
         self.prune_metadata = prune_metadata
         self.label_dict = label_dict
         self.fname_regex = fname_regex
+        if to_gpu and (not has_cp or not has_kvikio):
+            warnings.warn(
+                "PydicomReader: CuPy and/or Kvikio not installed for GPU loading, falling back to CPU loading."
+            )
+            to_gpu = False
+
+        if to_gpu:
+            self.warmup_kvikio()
+
+        self.to_gpu = to_gpu
+
+    def warmup_kvikio(self):
+        """
+        Warm up the Kvikio library to initialize the internal buffers, cuFile, GDS, etc.
+        This can accelerate the data loading process when `to_gpu` is set to True.
+        """
+        if has_cp and has_kvikio:
+            a = cp.arange(100)
+            with tempfile.NamedTemporaryFile() as tmp_file:
+                tmp_file_name = tmp_file.name
+                f = kvikio.CuFile(tmp_file_name, "w")
+                f.write(a)
+                f.close()
+
+                b = cp.empty_like(a)
+                f = kvikio.CuFile(tmp_file_name, "r")
+                f.read(b)
 
     def verify_suffix(self, filename: Sequence[PathLike] | PathLike) -> bool:
         """
@@ -475,19 +507,23 @@ def read(self, data: Sequence[PathLike] | PathLike, **kwargs):
         img_ = []
 
         filenames: Sequence[PathLike] = ensure_tuple(data)
+        self.filenames = filenames
         kwargs_ = self.kwargs.copy()
+        if self.to_gpu:
+            kwargs["defer_size"] = "100 KB"
         kwargs_.update(kwargs)
 
         self.has_series = False
 
-        for name in filenames:
+        for i, name in enumerate(filenames):
             name = f"{name}"
             if Path(name).is_dir():
                 # read DICOM series
                 if self.fname_regex is not None:
                     series_slcs = [slc for slc in glob.glob(os.path.join(name, "*")) if re.match(self.fname_regex, slc)]
                 else:
                     series_slcs = [slc for slc in glob.glob(os.path.join(name, "*")) if pydicom.misc.is_dicom(slc)]
+                self.filenames[i] = series_slcs
                 slices = []
                 for slc in series_slcs:
                     try:
@@ -502,7 +538,7 @@ def read(self, data: Sequence[PathLike] | PathLike, **kwargs):
                 img_.append(ds)
         return img_ if len(filenames) > 1 else img_[0]
 
-    def _combine_dicom_series(self, data: Iterable):
+    def _combine_dicom_series(self, data: Iterable, filenames: Sequence[PathLike]):
         """
         Combine dicom series (a list of pydicom dataset objects). Their data arrays will be stacked together at a new
         dimension as the last dimension.
@@ -522,25 +558,25 @@ def _combine_dicom_series(self, data: Iterable):
         """
         slices: list = []
         # for a dicom series
-        for slc_ds in data:
+        for slc_ds, filename in zip(data, filenames):
             if hasattr(slc_ds, "InstanceNumber"):
-                slices.append(slc_ds)
+                slices.append((slc_ds, filename))
             else:
-                warnings.warn(f"slice: {slc_ds.filename} does not have InstanceNumber tag, skip it.")
-        slices = sorted(slices, key=lambda s: s.InstanceNumber)
+                warnings.warn(f"slice: {filename} does not have InstanceNumber tag, skip it.")
+        slices = sorted(slices, key=lambda s: s[0].InstanceNumber)
 
         if len(slices) == 0:
             raise ValueError("the input does not have valid slices.")
 
-        first_slice = slices[0]
+        first_slice, first_filename = slices[0]
         average_distance = 0.0
-        first_array = self._get_array_data(first_slice)
+        first_array = self._get_array_data(first_slice, first_filename)
         shape = first_array.shape
         spacing = getattr(first_slice, "PixelSpacing", [1.0, 1.0, 1.0])
         prev_pos = getattr(first_slice, "ImagePositionPatient", (0.0, 0.0, 0.0))[2]
         stack_array = [first_array]
         for idx in range(1, len(slices)):
-            slc_array = self._get_array_data(slices[idx])
+            slc_array = self._get_array_data(slices[idx][0], slices[idx][1])
             slc_shape = slc_array.shape
             slc_spacing = getattr(slices[idx], "PixelSpacing", (1.0, 1.0, 1.0))
             slc_pos = getattr(slices[idx], "ImagePositionPatient", (0.0, 0.0, float(idx)))[2]
@@ -555,7 +591,10 @@ def _combine_dicom_series(self, data: Iterable):
         if len(slices) > 1:
             average_distance /= len(slices) - 1
             spacing.append(average_distance)
-            stack_array = np.stack(stack_array, axis=-1)
+            if self.to_gpu:
+                stack_array = cp.stack(stack_array, axis=-1)
+            else:
+                stack_array = np.stack(stack_array, axis=-1)
             stack_metadata = self._get_meta_dict(first_slice)
             stack_metadata["spacing"] = np.asarray(spacing)
             if hasattr(slices[-1], "ImagePositionPatient"):
@@ -597,29 +636,35 @@ def get_data(self, data) -> tuple[np.ndarray, dict]:
         if self.has_series is True:
             # a list, all objects within a list belong to one dicom series
             if not isinstance(data[0], list):
-                dicom_data.append(self._combine_dicom_series(data))
+                dicom_data.append(self._combine_dicom_series(data, self.filenames))
             # a list of list, each inner list represents a dicom series
             else:
-                for series in data:
-                    dicom_data.append(self._combine_dicom_series(series))
+                for i, series in enumerate(data):
+                    dicom_data.append(self._combine_dicom_series(series, self.filenames[i]))
         else:
             # a single pydicom dataset object
             if not isinstance(data, list):
                 data = [data]
-            for d in data:
+            for i, d in enumerate(data):
                 if hasattr(d, "SegmentSequence"):
-                    data_array, metadata = self._get_seg_data(d)
+                    data_array, metadata = self._get_seg_data(d, self.filenames[i])
                 else:
-                    data_array = self._get_array_data(d)
+                    data_array = self._get_array_data(d, self.filenames[i])
                     metadata = self._get_meta_dict(d)
                     metadata[MetaKeys.SPATIAL_SHAPE] = data_array.shape
                 dicom_data.append((data_array, metadata))
 
+
+        # TODO: the actual type is list[np.ndarray | cp.ndarray]
+        # should figure out how to define correct types without having cupy not found error
+        # https://github.com/Project-MONAI/MONAI/pull/8188#discussion_r1886645918
         img_array: list[np.ndarray] = []
         compatible_meta: dict = {}
 
         for data_array, metadata in ensure_tuple(dicom_data):
-            img_array.append(np.ascontiguousarray(np.swapaxes(data_array, 0, 1) if self.swap_ij else data_array))
+            if self.swap_ij:
+                data_array = cp.swapaxes(data_array, 0, 1) if self.to_gpu else np.swapaxes(data_array, 0, 1)
+            img_array.append(cp.ascontiguousarray(data_array) if self.to_gpu else np.ascontiguousarray(data_array))
             affine = self._get_affine(metadata, self.affine_lps_to_ras)
             metadata[MetaKeys.SPACE] = SpaceKeys.RAS if self.affine_lps_to_ras else SpaceKeys.LPS
             if self.swap_ij:
@@ -641,7 +686,7 @@ def get_data(self, data) -> tuple[np.ndarray, dict]:
 
             _copy_compatible_dict(metadata, compatible_meta)
 
-        return _stack_images(img_array, compatible_meta), compatible_meta
+        return _stack_images(img_array, compatible_meta, to_cupy=self.to_gpu), compatible_meta
 
     def _get_meta_dict(self, img) -> dict:
         """
@@ -713,7 +758,7 @@ def _get_affine(self, metadata: dict, lps_to_ras: bool = True):
             affine = orientation_ras_lps(affine)
         return affine
 
-    def _get_frame_data(self, img) -> Iterator:
+    def _get_frame_data(self, img, filename, array_data) -> Iterator:
         """
         yield frames and description from the segmentation image.
         This function is adapted from Highdicom:
@@ -752,47 +797,55 @@ def _get_frame_data(self, img) -> Iterator:
 
         if not hasattr(img, "PerFrameFunctionalGroupsSequence"):
             raise NotImplementedError(
-                f"To read dicom seg: {img.filename}, 'PerFrameFunctionalGroupsSequence' is required."
+                f"To read dicom seg: {filename}, 'PerFrameFunctionalGroupsSequence' is required."
             )
 
         frame_seg_nums = []
         for f in img.PerFrameFunctionalGroupsSequence:
             if not hasattr(f, "SegmentIdentificationSequence"):
                 raise NotImplementedError(
-                    f"To read dicom seg: {img.filename}, 'SegmentIdentificationSequence' is required for each frame."
+                    f"To read dicom seg: {filename}, 'SegmentIdentificationSequence' is required for each frame."
                 )
             frame_seg_nums.append(int(f.SegmentIdentificationSequence[0].ReferencedSegmentNumber))
 
-        frame_seg_nums_arr = np.array(frame_seg_nums)
+        frame_seg_nums_arr = cp.array(frame_seg_nums) if self.to_gpu else np.array(frame_seg_nums)
 
         seg_descriptions = {int(f.SegmentNumber): f for f in img.SegmentSequence}
 
-        for i in np.unique(frame_seg_nums_arr):
-            indices = np.where(frame_seg_nums_arr == i)[0]
-            yield (img.pixel_array[indices, ...], seg_descriptions[i])
+        for i in np.unique(frame_seg_nums_arr) if not self.to_gpu else cp.unique(frame_seg_nums_arr):
+            indices = np.where(frame_seg_nums_arr == i)[0] if not self.to_gpu else cp.where(frame_seg_nums_arr == i)[0]
+            yield (array_data[indices, ...], seg_descriptions[i])
 
-    def _get_seg_data(self, img):
+    def _get_seg_data(self, img, filename):
         """
         Get the array data and metadata of the segmentation image.
 
         Aegs:
             img: a Pydicom dataset object that has attribute "SegmentSequence".
+            filename: the file path of the image.
 
         """
 
         metadata = self._get_meta_dict(img)
         n_classes = len(img.SegmentSequence)
-        spatial_shape = list(img.pixel_array.shape)
+        array_data = self._get_array_data(img, filename)
+        spatial_shape = list(array_data.shape)
         spatial_shape[0] = spatial_shape[0] // n_classes
 
         if self.label_dict is not None:
             metadata["labels"] = self.label_dict
-            all_segs = np.zeros([*spatial_shape, len(self.label_dict)])
+            if self.to_gpu:
+                all_segs = cp.zeros([*spatial_shape, len(self.label_dict)], dtype=array_data.dtype)
+            else:
+                all_segs = np.zeros([*spatial_shape, len(self.label_dict)], dtype=array_data.dtype)
         else:
             metadata["labels"] = {}
-            all_segs = np.zeros([*spatial_shape, n_classes])
+            if self.to_gpu:
+                all_segs = cp.zeros([*spatial_shape, n_classes], dtype=array_data.dtype)
+            else:
+                all_segs = np.zeros([*spatial_shape, n_classes], dtype=array_data.dtype)
 
-        for i, (frames, description) in enumerate(self._get_frame_data(img)):
+        for i, (frames, description) in enumerate(self._get_frame_data(img, filename, array_data)):
             segment_label = getattr(description, "SegmentLabel", f"label_{i}")
             class_name = getattr(description, "SegmentDescription", segment_label)
             if class_name not in metadata["labels"].keys():
@@ -840,19 +893,51 @@ def _get_seg_data(self, img):
 
         return all_segs, metadata
 
-    def _get_array_data(self, img):
+    def _get_array_data(self, img, filename):
         """
         Get the array data of the image. If `RescaleSlope` and `RescaleIntercept` are available, the raw array data
-        will be rescaled. The output data has the dtype np.float32 if the rescaling is applied.
+        will be rescaled. The output data has the dtype float32 if the rescaling is applied.
 
         Args:
             img: a Pydicom dataset object.
+            filename: the file path of the image.
 
         """
         # process Dicom series
-        if not hasattr(img, "pixel_array"):
-            raise ValueError(f"dicom data: {img.filename} does not have pixel_array.")
-        data = img.pixel_array
+
+        if self.to_gpu:
+            rows = img.Rows
+            columns = img.Columns
+            bits_allocated = img.BitsAllocated
+            samples_per_pixel = img.SamplesPerPixel
+            number_of_frames = getattr(img, 'NumberOfFrames', 1)
+            pixel_representation = img.PixelRepresentation
+
+            if bits_allocated == 8:
+                dtype = cp.int8 if pixel_representation == 1 else cp.uint8
+            elif bits_allocated == 16:
+                dtype = cp.int16 if pixel_representation == 1 else cp.uint16
+            elif bits_allocated == 32:
+                dtype = cp.int32 if pixel_representation == 1 else cp.uint32
+            else:
+                raise ValueError("Unsupported BitsAllocated value")
+
+            bytes_per_pixel = bits_allocated // 8
+            total_pixels = rows * columns * samples_per_pixel * number_of_frames
+            expected_pixel_data_length = total_pixels * bytes_per_pixel
+
+            offset = img.get_item(0x7FE00010, keep_deferred=True).value_tell
+
+            with kvikio.CuFile(filename, "r") as f:
+                buffer = cp.empty(expected_pixel_data_length, dtype=cp.int8)
+                f.read(buffer, expected_pixel_data_length, offset)
+
+            data = buffer.view(dtype).reshape((number_of_frames, rows, columns))
+        
+        else:
+            if not hasattr(img, "pixel_array"):
+                raise ValueError(f"dicom data: {filename} does not have pixel_array.")
+            data = img.pixel_array
 
         slope, offset = 1.0, 0.0
         rescale_flag = False
@@ -862,8 +947,12 @@ def _get_array_data(self, img):
         if hasattr(img, "RescaleIntercept"):
             offset = img.RescaleIntercept
             rescale_flag = True
+
         if rescale_flag:
-            data = data.astype(np.float32) * slope + offset
+            if self.to_gpu:
+                data = data.astype(cp.float32) * slope + offset
+            else:
+                data = data.astype(np.float32) * slope + offset
 
         return data
 
@@ -884,8 +973,6 @@ class NibabelReader(ImageReader):
             Default is False. CuPy and Kvikio are required for this option.
             Note: For compressed NIfTI files, some operations may still be performed on CPU memory,
             and the acceleration may not be significant. In some cases, it may be slower than loading on CPU.
-            In practical use, it's recommended to add a warm up call before the actual loading.
-            A related tutorial will be prepared in the future, and the document will be updated accordingly.
         kwargs: additional args for `nibabel.load` API. more details about available args:
             https://github.com/nipy/nibabel/blob/master/nibabel/loadsave.py