feat: train fixed-size classifiers and export class probabilities
This commit is contained in:
@@ -801,36 +801,64 @@ ultralytics 自己还会在 `runs/detect/film/` 下写:
|
|||||||
配置支持 `_base_` 继承与命令行 `--set a.b.c=value` 覆盖,见 [utils/config.py](utils/config.py)。
|
配置支持 `_base_` 继承与命令行 `--set a.b.c=value` 覆盖,见 [utils/config.py](utils/config.py)。
|
||||||
新增骨架 / 数据集只需用注册表装饰器注册,再在 yaml 里按 `name` 引用。
|
新增骨架 / 数据集只需用注册表装饰器注册,再在 yaml 里按 `name` 引用。
|
||||||
|
|
||||||
## 分类模型复用 YOLO OpenVINO 接口
|
## 分类模型导出 OpenVINO
|
||||||
|
|
||||||
生产端只使用类别时,可以将分类器导出成与既有 YOLO IR 相同的输入输出接口:
|
分类器直接输出按训练类别顺序排列的 Softmax 概率,不适配 YOLO 检测输出:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
uv pip install openvino
|
uv pip install openvino
|
||||||
uv run python scripts/export_classifier_openvino.py \
|
uv run python scripts/export_classifier_openvino.py \
|
||||||
--checkpoint runs/classify/best.pth \
|
--checkpoint runs/classify/best.pth \
|
||||||
--template runs/detect/models/film_128_20261013_0902_openvino_model/film_128_20261013_0902.xml \
|
--fp32-weights --out runs/classify/models/resnet18_128_probabilities
|
||||||
--out runs/classify/models/resnet18_128_yolo_adapter
|
|
||||||
```
|
```
|
||||||
|
|
||||||
`--template` 请指定实际部署的 YOLO `.xml`,其同目录需要有 `metadata.yaml`。
|
|
||||||
脚本读取分类检查点的配置和同目录 `eval_report.json` 中的类别顺序;也可用
|
脚本读取分类检查点的配置和同目录 `eval_report.json` 中的类别顺序;也可用
|
||||||
`--classes classes.txt` 或 `--classes eval_report.json` 明确指定分类器的类别顺序。
|
`--classes classes.txt` 或 `--classes eval_report.json` 明确指定分类器的类别顺序。
|
||||||
两者类别集合不一致、训练预处理不受支持或输入尺寸不匹配时会拒绝导出。
|
指定顺序不能与检查点记录矛盾。输入尺寸从检查点读取,不依赖 YOLO 模型。
|
||||||
|
|
||||||
产出的 XML/BIN 使用模板文件名,输入为 RGB float32 BCHW、值域 `[0,1]`。
|
产出的 XML/BIN 按骨架命名,如 `resnet18.xml/.bin`、`mobilenet_v3_small.xml/.bin`。
|
||||||
逐图标准化按训练配置嵌入模型,分类概率按 YOLO 类别顺序排列。
|
输入为 RGB float32 BCHW、值域 `[0,1]`,输入名为 `images`。
|
||||||
输出保持 `[1,4+nc,anchors]`:仅第一个候选携带全图伪框及 Softmax 概率,其他候选为零。
|
逐图标准化仅在训练配置使用时嵌入模型。脚本不进行 resize,图片必须已切为目标尺寸。
|
||||||
伪框是像素 `cx,cy,w,h`,不表示缺陷位置。脚本不进行 resize,图片必须已切为目标尺寸。
|
输出名为 `probabilities`,形状 `[1,C]`:九类模型是 `[1,9]`,索引直接对应 `classes.txt`。
|
||||||
|
每类概率在 `[0,1]` 内,概率和为 1;normal 也返回真实概率,不清零、不做阈值过滤。
|
||||||
|
|
||||||
部署时成对替换 XML/BIN;使用元数据的程序也需同步 `metadata.yaml`。
|
部署端直接对 C 个概率取 argmax;normal 是否表示正常、是否应用阈值由业务层决定。
|
||||||
仓库的 `deploy/cpp` 解码器可直接读取适配模型,生产代码若有额外硬编码则需核对。
|
现有 `deploy/cpp` 是 YOLO 框解码器,不能用于此分类输出;需要按 `[1,C]` 读取,
|
||||||
Softmax 分数不同于 YOLO 置信度,需在实际数据上确认过滤阈值。
|
不要跳过前四个通道或遍历候选框。模型附带 `metadata.yaml`、`classes.txt` 和 `export.json`。
|
||||||
默认 top1 为 `normal` 时,所有候选类别分数清零,C++ 在正数置信度阈值下返回空列表,
|
|
||||||
与 YOLO 的背景行为一致。类别编号保持不变,`normal` 仍占其原通道。
|
|
||||||
用 `--background-class` 指定其他背景类别,或 `--keep-background-box` 恢复普通分类输出。
|
|
||||||
原图被切成多片时,每片独立分类,仅前景片返回伪框。
|
|
||||||
输出目录必须是新目录,原有模型不会被覆盖。
|
输出目录必须是新目录,原有模型不会被覆盖。
|
||||||
|
默认压缩权重为 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;模型本身仍返回全部真实概率。
|
||||||
|
报告同时记录各类支持数、总体/前景一致率和逐图结果。
|
||||||
|
|
||||||
## 开发
|
## 开发
|
||||||
|
|
||||||
|
|||||||
@@ -16,6 +16,7 @@ import json
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from datasets import FolderClassificationDataset, classification_collate
|
from datasets import FolderClassificationDataset, classification_collate
|
||||||
|
from datasets.tile_classification import YoloTileClassificationDataset
|
||||||
from engine import ClassificationTrainer
|
from engine import ClassificationTrainer
|
||||||
from models import build_model
|
from models import build_model
|
||||||
from utils import load_config, seed_everything, setup_logging
|
from utils import load_config, seed_everything, setup_logging
|
||||||
@@ -40,6 +41,13 @@ def build_loaders(cfg, split_cfg, classes=None):
|
|||||||
|
|
||||||
datasets = {}
|
datasets = {}
|
||||||
for split, transforms_key in (("train", "train_transforms"), ("val", "val_transforms")):
|
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(
|
datasets[split] = FolderClassificationDataset(
|
||||||
root=split_cfg["root"],
|
root=split_cfg["root"],
|
||||||
split=split,
|
split=split,
|
||||||
|
|||||||
@@ -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]
|
||||||
@@ -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]
|
||||||
@@ -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]}}
|
||||||
@@ -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()
|
||||||
@@ -1,10 +1,8 @@
|
|||||||
#!/usr/bin/env python
|
#!/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.
|
Training normalization is embedded. No resize, pseudo boxes or background gating.
|
||||||
Input must already be RGB float32 BCHW in [0, 1], exactly as in deploy/cpp.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
@@ -19,7 +17,6 @@ import yaml
|
|||||||
from torch import nn
|
from torch import nn
|
||||||
|
|
||||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||||
|
|
||||||
from models import build_model # noqa: E402
|
from models import build_model # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
@@ -36,63 +33,48 @@ def ordered_names(names) -> list[str]:
|
|||||||
return names
|
return names
|
||||||
|
|
||||||
|
|
||||||
def class_order(source: list[str], target: list[str]) -> list[int]:
|
def input_size(config: dict) -> list[int]:
|
||||||
source, target = ordered_names(source), ordered_names(target)
|
size = config["data"]["image_size"]
|
||||||
if set(source) != set(target):
|
size = [size, size] if isinstance(size, int) else list(size)
|
||||||
raise ValueError(f"Class sets differ: classifier={source}, YOLO={target}")
|
if len(size) != 2 or any(not isinstance(v, int) or v < 1 for v in size):
|
||||||
return [source.index(name) for name in target]
|
raise ValueError(f"Invalid training image_size: {size}")
|
||||||
|
return size
|
||||||
|
|
||||||
|
|
||||||
def preprocessing(config: dict, size: int) -> float | None:
|
def preprocessing(config: dict, size: int | list[int]) -> float | None:
|
||||||
"""Reject preprocessing that cannot be reproduced without resampling the input."""
|
shape = [size, size] if isinstance(size, int) else list(size)
|
||||||
transforms = config["data"]["val_transforms"]
|
transforms = config["data"]["val_transforms"]
|
||||||
sequence = [t["name"] for t in transforms]
|
sequence = [t["name"] for t in transforms]
|
||||||
if sequence not in (["Resize", "ToTensor", "PerImageStandardize"],
|
if sequence not in (["Resize", "ToTensor", "PerImageStandardize"],
|
||||||
["ToTensor", "PerImageStandardize"], ["Resize", "ToTensor"], ["ToTensor"]):
|
["ToTensor", "PerImageStandardize"], ["Resize", "ToTensor"], ["ToTensor"]):
|
||||||
raise ValueError(f"Unsupported validation preprocessing: {sequence}")
|
raise ValueError(f"Unsupported validation preprocessing: {sequence}")
|
||||||
|
eps = None
|
||||||
for transform in transforms:
|
for transform in transforms:
|
||||||
if transform["name"] == "Resize":
|
if transform["name"] == "Resize" and list(transform["size"]) != shape:
|
||||||
if transform["size"] != [size, size]:
|
raise ValueError("Training size differs from export input")
|
||||||
raise ValueError("Training size differs from template; retrain at the target size")
|
|
||||||
if transform["name"] == "PerImageStandardize":
|
if transform["name"] == "PerImageStandardize":
|
||||||
return float(transform.get("eps", 1e-5))
|
eps = float(transform.get("eps", 1e-5))
|
||||||
return None
|
if eps <= 0:
|
||||||
|
raise ValueError("Standardization eps must be positive")
|
||||||
|
return eps
|
||||||
|
|
||||||
|
|
||||||
class YoloClassifierAdapter(nn.Module):
|
class ClassifierProbabilities(nn.Module):
|
||||||
def __init__(self, classifier: nn.Module, size: int, anchors: int,
|
def __init__(self, classifier: nn.Module, eps: float | None = None):
|
||||||
order: list[int], eps: float | None = 1e-5,
|
|
||||||
background_index: int | None = None):
|
|
||||||
super().__init__()
|
super().__init__()
|
||||||
if size <= 0 or anchors <= 0:
|
self.classifier, self.eps = classifier, eps
|
||||||
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))
|
|
||||||
|
|
||||||
def normalize(self, x: torch.Tensor) -> torch.Tensor:
|
def normalize(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
if self.eps is None:
|
if self.eps is None:
|
||||||
return x
|
return x
|
||||||
dims = (1, 2, 3)
|
dims = (1, 2, 3)
|
||||||
mean = x.mean(dim=dims, keepdim=True)
|
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)
|
variance = ((x - mean) ** 2).sum(dim=dims, keepdim=True) / (x[0].numel() - 1)
|
||||||
return (x - mean) / (torch.sqrt(variance) + self.eps)
|
return (x - mean) / (torch.sqrt(variance) + self.eps)
|
||||||
|
|
||||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
scores = self.classifier(self.normalize(x)).softmax(dim=1)[:, self.order]
|
return self.classifier(self.normalize(x)).softmax(dim=1)
|
||||||
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)
|
|
||||||
|
|
||||||
|
|
||||||
def main() -> int:
|
def main() -> int:
|
||||||
@@ -100,97 +82,71 @@ def main() -> int:
|
|||||||
|
|
||||||
parser = argparse.ArgumentParser(description=__doc__)
|
parser = argparse.ArgumentParser(description=__doc__)
|
||||||
parser.add_argument("--checkpoint", default="runs/classify/best.pth")
|
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="Training class order: classes.txt or eval_report.json")
|
||||||
parser.add_argument("--classes", help="Classifier class order: classes.txt or eval_report.json")
|
parser.add_argument("--fp32-weights", action="store_true", help="Keep full-precision weights")
|
||||||
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("--out", required=True, help="New directory; existing directories are rejected")
|
parser.add_argument("--out", required=True, help="New directory; existing directories are rejected")
|
||||||
args = parser.parse_args()
|
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():
|
if out.exists():
|
||||||
raise FileExistsError(f"Refusing to overwrite {out}")
|
raise FileExistsError(f"Refusing to overwrite {out}")
|
||||||
state = torch.load(checkpoint, map_location="cpu", weights_only=False)
|
state = torch.load(checkpoint, map_location="cpu", weights_only=False)
|
||||||
config = state["config"]
|
config = state["config"]
|
||||||
|
names = config["data"].get("classes")
|
||||||
if args.classes:
|
if args.classes:
|
||||||
path = Path(args.classes)
|
path = Path(args.classes)
|
||||||
source = (json.loads(path.read_text())["classes"] if path.suffix == ".json"
|
supplied = (json.loads(path.read_text())["classes"] if path.suffix == ".json"
|
||||||
else path.read_text().splitlines())
|
else path.read_text().splitlines())
|
||||||
else:
|
if names and ordered_names(supplied) != ordered_names(names):
|
||||||
source = config["data"].get("classes")
|
raise ValueError("--classes contradicts checkpoint class order")
|
||||||
if not source:
|
names = supplied
|
||||||
|
if not names:
|
||||||
report = checkpoint.parent / "eval_report.json"
|
report = checkpoint.parent / "eval_report.json"
|
||||||
if not report.exists():
|
if not report.exists():
|
||||||
raise ValueError("No recorded class order; provide --classes (do not guess folder order)")
|
raise ValueError("No recorded class order; provide --classes (do not guess folder order)")
|
||||||
source = json.loads(report.read_text())["classes"]
|
names = json.loads(report.read_text())["classes"]
|
||||||
source = ordered_names(source)
|
names = ordered_names(names)
|
||||||
metadata = yaml.safe_load((template.parent / "metadata.yaml").read_text())
|
size = input_size(config)
|
||||||
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]
|
|
||||||
eps = preprocessing(config, size)
|
eps = preprocessing(config, size)
|
||||||
model_config = copy.deepcopy(config["model"])
|
model_config = copy.deepcopy(config["model"])
|
||||||
model_config["backbone"]["pretrained"] = False
|
model_config["backbone"]["pretrained"] = False
|
||||||
model_config["init_from"] = None
|
model_config["init_from"] = None
|
||||||
model_config["num_classes"] = len(source)
|
model_config["num_classes"] = len(names)
|
||||||
model = build_model(model_config).eval()
|
model = build_model(model_config).eval()
|
||||||
model.load_state_dict(state["model"], strict=True)
|
model.load_state_dict(state["model"], strict=True)
|
||||||
adapter = YoloClassifierAdapter(model, size, output_shape[2], order, eps,
|
wrapper = ClassifierProbabilities(model, eps).eval()
|
||||||
background_index=background_index).eval()
|
shape = [1, model_config["backbone"].get("in_channels", 3), *size]
|
||||||
torch.manual_seed(42)
|
torch.manual_seed(42)
|
||||||
example = torch.rand(input_shape)
|
example = torch.rand(shape)
|
||||||
converted = ov.convert_model(adapter, example_input=example, input=input_shape)
|
converted = ov.convert_model(wrapper, example_input=example, input=shape)
|
||||||
converted.input().get_tensor().set_names(reference.input().get_names())
|
converted.input().get_tensor().set_names({"images"})
|
||||||
converted.output().get_tensor().set_names(reference.output().get_names())
|
converted.output().get_tensor().set_names({"probabilities"})
|
||||||
# Preserve any interface metadata used by OpenVINO consumers, then mark pseudo boxes.
|
converted.set_rt_info("classify", ["model_info", "task"])
|
||||||
for key, value in reference.get_rt_info().items():
|
converted.set_rt_info(json.dumps(dict(enumerate(names))), ["model_info", "names"])
|
||||||
if key != "conversion_parameters":
|
|
||||||
converted.set_rt_info(value, [key])
|
|
||||||
converted.set_rt_info("classifier_pseudo_box", ["adapter_kind"])
|
|
||||||
out.mkdir(parents=True)
|
out.mkdir(parents=True)
|
||||||
xml = out / template.name
|
arch = model_config["backbone"].get("arch", model_config["backbone"]["name"])
|
||||||
ov.save_model(converted, str(xml), compress_to_fp16=True)
|
xml = out / f"{arch}.xml"
|
||||||
compiled = core.compile_model(str(xml), "CPU", {"PERFORMANCE_HINT": "LATENCY"})
|
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 = []
|
errors = []
|
||||||
with torch.inference_mode():
|
with torch.inference_mode():
|
||||||
for sample in (example, torch.zeros_like(example), torch.ones_like(example)):
|
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()]
|
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)
|
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 actual.argmax(axis=1).tolist() == expected.argmax(axis=1).tolist()
|
||||||
assert np.count_nonzero(actual[:, :, 1:]) == 0
|
|
||||||
errors.append(float(np.max(np.abs(expected - actual))))
|
errors.append(float(np.max(np.abs(expected - actual))))
|
||||||
metadata.update({"names": dict(enumerate(target)), "adapter_kind": "classifier_pseudo_box",
|
metadata = {"task": "classify", "names": dict(enumerate(names)), "imgsz": size,
|
||||||
"description": "Classifier with one full-image pseudo box; no localization",
|
"channels": shape[1], "layout": "BCHW", "input_range": [0, 1],
|
||||||
"background_class": None if background_index is None else args.background_class,
|
"output": "softmax_probabilities", "input_shape": shape,
|
||||||
"end2end": False, "args": {"nms": False}, "imgsz": [size, size]})
|
"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 / "metadata.yaml").write_text(yaml.safe_dump(metadata, allow_unicode=True))
|
||||||
(out / "classes.txt").write_text("\n".join(target) + "\n")
|
(out / "classes.txt").write_text("\n".join(names) + "\n")
|
||||||
summary = {"checkpoint": str(checkpoint.resolve()), "template": str(template.resolve()),
|
summary = {"checkpoint": str(checkpoint.resolve()), **metadata, "max_abs_errors": errors}
|
||||||
"classifier_classes": source, "output_classes": target, "class_order": order,
|
(out / "export.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2))
|
||||||
"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))
|
|
||||||
print(json.dumps({"model": str(xml), **summary}, ensure_ascii=False, indent=2))
|
print(json.dumps({"model": str(xml), **summary}, ensure_ascii=False, indent=2))
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
|
|||||||
@@ -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)
|
|
||||||
@@ -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)
|
||||||
@@ -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
|
||||||
@@ -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")
|
||||||
Reference in New Issue
Block a user