feat(datasets): YOLO 风格标注数据集与几何/像素增强
- yolo: images/ 与同层 labels/ 的路径约定、`class_id cx cy w h` 归一化行解析 （兼容分割多边形取外接矩形）、data.yaml 读写；缺失或空 txt 即无缺陷样本 - film_defect: 抛出图像、框、图像级标签；label_mode 决定多热向量还是互斥下标 - transforms: 图框同步的 Resize / Letterbox / CenterCrop / RandomResizedCrop / 翻转 / ColorJitter / GaussianBlur / PerImageStandardize - target_encoder: 高斯半径与 heatmap / wh / offset / mask 监督信号编码 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
3a36aad6c3
commit
af22dc639c
@@ -0,0 +1,189 @@
|
||||
"""YOLO 风格标注的读写。
|
||||
|
||||
目录约定（与 ultralytics 一致，`images` 与 `labels` 同层对应）：
|
||||
|
||||
```
|
||||
data/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
|
||||
Reference in New Issue
Block a user