File size: 4,047 Bytes
b08d258 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | # ------------------------------------------------------------------------
# RF-DETR
# Copyright (c) 2025 Roboflow. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
from pathlib import Path
import numpy as np
from PIL import Image
from rfdetr.utilities.logger import get_logger
logger = get_logger()
def save_gt_predictions_visualization(
scenario_name: str,
image_width: int,
image_height: int,
gt_boxes: list[list[float]],
gt_class_ids: list[int],
pred_boxes: list[list[float]],
pred_class_ids: list[int],
pred_confidences: list[float],
pred_ious: list[float | None],
save_dir: Path,
) -> None:
"""
Save a visualization image showing both GT and prediction boxes.
Boxes are labeled with class ID and confidence (for predictions).
For predictions with known IoU, the IoU value is also shown.
"""
import supervision as sv
save_dir.mkdir(exist_ok=True)
top_padding = 60
image = np.zeros((image_height + top_padding, image_width, 3), dtype=np.uint8)
gt_boxes_offset = [[x, y + top_padding, w, h] for x, y, w, h in gt_boxes]
pred_boxes_offset = [[x, y + top_padding, w, h] for x, y, w, h in pred_boxes]
gt_xyxy = sv.xywh_to_xyxy(np.array(gt_boxes_offset))
pred_xyxy = sv.xywh_to_xyxy(np.array(pred_boxes_offset))
gt_detections = None
pred_detections = None
if len(gt_xyxy) > 0:
gt_detections = sv.Detections(
xyxy=gt_xyxy,
class_id=np.array(gt_class_ids),
)
if len(pred_xyxy) > 0:
pred_detections = sv.Detections(
xyxy=pred_xyxy,
class_id=np.array(pred_class_ids),
confidence=np.array(pred_confidences),
)
# Index 0 is unused because class IDs start at 1
gt_colors = sv.ColorPalette(
[
sv.Color(128, 128, 128), # dummy color for index 0
sv.Color(0, 255, 100),
sv.Color(0, 200, 255),
]
)
pred_colors = sv.ColorPalette(
[
sv.Color(128, 128, 128), # dummy color for index 0
sv.Color(255, 100, 50),
sv.Color(255, 50, 200),
]
)
gt_box_annotator = sv.BoxAnnotator(color=gt_colors, thickness=3, color_lookup=sv.ColorLookup.CLASS)
pred_box_annotator = sv.BoxAnnotator(color=pred_colors, thickness=3, color_lookup=sv.ColorLookup.CLASS)
gt_label_annotator = sv.LabelAnnotator(
color=gt_colors,
text_color=sv.Color.BLACK,
text_scale=0.5,
text_padding=3,
text_position=sv.Position.TOP_LEFT,
color_lookup=sv.ColorLookup.CLASS,
)
pred_label_annotator = sv.LabelAnnotator(
color=pred_colors,
text_color=sv.Color.BLACK,
text_scale=0.5,
text_padding=3,
text_position=sv.Position.TOP_RIGHT,
color_lookup=sv.ColorLookup.CLASS,
)
gt_labels = [f"c{class_id}" for class_id in gt_class_ids]
pred_labels = []
for class_id, conf, iou in zip(pred_class_ids, pred_confidences, pred_ious):
if iou is not None:
pred_labels.append(f"c{class_id}\nconf={conf:.3f}\niou={iou:.3f}")
else:
pred_labels.append(f"c{class_id}\nconf={conf:.3f}")
if gt_detections is not None:
image = gt_box_annotator.annotate(scene=image, detections=gt_detections)
image = gt_label_annotator.annotate(scene=image, detections=gt_detections, labels=gt_labels)
if pred_detections is not None:
image = pred_box_annotator.annotate(scene=image, detections=pred_detections)
image = pred_label_annotator.annotate(scene=image, detections=pred_detections, labels=pred_labels)
Image.fromarray(image).save(save_dir / f"{scenario_name}.png")
logger.info(f"Saved visualization to {save_dir}/{scenario_name}.png")
|