YOLO Object Detection skill for Codex

Gives the Codex app a display name, a one-line description and a starter prompt for this skill.

interface:
  display_name: "YOLO Object Detection"
  short_description: "Custom detectors that survive the edge device"
  default_prompt: "$computer-vision-yolo Train a YOLO model to detect scratches and dents from our labelled factory photos"

#!/usr/bin/env python3
"""Check a YOLO-format dataset (detection, segmentation or pose labels) before training.

Usage:
    python3 check_yolo_labels.py path/to/data.yaml

Reads train/val/test image dirs from data.yaml (simple key: value YAML, names as a
mapping or inline list), finds the matching labels/ dirs, and reports errors. Exits 1 on errors.
"""
import ast
import os
import sys
from collections import Counter

IMG_EXT = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}


def parse_yaml(path):
    data, names, in_names = {}, {}, False
    with open(path, encoding="utf-8") as fh:
        for raw in fh:
            line = raw.split("#", 1)[0].rstrip()
            if not line.strip():
                continue
            indented = line[0] in " \t"
            key, _, val = line.strip().partition(":")
            val = val.strip().strip("'\"")
            if in_names and indented:
                names[int(key)] = val
                continue
            in_names = False
            if key == "names":
                if val.startswith("["):
                    items = [x.strip().strip("'\"") for x in val.strip("[]").split(",")]
                    names = dict(enumerate(x for x in items if x))
                else:
                    in_names = True
            elif not indented:
                data[key] = val
    data["names"] = names
    return data


def label_dir_for(img_dir):
    parts = img_dir.split(os.sep)
    if "images" in parts:
        idx = len(parts) - 1 - parts[::-1].index("images")
        parts[idx] = "labels"
        return os.sep.join(parts)
    return os.path.join(os.path.dirname(img_dir), "labels", os.path.basename(img_dir))


def check_split(split, img_dir, n_classes, kpt, errors, warnings):
    lbl_dir = label_dir_for(img_dir)
    if not os.path.isdir(img_dir):
        errors.append(f"{split}: image dir not found: {img_dir}")
        return Counter(), 0, 0
    images = {os.path.splitext(f)[0] for f in os.listdir(img_dir) if os.path.splitext(f)[1].lower() in IMG_EXT}
    labels = set()
    if os.path.isdir(lbl_dir):
        labels = {os.path.splitext(f)[0] for f in os.listdir(lbl_dir) if f.endswith(".txt")}
    for orphan in sorted(labels - images)[:20]:
        errors.append(f"{split}: label without image: {orphan}.txt")
    counts, background = Counter(), 0
    for stem in sorted(images):
        path = os.path.join(lbl_dir, stem + ".txt")
        if not os.path.exists(path):
            background += 1
            continue
        seen = set()
        with open(path, encoding="utf-8") as fh:
            rows = [r.split() for r in fh if r.strip()]
        if not rows:
            background += 1
        for i, parts in enumerate(rows, start=1):
            where = f"{split}/{stem}.txt:{i}"
            try:
                cls = int(parts[0])
                vals = [float(v) for v in parts[1:]]
            except (ValueError, IndexError):
                errors.append(f"{where}: non-numeric values")
                continue
            if not 0 <= cls < n_classes:
                errors.append(f"{where}: class {cls} not in names (0..{n_classes - 1})")
            if len(vals) < 4:
                errors.append(f"{where}: expected at least 4 coordinates, got {len(vals)}")
                continue
            coords = vals
            if kpt and kpt[1] == 3:
                vis = vals[6::3]
                coords = [v for j, v in enumerate(vals) if j < 4 or (j - 4) % 3 != 2]
                if any(v not in (0.0, 1.0, 2.0) for v in vis):
                    errors.append(f"{where}: keypoint visibility must be 0, 1 or 2")
            if any(v < 0 or v > 1.0001 for v in coords):
                errors.append(f"{where}: coordinates outside 0-1 (not normalized?)")
            if kpt:
                if len(vals) != 4 + kpt[0] * kpt[1]:
                    errors.append(f"{where}: expected {4 + kpt[0] * kpt[1]} values for kpt_shape {kpt}, got {len(vals)}")
            elif len(vals) == 4:
                if vals[2] <= 0 or vals[3] <= 0:
                    errors.append(f"{where}: zero or negative box size")
            elif len(vals) % 2 or len(vals) < 6:
                errors.append(f"{where}: polygon needs an even count of at least 6 values, got {len(vals)}")
            key = (cls, tuple(round(v, 4) for v in vals))
            if key in seen:
                warnings.append(f"{where}: duplicate annotation")
            seen.add(key)
            counts[cls] += 1
    return counts, len(images), background


def main():
    if len(sys.argv) != 2:
        print(__doc__)
        return 2
    yaml_path = sys.argv[1]
    cfg = parse_yaml(yaml_path)
    names = cfg["names"]
    if not names:
        print("error: no class names found in data.yaml")
        return 1
    kpt = ast.literal_eval(cfg["kpt_shape"]) if cfg.get("kpt_shape") else None
    root = cfg.get("path", "")
    if not os.path.isabs(root):
        root = os.path.join(os.path.dirname(os.path.abspath(yaml_path)), root) if root else os.path.dirname(os.path.abspath(yaml_path))
        if not os.path.isdir(root) and cfg.get("path"):
            root = os.path.abspath(cfg["path"])
    errors, warnings = [], []
    for split in ("train", "val", "test"):
        if not cfg.get(split):
            if split != "test":
                errors.append(f"data.yaml has no '{split}' entry")
            continue
        img_dir = os.path.normpath(os.path.join(root, cfg[split]))
        counts, n_img, bg = check_split(split, img_dir, len(names), kpt, errors, warnings)
        dist = ", ".join(f"{names.get(c, c)}={counts.get(c, 0)}" for c in sorted(names))
        print(f"{split}: {n_img} images, {bg} background, instances: {dist}")
        for c in sorted(names):
            if split != "train" and n_img and counts.get(c, 0) < 50:
                warnings.append(f"{split}: class '{names[c]}' has only {counts.get(c, 0)} instances")
    for w in warnings[:40]:
        print(f"warn: {w}")
    for e in errors[:100]:
        print(f"error: {e}")
    print(f"\n{len(errors)} error(s), {len(warnings)} warning(s)")
    return 1 if errors else 0


if __name__ == "__main__":
    sys.exit(main())