Unify visualisation

docs/source/metrics.rst

@@ -37,7 +37,3 @@ Metrics
 
 .. autoclass:: SurfaceDistanceMetric
     :members:
-
-`Occlusion sensitivity`
------------------------
-.. autofunction:: compute_occlusion_sensitivity

docs/source/visualize.rst

@@ -17,4 +17,10 @@ Class activation map
 --------------------
 
 .. automodule:: monai.visualize.class_activation_maps
-  :members:
+  :members:
+
+Occlusion sensitivity
+---------------------
+
+.. automodule:: monai.visualize.occlusion_sensitivity
+  :members:

monai/handlers/tensorboard_handlers.py

@@ -174,7 +174,7 @@ def _default_iteration_writer(self, engine: Engine, writer: SummaryWriter) -> No
 
 class TensorBoardImageHandler(object):
     """
-    TensorBoardImageHandler is an Ignite Event handler that can visualise images, labels and outputs as 2D/3D images.
+    TensorBoardImageHandler is an Ignite Event handler that can visualize images, labels and outputs as 2D/3D images.
     2D output (shape in Batch, channel, H, W) will be shown as simple image using the first element in the batch,
     for 3D to ND output (shape in Batch, channel, H, W, D) input, each of ``self.max_channels`` number of images'
     last three dimensions will be shown as animated GIF along the last axis (typically Depth).

monai/metrics/__init__.py

@@ -12,7 +12,6 @@
 from .confusion_matrix import ConfusionMatrixMetric, compute_confusion_matrix_metric, get_confusion_matrix
 from .hausdorff_distance import *
 from .meandice import DiceMetric, compute_meandice
-from .occlusion_sensitivity import compute_occlusion_sensitivity
 from .rocauc import compute_roc_auc
 from .surface_distance import SurfaceDistanceMetric, compute_average_surface_distance
 from .utils import *

monai/networks/utils.py

@@ -31,6 +31,7 @@
     "icnr_init",
     "pixelshuffle",
     "eval_mode",
+    "train_mode",
 ]
 
 
@@ -276,3 +277,37 @@ def eval_mode(*nets: nn.Module):
         # Return required networks to training
         for n in training:
             n.train()
+
+
+@contextmanager
+def train_mode(*nets: nn.Module):
+    """
+    Set network(s) to train mode and then return to original state at the end.
+
+    Args:
+        nets: Input network(s)
+
+    Examples
+
+    .. code-block:: python
+
+        t=torch.rand(1,1,16,16)
+        p=torch.nn.Conv2d(1,1,3)
+        p.eval()
+        print(p.training)  # False
+        with train_mode(p):
+            print(p.training)  # True
+            print(p(t).sum().backward())  # No exception
+    """
+
+    # Get original state of network(s)
+    eval = [n for n in nets if not n.training]
+
+    try:
+        # set to train mode
+        with torch.set_grad_enabled(True):
+            yield [n.train() for n in nets]
+    finally:
+        # Return required networks to eval
+        for n in eval:
+            n.eval()

monai/visualize/__init__.py

@@ -9,5 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .visualizer import default_normalizer, default_upsampler  # isort:skip
 from .class_activation_maps import *
 from .img2tensorboard import *
+from .occlusion_sensitivity import OcclusionSensitivity