From 6ccbdd2adf2644dafbd77872ffaf777a5e766dae Mon Sep 17 00:00:00 2001 From: "Dr. Cat" <47220664+catsmells@users.noreply.github.com> Date: Sun, 7 Sep 2025 13:42:28 -0400 Subject: [PATCH 1/2] Create prediction_visualizer.py Adding testing visualizer. --- rfdetr/util/prediction_visualizer.py | 56 ++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) create mode 100644 rfdetr/util/prediction_visualizer.py diff --git a/rfdetr/util/prediction_visualizer.py b/rfdetr/util/prediction_visualizer.py new file mode 100644 index 000000000..322f8214b --- /dev/null +++ b/rfdetr/util/prediction_visualizer.py @@ -0,0 +1,56 @@ +import argparse +import os +from typing import List +import cv2 +import supervision as sv +from PIL import Image +import numpy as np +from rfdetr import RFDETRBase + +def parse_args(): + parser = argparse.ArgumentParser(description="Visualize RF-DETR predictions on sample images.") + parser.add_argument("--weights", type=str, required=True, help="Path to pre-trained RF-DETR model weights.") + parser.add_argument("--input-dir", type=str, required=True, help="Directory containing input images.") + parser.add_argument("--output-dir", type=str, default="output", help="Directory to save annotated images.") + parser.add_argument("--confidence", type=float, default=0.5, help="Confidence threshold for predictions.") + return parser.parse_args() + +def load_images(input_dir: str) -> List[Image.Image]: + supported_extensions = (".jpg", ".jpeg", ".png") + return [Image.open(os.path.join(input_dir, f)) for f in os.listdir(input_dir) if f.lower().endswith(supported_extensions)] + +def save_annotated_image(image: Image.Image, detections: sv.Detections, output_path: str): + annotated_image = np.array(image) + annotated_image = sv.BoxAnnotator().annotate(annotated_image, detections) + labels = [f"{pred.class_name} {pred.confidence:.2f}" for pred in detections] + annotated_image = sv.LabelAnnotator().annotate(annotated_image, detections, labels) + cv2.imwrite(output_path, cv2.cvtColor(annotated_image, cv2.COLOR_RGB2BGR)) + +def main(): + args = parse_args() + + # Validate inputs + if not os.path.exists(args.input_dir): + raise ValueError(f"Input directory {args.input_dir} does not exist.") + os.makedirs(args.output_dir, exist_ok=True) + + # Load model + model = RFDETRBase(pretrain_weights=args.weights) + + # Load images + images = load_images(args.input_dir) + if not images: + raise ValueError(f"No supported images found in {args.input_dir}.") + + # Process each image + for idx, image in enumerate(images): + try: + detections = model.predict(image, threshold=args.confidence) + output_path = os.path.join(args.output_dir, f"annotated_{idx}.png") + save_annotated_image(image, detections, output_path) + print(f"Saved annotated image to {output_path}.") + except Exception as e: + print(f"Error processing image {idx}: {str(e)}.") + +if __name__ == "__main__": + main() From 3f86db367d1b7ccab9922084af9ec20cca4e8792 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 23 Mar 2026 21:13:20 +0000 Subject: [PATCH 2/2] =?UTF-8?q?fix(pre-commit):=20=F0=9F=8E=A8=20auto=20fo?= =?UTF-8?q?rmat=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- rfdetr/util/prediction_visualizer.py | 29 ++++++++++++++++++++++------ 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/rfdetr/util/prediction_visualizer.py b/rfdetr/util/prediction_visualizer.py index 322f8214b..9c0ad48a8 100644 --- a/rfdetr/util/prediction_visualizer.py +++ b/rfdetr/util/prediction_visualizer.py @@ -1,12 +1,21 @@ +# ------------------------------------------------------------------------ +# RF-DETR +# Copyright (c) 2025 Roboflow. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + import argparse import os from typing import List + import cv2 +import numpy as np import supervision as sv from PIL import Image -import numpy as np + from rfdetr import RFDETRBase + def parse_args(): parser = argparse.ArgumentParser(description="Visualize RF-DETR predictions on sample images.") parser.add_argument("--weights", type=str, required=True, help="Path to pre-trained RF-DETR model weights.") @@ -15,9 +24,15 @@ def parse_args(): parser.add_argument("--confidence", type=float, default=0.5, help="Confidence threshold for predictions.") return parser.parse_args() + def load_images(input_dir: str) -> List[Image.Image]: supported_extensions = (".jpg", ".jpeg", ".png") - return [Image.open(os.path.join(input_dir, f)) for f in os.listdir(input_dir) if f.lower().endswith(supported_extensions)] + return [ + Image.open(os.path.join(input_dir, f)) + for f in os.listdir(input_dir) + if f.lower().endswith(supported_extensions) + ] + def save_annotated_image(image: Image.Image, detections: sv.Detections, output_path: str): annotated_image = np.array(image) @@ -26,22 +41,23 @@ def save_annotated_image(image: Image.Image, detections: sv.Detections, output_p annotated_image = sv.LabelAnnotator().annotate(annotated_image, detections, labels) cv2.imwrite(output_path, cv2.cvtColor(annotated_image, cv2.COLOR_RGB2BGR)) + def main(): args = parse_args() - + # Validate inputs if not os.path.exists(args.input_dir): raise ValueError(f"Input directory {args.input_dir} does not exist.") os.makedirs(args.output_dir, exist_ok=True) - + # Load model model = RFDETRBase(pretrain_weights=args.weights) - + # Load images images = load_images(args.input_dir) if not images: raise ValueError(f"No supported images found in {args.input_dir}.") - + # Process each image for idx, image in enumerate(images): try: @@ -52,5 +68,6 @@ def main(): except Exception as e: print(f"Error processing image {idx}: {str(e)}.") + if __name__ == "__main__": main()