Training & Deploying a Detector on Custom Data

"The theory was lovely. Then they handed me four hundred photos of their specific kind of bolt, half of them blurry, and said: by Friday. That is when I learned that detection is ten percent architecture and ninety percent labels, learning rate, and remembering which folder the validation set lives in."

A Fine-Tuned Detector With a Deadline

Everything from Section 23.1 to Section 23.5 was about how detectors work. This section is about getting one to work for you. The honest truth of applied detection is that the architecture choice, the subject of four sections, is usually the smallest decision; far more of your time goes to labeling, to the transfer-learning recipe from Chapter 21, to reading the metrics of Section 23.1 without fooling yourself, and to exporting for the hardware of Chapter 28. We walk that whole route once, end to end.

1. Collecting and Labeling Data Beginner

A detector learns from images paired with box annotations, so the first task is to produce that pairing for your domain. Two practical rules dominate. First, label what you will deploy on: photograph the objects in the lighting, angles, backgrounds, and resolution your camera will actually see, because a detector trained on clean catalogue photos fails on a noisy warehouse feed. Second, label consistently: agree on exactly how tight a box should be, how to handle occluded or truncated objects, and which borderline cases count, then apply the rules uniformly, since inconsistent labels put a noise floor under your mAP that no architecture can lift.

Annotation tools (Label Studio, CVAT, Roboflow, and others) export to a handful of standard formats. The two you will meet most are COCO JSON (one file listing all images and annotations, corner-format boxes, used by torchvision and the DETR family) and the YOLO text format (one .txt per image, one line per object: class index and a center-format box normalized to $[0, 1]$, the format and conventions from Section 23.1). Knowing both, and being able to convert between them, saves more debugging hours than any modeling trick. Figure 23.6.1 shows the whole workflow this section follows.

2. Splitting Data Without Fooling Yourself Intermediate

Before training, split your images into training, validation, and test sets. The cardinal rule, the same one from Chapter 21, is that no information may leak from validation or test into training. In detection this rule has a sharp, easy-to-miss edge: if your images come from video or from burst captures, near-identical consecutive frames must all go to the same split. Put frame 100 in training and frame 101 in validation and your validation mAP will be gloriously high and completely meaningless, because the model has effectively seen the validation images. Split by source clip, by scene, or by capture session, never by random per-frame shuffle, whenever frames are correlated.

3. Fine-Tuning a Modern Detector Intermediate

With clean, split data you fine-tune. We use Ultralytics YOLO because it is the most widely deployed toolkit and reduces the whole training loop to a few lines, but the principles transfer to torchvision and the DETR family. The data is described by a small YAML file naming the train and validation image folders and the class names; the model loads pretrained COCO weights (transfer learning from Chapter 21, since COCO features are an excellent starting point for almost any natural-image domain); and a single train call runs the schedule with sensible default augmentation. The code below is a complete, runnable fine-tuning script.

# Fine-tune a COCO-pretrained YOLO on a custom dataset described by a YAML file.
# Transfer learning from COCO means the backbone already sees natural-image
# features, so a short schedule on your own classes is usually enough.
from ultralytics import YOLO

# data.yaml lists: path, train: images/train, val: images/val, names: {0: bolt, 1: nut}
model = YOLO("yolo11n.pt")          # nano model, COCO-pretrained; start small

results = model.train(
    data="data.yaml",
    epochs=100,
    imgsz=640,                      # train/infer resolution; match your deployment
    batch=16,
    lr0=0.01,                       # initial LR; YOLO cosine-decays it automatically
    patience=20,                    # early-stop if val mAP stalls for 20 epochs
    augment=True,                   # mosaic, flips, HSV jitter (Chapter 21 policies)
)
metrics = model.val()               # evaluates on the val set
print(metrics.box.map)              # COCO mAP@[0.50:0.95]
print(metrics.box.map50)            # mAP@0.50

Two recipe choices matter most. Start with the smallest model (the "n" nano variant) to get a fast signal that your pipeline and labels are correct; only scale up to "s", "m", or larger once you have confirmed the data is sound and you need more accuracy. And match the training image size to your deployment resolution: train at $640$ and deploy at $1280$ and the model sees objects at unfamiliar scales. The augmentation defaults (mosaic, horizontal flip, HSV jitter) implement the policies of Chapter 21 and are usually a good starting point, though mosaic is often disabled for the final few epochs so the model finishes on un-augmented images.

Practical Example: The Detector That Learned the Watermark

Who: a retail-analytics startup training a detector to find specific product packages on store shelves from phone photos, 2024. Situation: their training images came from the brand's marketing team and their validation images from a separate field-collection effort. The first model scored an excellent validation mAP and they prepared to ship. Problem: in a pre-launch field test the detector failed badly on ordinary store photos. Investigating, they found the marketing training images all carried a faint corner watermark and a consistent studio background, and the model had partly keyed on those cues; the validation set happened to share the studio look, so validation mAP was inflated. Decision: they re-collected training data in real stores, re-split by store location to prevent any single store's shelf from spanning train and val, and added heavy background and color augmentation to break the spurious cues. Result: validation mAP dropped to a lower but honest number that matched field performance, and the shipped model worked. Lesson: a detector will exploit any shortcut your data offers, including watermarks and backgrounds; an mAP is only as trustworthy as the realism and independence of the set it is measured on, exactly the leakage and easy-set warnings of subsection 2. The illustration below shows the failure: a proud detector boxing the watermark instead of the product.

4. Diagnosing Training Failures Advanced

When a custom detector trains poorly, the cause is almost never the architecture and almost always one of a short list of practical faults. Table 23.6.1 is the checklist experienced practitioners run through; the column order roughly matches how often each cause appears. The most common by far is a data or label problem: a class-index off-by-one, boxes in the wrong format (center versus corner, normalized versus pixel, the conversions from Section 23.1), or annotations that are simply inconsistent. The single most useful debugging action is to render a handful of training images with their loaded labels drawn on top, before any training, and confirm the boxes sit on the objects; this five-minute check catches the majority of "the model will not learn" reports.

# Load the best checkpoint from training, run inference, and export the model
# to deployment runtimes. Each of predict and export is a single call that
# hides the tracing and NMS-graph surgery you would otherwise write by hand.
model = YOLO("runs/detect/train/weights/best.pt")

# Inference on new images: returns boxes, classes, scores per image.
preds = model.predict("new_shelf.jpg", conf=0.4, imgsz=640)
preds[0].show()                       # draw the boxes

# Export for deployment. One line each:
model.export(format="onnx")           # portable ONNX graph, runs anywhere
model.export(format="engine")         # TensorRT engine for NVIDIA edge devices

5. The Whole Chapter, in One Pipeline Advanced

Step back and notice how this section reuses everything before it. The boxes and mAP of Section 23.1 are the labels you create and the metric you read. The architecture you fine-tune is a member of one of the families from Sections 23.2 to 23.5, chosen by the speed-accuracy trade-off those sections laid out: a YOLO for the edge, a Faster R-CNN for small-object recall, a DINO-DETR for top accuracy. The augmentation and transfer learning come from Chapter 21, the export targets are the edge devices of Chapter 28, and the moment you need pixel-precise outlines instead of boxes you will extend this exact workflow into the segmentation of Chapter 24. Detection is the hub, and you have now driven the full loop from raw images to a deployable model. Put the whole chapter into practice in the Hands-On Lab at the end of this section, which walks the entire pipeline once and leaves you with a trained, validated, and exported detector you can show.

pip install ultralytics onnx onnxruntime

# Expected layout:
#   dataset/images/train/*.jpg   dataset/labels/train/*.txt
#   dataset/images/val/*.jpg     dataset/labels/val/*.txt
# Each label .txt has one line per object: "class_idx cx cy w h"
# with cx, cy, w, h normalized to [0, 1] (the YOLO center-format of Section 23.1).
from pathlib import Path

root = Path("dataset")
# TODO: write dataset/data.yaml as text with these keys:
#   path: ./dataset
#   train: images/train
#   val: images/val
#   names: {0: your_class_name}
# Hint: build a short multi-line string and Path("dataset/data.yaml").write_text(...)
yaml_text = ...

import cv2, matplotlib.pyplot as plt

def draw_yolo_labels(img_path, label_path):
    img = cv2.cvtColor(cv2.imread(str(img_path)), cv2.COLOR_BGR2RGB)
    h, w = img.shape[:2]
    for line in open(label_path):
        c, cx, cy, bw, bh = map(float, line.split())
        # TODO: convert the normalized center-format box (cx, cy, bw, bh)
        # to pixel corners (x1, y1, x2, y2) and draw it with cv2.rectangle.
        # Hint: x1 = (cx - bw/2) * w ; y1 = (cy - bh/2) * h ; and so on.
        ...
    return img

plt.imshow(draw_yolo_labels("dataset/images/train/ex.jpg",
                            "dataset/labels/train/ex.txt"))
plt.axis("off"); plt.show()

from ultralytics import YOLO

model = YOLO("yolo11n.pt")   # nano, COCO-pretrained; start small on purpose
# TODO: call model.train(...) with data="dataset/data.yaml", a modest epochs
# count (50 is plenty for a tiny set), imgsz=640, and patience=20 early-stopping.
results = ...

metrics = model.val()
# TODO: print metrics.box.map (COCO mAP@[0.50:0.95]) and metrics.box.map50.
# Then run model.predict on five val images at conf=0.4 and view the boxes.
...

best = YOLO("runs/detect/train/weights/best.pt")
# TODO: export to ONNX with best.export(format="onnx"), then confirm the file
# exists and re-load it with YOLO(...) to run one inference, proving it works.
...

# Complete custom-detector pipeline: data.yaml, label check, train, eval, export.
from pathlib import Path
import cv2, matplotlib.pyplot as plt
from ultralytics import YOLO

# --- Step 1: write data.yaml ---
yaml_text = """\
path: ./dataset
train: images/train
val: images/val
names:
  0: widget
"""
Path("dataset/data.yaml").write_text(yaml_text)

# --- Step 2: render loaded labels to verify format before training ---
def draw_yolo_labels(img_path, label_path):
    img = cv2.cvtColor(cv2.imread(str(img_path)), cv2.COLOR_BGR2RGB)
    h, w = img.shape[:2]
    for line in open(label_path):
        c, cx, cy, bw, bh = map(float, line.split())
        x1 = int((cx - bw / 2) * w); y1 = int((cy - bh / 2) * h)
        x2 = int((cx + bw / 2) * w); y2 = int((cy + bh / 2) * h)
        cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 2)
    return img

plt.imshow(draw_yolo_labels("dataset/images/train/ex.jpg",
                            "dataset/labels/train/ex.txt"))
plt.axis("off"); plt.show()
# (Splitting rule: if images are correlated frames, assign whole clips to one split.)

# --- Step 3: fine-tune the COCO-pretrained nano model ---
model = YOLO("yolo11n.pt")
model.train(
    data="dataset/data.yaml",
    epochs=50,
    imgsz=640,
    batch=16,
    patience=20,      # early-stop if val mAP stalls
)

# --- Step 4: read mAP honestly, then check predictions by eye ---
metrics = model.val()
print("mAP@0.50        :", metrics.box.map50)
print("mAP@[0.50:0.95] :", metrics.box.map)
val_imgs = sorted(Path("dataset/images/val").glob("*.jpg"))[:5]
for r in model.predict(val_imgs, conf=0.4, imgsz=640):
    r.show()

# --- Step 5: export to ONNX and verify the exported model reloads and runs ---
best = YOLO("runs/detect/train/weights/best.pt")
onnx_path = best.export(format="onnx")     # one line; bakes decode + NMS into the graph
print("exported:", onnx_path)
reloaded = YOLO(onnx_path)                  # runs through ONNX Runtime
reloaded.predict(val_imgs[0], conf=0.4)[0].show()