|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from annotator.uniformer.mmcv.image import imread, imwrite
|
|
from .color import color_val
|
|
|
|
|
|
def imshow(img, win_name='', wait_time=0):
|
|
"""Show an image.
|
|
|
|
Args:
|
|
img (str or ndarray): The image to be displayed.
|
|
win_name (str): The window name.
|
|
wait_time (int): Value of waitKey param.
|
|
"""
|
|
cv2.imshow(win_name, imread(img))
|
|
if wait_time == 0:
|
|
while True:
|
|
ret = cv2.waitKey(1)
|
|
|
|
closed = cv2.getWindowProperty(win_name, cv2.WND_PROP_VISIBLE) < 1
|
|
|
|
if closed or ret != -1:
|
|
break
|
|
else:
|
|
ret = cv2.waitKey(wait_time)
|
|
|
|
|
|
def imshow_bboxes(img,
|
|
bboxes,
|
|
colors='green',
|
|
top_k=-1,
|
|
thickness=1,
|
|
show=True,
|
|
win_name='',
|
|
wait_time=0,
|
|
out_file=None):
|
|
"""Draw bboxes on an image.
|
|
|
|
Args:
|
|
img (str or ndarray): The image to be displayed.
|
|
bboxes (list or ndarray): A list of ndarray of shape (k, 4).
|
|
colors (list[str or tuple or Color]): A list of colors.
|
|
top_k (int): Plot the first k bboxes only if set positive.
|
|
thickness (int): Thickness of lines.
|
|
show (bool): Whether to show the image.
|
|
win_name (str): The window name.
|
|
wait_time (int): Value of waitKey param.
|
|
out_file (str, optional): The filename to write the image.
|
|
|
|
Returns:
|
|
ndarray: The image with bboxes drawn on it.
|
|
"""
|
|
img = imread(img)
|
|
img = np.ascontiguousarray(img)
|
|
|
|
if isinstance(bboxes, np.ndarray):
|
|
bboxes = [bboxes]
|
|
if not isinstance(colors, list):
|
|
colors = [colors for _ in range(len(bboxes))]
|
|
colors = [color_val(c) for c in colors]
|
|
assert len(bboxes) == len(colors)
|
|
|
|
for i, _bboxes in enumerate(bboxes):
|
|
_bboxes = _bboxes.astype(np.int32)
|
|
if top_k <= 0:
|
|
_top_k = _bboxes.shape[0]
|
|
else:
|
|
_top_k = min(top_k, _bboxes.shape[0])
|
|
for j in range(_top_k):
|
|
left_top = (_bboxes[j, 0], _bboxes[j, 1])
|
|
right_bottom = (_bboxes[j, 2], _bboxes[j, 3])
|
|
cv2.rectangle(
|
|
img, left_top, right_bottom, colors[i], thickness=thickness)
|
|
|
|
if show:
|
|
imshow(img, win_name, wait_time)
|
|
if out_file is not None:
|
|
imwrite(img, out_file)
|
|
return img
|
|
|
|
|
|
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.
|
|
|
|
Returns:
|
|
ndarray: The image with bboxes drawn on it.
|
|
"""
|
|
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)
|
|
img = np.ascontiguousarray(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 f'cls {label}'
|
|
if len(bbox) > 4:
|
|
label_text += f'|{bbox[-1]:.02f}'
|
|
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)
|
|
return img
|
|
|
