Shortcuts

mmengine.hooks.naive_visualization_hook 源代码

# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
from typing import Optional, Sequence, Tuple, Union

import cv2
import numpy as np

from mmengine.hooks import Hook
from mmengine.registry import HOOKS
from mmengine.utils.dl_utils import tensor2imgs

DATA_BATCH = Optional[Union[dict, tuple, list]]


# TODO: Due to interface changes, the current class
#  functions incorrectly
[文档]@HOOKS.register_module() class NaiveVisualizationHook(Hook): """Show or Write the predicted results during the process of testing. Args: interval (int): Visualization interval. Defaults to 1. draw_gt (bool): Whether to draw the ground truth. Defaults to True. draw_pred (bool): Whether to draw the predicted result. Defaults to True. """ priority = 'NORMAL' def __init__(self, interval: int = 1, draw_gt: bool = True, draw_pred: bool = True): self.draw_gt = draw_gt self.draw_pred = draw_pred self._interval = interval def _unpad(self, input: np.ndarray, unpad_shape: Tuple[int, int]) -> np.ndarray: """Unpad the input image. Args: input (np.ndarray): The image to unpad. unpad_shape (tuple): The shape of image before padding. Returns: np.ndarray: The image before padding. """ unpad_width, unpad_height = unpad_shape unpad_image = input[:unpad_height, :unpad_width] return unpad_image
[文档] def before_train(self, runner) -> None: """Call add_graph method of visualizer. Args: runner (Runner): The runner of the training process. """ runner.visualizer.add_graph(runner.model, None)
[文档] def after_test_iter(self, runner, batch_idx: int, data_batch: DATA_BATCH = None, outputs: Optional[Sequence] = None) -> None: """Show or Write the predicted results. Args: runner (Runner): The runner of the training process. batch_idx (int): The index of the current batch in the test loop. data_batch (dict or tuple or list, optional): Data from dataloader. outputs (Sequence, optional): Outputs from model. """ if self.every_n_inner_iters(batch_idx, self._interval): for data, output in zip(data_batch, outputs): # type: ignore input = data['inputs'] data_sample = data['data_sample'] input = tensor2imgs(input, **data_sample.get('img_norm_cfg', dict()))[0] # TODO We will implement a function to revert the augmentation # in the future. ori_shape = (data_sample.ori_width, data_sample.ori_height) if 'pad_shape' in data_sample: input = self._unpad(input, data_sample.get('scale', ori_shape)) origin_image = cv2.resize(input, ori_shape) name = osp.basename(data_sample.img_path) runner.visualizer.add_datasample(name, origin_image, data_sample, output, self.draw_gt, self.draw_pred)