首页 > 技术文章 > 基于YOLO3对图像加框的函数draw_image()

tangjunjun 2019-10-30 16:00 原文

 

 

def draw_bbox(image, bboxes, class_i, show_label=True):
# 将中心点坐标与w,h通过变化为左上角与右下角坐标
bboxes_change = np.copy(bboxes)
bboxes[:,0:2]=bboxes_change[:,0:2]-0.5*bboxes_change[:,2:4]
bboxes[:, 0:2] = bboxes_change[:, 0:2] + 0.5 * bboxes_change[:, 2:4]

"""
bboxes: [x_min, y_min, x_max, y_max] format coordinates.
"""
image_h, image_w, _ = image.shape
hsv_tuples = [(1.0 * x / 90, 1., 1.) for x in range(90)]
colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), colors))
for i, bbox in enumerate(bboxes):
coor = np.array(bbox[:4], dtype=np.int32)
fontScale = 0.5
class_ind = int(class_i)
bbox_color = colors[class_ind]
bbox_thick = int(0.6 * (image_h + image_w) / 600)
# bbox_thick = int(400 / 600)
c1, c2 = (coor[0], coor[1]), (coor[2], coor[3])
cv2.rectangle(image, c1, c2, bbox_color, bbox_thick) # 已经在图片上画了矩形

if show_label:
bbox_mess = '%s: %d' % ('class = ', class_ind)
t_size = cv2.getTextSize(bbox_mess, 0, fontScale, thickness=bbox_thick//2)[0]
cv2.rectangle(image, c1, (c1[0] + t_size[0], c1[1] - t_size[1] - 3), bbox_color, -1) # filled
cv2.putText(image, bbox_mess, (c1[0], c1[1]-2), cv2.FONT_HERSHEY_SIMPLEX, fontScale, (0, 0, 0), bbox_thick//2, lineType=cv2.LINE_AA) # 添加文字
# 图片 添加的文字 左上角坐标 字体 字体大小 颜色 字体粗细
return image

推荐阅读