changes after code review

rijobro · rijobro · commit 9e7b3d68022d · 2022-04-29T11:28:49.000+01:00
Signed-off-by: Richard Brown &lt;33289025+rijobro@users.noreply.github.com&gt;
diff --git a/monai/data/meta_tensor.py b/monai/data/meta_tensor.py
@@ -17,8 +17,10 @@
 
 import torch
 
+from monai.config.type_definitions import NdarrayTensor
 from monai.data.meta_obj import MetaObj, get_track_meta, get_track_transforms
 from monai.data.utils import decollate_batch, list_data_collate
+from monai.transforms.utils import remove_extra_metadata
 from monai.utils.enums import PostFix
 
 __all__ = ["MetaTensor"]
@@ -232,3 +234,33 @@ def affine(self) -> torch.Tensor:
     def affine(self, d: torch.Tensor) -> None:
         """Set the affine."""
         self.meta["affine"] = d
+
+    @staticmethod
+    def ensure_torch_and_prune_meta(im: NdarrayTensor, meta: dict):
+        """
+        Convert the image to `torch.Tensor`. If `affine` is in the `meta` dictionary,
+        convert that to `torch.Tensor`, too. Remove any superfluous metadata.
+
+        Args:
+            im: Input image (`np.ndarray` or `torch.Tensor`)
+            meta: Metadata dictionary.
+
+        Returns:
+            By default, a `MetaTensor` is returned.
+            However, if `get_track_meta()` is `False`, a `torch.Tensor` is returned.
+        """
+        img = torch.as_tensor(im)
+
+        # if not tracking metadata, return `torch.Tensor`
+        if not get_track_meta() or meta is None:
+            return img
+
+        # ensure affine is of type `torch.Tensor`
+        if "affine" in meta:
+            meta["affine"] = torch.as_tensor(meta["affine"])
+
+        # remove any superfluous metadata.
+        remove_extra_metadata(meta)
+
+        # return the `MetaTensor`
+        return MetaTensor(img, meta=meta)
diff --git a/monai/transforms/__init__.py b/monai/transforms/__init__.py
@@ -569,6 +569,7 @@
     generate_label_classes_crop_centers,
     generate_pos_neg_label_crop_centers,
     generate_spatial_bounding_box,
+    get_extra_metadata_keys,
     get_extreme_points,
     get_largest_connected_component_mask,
     get_number_image_type_conversions,
@@ -582,6 +583,8 @@
     map_spatial_axes,
     print_transform_backends,
     rand_choice,
+    remove_extra_metadata,
+    remove_keys,
     rescale_array,
     rescale_array_int_max,
     rescale_instance_array,
diff --git a/monai/transforms/io/array.py b/monai/transforms/io/array.py
@@ -29,7 +29,6 @@
 from monai.data import image_writer
 from monai.data.folder_layout import FolderLayout
 from monai.data.image_reader import ImageReader, ITKReader, NibabelReader, NumpyReader, PILReader
-from monai.data.meta_obj import get_track_meta
 from monai.data.meta_tensor import MetaTensor
 from monai.transforms.transform import Transform
 from monai.transforms.utility.array import EnsureChannelFirst
@@ -247,45 +246,11 @@ def __call__(self, filename: Union[Sequence[PathLike], PathLike], reader: Option
         meta_data = switch_endianness(meta_data, "<")
 
         meta_data[Key.FILENAME_OR_OBJ] = f"{ensure_tuple(filename)[0]}"  # Path obj should be strings for data loader
-        img = self.join_im_and_meta(img_array, meta_data)
+        img = MetaTensor.ensure_torch_and_prune_meta(img_array, meta_data)
         if self.ensure_channel_first:
             img = EnsureChannelFirst()(img)
         return img
 
-    @staticmethod
-    def join_im_and_meta(im, meta: dict):
-        img = torch.as_tensor(im)
-
-        # if not tracking metadata, return torch.Tensor
-        if not get_track_meta() or meta is None:
-            return img
-
-        if "affine" in meta:
-            meta["affine"] = torch.as_tensor(meta["affine"])
-
-        # TODO: delete extra metadata
-        for i in range(8):
-            for k in ("dim", "pixdim"):
-                if f"{k}[{i}]" in meta:
-                    del meta[f"{k}[{i}]"]
-        for k in (
-            #     "original_affine",
-            #     "spatial_shape",
-            #     "spacing",
-            "srow_x",
-            "srow_y",
-            "srow_z",
-            "quatern_b",
-            "quatern_c",
-            "quatern_d",
-            "qoffset_x",
-            "qoffset_y",
-            "qoffset_z",
-        ):
-            if k in meta:
-                del meta[k]
-        return MetaTensor(img, meta=meta)
-
 
 class SaveImage(Transform):
     """
diff --git a/monai/transforms/utility/dictionary.py b/monai/transforms/utility/dictionary.py
@@ -427,11 +427,11 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
                         if dim > 0:  # don't update affine if channel dim
                             affine = d[split_meta_key]["affine"]  # type: ignore
                             ndim = len(affine)
-                            shift = (
-                                torch.eye(ndim, device=affine.device, dtype=affine.dtype)
-                                if isinstance(affine, torch.Tensor)
-                                else np.eye(ndim)
-                            )
+                            shift: NdarrayOrTensor
+                            if isinstance(affine, torch.Tensor):
+                                shift = torch.eye(ndim, device=affine.device, dtype=affine.dtype)
+                            else:
+                                shift = np.eye(ndim)
                             shift[dim - 1, -1] = i  # type: ignore
                             d[split_meta_key]["affine"] = d[split_meta_key]["affine"] @ shift  # type: ignore
 
diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py
@@ -105,6 +105,9 @@
     "convert_pad_mode",
     "convert_to_contiguous",
     "get_unique_labels",
+    "remove_keys",
+    "remove_extra_metadata",
+    "get_extra_metadata_keys",
 ]
 
 
@@ -1573,5 +1576,68 @@ def convert_to_contiguous(data, **kwargs):
     return data
 
 
+def remove_keys(data: dict, keys: List[str]) -> None:
+    """
+    Remove keys from a dictionary. Operates in-place so nothing is returned.
+
+    Args:
+        data: dictionary to be modified.
+        keys: keys to be deleted from dictionary.
+
+    Returns:
+        `None`
+    """
+    for k in keys:
+        _ = data.pop(k, None)
+
+
+def remove_extra_metadata(meta: dict) -> None:
+    """
+    Remove extra metadata from the dictionary. Operates in-place so nothing is returned.
+
+    Args:
+        meta: dictionary containing metadata to be modified.
+
+    Returns:
+        `None`
+    """
+    keys = get_extra_metadata_keys()
+    remove_keys(data=meta, keys=keys)
+
+
+def get_extra_metadata_keys() -> List[str]:
+    """
+    Get a list of unnecessary keys for metadata that can be removed.
+
+    Returns:
+        List of keys to be removed.
+    """
+    keys = [
+        "srow_x",
+        "srow_y",
+        "srow_z",
+        "quatern_b",
+        "quatern_c",
+        "quatern_d",
+        "qoffset_x",
+        "qoffset_y",
+        "qoffset_z",
+        "dim",
+        "pixdim",
+        *[f"dim[{i}]" for i in range(8)],
+        *[f"pixdim[{i}]" for i in range(8)],
+    ]
+
+    # TODO: it would be good to remove these, but they are currently being used in the
+    # codebase.
+    # keys += [
+    #     "original_affine",
+    #     "spatial_shape",
+    #     "spacing",
+    # ]
+
+    return keys
+
+
 if __name__ == "__main__":
     print_transform_backends()
diff --git a/tests/test_load_image.py b/tests/test_load_image.py
@@ -242,7 +242,7 @@ def test_kwargs(self):
             reader = ITKReader()
             img = reader.read(filename, fallback_only=False)
             result_raw = reader.get_data(img)
-            result_raw = LoadImage.join_im_and_meta(*result_raw)
+            result_raw = MetaTensor.ensure_torch_and_prune_meta(*result_raw)
             self.assertTupleEqual(result.shape, result_raw.shape)
 
     def test_my_reader(self):
