ModelTrainEval/datasets/yolo.py

"""YOLO 风格标注的读写。

目录约定（与 ultralytics 一致，`images` 与 `labels` 同层对应）：

```
data/raw/film/
├── images/train/0001.png
├── labels/train/0001.txt      # 与图片同名，后缀换成 .txt
├── images/val/...
├── labels/val/...
└── data.yaml                  # path / train / val / names
```

标注文件每行一个目标，坐标是**相对图像宽高归一化**到 [0, 1] 的中心点格式：

```
<class_id> <cx> <cy> <w> <h>
```

也接受 YOLO 分割格式的多边形（`<class_id> x1 y1 x2 y2 ...`，偶数个坐标），
读取时取其外接矩形。空文件或缺失文件表示该图无缺陷（纯负样本）。
"""

from __future__ import annotations

from collections.abc import Sequence
from pathlib import Path

import numpy as np

IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp")


def label_path_for(image_path: Path, img_dir: Path, label_dir: Path | None = None) -> Path:
    """图片路径 -> 标注路径。

    缺省按 YOLO 约定把路径里最后一个 `images` 目录名换成 `labels`；
    显式给了 `label_dir` 时，保持图片相对 `img_dir` 的子目录结构。
    """
    relative = image_path.relative_to(img_dir).with_suffix(".txt")
    if label_dir is not None:
        return Path(label_dir) / relative

    parts = list(img_dir.parts)
    for i in range(len(parts) - 1, -1, -1):
        if parts[i] == "images":
            parts[i] = "labels"
            return Path(*parts) / relative
    # 没有 images 目录时退化为同级的 labels/
    return img_dir.parent / "labels" / relative


def list_images(img_dir: Path, list_file: Path | None = None) -> list[Path]:
    """列出图片；给了 `list_file` 则按文件里的路径顺序（YOLO 的 train.txt 风格）。"""
    if list_file is not None:
        list_file = Path(list_file)
        if not list_file.exists():
            raise FileNotFoundError(f"图片清单不存在: {list_file}")
        images = []
        for line in list_file.read_text(encoding="utf-8").splitlines():
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            path = Path(line)
            images.append(path if path.is_absolute() else (list_file.parent / path).resolve())
        return images

    if not img_dir.is_dir():
        raise FileNotFoundError(f"图片目录不存在: {img_dir}")
    return sorted(p for p in img_dir.rglob("*") if p.suffix.lower() in IMAGE_SUFFIXES)


def parse_label_file(path: Path, num_classes: int | None = None) -> tuple[np.ndarray, np.ndarray]:
    """读取一个 .txt，返回归一化的 xyxy 框 (N, 4) 与类别下标 (N,)。

    缺失文件按无缺陷处理（YOLO 的惯例），格式错误则带上 `文件:行号` 抛出。
    """
    boxes: list[list[float]] = []
    labels: list[int] = []
    if not path.exists():
        return np.zeros((0, 4), dtype=np.float32), np.zeros((0,), dtype=np.int64)

    for line_no, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        parts = line.split()
        try:
            class_id = int(float(parts[0]))
            values = [float(v) for v in parts[1:]]
        except ValueError as exc:
            raise ValueError(f"{path}:{line_no} 不是合法的 YOLO 标注行: {line!r}") from exc

        if num_classes is not None and not 0 <= class_id < num_classes:
            raise ValueError(
                f"{path}:{line_no} 类别下标 {class_id} 越界，classes 只有 {num_classes} 个"
            )

        if len(values) == 4:
            cx, cy, width, height = values
            box = [cx - width / 2, cy - height / 2, cx + width / 2, cy + height / 2]
        elif len(values) >= 6 and len(values) % 2 == 0:
            # 分割格式的多边形，取外接矩形
            xs, ys = values[0::2], values[1::2]
            box = [min(xs), min(ys), max(xs), max(ys)]
        else:
            raise ValueError(
                f"{path}:{line_no} 期望 `class cx cy w h` 或分割多边形，实际有 {len(values)} 个坐标"
            )

        box = [min(max(v, 0.0), 1.0) for v in box]
        if box[2] <= box[0] or box[3] <= box[1]:
            continue  # 退化框直接丢弃
        boxes.append(box)
        labels.append(class_id)

    return (
        np.asarray(boxes, dtype=np.float32).reshape(-1, 4),
        np.asarray(labels, dtype=np.int64),
    )


def denormalize(boxes: np.ndarray, width: int, height: int) -> np.ndarray:
    """归一化 xyxy -> 像素 xyxy。"""
    if not len(boxes):
        return np.zeros((0, 4), dtype=np.float32)
    scale = np.array([width, height, width, height], dtype=np.float32)
    return boxes.astype(np.float32) * scale


def format_label_lines(boxes: np.ndarray, labels: Sequence[int], width: int, height: int) -> list[str]:
    """像素 xyxy -> YOLO 文本行（归一化 cx cy w h，保留 6 位小数）。"""
    lines = []
    for (x1, y1, x2, y2), label in zip(np.asarray(boxes, dtype=np.float64), labels, strict=True):
        x1, x2 = sorted((max(x1, 0.0), min(x2, float(width))))
        y1, y2 = sorted((max(y1, 0.0), min(y2, float(height))))
        box_w, box_h = (x2 - x1) / width, (y2 - y1) / height
        if box_w <= 0 or box_h <= 0:
            continue
        cx, cy = (x1 + x2) / 2 / width, (y1 + y2) / 2 / height
        lines.append(f"{int(label)} {cx:.6f} {cy:.6f} {box_w:.6f} {box_h:.6f}")
    return lines


def load_classes(spec: Sequence[str] | str | Path) -> list[str]:
    """类别名可以直接写成列表，也可以指向 YOLO 的 data.yaml（读其中的 `names`）。"""
    if not isinstance(spec, str | Path):
        return list(spec)

    import yaml

    path = Path(spec)
    if not path.exists():
        raise FileNotFoundError(f"data.yaml 不存在: {path}")
    names = (yaml.safe_load(path.read_text(encoding="utf-8")) or {}).get("names")
    if names is None:
        raise KeyError(f"{path} 里没有 `names` 字段")
    if isinstance(names, dict):  # {0: scratch, 1: bubble}
        return [names[key] for key in sorted(names, key=int)]
    return list(names)


def write_data_yaml(
    path: str | Path,
    classes: Sequence[str],
    root: str | Path = ".",
    train: str = "images/train",
    val: str = "images/val",
) -> Path:
    """写一个 YOLO 风格的 data.yaml。"""
    import yaml

    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        yaml.safe_dump(
            {
                "path": str(root),
                "train": train,
                "val": val,
                "nc": len(classes),
                "names": {i: name for i, name in enumerate(classes)},
            },
            allow_unicode=True,
            sort_keys=False,
        ),
        encoding="utf-8",
    )
    return path