#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File : box.py
# Author : Jiayuan Mao
# Email : maojiayuan@gmail.com
# Date : 03/03/2018
#
# This file is part of Jacinle.
# Distributed under terms of the MIT license.
import itertools
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from typing import List, Union, Optional
__all__ = ['add_bbox_patches', 'vis_bboxes']
[docs]
def add_bbox_patches(ax, boxes: np.ndarray, class_name: Union[str, List[str]], add_text=True, legends: Optional[Union[str, List[str]]] = None, fontsize: Optional[int] = 14):
if legends is not None:
assert len(legends) == len(boxes)
else:
legends = ['' for i in range(len(boxes))]
if isinstance(class_name, str):
class_name = itertools.repeat(class_name, len(boxes))
for box, cls, leg in zip(boxes, class_name, legends):
ax.add_patch(plt.Rectangle(
(box[0], box[1]), box[2] - box[0], box[3] - box[1],
fill=False, edgecolor='green', linewidth=3.5
))
if add_text:
ax.text(
box[0], box[1] - 2,
'{:s}{}'.format(cls, str(leg)),
bbox=dict(facecolor='blue', alpha=0.5), fontsize=fontsize, color='white'
)
return ax
[docs]
def vis_bboxes(image: Image.Image, boxes: np.ndarray, class_name: Union[str, List[str]], add_text=True, legends: Optional[Union[str, List[str]]] = None, title: Optional[str] = None, fontsize: Optional[int] = 14):
fig, ax = plt.subplots(figsize=(12, int(12 / image.width * image.height)))
fig.tight_layout()
ax.imshow(image, aspect='equal')
if title is not None:
ax.set_title(title, fontsize=fontsize)
add_bbox_patches(ax, boxes, class_name, add_text=add_text, legends=legends, fontsize=fontsize)
ax.axis('off')
return fig, ax