计算YOLO数据集中每个类的目标数

这里写自定义目录标题

运行命令：

bash 复制代码

python count_labels.py --data data.yaml --labels labels

计算YOLO数据集中每个类的目标数。类名从data.yaml读取。

bash 复制代码

import argparse
from collections import Counter
from pathlib import Path

import yaml


def count_dir(label_dir: Path):
    """Count (class_id -> object count) and (class_id -> image count) under a dir."""
    obj_count = Counter()
    img_count = Counter()
    files = sorted(label_dir.rglob('*.txt'))
    for f in files:
        cls_in_img = set()
        if f.stat().st_size == 0:
            continue
        for line in f.read_text().strip().splitlines():
            parts = line.split()
            if len(parts) != 5:
                continue
            cls = int(float(parts[0]))
            obj_count[cls] += 1
            cls_in_img.add(cls)
        for c in cls_in_img:
            img_count[c] += 1
    return obj_count, img_count, len(files)


def main():
    p = argparse.ArgumentParser(description='Count YOLO objects per class from data.yaml')
    p.add_argument('--data', required=True, help='/mnt/sda/rdd/datasets_big2small/data640_640/data.yaml')
    p.add_argument('--labels', required=True, help='/mnt/sda/rdd/datasets_big2small/data640_640/labels/train')
    args = p.parse_args()

    with open(args.data) as f:
        cfg = yaml.safe_load(f)
    names = cfg.get('names') or [str(i) for i in range(cfg.get('nc', 1))]
    print(f'classes (nc={len(names)}): {names}\n')

    root = Path(args.labels)
    splits = [d for d in root.iterdir() if d.is_dir()] if root.is_dir() else []
    if splits:
        targets = sorted(splits)
        print(f'found splits: {[d.name for d in targets]}')
        totals = Counter()
        img_totals = Counter()
        for d in targets:
            oc, ic, nf = count_dir(d)
            print(f'\n--- {d.name} ({nf} label files) ---')
            for i, name in enumerate(names):
                print(f'  {i:>2} {name:<12} objects={oc[i]:>6}  images={ic[i]:>6}')
            totals += oc
            img_totals += ic
        print('\n--- TOTAL ---')
        for i, name in enumerate(names):
            print(f'  {i:>2} {name:<12} objects={totals[i]:>6}  images={img_totals[i]:>6}')
    else:
        oc, ic, nf = count_dir(root)
        print(f'--- {root} ({nf} label files) ---')
        for i, name in enumerate(names):
            print(f'  {i:>2} {name:<12} objects={oc[i]:>6}  images={ic[i]:>6}')


if __name__ == '__main__':
    main()