what.models.detection.utils.box_utils API documentation

# what.models.detection.utils.box_utils

View Source

     1import cv2
     2import numpy as np
     3from .array_utils import to_numpy
     4
     5def draw_bounding_boxes(image, boxes, labels, class_names, probs):
     6    if len(boxes) > 0:
     7        assert(boxes.shape[1] == 4)
     8        boxes = to_numpy(boxes)
     9
    10    # (x, y, w, h) --> (x1, y1, x2, y2)
    11    height, width, _ = image.shape
    12    for box in boxes:
    13        box[0] *= width
    14        box[1] *= height
    15        box[2] *= width
    16        box[3] *= height
    17
    18        # From center to top left
    19        box[0] -= box[2] / 2
    20        box[1] -= box[3] / 2
    21
    22        # From width and height to x2 and y2
    23        box[2] += box[0]
    24        box[3] += box[1]
    25
    26    # Draw bounding boxes and labels
    27    for i in range(boxes.shape[0]):
    28        box = boxes[i]
    29        label = f"{class_names[labels[i]]}: {probs[i]:.2f}"
    30        # print(label)
    31
    32        # Draw bounding boxes
    33        cv2.rectangle(image, (int(box[0].item()), int(box[1].item())), (int(box[2].item()), int(box[3].item())), (255, 255, 0), 4)
    34
    35        # Draw labels
    36        cv2.putText(image, label,
    37                    (int(box[0]+20), int(box[1]+40)),
    38                    cv2.FONT_HERSHEY_SIMPLEX,
    39                    1,  # font scale
    40                    (255, 0, 255),
    41                    2)  # line type
    42    return image

def draw_bounding_boxes(image, boxes, labels, class_names, probs): View Source

     6def draw_bounding_boxes(image, boxes, labels, class_names, probs):
     7    if len(boxes) > 0:
     8        assert(boxes.shape[1] == 4)
     9        boxes = to_numpy(boxes)
    10
    11    # (x, y, w, h) --> (x1, y1, x2, y2)
    12    height, width, _ = image.shape
    13    for box in boxes:
    14        box[0] *= width
    15        box[1] *= height
    16        box[2] *= width
    17        box[3] *= height
    18
    19        # From center to top left
    20        box[0] -= box[2] / 2
    21        box[1] -= box[3] / 2
    22
    23        # From width and height to x2 and y2
    24        box[2] += box[0]
    25        box[3] += box[1]
    26
    27    # Draw bounding boxes and labels
    28    for i in range(boxes.shape[0]):
    29        box = boxes[i]
    30        label = f"{class_names[labels[i]]}: {probs[i]:.2f}"
    31        # print(label)
    32
    33        # Draw bounding boxes
    34        cv2.rectangle(image, (int(box[0].item()), int(box[1].item())), (int(box[2].item()), int(box[3].item())), (255, 255, 0), 4)
    35
    36        # Draw labels
    37        cv2.putText(image, label,
    38                    (int(box[0]+20), int(box[1]+40)),
    39                    cv2.FONT_HERSHEY_SIMPLEX,
    40                    1,  # font scale
    41                    (255, 0, 255),
    42                    2)  # line type
    43    return image