Skip to content

Create prediction_visualizer.py #345

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
catsmells wants to merge 4 commits into
base: develop
Choose a base branch
from
Draft

Create prediction_visualizer.py #345

Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
6ccbdd2
Create prediction_visualizer.py
catsmells Sep 7, 2025
5f4791a
Merge branch 'develop' into patch-1
Borda Mar 23, 2026
3f86db3
fix(pre-commit): 🎨 auto format pre-commit hooks
pre-commit-ci[bot] Mar 23, 2026
65258c8
Merge branch 'develop' into patch-1
Borda Apr 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 73 additions & 0 deletions rfdetr/util/prediction_visualizer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
# ------------------------------------------------------------------------
# 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
Comment on lines +1 to +8
Copy link

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copilot uses AI. Check for mistakes.
from typing import List

import cv2
import numpy as np
import supervision as sv
from PIL import Image

Comment on lines +7 to +15
Copy link

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

Copilot uses AI. Check for mistakes.
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)
]
Comment on lines +30 to +34
Copy link

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

Suggested change
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

Copilot uses AI. Check for mistakes.


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)
Comment on lines +39 to +41
Copy link

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
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)

Copilot uses AI. Check for mistakes.
cv2.imwrite(output_path, cv2.cvtColor(annotated_image, cv2.COLOR_RGB2BGR))
Comment on lines +37 to +42
Copy link

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copilot uses AI. Check for mistakes.


def main():
args = parse_args()

Comment on lines +19 to +47
Copy link

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copilot uses AI. Check for mistakes.
# 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()
Loading