From f9a40327839ec3e4df4a2031e7461af471579956 Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Wed, 14 Oct 2026 17:08:14 +0800 Subject: [PATCH] feat: train fixed-size classifiers and export class probabilities --- README.md | 62 +++++-- classify.py | 8 + .../film_tiles128_mobilenet_v3_small.yaml | 8 + configs/classify/film_tiles128_resnet18.yaml | 20 ++ datasets/tile_classification.py | 81 +++++++++ scripts/compare_openvino_classes.py | 134 ++++++++++++++ scripts/export_classifier_openvino.py | 172 +++++++----------- tests/test_classifier_adapter.py | 83 --------- tests/test_classifier_export.py | 52 ++++++ tests/test_compare_openvino_classes.py | 26 +++ tests/test_tile_classification.py | 40 ++++ 11 files changed, 478 insertions(+), 208 deletions(-) create mode 100644 configs/classify/film_tiles128_mobilenet_v3_small.yaml create mode 100644 configs/classify/film_tiles128_resnet18.yaml create mode 100644 datasets/tile_classification.py create mode 100644 scripts/compare_openvino_classes.py delete mode 100644 tests/test_classifier_adapter.py create mode 100644 tests/test_classifier_export.py create mode 100644 tests/test_compare_openvino_classes.py create mode 100644 tests/test_tile_classification.py diff --git a/README.md b/README.md index 6d758bb..b5c4ff1 100644 --- a/README.md +++ b/README.md @@ -801,36 +801,64 @@ ultralytics 自己还会在 `runs/detect/film/` 下写： 配置支持 `_base_` 继承与命令行 `--set a.b.c=value` 覆盖，见 [utils/config.py](utils/config.py)。 新增骨架 / 数据集只需用注册表装饰器注册，再在 yaml 里按 `name` 引用。 -## 分类模型复用 YOLO OpenVINO 接口 +## 分类模型导出 OpenVINO -生产端只使用类别时，可以将分类器导出成与既有 YOLO IR 相同的输入输出接口： +分类器直接输出按训练类别顺序排列的 Softmax 概率，不适配 YOLO 检测输出： ```bash uv pip install openvino uv run python scripts/export_classifier_openvino.py \ --checkpoint runs/classify/best.pth \ - --template runs/detect/models/film_128_20261013_0902_openvino_model/film_128_20261013_0902.xml \ - --out runs/classify/models/resnet18_128_yolo_adapter + --fp32-weights --out runs/classify/models/resnet18_128_probabilities ``` -`--template` 请指定实际部署的 YOLO `.xml`，其同目录需要有 `metadata.yaml`。 脚本读取分类检查点的配置和同目录 `eval_report.json` 中的类别顺序；也可用 `--classes classes.txt` 或 `--classes eval_report.json` 明确指定分类器的类别顺序。 -两者类别集合不一致、训练预处理不受支持或输入尺寸不匹配时会拒绝导出。 +指定顺序不能与检查点记录矛盾。输入尺寸从检查点读取，不依赖 YOLO 模型。 -产出的 XML/BIN 使用模板文件名，输入为 RGB float32 BCHW、值域 `[0,1]`。 -逐图标准化按训练配置嵌入模型，分类概率按 YOLO 类别顺序排列。 -输出保持 `[1,4+nc,anchors]`：仅第一个候选携带全图伪框及 Softmax 概率，其他候选为零。 -伪框是像素 `cx,cy,w,h`，不表示缺陷位置。脚本不进行 resize，图片必须已切为目标尺寸。 +产出的 XML/BIN 按骨架命名，如 `resnet18.xml/.bin`、`mobilenet_v3_small.xml/.bin`。 +输入为 RGB float32 BCHW、值域 `[0,1]`，输入名为 `images`。 +逐图标准化仅在训练配置使用时嵌入模型。脚本不进行 resize，图片必须已切为目标尺寸。 +输出名为 `probabilities`，形状 `[1,C]`：九类模型是 `[1,9]`，索引直接对应 `classes.txt`。 +每类概率在 `[0,1]` 内，概率和为 1；normal 也返回真实概率，不清零、不做阈值过滤。 -部署时成对替换 XML/BIN；使用元数据的程序也需同步 `metadata.yaml`。 -仓库的 `deploy/cpp` 解码器可直接读取适配模型，生产代码若有额外硬编码则需核对。 -Softmax 分数不同于 YOLO 置信度，需在实际数据上确认过滤阈值。 -默认 top1 为 `normal` 时，所有候选类别分数清零，C++ 在正数置信度阈值下返回空列表， -与 YOLO 的背景行为一致。类别编号保持不变，`normal` 仍占其原通道。 -用 `--background-class` 指定其他背景类别，或 `--keep-background-box` 恢复普通分类输出。 -原图被切成多片时，每片独立分类，仅前景片返回伪框。 +部署端直接对 C 个概率取 argmax；normal 是否表示正常、是否应用阈值由业务层决定。 +现有 `deploy/cpp` 是 YOLO 框解码器，不能用于此分类输出；需要按 `[1,C]` 读取， +不要跳过前四个通道或遍历候选框。模型附带 `metadata.yaml`、`classes.txt` 和 `export.json`。 输出目录必须是新目录，原有模型不会被覆盖。 +默认压缩权重为 FP16；需要优先核对转换前后判别一致性时，加 `--fp32-weights` +保留 FP32 权重，减少分类边界处的舍入差异。输入输出接口仍是 float32。 + +### 在固定 128×128 YOLO 裁片上监督训练分类器 + +```bash +uv run python classify.py --config configs/classify/film_tiles128_resnet18.yaml \ + --work-dir runs/classify_tiles128_resnet18 +uv run python classify.py --config configs/classify/film_tiles128_mobilenet_v3_small.yaml \ + --work-dir runs/classify_tiles128_mobilenet_v3_small +``` + +数据从 `data/film_yolo_128/data.yaml` 读取，保留现有 train/val 划分和类别编号。 +监督标签来自 YOLO 标注文件，空标签表示 normal；多类别裁片取面积最大的框对应类别。 +缺失标签文件会报错，不会自动当背景。图片必须已经是 128×128 PNG，禁止 resize。 +这两份配置使用 RGB / 255，与检测部署输入一致；不进行逐图标准化。 +最佳权重按验证 macro-F1 保存，可用上述导出脚本生成 OpenVINO 分类概率模型。 +分类准确率、转换前后一致率和与 YOLO 的一致率是不同指标，应分别报告。 + +导出后，在同一固定裁片验证集上检查这三个指标（示例）： + +```bash +uv run python scripts/compare_openvino_classes.py \ + --teacher runs/detect/models/film_128_20261013_0902_openvino_model/film_128_20261013_0902.xml \ + --student runs/classify/models/resnet18_tiles128/resnet18.xml \ + --checkpoint runs/classify_tiles128_resnet18/best.pth \ + --conf 0.25 --report runs/classify_tiles128_resnet18/comparison.json +``` + +`--student` 指向实际导出的 XML。这里 `--teacher` 只用于比较，不参与训练。 +YOLO 有有效框时取最高分候选的类别，无有效框记为 normal。比较工具在业务层对分类概率 +使用同一置信度阈值，低于阈值时记为 normal；模型本身仍返回全部真实概率。 +报告同时记录各类支持数、总体/前景一致率和逐图结果。 ## 开发 diff --git a/classify.py b/classify.py index 4e9ccdd..73779a8 100755 --- a/classify.py +++ b/classify.py @@ -16,6 +16,7 @@ import json from pathlib import Path from datasets import FolderClassificationDataset, classification_collate +from datasets.tile_classification import YoloTileClassificationDataset from engine import ClassificationTrainer from models import build_model from utils import load_config, seed_everything, setup_logging @@ -40,6 +41,13 @@ def build_loaders(cfg, split_cfg, classes=None): datasets = {} for split, transforms_key in (("train", "train_transforms"), ("val", "val_transforms")): + if split_cfg.get("format", "folder") == "yolo_tiles": + datasets[split] = YoloTileClassificationDataset( + yaml_path=split_cfg["yaml"], split=split, + image_size=split_cfg["image_size"], transforms=split_cfg[transforms_key], + classes=classes or split_cfg.get("classes"), + ) + continue datasets[split] = FolderClassificationDataset( root=split_cfg["root"], split=split, diff --git a/configs/classify/film_tiles128_mobilenet_v3_small.yaml b/configs/classify/film_tiles128_mobilenet_v3_small.yaml new file mode 100644 index 0000000..4fb650d --- /dev/null +++ b/configs/classify/film_tiles128_mobilenet_v3_small.yaml @@ -0,0 +1,8 @@ +_base_: film_tiles128_resnet18.yaml +work_dir: runs/classify_tiles128_mobilenet_v3_small +model: + backbone: + arch: mobilenet_v3_small +train: + swanlab: + tags: [classify, mobilenet_v3_small, tiles128] diff --git a/configs/classify/film_tiles128_resnet18.yaml b/configs/classify/film_tiles128_resnet18.yaml new file mode 100644 index 0000000..1bb90c3 --- /dev/null +++ b/configs/classify/film_tiles128_resnet18.yaml @@ -0,0 +1,20 @@ +# Train on human YOLO labels, preserving the existing fixed-tile train/val split. +_base_: film.yaml +work_dir: runs/classify_tiles128_resnet18 +data: + root: data/film_yolo_128 + format: yolo_tiles + yaml: data/film_yolo_128/data.yaml + image_size: [128, 128] + classes: [black, blackdot, grease, mosquito, normal, obstacle, ripple, uneven, wrinkle] + train_transforms: + - {name: RandomHorizontalFlip, p: 0.5} + - {name: RandomVerticalFlip, p: 0.5} + - {name: ToTensor} + val_transforms: + - {name: ToTensor} +train: + # Strong inverse-frequency weights plus uniform smoothing bias scores toward rare classes. + label_smoothing: 0.0 + swanlab: + tags: [classify, resnet18, tiles128] diff --git a/datasets/tile_classification.py b/datasets/tile_classification.py new file mode 100644 index 0000000..4c51e57 --- /dev/null +++ b/datasets/tile_classification.py @@ -0,0 +1,81 @@ +"""Supervised classification from fixed-size YOLO tiles and human box labels.""" +from __future__ import annotations + +import math +from pathlib import Path + +import yaml +from PIL import Image + +from datasets.classification import FolderClassificationDataset +from datasets.transforms import build_transforms + + +def tile_label(text: str, num_classes: int, background: int) -> int: + boxes = [] + for line in text.splitlines(): + if not line.strip(): + continue + parts = line.split() + if len(parts) != 5: + raise ValueError(f"Malformed YOLO label: {line}") + cls, cx, cy, w, h = map(float, parts) + if (not all(math.isfinite(v) for v in (cls, cx, cy, w, h)) + or not cls.is_integer() or not 0 <= cls < num_classes + or not 0 <= cx <= 1 or not 0 <= cy <= 1 or not 0 < w <= 1 or not 0 < h <= 1): + raise ValueError(f"Invalid YOLO label: {line}") + boxes.append((int(cls), w * h)) + return max(boxes, key=lambda b: b[1])[0] if boxes else background + + +class YoloTileClassificationDataset(FolderClassificationDataset): + """Keep the existing YOLO split. Empty labels are normal; largest box wins. + + No dataset transform may resample the image. The size is checked on every read. + """ + + def __init__(self, yaml_path, split, image_size=(128, 128), transforms=(), classes=None): + path = Path(yaml_path).resolve() + cfg = yaml.safe_load(path.read_text()) + names = cfg["names"] + self.classes = ([names.get(i, names.get(str(i))) for i in range(len(names))] + if isinstance(names, dict) else list(names)) + if not self.classes or any(not isinstance(n, str) for n in self.classes): + raise ValueError("Invalid YOLO names mapping") + if classes is not None and list(classes) != self.classes: + raise ValueError("Classifier class order must match YOLO names") + background = self.classes.index("normal") + self.class_to_idx = dict(zip(self.classes, range(len(self.classes)), strict=True)) + self.root = Path(cfg.get("path", path.parent)) + if not self.root.is_absolute(): + self.root = path.parent / self.root + folder = self.root / cfg[split] + if not folder.is_dir(): + raise FileNotFoundError(folder) + relative = folder.relative_to(self.root) + if relative.parts[0] != "images": + raise ValueError("YOLO split must be beneath images/") + labels = self.root / "labels" / Path(*relative.parts[1:]) + allowed = {"ToTensor", "RandomHorizontalFlip", "RandomVerticalFlip", "PerImageStandardize"} + if any(t["name"] not in allowed for t in transforms): + raise ValueError("Fixed tiles forbid resampling transforms") + self.transforms = build_transforms(transforms) + self.image_size = tuple(image_size) + self.samples = [] + for image in sorted(folder.rglob("*.png")): + label = labels / image.relative_to(folder).with_suffix(".txt") + if not label.is_file(): + raise FileNotFoundError(f"Missing label (not an implicit background): {label}") + self.samples.append((image, tile_label(label.read_text(), len(self.classes), background))) + if not self.samples: + raise ValueError(f"No PNG tiles in {folder}") + + def __getitem__(self, index): + path, label = self.samples[index] + with Image.open(path) as raw: + if (raw.height, raw.width) != self.image_size: + raise ValueError(f"{path}: size {raw.size}, expected {self.image_size}; resize is forbidden") + image = raw.convert("RGB") + sample = self.transforms({"image": image}) + return {"image": sample["image"], "label": label, + "meta": {"path": str(path), "index": index, "class": self.classes[label]}} diff --git a/scripts/compare_openvino_classes.py b/scripts/compare_openvino_classes.py new file mode 100644 index 0000000..1465294 --- /dev/null +++ b/scripts/compare_openvino_classes.py @@ -0,0 +1,134 @@ +#!/usr/bin/env python +"""Compare supervised classification accuracy, export parity and YOLO agreement. + +Uses fixed-size human-labelled validation tiles, never training data. No resize. +Empty detector output is normal; the strongest valid candidate supplies top1. +""" +from __future__ import annotations + +import argparse +import copy +import json +import sys +from pathlib import Path + +import numpy as np +import torch +import yaml + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from datasets.tile_classification import YoloTileClassificationDataset # noqa: E402 +from models import build_model # noqa: E402 + + +def output_class(output: np.ndarray, conf: float, background: int) -> int: + scores = output[0, 4:, :] + index = int(scores.max(axis=0).argmax()) + return int(scores[:, index].argmax()) if scores[:, index].max() >= conf else background + + +def probability_class(output: np.ndarray, conf: float, background: int) -> int: + if output.ndim != 2 or output.shape[0] != 1: + raise ValueError("Expected classifier probabilities [1,C]") + return int(output[0].argmax()) if output[0].max() >= conf else background + + +def matrix_report(matrix, names, background): + counts = matrix.sum(1) + foreground = int(matrix.sum() - counts[background]) + return {"agreement": float(np.trace(matrix) / matrix.sum()), + "foreground_agreement": float((np.trace(matrix) - matrix[background, background]) + / foreground) if foreground else None, + "per_reference_class": {name: {"support": int(counts[i]), + "agreement": float(matrix[i, i] / counts[i]) + if counts[i] else None} + for i, name in enumerate(names)}, + "confusion_reference_rows_prediction_columns": matrix.tolist()} + + +def main(): + import openvino as ov + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--teacher", required=True, help="Reference YOLO .xml; used only for comparison") + parser.add_argument("--student", required=True, help="Classifier probability .xml, output [1,C]") + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--report", required=True) + parser.add_argument("--conf", type=float, default=0.25) + parser.add_argument("--device", default="cuda:0") + args = parser.parse_args() + if not 0 < args.conf < 1: + raise ValueError("Positive confidence threshold required") + state = torch.load(args.checkpoint, map_location="cpu", weights_only=False) + cfg = state["config"] + names = cfg["data"]["classes"] + background = names.index("normal") + for xml in (args.teacher, args.student): + meta = yaml.safe_load((Path(xml).parent / "metadata.yaml").read_text()) + if [meta["names"][i] for i in range(len(meta["names"]))] != names: + raise ValueError(f"Class order mismatch: {xml}") + if cfg["data"]["val_transforms"] != [{"name": "ToTensor"}]: + raise ValueError("This comparison expects the new raw RGB / 255 tile classifiers") + ds = YoloTileClassificationDataset(cfg["data"]["yaml"], "val", cfg["data"]["image_size"], + transforms=[{"name": "ToTensor"}], classes=names) + model_cfg = copy.deepcopy(cfg["model"]) + model_cfg["backbone"]["pretrained"] = False + model_cfg["init_from"] = None + model_cfg["num_classes"] = len(names) + model = build_model(model_cfg).to(args.device).eval() + model.load_state_dict(state["model"]) + core = ov.Core() + settings = {"PERFORMANCE_HINT": "LATENCY", "INFERENCE_NUM_THREADS": 4} + teacher = core.compile_model(args.teacher, "CPU", settings) + student = core.compile_model(args.student, "CPU", settings) + for compiled in (teacher, student): + if list(compiled.input().shape) != [1, 3, *cfg["data"]["image_size"]]: + raise ValueError("Unexpected model interface") + if (list(teacher.output().shape)[:2] != [1, 4 + len(names)] + or list(student.output().shape) != [1, len(names)]): + raise ValueError("Expected raw YOLO output and plain classifier probabilities") + matrices = {key: np.zeros((len(names), len(names)), dtype=int) + for key in ["pt_accuracy", "ov_accuracy", "ov_deployment_accuracy", + "yolo_deployment_accuracy", + "export_top1_parity", "export_deployment_parity", "yolo_student_agreement"]} + records = [] + with torch.inference_mode(): + for start in range(0, len(ds), 64): + batch = [ds[i] for i in range(start, min(start + 64, len(ds)))] + images = torch.stack([b["image"] for b in batch]) + probabilities = model(images.to(args.device)).softmax(1).cpu().numpy() + for i, item in enumerate(batch): + x = images[i:i+1].numpy() + t = output_class(teacher(x)[teacher.output()], args.conf, background) + raw = student(x)[student.output()] + if (not np.isfinite(raw).all() or (raw < 0).any() or (raw > 1).any() + or not np.allclose(raw.sum(1), 1, atol=1e-5)): + raise ValueError("Invalid classifier probability output") + s = probability_class(raw, args.conf, background) + raw_s = int(raw[0].argmax()) + p = int(probabilities[i].argmax()) + deployed_p = p if probabilities[i].max() >= args.conf else background + y = item["label"] + for key, row, col in [("pt_accuracy", y, p), ("ov_accuracy", y, raw_s), + ("ov_deployment_accuracy", y, s), + ("yolo_deployment_accuracy", y, t), + ("export_top1_parity", p, raw_s), + ("export_deployment_parity", deployed_p, s), + ("yolo_student_agreement", t, s)]: + matrices[key][row, col] += 1 + records.append({"path": item["meta"]["path"], "truth": y, "pytorch": p, + "pytorch_deployed": deployed_p, "openvino": s, + "openvino_top1": raw_s, "yolo": t, + "student_max_score": float(raw.max())}) + report = {"classes": names, "images": len(ds), "conf": args.conf, + "teacher": args.teacher, "student": args.student, "checkpoint": args.checkpoint, + "metrics": {k: matrix_report(m, names, background) for k, m in matrices.items()}} + out = Path(args.report) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(json.dumps(report, indent=2)) + out.with_suffix(".predictions.json").write_text(json.dumps(records, indent=2)) + print(json.dumps(report, indent=2), flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/export_classifier_openvino.py b/scripts/export_classifier_openvino.py index 0490763..21af73d 100644 --- a/scripts/export_classifier_openvino.py +++ b/scripts/export_classifier_openvino.py @@ -1,10 +1,8 @@ #!/usr/bin/env python -"""Export a classifier with the input/output contract of a YOLO OpenVINO IR. +"""Export RGB / 255 BCHW -> Softmax probabilities [1, C], in training class order. -The sole active candidate is a full-image pseudo box, not a localization result. -Input must already be RGB float32 BCHW in [0, 1], exactly as in deploy/cpp. +Training normalization is embedded. No resize, pseudo boxes or background gating. """ - from __future__ import annotations import argparse @@ -19,7 +17,6 @@ import yaml from torch import nn sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - from models import build_model # noqa: E402 @@ -36,63 +33,48 @@ def ordered_names(names) -> list[str]: return names -def class_order(source: list[str], target: list[str]) -> list[int]: - source, target = ordered_names(source), ordered_names(target) - if set(source) != set(target): - raise ValueError(f"Class sets differ: classifier={source}, YOLO={target}") - return [source.index(name) for name in target] +def input_size(config: dict) -> list[int]: + size = config["data"]["image_size"] + size = [size, size] if isinstance(size, int) else list(size) + if len(size) != 2 or any(not isinstance(v, int) or v < 1 for v in size): + raise ValueError(f"Invalid training image_size: {size}") + return size -def preprocessing(config: dict, size: int) -> float | None: - """Reject preprocessing that cannot be reproduced without resampling the input.""" +def preprocessing(config: dict, size: int | list[int]) -> float | None: + shape = [size, size] if isinstance(size, int) else list(size) transforms = config["data"]["val_transforms"] sequence = [t["name"] for t in transforms] if sequence not in (["Resize", "ToTensor", "PerImageStandardize"], ["ToTensor", "PerImageStandardize"], ["Resize", "ToTensor"], ["ToTensor"]): raise ValueError(f"Unsupported validation preprocessing: {sequence}") + eps = None for transform in transforms: - if transform["name"] == "Resize": - if transform["size"] != [size, size]: - raise ValueError("Training size differs from template; retrain at the target size") + if transform["name"] == "Resize" and list(transform["size"]) != shape: + raise ValueError("Training size differs from export input") if transform["name"] == "PerImageStandardize": - return float(transform.get("eps", 1e-5)) - return None + eps = float(transform.get("eps", 1e-5)) + if eps <= 0: + raise ValueError("Standardization eps must be positive") + return eps -class YoloClassifierAdapter(nn.Module): - def __init__(self, classifier: nn.Module, size: int, anchors: int, - order: list[int], eps: float | None = 1e-5, - background_index: int | None = None): +class ClassifierProbabilities(nn.Module): + def __init__(self, classifier: nn.Module, eps: float | None = None): super().__init__() - if size <= 0 or anchors <= 0: - raise ValueError("size and anchors must be positive") - self.classifier = classifier - self.eps = eps - if background_index is not None and not 0 <= background_index < len(order): - raise ValueError("background_index must be an output class index") - self.background_index = background_index - self.register_buffer("order", torch.tensor(order, dtype=torch.long)) - self.register_buffer("box", torch.tensor([size / 2, size / 2, size, size]) - .reshape(1, 4, 1)) - # Only candidate zero is active; do not duplicate predictions across anchors. - self.register_buffer("inactive", torch.zeros(1, 4 + len(order), anchors - 1)) + self.classifier, self.eps = classifier, eps def normalize(self, x: torch.Tensor) -> torch.Tensor: if self.eps is None: return x dims = (1, 2, 3) mean = x.mean(dim=dims, keepdim=True) - # Match datasets.transforms.PerImageStandardize: torch.std correction=1. + # Match training torch.std(correction=1), separately for each image. variance = ((x - mean) ** 2).sum(dim=dims, keepdim=True) / (x[0].numel() - 1) return (x - mean) / (torch.sqrt(variance) + self.eps) def forward(self, x: torch.Tensor) -> torch.Tensor: - scores = self.classifier(self.normalize(x)).softmax(dim=1)[:, self.order] - if self.background_index is not None: - foreground = (scores.argmax(dim=1) != self.background_index).unsqueeze(1) - scores = scores * foreground.to(scores.dtype) - active = torch.cat((self.box.expand(x.shape[0], -1, -1), scores.unsqueeze(2)), dim=1) - return torch.cat((active, self.inactive.expand(x.shape[0], -1, -1)), dim=2) + return self.classifier(self.normalize(x)).softmax(dim=1) def main() -> int: @@ -100,97 +82,71 @@ def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--checkpoint", default="runs/classify/best.pth") - parser.add_argument("--template", required=True, help="Production YOLO .xml (metadata.yaml beside it)") - parser.add_argument("--classes", help="Classifier class order: classes.txt or eval_report.json") - parser.add_argument("--background-class", default="normal", help="Top1 class to emit no boxes for") - parser.add_argument("--keep-background-box", action="store_true", - help="Opt out of background suppression (legacy classification behavior)") + parser.add_argument("--classes", help="Training class order: classes.txt or eval_report.json") + parser.add_argument("--fp32-weights", action="store_true", help="Keep full-precision weights") parser.add_argument("--out", required=True, help="New directory; existing directories are rejected") args = parser.parse_args() - checkpoint, template, out = Path(args.checkpoint), Path(args.template), Path(args.out) + checkpoint, out = Path(args.checkpoint), Path(args.out) if out.exists(): raise FileExistsError(f"Refusing to overwrite {out}") state = torch.load(checkpoint, map_location="cpu", weights_only=False) config = state["config"] + names = config["data"].get("classes") if args.classes: path = Path(args.classes) - source = (json.loads(path.read_text())["classes"] if path.suffix == ".json" - else path.read_text().splitlines()) - else: - source = config["data"].get("classes") - if not source: - report = checkpoint.parent / "eval_report.json" - if not report.exists(): - raise ValueError("No recorded class order; provide --classes (do not guess folder order)") - source = json.loads(report.read_text())["classes"] - source = ordered_names(source) - metadata = yaml.safe_load((template.parent / "metadata.yaml").read_text()) - if metadata.get("end2end") or metadata.get("args", {}).get("nms"): - raise ValueError("Expected raw [1, 4+nc, anchors] output, not end-to-end/NMS output") - target = ordered_names(metadata["names"]) - order = class_order(source, target) - if not args.keep_background_box and args.background_class not in target: - raise ValueError(f"Background class {args.background_class!r} is not in {target}") - background_index = None if args.keep_background_box else target.index(args.background_class) - core = ov.Core() - reference = core.read_model(str(template)) - if len(reference.inputs) != 1 or len(reference.outputs) != 1: - raise ValueError("Expected one input and one output") - input_shape, output_shape = list(reference.input().shape), list(reference.output().shape) - if (len(input_shape) != 4 or input_shape[:2] != [1, 3] - or input_shape[2] != input_shape[3] or len(output_shape) != 3 - or output_shape[:2] != [1, 4 + len(target)] or output_shape[2] < 1): - raise ValueError(f"Unsupported template shapes: {input_shape}, {output_shape}") - if (reference.input().element_type != ov.Type.f32 - or reference.output().element_type != ov.Type.f32): - raise ValueError("Only float32 input/output are supported") - size = input_shape[2] + supplied = (json.loads(path.read_text())["classes"] if path.suffix == ".json" + else path.read_text().splitlines()) + if names and ordered_names(supplied) != ordered_names(names): + raise ValueError("--classes contradicts checkpoint class order") + names = supplied + if not names: + report = checkpoint.parent / "eval_report.json" + if not report.exists(): + raise ValueError("No recorded class order; provide --classes (do not guess folder order)") + names = json.loads(report.read_text())["classes"] + names = ordered_names(names) + size = input_size(config) eps = preprocessing(config, size) model_config = copy.deepcopy(config["model"]) model_config["backbone"]["pretrained"] = False model_config["init_from"] = None - model_config["num_classes"] = len(source) + model_config["num_classes"] = len(names) model = build_model(model_config).eval() model.load_state_dict(state["model"], strict=True) - adapter = YoloClassifierAdapter(model, size, output_shape[2], order, eps, - background_index=background_index).eval() + wrapper = ClassifierProbabilities(model, eps).eval() + shape = [1, model_config["backbone"].get("in_channels", 3), *size] torch.manual_seed(42) - example = torch.rand(input_shape) - converted = ov.convert_model(adapter, example_input=example, input=input_shape) - converted.input().get_tensor().set_names(reference.input().get_names()) - converted.output().get_tensor().set_names(reference.output().get_names()) - # Preserve any interface metadata used by OpenVINO consumers, then mark pseudo boxes. - for key, value in reference.get_rt_info().items(): - if key != "conversion_parameters": - converted.set_rt_info(value, [key]) - converted.set_rt_info("classifier_pseudo_box", ["adapter_kind"]) + example = torch.rand(shape) + converted = ov.convert_model(wrapper, example_input=example, input=shape) + converted.input().get_tensor().set_names({"images"}) + converted.output().get_tensor().set_names({"probabilities"}) + converted.set_rt_info("classify", ["model_info", "task"]) + converted.set_rt_info(json.dumps(dict(enumerate(names))), ["model_info", "names"]) out.mkdir(parents=True) - xml = out / template.name - ov.save_model(converted, str(xml), compress_to_fp16=True) - compiled = core.compile_model(str(xml), "CPU", {"PERFORMANCE_HINT": "LATENCY"}) + arch = model_config["backbone"].get("arch", model_config["backbone"]["name"]) + xml = out / f"{arch}.xml" + ov.save_model(converted, str(xml), compress_to_fp16=not args.fp32_weights) + compiled = ov.Core().compile_model(str(xml), "CPU", {"PERFORMANCE_HINT": "LATENCY"}) errors = [] with torch.inference_mode(): for sample in (example, torch.zeros_like(example), torch.ones_like(example)): - expected = adapter(sample).numpy() + expected = wrapper(sample).numpy() actual = compiled(sample.numpy())[compiled.output()] + assert actual.shape == (1, len(names)) + assert np.isfinite(actual).all() and (actual >= 0).all() and (actual <= 1).all() + np.testing.assert_allclose(actual.sum(axis=1), 1, atol=1e-5) np.testing.assert_allclose(actual, expected, atol=0.02, rtol=0.02) - assert np.argmax(actual[0, 4:, 0]) == np.argmax(expected[0, 4:, 0]) - assert np.count_nonzero(actual[:, :, 1:]) == 0 + assert actual.argmax(axis=1).tolist() == expected.argmax(axis=1).tolist() errors.append(float(np.max(np.abs(expected - actual)))) - metadata.update({"names": dict(enumerate(target)), "adapter_kind": "classifier_pseudo_box", - "description": "Classifier with one full-image pseudo box; no localization", - "background_class": None if background_index is None else args.background_class, - "end2end": False, "args": {"nms": False}, "imgsz": [size, size]}) + metadata = {"task": "classify", "names": dict(enumerate(names)), "imgsz": size, + "channels": shape[1], "layout": "BCHW", "input_range": [0, 1], + "output": "softmax_probabilities", "input_shape": shape, + "output_shape": [1, len(names)], "standardize_eps": eps, + "weight_storage": "fp32" if args.fp32_weights else "fp16"} (out / "metadata.yaml").write_text(yaml.safe_dump(metadata, allow_unicode=True)) - (out / "classes.txt").write_text("\n".join(target) + "\n") - summary = {"checkpoint": str(checkpoint.resolve()), "template": str(template.resolve()), - "classifier_classes": source, "output_classes": target, "class_order": order, - "input_shape": input_shape, "output_shape": output_shape, - "standardize_eps": eps, "max_abs_errors": errors, - "background_class": None if background_index is None else args.background_class, - "background_index": background_index, - "pseudo_box": [size / 2, size / 2, size, size]} - (out / "adapter.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2)) + (out / "classes.txt").write_text("\n".join(names) + "\n") + summary = {"checkpoint": str(checkpoint.resolve()), **metadata, "max_abs_errors": errors} + (out / "export.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2)) print(json.dumps({"model": str(xml), **summary}, ensure_ascii=False, indent=2)) return 0 diff --git a/tests/test_classifier_adapter.py b/tests/test_classifier_adapter.py deleted file mode 100644 index a1868dd..0000000 --- a/tests/test_classifier_adapter.py +++ /dev/null @@ -1,83 +0,0 @@ -import pytest -import torch -from torch import nn - -from scripts.export_classifier_openvino import YoloClassifierAdapter, class_order, preprocessing - - -class FixedClassifier(nn.Module): - def forward(self, x): - return torch.tensor([1.0, 3.0, -1.0]).expand(x.shape[0], -1) - - -def test_yolo_contract_and_class_reordering(): - adapter = YoloClassifierAdapter(FixedClassifier(), 128, 336, [2, 0, 1]) - result = adapter(torch.rand(1, 3, 128, 128)) - assert result.shape == (1, 7, 336) - torch.testing.assert_close(result[0, :4, 0], torch.tensor([64., 64., 128., 128.])) - expected = torch.tensor([1., 3., -1.]).softmax(0)[[2, 0, 1]] - torch.testing.assert_close(result[0, 4:, 0], expected) - assert result[0, 4:, 0].argmax().item() == 2 - assert torch.count_nonzero(result[:, :, 1:]) == 0 - - -def test_normalization_matches_training_for_each_image(): - adapter = YoloClassifierAdapter(FixedClassifier(), 128, 1, [0, 1, 2]) - images = torch.stack((torch.rand(3, 128, 128), torch.ones(3, 128, 128))) - expected = torch.stack([(x - x.mean()) / (x.std() + 1e-5) for x in images]) - torch.testing.assert_close(adapter.normalize(images), expected) - assert torch.isfinite(adapter(images)).all() - - -def test_class_sets_must_match(): - assert class_order(["normal", "mosquito"], ["mosquito", "normal"]) == [1, 0] - with pytest.raises(ValueError, match="Class sets differ"): - class_order(["normal", "mosquito"], ["normal", "blackdot"]) - with pytest.raises(ValueError, match="Duplicate"): - class_order(["normal", "normal"], ["normal"]) - - -def test_background_suppression_uses_reordered_output_index(): - # Source class 1 wins, but it is output class 2 after reordering. - adapter = YoloClassifierAdapter(FixedClassifier(), 128, 336, [2, 0, 1], background_index=2) - result = adapter(torch.rand(1, 3, 128, 128)) - assert result.shape == (1, 7, 336) - assert torch.count_nonzero(result[:, 4:, :]) == 0 - # The existing C++ decoder returns zero candidates at any positive threshold. - assert not (result[0, 4:, :].max(dim=0).values >= 0.25).any() - - -def test_foreground_class_indices_and_scores_are_preserved(): - foreground = YoloClassifierAdapter(FixedClassifier(), 128, 336, [0, 1, 2], - background_index=0) - legacy = YoloClassifierAdapter(FixedClassifier(), 128, 336, [0, 1, 2]) - x = torch.rand(1, 3, 128, 128) - torch.testing.assert_close(foreground(x), legacy(x)) - assert foreground(x)[0, 4:, 0].argmax() == 1 - - -def test_background_gate_is_per_image_in_batch(): - class BatchClassifier(nn.Module): - def forward(self, x): - return x[:, :, 0, 0] - - x = torch.zeros(2, 3, 4, 4) - x[0, 0, 0, 0] = 3 # normal - x[1, 1, 0, 0] = 3 # blackdot - adapter = YoloClassifierAdapter(BatchClassifier(), 4, 10, [0, 1, 2], - eps=None, background_index=0) - result = adapter(x) - assert not torch.count_nonzero(result[0, 4:, :]) - assert result[1, 4:, 0].argmax() == 1 - assert result[1, 5, 0] > 0.9 - - -def test_refuses_resampling_and_unknown_preprocessing(): - config = {"data": {"val_transforms": [ - {"name": "Resize", "size": [160, 160]}, {"name": "ToTensor"}, - {"name": "PerImageStandardize"}]}} - with pytest.raises(ValueError, match="Training size differs"): - preprocessing(config, 128) - config["data"]["val_transforms"] = [{"name": "Normalize"}] - with pytest.raises(ValueError, match="Unsupported"): - preprocessing(config, 128) diff --git a/tests/test_classifier_export.py b/tests/test_classifier_export.py new file mode 100644 index 0000000..870e9cb --- /dev/null +++ b/tests/test_classifier_export.py @@ -0,0 +1,52 @@ +import pytest +import torch +from torch import nn + +from scripts.export_classifier_openvino import ClassifierProbabilities, ordered_names, preprocessing + + +class FixedClassifier(nn.Module): + def forward(self, x): + return torch.tensor([1.0, 3.0, -1.0]).expand(x.shape[0], -1) + + +def test_probabilities_keep_training_class_order_and_sum_to_one(): + result = ClassifierProbabilities(FixedClassifier())(torch.rand(2, 3, 128, 128)) + assert result.shape == (2, 3) + torch.testing.assert_close(result[0], torch.tensor([1., 3., -1.]).softmax(0)) + torch.testing.assert_close(result.sum(1), torch.ones(2)) + assert result.argmax(1).tolist() == [1, 1] + + +def test_normal_probability_is_not_suppressed(): + # Suppose class 1 is normal: its probability is returned like any other class. + result = ClassifierProbabilities(FixedClassifier())(torch.zeros(1, 3, 128, 128)) + assert result[0, 1] > .8 + assert (result > 0).all() + + +def test_standardization_matches_training_per_image(): + exported = ClassifierProbabilities(FixedClassifier(), eps=1e-5) + images = torch.stack((torch.rand(3, 128, 128), torch.ones(3, 128, 128))) + expected = torch.stack([(x-x.mean())/(x.std()+1e-5) for x in images]) + torch.testing.assert_close(exported.normalize(images), expected) + assert torch.isfinite(exported(images)).all() + + +def test_classes_validate_index_order_and_duplicates(): + assert ordered_names({1: "blackdot", 0: "normal"}) == ["normal", "blackdot"] + with pytest.raises(ValueError, match="Duplicate"): + ordered_names(["normal", "normal"]) + with pytest.raises(ValueError, match="contiguous"): + ordered_names({1: "normal"}) + + +def test_refuses_resampling_and_unknown_preprocessing(): + config = {"data": {"val_transforms": [ + {"name": "Resize", "size": [160, 160]}, {"name": "ToTensor"}, + {"name": "PerImageStandardize"}]}} + with pytest.raises(ValueError, match="Training size differs"): + preprocessing(config, 128) + config["data"]["val_transforms"] = [{"name": "Normalize"}] + with pytest.raises(ValueError, match="Unsupported"): + preprocessing(config, 128) diff --git a/tests/test_compare_openvino_classes.py b/tests/test_compare_openvino_classes.py new file mode 100644 index 0000000..34f9797 --- /dev/null +++ b/tests/test_compare_openvino_classes.py @@ -0,0 +1,26 @@ +import numpy as np + +from scripts.compare_openvino_classes import matrix_report, output_class, probability_class + + +def test_detector_output_channel_offset_and_background(): + output = np.zeros((1, 13, 336), dtype=np.float32) + output[0, :4, 0] = [64, 64, 128, 128] + assert output_class(output, 0.25, 4) == 4 + output[0, 5, 0] = 0.6 + assert output_class(output, 0.25, 4) == 1 # blackdot, not black + output[0, 4, 50] = 0.9 + assert output_class(output, 0.25, 4) == 0 + + +def test_report_excludes_background_from_foreground_agreement(): + matrix = np.array([[1, 1], [0, 8]]) + report = matrix_report(matrix, ["blackdot", "normal"], 1) + assert report["agreement"] == 0.9 + assert report["foreground_agreement"] == 0.5 + + +def test_probability_output_has_no_coordinate_offset(): + probabilities = np.array([[0.1, 0.8, 0.1]]) + assert probability_class(probabilities, .25, 2) == 1 + assert probability_class(np.array([[.3, .4, .3]]), .5, 2) == 2 diff --git a/tests/test_tile_classification.py b/tests/test_tile_classification.py new file mode 100644 index 0000000..13f5687 --- /dev/null +++ b/tests/test_tile_classification.py @@ -0,0 +1,40 @@ +import pytest +from PIL import Image + +from datasets.tile_classification import YoloTileClassificationDataset, tile_label +from datasets.yolo import write_data_yaml + + +def test_label_uses_human_class_ids_and_empty_background(): + assert tile_label("", 9, 4) == 4 + assert tile_label("1 0.5 0.5 0.2 0.2", 9, 4) == 1 + assert tile_label("1 0.5 0.5 0.2 0.2\n3 0.5 0.5 0.4 0.4", 9, 4) == 3 + with pytest.raises(ValueError): + tile_label("1 0.5 0.5 nan 0.2", 9, 4) + + +def test_existing_split_is_preserved_and_no_resize(tmp_path): + for split in ("train", "val"): + (tmp_path / "images" / split).mkdir(parents=True) + (tmp_path / "labels" / split).mkdir(parents=True) + Image.new("RGB", (16, 16)).save(tmp_path / "images" / split / "sample.png") + (tmp_path / "labels" / split / "sample.txt").write_text( + "1 0.5 0.5 0.2 0.2" if split == "train" else "") + yaml = write_data_yaml(tmp_path / "data.yaml", ["normal", "blackdot"], root=tmp_path) + kwargs = {"yaml_path": yaml, "image_size": [16, 16], "transforms": [{"name": "ToTensor"}]} + train = YoloTileClassificationDataset(split="train", **kwargs) + val = YoloTileClassificationDataset(split="val", **kwargs) + assert train.class_counts() == [0, 1] + assert val.class_counts() == [1, 0] + assert train[0]["image"].shape == (3, 16, 16) + Image.new("RGB", (12, 16)).save(train.samples[0][0]) + with pytest.raises(ValueError, match="resize is forbidden"): + train[0] + + +def test_missing_labels_are_not_silent_background(tmp_path): + (tmp_path / "images" / "train").mkdir(parents=True) + Image.new("RGB", (16, 16)).save(tmp_path / "images" / "train" / "sample.png") + yaml = write_data_yaml(tmp_path / "data.yaml", ["normal", "blackdot"], root=tmp_path) + with pytest.raises(FileNotFoundError, match="Missing label"): + YoloTileClassificationDataset(yaml, "train")