不知道大家在可视化检测结果时是否会有设置的线宽太小在图片上画出来的框看不清楚,写的字符串看不清楚的烦恼。我发现不论是什么大小的图片,使用yolov5可视化线宽和字符串都能很清晰的显示出来,我根据yolov5上面的可视化代码,改写了这个draw_box
函数能实现自适应图片大小选择线宽和字体大小,从此不需要调整线宽和字体大小了;
def get_color(index):index = int(index) ** 11 + 11# L = R * 299 / 1000 + G * 587 / 1000 + B * 114 / 1000return index * 299 % 255, index * 587 % 255, index * 114 % 255def draw_box(image_bgr, preds, names):"""@param image_bgr: 原图@param preds: 预测结果,[[x1, y1, x2, y2, conf, cls], ...]@param names: 对应preds的类别名[...]@return:"""if len(preds) == 0:return image_bgr# 自适应图片大小计算线宽和文字大小# 图片大小img_h, img_w = image_bgr.shape[:2]# box线宽box_thickness = max(round(sum((img_h, img_w)) / 2 * 0.003), 2)# font_size = max(round(sum((img_h, img_w)) / 2 * 0.035), 12)# text线宽text_thickness = max(box_thickness - 1, 1)# text字体大小fontScale = box_thickness / 4. # NOTE:如果嫌姿态太大把这个4修改为更大值就行,反之字体太小就改小# 字体fontFace = 0# 画框for i, (x1, y1, x2, y2, conf, cls) in enumerate(preds):x1, y1, x2, y2, cls = np.array([x1, y1, x2, y2, cls], np.int16)# 画框cv2.rectangle(image_bgr, pt1=(x1, y1),pt2=(x2, y2),color=get_color(cls),thickness=box_thickness, lineType=cv2.LINE_AA)# 在框上显示的texttext = f"{names[i]}:{conf:.2f}"# 用于计算特定文本字符串在给定字体和大小下的尺寸text_w, text_h = cv2.getTextSize(text=text,fontFace=fontFace,fontScale=fontScale,thickness=text_thickness)[0]# 防止写的text超过了上边界,导致看不到,如果能写外面就写外面outside_x = img_w - x1 - text_w # 如果横着会导致文本超过图像边界,text向左挪一挪outside_y = y1 - text_h >= 5# 写文字cv2.putText(image_bgr, text=text,org=(x1 if outside_x >= 5 else x1 + outside_x - 5, # 防止超过右边界y1 - 3 if outside_y else y1 + text_h + 5), # 防止超过上边界fontFace=fontFace, fontScale=fontScale, thickness=text_thickness,color=get_color(cls))return image_bgr