add tests

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

Author: yiheng-wang-nv

monai/data/image_reader.py

@@ -942,7 +942,8 @@ def _get_array_data_from_gpu(self, img, filename):
             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))
+        new_shape = (number_of_frames, rows, columns) if number_of_frames > 1 else (rows, columns)
+        data = buffer.view(dtype).reshape(new_shape)
 
         return data
 

tests/test_load_image.py

@@ -168,6 +168,16 @@ def get_data(self, _obj):
 # test reader consistency between PydicomReader and ITKReader on dicom data
 TEST_CASE_22 = ["tests/testing_data/CT_DICOM"]
 
+# test pydicom gpu reader
+TEST_CASE_GPU_5 = [{"reader": "PydicomReader", "to_gpu": True}, "tests/testing_data/CT_DICOM", (16, 16, 4), (16, 16, 4)]
+
+TEST_CASE_GPU_6 = [
+    {"reader": "PydicomReader", "ensure_channel_first": True, "force": True, "to_gpu": True},
+    "tests/testing_data/CT_DICOM",
+    (16, 16, 4),
+    (1, 16, 16, 4),
+]
+
 TESTS_META = []
 for track_meta in (False, True):
     TESTS_META.append([{}, (128, 128, 128), track_meta])
@@ -272,6 +282,26 @@ def test_itk_dicom_series_reader(self, input_param, filenames, expected_shape, e
         )
         self.assertTupleEqual(result.shape, expected_np_shape)
 
+    @SkipIfNoModule("pydicom")
+    @SkipIfNoModule("cupy")
+    @SkipIfNoModule("kvikio")
+    @parameterized.expand([TEST_CASE_GPU_5, TEST_CASE_GPU_6])
+    def test_pydicom_gpu_reader(self, input_param, filenames, expected_shape, expected_np_shape):
+        result = LoadImage(image_only=True, **input_param)(filenames)
+        self.assertEqual(result.meta["filename_or_obj"], f"{Path(filenames)}")
+        assert_allclose(
+            result.affine,
+            torch.tensor(
+                [
+                    [-0.488281, 0.0, 0.0, 125.0],
+                    [0.0, -0.488281, 0.0, 128.100006],
+                    [0.0, 0.0, 68.33333333, -99.480003],
+                    [0.0, 0.0, 0.0, 1.0],
+                ]
+            ),
+        )
+        self.assertTupleEqual(result.shape, expected_np_shape)
+
     def test_no_files(self):
         with self.assertRaisesRegex(RuntimeError, "list index out of range"):  # fname_regex excludes everything
             LoadImage(image_only=True, reader="PydicomReader", fname_regex=r"^(?!.*).*")("tests/testing_data/CT_DICOM")
@@ -318,6 +348,21 @@ def test_dicom_reader_consistency(self, filenames):
             np.testing.assert_allclose(pydicom_result, itk_result)
             np.testing.assert_allclose(pydicom_result.affine, itk_result.affine)
 
+    @SkipIfNoModule("pydicom")
+    @SkipIfNoModule("cupy")
+    @SkipIfNoModule("kvikio")
+    @parameterized.expand([TEST_CASE_22])
+    def test_pydicom_reader_gpu_cpu_consistency(self, filenames):
+        gpu_param = {"reader": "PydicomReader", "to_gpu": True}
+        cpu_param = {"reader": "PydicomReader", "to_gpu": False}
+        for affine_flag in [True, False]:
+            gpu_param["affine_lps_to_ras"] = affine_flag
+            cpu_param["affine_lps_to_ras"] = affine_flag
+            gpu_result = LoadImage(image_only=True, **gpu_param)(filenames)
+            cpu_result = LoadImage(image_only=True, **cpu_param)(filenames)
+            np.testing.assert_allclose(gpu_result.cpu(), cpu_result)
+            np.testing.assert_allclose(gpu_result.affine.cpu(), cpu_result.affine)
+
     def test_dicom_reader_consistency_single(self):
         itk_param = {"reader": "ITKReader"}
         pydicom_param = {"reader": "PydicomReader"}