路径
:mmdetection/mmdet/core/visualization/image.py
更改的地方是在imshow_det_bboxes
这个函数中
更改后的image.py文件
# Copyright (c) OpenMMLab. All rights reserved.
import matplotlib.pyplot as plt
import mmcv
import numpy as np
import pycocotools.mask as mask_util
from matplotlib.collections import PatchCollection
from matplotlib.patches import Polygon
from ..utils import mask2ndarray
from .palette import get_palette, palette_val
EPS = 1e-2
def color_val_matplotlib(color):
"""Convert various input in BGR order to normalized RGB matplotlib color
tuples,
Args:
color (:obj:`Color`/str/tuple/int/ndarray): Color inputs
Returns:
tuple[float]: A tuple of 3 normalized floats indicating RGB channels.
"""
color = mmcv.color_val(color)
color = [color / 255 for color in color[::-1]]
return tuple(color)
def imshow_det_bboxes(img,
bboxes,
labels,
segms=None,
class_names=None,
score_thr=0,
bbox_color='green',
text_color='green',
mask_color=None,
thickness=2,
font_size=13,
win_name='',
show=True,
wait_time=0,
out_file=None):
"""Draw bboxes and class labels (with scores) on an image.
Args:
img (str | ndarray): The image to be displayed.
bboxes (ndarray): Bounding boxes (with scores), shaped (n, 4) or
(n, 5).
labels (ndarray): Labels of bboxes.
segms (ndarray | None): Masks, shaped (n,h,w) or None.
class_names (list[str]): Names of each classes.
score_thr (float): Minimum score of bboxes to be shown. Default: 0.
bbox_color (list[tuple] | tuple | str | None): Colors of bbox lines.
If a single color is given, it will be applied to all classes.
The tuple of color should be in RGB order. Default: 'green'.
text_color (list[tuple] | tuple | str | None): Colors of texts.
If a single color is given, it will be applied to all classes.
The tuple of color should be in RGB order. Default: 'green'.
mask_color (list[tuple] | tuple | str | None, optional): Colors of
masks. If a single color is given, it will be applied to all
classes. The tuple of color should be in RGB order.
Default: None.
thickness (int): Thickness of lines. Default: 2.
font_size (int): Font size of texts. Default: 13.
show (bool): Whether to show the image. Default: True.
win_name (str): The window name. Default: ''.
wait_time (float): Value of waitKey param. Default: 0.
out_file (str, optional): The filename to write the image.
Default: None.
Returns:
ndarray: The image with bboxes drawn on it.
"""
assert bboxes.ndim == 2, \
f' bboxes ndim should be 2, but its ndim is {bboxes.ndim}.'
assert labels.ndim == 1, \
f' labels ndim should be 1, but its ndim is {labels.ndim}.'
assert bboxes.shape[0] == labels.shape[0], \
'bboxes.shape[0] and labels.shape[0] should have the same length.'
assert bboxes.shape[1] == 4 or bboxes.shape[1] == 5, \
f' bboxes.shape[1] should be 4 or 5, but its {bboxes.shape[1]}.'
img = mmcv.imread(img).astype(np.uint8)
if score_thr > 0:
assert bboxes.shape[1] == 5
scores = bboxes[:, -1]
inds = scores > score_thr
bboxes = bboxes[inds, :]
labels = labels[inds]
if segms is not None:
segms = segms[inds, ...]
max_label = int(max(labels)) if labels.shape[0] > 0 else -1
bbox_color = palette_val(get_palette(bbox_color, max_label + 1))
text_color = palette_val(get_palette(text_color, max_label + 1))
mask_color = get_palette(mask_color, max_label + 1)
mask_color = np.array(mask_color, dtype=np.uint8)
img = mmcv.bgr2rgb(img)
width, height = img.shape[1], img.shape[0]
img = np.ascontiguousarray(img)
import random
bbox_color = [[random.randint(0, 255)/255 for _ in range(3)] for _ in range(10)]
fig = plt.figure(win_name, frameon=False)
plt.title(win_name)
canvas = fig.canvas
dpi = fig.get_dpi()
# add a small EPS to avoid precision lost due to matplotlib's truncation
# (https://github.com/matplotlib/matplotlib/issues/15363)
fig.set_size_inches((width + EPS) / dpi, (height + EPS) / dpi)
# remove white edges by set subplot margin
plt.subplots_adjust(left=0, right=1, bottom=0, top=1)
ax = plt.gca()
ax.axis('off')
polygons = []
color = []
for i, (bbox, label) in enumerate(zip(bboxes, labels)):
bbox_int = bbox.astype(np.int32)
poly = [[bbox_int[0], bbox_int[1]], [bbox_int[0], bbox_int[3]],
[bbox_int[2], bbox_int[3]], [bbox_int[2], bbox_int[1]]]
np_poly = np.array(poly).reshape((4, 2))
polygons.append(Polygon(np_poly))
color.append(bbox_color[label])
label_text = class_names[
label] if class_names is not None else f'class {label}'
if len(bbox) > 4:
label_text += f'|{bbox[-1]:.02f}'
# ax.text(
# bbox_int[0],
# bbox_int[1],
# f'{label_text}',
# bbox={
# 'facecolor': 'black',
# 'alpha': 0.8,
# 'pad': 0.7,
# 'edgecolor': 'none'
# },
# color=text_color[label],
# fontsize=font_size,
# verticalalignment='top',
# horizontalalignment='left')
if segms is not None:
color_mask = mask_color[labels[i]]
mask = segms[i].astype(bool)
img[mask] = img[mask] * 0.5 + color_mask * 0.5
plt.imshow(img)
p = PatchCollection(
polygons, facecolor='none', edgecolors=color, linewidths=thickness)
ax.add_collection(p)
stream, _ = canvas.print_to_buffer()
buffer = np.frombuffer(stream, dtype='uint8')
img_rgba = buffer.reshape(height, width, 4)
rgb, alpha = np.split(img_rgba, [3], axis=2)
img = rgb.astype('uint8')
img = mmcv.rgb2bgr(img)
if show:
# We do not use cv2 for display because in some cases, opencv will
# conflict with Qt, it will output a warning: Current thread
# is not the object's thread. You can refer to
# https://github.com/opencv/opencv-python/issues/46 for details
if wait_time == 0:
plt.show()
else:
plt.show(block=False)
plt.pause(wait_time)
if out_file is not None:
mmcv.imwrite(img, out_file)
plt.close()
return img
def imshow_gt_det_bboxes(img,
annotation,
result,
class_names=None,
score_thr=0,
gt_bbox_color=(255, 102, 61),
gt_text_color=(255, 102, 61),
gt_mask_color=(255, 102, 61),
det_bbox_color=(72, 101, 241),
det_text_color=(72, 101, 241),
det_mask_color=(72, 101, 241),
thickness=2,
font_size=13,
win_name='',
show=True,
wait_time=0,
out_file=None):
"""General visualization GT and result function.
Args:
img (str | ndarray): The image to be displayed.
annotation (dict): Ground truth annotations where contain keys of
'gt_bboxes' and 'gt_labels' or 'gt_masks'.
result (tuple[list] | list): The detection result, can be either
(bbox, segm) or just bbox.
class_names (list[str]): Names of each classes.
score_thr (float): Minimum score of bboxes to be shown. Default: 0.
gt_bbox_color (list[tuple] | tuple | str | None): Colors of bbox lines.
If a single color is given, it will be applied to all classes.
The tuple of color should be in RGB order. Default: (255, 102, 61).
gt_text_color (list[tuple] | tuple | str | None): Colors of texts.
If a single color is given, it will be applied to all classes.
The tuple of color should be in RGB order. Default: (255, 102, 61).
gt_mask_color (list[tuple] | tuple | str | None, optional): Colors of
masks. If a single color is given, it will be applied to all classes.
The tuple of color should be in RGB order. Default: (255, 102, 61).
det_bbox_color (list[tuple] | tuple | str | None):Colors of bbox lines.
If a single color is given, it will be applied to all classes.
The tuple of color should be in RGB order. Default: (72, 101, 241).
det_text_color (list[tuple] | tuple | str | None):Colors of texts.
If a single color is given, it will be applied to all classes.
The tuple of color should be in RGB order. Default: (72, 101, 241).
det_mask_color (list[tuple] | tuple | str | None, optional): Color of
masks. If a single color is given, it will be applied to all classes.
The tuple of color should be in RGB order. Default: (72, 101, 241).
thickness (int): Thickness of lines. Default: 2.
font_size (int): Font size of texts. Default: 13.
win_name (str): The window name. Default: ''.
show (bool): Whether to show the image. Default: True.
wait_time (float): Value of waitKey param. Default: 0.
out_file (str, optional): The filename to write the image.
Default: None.
Returns:
ndarray: The image with bboxes or masks drawn on it.
"""
assert 'gt_bboxes' in annotation
assert 'gt_labels' in annotation
assert isinstance(
result,
(tuple, list)), f'Expected tuple or list, but get {type(result)}'
gt_masks = annotation.get('gt_masks', None)
if gt_masks is not None:
gt_masks = mask2ndarray(gt_masks)
img = mmcv.imread(img)
img = imshow_det_bboxes(
img,
annotation['gt_bboxes'],
annotation['gt_labels'],
gt_masks,
class_names=class_names,
bbox_color=gt_bbox_color,
text_color=gt_text_color,
mask_color=gt_mask_color,
thickness=thickness,
font_size=font_size,
win_name=win_name,
show=False)
if isinstance(result, tuple):
bbox_result, segm_result = result
if isinstance(segm_result, tuple):
segm_result = segm_result[0] # ms rcnn
else:
bbox_result, segm_result = result, None
bboxes = np.vstack(bbox_result)
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
segms = None
if segm_result is not None and len(labels) > 0: # non empty
segms = mmcv.concat_list(segm_result)
segms = mask_util.decode(segms)
segms = segms.transpose(2, 0, 1)
img = imshow_det_bboxes(
img,
bboxes,
labels,
segms=segms,
class_names=class_names,
score_thr=score_thr,
bbox_color=det_bbox_color,
text_color=det_text_color,
mask_color=det_mask_color,
thickness=thickness,
font_size=font_size,
win_name=win_name,
show=show,
wait_time=wait_time,
out_file=out_file)
return img
更改说明
添加不同类别显示不同颜色的代码
import random
bbox_color = [[random.randint(0, 255)/255 for _ in range(3)] for _ in range(10)]
注意:bbox_color 要覆盖掉之前的bbox_color,详情可以看更改后的代码
去掉标签和置信度
将imshow_det_bboxes
中下面的代码注释掉即可。
# ax.text(
# bbox_int[0],
# bbox_int[1],
# f'{label_text}',
# bbox={
# 'facecolor': 'black',
# 'alpha': 0.8,
# 'pad': 0.7,
# 'edgecolor': 'none'
# },
# color=text_color[label],
# fontsize=font_size,
# verticalalignment='top',
# horizontalalignment='left')