|
- # 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)
|