Skip to content

Commit

Permalink
[Style] add font_families argument to fn
Browse files Browse the repository at this point in the history
  • Loading branch information
ProtossDragoon committed Dec 7, 2022
1 parent 19810d6 commit 7a9be75
Show file tree
Hide file tree
Showing 3 changed files with 24 additions and 8 deletions.
8 changes: 6 additions & 2 deletions mmocr/visualization/base_visualizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,9 @@ def get_labels_image(self,
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:
Expand All @@ -80,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))
Expand All @@ -96,7 +100,7 @@ def get_labels_image(self,
horizontal_alignments='center',
colors='k',
font_sizes=font_size,
font_families=self.font_families)
font_families=font_families)
return self.get_image()

def get_polygons_image(self,
Expand Down
19 changes: 14 additions & 5 deletions mmocr/visualization/kie_visualizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -152,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(
Expand All @@ -176,8 +180,13 @@ def _draw_instances(
cat_image = [image, text_image, classes_image]
if is_openset:
edge_image = np.full(empty_shape, 255, dtype=np.uint8)
edge_image = self._draw_edge_label(edge_image, edge_labels, bboxes,
texts, arrow_colors)
edge_image = self._draw_edge_label(
edge_image,
edge_labels,
bboxes,
texts,
arrow_colors,
font_families=self.font_families)
cat_image.append(edge_image)
return self._cat_image(cat_image, axis=1)

Expand Down
5 changes: 4 additions & 1 deletion mmocr/visualization/textspotting_visualizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down

0 comments on commit 7a9be75

Please sign in to comment.