深度学习——自适应图片大小选择线宽和字体大小

不知道大家在可视化检测结果时是否会有设置的线宽太小在图片上画出来的框看不清楚，写的字符串看不清楚的烦恼。我发现不论是什么大小的图片，使用yolov5可视化线宽和字符串都能很清晰的显示出来，我根据yolov5上面的可视化代码，改写了这个draw_box函数能实现自适应图片大小选择线宽和字体大小，从此不需要调整线宽和字体大小了；

def get_color(index):
    index = int(index) ** 11 + 11
    # L = R * 299 / 1000 + G * 587 / 1000 + B * 114 / 1000
    return index * 299 % 255, index * 587 % 255, index * 114 % 255


def 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线宽，线宽最小值是2，如果感觉线宽太小将这个值增大即可
    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)
        # 在框上显示的text
        text = 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