Conversation
Adding testing visualizer.
catsmells
requested review from
Matvezy,
SkalskiP,
isaacrob-roboflow and
probicheaux
as code owners
|
|
|
I have read the CLA Document and I sign the CLA. |
Borda
force-pushed
the
develop
branch
4 times, most recently
from
60b16c1 to
523f9df
Compare
There was a problem hiding this comment.
Pull request overview
Adds a helper script intended to visualize RF-DETR prediction outputs on local images to aid onboarding/debugging.
Changes:
- Introduces a new prediction-visualization script that loads weights, runs
RFDETRBase.predict()over a directory of images, and writes annotated images. - Implements basic CLI argument parsing for weights/input/output/confidence.
| # ------------------------------------------------------------------------ | ||
| # 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 |
There was a problem hiding this comment.
This file is added under the top-level rfdetr/ directory, but the package is configured as a src/ layout ([tool.setuptools.packages.find] where = ["src"]), so rfdetr/util/prediction_visualizer.py will not be included in the installed distribution or importable via python -m rfdetr. If this is meant to be part of the library, it should live under src/rfdetr/...; if it’s meant to be a runnable helper script (as the PR description suggests), consider placing it in a scripts/ (or similar) directory and documenting that path instead. Also note the PR description references visualize_predictions.py, but this change introduces prediction_visualizer.py.
| import argparse | ||
| import os | ||
| from typing import List | ||
|
|
||
| import cv2 | ||
| import numpy as np | ||
| import supervision as sv | ||
| from PIL import Image | ||
|
|
There was a problem hiding this comment.
This module imports cv2 at import time, but opencv-python is not a declared core dependency in pyproject.toml. If this file is moved under src/rfdetr (so it ships), importing it will raise ImportError in minimal installs and potentially in CI. Consider moving the cv2 dependency behind an optional import inside main()/save_annotated_image, or avoid OpenCV entirely (PIL can save NumPy arrays).
| 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() | ||
|
|
There was a problem hiding this comment.
Several new functions are missing return type hints and docstrings (parse_args, save_annotated_image, main). The repository generally uses full typing + docstrings for functions; adding them here will keep this consistent and help downstream users understand expected types/behavior.
| return [ | ||
| Image.open(os.path.join(input_dir, f)) | ||
| for f in os.listdir(input_dir) | ||
| if f.lower().endswith(supported_extensions) | ||
| ] |
There was a problem hiding this comment.
load_images() returns PIL.Image objects in whatever mode the source file has (e.g. "RGBA", "L"). RFDETR.predict() expects RGB (3 channels) and will raise for non-3-channel images. Consider converting each loaded image to RGB (and forcing image data to be loaded so file handles don’t remain open).
| return [ | |
| Image.open(os.path.join(input_dir, f)) | |
| for f in os.listdir(input_dir) | |
| if f.lower().endswith(supported_extensions) | |
| ] | |
| images: List[Image.Image] = [] | |
| for filename in os.listdir(input_dir): | |
| if not filename.lower().endswith(supported_extensions): | |
| continue | |
| image_path = os.path.join(input_dir, filename) | |
| with Image.open(image_path) as img: | |
| rgb_image = img.convert("RGB") | |
| rgb_image.load() | |
| images.append(rgb_image) | |
| return images |
| 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)) |
There was a problem hiding this comment.
labels = [f"{pred.class_name} {pred.confidence:.2f}" for pred in detections] is unlikely to work with sv.Detections returned by RFDETR.predict(). That method constructs sv.Detections with class_id and confidence arrays (no class_name field), so this will raise at runtime or produce incorrect labels. Build labels from detections.class_id/detections.confidence and map class IDs via model.class_names if you want human-readable names.
| 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) |
There was a problem hiding this comment.
sv.BoxAnnotator().annotate(annotated_image, detections) and LabelAnnotator().annotate(annotated_image, detections, labels) are called positionally here, but in the codebase the supervision annotators are used with keyword arguments (scene=..., detections=..., labels=...). Using keywords avoids accidental argument-order mismatches across supervision versions and improves readability.
| 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) | |
| annotated_image = sv.BoxAnnotator().annotate(scene=annotated_image, detections=detections) | |
| labels = [f"{pred.class_name} {pred.confidence:.2f}" for pred in detections] | |
| annotated_image = sv.LabelAnnotator().annotate(scene=annotated_image, detections=detections, labels=labels) |
Codecov Report✅ All modified and coverable lines are covered by tests. ❌ Your project check has failed because the head coverage (0%) is below the target coverage (95%). You can increase the head coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## develop #345 +/- ##
========================================
- Coverage 76% 0 -76%
========================================
Files 97 0 -97
Lines 7368 0 -7368
========================================
- Hits 5609 0 -5609
+ Misses 1759 0 -1759 🚀 New features to boost your workflow:
|
Description
Added
visualize_predictions.pyto the RF-DETR repository (https://github.com/roboflow/rf-detr/tree/develop/rfdetr) to enable visualization of model predictions on sample images, improving user onboarding and debugging. The script loads a pre-trained RF-DETR model, runs inference on images in a specified directory, and saves annotated outputs with bounding boxes and labels usingsupervision. Motivation: enhance usability for developers testing RF-DETR performance. Dependencies:numpy,opencv-python,supervision,pillow,rfdetr.Type of change
How has this change been tested, please provide a testcase or example of how you tested the change?
Tested by running
python visualize_predictions.py --weights path/to/weights.pth --input-dir sample_images --output-dir output --confidence 0.5with a COCO sample dataset. Verified annotated images in the output directory displayed correct bounding boxes and labels for detected objects.Any specific deployment considerations
Requires pre-trained RF-DETR model weights and a directory with supported image formats (.jpg, .jpeg, .png). No additional costs or secrets required.
Docs
Added section for
visualize_predictions.py, detailing usage, command-line arguments (--weights,--input-dir,--output-dir,--confidence), and example command:python visualize_predictions.py --weights weights.pth --input-dir images/ --output-dir output/.I have read the CLA Document and I sign the CLA.