Repository: Project-MONAI/MONAI

Issue:
refactor event handlers for flexibility
**Is your feature request related to a problem? Please describe.**
Currently, we implemented many logics in the event handlers directly, it's not flexible to use.

**Describe the solution you'd like**
Split handlers code structure to 2 levels then we can call the same functions even without ignite.
For example:
(1) save data as an image to TensorBoard.
(2) save segmentation model output as Nifti format image file.
(3) save classification model output as CSV file.


Discussion:
Nic-Ma: @wyli , when this PR merged, I will submit another PR to implement all the 8 existing examples(3D segmentation training/evaluation, 3D classification training/evaluation) without ignite.
Thanks.
Nic-Ma: Hi @wyli ,

Thanks for your review.
I updated the code, could you please help review it again?
Thanks.

codecov-io: # [Codecov](https://codecov.io/gh/Project-MONAI/MONAI/pull/202?src=pr&el=h1) Report
> Merging [#202](https://codecov.io/gh/Project-MONAI/MONAI/pull/202?src=pr&el=desc) into [master](https://codecov.io/gh/Project-MONAI/MONAI/commit/c881c2dd7b94ffa7e72e9f8063570d10be9f9b2d&el=desc) will **increase** coverage by `1.00%`.
> The diff coverage is `93.48%`.

[![Impacted file tree graph](https://codecov.io/gh/Project-MONAI/MONAI/pull/202/graphs/tree.svg?width=650&height=150&src=pr&token=6FTC7U1JJ4)](https://codecov.io/gh/Project-MONAI/MONAI/pull/202?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master     #202      +/-   ##
==========================================
+ Coverage   85.41%   86.41%   +1.00%     
==========================================
  Files         119      125       +6     
  Lines        4855     4970     +115     
==========================================
+ Hits         4147     4295     +148     
+ Misses        708      675      -33     
```


| [Impacted Files](https://codecov.io/gh/Project-MONAI/MONAI/pull/202?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [monai/transforms/transforms.py](https://codecov.io/gh/Project-MONAI/MONAI/pull/202/diff?src=pr&el=tree#diff-bW9uYWkvdHJhbnNmb3Jtcy90cmFuc2Zvcm1zLnB5) | `93.36% <ø> (ø)` | |
| [tests/integration\_sliding\_window.py](https://codecov.io/gh/Project-MONAI/MONAI/pull/202/diff?src=pr&el=tree#diff-dGVzdHMvaW50ZWdyYXRpb25fc2xpZGluZ193aW5kb3cucHk=) | `0.00% <0.00%> (ø)` | |
| [monai/eval\_savers/csv\_saver.py](https://codecov.io/gh/Project-MONAI/MONAI/pull/202/diff?src=pr&el=tree#diff-bW9uYWkvZXZhbF9zYXZlcnMvY3N2X3NhdmVyLnB5) | `88.88% <88.88%> (ø)` | |
| [monai/visualize/img2tensorboard.py](https://codecov.io/gh/Project-MONAI/MONAI/pull/202/diff?src=pr&el=tree#diff-bW9uYWkvdmlzdWFsaXplL2ltZzJ0ZW5zb3Jib2FyZC5weQ==) | `90.00% <89.65%> (-0.25%)` | :arrow_down: |
| [tests/test\_nifti\_saver.py](https://codecov.io/gh/Project-MONAI/MONAI/pull/202/diff?src=pr&el=tree#diff-dGVzdHMvdGVzdF9uaWZ0aV9zYXZlci5weQ==) | `94.44% <94.44%> (ø)` | |
| [monai/eval\_savers/nifti\_saver.py](https://codecov.io/gh/Project-MONAI/MONAI/pull/202/diff?src=pr&el=tree#diff-bW9uYWkvZXZhbF9zYXZlcnMvbmlmdGlfc2F2ZXIucHk=) | `94.73% <94.73%> (ø)` | |
| [tests/test\_handler\_segmentation\_saver.py](https://codecov.io/gh/Project-MONAI/MONAI/pull/202/diff?src=pr&el=tree#diff-dGVzdHMvdGVzdF9oYW5kbGVyX3NlZ21lbnRhdGlvbl9zYXZlci5weQ==) | `95.65% <95.65%> (ø)` | |
| [tests/test\_plot\_2d\_or\_3d\_image.py](https://codecov.io/gh/Project-MONAI/MONAI/pull/202/diff?src=pr&el=tree#diff-dGVzdHMvdGVzdF9wbG90XzJkX29yXzNkX2ltYWdlLnB5) | `95.65% <95.65%> (ø)` | |
| [tests/test\_csv\_saver.py](https://codecov.io/gh/Project-MONAI/MONAI/pull/202/diff?src=pr&el=tree#diff-dGVzdHMvdGVzdF9jc3Zfc2F2ZXIucHk=) | `96.42% <96.42%> (ø)` | |
| [monai/data/nifti\_writer.py](https://codecov.io/gh/Project-MONAI/MONAI/pull/202/diff?src=pr&el=tree#diff-bW9uYWkvZGF0YS9uaWZ0aV93cml0ZXIucHk=) | `100.00% <100.00%> (+7.14%)` | :arrow_up: |
| ... and [12 more](https://codecov.io/gh/Project-MONAI/MONAI/pull/202/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/Project-MONAI/MONAI/pull/202?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/Project-MONAI/MONAI/pull/202?src=pr&el=footer). Last update [c881c2d...ad59771](https://codecov.io/gh/Project-MONAI/MONAI/pull/202?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

wyli: instead of creating a new folder, I think classes from both `csv_saver.py` and `nifti_saver.py` should go to the monai/data folder, e.g. `class NiftiSaver` goes into `monai/data/nifti_writer.py`, `class CSVSaver` goes into a new file `monai/data/csv_writer.py`. apart from these, the PR looks good, thanks!

Patch:
diff --git a/examples/segmentation_3d/unet_evaluation_array.py b/examples/segmentation_3d/unet_evaluation_array.py
index 305a513c..60d81cf8 100644
--- a/examples/segmentation_3d/unet_evaluation_array.py
+++ b/examples/segmentation_3d/unet_evaluation_array.py
@@ -95,7 +95,7 @@ val_stats_handler.attach(evaluator)
 
 # for the arrary data format, assume the 3rd item of batch data is the meta_data
 file_saver = SegmentationSaver(
-    output_path='tempdir', output_ext='.nii.gz', output_postfix='seg', name='evaluator',
+    output_dir='tempdir', output_ext='.nii.gz', output_postfix='seg', name='evaluator',
     batch_transform=lambda x: x[2], output_transform=lambda output: predict_segmentation(output[0]))
 file_saver.attach(evaluator)
 
diff --git a/examples/segmentation_3d/unet_evaluation_dict.py b/examples/segmentation_3d/unet_evaluation_dict.py
index 5a906af1..86766fb3 100644
--- a/examples/segmentation_3d/unet_evaluation_dict.py
+++ b/examples/segmentation_3d/unet_evaluation_dict.py
@@ -98,7 +98,7 @@ val_stats_handler = StatsHandler(
 val_stats_handler.attach(evaluator)
 
 # convert the necessary metadata from batch data
-SegmentationSaver(output_path='tempdir', output_ext='.nii.gz', output_postfix='seg', name='evaluator',
+SegmentationSaver(output_dir='tempdir', output_ext='.nii.gz', output_postfix='seg', name='evaluator',
                   batch_transform=lambda batch: {'filename_or_obj': batch['img.filename_or_obj'],
                                                  'original_affine': batch['img.original_affine'],
                                                  'affine': batch['img.affine'],
diff --git a/monai/data/csv_saver.py b/monai/data/csv_saver.py
new file mode 100644
index 00000000..98251738
--- /dev/null
+++ b/monai/data/csv_saver.py
@@ -0,0 +1,88 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import csv
+import numpy as np
+import torch
+from collections import OrderedDict
+
+
+class CSVSaver:
+    """
+    save the data in a dictionary format cache, and write to a CSV file finally.
+    Typically, the data can be classification predictions, call `save` for single data
+    or call `save_batch` to save a batch of data together, and call `finalize` to write
+    the cached data into CSV file. If no meta data provided, use index from 0 to save data.
+    """
+
+    def __init__(self, output_dir='./', filename='predictions.csv', overwrite=True):
+        """
+        Args:
+            output_dir (str): output CSV file directory.
+            filename (str): name of the saved CSV file name.
+            overwrite (bool): whether to overwriting existing CSV file content. If we are not overwriting,
+                then we check if the results have been previously saved, and load them to the prediction_dict.
+
+        """
+        self.output_dir = output_dir
+        self._cache_dict = OrderedDict()
+        assert isinstance(filename, str) and filename[-4:] == '.csv', 'filename must be a string with CSV format.'
+        self._filepath = os.path.join(output_dir, filename)
+        self.overwrite = overwrite
+        self._data_index = 0
+
+    def finalize(self):
+        """
+        Writes the cahced dict to a csv
+
+        """
+        if not self.overwrite and os.path.exists(self._filepath):
+            with open(self._filepath, 'r') as f:
+                reader = csv.reader(f)
+                for row in reader:
+                    self._cache_dict[row[0]] = np.array(row[1:]).astype(np.float32)
+
+        if not os.path.exists(self.output_dir):
+            os.makedirs(self.output_dir)
+        with open(self._filepath, 'w') as f:
+            for k, v in self._cache_dict.items():
+                f.write(k)
+                for result in v.flatten():
+                    f.write("," + str(result))
+                f.write("\n")
+
+    def save(self, data, meta_data=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 detault index from 0 to save data instead.
+
+        args:
+            data (Tensor or ndarray): target data content that save into cache.
+            meta_data (dict): the meta data information corresponding to the data.
+
+        """
+        save_key = meta_data['filename_or_obj'] if meta_data else str(self._data_index)
+        self._data_index += 1
+        if torch.is_tensor(data):
+            data = data.detach().cpu().numpy()
+        self._cache_dict[save_key] = data.astype(np.float32)
+
+    def save_batch(self, batch_data, meta_data=None):
+        """Save a batch of data into the cache dictionary.
+
+        args:
+            batch_data (Tensor or ndarray): target batch data content that save into cache.
+            meta_data (dict): every key-value in the meta_data is corresponding to 1 batch of data.
+
+        """
+        for i, data in enumerate(batch_data):  # save a batch of files
+            self.save(data, {k: meta_data[k][i] for k in meta_data} if meta_data else None)
diff --git a/monai/data/nifti_saver.py b/monai/data/nifti_saver.py
new file mode 100644
index 00000000..924f9983
--- /dev/null
+++ b/monai/data/nifti_saver.py
@@ -0,0 +1,112 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import numpy as np
+import torch
+from monai.data.nifti_writer import write_nifti
+
+
+class NiftiSaver:
+    """
+    Save the data as nifti file, it can support single data content or a batch of data.
+    Typically, the data can be segmentation predictions, call `save` for single data
+    or call `save_batch` to save a batch of data together. If no meta data provided,
+    use index from 0 as the filename prefix.
+    """
+
+    def __init__(self, output_dir='./', output_postfix='seg', output_ext='.nii.gz', dtype=None):
+        """
+        Args:
+            output_dir (str): output image directory.
+            output_postfix (str): a string appended to all output file names.
+            output_ext (str): output file extension name.
+            dtype (np.dtype, optional): convert the image data to save to this data type.
+                If None, keep the original type of data.
+
+        """
+        self.output_dir = output_dir
+        self.output_postfix = output_postfix
+        self.output_ext = output_ext
+        self.dtype = dtype
+        self._data_index = 0
+
+    @staticmethod
+    def _create_file_basename(postfix, input_file_name, folder_path, data_root_dir=""):
+        """
+        Utility function to create the path to the output file based on the input
+        filename (extension is added by lib level writer before writing the file)
+
+        Args:
+            postfix (str): output name's postfix
+            input_file_name (str): path to the input image file
+            folder_path (str): path for the output file
+            data_root_dir (str): if not empty, it specifies the beginning parts of the input file's
+          absolute path. This is used to compute `input_file_rel_path`, the relative path to the file from
+          `data_root_dir` to preserve folder structure when saving in case there are files in different
+          folders with the same file names.
+        """
+
+        # get the filename and directory
+        filedir, filename = os.path.split(input_file_name)
+
+        # jettison the extension to have just filename
+        filename, ext = os.path.splitext(filename)
+        while ext != "":
+            filename, ext = os.path.splitext(filename)
+
+        # use data_root_dir to find relative path to file
+        filedir_rel_path = ""
+        if data_root_dir:
+            filedir_rel_path = os.path.relpath(filedir, data_root_dir)
+
+        # sub-folder path will be original name without the extension
+        subfolder_path = os.path.join(folder_path, filedir_rel_path, filename)
+        if not os.path.exists(subfolder_path):
+            os.makedirs(subfolder_path)
+
+        # add the sub-folder plus the postfix name to become the file basename in the output path
+        return os.path.join(subfolder_path, filename + "_" + postfix)
+
+    def save(self, data, meta_data=None):
+        """Save data into Nifti format file. The metadata should have the following keys:
+            - ``'filename_or_obj'`` -- for output file name creation, corresponding to filename or object
+            - ``'original_affine'`` (optional) for data orientation handling
+            - ``'affine'`` (optional) for data output affine.
+        If meta_data is None, use the detault index from 0 to save data instead.
+
+        args:
+            data (Tensor or ndarray): target data content that save into Nifti format file.
+                Assuming the data shape is `NCHW[D]`.
+            meta_data (dict): the meta data information corresponding to the data.
+        """
+        filename = meta_data['filename_or_obj'] if meta_data else str(self._data_index)
+        self._data_index += 1
+        original_affine = meta_data.get('original_affine', None) if meta_data else None
+        affine = meta_data.get('affine', None) if meta_data else None
+
+        if torch.is_tensor(data):
+            data = data.detach().cpu().numpy()
+        filename = self._create_file_basename(self.output_postfix, filename, self.output_dir)
+        filename = '{}{}'.format(filename, self.output_ext)
+        # change data to "channel last" format and write to nifti format file
+        data = np.moveaxis(data, 0, -1)
+        write_nifti(data, affine, filename, original_affine, dtype=self.dtype or data.dtype)
+
+    def save_batch(self, batch_data, meta_data=None):
+        """Save a batch of data into Nifti format files.
+
+        args:
+            batch_data (Tensor or ndarray): target batch data content that save into Nifti format files.
+            meta_data (dict): every key-value in the meta_data is corresponding to 1 batch of data.
+        """
+        for i, data in enumerate(batch_data):  # save a batch of files
+            self.save(data, {k: meta_data[k][i] for k in meta_data} if meta_data else None)
diff --git a/monai/data/nifti_writer.py b/monai/data/nifti_writer.py
index 44e454d9..ac1d5b48 100644
--- a/monai/data/nifti_writer.py
+++ b/monai/data/nifti_writer.py
@@ -13,7 +13,7 @@ import numpy as np
 import nibabel as nib
 
 
-def write_nifti(data, affine, file_name, target_affine=None, dtype="float32"):
+def write_nifti(data, affine, file_name, target_affine=None, dtype=np.float32):
     """Write numpy data into nifti files to disk.
 
     Args:
diff --git a/monai/handlers/classification_saver.py b/monai/handlers/classification_saver.py
index 501dce81..50746187 100644
--- a/monai/handlers/classification_saver.py
+++ b/monai/handlers/classification_saver.py
@@ -9,12 +9,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-import csv
-import numpy as np
-import torch
 from ignite.engine import Events
 import logging
+from monai.data.csv_saver import CSVSaver
 
 
 class ClassificationSaver:
@@ -22,11 +19,12 @@ class ClassificationSaver:
     Event handler triggered on completing every iteration to save the classification predictions as CSV file.
     """
 
-    def __init__(self, output_dir='./', overwrite=True,
+    def __init__(self, output_dir='./', filename='predictions.csv', overwrite=True,
                  batch_transform=lambda x: x, output_transform=lambda x: x, name=None):
         """
         Args:
             output_dir (str): output CSV file directory.
+            filename (str): name of the saved CSV file name.
             overwrite (bool): whether to overwriting existing CSV file content. If we are not overwriting,
                 then we check if the results have been previously saved, and load them to the prediction_dict.
             batch_transform (Callable): a callable that is used to transform the
@@ -38,10 +36,7 @@ class ClassificationSaver:
             name (str): identifier of logging.logger to use, defaulting to `engine.logger`.
 
         """
-        self.output_dir = output_dir
-        self._prediction_dict = {}
-        self._preds_filepath = os.path.join(output_dir, 'predictions.csv')
-        self.overwrite = overwrite
+        self.saver = CSVSaver(output_dir, filename, overwrite)
         self.batch_transform = batch_transform
         self.output_transform = output_transform
 
@@ -52,44 +47,14 @@ class ClassificationSaver:
             self.logger = engine.logger
         if not engine.has_event_handler(self, Events.ITERATION_COMPLETED):
             engine.add_event_handler(Events.ITERATION_COMPLETED, self)
-        if not engine.has_event_handler(self.finalize, Events.COMPLETED):
-            engine.add_event_handler(Events.COMPLETED, self.finalize)
-
-    def finalize(self, _engine=None):
-        """
-        Writes the prediction dict to a csv
-
-        """
-        if not self.overwrite and os.path.exists(self._preds_filepath):
-            with open(self._preds_filepath, 'r') as f:
-                reader = csv.reader(f)
-                for row in reader:
-                    self._prediction_dict[row[0]] = np.array(row[1:]).astype(np.float32)
-
-        if not os.path.exists(self.output_dir):
-            os.makedirs(self.output_dir)
-        with open(self._preds_filepath, 'w') as f:
-            for k, v in sorted(self._prediction_dict.items()):
-                f.write(k)
-                for result in v.flatten():
-                    f.write("," + str(result))
-                f.write("\n")
-        self.logger.info('saved classification predictions into: {}'.format(self._preds_filepath))
+        if not engine.has_event_handler(self.saver.finalize, Events.COMPLETED):
+            engine.add_event_handler(Events.COMPLETED, lambda engine: self.saver.finalize())
 
     def __call__(self, engine):
         """
-        This method assumes self.batch_transform will extract Metadata from the input batch.
-        Metadata should have the following keys:
-
-            - ``'filename_or_obj'`` -- save the prediction corresponding to file name.
+        This method assumes self.batch_transform will extract metadata from the input batch.
 
         """
         meta_data = self.batch_transform(engine.state.batch)
-        filenames = meta_data['filename_or_obj']
-
         engine_output = self.output_transform(engine.state.output)
-        for batch_id, filename in enumerate(filenames):  # save a batch of files
-            output = engine_output[batch_id]
-            if isinstance(output, torch.Tensor):
-                output = output.detach().cpu().numpy()
-            self._prediction_dict[filename] = output.astype(np.float32)
+        self.saver.save_batch(engine_output, meta_data)
diff --git a/monai/handlers/segmentation_saver.py b/monai/handlers/segmentation_saver.py
index 98eb972d..36256596 100644
--- a/monai/handlers/segmentation_saver.py
+++ b/monai/handlers/segmentation_saver.py
@@ -9,12 +9,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-import numpy as np
-import torch
 from ignite.engine import Events
 import logging
-from monai.data.nifti_writer import write_nifti
+from monai.data.nifti_saver import NiftiSaver
 
 
 class SegmentationSaver:
@@ -22,14 +19,15 @@ class SegmentationSaver:
     Event handler triggered on completing every iteration to save the segmentation predictions as nifti files.
     """
 
-    def __init__(self, output_path='./', dtype='float32', output_postfix='seg', output_ext='.nii.gz',
+    def __init__(self, output_dir='./', output_postfix='seg', output_ext='.nii.gz', dtype=None,
                  batch_transform=lambda x: x, output_transform=lambda x: x, name=None):
         """
         Args:
-            output_path (str): output image directory.
-            dtype (str): to convert the image to save to this datatype.
+            output_dir (str): output image directory.
             output_postfix (str): a string appended to all output file names.
             output_ext (str): output file extension name.
+            dtype (np.dtype, optional): convert the image data to save to this data type.
+                If None, keep the original type of data.
             batch_transform (Callable): a callable that is used to transform the
                 ignite.engine.batch into expected format to extract the meta_data dictionary.
             output_transform (Callable): a callable that is used to transform the
@@ -37,11 +35,9 @@ class SegmentationSaver:
                 The first dimension of this transform's output will be treated as the
                 batch dimension. Each item in the batch will be saved individually.
             name (str): identifier of logging.logger to use, defaulting to `engine.logger`.
+
         """
-        self.output_path = output_path
-        self.dtype = dtype
-        self.output_postfix = output_postfix
-        self.output_ext = output_ext
+        self.saver = NiftiSaver(output_dir, output_postfix, output_ext, dtype)
         self.batch_transform = batch_transform
         self.output_transform = output_transform
 
@@ -53,69 +49,13 @@ class SegmentationSaver:
         if not engine.has_event_handler(self, Events.ITERATION_COMPLETED):
             engine.add_event_handler(Events.ITERATION_COMPLETED, self)
 
-    @staticmethod
-    def _create_file_basename(postfix, input_file_name, folder_path, data_root_dir=""):
-        """
-        Utility function to create the path to the output file based on the input
-        filename (extension is added by lib level writer before writing the file)
-
-        Args:
-            postfix (str): output name's postfix
-            input_file_name (str): path to the input image file
-            folder_path (str): path for the output file
-            data_root_dir (str): if not empty, it specifies the beginning parts of the input file's
-          absolute path. This is used to compute `input_file_rel_path`, the relative path to the file from
-          `data_root_dir` to preserve folder structure when saving in case there are files in different
-          folders with the same file names.
-        """
-
-        # get the filename and directory
-        filedir, filename = os.path.split(input_file_name)
-
-        # jettison the extension to have just filename
-        filename, ext = os.path.splitext(filename)
-        while ext != "":
-            filename, ext = os.path.splitext(filename)
-
-        # use data_root_dir to find relative path to file
-        filedir_rel_path = ""
-        if data_root_dir:
-            filedir_rel_path = os.path.relpath(filedir, data_root_dir)
-
-        # sub-folder path will be original name without the extension
-        subfolder_path = os.path.join(folder_path, filedir_rel_path, filename)
-        if not os.path.exists(subfolder_path):
-            os.makedirs(subfolder_path)
-
-        # add the sub-folder plus the postfix name to become the file basename in the output path
-        return os.path.join(subfolder_path, filename + "_" + postfix)
-
     def __call__(self, engine):
         """
-        This method assumes self.batch_transform will extract Metadata from the input batch.
-        Metadata should have the following keys:
-
-            - ``'filename_or_obj'`` -- for output file name creation
-            - ``'original_affine'`` (optional) for data orientation handling
-            - ``'affine'`` (optional) for data output affine.
-
+        This method assumes self.batch_transform will extract metadata from the input batch.
         output file datatype is determined from ``engine.state.output.dtype``.
+
         """
         meta_data = self.batch_transform(engine.state.batch)
-        filenames = meta_data['filename_or_obj']
-        original_affine = meta_data.get('original_affine', None)
-        affine = meta_data.get('affine', None)
-
         engine_output = self.output_transform(engine.state.output)
-        for batch_id, filename in enumerate(filenames):  # save a batch of files
-            seg_output = engine_output[batch_id]
-            affine_ = affine[batch_id]
-            original_affine_ = original_affine[batch_id]
-            if isinstance(seg_output, torch.Tensor):
-                seg_output = seg_output.detach().cpu().numpy()
-            output_filename = self._create_file_basename(self.output_postfix, filename, self.output_path)
-            output_filename = '{}{}'.format(output_filename, self.output_ext)
-            # change output to "channel last" format and write to nifti format file
-            to_save = np.moveaxis(seg_output, 0, -1)
-            write_nifti(to_save, affine_, output_filename, original_affine_, dtype=seg_output.dtype)
-            self.logger.info('saved: {}'.format(output_filename))
+        self.saver.save_batch(engine_output, meta_data)
+        self.logger.info('saved all the model outputs as nifti files.')
diff --git a/monai/handlers/tensorboard_handlers.py b/monai/handlers/tensorboard_handlers.py
index fb5eb34e..3a3e7a06 100644
--- a/monai/handlers/tensorboard_handlers.py
+++ b/monai/handlers/tensorboard_handlers.py
@@ -14,9 +14,8 @@ import warnings
 import torch
 from torch.utils.tensorboard import SummaryWriter
 from ignite.engine import Engine, Events
-from monai.visualize import img2tensorboard
+from monai.visualize.img2tensorboard import plot_2d_or_3d_image
 from monai.utils.misc import is_scalar
-from monai.transforms.utils import rescale_array
 
 DEFAULT_TAG = 'Loss'
 
@@ -176,6 +175,7 @@ class TensorBoardImageHandler(object):
                  batch_transform=lambda x: x,
                  output_transform=lambda x: x,
                  global_iter_transform=lambda x: x,
+                 index=0,
                  max_channels=1,
                  max_frames=64):
         """
@@ -188,6 +188,7 @@ class TensorBoardImageHandler(object):
                 ``ignite.engine.output`` into expected format to extract several output data.
             global_iter_transform (Callable): a callable that is used to customize global step number for TensorBoard.
                 For example, in evaluation, the evaluator engine needs to know current epoch from trainer.
+            index (int): plot which element in a data batch, default is the first element.
             max_channels (int): number of channels to plot.
             max_frames (int): number of frames for 2D-t plot.
         """
@@ -195,7 +196,7 @@ class TensorBoardImageHandler(object):
         self.batch_transform = batch_transform
         self.output_transform = output_transform
         self.global_iter_transform = global_iter_transform
-
+        self.index = index
         self.max_frames = max_frames
         self.max_channels = max_channels
 
@@ -208,7 +209,8 @@ class TensorBoardImageHandler(object):
         if show_images is not None:
             if not isinstance(show_images, np.ndarray):
                 raise ValueError('output_transform(engine.state.output)[0] must be an ndarray or tensor.')
-            self._add_2_or_3_d(show_images, step, 'input_0')
+            plot_2d_or_3d_image(show_images, step, self._writer, self.index,
+                                self.max_channels, self.max_frames, 'input_0')
 
         show_labels = self.batch_transform(engine.state.batch)[1]
         if torch.is_tensor(show_labels):
@@ -216,7 +218,8 @@ class TensorBoardImageHandler(object):
         if show_labels is not None:
             if not isinstance(show_labels, np.ndarray):
                 raise ValueError('batch_transform(engine.state.batch)[1] must be an ndarray or tensor.')
-            self._add_2_or_3_d(show_labels, step, 'input_1')
+            plot_2d_or_3d_image(show_labels, step, self._writer, self.index,
+                                self.max_channels, self.max_frames, 'input_1')
 
         show_outputs = self.output_transform(engine.state.output)
         if torch.is_tensor(show_outputs):
@@ -224,35 +227,7 @@ class TensorBoardImageHandler(object):
         if show_outputs is not None:
             if not isinstance(show_outputs, np.ndarray):
                 raise ValueError('output_transform(engine.state.output) must be an ndarray or tensor.')
-            self._add_2_or_3_d(show_outputs, step, 'output')
+            plot_2d_or_3d_image(show_outputs, step, self._writer, self.index,
+                                self.max_channels, self.max_frames, 'output')
 
         self._writer.flush()
-
-    def _add_2_or_3_d(self, data, step, tag='output'):
-        # for i, d in enumerate(data):  # go through a batch of images
-        d = data[0]  # show the first element in a batch
-
-        if d.ndim == 2:
-            d = rescale_array(d, 0, 1)
-            dataformats = 'HW'
-            self._writer.add_image('{}_{}'.format(tag, dataformats), d, step, dataformats=dataformats)
-            return
-
-        if d.ndim == 3:
-            if d.shape[0] == 3 and self.max_channels == 3:  # RGB
-                dataformats = 'CHW'
-                self._writer.add_image('{}_{}'.format(tag, dataformats), d, step, dataformats=dataformats)
-                return
-            for j, d2 in enumerate(d[:self.max_channels]):
-                d2 = rescale_array(d2, 0, 1)
-                dataformats = 'HW'
-                self._writer.add_image('{}_{}_{}'.format(tag, dataformats, j), d2, step, dataformats=dataformats)
-                return
-
-        if d.ndim >= 4:
-            spatial = d.shape[-3:]
-            for j, d3 in enumerate(d.reshape([-1] + list(spatial))[:self.max_channels]):
-                d3 = rescale_array(d3, 0, 255)
-                img2tensorboard.add_animated_gif(
-                    self._writer, '{}_HWD_{}'.format(tag, j), d3[None], self.max_frames, 1.0, step)
-            return
diff --git a/monai/transforms/transforms.py b/monai/transforms/transforms.py
index f0ee4a93..e3f76771 100644
--- a/monai/transforms/transforms.py
+++ b/monai/transforms/transforms.py
@@ -520,7 +520,6 @@ class NormalizeIntensity:
     """Normalize input based on provided args, using calculated mean and std if not provided
     (shape of subtrahend and divisor must match. if 0, entire volume uses same subtrahend and
     divisor, otherwise the shape can have dimension 1 for channels).
-    Current implementation can only support 'channel_last' format data.
 
     Args:
         subtrahend (ndarray): the amount to subtract by (usually the mean)
diff --git a/monai/visualize/img2tensorboard.py b/monai/visualize/img2tensorboard.py
index 7498ce1c..8d833223 100644
--- a/monai/visualize/img2tensorboard.py
+++ b/monai/visualize/img2tensorboard.py
@@ -9,15 +9,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import torch
 import numpy as np
 import PIL
 from PIL.GifImagePlugin import Image as GifImage
+from torch.utils.tensorboard import SummaryWriter
 from tensorboard.compat.proto import summary_pb2
+from monai.transforms.utils import rescale_array
 
 
 def _image3_animated_gif(imp, scale_factor=1):
-    """
-    Function to actually create the animated gif.
+    """Function to actually create the animated gif.
     Args:
         imp: tuple of tag and a list of image tensors
         scale_factor: amount to multiply values by. if the image data is between 0 and 1, using 255 for this value will
@@ -52,8 +54,7 @@ def make_animated_gif_summary(tag,
                               image_axes=(1, 2),
                               other_indices=None,
                               scale_factor=1):
-    """
-    Creates an animated gif out of an image tensor and returns Summary.
+    """Creates an animated gif out of an image tensor and returns Summary.
 
     Args:
         tag: Data identifier
@@ -92,8 +93,7 @@ def make_animated_gif_summary(tag,
 
 
 def add_animated_gif(writer, tag, image_tensor, max_out, scale_factor, global_step=None):
-    """
-    Creates an animated gif out of an image tensor and writes it with SummaryWriter.
+    """Creates an animated gif out of an image tensor and writes it with SummaryWriter.
 
     Args:
         writer: Tensorboard SummaryWriter to write to
@@ -110,8 +110,7 @@ def add_animated_gif(writer, tag, image_tensor, max_out, scale_factor, global_st
 
 
 def add_animated_gif_no_channels(writer, tag, image_tensor, max_out, scale_factor, global_step=None):
-    """
-    Creates an animated gif out of an image tensor and writes it with SummaryWriter.
+    """Creates an animated gif out of an image tensor and writes it with SummaryWriter.
 
     Args:
         writer: Tensorboard SummaryWriter to write to
@@ -126,3 +125,49 @@ def add_animated_gif_no_channels(writer, tag, image_tensor, max_out, scale_facto
                                                                     max_out=max_out, animation_axes=[1],
                                                                     image_axes=[2, 3], scale_factor=scale_factor),
                                           global_step)
+
+
+def plot_2d_or_3d_image(data, step, writer, index=0, max_channels=1, max_frames=64, tag='output'):
+    """Plot 2D or 3D image on the TensorBoard, 3D image will be converted to GIF image.
+
+    Note:
+        Plot 3D or 2D image(with more than 3 channels) as separate images.
+
+    Args:
+        data (Tensor or ndarray): target data to be plotted as image on the TensorBoard.
+            The data is expected to have 'NCHW[D]' dimensions, and only plot the first in the batch.
+        step (int): current step to plot in a chart.
+        writer (SummaryWriter): specify TensorBoard SummaryWriter to plot the image.
+        index (int): plot which element in the input data batch, default is the first element.
+        max_channels (int): number of channels to plot.
+        max_frames (int): number of frames for 2D-t plot.
+        tag (str): tag of the plotted image on TensorBoard.
+    """
+    assert isinstance(writer, SummaryWriter) is True, 'must provide a TensorBoard SummaryWriter.'
+    d = data[index]
+    if torch.is_tensor(d):
+        d = d.detach().cpu().numpy()
+
+    if d.ndim == 2:
+        d = rescale_array(d, 0, 1)
+        dataformats = 'HW'
+        writer.add_image('{}_{}'.format(tag, dataformats), d, step, dataformats=dataformats)
+        return
+
+    if d.ndim == 3:
+        if d.shape[0] == 3 and max_channels == 3:  # RGB
+            dataformats = 'CHW'
+            writer.add_image('{}_{}'.format(tag, dataformats), d, step, dataformats=dataformats)
+            return
+        for j, d2 in enumerate(d[:max_channels]):
+            d2 = rescale_array(d2, 0, 1)
+            dataformats = 'HW'
+            writer.add_image('{}_{}_{}'.format(tag, dataformats, j), d2, step, dataformats=dataformats)
+            return
+
+    if d.ndim >= 4:
+        spatial = d.shape[-3:]
+        for j, d3 in enumerate(d.reshape([-1] + list(spatial))[:max_channels]):
+            d3 = rescale_array(d3, 0, 255)
+            add_animated_gif(writer, '{}_HWD_{}'.format(tag, j), d3[None], max_frames, 1.0, step)
+        return


Test patch:
diff --git a/tests/integration_sliding_window.py b/tests/integration_sliding_window.py
index 91fd2994..4cc1e6b5 100644
--- a/tests/integration_sliding_window.py
+++ b/tests/integration_sliding_window.py
@@ -58,7 +58,7 @@ def run_test(batch_size=2, device=torch.device("cuda:0")):
     infer_engine = Engine(_sliding_window_processor)
 
     with tempfile.TemporaryDirectory() as temp_dir:
-        SegmentationSaver(output_path=temp_dir, output_ext='.nii.gz', output_postfix='seg',
+        SegmentationSaver(output_dir=temp_dir, output_ext='.nii.gz', output_postfix='seg',
                           batch_transform=lambda x: x[2]).attach(infer_engine)
 
         infer_engine.run(loader)
diff --git a/tests/test_csv_saver.py b/tests/test_csv_saver.py
new file mode 100644
index 00000000..e16bc4f0
--- /dev/null
+++ b/tests/test_csv_saver.py
@@ -0,0 +1,48 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import csv
+import shutil
+import unittest
+import numpy as np
+import torch
+
+
+from monai.data.csv_saver import CSVSaver
+
+
+class TestCSVSaver(unittest.TestCase):
+
+    def test_saved_content(self):
+        default_dir = os.path.join('.', 'tempdir')
+        shutil.rmtree(default_dir, ignore_errors=True)
+
+        saver = CSVSaver(output_dir=default_dir, filename='predictions.csv')
+
+        meta_data = {'filename_or_obj': ['testfile' + str(i) for i in range(8)]}
+        saver.save_batch(torch.zeros(8), meta_data)
+        saver.finalize()
+        filepath = os.path.join(default_dir, 'predictions.csv')
+        self.assertTrue(os.path.exists(filepath))
+        with open(filepath, 'r') as f:
+            reader = csv.reader(f)
+            i = 0
+            for row in reader:
+                self.assertEqual(row[0], 'testfile' + str(i))
+                self.assertEqual(np.array(row[1:]).astype(np.float32), 0.0)
+                i += 1
+            self.assertEqual(i, 8)
+        shutil.rmtree(default_dir)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/test_handler_classification_saver.py b/tests/test_handler_classification_saver.py
index 3eea9d86..f599c0b6 100644
--- a/tests/test_handler_classification_saver.py
+++ b/tests/test_handler_classification_saver.py
@@ -33,11 +33,11 @@ class TestHandlerClassificationSaver(unittest.TestCase):
         engine = Engine(_train_func)
 
         # set up testing handler
-        saver = ClassificationSaver(output_dir=default_dir)
+        saver = ClassificationSaver(output_dir=default_dir, filename='predictions.csv')
         saver.attach(engine)
 
         data = [{'filename_or_obj': ['testfile' + str(i) for i in range(8)]}]
-        engine.run(data, epoch_length=2, max_epochs=1)
+        engine.run(data, max_epochs=1)
         filepath = os.path.join(default_dir, 'predictions.csv')
         self.assertTrue(os.path.exists(filepath))
         with open(filepath, 'r') as f:
diff --git a/tests/test_handler_segmentation_saver.py b/tests/test_handler_segmentation_saver.py
new file mode 100644
index 00000000..6f3f2c0c
--- /dev/null
+++ b/tests/test_handler_segmentation_saver.py
@@ -0,0 +1,46 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import shutil
+import unittest
+import torch
+from ignite.engine import Engine
+
+from monai.handlers.segmentation_saver import SegmentationSaver
+
+
+class TestHandlerSegmentationSaver(unittest.TestCase):
+
+    def test_saved_content(self):
+        default_dir = os.path.join('.', 'tempdir')
+        shutil.rmtree(default_dir, ignore_errors=True)
+
+        # set up engine
+        def _train_func(engine, batch):
+            return torch.zeros(8, 1, 2, 2)
+
+        engine = Engine(_train_func)
+
+        # set up testing handler
+        saver = SegmentationSaver(output_dir=default_dir, output_postfix='seg', output_ext='.nii.gz')
+        saver.attach(engine)
+
+        data = [{'filename_or_obj': ['testfile' + str(i) for i in range(8)]}]
+        engine.run(data, max_epochs=1)
+        for i in range(8):
+            filepath = os.path.join('testfile' + str(i), 'testfile' + str(i) + '_seg.nii.gz')
+            self.assertTrue(os.path.exists(os.path.join(default_dir, filepath)))
+        shutil.rmtree(default_dir)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/test_nifti_saver.py b/tests/test_nifti_saver.py
new file mode 100644
index 00000000..472c8329
--- /dev/null
+++ b/tests/test_nifti_saver.py
@@ -0,0 +1,38 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import shutil
+import unittest
+import torch
+
+
+from monai.data.nifti_saver import NiftiSaver
+
+
+class TestNiftiSaver(unittest.TestCase):
+
+    def test_saved_content(self):
+        default_dir = os.path.join('.', 'tempdir')
+        shutil.rmtree(default_dir, ignore_errors=True)
+
+        saver = NiftiSaver(output_dir=default_dir, output_postfix='seg', output_ext='.nii.gz')
+
+        meta_data = {'filename_or_obj': ['testfile' + str(i) for i in range(8)]}
+        saver.save_batch(torch.zeros(8, 1, 2, 2), meta_data)
+        for i in range(8):
+            filepath = os.path.join('testfile' + str(i), 'testfile' + str(i) + '_seg.nii.gz')
+            self.assertTrue(os.path.exists(os.path.join(default_dir, filepath)))
+        shutil.rmtree(default_dir)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/test_plot_2d_or_3d_image.py b/tests/test_plot_2d_or_3d_image.py
new file mode 100644
index 00000000..483eb051
--- /dev/null
+++ b/tests/test_plot_2d_or_3d_image.py
@@ -0,0 +1,57 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import glob
+import os
+import shutil
+import unittest
+from torch.utils.tensorboard import SummaryWriter
+import torch
+from parameterized import parameterized
+from monai.visualize.img2tensorboard import plot_2d_or_3d_image
+
+TEST_CASE_1 = [
+    (1, 1, 10, 10)
+]
+
+TEST_CASE_2 = [
+    (1, 3, 10, 10)
+]
+
+TEST_CASE_3 = [
+    (1, 4, 10, 10)
+]
+
+TEST_CASE_4 = [
+    (1, 1, 10, 10, 10)
+]
+
+TEST_CASE_5 = [
+    (1, 3, 10, 10, 10)
+]
+
+
+class TestPlot2dOr3dImage(unittest.TestCase):
+
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5])
+    def test_tb_image_shape(self, shape):
+        default_dir = os.path.join('.', 'runs')
+        shutil.rmtree(default_dir, ignore_errors=True)
+
+        plot_2d_or_3d_image(torch.zeros(shape), 0, SummaryWriter())
+
+        self.assertTrue(os.path.exists(default_dir))
+        self.assertTrue(len(glob.glob(default_dir)) > 0)
+        shutil.rmtree(default_dir)
+
+
+if __name__ == '__main__':
+    unittest.main()