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
43    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
44    return image