draw_BBox.py

import cv2
import random
import numpy as np

'''
    Draw AABB in given image with labels
    Inputs: img - image file path or cv2 image
          : boxes_list - 2d list or 2d array of BBox (4 columns) with bounding boxes with top-left (x_min, y_min) and bottom right (x_max, y_max) pixel value. PS: maintain order xyxy
          : label list - list(str) or 1-d array having name of the labels. Use cls = [classes[l] for l in list(labels)] # index to str labels
    Return: CV2 image in BGR format
''' 
def draw_boxes(img, boxes_list, labels_list):
    obj_color={}
    if isinstance(img, str):  #If image path, then read image first
        assert os.path.isfile(img), "Image path doesn't exist"
        img = cv2.imread(img)
        # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
    for label in labels_list: # Create BBox color for each object category
        if label not in obj_color.keys():
            obj_color[label]=(random.randint(0,255),random.randint(0,255),random.randint(0,255))
    
    for i,(box,label) in enumerate(zip(boxes_list,labels_list)):
        if np.size(box) ==4:
            plot_one_bbox(img, box, obj_color[label], label)
    cv2.imwrite('output.png',img)
    return img

def plot_one_bbox(img, box, color=None, label=None, line_thickness=None):
    tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2)  # line/font thickness
    color = color or [random.randint(0, 255) for _ in range(3)]
    c1, c2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[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]
        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)

Input Image

Output Image