python

Draw bounding box

Code to draw bbox

import matplotlib as plt
import numpy as np
import random
import cv2
import torch
#_______________________code draw predicted bbox from yolov7______________________________________#

#input
cls=torch.tensor(1.0)
conf= torch.tensor(1.0)
names=['blister_on_hand', 'blister']
# Create random color according to label
colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
box=[torch.tensor(100),torch.tensor(100),torch.tensor(100),torch.tensor(100)]
det=torch.tensor([[100,230,150,256,0.83,1.],
                  [300,130,450,256,0.83,1.],
                  [200,330,150,256,0.9,0.1],
                ])
label = f'{names[int(cls)]} {conf:.2f}'
im0=cv2.imread('/home/tonyhuy/yolov7/runs/detect/exp2/test_0.png')
# to get those input vars, loop over prediction
#function
def plot_one_box(x, img, color=None, label=None, line_thickness=3):
    # Plots one bounding box on image img
    tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1  # line/font thickness
    color = color or [random.randint(0, 255) for _ in range(3)]
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
    if label:
        tf = max(tl - 1, 1)  # font thickness
        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
        cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled
        cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
        cv2.imwrite(f'/home/tonyhuy/yolov7/draw_mask/img.jpg',img)
Output
draw bounding box
Was this helpful?