RF-DETR object detection fine-tuning on Roboflow Universe datasets¶
Select one dataset with DATASET_KEY, then run the same fine-tuning, plotting, checkpoint loading, and inference cells.
What is object detection?¶
Object detection answers two questions simultaneously: what objects are in an image and where they are. Each detected object gets a bounding box (four coordinates) and a class label (e.g. "car", "person"). This is the most common computer vision task — it powers traffic monitoring, safety compliance, retail analytics, medical imaging, environmental surveys, and virtually every real-time vision application.
Why fine-tune instead of train from scratch?¶
RF-DETR ships pre-trained on the COCO dataset (80 common object classes). Fine-tuning reuses the backbone's rich visual representations — edges, textures, object parts — learned from millions of images. Starting from these weights instead of random initialisation typically:
- Reduces training time from days to minutes or hours
- Requires far fewer labelled images (hundreds instead of tens of thousands)
- Achieves higher accuracy than training from scratch on small datasets
The process is straightforward: swap out the final classification head for your number of classes, then continue training with a small learning rate so the pre-trained features are preserved.
Dataset overview¶
Five datasets covering different domains are pre-configured. Set DATASET_KEY to one of:
| Key | Dataset | Domain | Classes |
|---|---|---|---|
"hard_hat_ppes" |
Hard Hat Worker Safety | Construction / PPE | hat, person, vest |
"road_damage" |
Pavement Distress Detection | Infrastructure Inspection | pot_hole, cracking, ravelling, ... |
"traffic" |
Traffic Detection | Autonomous Driving | car, bus, truck, motorbike, bike, traffic lights |
"football" |
Football Player Detection | Sports Analytics | player, ball, goalkeeper, referee |
"brackish_underwater" |
Brackish Underwater | Marine Wildlife | fish, crab, jellyfish, shrimp, starfish, small_fish |
Every subsequent cell (fine-tuning, plotting, checkpoint loading, inference) adapts automatically. By the end you will have a fine-tuned model, training curves, and annotated inference images with bounding boxes.
Setup¶
Install rfdetr with the train and visual extras. The train extra pulls in the training loop dependencies
(PyTorch Lightning, COCO evaluation tools, and data augmentation libraries), while visual adds the visualisation
helpers used later in the notebook. The roboflow package handles dataset download; pandas and seaborn are
needed for the metrics table and optional plot styling. If you hit an ImportError after running this cell,
the most likely cause is a stale in-memory import — restart the Python runtime once and re-run from the top.
GPU recommended. Fine-tuning 50 epochs on a ~1,500-image dataset takes roughly 10–25 minutes on a modern GPU (RTX 3090, A10, T4). CPU-only is possible for a quick test but will be much slower.
!pip install -q "rfdetr[train,visual]>=1.8.0" roboflow pandas seaborn
"""Shared RF-DETR object detection fine-tuning demo for several Roboflow COCO detection exports."""
import json
import os
from pathlib import Path
from typing import Any
import numpy as np
import pandas as pd
import supervision as sv
import torch
from IPython.display import display
from matplotlib import pyplot as plt
from matplotlib.figure import Figure
from PIL import Image
from roboflow import Roboflow
from rfdetr import RFDETRSmall
from rfdetr.config import TrainConfig
from rfdetr.training import RFDETRDataModule, RFDETRModelModule, build_trainer
from rfdetr.training.callbacks.best_model import BestModelCallback
from rfdetr.utilities.reproducibility import seed_all
from rfdetr.visualize.training import plot_loss_metrics, plot_map_metrics
PROJECT_ROOT = Path(__file__).resolve().parent if "__file__" in globals() else Path.cwd()
DATASETS_DIR = PROJECT_ROOT / "datasets"
# Resolution must be divisible by patch_size × num_windows (32 for RFDETRSmall: patch_size=16, num_windows=2).
# Higher resolution improves detection of small objects at the cost of GPU memory and training time.
# Practical range: 384–640. The default 512 suits most fine-tuning runs.
DATASETS: dict[str, dict[str, Any]] = {
"hard_hat_ppes": {
"name": "Hard Hat Worker Safety",
"workspace": "safety-first",
"project": "hard-hat-worker-safety-equipments",
"version": 9,
"source_url": "https://universe.roboflow.com/safety-first/hard-hat-worker-safety-equipments",
"output_name": "det_hard_hat_ppes",
# 512 — outdoor construction scenes; workers appear at medium scale so the
# default resolution is sufficient to distinguish helmet vs no-helmet.
"resolution": 512,
},
"road_damage": {
"name": "Pavement Distress Detection",
"workspace": "pavement-distresses",
"project": "distress-detection",
"version": 2,
"source_url": "https://universe.roboflow.com/pavement-distresses/distress-detection",
"output_name": "det_road_damage",
# 544 — road-surface images captured from dashcams and drones; fine cracks
# and potholes are small relative to the frame, so the extra resolution
# helps the model distinguish closely spaced defect categories.
"resolution": 544,
},
"traffic": {
"name": "Traffic Detection",
"workspace": "redlightrunningdection",
"project": "traffic-detection-sutq6",
"version": 36,
"source_url": "https://universe.roboflow.com/redlightrunningdection/traffic-detection-sutq6",
"output_name": "det_traffic",
# 512 — street-level traffic scenes; vehicles span a wide size range
# (large buses to distant motorbikes); 512 gives a good speed/accuracy trade-off.
"resolution": 512,
},
"football": {
"name": "Football Player Detection",
"workspace": "football-gozni",
"project": "football-player-detection-bfswn",
"version": 1,
"source_url": "https://universe.roboflow.com/football-gozni/football-player-detection-bfswn",
"output_name": "det_football",
# 544 — broadcast and stadium-angle shots; the ball is a tiny object relative
# to the field; the extra resolution meaningfully reduces false negatives.
"resolution": 544,
},
"brackish_underwater": {
"name": "Brackish Underwater",
"workspace": "brad-dwyer",
"project": "brackish-underwater",
"version": 2,
"source_url": "https://universe.roboflow.com/brad-dwyer/brackish-underwater",
"output_name": "det_brackish_underwater",
# 512 — sonar and underwater camera frames; marine creatures appear at
# medium-to-large scale; the default resolution captures enough detail.
"resolution": 512,
},
}
DATASET_KEY = "hard_hat_ppes"
DATASET_INFO = DATASETS[DATASET_KEY]
OUTPUT_DIR = PROJECT_ROOT / "output" / str(DATASET_INFO["output_name"])
METRICS_CSV = OUTPUT_DIR / "metrics.csv"
VALIDATION_METRICS_JSON = OUTPUT_DIR / "validation_metrics.json"
FINAL_CHECKPOINT_PATH = OUTPUT_DIR / "checkpoint_final_demo.pth"
RESOLUTION = int(DATASET_INFO["resolution"])
SEED = 7
EPOCHS = 50
BATCH_SIZE = 8
GRAD_ACCUM_STEPS = 2
NUM_WORKERS = 8
LR = 1e-4
LR_ENCODER = 1e-4
SAMPLE_PREVIEW_COUNT = 6
SAMPLE_PREVIEW_COLUMNS = 3
SAMPLE_PREVIEW_FIGURE_SIZE: tuple[float, float] = (15.0, 10.0)
INFERENCE_COUNT = 6
INFERENCE_COLUMNS = 3
INFERENCE_THRESHOLD = 0.3
PLOT_LOSS_LOG_SCALE = False
print(f"dataset_key={DATASET_KEY}")
print(f"dataset={DATASET_INFO['name']}")
print(f"source_url={DATASET_INFO['source_url']}")
print(f"resolution={RESOLUTION}")
Notebook display¶
This helper registers the %matplotlib inline magic so figures render directly beneath each cell when you run the
notebook in Jupyter or Google Colab. If you are running this file as a plain Python script the function detects the
absence of an IPython kernel and exits silently, so it is always safe to call.
def _enable_notebook_inline_matplotlib() -> None:
"""Enable inline matplotlib figures when running in IPython."""
get_ipython_func = globals().get("get_ipython")
if not callable(get_ipython_func):
return
ipython = get_ipython_func()
if ipython is not None:
ipython.run_line_magic("matplotlib", "inline")
ipython.run_line_magic("config", "InlineBackend.close_figures = True")
_enable_notebook_inline_matplotlib()
1 - Download dataset¶
Roboflow exports object detection datasets in COCO format: a JSON file containing image metadata and per-image
annotations, each with a bbox field (x, y, width, height in pixels) and a category_id referencing the class.
The download cell fetches the exact dataset version listed in DATASETS and places it under
datasets/<DATASET_KEY>/. You need a Roboflow API key — get one for free at app.roboflow.com/settings/api
and set it as the ROBOFLOW_API_KEY environment variable (or as a Colab secret with the same name).
The download is idempotent: if the target directory already exists Roboflow skips the network transfer, so re-running this cell after a successful download is fast.
Tip: All five datasets in this notebook are open-access on Roboflow Universe. You can browse them in your browser to explore image samples and annotation statistics before downloading.
try:
from google.colab import userdata
try:
ROBOFLOW_API_KEY = userdata.get("ROBOFLOW_API_KEY") or ""
except Exception:
ROBOFLOW_API_KEY = ""
except ImportError:
ROBOFLOW_API_KEY = ""
if not ROBOFLOW_API_KEY:
ROBOFLOW_API_KEY = os.environ.get("ROBOFLOW_API_KEY", "")
if not ROBOFLOW_API_KEY:
raise RuntimeError(
"ROBOFLOW_API_KEY not found. "
"In Colab: add it via Secrets (key icon). "
"Locally: set the environment variable before running."
)
rf = Roboflow(api_key=ROBOFLOW_API_KEY)
dataset = (
rf.workspace(str(DATASET_INFO["workspace"]))
.project(str(DATASET_INFO["project"]))
.version(int(DATASET_INFO["version"]))
.download("coco", location=str(DATASETS_DIR / DATASET_KEY))
)
DATASET_DIR = Path(dataset.location)
TRAIN_ANNOTATIONS = DATASET_DIR / "train" / "_annotations.coco.json"
print(f"dataset_dir={DATASET_DIR}")
2 - Infer class names¶
The COCO annotation file defines a categories list mapping integer IDs to human-readable names. RF-DETR uses
0-based class indices internally, so we sort the categories by their original id field and read names in order.
The resulting CLASS_NAMES list maps each index back to a label, which is used during inference visualisation.
If your own dataset has a single category, this list will have one element — that is expected and valid.
with TRAIN_ANNOTATIONS.open() as _f:
_coco = json.load(_f)
CLASS_NAMES: list[str] = [cat["name"] for cat in sorted(_coco["categories"], key=lambda c: c["id"])]
NUM_CLASSES = len(CLASS_NAMES)
print(f"class_names={CLASS_NAMES}")
print(f"num_classes={NUM_CLASSES}")
_save_final_checkpoint writes a self-contained .pth file that bundles model weights with the full training and
model config. This is separate from the PTL checkpoint because it can be loaded with a single from_checkpoint
call on any machine, without reconstructing the original config.
def _save_final_checkpoint(
module: RFDETRModelModule,
trainer: Any,
train_config: TrainConfig,
model_config: Any,
output_path: Path,
) -> Path:
"""Save a final RF-DETR .pth checkpoint that can be loaded with ``from_checkpoint``."""
raw_model: Any = getattr(module.model, "_orig_mod", module.model)
output_path.parent.mkdir(parents=True, exist_ok=True)
torch.save(
BestModelCallback._build_checkpoint_payload(
raw_model.state_dict(),
train_config.model_dump(),
trainer,
model_name="RFDETRSmall",
model_config_dict=model_config.model_dump(),
),
output_path,
)
return output_path
_detection_grid_figure arranges multiple annotated images into a fixed-column matplotlib grid.
Bounding boxes are drawn with sv.BoxAnnotator and labels with sv.LabelAnnotator, showing the class name and
confidence score for each detection. Using a grid keeps the visual summary compact and makes it easy to compare
predictions across images at a glance.
def _detection_grid_figure(
items: list[tuple[str, np.ndarray, sv.Detections]],
columns: int,
class_names: list[str],
) -> Figure:
"""Render bounding box detections in a fixed-column subplot grid."""
if columns <= 0:
raise ValueError(f"columns must be positive, got {columns}.")
box_annotator = sv.BoxAnnotator(thickness=2)
label_annotator = sv.LabelAnnotator(text_scale=0.5, text_thickness=1, text_padding=4)
rows = max(1, (len(items) + columns - 1) // columns)
figure, axes = plt.subplots(rows, columns, figsize=(5 * columns, 5 * rows))
axes_array = np.asarray(axes, dtype=object).reshape(-1)
for axis in axes_array:
axis.axis("off")
for axis, (title, image, detections) in zip(axes_array, items, strict=False):
scene = image.copy()
scene = box_annotator.annotate(scene=scene, detections=detections)
if len(detections) > 0 and detections.class_id is not None:
conf_list = (
detections.confidence.tolist() if detections.confidence is not None else [None] * len(detections)
)
labels = [
f"{class_names[cid] if cid < len(class_names) else cid}" + (f" {conf:.2f}" if conf is not None else "")
for cid, conf in zip(detections.class_id.tolist(), conf_list)
]
scene = label_annotator.annotate(scene=scene, detections=detections, labels=labels)
axis.imshow(scene)
axis.set_title(title, fontsize=10)
axis.axis("off")
figure.tight_layout()
return figure
3 - Configure training¶
TrainConfig centralises every hyperparameter the training loop needs. The most important fields to understand
before customising them for your own dataset:
Learning rates — lr and lr_encoder are both set to 1e-4 for fine-tuning. The lr controls the
detection head and decoder; lr_encoder controls the vision backbone (DINOv2). For fine-tuning, both are usually
equal. If you notice the backbone overfitting early — validation loss rising while training loss keeps falling —
halve lr_encoder to protect the pre-trained features while the head continues to adapt.
Batch size and gradient accumulation — batch_size=8 with grad_accum_steps=2 gives an effective batch size
of 16 without requiring extra GPU memory. Effective batch size affects how smoothly gradients are estimated each
update; too small (< 8) leads to noisy updates; too large (> 64) may over-smooth and slow convergence.
EMA — use_ema=False keeps this demo fast. For a production run you can enable EMA
(use_ema=True) to maintain an exponential moving average of the weights, which typically adds 0.5–1.0 mAP points
on small datasets by reducing the impact of noisy late-epoch updates.
Multi-scale training — multi_scale=False and expanded_scales=False disable the multi-resolution
augmentation used during pre-training. Turning them off shortens epoch time by up to 40 % and is fine for
most fine-tuning runs on small custom datasets. Re-enable if your validation mAP plateaus early.
Notes dict — stored verbatim in the checkpoint file alongside the weights, giving you a lightweight experiment log that travels with the model file.
seed_all(SEED)
variant = RFDETRSmall( # type: ignore[no-untyped-call]
num_classes=NUM_CLASSES,
resolution=RESOLUTION,
)
variant.model_config.model_name = type(variant).__name__
train_config = TrainConfig(
dataset_file="roboflow",
dataset_dir=str(DATASET_DIR),
output_dir=str(OUTPUT_DIR),
epochs=EPOCHS,
batch_size=BATCH_SIZE,
grad_accum_steps=GRAD_ACCUM_STEPS,
num_workers=NUM_WORKERS,
lr=LR,
lr_encoder=LR_ENCODER,
use_ema=False,
run_test=False,
compute_train_metrics=True,
compute_val_loss=True,
multi_scale=False,
expanded_scales=False,
do_random_resize_via_padding=False,
tensorboard=False,
wandb=False,
mlflow=False,
clearml=False,
class_names=CLASS_NAMES,
notes={
"demo": f"detection PTL fine-tune on Roboflow Universe {DATASET_INFO['project']}",
"source_url": DATASET_INFO["source_url"],
"roboflow_workspace": DATASET_INFO["workspace"],
"roboflow_project": DATASET_INFO["project"],
"roboflow_version": DATASET_INFO["version"],
"num_classes": NUM_CLASSES,
"class_names": CLASS_NAMES,
},
progress_bar="tqdm",
)
datamodule = RFDETRDataModule(variant.model_config, train_config)
model = RFDETRModelModule(variant.model_config, train_config)
trainer = build_trainer(train_config, variant.model_config)
4 - Preview dataset inputs¶
Always look at your data before starting a long training run. This cell renders a grid of annotated training images drawn from the augmented pipeline — the same view the model will see during training.
What to check here:
- Box alignment — do bounding boxes tightly surround the objects, or are they shifted or too large?
- Class correctness — are labels correct for each box? Mixed-up classes are one of the most common annotation errors and are hard to recover from after training.
- Augmentation sanity — horizontal flips, colour jitter, and random crops should look plausible; cropped boxes at image edges are normal.
- Class imbalance — if one class dominates the preview, the model will likely be biased toward it; rebalancing or oversampling rare classes can help.
Catching label noise here costs a few seconds; catching it after 50 epochs of training costs much more.
sample_figure = datamodule._show_samples(
SAMPLE_PREVIEW_COUNT,
split="train",
columns=SAMPLE_PREVIEW_COLUMNS,
figure_size=SAMPLE_PREVIEW_FIGURE_SIZE,
)
display(sample_figure)
plt.close(sample_figure)
5 - Fine-tune¶
trainer.fit hands control to PyTorch Lightning, which drives the full training loop: forward pass, bipartite
matching loss (classification + bounding-box L1 + GIoU), gradient accumulation, weight updates,
learning-rate scheduling, and periodic validation with COCO bounding-box mAP.
After each epoch the trainer appends a row to metrics.csv and, if validation mAP improves,
saves a new best checkpoint via BestModelCallback.
Typical training times on a single GPU (RTX 3090 or similar):
- PPE / hard hat dataset (~2,600 images): ~20–30 min for 50 epochs
- Traffic dataset (~1,400 images): ~15–20 min for 50 epochs
- Underwater dataset (~12,000 images): ~60–90 min for 50 epochs
To resume an interrupted run, set train_config.resume to the path of the last PTL checkpoint
(e.g. OUTPUT_DIR / "last.ckpt") before calling this cell again.
trainer.fit(model, datamodule=datamodule, ckpt_path=train_config.resume or None)
print(f"output_dir={OUTPUT_DIR}")
6 - Save checkpoint/model¶
PTL saves its own checkpoints during training (optimizer state, scheduler state, epoch counter), but those files
are not directly portable — they require the same class hierarchy to load. The .pth file written here is a
self-contained RF-DETR checkpoint: it bundles the model weights together with the full training and model
configs, including the class names. You can share the file with a colleague and they can run inference with a
single RFDETRSmall.from_checkpoint(path) call, with no need to reconstruct the original config.
For full reproducibility, keep this checkpoint alongside the dataset version number printed in section 1.
final_checkpoint = _save_final_checkpoint(model, trainer, train_config, variant.model_config, FINAL_CHECKPOINT_PATH)
print(f"saved_checkpoint={final_checkpoint}")
7 - Validate metrics¶
This cell runs a clean post-training validation: no augmentation, the best checkpoint loaded from
checkpoint_best_total.pth written by BestModelCallback, and results serialised to JSON for downstream
comparison or reporting.
Key metrics to inspect:
bbox/map— the primary metric: standard COCO bounding-box mAP averaged across IoU thresholds 0.50–0.95 in steps of 0.05. A value of 0.50 means the model correctly localises and classifies 50 % of objects at strict IoU thresholds. This is the most commonly reported number in object detection benchmarks.bbox/map_50— mAP at IoU ≥ 0.50 (a box is "correct" if it overlaps the ground-truth box by at least 50 %). Rises fastest and is the clearest early signal of whether the model is learning at all.bbox/map_75— mAP at the stricter IoU ≥ 0.75 threshold. Reflects how precisely the model localises objects, not just whether it finds them.
A model with high bbox/map_50 but low bbox/map_75 finds objects reliably but draws imprecise boxes.
This usually improves with more training epochs or by increasing RESOLUTION.
Note:
checkpoint_best_total.pthis written after the first completed validation epoch. If training was interrupted before any epoch finished, this file will not exist and the cell below will raiseFileNotFoundError.
# checkpoint_best_total.pth stores plain model weights under key "model", not a full
# PTL checkpoint. Passing it as ckpt_path to trainer.validate() triggers a KeyError
# on "validate_loop" because PTL tries to restore loop state that isn't present.
# Load the weights directly and run validate without ckpt_path instead.
_ckpt = torch.load(OUTPUT_DIR / "checkpoint_best_total.pth", map_location="cpu", weights_only=False)
model.model.load_state_dict(_ckpt["model"], strict=True)
del _ckpt
validation_results = trainer.validate(model, datamodule=datamodule)
validation_metrics = {key: float(value) for key, value in validation_results[0].items()} if validation_results else {}
VALIDATION_METRICS_JSON.write_text(json.dumps(validation_metrics, indent=2, sort_keys=True), encoding="utf-8")
print(f"validation_metrics={validation_metrics}")
print(f"validation_metrics_json={VALIDATION_METRICS_JSON}")
8 - Plot CSVLogger metrics¶
Reading loss and mAP curves is the fastest way to diagnose training problems.
Loss curves (first plot):
- A healthy run shows both training loss and validation loss decreasing together.
- A growing gap between train and val loss is a classic overfitting signal — the model is memorising the
training set. Add augmentation (
multi_scale=True, or custom Albumentations), reduce learning rate, or collect more data. - A flat or rising loss from epoch 1 usually means the learning rate is too high. Try reducing
LRto5e-5.
mAP curves (second plot):
bbox_map_50(solid line) rises fastest and is the clearest signal of learning.- On small datasets (a few hundred images) you typically see rapid improvement in epochs 1–20 followed by a plateau around epoch 30–50.
- If mAP is still rising at epoch 50, extend
EPOCHSto 75 or 100. - If mAP never rises above 0.05, check class names, annotation quality, and whether
RESOLUTIONis appropriate for the objects' scale in your dataset.
Set PLOT_LOSS_LOG_SCALE = True if the loss drops by an order of magnitude in the first few epochs and the later,
more meaningful portion of the curve gets compressed into a flat line.
print(f"metrics_csv={METRICS_CSV}")
loss_figure = plot_loss_metrics(str(METRICS_CSV), loss_log_scale=PLOT_LOSS_LOG_SCALE)
display(loss_figure)
plt.close(loss_figure)
map_figure = plot_map_metrics(str(METRICS_CSV))
display(map_figure)
plt.close(map_figure)
9 - Load checkpoint/model¶
from_checkpoint is the standard entry point for loading a saved RF-DETR model. It reads both the weights and the
stored config from the .pth file, so the reconstructed model has the correct number of classes and resolution
without you having to pass them explicitly.
This cell also confirms that the save–load round trip works before you proceed to inference, so any file corruption or version mismatch surfaces here rather than silently producing wrong predictions later.
To deploy this model in a production pipeline, you only need to ship the .pth file and call from_checkpoint.
loaded_model = RFDETRSmall.from_checkpoint(FINAL_CHECKPOINT_PATH)
10 - Select inference images¶
This cell implements a fallback chain — test split first, then validation, then train — so the notebook always finds images to run inference on even when the dataset has no dedicated test split.
Using test images gives an unbiased view of model performance because those images were never seen during
training or used to select the best checkpoint. If your deployment images look different from the training set
(different lighting, camera angle, resolution), replace inference_image_paths with a list of Path objects
pointing to your own files to assess real-world performance before shipping the model.
_IMAGE_EXTS = {".jpg", ".jpeg", ".png"}
def _find_images(directory: Path) -> list[Path]:
if not directory.is_dir():
return []
return sorted(p for p in directory.glob("**/*") if p.is_file() and p.suffix.lower() in _IMAGE_EXTS)
inference_image_paths = _find_images(DATASET_DIR / "test")
if not inference_image_paths:
inference_image_paths = _find_images(DATASET_DIR / "valid")
if not inference_image_paths:
inference_image_paths = _find_images(DATASET_DIR / "train")
if not inference_image_paths:
raise FileNotFoundError(f"No images ({', '.join(sorted(_IMAGE_EXTS))}) found under {DATASET_DIR}")
inference_image_paths = inference_image_paths[:INFERENCE_COUNT]
print(f"inference_images={[str(p) for p in inference_image_paths]}")
11 - Visualize inference¶
INFERENCE_THRESHOLD is a detection confidence threshold: only detections whose bounding-box score exceeds
this value are shown. The default 0.3 is a reasonable starting point; adjust it based on what you see:
- Too many false positives (boxes on background or wrong objects) → raise the threshold towards
0.5–0.7. - Missing true objects (objects in the image not detected) → lower the threshold towards
0.1–0.2. - Zero detections on most images → lower the threshold; the model may be well-calibrated to lower confidence scores on a custom dataset than the COCO-pretrained default.
The confidence score reflects how certain the model is about each box. On a well-fine-tuned model you will typically see scores above 0.7 for clear, unoccluded objects of the target classes.
inference_grid_items: list[tuple[str, np.ndarray, sv.Detections]] = []
for image_path in inference_image_paths:
with Image.open(image_path) as image:
image_rgb = image.convert("RGB")
detections = loaded_model.predict(image_rgb, threshold=INFERENCE_THRESHOLD)
if not isinstance(detections, sv.Detections):
raise RuntimeError(f"Expected RFDETRSmall.predict() to return sv.Detections, got {type(detections)!r}.")
inference_grid_items.append((image_path.name, np.array(image_rgb), detections))
if inference_grid_items:
figure = _detection_grid_figure(inference_grid_items, columns=INFERENCE_COLUMNS, class_names=CLASS_NAMES)
display(figure)
plt.close(figure)
12 - Detection summary¶
Print per-image detection counts and confidence scores as a quick sanity check. This table is useful for:
- Spotting images with zero detections — if many images produce no detections at your chosen threshold,
the model may be under-confident; try lowering
INFERENCE_THRESHOLDfirst. - Checking average confidence — a well-calibrated fine-tuned model typically reports average confidence between 0.5 and 0.9 on in-distribution images. Scores consistently below 0.3 suggest the model is uncertain and may benefit from more training epochs or additional labelled data.
- Class distribution — if certain classes never appear in inference results but were present in training data, they may be under-represented in the training set. Consider adding more examples for those classes.
summary_rows = []
for image_path, (_, _, detections) in zip(inference_image_paths, inference_grid_items):
num_det = len(detections)
avg_conf = (
float(np.mean(detections.confidence)) if detections.confidence is not None and num_det > 0 else float("nan")
)
class_counts: dict[str, int] = {}
if num_det > 0 and detections.class_id is not None:
for cid in detections.class_id.tolist():
label = CLASS_NAMES[cid] if cid < len(CLASS_NAMES) else str(cid)
class_counts[label] = class_counts.get(label, 0) + 1
summary_rows.append(
{
"image": image_path.name,
"detections": num_det,
"avg_confidence": round(avg_conf, 3),
"classes": ", ".join(f"{cls}×{cnt}" for cls, cnt in sorted(class_counts.items())),
}
)
summary_df = pd.DataFrame(summary_rows)
display(summary_df)