RF-DETR instance segmentation fine-tuning on Roboflow Universe datasets¶
Select one dataset with DATASET_KEY, then run the same fine-tuning, plotting, checkpoint loading, and inference cells.
RF-DETR extends bounding-box detection with a sparse segmentation head that predicts per-object binary masks — useful for facade analysis, infrastructure damage inspection, vehicle part assessment, dental imaging, litter mapping, and any task where what shape matters as much as what class.
Datasets — set DATASET_KEY to one of:
| Key | Dataset | Classes |
|---|---|---|
"building_facade" |
Building Facade Segmentation | facade, window, street, vegetation, car, fence, ... |
"spalling_rebar" |
Spalling and Exposed Rebar | exposed_rebar, spalling |
"car_damage_parts" |
Car Part Damage Segmentation | bumper, wheel, door, headlight, dent, scratch, ... |
"teeth_numbering" |
Teeth Numbering Segmentation | tooth IDs 11-48 |
"taco_trash" |
TACO Trash | bottle, can, carton, cup, lid, straw, wrapper, ... |
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 segmentation masks.
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.
!pip install -q "rfdetr[train,visual]>=1.8.2" roboflow pandas seaborn
"""Shared RF-DETR instance segmentation fine-tuning demo for several Roboflow COCO segmentation 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 RFDETRSegSmall
from rfdetr.config import SegmentationTrainConfig
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 (24 for RFDETRSegSmall).
# Larger resolutions capture fine-grained mask boundaries at the cost of GPU memory and time.
DATASETS: dict[str, dict[str, Any]] = {
"building_facade": {
"name": "Building Facade Segmentation",
"workspace": "building-facade",
"project": "building-facade-segmentation-instance",
"version": 4,
"source_url": "https://universe.roboflow.com/building-facade/building-facade-segmentation-instance",
"output_name": "seg_building_facade",
# 432 — structured street scenes with medium-to-large facade regions plus
# smaller windows, cars, vegetation, and infrastructure boundaries.
"resolution": 432,
},
"spalling_rebar": {
"name": "Spalling and Exposed Rebar",
"workspace": "uni-of-birmingham",
"project": "spalling-and-exposed-rebar",
"version": 7,
"source_url": "https://universe.roboflow.com/uni-of-birmingham/spalling-and-exposed-rebar",
"output_name": "seg_spalling_rebar",
# 432 — civil-infrastructure damage masks with irregular thin regions and
# exposed-rebar details that stress boundary quality.
"resolution": 432,
},
"car_damage_parts": {
"name": "Car Part Damage Segmentation",
"workspace": "car-damaged-detection-e66m0",
"project": "car-part-detection-with-damage-part",
"version": 1,
"source_url": "https://universe.roboflow.com/car-damaged-detection-e66m0/car-part-detection-with-damage-part",
"output_name": "seg_car_damage_parts",
# 432 — fine-grained automotive part masks with many adjacent classes and
# damage regions, useful for checking class-specific mask boundaries.
"resolution": 432,
},
"teeth_numbering": {
"name": "Teeth Numbering Segmentation",
"workspace": "godento2",
"project": "teeth-seg-3537-iaky1",
"version": 1,
"source_url": "https://universe.roboflow.com/godento2/teeth-seg-3537-iaky1",
"output_name": "seg_teeth_numbering",
# 432 — dental images with many repeated small instances whose labels are
# tied to position, stressing class confusion and small-mask quality.
"resolution": 432,
},
"taco_trash": {
"name": "TACO Trash",
"workspace": "mohamed-traore-2ekkp",
"project": "taco-trash-annotations-in-context",
"version": 13,
"source_url": "https://universe.roboflow.com/mohamed-traore-2ekkp/taco-trash-annotations-in-context",
"output_name": "seg_taco_trash",
# 432 — long-tail litter segmentation with many small, cluttered object
# categories. Use as a harder stress test after the simpler datasets.
"resolution": 432,
},
}
DATASET_KEY = "building_facade"
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 instance segmentation datasets in COCO format: a JSON file where each annotation contains a
segmentation field with polygon coordinates that trace the precise object boundary. 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.
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¶
Read the class names directly from the COCO annotation file. RF-DETR uses 0-based class indices internally; the
class_names list maps each index to a human-readable label so the inference visualisation shows category names
rather than numbers. If your dataset has a single category the list will have one element — that is expected.
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: SegmentationTrainConfig,
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="RFDETRSegSmall",
model_config_dict=model_config.model_dump(),
),
output_path,
)
return output_path
_segmentation_grid_figure arranges multiple annotated images into a fixed-column matplotlib grid.
Masks are drawn with sv.MaskAnnotator (translucent colour fill) and bounding boxes with
sv.BoxAnnotator. Labels show the class name and confidence score.
def _segmentation_grid_figure(
items: list[tuple[str, np.ndarray, sv.Detections]],
columns: int,
class_names: list[str],
) -> Figure:
"""Render segmentation masks and bounding boxes in a fixed-column subplot grid."""
if columns <= 0:
raise ValueError(f"columns must be positive, got {columns}.")
mask_annotator = sv.MaskAnnotator()
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()
if detections.mask is not None and len(detections) > 0:
scene = mask_annotator.annotate(scene=scene, detections=detections)
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¶
SegmentationTrainConfig centralises every hyperparameter the training loop needs, including mask-specific loss
coefficients (mask_ce_loss_coef and mask_dice_loss_coef). The defaults (5.0 each) work well for most datasets;
if the segmentation head overfits early you can reduce them towards 2.0 while keeping cls_loss_coef=5.0.
lr and lr_encoder control the head and backbone learning rates respectively. For fine-tuning both are usually
set to the same value; if you notice the backbone overfitting early, halve lr_encoder to protect the pre-trained
features. grad_accum_steps=2 doubles the effective batch size without extra GPU memory.
multi_scale=False and expanded_scales=False disable multi-resolution augmentation to shorten epoch time during
fine-tuning — they rarely help on small custom datasets. The notes dict is stored alongside the weights in the
checkpoint, giving you a lightweight experiment log that travels with the model file.
seed_all(SEED)
variant = RFDETRSegSmall( # type: ignore[no-untyped-call]
num_classes=NUM_CLASSES,
resolution=RESOLUTION,
)
variant.model_config.model_name = type(variant).__name__
train_config = SegmentationTrainConfig(
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"segmentation 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 you start a long training run. This cell renders a grid of annotated training images so you can confirm that segmentation masks align with objects, colour jitter and flips look reasonable, and there are no systematic labelling errors such as masks placed on the wrong object or masks that bleed outside the object boundary. Catching label noise here costs a few seconds; catching it after 50 epochs 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, loss
computation (cross-entropy + Dice on sampled mask points), gradient accumulation, weight updates,
learning-rate scheduling, and periodic validation with COCO mask mAP. After each epoch the trainer appends a row to
metrics.csv and, if validation mAP improves, saves a new best checkpoint via BestModelCallback. On a single
consumer GPU (RTX 3090 or similar) 50 epochs on the crack dataset takes roughly 20–35 minutes depending on
BATCH_SIZE and NUM_WORKERS.
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
RFDETRSegSmall.from_checkpoint(path) call. 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 with no augmentation, using the best checkpoint loaded from
checkpoint_best_total.pth written by BestModelCallback, and serialises results to JSON. Two metrics are most
important to inspect. bbox/map is the standard COCO bounding-box mAP at IoU 0.50:0.95. segm/map (when logged)
is the mask mAP — it measures how precisely the predicted binary masks overlap with the annotated polygons. A model
can have high bbox/map but low segm/map if the detector locates objects well but the mask boundaries are coarse;
this usually improves with more training epochs or by increasing resolution.
Note: checkpoint_best_total.pth is 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 raise FileNotFoundError.
# 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 curves is one of the fastest ways to diagnose training problems. Healthy segmentation runs show the
combined detection + mask loss decreasing together for train and val. A loss that rises or plateaus from epoch 1 is
usually a sign that the learning rate is too high or the dataset is too small for the chosen BATCH_SIZE. For the
mAP curves, segm mAP 50 (mask IoU ≥ 0.50) rises fastest and is the clearest signal of whether the model is
learning to segment objects; segm mAP 50:95 (averaged across stricter thresholds) follows more slowly.
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. You can share this file with teammates and they can run inference
immediately — the class names and architecture are all embedded alongside the weights.
loaded_model = RFDETRSegSmall.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 images from the test set gives
you an unbiased view of model performance because those images were never seen during training or used to pick the
best checkpoint. Replace inference_image_paths with a list of Path objects pointing to your own files to run
inference on custom images.
_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)
validation_image_paths = _find_images(DATASET_DIR / "test")
if not validation_image_paths:
validation_image_paths = _find_images(DATASET_DIR / "valid")
if not validation_image_paths:
validation_image_paths = _find_images(DATASET_DIR / "train")
if not validation_image_paths:
raise FileNotFoundError(f"No images ({', '.join(sorted(_IMAGE_EXTS))}) found under {DATASET_DIR}")
inference_image_paths = validation_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. Raise it to suppress false positives; lower it to surface low-confidence detections. Each accepted
detection includes a binary mask resized to the original image dimensions — sv.MaskAnnotator renders it as a
translucent colour overlay so you can see both the mask shape and the image beneath it.
If the mask boundaries look jagged, consider increasing RESOLUTION (in multiples of 24) and retraining; higher
resolution gives the segmentation head more spatial detail to work with.
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 RFDETRSegSmall.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 = _segmentation_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. If you see zero detections on most
images, try lowering INFERENCE_THRESHOLD (towards 0.1) — the model may be calibrated to lower confidence scores
on a custom dataset than the COCO-pretrained default. If you see too many false positives, raise the threshold or
add more negative examples to your training set.
summary_rows = []
for image_path, (_, _, detections) in zip(inference_image_paths, inference_grid_items):
num_det = len(detections)
has_masks = detections.mask is not None and detections.mask.shape[0] > 0
avg_conf = (
float(np.mean(detections.confidence)) if detections.confidence is not None and num_det > 0 else float("nan")
)
summary_rows.append(
{
"image": image_path.name,
"detections": num_det,
"has_masks": has_masks,
"avg_confidence": round(avg_conf, 3),
}
)
summary_df = pd.DataFrame(summary_rows)
display(summary_df)