Spaces:
Build error
Build error
File size: 5,825 Bytes
28c256d |
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 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
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_bboxes(img: Union[str, np.ndarray],
bboxes: Union[list, np.ndarray],
colors: ColorType = 'green',
top_k: int = -1,
thickness: int = 1,
show: bool = True,
win_name: str = '',
wait_time: int = 0,
out_file: Optional[str] = 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 (Color or str or tuple or int or ndarray): 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: 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 = 0.5,
show: bool = True,
win_name: str = '',
wait_time: int = 0,
out_file: Optional[str] = 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.
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
|