train_visdrone.py

"""Exact training recipe for LibreYOLO9P2s-visdrone.

Reproduces the released checkpoint: yolo9_p2-s, VisDrone2019-DET, 768 px,
60 epochs, transfer init from stock LibreYOLO9s.pt.

Prepare the dataset first with build_visdrone.py, then:

    python train_visdrone.py --data /abs/path/to/visdrone/visdrone.yaml

Notes discovered the hard way:
- lr0=0.01 (the family default) DIVERGES on transfer init; 0.005 is stable.
- Mosaic/mixup HURT on VisDrone (tiling shrinks tiny objects below
  detectability); mild hsv + horizontal flip help.
- max_labels must be raised: dense aerial frames exceed the default 100-box
  cap, silently dropping ground truth.
- Pass the dataset yaml by ABSOLUTE path.
"""

import argparse

from libreyolo import LibreYOLO9P2


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data", required=True, help="Absolute path to visdrone.yaml")
    parser.add_argument("--epochs", type=int, default=60)
    parser.add_argument("--batch", type=int, default=2)
    parser.add_argument("--imgsz", type=int, default=768)
    parser.add_argument("--workers", type=int, default=2)
    parser.add_argument("--name", default="visdrone_p2s_768")
    args = parser.parse_args()

    model = LibreYOLO9P2(None, size="s")
    model.train(
        data=args.data,
        epochs=args.epochs,
        batch=args.batch,
        nbs=16,  # effective batch 16 via gradient accumulation
        imgsz=args.imgsz,
        workers=args.workers,
        lr0=0.005,
        warmup_epochs=5,
        mosaic_prob=0.0,
        mixup_prob=0.0,
        hsv_prob=1.0,
        flip_prob=0.5,
        max_labels=600,
        pretrained="LibreYOLO9s.pt",
        name=args.name,
        exist_ok=True,
        eval_interval=5,
        save_period=5,
    )


if __name__ == "__main__":
    main()