mirror of
https://github.com/open-mmlab/mmengine.git
synced 2025-06-03 21:54:44 +08:00
5f8f36e6a5
* [WIP] add inline * refactor vis module * [Refactor] according review * [Fix] fix comment * fix some error * Get sub visualizer be Visualizer.get_instance * fix conflict * fix lint * fix unit test * fix mypy * fix comment * fix lint * update docstr * update * update instancedata * remove replace __mro__ with issubclass Co-authored-by: PJLAB\huanghaian <1286304229@qq.com> Co-authored-by: HAOCHENYE <21724054@zju.edu.cn>
476 lines
18 KiB
Python
476 lines
18 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import copy
|
|
from typing import Any, List, Optional, Union
|
|
from unittest import TestCase
|
|
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from mmengine import VISBACKENDS
|
|
from mmengine.visualization import Visualizer
|
|
|
|
|
|
@VISBACKENDS.register_module()
|
|
class MockVisBackend:
|
|
|
|
def __init__(self, save_dir: Optional[str] = None):
|
|
self._save_dir = save_dir
|
|
self._close = False
|
|
|
|
@property
|
|
def experiment(self) -> Any:
|
|
return self
|
|
|
|
def add_config(self, params_dict: dict, **kwargs) -> None:
|
|
self._add_config = True
|
|
|
|
def add_graph(self, model: torch.nn.Module,
|
|
input_tensor: Union[torch.Tensor,
|
|
List[torch.Tensor]], **kwargs) -> None:
|
|
|
|
self._add_graph = True
|
|
|
|
def add_image(self,
|
|
name: str,
|
|
image: np.ndarray,
|
|
step: int = 0,
|
|
**kwargs) -> None:
|
|
self._add_image = True
|
|
|
|
def add_scalar(self,
|
|
name: str,
|
|
value: Union[int, float],
|
|
step: int = 0,
|
|
**kwargs) -> None:
|
|
self._add_scalar = True
|
|
|
|
def add_scalars(self,
|
|
scalar_dict: dict,
|
|
step: int = 0,
|
|
file_path: Optional[str] = None,
|
|
**kwargs) -> None:
|
|
self._add_scalars = True
|
|
|
|
def close(self) -> None:
|
|
"""close an opened object."""
|
|
self._close = True
|
|
|
|
|
|
class TestVisualizer(TestCase):
|
|
|
|
def setUp(self):
|
|
"""Setup the demo image in every test method.
|
|
|
|
TestCase calls functions in this order: setUp() -> testMethod() ->
|
|
tearDown() -> cleanUp()
|
|
"""
|
|
self.image = np.random.randint(
|
|
0, 256, size=(10, 10, 3)).astype('uint8')
|
|
self.vis_backend_cfg = [
|
|
dict(type='MockVisBackend', name='mock1', save_dir='tmp'),
|
|
dict(type='MockVisBackend', name='mock2', save_dir='tmp')
|
|
]
|
|
|
|
def test_init(self):
|
|
visualizer = Visualizer(image=self.image)
|
|
visualizer.get_image()
|
|
|
|
visualizer = Visualizer(
|
|
vis_backends=copy.deepcopy(self.vis_backend_cfg))
|
|
assert isinstance(visualizer.get_backend('mock1'), MockVisBackend)
|
|
assert len(visualizer._vis_backends) == 2
|
|
|
|
# test global
|
|
visualizer = Visualizer.get_instance(
|
|
'visualizer', vis_backends=copy.deepcopy(self.vis_backend_cfg))
|
|
assert len(visualizer._vis_backends) == 2
|
|
visualizer_any = Visualizer.get_instance('visualizer')
|
|
assert visualizer_any == visualizer
|
|
|
|
def test_set_image(self):
|
|
visualizer = Visualizer()
|
|
visualizer.set_image(self.image)
|
|
with pytest.raises(AssertionError):
|
|
visualizer.set_image(None)
|
|
|
|
def test_get_image(self):
|
|
visualizer = Visualizer(image=self.image)
|
|
visualizer.get_image()
|
|
|
|
def test_draw_bboxes(self):
|
|
visualizer = Visualizer(image=self.image)
|
|
|
|
# only support 4 or nx4 tensor and numpy
|
|
visualizer.draw_bboxes(torch.tensor([1, 1, 2, 2]))
|
|
# valid bbox
|
|
visualizer.draw_bboxes(torch.tensor([1, 1, 1, 2]))
|
|
bboxes = torch.tensor([[1, 1, 2, 2], [1, 2, 2, 2.5]])
|
|
visualizer.draw_bboxes(
|
|
bboxes, alpha=0.5, edge_colors=(255, 0, 0), line_styles='-')
|
|
bboxes = bboxes.numpy()
|
|
visualizer.draw_bboxes(bboxes)
|
|
|
|
# test invalid bbox
|
|
with pytest.raises(AssertionError):
|
|
# x1 > x2
|
|
visualizer.draw_bboxes(torch.tensor([5, 1, 2, 2]))
|
|
|
|
# test out of bounds
|
|
with pytest.warns(
|
|
UserWarning,
|
|
match='Warning: The bbox is out of bounds,'
|
|
' the drawn bbox may not be in the image'):
|
|
visualizer.draw_bboxes(torch.tensor([1, 1, 20, 2]))
|
|
|
|
# test incorrect bbox format
|
|
with pytest.raises(TypeError):
|
|
visualizer.draw_bboxes([1, 1, 2, 2])
|
|
|
|
def test_close(self):
|
|
visualizer = Visualizer(
|
|
image=self.image, vis_backends=copy.deepcopy(self.vis_backend_cfg))
|
|
fig_num = visualizer.fig_save_num
|
|
assert fig_num in plt.get_fignums()
|
|
for name in ['mock1', 'mock2']:
|
|
assert visualizer.get_backend(name)._close is False
|
|
visualizer.close()
|
|
assert fig_num not in plt.get_fignums()
|
|
for name in ['mock1', 'mock2']:
|
|
assert visualizer.get_backend(name)._close is True
|
|
|
|
def test_draw_texts(self):
|
|
visualizer = Visualizer(image=self.image)
|
|
|
|
# only support tensor and numpy
|
|
visualizer.draw_texts(
|
|
'text1', positions=torch.tensor([5, 5]), colors=(0, 255, 0))
|
|
visualizer.draw_texts(['text1', 'text2'],
|
|
positions=torch.tensor([[5, 5], [3, 3]]),
|
|
colors=[(255, 0, 0), (255, 0, 0)])
|
|
visualizer.draw_texts('text1', positions=np.array([5, 5]))
|
|
visualizer.draw_texts(['text1', 'text2'],
|
|
positions=np.array([[5, 5], [3, 3]]))
|
|
visualizer.draw_texts(
|
|
'text1',
|
|
positions=torch.tensor([5, 5]),
|
|
bboxes=dict(facecolor='r', alpha=0.6))
|
|
# test out of bounds
|
|
with pytest.warns(
|
|
UserWarning,
|
|
match='Warning: The text is out of bounds,'
|
|
' the drawn text may not be in the image'):
|
|
visualizer.draw_texts('text1', positions=torch.tensor([15, 5]))
|
|
|
|
# test incorrect format
|
|
with pytest.raises(TypeError):
|
|
visualizer.draw_texts('text', positions=[5, 5])
|
|
|
|
# test length mismatch
|
|
with pytest.raises(AssertionError):
|
|
visualizer.draw_texts(['text1', 'text2'],
|
|
positions=torch.tensor([5, 5]))
|
|
with pytest.raises(AssertionError):
|
|
visualizer.draw_texts(
|
|
'text1', positions=torch.tensor([[5, 5], [3, 3]]))
|
|
with pytest.raises(AssertionError):
|
|
visualizer.draw_texts(['text1', 'test2'],
|
|
positions=torch.tensor([[5, 5], [3, 3]]),
|
|
colors=['r'])
|
|
with pytest.raises(AssertionError):
|
|
visualizer.draw_texts(['text1', 'test2'],
|
|
positions=torch.tensor([[5, 5], [3, 3]]),
|
|
vertical_alignments=['top'])
|
|
with pytest.raises(AssertionError):
|
|
visualizer.draw_texts(['text1', 'test2'],
|
|
positions=torch.tensor([[5, 5], [3, 3]]),
|
|
horizontal_alignments=['left'])
|
|
with pytest.raises(AssertionError):
|
|
visualizer.draw_texts(['text1', 'test2'],
|
|
positions=torch.tensor([[5, 5], [3, 3]]),
|
|
font_sizes=[1])
|
|
|
|
# test type valid
|
|
with pytest.raises(TypeError):
|
|
visualizer.draw_texts(['text1', 'test2'],
|
|
positions=torch.tensor([[5, 5], [3, 3]]),
|
|
font_sizes='b')
|
|
|
|
def test_draw_lines(self):
|
|
visualizer = Visualizer(image=self.image)
|
|
|
|
# only support tensor and numpy
|
|
visualizer.draw_lines(
|
|
x_datas=torch.tensor([1, 5]), y_datas=torch.tensor([2, 6]))
|
|
visualizer.draw_lines(
|
|
x_datas=np.array([[1, 5], [2, 4]]),
|
|
y_datas=np.array([[2, 6], [4, 7]]))
|
|
visualizer.draw_lines(
|
|
x_datas=np.array([[1, 5], [2, 4]]),
|
|
y_datas=np.array([[2, 6], [4, 7]]),
|
|
colors='r',
|
|
line_styles=['-', '-.'],
|
|
line_widths=[1, 2])
|
|
# test out of bounds
|
|
with pytest.warns(
|
|
UserWarning,
|
|
match='Warning: The line is out of bounds,'
|
|
' the drawn line may not be in the image'):
|
|
visualizer.draw_lines(
|
|
x_datas=torch.tensor([12, 5]), y_datas=torch.tensor([2, 6]))
|
|
|
|
# test incorrect format
|
|
with pytest.raises(TypeError):
|
|
visualizer.draw_lines(x_datas=[5, 5], y_datas=torch.tensor([2, 6]))
|
|
with pytest.raises(TypeError):
|
|
visualizer.draw_lines(y_datas=[5, 5], x_datas=torch.tensor([2, 6]))
|
|
|
|
# test length mismatch
|
|
with pytest.raises(AssertionError):
|
|
visualizer.draw_lines(
|
|
x_datas=torch.tensor([1, 5]),
|
|
y_datas=torch.tensor([[2, 6], [4, 7]]))
|
|
|
|
def test_draw_circles(self):
|
|
visualizer = Visualizer(image=self.image)
|
|
|
|
# only support tensor and numpy
|
|
visualizer.draw_circles(torch.tensor([1, 5]), torch.tensor([1]))
|
|
visualizer.draw_circles(np.array([1, 5]), np.array([1]))
|
|
visualizer.draw_circles(
|
|
torch.tensor([[1, 5], [2, 6]]), radius=torch.tensor([1, 2]))
|
|
|
|
# test face_colors
|
|
visualizer.draw_circles(
|
|
torch.tensor([[1, 5], [2, 6]]),
|
|
radius=torch.tensor([1, 2]),
|
|
face_colors=(255, 0, 0),
|
|
edge_colors=(255, 0, 0))
|
|
|
|
# test config
|
|
visualizer.draw_circles(
|
|
torch.tensor([[1, 5], [2, 6]]),
|
|
radius=torch.tensor([1, 2]),
|
|
edge_colors=['g', 'r'],
|
|
line_styles=['-', '-.'],
|
|
line_widths=[1, 2])
|
|
|
|
# test out of bounds
|
|
with pytest.warns(
|
|
UserWarning,
|
|
match='Warning: The circle is out of bounds,'
|
|
' the drawn circle may not be in the image'):
|
|
visualizer.draw_circles(
|
|
torch.tensor([12, 5]), radius=torch.tensor([1]))
|
|
visualizer.draw_circles(
|
|
torch.tensor([1, 5]), radius=torch.tensor([10]))
|
|
|
|
# test incorrect format
|
|
with pytest.raises(TypeError):
|
|
visualizer.draw_circles([1, 5], radius=torch.tensor([1]))
|
|
with pytest.raises(TypeError):
|
|
visualizer.draw_circles(np.array([1, 5]), radius=10)
|
|
|
|
# test length mismatch
|
|
with pytest.raises(AssertionError):
|
|
visualizer.draw_circles(
|
|
torch.tensor([[1, 5]]), radius=torch.tensor([1, 2]))
|
|
|
|
def test_draw_polygons(self):
|
|
visualizer = Visualizer(image=self.image)
|
|
# shape Nx2 or list[Nx2]
|
|
visualizer.draw_polygons(torch.tensor([[1, 1], [2, 2], [3, 4]]))
|
|
visualizer.draw_polygons(np.array([[1, 1], [2, 2], [3, 4]]))
|
|
visualizer.draw_polygons([
|
|
np.array([[1, 1], [2, 2], [3, 4]]),
|
|
torch.tensor([[1, 1], [2, 2], [3, 4]])
|
|
])
|
|
visualizer.draw_polygons(
|
|
polygons=[
|
|
np.array([[1, 1], [2, 2], [3, 4]]),
|
|
torch.tensor([[1, 1], [2, 2], [3, 4]])
|
|
],
|
|
face_colors=(255, 0, 0),
|
|
edge_colors=(255, 0, 0))
|
|
visualizer.draw_polygons(
|
|
polygons=[
|
|
np.array([[1, 1], [2, 2], [3, 4]]),
|
|
torch.tensor([[1, 1], [2, 2], [3, 4]])
|
|
],
|
|
edge_colors=['r', 'g'],
|
|
line_styles='-',
|
|
line_widths=[2, 1])
|
|
|
|
# test out of bounds
|
|
with pytest.warns(
|
|
UserWarning,
|
|
match='Warning: The polygon is out of bounds,'
|
|
' the drawn polygon may not be in the image'):
|
|
visualizer.draw_polygons(torch.tensor([[1, 1], [2, 2], [16, 4]]))
|
|
|
|
def test_draw_binary_masks(self):
|
|
binary_mask = np.random.randint(0, 2, size=(10, 10)).astype(np.bool)
|
|
visualizer = Visualizer(image=self.image)
|
|
visualizer.draw_binary_masks(binary_mask)
|
|
visualizer.draw_binary_masks(torch.from_numpy(binary_mask))
|
|
# multi binary
|
|
binary_mask = np.random.randint(0, 2, size=(2, 10, 10)).astype(np.bool)
|
|
visualizer = Visualizer(image=self.image)
|
|
visualizer.draw_binary_masks(binary_mask, colors=['r', (0, 255, 0)])
|
|
# test the error that the size of mask and image are different.
|
|
with pytest.raises(AssertionError):
|
|
binary_mask = np.random.randint(0, 2, size=(8, 10)).astype(np.bool)
|
|
visualizer.draw_binary_masks(binary_mask)
|
|
|
|
# test non binary mask error
|
|
binary_mask = np.random.randint(0, 2, size=(10, 10, 3)).astype(np.bool)
|
|
with pytest.raises(AssertionError):
|
|
visualizer.draw_binary_masks(binary_mask)
|
|
|
|
# test color dim
|
|
with pytest.raises(AssertionError):
|
|
visualizer.draw_binary_masks(
|
|
binary_mask, colors=np.array([1, 22, 4, 45]))
|
|
binary_mask = np.random.randint(0, 2, size=(10, 10))
|
|
with pytest.raises(AssertionError):
|
|
visualizer.draw_binary_masks(binary_mask)
|
|
|
|
def test_draw_featmap(self):
|
|
visualizer = Visualizer()
|
|
image = np.random.randint(0, 256, size=(3, 3, 3), dtype='uint8')
|
|
# test tensor format
|
|
with pytest.raises(AssertionError, match='Input dimension must be 3'):
|
|
visualizer.draw_featmap(torch.randn(1, 1, 3, 3))
|
|
|
|
# test mode parameter
|
|
# mode only supports 'mean' and 'max'
|
|
with pytest.raises(AssertionError):
|
|
visualizer.draw_featmap(torch.randn(2, 3, 3), mode='xx')
|
|
# test tensor_chw and img have difference height and width
|
|
with pytest.raises(AssertionError):
|
|
visualizer.draw_featmap(torch.randn(2, 3, 3), mode='xx')
|
|
|
|
# test topk parameter
|
|
with pytest.raises(
|
|
AssertionError,
|
|
match='The input tensor channel dimension must be 1 or 3 '
|
|
'when topk is less than 1, but the channel '
|
|
'dimension you input is 6, you can use the '
|
|
'mode parameter or set topk greater than 0 to solve '
|
|
'the error'):
|
|
visualizer.draw_featmap(torch.randn(6, 3, 3), mode=None, topk=0)
|
|
|
|
visualizer.draw_featmap(torch.randn(6, 3, 3), mode='mean')
|
|
visualizer.draw_featmap(torch.randn(1, 3, 3), mode='mean')
|
|
visualizer.draw_featmap(torch.randn(6, 3, 3), mode='max')
|
|
visualizer.draw_featmap(torch.randn(6, 3, 3), mode='max', topk=10)
|
|
visualizer.draw_featmap(torch.randn(1, 3, 3), mode=None, topk=-1)
|
|
visualizer.draw_featmap(
|
|
torch.randn(3, 3, 3), image=image, mode=None, topk=-1)
|
|
visualizer.draw_featmap(torch.randn(6, 3, 3), mode=None, topk=4)
|
|
visualizer.draw_featmap(
|
|
torch.randn(6, 3, 3), image=image, mode=None, topk=8)
|
|
|
|
# test gray
|
|
visualizer.draw_featmap(
|
|
torch.randn(6, 3, 3),
|
|
image=np.random.randint(0, 256, size=(3, 3), dtype='uint8'),
|
|
mode=None,
|
|
topk=8)
|
|
|
|
# test arrangement
|
|
with pytest.raises(AssertionError):
|
|
visualizer.draw_featmap(
|
|
torch.randn(10, 3, 3),
|
|
image=image,
|
|
mode=None,
|
|
topk=8,
|
|
arrangement=(2, 2))
|
|
|
|
def test_chain_call(self):
|
|
visualizer = Visualizer(image=self.image)
|
|
binary_mask = np.random.randint(0, 2, size=(10, 10)).astype(np.bool)
|
|
visualizer.draw_bboxes(torch.tensor([1, 1, 2, 2])). \
|
|
draw_texts('test', torch.tensor([5, 5])). \
|
|
draw_lines(x_datas=torch.tensor([1, 5]),
|
|
y_datas=torch.tensor([2, 6])). \
|
|
draw_circles(torch.tensor([1, 5]), radius=torch.tensor([2])). \
|
|
draw_polygons(torch.tensor([[1, 1], [2, 2], [3, 4]])). \
|
|
draw_binary_masks(binary_mask)
|
|
|
|
def test_get_backend(self):
|
|
visualizer = Visualizer(
|
|
image=self.image, vis_backends=copy.deepcopy(self.vis_backend_cfg))
|
|
for name in ['mock1', 'mock2']:
|
|
assert isinstance(visualizer.get_backend(name), MockVisBackend)
|
|
|
|
def test_add_config(self):
|
|
visualizer = Visualizer(
|
|
vis_backends=copy.deepcopy(self.vis_backend_cfg))
|
|
|
|
params_dict = dict(lr=0.1, wd=0.2, mode='linear')
|
|
visualizer.add_config(params_dict)
|
|
for name in ['mock1', 'mock2']:
|
|
assert visualizer.get_backend(name)._add_config is True
|
|
|
|
def test_add_graph(self):
|
|
visualizer = Visualizer(
|
|
vis_backends=copy.deepcopy(self.vis_backend_cfg))
|
|
|
|
class Model(nn.Module):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv = nn.Conv2d(1, 2, 1)
|
|
|
|
def forward(self, x, y=None):
|
|
return self.conv(x)
|
|
|
|
visualizer.add_graph(Model(), np.zeros([1, 1, 3, 3]))
|
|
for name in ['mock1', 'mock2']:
|
|
assert visualizer.get_backend(name)._add_graph is True
|
|
|
|
def test_add_image(self):
|
|
image = np.random.randint(0, 256, size=(10, 10, 3)).astype(np.uint8)
|
|
visualizer = Visualizer(
|
|
vis_backends=copy.deepcopy(self.vis_backend_cfg))
|
|
|
|
visualizer.add_image('img', image)
|
|
for name in ['mock1', 'mock2']:
|
|
assert visualizer.get_backend(name)._add_image is True
|
|
|
|
def test_add_scalar(self):
|
|
visualizer = Visualizer(
|
|
vis_backends=copy.deepcopy(self.vis_backend_cfg))
|
|
visualizer.add_scalar('map', 0.9, step=0)
|
|
for name in ['mock1', 'mock2']:
|
|
assert visualizer.get_backend(name)._add_scalar is True
|
|
|
|
def test_add_scalars(self):
|
|
visualizer = Visualizer(
|
|
vis_backends=copy.deepcopy(self.vis_backend_cfg))
|
|
input_dict = {'map': 0.7, 'acc': 0.9}
|
|
visualizer.add_scalars(input_dict)
|
|
for name in ['mock1', 'mock2']:
|
|
assert visualizer.get_backend(name)._add_scalars is True
|
|
|
|
def test_get_instance(self):
|
|
|
|
class DetLocalVisualizer(Visualizer):
|
|
|
|
def __init__(self, name):
|
|
super().__init__(name)
|
|
|
|
visualizer1 = DetLocalVisualizer.get_instance('name1')
|
|
visualizer2 = Visualizer.get_current_instance()
|
|
visualizer3 = DetLocalVisualizer.get_current_instance()
|
|
assert id(visualizer1) == id(visualizer2) == id(visualizer3)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
t = TestVisualizer()
|
|
t.setUp()
|
|
t.test_init()
|