"server/vscode:/vscode.git/clone" did not exist on "ba291dad9f5cd0061e5ee4cab332db42c434e369"
Commit 97f9febd authored by Kai Chen's avatar Kai Chen
Browse files

add a method to draw bboxes and labels

parent c3f87f48
from .color import Color, color_val
from .image import imshow, imshow_bboxes
from .image import imshow, imshow_bboxes, imshow_det_bboxes
from .optflow import flowshow, flow2rgb, make_color_wheel
__all__ = [
'Color', 'color_val', 'imshow', 'imshow_bboxes', 'flowshow', 'flow2rgb',
'make_color_wheel'
'Color', 'color_val', 'imshow', 'imshow_bboxes', 'imshow_det_bboxes',
'flowshow', 'flow2rgb', 'make_color_wheel'
]
......@@ -64,3 +64,68 @@ def imshow_bboxes(img,
imshow(img, win_name, wait_time)
if out_file is not None:
imwrite(img, out_file)
def imshow_det_bboxes(img,
bboxes,
labels,
class_names=None,
score_thr=0,
bbox_color='green',
text_color='green',
thickness=1,
font_scale=0.5,
show=True,
win_name='',
wait_time=0,
out_file=None):
"""Draw bboxes and class labels (with scores) on an image.
Args:
img (str or ndarray): The image to be displayed.
bboxes (ndarray): Bounding boxes (with scores), shaped (n, 4) or (n, 5).
labels (ndarray): Labels of bboxes.
class_names (list[str]): Names of each classes.
score_thr (float): Minimum score of bboxes to be shown.
bbox_color (str or tuple or :obj:`Color`): Color of bbox lines.
text_color (str or tuple or :obj:`Color`): Color of texts.
thickness (int): Thickness of lines.
font_scale (float): Font scales of texts.
show (bool): Whether to show the image.
win_name (str): The window name.
wait_time (int): Value of waitKey param.
out_file (str or None): The filename to write the image.
"""
assert bboxes.ndim == 2
assert labels.ndim == 1
assert bboxes.shape[0] == labels.shape[0]
assert bboxes.shape[1] == 4 or bboxes.shape[1] == 5
img = imread(img)
if score_thr > 0:
assert bboxes.shape[1] == 5
scores = bboxes[:, -1]
inds = scores > score_thr
bboxes = bboxes[inds, :]
labels = labels[inds, :]
bbox_color = color_val(bbox_color)
text_color = color_val(text_color)
for bbox, label in zip(bboxes, labels):
bbox_int = bbox.astype(np.int32)
left_top = (bbox_int[0], bbox_int[1])
right_bottom = (bbox_int[2], bbox_int[3])
cv2.rectangle(
img, left_top, right_bottom, bbox_color, thickness=thickness)
label_text = class_names[
label] if class_names is not None else 'cls {}'.format(label)
if bbox.shape[1] > 4:
label_text += '|{:.02f}'.format(bbox[-1])
cv2.putText(img, label_text, (bbox_int[0], bbox_int[1] - 2),
cv2.FONT_HERSHEY_COMPLEX, font_scale, text_color)
if show:
imshow(img, win_name, wait_time)
if out_file is not None:
imwrite(img, out_file)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment