Unify all tests to use tempdir (#450)

* [DLMED] fix flake8 error

* [DLMED] revert

* [DLMED] update all tests to use tempdir

* [DLMED] fix flake8 issue

Co-authored-by: Wenqi Li <[email protected]>

Author: Nic-Ma

monai/handlers/tensorboard_handlers.py

@@ -36,6 +36,7 @@ class TensorBoardStatsHandler(object):
     def __init__(
         self,
         summary_writer=None,
+        log_dir="./runs",
         epoch_event_writer=None,
         iteration_event_writer=None,
         output_transform=lambda x: x,
@@ -46,6 +47,7 @@ def __init__(
         Args:
             summary_writer (SummaryWriter): user can specify TensorBoard SummaryWriter,
                 default to create a new writer.
+            log_dir (str): if using default SummaryWriter, write logs to this directory, default is `./runs`.
             epoch_event_writer (Callable): customized callable TensorBoard writer for epoch level.
                 must accept parameter "engine" and "summary_writer", use default event writer if None.
             iteration_event_writer (Callable): customized callable TensorBoard writer for iteration level.
@@ -59,7 +61,7 @@ def __init__(
                 when plotting epoch vs metric curves.
             tag_name (string): when iteration output is a scalar, tag_name is used to plot, defaults to ``'Loss'``.
         """
-        self._writer = SummaryWriter() if summary_writer is None else summary_writer
+        self._writer = SummaryWriter(log_dir=log_dir) if summary_writer is None else summary_writer
         self.epoch_event_writer = epoch_event_writer
         self.iteration_event_writer = iteration_event_writer
         self.output_transform = output_transform
@@ -179,6 +181,7 @@ class TensorBoardImageHandler(object):
     def __init__(
         self,
         summary_writer=None,
+        log_dir="./runs",
         batch_transform=lambda x: x,
         output_transform=lambda x: x,
         global_iter_transform=lambda x: x,
@@ -190,6 +193,7 @@ def __init__(
         Args:
             summary_writer (SummaryWriter): user can specify TensorBoard SummaryWriter,
                 default to create a new writer.
+            log_dir (str): if using default SummaryWriter, write logs to this directory, default is `./runs`.
             batch_transform (Callable): a callable that is used to transform the
                 ``ignite.engine.batch`` into expected format to extract several label data.
             output_transform (Callable): a callable that is used to transform the
@@ -200,7 +204,7 @@ def __init__(
             max_channels (int): number of channels to plot.
             max_frames (int): number of frames for 2D-t plot.
         """
-        self._writer = SummaryWriter() if summary_writer is None else summary_writer
+        self._writer = SummaryWriter(log_dir=log_dir) if summary_writer is None else summary_writer
         self.batch_transform = batch_transform
         self.output_transform = output_transform
         self.global_iter_transform = global_iter_transform

tests/test_handler_tb_image.py

@@ -13,7 +13,7 @@
 import os
 import shutil
 import unittest
-
+import tempfile
 import numpy as np
 import torch
 from ignite.engine import Engine, Events
@@ -34,8 +34,8 @@
 class TestHandlerTBImage(unittest.TestCase):
     @parameterized.expand(TEST_CASES)
     def test_tb_image_shape(self, shape):
-        default_dir = os.path.join(".", "runs")
-        shutil.rmtree(default_dir, ignore_errors=True)
+        tempdir = tempfile.mkdtemp()
+        shutil.rmtree(tempdir, ignore_errors=True)
 
         # set up engine
         def _train_func(engine, batch):
@@ -44,15 +44,15 @@ def _train_func(engine, batch):
         engine = Engine(_train_func)
 
         # set up testing handler
-        stats_handler = TensorBoardImageHandler()
+        stats_handler = TensorBoardImageHandler(log_dir=tempdir)
         engine.add_event_handler(Events.ITERATION_COMPLETED, stats_handler)
 
         data = zip(np.random.normal(size=(10, 4, *shape)), np.random.normal(size=(10, 4, *shape)))
         engine.run(data, epoch_length=10, max_epochs=1)
 
-        self.assertTrue(os.path.exists(default_dir))
-        self.assertTrue(len(glob.glob(default_dir)) > 0)
-        shutil.rmtree(default_dir)
+        self.assertTrue(os.path.exists(tempdir))
+        self.assertTrue(len(glob.glob(tempdir)) > 0)
+        shutil.rmtree(tempdir)
 
 
 if __name__ == "__main__":

tests/test_handler_tb_stats.py

@@ -23,8 +23,8 @@
 
 class TestHandlerTBStats(unittest.TestCase):
     def test_metrics_print(self):
-        default_dir = os.path.join(".", "runs")
-        shutil.rmtree(default_dir, ignore_errors=True)
+        tempdir = tempfile.mkdtemp()
+        shutil.rmtree(tempdir, ignore_errors=True)
 
         # set up engine
         def _train_func(engine, batch):
@@ -39,41 +39,41 @@ def _update_metric(engine):
             engine.state.metrics["acc"] = current_metric + 0.1
 
         # set up testing handler
-        stats_handler = TensorBoardStatsHandler()
+        stats_handler = TensorBoardStatsHandler(log_dir=tempdir)
         stats_handler.attach(engine)
         engine.run(range(3), max_epochs=2)
         # check logging output
 
-        self.assertTrue(os.path.exists(default_dir))
-        shutil.rmtree(default_dir)
+        self.assertTrue(os.path.exists(tempdir))
+        shutil.rmtree(tempdir)
 
     def test_metrics_writer(self):
-        default_dir = os.path.join(".", "runs")
-        shutil.rmtree(default_dir, ignore_errors=True)
-        with tempfile.TemporaryDirectory() as temp_dir:
-
-            # set up engine
-            def _train_func(engine, batch):
-                return batch + 1.0
-
-            engine = Engine(_train_func)
-
-            # set up dummy metric
-            @engine.on(Events.EPOCH_COMPLETED)
-            def _update_metric(engine):
-                current_metric = engine.state.metrics.get("acc", 0.1)
-                engine.state.metrics["acc"] = current_metric + 0.1
-
-            # set up testing handler
-            writer = SummaryWriter(log_dir=temp_dir)
-            stats_handler = TensorBoardStatsHandler(
-                writer, output_transform=lambda x: {"loss": x * 2.0}, global_epoch_transform=lambda x: x * 3.0
-            )
-            stats_handler.attach(engine)
-            engine.run(range(3), max_epochs=2)
-            # check logging output
-            self.assertTrue(len(glob.glob(temp_dir)) > 0)
-            self.assertTrue(not os.path.exists(default_dir))
+        tempdir = tempfile.mkdtemp()
+        shutil.rmtree(tempdir, ignore_errors=True)
+
+        # set up engine
+        def _train_func(engine, batch):
+            return batch + 1.0
+
+        engine = Engine(_train_func)
+
+        # set up dummy metric
+        @engine.on(Events.EPOCH_COMPLETED)
+        def _update_metric(engine):
+            current_metric = engine.state.metrics.get("acc", 0.1)
+            engine.state.metrics["acc"] = current_metric + 0.1
+
+        # set up testing handler
+        writer = SummaryWriter(log_dir=tempdir)
+        stats_handler = TensorBoardStatsHandler(
+            writer, output_transform=lambda x: {"loss": x * 2.0}, global_epoch_transform=lambda x: x * 3.0
+        )
+        stats_handler.attach(engine)
+        engine.run(range(3), max_epochs=2)
+        # check logging output
+        self.assertTrue(os.path.exists(tempdir))
+        self.assertTrue(len(glob.glob(tempdir)) > 0)
+        shutil.rmtree(tempdir)
 
 
 if __name__ == "__main__":

tests/test_integration_sliding_window.py

@@ -10,6 +10,7 @@
 # limitations under the License.
 
 import os
+import shutil
 import tempfile
 import unittest
 
@@ -76,13 +77,14 @@ def tearDown(self):
             os.remove(self.seg_name)
 
     def test_training(self):
-        with tempfile.TemporaryDirectory() as temp_dir:
-            output_file = run_test(
-                batch_size=2, img_name=self.img_name, seg_name=self.seg_name, output_dir=temp_dir, device=self.device
-            )
-            output_image = nib.load(output_file).get_fdata()
-            np.testing.assert_allclose(np.sum(output_image), 34070)
-            np.testing.assert_allclose(output_image.shape, (28, 25, 63, 1))
+        tempdir = tempfile.mkdtemp()
+        output_file = run_test(
+            batch_size=2, img_name=self.img_name, seg_name=self.seg_name, output_dir=tempdir, device=self.device
+        )
+        output_image = nib.load(output_file).get_fdata()
+        np.testing.assert_allclose(np.sum(output_image), 34070)
+        np.testing.assert_allclose(output_image.shape, (28, 25, 63, 1))
+        shutil.rmtree(tempdir)
 
 
 if __name__ == "__main__":

tests/test_plot_2d_or_3d_image.py

@@ -11,6 +11,7 @@
 
 import glob
 import os
+import tempfile
 import shutil
 import unittest
 from torch.utils.tensorboard import SummaryWriter
@@ -32,14 +33,14 @@
 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)
+        tempdir = tempfile.mkdtemp()
+        shutil.rmtree(tempdir, ignore_errors=True)
 
-        plot_2d_or_3d_image(torch.zeros(shape), 0, SummaryWriter())
+        plot_2d_or_3d_image(torch.zeros(shape), 0, SummaryWriter(log_dir=tempdir))
 
-        self.assertTrue(os.path.exists(default_dir))
-        self.assertTrue(len(glob.glob(default_dir)) > 0)
-        shutil.rmtree(default_dir, ignore_errors=True)
+        self.assertTrue(os.path.exists(tempdir))
+        self.assertTrue(len(glob.glob(tempdir)) > 0)
+        shutil.rmtree(tempdir, ignore_errors=True)
 
 
 if __name__ == "__main__":