Spaces:
Runtime error
Runtime error
# THE ORIGINAL mmcv.imshow_det_bboxes | |
# Copyright (c) OpenMMLab. All rights reserved. | |
from typing import List, Optional, Union | |
import cv2 | |
import numpy as np | |
from mmcv.image import imread, imwrite | |
from .color import Color, color_val | |
# a type alias declares the optional types of color argument | |
ColorType = Union[Color, str, tuple, int, np.ndarray] | |
def imshow(img: Union[str, np.ndarray], win_name: str = "", wait_time: int = 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: # prevent from hanging if windows was closed | |
while True: | |
ret = cv2.waitKey(1) | |
closed = cv2.getWindowProperty(win_name, cv2.WND_PROP_VISIBLE) < 1 | |
# if user closed window or if some key pressed | |
if closed or ret != -1: | |
break | |
else: | |
ret = cv2.waitKey(wait_time) | |
def imshow_det_bboxes( | |
img: Union[str, np.ndarray], | |
bboxes: np.ndarray, | |
labels: np.ndarray, | |
class_names: List[str] = None, | |
score_thr: float = 0, | |
bbox_color: ColorType = "green", | |
text_color: ColorType = "green", | |
thickness: int = 1, | |
font_scale: float = 1, | |
show: bool = True, | |
win_name: str = "", | |
wait_time: int = 0, | |
out_file: Optional[str] = None, | |
colors: np.ndarray = 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 (Color or str or tuple or int or ndarray): Color | |
of bbox lines. | |
text_color (Color or str or tuple or int or ndarray): 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. | |
colors (array of tuple RGB int): the color of bbox and label of each class | |
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) | |
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]) | |
if colors is not None: | |
bbox_color = text_color = color_val(colors[label]) | |
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_TRIPLEX, | |
font_scale, | |
text_color, | |
4, | |
) | |
if show: | |
imshow(img, win_name, wait_time) | |
if out_file is not None: | |
imwrite(img, out_file) | |
return img | |