102 lines
4.0 KiB
Python
102 lines
4.0 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import os.path as osp
|
|
import warnings
|
|
from typing import Sequence
|
|
|
|
import mmcv
|
|
from mmengine.hooks import Hook
|
|
from mmengine.runner import Runner
|
|
|
|
from mmseg.data import SegDataSample
|
|
from mmseg.registry import HOOKS
|
|
from mmseg.visualization import SegLocalVisualizer
|
|
|
|
|
|
@HOOKS.register_module()
|
|
class SegVisualizationHook(Hook):
|
|
"""Segmentation Visualization Hook. Used to visualize validation and
|
|
testing process prediction results.
|
|
|
|
In the testing phase:
|
|
|
|
1. If ``show`` is True, it means that only the prediction results are
|
|
visualized without storing data, so ``vis_backends`` needs to
|
|
be excluded.
|
|
|
|
Args:
|
|
draw (bool): whether to draw prediction results. If it is False,
|
|
it means that no drawing will be done. Defaults to False.
|
|
interval (int): The interval of visualization. Defaults to 50.
|
|
show (bool): Whether to display the drawn image. Default to False.
|
|
wait_time (float): The interval of show (s). Defaults to 0.
|
|
file_client_args (dict): Arguments to instantiate a FileClient.
|
|
See :class:`mmcv.fileio.FileClient` for details.
|
|
Defaults to ``dict(backend='disk')``.
|
|
"""
|
|
|
|
def __init__(self,
|
|
draw: bool = False,
|
|
interval: int = 50,
|
|
show: bool = False,
|
|
wait_time: float = 0.,
|
|
file_client_args: dict = dict(backend='disk')):
|
|
self._visualizer: SegLocalVisualizer = \
|
|
SegLocalVisualizer.get_current_instance()
|
|
self.interval = interval
|
|
self.show = show
|
|
if self.show:
|
|
# No need to think about vis backends.
|
|
self._visualizer._vis_backends = {}
|
|
warnings.warn('The show is True, it means that only '
|
|
'the prediction results are visualized '
|
|
'without storing data, so vis_backends '
|
|
'needs to be excluded.')
|
|
|
|
self.wait_time = wait_time
|
|
self.file_client_args = file_client_args.copy()
|
|
self.file_client = None
|
|
self.draw = draw
|
|
if not self.draw:
|
|
warnings.warn('The draw is False, it means that the '
|
|
'hook for visualization will not take '
|
|
'effect. The results will NOT be '
|
|
'visualized or stored.')
|
|
|
|
def after_iter(self,
|
|
runner: Runner,
|
|
batch_idx: int,
|
|
data_batch: Sequence[dict],
|
|
outputs: Sequence[SegDataSample],
|
|
mode: str = 'val') -> None:
|
|
"""Run after every ``self.interval`` validation iterations.
|
|
|
|
Args:
|
|
runner (:obj:`Runner`): The runner of the validation process.
|
|
batch_idx (int): The index of the current batch in the val loop.
|
|
data_batch (Sequence[dict]): Data from dataloader.
|
|
outputs (Sequence[:obj:`SegDataSample`]): Outputs from model.
|
|
mode (str): mode (str): Current mode of runner. Defaults to 'val'.
|
|
"""
|
|
if self.draw is False or mode == 'train':
|
|
return
|
|
|
|
if self.file_client is None:
|
|
self.file_client = mmcv.FileClient(**self.file_client_args)
|
|
|
|
if self.every_n_inner_iters(batch_idx, self.interval):
|
|
for input_data, output in zip(data_batch, outputs):
|
|
img_path = input_data['data_sample'].img_path
|
|
img_bytes = self.file_client.get(img_path)
|
|
img = mmcv.imfrombytes(img_bytes, channel_order='rgb')
|
|
window_name = f'{mode}_{osp.basename(img_path)}'
|
|
|
|
gt_sample = input_data['data_sample']
|
|
self._visualizer.add_datasample(
|
|
window_name,
|
|
img,
|
|
gt_sample=gt_sample,
|
|
pred_sample=output,
|
|
show=self.show,
|
|
wait_time=self.wait_time,
|
|
step=runner.iter)
|