Class API for Visualize Module

[abpwrs]

Is your feature request related to a problem? Please describe.

The current monai.visualize module is highly coupled with Tensorboard's SummaryWriter. It'd be convenient to have a set of classes designed to process tensors/arrays both inside and outside the training loop to generate common visualizations, like a segmentation overlay.

Describe the solution you'd like

A restructuring of monai.visualize to support out of training loop visualization, and easier creation of custom visualization. I could see this including the following:

• Create an abstract Visualization base class
• Refactoring current 2D/3D visualizers to classes that implement Visualization
• Refactor TensorBoardImageHandler to support a list of Visualization objects
• Add an add(viz : Visualization) method, or the like, to TensorBoardImageHandler

Describe alternatives you've considered

The monai.visualize module could also be restructured to be a functional API instead of using classes.

Additional context

Example usages:

Outside training loop

image_data = itk.GetArrayFromImage(itk.imread("image1.nii.gz")
seg_data = itk.GetArrayFromImage(itk.imread("seg1.nii.gz")
segmentation_overlay_visualizer = monai.visualize.SegmentationOverlayVisualization()
viz = segmentation_overlay_visualizer(image_data, seg_data)

Inside training loop

train_tensorboard_image_handler = TensorBoardImageHandler(
    batch_transform=lambda batch: (batch[0], batch[1]),
    output_transform=lambda output: predict_segmentation(output[0]),
    global_iter_transform=lambda x: trainer.state.epoch
)

# the add function could take in any class that implements the Visualization class
train_tensorboard_image_handler.add(monai.visualize.SegmentationOverlayVisualization())

trainer.add_event_handler(event_name=Events.EPOCH_COMPLETED, handler=train_tensorboard_image_handler)

[abpwrs]

@hjmjohnson

[wyli]

thanks, is the variableviz an image array in viz = segmentation_overlay_visualizer(image_data, seg_data)? if yes then I guess segmentation_overlay_visualizer is actually a data transform which does some postprocessing and fuses image/seg data into one?

[abpwrs]

Correct, this would be post-processing of the data to overlay the outline of the segmentation onto the image.

[wyli]

ok, i think it would be great to have generic post-processing layers inmonai.transforms for alpha blending or mixup https://forums.fast.ai/t/mixup-data-augmentation/22764

then in the TensorBoardImageHandler we could use it as an output_transform:

train_tensorboard_image_handler = TensorBoardImageHandler(
    batch_transform=lambda batch: (batch[0], batch[1]),
    output_transform=lambda output: segmentation_overlay_transform(
          output[0], predict_segmentation(output[1])),
    global_iter_transform=lambda x: trainer.state.epoch
)

but I agree with you about the idea of decoupling tensorboard writing and gif/visuals creations

[abpwrs]

If a user wants to use multiple of these post-processing layers/transforms, would they then need to create multiple TensorBoardImageHandler objects?

Otherwise, this solution sounds great!
I can work on this feature and start a PR.

[wyli]

If a user wants to use multiple of these post-processing layers/transforms, would they then need to create multiple TensorBoardImageHandler objects?

Otherwise, this solution sounds great!
I can work on this feature and start a PR.

great thanks! I think as a starting point we focus on single postprocessing to tensorboard, so that the interface is intuitive. Once we identified commonly used combinations of multiple post-processing visuals, we could consider a slightly high-level interface.

[vikashg]

close as merged