Feature/2231 transforms

monai/apps/deepgrow/dataset.py

@@ -17,6 +17,7 @@
 
 from monai.transforms import AsChannelFirstd, Compose, LoadImaged, Orientationd, Spacingd
 from monai.utils import GridSampleMode
+from monai.utils.misc import convert_data_type
 
 
 def create_dataset(
@@ -82,7 +83,7 @@ def create_dataset(
     if not len(datalist):
         raise ValueError("Input datalist is empty")
 
-    transforms = _default_transforms(image_key, label_key, pixdim) if transforms is None else transforms
+    transforms = transforms or _default_transforms(image_key, label_key, pixdim)
     new_datalist = []
     for idx in range(len(datalist)):
         if limit and idx >= limit:
@@ -133,25 +134,34 @@ def _default_transforms(image_key, label_key, pixdim):
 
 
 def _save_data_2d(vol_idx, vol_image, vol_label, dataset_dir, relative_path):
+    if vol_image is not None:
+        vol_image_np, *_ = convert_data_type(vol_image, np.ndarray)
+    else:
+        vol_image_np = vol_image
+    if vol_label is not None:
+        vol_label_np, *_ = convert_data_type(vol_label, np.ndarray)
+    else:
+        vol_label_np = vol_label
+
     data_list = []
 
-    if len(vol_image.shape) == 4:
+    if len(vol_image_np.shape) == 4:
         logging.info(
             "4D-Image, pick only first series; Image: {}; Label: {}".format(
-                vol_image.shape, vol_label.shape if vol_label is not None else None
+                vol_image_np.shape, vol_label_np.shape if vol_label_np is not None else None
             )
         )
-        vol_image = vol_image[0]
-        vol_image = np.moveaxis(vol_image, -1, 0)
+        vol_image_np = vol_image_np[0]
+        vol_image_np = np.moveaxis(vol_image_np, -1, 0)
 
     image_count = 0
     label_count = 0
     unique_labels_count = 0
-    for sid in range(vol_image.shape[0]):
-        image = vol_image[sid, ...]
-        label = vol_label[sid, ...] if vol_label is not None else None
+    for sid in range(vol_image_np.shape[0]):
+        image = vol_image_np[sid, ...]
+        label = vol_label_np[sid, ...] if vol_label_np is not None else None
 
-        if vol_label is not None and np.sum(label) == 0:
+        if vol_label_np is not None and np.sum(label) == 0:
             continue
 
         image_file_prefix = "vol_idx_{:0>4d}_slice_{:0>3d}".format(vol_idx, sid)
@@ -163,7 +173,7 @@ def _save_data_2d(vol_idx, vol_image, vol_label, dataset_dir, relative_path):
         image_count += 1
 
         # Test Data
-        if vol_label is None:
+        if vol_label_np is None:
             data_list.append(
                 {
                     "image": image_file.replace(dataset_dir + os.pathsep, "") if relative_path else image_file,
@@ -200,9 +210,9 @@ def _save_data_2d(vol_idx, vol_image, vol_label, dataset_dir, relative_path):
     logging.info(
         "{} => Image Shape: {} => {}; Label Shape: {} => {}; Unique Labels: {}".format(
             vol_idx,
-            vol_image.shape,
+            vol_image_np.shape,
             image_count,
-            vol_label.shape if vol_label is not None else None,
+            vol_label_np.shape if vol_label_np is not None else None,
             label_count,
             unique_labels_count,
         )
@@ -211,16 +221,25 @@ def _save_data_2d(vol_idx, vol_image, vol_label, dataset_dir, relative_path):
 
 
 def _save_data_3d(vol_idx, vol_image, vol_label, dataset_dir, relative_path):
+    if vol_image is not None:
+        vol_image_np, *_ = convert_data_type(vol_image, np.ndarray)
+    else:
+        vol_image_np = vol_image
+    if vol_label is not None:
+        vol_label_np, *_ = convert_data_type(vol_label, np.ndarray)
+    else:
+        vol_label_np = vol_label
+
     data_list = []
 
-    if len(vol_image.shape) == 4:
+    if len(vol_image_np.shape) == 4:
         logging.info(
             "4D-Image, pick only first series; Image: {}; Label: {}".format(
-                vol_image.shape, vol_label.shape if vol_label is not None else None
+                vol_image_np.shape, vol_label_np.shape if vol_label_np is not None else None
             )
         )
-        vol_image = vol_image[0]
-        vol_image = np.moveaxis(vol_image, -1, 0)
+        vol_image_np = vol_image_np[0]
+        vol_image_np = np.moveaxis(vol_image_np, -1, 0)
 
     image_count = 0
     label_count = 0
@@ -231,19 +250,19 @@ def _save_data_3d(vol_idx, vol_image, vol_label, dataset_dir, relative_path):
     image_file += ".npy"
 
     os.makedirs(os.path.join(dataset_dir, "images"), exist_ok=True)
-    np.save(image_file, vol_image)
+    np.save(image_file, vol_image_np)
     image_count += 1
 
     # Test Data
-    if vol_label is None:
+    if vol_label_np is None:
         data_list.append(
             {
                 "image": image_file.replace(dataset_dir + os.pathsep, "") if relative_path else image_file,
             }
         )
     else:
         # For all Labels
-        unique_labels = np.unique(vol_label.flatten())
+        unique_labels = np.unique(vol_label_np.flatten())
         unique_labels = unique_labels[unique_labels != 0]
         unique_labels_count = max(unique_labels_count, len(unique_labels))
 
@@ -252,7 +271,7 @@ def _save_data_3d(vol_idx, vol_image, vol_label, dataset_dir, relative_path):
             label_file = os.path.join(dataset_dir, "labels", label_file_prefix)
             label_file += ".npy"
 
-            curr_label = (vol_label == idx).astype(np.float32)
+            curr_label = (vol_label_np == idx).astype(np.float32)
             os.makedirs(os.path.join(dataset_dir, "labels"), exist_ok=True)
             np.save(label_file, curr_label)
 
@@ -271,9 +290,9 @@ def _save_data_3d(vol_idx, vol_image, vol_label, dataset_dir, relative_path):
     logging.info(
         "{} => Image Shape: {} => {}; Label Shape: {} => {}; Unique Labels: {}".format(
             vol_idx,
-            vol_image.shape,
+            vol_image_np.shape,
             image_count,
-            vol_label.shape if vol_label is not None else None,
+            vol_label_np.shape if vol_label_np is not None else None,
             label_count,
             unique_labels_count,
         )

monai/apps/pathology/utils.py

@@ -9,13 +9,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Union
+from typing import List
 
 import numpy as np
-import torch
 
 from monai.transforms.post.array import ProbNMS
 from monai.utils import optional_import
+from monai.utils.enums import DataObjects
 
 measure, _ = optional_import("skimage.measure")
 ndimage, _ = optional_import("scipy.ndimage")
@@ -67,7 +67,7 @@ class PathologyProbNMS(ProbNMS):
 
     def __call__(
         self,
-        probs_map: Union[np.ndarray, torch.Tensor],
+        probs_map: DataObjects.Images,
         resolution_level: int = 0,
     ):
         """

monai/data/csv_saver.py

@@ -12,12 +12,13 @@
 import os
 import warnings
 from collections import OrderedDict
-from typing import Dict, Optional, Union
+from typing import Dict, Optional
 
 import numpy as np
 import torch
 
 from monai.utils import ImageMetaKey as Key
+from monai.utils.enums import DataObjects
 
 
 class CSVSaver:
@@ -75,7 +76,7 @@ def finalize(self) -> None:
         # clear cache content after writing
         self.reset_cache()
 
-    def save(self, data: Union[torch.Tensor, np.ndarray], meta_data: Optional[Dict] = None) -> None:
+    def save(self, data: DataObjects.Images, meta_data: Optional[Dict] = None) -> None:
         """Save data into the cache dictionary. The metadata should have the following key:
             - ``'filename_or_obj'`` -- save the data corresponding to file name or object.
         If meta_data is None, use the default index from 0 to save data instead.
@@ -92,7 +93,7 @@ def save(self, data: Union[torch.Tensor, np.ndarray], meta_data: Optional[Dict]
             data = data.detach().cpu().numpy()
         self._cache_dict[save_key] = np.asarray(data, dtype=float)
 
-    def save_batch(self, batch_data: Union[torch.Tensor, np.ndarray], meta_data: Optional[Dict] = None) -> None:
+    def save_batch(self, batch_data: DataObjects.Images, meta_data: Optional[Dict] = None) -> None:
         """Save a batch of data into the cache dictionary.
 
         Args:

monai/data/nifti_saver.py

@@ -19,6 +19,7 @@
 from monai.data.utils import create_file_basename
 from monai.utils import GridSampleMode, GridSamplePadMode
 from monai.utils import ImageMetaKey as Key
+from monai.utils.enums import DataObjects
 
 
 class NiftiSaver:
@@ -104,7 +105,7 @@ def __init__(
         self.separate_folder = separate_folder
         self.print_log = print_log
 
-    def save(self, data: Union[torch.Tensor, np.ndarray], meta_data: Optional[Dict] = None) -> None:
+    def save(self, data: DataObjects.Images, meta_data: Optional[Dict] = None) -> None:
         """
         Save data into a Nifti file.
         The meta_data could optionally have the following keys:
@@ -175,7 +176,7 @@ def save(self, data: Union[torch.Tensor, np.ndarray], meta_data: Optional[Dict]
         if self.print_log:
             print(f"file written: {path}.")
 
-    def save_batch(self, batch_data: Union[torch.Tensor, np.ndarray], meta_data: Optional[Dict] = None) -> None:
+    def save_batch(self, batch_data: DataObjects.Images, meta_data: Optional[Dict] = None) -> None:
         """
         Save a batch of data into Nifti format files.
 

monai/data/nifti_writer.py

@@ -18,12 +18,14 @@
 from monai.data.utils import compute_shape_offset, to_affine_nd
 from monai.networks.layers import AffineTransform
 from monai.utils import GridSampleMode, GridSamplePadMode, optional_import
+from monai.utils.enums import DataObjects
+from monai.utils.misc import convert_data_type
 
 nib, _ = optional_import("nibabel")
 
 
 def write_nifti(
-    data: np.ndarray,
+    data: DataObjects.Images,
     file_name: str,
     affine: Optional[np.ndarray] = None,
     target_affine: Optional[np.ndarray] = None,
@@ -36,7 +38,7 @@ def write_nifti(
     output_dtype: DtypeLike = np.float32,
 ) -> None:
     """
-    Write numpy data into NIfTI files to disk.  This function converts data
+    Write numpy or torch data into NIfTI files to disk.  This function converts data
     into the coordinate system defined by `target_affine` when `target_affine`
     is specified.
 
@@ -96,33 +98,39 @@ def write_nifti(
             If None, use the data type of input data.
         output_dtype: data type for saving data. Defaults to ``np.float32``.
     """
-    if not isinstance(data, np.ndarray):
-        raise AssertionError("input data must be numpy array.")
-    dtype = dtype or data.dtype
-    sr = min(data.ndim, 3)
+    if not isinstance(data, (np.ndarray, torch.Tensor)):
+        raise AssertionError("input data must be numpy array or torch.Tensor.")
+    # if torch, convert to numpy
+    data_np: np.ndarray
+    data_np, *_ = convert_data_type(data, np.ndarray)  # type: ignore
+    if target_affine is not None:
+        target_affine, *_ = convert_data_type(target_affine, np.ndarray)  # type: ignore
+
+    dtype = dtype or data_np.dtype
+    sr = min(data_np.ndim, 3)
     if affine is None:
         affine = np.eye(4, dtype=np.float64)
-    affine = to_affine_nd(sr, affine)
+    affine = to_affine_nd(sr, affine)  # type: ignore
 
     if target_affine is None:
         target_affine = affine
-    target_affine = to_affine_nd(sr, target_affine)
+    target_affine = to_affine_nd(sr, target_affine)  # type: ignore
 
     if np.allclose(affine, target_affine, atol=1e-3):
         # no affine changes, save (data, affine)
-        results_img = nib.Nifti1Image(data.astype(output_dtype), to_affine_nd(3, target_affine))
+        results_img = nib.Nifti1Image(data_np.astype(output_dtype), to_affine_nd(3, target_affine))
         nib.save(results_img, file_name)
         return
 
     # resolve orientation
     start_ornt = nib.orientations.io_orientation(affine)
     target_ornt = nib.orientations.io_orientation(target_affine)
     ornt_transform = nib.orientations.ornt_transform(start_ornt, target_ornt)
-    data_shape = data.shape
-    data = nib.orientations.apply_orientation(data, ornt_transform)
+    data_shape = data_np.shape
+    data_np = nib.orientations.apply_orientation(data_np, ornt_transform)
     _affine = affine @ nib.orientations.inv_ornt_aff(ornt_transform, data_shape)
     if np.allclose(_affine, target_affine, atol=1e-3) or not resample:
-        results_img = nib.Nifti1Image(data.astype(output_dtype), to_affine_nd(3, _affine))
+        results_img = nib.Nifti1Image(data_np.astype(output_dtype), to_affine_nd(3, _affine))
         nib.save(results_img, file_name)
         return
 
@@ -132,13 +140,13 @@ def write_nifti(
     )
     transform = np.linalg.inv(_affine) @ target_affine
     if output_spatial_shape is None:
-        output_spatial_shape, _ = compute_shape_offset(data.shape, _affine, target_affine)
+        output_spatial_shape, _ = compute_shape_offset(data_np.shape, _affine, target_affine)
     output_spatial_shape_ = list(output_spatial_shape) if output_spatial_shape is not None else []
-    if data.ndim > 3:  # multi channel, resampling each channel
+    if data_np.ndim > 3:  # multi channel, resampling each channel
         while len(output_spatial_shape_) < 3:
             output_spatial_shape_ = output_spatial_shape_ + [1]
-        spatial_shape, channel_shape = data.shape[:3], data.shape[3:]
-        data_np = data.reshape(list(spatial_shape) + [-1])
+        spatial_shape, channel_shape = data_np.shape[:3], data_np.shape[3:]
+        data_np = data_np.reshape(list(spatial_shape) + [-1])
         data_np = np.moveaxis(data_np, -1, 0)  # channel first for pytorch
         data_torch = affine_xform(
             torch.as_tensor(np.ascontiguousarray(data_np).astype(dtype)).unsqueeze(0),
@@ -149,12 +157,12 @@ def write_nifti(
         data_np = np.moveaxis(data_np, 0, -1)  # channel last for nifti
         data_np = data_np.reshape(list(data_np.shape[:3]) + list(channel_shape))
     else:  # single channel image, need to expand to have batch and channel
-        while len(output_spatial_shape_) < len(data.shape):
+        while len(output_spatial_shape_) < len(data_np.shape):
             output_spatial_shape_ = output_spatial_shape_ + [1]
         data_torch = affine_xform(
-            torch.as_tensor(np.ascontiguousarray(data).astype(dtype)[None, None]),
+            torch.as_tensor(np.ascontiguousarray(data_np).astype(dtype)[None, None]),
             torch.as_tensor(np.ascontiguousarray(transform).astype(dtype)),
-            spatial_size=output_spatial_shape_[: len(data.shape)],
+            spatial_size=output_spatial_shape_[: len(data_np.shape)],
         )
         data_np = data_torch.squeeze(0).squeeze(0).detach().cpu().numpy()
 

monai/data/png_saver.py

@@ -18,6 +18,7 @@
 from monai.data.utils import create_file_basename
 from monai.utils import ImageMetaKey as Key
 from monai.utils import InterpolateMode, look_up_option
+from monai.utils.enums import DataObjects
 
 
 class PNGSaver:
@@ -82,7 +83,7 @@ def __init__(
 
         self._data_index = 0
 
-    def save(self, data: Union[torch.Tensor, np.ndarray], meta_data: Optional[Dict] = None) -> None:
+    def save(self, data: DataObjects.Images, meta_data: Optional[Dict] = None) -> None:
         """
         Save data into a png file.
         The meta_data could optionally have the following keys:
@@ -143,7 +144,7 @@ def save(self, data: Union[torch.Tensor, np.ndarray], meta_data: Optional[Dict]
         if self.print_log:
             print(f"file written: {path}.")
 
-    def save_batch(self, batch_data: Union[torch.Tensor, np.ndarray], meta_data: Optional[Dict] = None) -> None:
+    def save_batch(self, batch_data: DataObjects.Images, meta_data: Optional[Dict] = None) -> None:
         """Save a batch of data into png format files.
 
         Args: