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