Skip to content

Data & convert

Bring a YOLO detection dataset into the COCO layout DFINE.train(data=...) consumes.

dfine.convert.yolo_to_coco

yolo_to_coco(yolo_root: str | PathLike, output_dir: str | PathLike, *, class_names: list[str] | None = None, splits: dict[str, str] | None = None, copy_images: bool = True, split_names: dict[str, str] | None = None) -> dict[str, str]

Convert a YOLO detection dataset to the COCO layout under output_dir.

Parameters:

Name Type Description Default
yolo_root str | PathLike

dataset root (with images/<split> + labels/<split>, and an optional data.yaml).

required
output_dir str | PathLike

where the COCO train2017//val2017/ + annotations/ are written (consumable directly by DFINE.train(data=output_dir)).

required
class_names list[str] | None

class names (index = class id). Falls back to data.yaml's names, then to inferred class_<i> if neither is available.

None
splits dict[str, str] | None

explicit {split: image_dir} (relative to yolo_root or absolute). Defaults to auto-detecting images/<split> and <split>/images.

None
copy_images bool

copy images (default) or symlink them into the output.

True
split_names dict[str, str] | None

override the split→folder map (default train→train2017, val→val2017, test→test2017).

None

Returns:

Type Description
dict[str, str]

{output_split_name: annotation_json_path} for each converted split.

Source code in dfine/convert.py
def yolo_to_coco(
    yolo_root: str | os.PathLike,
    output_dir: str | os.PathLike,
    *,
    class_names: list[str] | None = None,
    splits: dict[str, str] | None = None,
    copy_images: bool = True,
    split_names: dict[str, str] | None = None,
) -> dict[str, str]:
    """Convert a YOLO detection dataset to the COCO layout under ``output_dir``.

    Args:
        yolo_root: dataset root (with ``images/<split>`` + ``labels/<split>``, and an
            optional ``data.yaml``).
        output_dir: where the COCO ``train2017/``/``val2017/`` + ``annotations/`` are
            written (consumable directly by ``DFINE.train(data=output_dir)``).
        class_names: class names (index = class id). Falls back to ``data.yaml``'s
            ``names``, then to inferred ``class_<i>`` if neither is available.
        splits: explicit ``{split: image_dir}`` (relative to ``yolo_root`` or absolute).
            Defaults to auto-detecting ``images/<split>`` and ``<split>/images``.
        copy_images: copy images (default) or symlink them into the output.
        split_names: override the split→folder map (default ``train→train2017``,
            ``val→val2017``, ``test→test2017``).

    Returns:
        ``{output_split_name: annotation_json_path}`` for each converted split.
    """
    yolo_root = Path(yolo_root)
    output_dir = Path(output_dir)
    split_names = split_names or _DEFAULT_SPLIT_NAMES

    names = _resolve_class_names(yolo_root, class_names)
    split_dirs = _detect_splits(yolo_root, splits)
    (output_dir / "annotations").mkdir(parents=True, exist_ok=True)

    num_classes_seen: list[int] = []
    coco_splits: dict[str, dict] = {}
    for split, image_dir in split_dirs.items():
        out_name = split_names.get(split, split)
        coco_splits[out_name] = _convert_split(
            image_dir, output_dir / out_name, num_classes_seen, copy_images
        )
        logger.info(
            "converted %s: %d images, %d annotations",
            split,
            len(coco_splits[out_name]["images"]),
            len(coco_splits[out_name]["annotations"]),
        )

    # Resolve the category table: explicit/yaml names, else infer from the labels.
    if names is None:
        n = (max(num_classes_seen) + 1) if num_classes_seen else 0
        names = [f"class_{i}" for i in range(n)]
    elif num_classes_seen and max(num_classes_seen) >= len(names):
        # Guard against a names list shorter than the labels: otherwise annotations would
        # carry a category_id with no `categories` entry and COCO eval fails cryptically.
        raise ValueError(
            f"labels reference class id {max(num_classes_seen)} but only {len(names)} class "
            f"name(s) were provided — pass class_names with at least "
            f"{max(num_classes_seen) + 1} entries, or omit it to infer names from the labels."
        )
    categories = [{"id": i, "name": name} for i, name in enumerate(names)]

    written: dict[str, str] = {}
    for out_name, coco in coco_splits.items():
        coco["categories"] = categories
        ann_path = output_dir / "annotations" / f"instances_{out_name}.json"
        ann_path.write_text(json.dumps(coco))
        written[out_name] = str(ann_path)
    return written