Spaces:
Runtime error
Runtime error
File size: 3,916 Bytes
7bf08cb 4986f6d 7bf08cb 4986f6d 7bf08cb 4986f6d 7bf08cb 4986f6d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
# 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
|