mmocr.visualization.base_visualizer 源代码
# Copyright (c) OpenMMLab. All rights reserved.
import math
from typing import List, Optional, Sequence, Union
import numpy as np
import torch
from matplotlib.font_manager import FontProperties
from mmengine.visualization import Visualizer
from mmocr.registry import VISUALIZERS
[文档]@VISUALIZERS.register_module()
class BaseLocalVisualizer(Visualizer):
"""The MMOCR Text Detection Local Visualizer.
Args:
name (str): Name of the instance. Defaults to 'visualizer'.
image (np.ndarray, optional): the origin image to draw. The format
should be RGB. Defaults to None.
vis_backends (list, optional): Visual backend config list.
Default to None.
save_dir (str, optional): Save file dir for all storage backends.
If it is None, the backend storage will not save any data.
fig_save_cfg (dict): Keyword parameters of figure for saving.
Defaults to empty dict.
fig_show_cfg (dict): Keyword parameters of figure for showing.
Defaults to empty dict.
is_openset (bool, optional): Whether the visualizer is used in
OpenSet. Defaults to False.
font_families (Union[str, List[str]]): The font families of labels.
Defaults to 'sans-serif'.
font_properties (Union[str, FontProperties], optional):
The font properties of texts. The format should be a path str
to font file or a `font_manager.FontProperties()` object.
If you want to draw Chinese texts, you need to prepare
a font file that can show Chinese characters properly.
For example: `simhei.ttf`,`simsun.ttc`,`simkai.ttf` and so on.
Then set font_properties=matplotlib.font_manager.FontProperties
(fname='path/to/font_file') or font_properties='path/to/font_file'
This function need mmengine version >=0.6.0.
Defaults to None.
"""
PALETTE = [(220, 20, 60), (119, 11, 32), (0, 0, 142), (0, 0, 230),
(106, 0, 228), (0, 60, 100), (0, 80, 100), (0, 0, 70),
(0, 0, 192), (250, 170, 30), (100, 170, 30), (220, 220, 0),
(175, 116, 175), (250, 0, 30), (165, 42, 42), (255, 77, 255),
(0, 226, 252), (182, 182, 255), (0, 82, 0), (120, 166, 157),
(110, 76, 0), (174, 57, 255), (199, 100, 0), (72, 0, 118),
(255, 179, 240), (0, 125, 92), (209, 0, 151), (188, 208, 182),
(0, 220, 176), (255, 99, 164), (92, 0, 73), (133, 129, 255),
(78, 180, 255), (0, 228, 0), (174, 255, 243), (45, 89, 255),
(134, 134, 103), (145, 148, 174), (255, 208, 186),
(197, 226, 255), (171, 134, 1), (109, 63, 54), (207, 138, 255),
(151, 0, 95), (9, 80, 61), (84, 105, 51), (74, 65, 105),
(166, 196, 102), (208, 195, 210), (255, 109, 65), (0, 143, 149),
(179, 0, 194), (209, 99, 106), (5, 121, 0), (227, 255, 205),
(147, 186, 208), (153, 69, 1), (3, 95, 161), (163, 255, 0),
(119, 0, 170), (0, 182, 199), (0, 165, 120), (183, 130, 88),
(95, 32, 0), (130, 114, 135), (110, 129, 133), (166, 74, 118),
(219, 142, 185), (79, 210, 114), (178, 90, 62), (65, 70, 15),
(127, 167, 115), (59, 105, 106), (142, 108, 45), (196, 172, 0),
(95, 54, 80), (128, 76, 255), (201, 57, 1), (246, 0, 122),
(191, 162, 208)]
def __init__(self,
name: str = 'visualizer',
font_families: Union[str, List[str]] = 'sans-serif',
font_properties: Optional[Union[str, FontProperties]] = None,
**kwargs) -> None:
super().__init__(name=name, **kwargs)
self.font_families = font_families
self.font_properties = self._set_font_properties(font_properties)
def _set_font_properties(self,
fp: Optional[Union[str, FontProperties]] = None):
if fp is None:
return None
elif isinstance(fp, str):
return FontProperties(fname=fp)
elif isinstance(fp, FontProperties):
return fp
else:
raise ValueError(
'font_properties argument type should be'
' `str` or `matplotlib.font_manager.FontProperties`')
[文档] def get_labels_image(
self,
image: np.ndarray,
labels: Union[np.ndarray, torch.Tensor],
bboxes: Union[np.ndarray, torch.Tensor],
colors: Union[str, Sequence[str]] = 'k',
font_size: Union[int, float] = 10,
auto_font_size: bool = False,
font_families: Union[str, List[str]] = 'sans-serif',
font_properties: Optional[Union[str, FontProperties]] = None
) -> np.ndarray:
"""Draw labels on image.
Args:
image (np.ndarray): The origin image to draw. The format
should be RGB.
labels (Union[np.ndarray, torch.Tensor]): The labels to draw.
bboxes (Union[np.ndarray, torch.Tensor]): The bboxes to draw.
colors (Union[str, Sequence[str]]): The colors of labels.
``colors`` can have the same length with labels or just single
value. If ``colors`` is single value, all the labels will have
the same colors. Refer to `matplotlib.colors` for full list of
formats that are accepted. Defaults to 'k'.
font_size (Union[int, float]): The font size of labels. Defaults
to 10.
auto_font_size (bool): Whether to automatically adjust font size.
Defaults to False.
font_families (Union[str, List[str]]): The font families of labels.
Defaults to 'sans-serif'.
font_properties (Union[str, FontProperties], optional):
The font properties of texts. The format should be a path str
to font file or a `font_manager.FontProperties()` object.
If you want to draw Chinese texts, you need to prepare
a font file that can show Chinese characters properly.
For example: `simhei.ttf`,`simsun.ttc`,`simkai.ttf` and so on.
Then set font_properties=matplotlib.font_manager.FontProperties
(fname='path/to/font_file') or
font_properties='path/to/font_file'.
This function need mmengine version >=0.6.0.
Defaults to None.
"""
if not labels and not bboxes:
return image
if colors is not None and isinstance(colors, (list, tuple)):
size = math.ceil(len(labels) / len(colors))
colors = (colors * size)[:len(labels)]
if auto_font_size:
assert font_size is not None and isinstance(
font_size, (int, float))
font_size = (bboxes[:, 2:] - bboxes[:, :2]).min(-1) * font_size
font_size = font_size.tolist()
self.set_image(image)
self.draw_texts(
labels, (bboxes[:, :2] + bboxes[:, 2:]) / 2,
vertical_alignments='center',
horizontal_alignments='center',
colors='k',
font_sizes=font_size,
font_families=font_families,
font_properties=font_properties)
return self.get_image()
[文档] def get_polygons_image(self,
image: np.ndarray,
polygons: Sequence[np.ndarray],
colors: Union[str, Sequence[str]] = 'g',
filling: bool = False,
line_width: Union[int, float] = 0.5,
alpha: float = 0.5) -> np.ndarray:
"""Draw polygons on image.
Args:
image (np.ndarray): The origin image to draw. The format
should be RGB.
polygons (Sequence[np.ndarray]): The polygons to draw. The shape
should be (N, 2).
colors (Union[str, Sequence[str]]): The colors of polygons.
``colors`` can have the same length with polygons or just
single value. If ``colors`` is single value, all the polygons
will have the same colors. Refer to `matplotlib.colors` for
full list of formats that are accepted. Defaults to 'g'.
filling (bool): Whether to fill the polygons. Defaults to False.
line_width (Union[int, float]): The line width of polygons.
Defaults to 0.5.
alpha (float): The alpha of polygons. Defaults to 0.5.
Returns:
np.ndarray: The image with polygons drawn.
"""
if colors is not None and isinstance(colors, (list, tuple)):
size = math.ceil(len(polygons) / len(colors))
colors = (colors * size)[:len(polygons)]
self.set_image(image)
if filling:
self.draw_polygons(
polygons,
face_colors=colors,
edge_colors=colors,
line_widths=line_width,
alpha=alpha)
else:
self.draw_polygons(
polygons,
edge_colors=colors,
line_widths=line_width,
alpha=alpha)
return self.get_image()
[文档] def get_bboxes_image(self: Visualizer,
image: np.ndarray,
bboxes: Union[np.ndarray, torch.Tensor],
colors: Union[str, Sequence[str]] = 'g',
filling: bool = False,
line_width: Union[int, float] = 0.5,
alpha: float = 0.5) -> np.ndarray:
"""Draw bboxes on image.
Args:
image (np.ndarray): The origin image to draw. The format
should be RGB.
bboxes (Union[np.ndarray, torch.Tensor]): The bboxes to draw.
colors (Union[str, Sequence[str]]): The colors of bboxes.
``colors`` can have the same length with bboxes or just single
value. If ``colors`` is single value, all the bboxes will have
the same colors. Refer to `matplotlib.colors` for full list of
formats that are accepted. Defaults to 'g'.
filling (bool): Whether to fill the bboxes. Defaults to False.
line_width (Union[int, float]): The line width of bboxes.
Defaults to 0.5.
alpha (float): The alpha of bboxes. Defaults to 0.5.
Returns:
np.ndarray: The image with bboxes drawn.
"""
if colors is not None and isinstance(colors, (list, tuple)):
size = math.ceil(len(bboxes) / len(colors))
colors = (colors * size)[:len(bboxes)]
self.set_image(image)
if filling:
self.draw_bboxes(
bboxes,
face_colors=colors,
edge_colors=colors,
line_widths=line_width,
alpha=alpha)
else:
self.draw_bboxes(
bboxes,
edge_colors=colors,
line_widths=line_width,
alpha=alpha)
return self.get_image()
def _draw_instances(self) -> np.ndarray:
raise NotImplementedError
def _cat_image(self, imgs: Sequence[np.ndarray], axis: int) -> np.ndarray:
"""Concatenate images.
Args:
imgs (Sequence[np.ndarray]): The images to concatenate.
axis (int): The axis to concatenate.
Returns:
np.ndarray: The concatenated image.
"""
cat_image = list()
for img in imgs:
if img is not None:
cat_image.append(img)
if len(cat_image):
return np.concatenate(cat_image, axis=axis)
else:
return None