Training a YOLO object detector on a custom dataset
This post used to be a Darknet tutorial. It walked through training YOLOv3 on a custom dataset, and most of its length went on editing a .cfg file: work out filters = (classes + 5) * 3 for each of the three YOLO layers, set max_batches to classes * 2000, then set steps to 80% and 90% of that. Reading it back, the arithmetic is laid out clearly enough. It is also the single least useful thing I could have spent that many words on.
None of it is worth reading now. pjreddie/darknet has not taken a commit since July 2022, and the repo is not archived so it still looks alive if you only glance. HyperLabel, which I used to draw the boxes, does not resolve at all any more, which means the very first instruction in the old post sends you to a dead domain.
The data outlived the tooling. Same 750 frames from a shipping-port video, same five classes, still sitting on this site where I put them in 2019: yolo-dataset.zip, 20 MB, 750 JPEGs with a matching .txt label each. So I kept the dataset, threw away everything else, and retrained.
Everything below is output from a run I did on a Colab T4. The notebook runs the whole thing top to bottom.
Getting the old data into the new tool
This part turned out to be free, which I did not expect. Darknet writes labels as one row per box, class cx cy w h, normalised to 0-1, and that is exactly the format Ultralytics reads. The 2019 annotations load into a 2026 trainer with no conversion at all. All I had to do was shuffle the files into images/ and labels/ folders and write a small YAML.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import urllib.request, zipfile, pathlib, shutil, random
urllib.request.urlretrieve(
"https://tejashwi.in/assets/files/yolo/yolo-dataset.zip", "yolo-dataset.zip")
with zipfile.ZipFile("yolo-dataset.zip") as z:
z.extractall("raw")
names = [n.strip() for n in open("raw/obj.names") if n.strip()]
src = pathlib.Path("raw/img")
pairs = [(j, j.with_suffix(".txt")) for j in sorted(src.glob("*.jpg"))
if j.with_suffix(".txt").exists()]
root = pathlib.Path("boats")
random.seed(42); random.shuffle(pairs)
cut = int(len(pairs) * 0.8)
for split, items in {"train": pairs[:cut], "val": pairs[cut:]}.items():
(root/"images"/split).mkdir(parents=True, exist_ok=True)
(root/"labels"/split).mkdir(parents=True, exist_ok=True)
for jpg, txt in items:
shutil.copy(jpg, root/"images"/split/jpg.name)
shutil.copy(txt, root/"labels"/split/txt.name)
(root/"data.yaml").write_text(
f"path: {root.resolve()}\ntrain: images/train\nval: images/val\n\n"
f"nc: {len(names)}\nnames: {names}\n")
600 images to train on, 150 held back.
Before starting anything I counted the boxes per class, which is a habit I would recommend to anyone and which I am about to spend most of this post justifying:
1
{'Tug': 186, 'Ferry': 38, 'Sailboat': 364, 'Cruise': 1518, 'Unknown': 83} malformed: 0
Then the training, which is the bit that replaced two thousand words of configuration advice:
1
2
3
4
5
from ultralytics import YOLO
model = YOLO("yolo26n.pt")
model.train(data="boats/data.yaml", epochs=60, imgsz=640,
device=0, batch=16, patience=20)
That is all of it. The class count comes out of the YAML, the detection head gets rebuilt to match, and the COCO weights carry into everything else. Sixty epochs took 773 seconds, so call it thirteen minutes. I had run the same thing on my laptop first, an M3 Pro through Metal, where an epoch took about 70 seconds against the T4’s 13, and I would not bother with the laptop again for a job this size.
The result:
1
2
3
4
5
mAP50 0.778
mAP50-95 0.444
precision 0.617
recall 0.796
inference 4.65 ms per image
Which is a perfectly decent detector for thirteen minutes of work. The 2019 version of this post reported no mAP at all. It showed a photo with boxes on it, and I apparently considered that a result.
The number that matters is not that one
Break 0.778 apart by class and it stops being a single reassuring figure:
| Class | Training boxes | mAP50 |
|---|---|---|
| Cruise | 1518 | 0.983 |
| Sailboat | 364 | 0.929 |
| Unknown | 83 | 0.792 |
| Tug | 186 | 0.689 |
| Ferry | 38 | 0.498 |
The tempting story is that accuracy follows the box count, and Cruise at 1518 against Ferry at 38 looks like a clean demonstration of it. I believed that for about ten minutes, then measured the box sizes and found the classes differ on that too: Cruise objects have a median area of 3.79% of the frame, roughly 249 pixels across, while everything else sits between 0.22% and 0.43%, around 60 to 84 pixels. So Cruise ships are both the most common thing in the harbour and by far the largest, and some unknown share of that 0.983 is just that a 249 pixel object is easy to find.
Among the four small classes the comparison is fairer, since they are all within about 25 pixels of each other. There the counts do mostly govern: Sailboat with 364 examples gets 0.929, Ferry with 38 gets 0.498 despite having the largest median box of the four. Tug and Unknown come out the wrong way round though, so it is a tendency and not a law, and I am not going to pretend the data is tidier than it is.
I did not construct any of this. I labelled frames from a video of a working port, and a working port contains a great many enormous cruise ships and almost no ferries. Both skews arrived with the footage.
That is the thing I would want someone to take from this post. Real data hands you several confounded problems at once and gives you one aggregate number to notice them through, and 0.778 does not look like a model with a class in it scoring 0.498.
There is no clever fix, by the way. Label more ferries. If you cannot, weight the loss towards the rare classes, oversample them, or accept that you are merging Ferry into Cruise whether you write it down or not. I wrote a companion post about the accuracy-first end of this, where a much heavier model gets Ferry to 0.995 and where I go into why I am not willing to say exactly what fixed it.
Running it
1
2
3
4
5
6
7
8
model = YOLO("runs/detect/train/weights/best.pt")
r = model.predict("boats/images/val/output-000000197.jpg", conf=0.35)[0]
counts = {}
for c in r.boxes.cls.tolist():
counts[model.names[int(c)]] = counts.get(model.names[int(c)], 0) + 1
print(counts) # {'Cruise': 2, 'Sailboat': 1}
r.save(filename="prediction.jpg")
On a validation frame it never saw:
Three ships at 0.95, 0.96 and 0.98, which looks like a clean sweep until you check the labels for that frame and find four objects. The missing one is a sailboat about 44 pixels wide, tucked against the left edge, and the model does not return it at any confidence I would accept. Same problem as the table, in a single picture.
If you came here from the old version: ignore the Darknet material entirely. Label Studio and CVAT both do what HyperLabel used to. The dataset still downloads and the labels in it are still good, so you can skip the labelling step entirely and go straight to training.
Discussion
Powered by Disqus