Computer Vision skill for Claude Code

#!/usr/bin/env python3
"""Validate a YOLO-format detection/segmentation dataset (stdlib only).

Usage:
    python3 check_yolo_dataset.py <dataset_dir> [--data data.yaml] [--min-size 0.0005]
                                  [--splits train,val,test] [--show 10]

Expected layout (Ultralytics convention):
    <dataset_dir>/data.yaml
    <dataset_dir>/images/<split>/*.jpg|png|...
    <dataset_dir>/labels/<split>/*.txt     one object per line: class cx cy w h  (normalized)
                                            segmentation lines (class x1 y1 x2 y2 ...) are also checked
data.yaml: reads `names:` as a list ([a, b] or "- a" lines) or a mapping ("0: a" lines),
or `nc:` alone. Only this simple subset of YAML is parsed.

Errors (exit 1): label without image, class id out of range, malformed line,
coordinates outside [0, 1], zero-area boxes, identical image in more than one split.
Warnings: image without label file (allowed for negatives but reported), empty label files,
tiny boxes (w*h < --min-size), duplicate boxes within one file, classes with no instances.
Exits 2 if the dataset cannot be read.
"""
import argparse
import hashlib
import os
import re
import sys
from collections import Counter, defaultdict

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


def parse_names(path):
    if not os.path.exists(path):
        return None
    lines = open(path, encoding="utf-8").read().splitlines()
    names, nc, in_names = None, None, False
    for raw in lines:
        line = raw.split("#", 1)[0].rstrip()
        if not line.strip():
            continue
        m = re.match(r"^nc:\s*(\d+)", line)
        if m:
            nc = int(m.group(1))
            in_names = False
            continue
        m = re.match(r"^names:\s*(.*)$", line)
        if m:
            rest = m.group(1).strip()
            if rest.startswith("["):
                names = [x.strip().strip("'\"") for x in rest.strip("[]").split(",") if x.strip()]
                in_names = False
            else:
                names, in_names = {}, True
            continue
        if in_names and raw.startswith((" ", "\t")):
            item = line.strip()
            m = re.match(r"^(\d+):\s*(.+)$", item)
            if m:
                names[int(m.group(1))] = m.group(2).strip().strip("'\"")
            elif item.startswith("- "):
                names[len(names)] = item[2:].strip().strip("'\"")
            continue
        in_names = False
    if isinstance(names, dict):
        names = [names[k] for k in sorted(names)]
    if names is None and nc is not None:
        names = [str(i) for i in range(nc)]
    return names


def sha1(path):
    h = hashlib.sha1()
    with open(path, "rb") as fh:
        for block in iter(lambda: fh.read(1 << 20), b""):
            h.update(block)
    return h.hexdigest()


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("dataset")
    ap.add_argument("--data", default="data.yaml")
    ap.add_argument("--min-size", type=float, default=0.0005, help="warn when w*h is below this")
    ap.add_argument("--splits", default="train,val,test")
    ap.add_argument("--show", type=int, default=10)
    args = ap.parse_args()

    root = args.dataset
    if not os.path.isdir(os.path.join(root, "images")):
        print(f"error: {root}/images not found", file=sys.stderr)
        sys.exit(2)
    names = parse_names(os.path.join(root, args.data))
    if names is None:
        print(f"warning: could not read class names from {args.data}; class range not checked")
    nclass = len(names) if names else None

    errors, warnings = defaultdict(list), defaultdict(list)
    per_split = {}
    hashes = defaultdict(set)

    for split in [s for s in args.splits.split(",") if s]:
        img_dir = os.path.join(root, "images", split)
        lbl_dir = os.path.join(root, "labels", split)
        if not os.path.isdir(img_dir):
            continue
        imgs = {os.path.splitext(f)[0]: f for f in os.listdir(img_dir)
                if os.path.splitext(f)[1].lower() in IMG_EXT}
        lbls = {os.path.splitext(f)[0]: f for f in os.listdir(lbl_dir)
                if f.endswith(".txt")} if os.path.isdir(lbl_dir) else {}
        counts = Counter()
        n_boxes = 0
        for stem, f in imgs.items():
            hashes[sha1(os.path.join(img_dir, f))].add(f"{split}/{f}")
            if stem not in lbls:
                warnings["image without label (negative?)"].append(f"{split}/{f}")
        for stem, f in lbls.items():
            rel = f"{split}/{f}"
            if stem not in imgs:
                errors["label without image"].append(rel)
            lines = [l for l in open(os.path.join(lbl_dir, f), encoding="utf-8").read().splitlines() if l.strip()]
            if not lines:
                warnings["empty label file"].append(rel)
            seen = set()
            for n, line in enumerate(lines, 1):
                parts = line.split()
                where = f"{rel}:{n}"
                try:
                    cls = int(parts[0])
                    vals = [float(x) for x in parts[1:]]
                except (ValueError, IndexError):
                    errors["malformed line"].append(where)
                    continue
                if len(vals) < 4 or (len(vals) > 4 and len(vals) % 2):
                    errors["malformed line"].append(where)
                    continue
                if cls < 0 or (nclass is not None and cls >= nclass):
                    errors["class id out of range"].append(f"{where} (class {cls})")
                if any(v < 0 or v > 1 for v in vals):
                    errors["coordinate outside [0,1]"].append(where)
                if len(vals) == 4:
                    w, h = vals[2], vals[3]
                    if w <= 0 or h <= 0:
                        errors["zero-area box"].append(where)
                    elif w * h < args.min_size:
                        warnings["tiny box"].append(where)
                key = (cls, tuple(round(v, 4) for v in vals))
                if key in seen:
                    warnings["duplicate box in file"].append(where)
                seen.add(key)
                counts[cls] += 1
                n_boxes += 1
        per_split[split] = (len(imgs), len(lbls), n_boxes, counts)

    for h, files in hashes.items():
        splits = {f.split("/", 1)[0] for f in files}
        if len(splits) > 1:
            errors["identical image in multiple splits"].append(" == ".join(sorted(files)))

    if not per_split:
        print("error: no split directories found under images/", file=sys.stderr)
        sys.exit(2)

    print(f"classes: {nclass if nclass is not None else '?'}" + (f" {names}" if names and len(names) <= 20 else ""))
    print(f"\n{'split':8} {'images':>7} {'labels':>7} {'boxes':>7}")
    for split, (ni, nl, nb, _) in per_split.items():
        print(f"{split:8} {ni:7} {nl:7} {nb:7}")
    all_classes = range(nclass) if nclass else sorted({c for *_, cnt in per_split.values() for c in cnt})
    print("\ninstances per class:")
    header = "  " + f"{'class':20}" + "".join(f"{s:>8}" for s in per_split)
    print(header)
    for c in all_classes:
        label = names[c] if names and c < len(names) else str(c)
        row = [per_split[s][3].get(c, 0) for s in per_split]
        print("  " + f"{label[:20]:20}" + "".join(f"{v:8}" for v in row))
        if per_split.get("train") and per_split["train"][3].get(c, 0) == 0:
            warnings["class with no training instances"].append(label)

    for title, bucket in (("ERRORS", errors), ("WARNINGS", warnings)):
        if bucket:
            print(f"\n{title}:")
            for kind, items in bucket.items():
                print(f"  {kind}: {len(items)}")
                for it in items[: args.show]:
                    print(f"    - {it}")
    if errors:
        print(f"\nFAIL: {sum(len(v) for v in errors.values())} error(s)")
        sys.exit(1)
    print("\nPASS: no errors")


if __name__ == "__main__":
    main()