feat: 新增缺陷分类训练与评估链路
在聚类 + 人工归类得到的数据集（data/<集>/<类别名>/*.png）上训练分类模型。 - datasets/classification.py：按目录组织的分类数据集，**逐类分层**切分 train/val （others 只有 65 张，整体随机切很容易让它在验证集里挂零），固定种子可复现 - models/classifier.py：卷积骨架 + 分类头，与 DefectNet 共用 backbone 注册表， 因此自监督产出的 simclr_backbone.pth 可用 init_from 直接初始化； frozen_backbone 可做线性探针 - engine/classification.py：训练循环 + 指标。主指标用 macro-F1 而非准确率—— 数据从 normal 2679 张到 others 65 张差 41 倍，全预测成 normal 也有 24% 准确率； 损失按类别频率加权，避免小类被淹没；输出逐类 P/R/F1 与混淆矩阵 - classify.py：训练/评估入口，--eval-only 可单独评估并导出 eval_report.json - configs/classify/film.yaml：预处理沿用聚类阶段的结论（拉伸 + 逐图标准化）。 抹掉亮度在这里尤其重要：有几类几乎只来自单一原始数据集，不抹掉的话模型可能 靠"认出这是哪个数据集"作弊 - 续训语义与检测/自监督一致：检查点只提供进度，超参以配置为准 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
33d28dcc7c
commit
33fdf396ee
@@ -1,6 +1,7 @@
|
||||
"""数据集：抛出样本的 Dataset 类与相应的数据变换、标签编码。"""
|
||||
|
||||
from datasets.build import build_dataloader, build_dataset
|
||||
from datasets.classification import FolderClassificationDataset, classification_collate
|
||||
from datasets.film_defect import FilmDefectDataset, collate_fn
|
||||
from datasets.registry import DATASETS
|
||||
from datasets.target_encoder import CenterTargetEncoder
|
||||
@@ -10,6 +11,8 @@ from datasets.yolo import format_label_lines, load_classes, parse_label_file, wr
|
||||
|
||||
__all__ = [
|
||||
"build_dataset",
|
||||
"FolderClassificationDataset",
|
||||
"classification_collate",
|
||||
"build_dataloader",
|
||||
"FilmDefectDataset",
|
||||
"collate_fn",
|
||||
|
||||
@@ -0,0 +1,122 @@
|
||||
"""按目录组织的分类数据集：`root/<类别名>/*.png`，子目录名即标签。
|
||||
|
||||
这是 scripts/export_clusters.py 的产物结构——聚类 + 人工归类之后得到的数据集。
|
||||
没有预先划分 train/val，这里按类别分层、用固定随机种子切分，保证两次运行拿到同一份划分。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from datasets.registry import DATASETS
|
||||
from datasets.transforms import build_transforms
|
||||
from datasets.yolo import IMAGE_SUFFIXES
|
||||
|
||||
SPLITS = ("train", "val", "all")
|
||||
|
||||
|
||||
@DATASETS.register("folder_classification")
|
||||
class FolderClassificationDataset(Dataset):
|
||||
"""Args:
|
||||
root: 数据集根目录，其下每个子目录是一个类别。
|
||||
split: "train" / "val" / "all"。
|
||||
val_ratio: 验证集比例，**按类别分别切分**——小类只有几十张，随机整体切分
|
||||
很容易让某个类在验证集里一张都没有。
|
||||
seed: 划分用的随机种子；同一个 seed 两次运行得到完全相同的划分。
|
||||
classes: 显式指定类别顺序；缺省按目录名排序。训练与评估必须一致，
|
||||
否则标签下标对不上。
|
||||
transforms: transform 配置列表，见 `datasets.transforms.build_transforms`。
|
||||
exclude: 要跳过的子目录名。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
root: str | Path,
|
||||
split: str = "train",
|
||||
val_ratio: float = 0.2,
|
||||
seed: int = 42,
|
||||
classes: Sequence[str] | None = None,
|
||||
transforms: Sequence[Mapping[str, Any]] = (),
|
||||
gray: bool = False,
|
||||
exclude: Sequence[str] = (),
|
||||
) -> None:
|
||||
if split not in SPLITS:
|
||||
raise ValueError(f"未知 split: {split}，可选 {SPLITS}")
|
||||
self.root = Path(root)
|
||||
if not self.root.is_dir():
|
||||
raise FileNotFoundError(f"数据集目录不存在: {self.root}")
|
||||
|
||||
self.split = split
|
||||
self.gray = gray
|
||||
self.transforms = build_transforms(transforms)
|
||||
|
||||
excluded = set(exclude)
|
||||
found = sorted(
|
||||
p.name for p in self.root.iterdir() if p.is_dir() and p.name not in excluded
|
||||
)
|
||||
self.classes = list(classes) if classes else found
|
||||
if not self.classes:
|
||||
raise FileNotFoundError(f"{self.root} 下没有找到类别子目录")
|
||||
missing = [c for c in self.classes if not (self.root / c).is_dir()]
|
||||
if missing:
|
||||
raise FileNotFoundError(f"以下类别在 {self.root} 下不存在: {missing}")
|
||||
self.class_to_idx = {name: i for i, name in enumerate(self.classes)}
|
||||
|
||||
self.samples = self._collect(val_ratio, seed)
|
||||
if not self.samples:
|
||||
raise ValueError(f"{self.root} 的 {split} 划分里没有样本")
|
||||
|
||||
def _collect(self, val_ratio: float, seed: int) -> list[tuple[Path, int]]:
|
||||
samples: list[tuple[Path, int]] = []
|
||||
for label, name in enumerate(self.classes):
|
||||
files = sorted(
|
||||
p
|
||||
for p in (self.root / name).rglob("*")
|
||||
if p.is_file() and p.suffix.lower() in IMAGE_SUFFIXES
|
||||
)
|
||||
if not files:
|
||||
continue
|
||||
# 逐类切分：小类只有几十张，整体随机切很容易让它在验证集里挂零
|
||||
generator = torch.Generator().manual_seed(seed + label)
|
||||
order = torch.randperm(len(files), generator=generator).tolist()
|
||||
n_val = max(1, round(len(files) * val_ratio)) if len(files) > 1 else 0
|
||||
val_idx = set(order[:n_val])
|
||||
for index, path in enumerate(files):
|
||||
in_val = index in val_idx
|
||||
if self.split == "all" or (self.split == "val") == in_val:
|
||||
samples.append((path, label))
|
||||
return samples
|
||||
|
||||
def class_counts(self) -> list[int]:
|
||||
counts = [0] * len(self.classes)
|
||||
for _, label in self.samples:
|
||||
counts[label] += 1
|
||||
return counts
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, index: int) -> dict:
|
||||
path, label = self.samples[index]
|
||||
with Image.open(path) as raw:
|
||||
image = raw.convert("L" if self.gray else "RGB")
|
||||
sample = self.transforms({"image": image})
|
||||
return {
|
||||
"image": sample["image"],
|
||||
"label": label,
|
||||
"meta": {"path": str(path), "index": index, "class": self.classes[label]},
|
||||
}
|
||||
|
||||
|
||||
def classification_collate(batch: list[dict]) -> dict:
|
||||
return {
|
||||
"images": torch.stack([item["image"] for item in batch]),
|
||||
"labels": torch.tensor([item["label"] for item in batch], dtype=torch.long),
|
||||
"meta": [item["meta"] for item in batch],
|
||||
}
|
||||
Reference in New Issue
Block a user