diff --git a/mmocr/apis/inferencers/mmocr_inferencer.py b/mmocr/apis/inferencers/mmocr_inferencer.py index 826f4f8be..a8b7bba47 100644 --- a/mmocr/apis/inferencers/mmocr_inferencer.py +++ b/mmocr/apis/inferencers/mmocr_inferencer.py @@ -46,7 +46,9 @@ def __init__(self, self.visualizer = VISUALIZERS.build( dict( type='TextSpottingLocalVisualizer', - name=f'inferencer{ts}')) + name=f'inferencer{ts}', + font_families=self.textrec_inferencer.visualizer. + font_families)) else: self.mode = 'rec' if kie_config is not None: diff --git a/mmocr/visualization/base_visualizer.py b/mmocr/visualization/base_visualizer.py index 1501c6cb9..b53972cd1 100644 --- a/mmocr/visualization/base_visualizer.py +++ b/mmocr/visualization/base_visualizer.py @@ -1,6 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. import math -from typing import Sequence, Union +from typing import List, Sequence, Union import numpy as np import torch @@ -50,13 +50,22 @@ class BaseLocalVisualizer(Visualizer): (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', + **kwargs) -> None: + super().__init__(name=name, **kwargs) + self.font_families = font_families + 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) -> np.ndarray: + auto_font_size: bool = False, + font_families: Union[str, List[str]] = 'sans-serif' + ) -> np.ndarray: """Draw labels on image. Args: @@ -73,6 +82,8 @@ def get_labels_image(self, 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'. """ if colors is not None and isinstance(colors, (list, tuple)): size = math.ceil(len(labels) / len(colors)) @@ -88,7 +99,8 @@ def get_labels_image(self, vertical_alignments='center', horizontal_alignments='center', colors='k', - font_sizes=font_size) + font_sizes=font_size, + font_families=font_families) return self.get_image() def get_polygons_image(self, diff --git a/mmocr/visualization/kie_visualizer.py b/mmocr/visualization/kie_visualizer.py index b29cceb95..fb69e3440 100644 --- a/mmocr/visualization/kie_visualizer.py +++ b/mmocr/visualization/kie_visualizer.py @@ -38,7 +38,7 @@ class KIELocalVisualizer(BaseLocalVisualizer): """ def __init__(self, - name='kie_visualizer', + name: str = 'kie_visualizer', is_openset: bool = False, **kwargs) -> None: super().__init__(name=name, **kwargs) @@ -89,13 +89,15 @@ def _draw_edge_label(self, key_texts, (bboxes[key_index, :2] + bboxes[key_index, 2:]) / 2, colors='k', horizontal_alignments='center', - vertical_alignments='center') + vertical_alignments='center', + font_families=self.font_families) if val_texts: self.draw_texts( val_texts, (bboxes[val_index, :2] + bboxes[val_index, 2:]) / 2, colors='k', horizontal_alignments='center', - vertical_alignments='center') + vertical_alignments='center', + font_families=self.font_families) self.draw_arrows( x_data, y_data, @@ -150,12 +152,16 @@ def _draw_instances( empty_shape = (img_shape[0], img_shape[1], 3) text_image = np.full(empty_shape, 255, dtype=np.uint8) - text_image = self.get_labels_image(text_image, texts, bboxes) + text_image = self.get_labels_image( + text_image, texts, bboxes, font_families=self.font_families) classes_image = np.full(empty_shape, 255, dtype=np.uint8) bbox_classes = [class_names[int(i)]['name'] for i in bbox_labels] - classes_image = self.get_labels_image(classes_image, bbox_classes, - bboxes) + classes_image = self.get_labels_image( + classes_image, + bbox_classes, + bboxes, + font_families=self.font_families) if polygons: polygons = [polygon.reshape(-1, 2) for polygon in polygons] image = self.get_polygons_image( diff --git a/mmocr/visualization/textrecog_visualizer.py b/mmocr/visualization/textrecog_visualizer.py index 623bf7642..9cbb2c57b 100644 --- a/mmocr/visualization/textrecog_visualizer.py +++ b/mmocr/visualization/textrecog_visualizer.py @@ -36,13 +36,14 @@ def __init__(self, vis_backends: Optional[Dict] = None, save_dir: Optional[str] = None, gt_color: Optional[Union[str, Tuple[int, int, int]]] = 'g', - pred_color: Optional[Union[str, Tuple[int, int, - int]]] = 'r') -> None: + pred_color: Optional[Union[str, Tuple[int, int, int]]] = 'r', + **kwargs) -> None: super().__init__( name=name, image=image, vis_backends=vis_backends, - save_dir=save_dir) + save_dir=save_dir, + **kwargs) self.gt_color = gt_color self.pred_color = pred_color @@ -66,7 +67,8 @@ def _draw_instances(self, image: np.ndarray, text: str) -> np.ndarray: colors=self.gt_color, font_sizes=font_size, vertical_alignments='center', - horizontal_alignments='center') + horizontal_alignments='center', + font_families=self.font_families) text_image = self.get_image() return text_image diff --git a/mmocr/visualization/textspotting_visualizer.py b/mmocr/visualization/textspotting_visualizer.py index 19a5e4ad3..6371b063f 100644 --- a/mmocr/visualization/textspotting_visualizer.py +++ b/mmocr/visualization/textspotting_visualizer.py @@ -45,7 +45,10 @@ def _draw_instances( empty_shape = (img_shape[0], img_shape[1], 3) text_image = np.full(empty_shape, 255, dtype=np.uint8) text_image = self.get_labels_image( - text_image, labels=texts, bboxes=bboxes) + text_image, + labels=texts, + bboxes=bboxes, + font_families=self.font_families) if polygons: polygons = [polygon.reshape(-1, 2) for polygon in polygons] image = self.get_polygons_image(