| import numpy as np |
| import matplotlib.pyplot as plt |
| from matplotlib.patches import Rectangle |
|
|
| COLORS = [ |
| "#003EFF", |
| "#FF8F00", |
| "#079700", |
| "#A123FF", |
| "#87CEEB", |
| "#FF5733", |
| "#C70039", |
| "#900C3F", |
| "#581845", |
| "#11998E", |
| ] |
|
|
|
|
| def reformat_for_plotting(labels, bboxes, scores, shape, num_classes): |
| """ |
| Reformat YOLOX predictions for plotting. |
| |
| Args: |
| labels (np.ndarray): Array of labels. |
| bboxes (np.ndarray): Array of bounding boxes. |
| scores (np.ndarray): Array of confidence scores. |
| shape (tuple): Shape of the image. |
| num_classes (int): Number of classes. |
| |
| Returns: |
| list[np.ndarray]: List of box bounding boxes per class. |
| list[np.ndarray]: List of confidence scores per class. |
| """ |
| boxes_plot = bboxes.copy() |
| boxes_plot[:, [0, 2]] *= shape[1] |
| boxes_plot[:, [1, 3]] *= shape[0] |
| boxes_plot = boxes_plot.astype(int) |
| boxes_plot[:, 2] -= boxes_plot[:, 0] |
| boxes_plot[:, 3] -= boxes_plot[:, 1] |
| boxes_plot = [boxes_plot[labels == c] for c in range(num_classes)] |
| confs = [scores[labels == c] for c in range(num_classes)] |
| return boxes_plot, confs |
|
|
|
|
| def plot_sample(img, boxes_list, confs_list, labels): |
| """ |
| Plots an image with bounding boxes. |
| Coordinates are expected in format [x_min, y_min, width, height]. |
| |
| Args: |
| img (numpy.ndarray): The input image to be plotted. |
| boxes_list (list[np.ndarray]): List of box bounding boxes per class. |
| confs_list (list[np.ndarray]): List of confidence scores per class. |
| labels (list): List of class labels. |
| """ |
| plt.imshow(img, cmap="gray") |
| plt.axis(False) |
|
|
| for boxes, confs, col, l in zip(boxes_list, confs_list, COLORS, labels): |
| for box_idx, box in enumerate(boxes): |
| |
| h, w, _ = img.shape |
| box = np.copy(box) |
| box[:2] = np.clip(box[:2], 2, max(h, w)) |
| box[2] = min(box[2], w - 2 - box[0]) |
| box[3] = min(box[3], h - 2 - box[1]) |
|
|
| rect = Rectangle( |
| (box[0], box[1]), |
| box[2], |
| box[3], |
| linewidth=2, |
| facecolor="none", |
| edgecolor=col, |
| ) |
| plt.gca().add_patch(rect) |
|
|
| |
| plt.text( |
| box[0], box[1], |
| f"{l}_{box_idx} conf={confs[box_idx]:.3f}", |
| color='white', |
| fontsize=8, |
| bbox=dict(facecolor=col, alpha=1, edgecolor=col, pad=0, linewidth=2), |
| verticalalignment='bottom', |
| horizontalalignment='left' |
| ) |
|
|
|
|
| def postprocess_preds_page_element(preds, thresholds_per_class, class_labels): |
| """ |
| Post process predictions for the page element task. |
| - Applies thresholding |
| |
| Args: |
| preds (dict): Predictions. Keys are "scores", "boxes", "labels". |
| thresholds_per_class (dict): Thresholds per class. |
| labels (list): List of class labels. |
| |
| Returns: |
| labels (numpy.ndarray): Array of labels. |
| bboxes (numpy.ndarray): Array of bounding boxes. |
| scores (numpy.ndarray): Array of scores. |
| """ |
| labels = preds["labels"].cpu().numpy() |
| boxes = preds["boxes"].cpu().numpy() |
| scores = preds["scores"].cpu().numpy() |
|
|
| |
| thresholds = np.array( |
| [thresholds_per_class[class_labels[int(x)]] for x in labels] |
| ) |
| labels = labels[scores > thresholds] |
| boxes = boxes[scores > thresholds] |
| scores = scores[scores > thresholds] |
|
|
| return labels, boxes, scores |
|
|
