what.models.detection.utils.box_utils

 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):
 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