Shaun ZhouandClaude Opus 5
c2bd5301b4
之前 160 的 OpenVINO 导出直接从 current.pt 导，文件名不带边长也不带时间戳， 同时留着 128 和 160 两套时分不清谁是谁，软链接又会被下次训练改写。 - 新增 scripts/export_openvino.py：把权重归档成 film_<边长>_<时间戳>.pt 并导出 同名的 OpenVINO 目录。边长从 checkpoint 读、时间戳从训练目录名取，所以同一次 训练无论何时归档名字都一样，也不会出现文件名与实际输入尺寸对不上。 - 重训脚本第 5 步自动归档；训练目录与切片数据集的名字也带上边长。 - openvino_predict.py 的示例路径跟着更新。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
247 lines
10 KiB
Python
247 lines
10 KiB
Python
#!/usr/bin/env python
|
||
"""用导出的 OpenVINO 模型推理，不依赖 ultralytics 与 torch。
|
||
|
||
# 已经切好的 160 片：直接推理
|
||
uv run python scripts/openvino_predict.py \\
|
||
--model runs/detect/models/film_160_20260924_0817_openvino_model \\
|
||
--images data/film_yolo_160/images/val --out /tmp/ov_pred
|
||
|
||
# 任意尺寸的原图：加 --tile-size，脚本自己切片再拼回
|
||
uv run python scripts/openvino_predict.py \\
|
||
--model runs/detect/models/film_160_20260924_0817_openvino_model \\
|
||
--images runs/boxes_shap_label_studio/images --tile-size 160 --out /tmp/ov_pred
|
||
|
||
产出 YOLO 格式的标签（`class cx cy w h conf`，坐标归一化），`--labels` 指向真值目录时
|
||
还会给出逐类的 TP/FP/FN。
|
||
|
||
**拷到别的设备上跑，需要这三个文件**：本脚本、`engine/tiling.py`、`engine/detection.py`
|
||
（后两者只依赖 numpy 与标准库），外加 `pip install openvino numpy pillow`。
|
||
|
||
推理契约（实测确认）：
|
||
|
||
- 输入 `[1, 3, N, N]` float32，**RGB** 顺序，除以 255，无 mean/std 标准化；
|
||
- 输出 `[1, 4 + 类别数, anchor 数]`：前 4 行是 **xywh 中心点、像素单位**（可能超出边界，
|
||
要裁剪），后面每行是一个类别的分数，已过 sigmoid；
|
||
- 导出的 `metadata.yaml` 里 `end2end: false`，**必须自己做 NMS**——实测同一个目标会出
|
||
3 个候选框。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import sys
|
||
from collections import Counter
|
||
from pathlib import Path
|
||
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||
|
||
import numpy as np # noqa: E402
|
||
|
||
from engine.detection import Box, match_boxes, nms, parse_label # noqa: E402
|
||
from engine.tiling import cut, merge_tiles, plan_tiles # noqa: E402
|
||
from utils import setup_logging # noqa: E402
|
||
|
||
logger = None
|
||
|
||
IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff")
|
||
|
||
|
||
class OpenVinoDetector:
|
||
"""加载导出目录里的 .xml/.bin 与 metadata.yaml。"""
|
||
|
||
def __init__(self, model_dir: Path, device: str = "CPU") -> None:
|
||
import openvino as ov
|
||
import yaml
|
||
|
||
xml = next(model_dir.glob("*.xml"), None)
|
||
if xml is None:
|
||
raise FileNotFoundError(f"{model_dir} 下没有 .xml，这不像 OpenVINO 的导出目录")
|
||
core = ov.Core()
|
||
self.compiled = core.compile_model(core.read_model(str(xml)), device)
|
||
self.output = self.compiled.output(0)
|
||
|
||
meta_file = model_dir / "metadata.yaml"
|
||
meta = yaml.safe_load(meta_file.read_text(encoding="utf-8")) if meta_file.exists() else {}
|
||
names = meta.get("names") or {}
|
||
self.classes = [names[i] for i in sorted(names)] if names else []
|
||
imgsz = meta.get("imgsz") or []
|
||
# 以网络自身的输入形状为准，metadata 只是参考
|
||
shape = self.compiled.input(0).get_partial_shape()
|
||
self.size = int(shape[2].get_length()) if shape[2].is_static else (imgsz[0] if imgsz else 0)
|
||
|
||
def __call__(self, patches, conf: float, iou: float) -> list[list[Box]]:
|
||
"""一批 size×size 的图 -> 每张图的归一化框（已做 NMS）。"""
|
||
out = []
|
||
for patch in patches:
|
||
raw = self.compiled([self._preprocess(patch)])[self.output][0] # (4+nc, anchors)
|
||
out.append(self._decode(raw, conf, iou))
|
||
return out
|
||
|
||
def _preprocess(self, patch: np.ndarray) -> np.ndarray:
|
||
"""HWC uint8 -> BCHW float32 RGB，除以 255。灰度图补成三通道。"""
|
||
if patch.ndim == 2:
|
||
patch = np.repeat(patch[..., None], 3, axis=2)
|
||
elif patch.shape[2] == 1:
|
||
patch = np.repeat(patch, 3, axis=2)
|
||
elif patch.shape[2] == 4:
|
||
patch = patch[..., :3]
|
||
return patch.astype(np.float32).transpose(2, 0, 1)[None] / 255.0
|
||
|
||
def _decode(self, raw: np.ndarray, conf: float, iou: float) -> list[Box]:
|
||
"""原始输出 -> 归一化 xyxy 框。坐标是像素 xywh，分数已过 sigmoid。"""
|
||
xywh, scores = raw[:4], raw[4:]
|
||
class_ids, confidences = scores.argmax(0), scores.max(0)
|
||
keep = confidences >= conf
|
||
if not keep.any():
|
||
return []
|
||
cx, cy, w, h = xywh[:, keep]
|
||
size = self.size
|
||
boxes = [
|
||
Box(int(class_id),
|
||
float(np.clip((x - width / 2) / size, 0.0, 1.0)),
|
||
float(np.clip((y - height / 2) / size, 0.0, 1.0)),
|
||
float(np.clip((x + width / 2) / size, 0.0, 1.0)),
|
||
float(np.clip((y + height / 2) / size, 0.0, 1.0)),
|
||
float(score))
|
||
for x, y, width, height, class_id, score in zip(
|
||
cx, cy, w, h, class_ids[keep], confidences[keep], strict=True,
|
||
)
|
||
]
|
||
return [boxes[i] for i in nms(boxes, iou)]
|
||
|
||
|
||
def load_image(path: Path) -> np.ndarray:
|
||
from PIL import Image
|
||
|
||
with Image.open(path) as handle:
|
||
image = handle.convert("L") if handle.mode in ("L", "1", "I;16") else handle.convert("RGB")
|
||
return np.asarray(image)
|
||
|
||
|
||
def main() -> int:
|
||
global logger
|
||
parser = argparse.ArgumentParser(
|
||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
|
||
)
|
||
parser.add_argument("--model", required=True, help="OpenVINO 导出目录（含 .xml/.bin）")
|
||
parser.add_argument("--images", required=True, help="图片目录")
|
||
parser.add_argument("--out", default=None, help="把预测写成 YOLO 标签的目录")
|
||
parser.add_argument("--labels", default=None, help="真值标签目录，给了就统计 TP/FP/FN")
|
||
parser.add_argument("--device", default="CPU", help="CPU / GPU / AUTO")
|
||
parser.add_argument("--conf", type=float, default=0.25)
|
||
parser.add_argument("--iou", type=float, default=0.5, help="NMS 的 IoU 阈值")
|
||
parser.add_argument(
|
||
"--tile-size", type=int, default=0,
|
||
help="图片大于模型输入时需要切片；图片本身就是模型输入尺寸时留 0",
|
||
)
|
||
parser.add_argument("--tile-overlap", type=int, default=32)
|
||
parser.add_argument("--limit", type=int, default=0)
|
||
args = parser.parse_args()
|
||
|
||
logger = setup_logging()
|
||
images = sorted(
|
||
p for p in Path(args.images).iterdir() if p.suffix.lower() in IMAGE_SUFFIXES
|
||
)
|
||
if args.limit:
|
||
images = images[: args.limit]
|
||
if not images:
|
||
logger.error("%s 下没有图片", args.images)
|
||
return 1
|
||
|
||
try:
|
||
detector = OpenVinoDetector(Path(args.model), args.device)
|
||
except FileNotFoundError as exc:
|
||
logger.error("%s", exc)
|
||
return 1
|
||
logger.info(
|
||
"模型 %s（输入 %d，%d 类）| 设备 %s | 图片 %d 张",
|
||
args.model, detector.size, len(detector.classes), args.device, len(images),
|
||
)
|
||
|
||
out_dir = Path(args.out) if args.out else None
|
||
if out_dir:
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
labels_dir = Path(args.labels) if args.labels else None
|
||
|
||
per_class: Counter[str] = Counter()
|
||
stats = {"tp": Counter(), "fp": Counter(), "fn": Counter()}
|
||
with_boxes = 0
|
||
mismatched = 0
|
||
|
||
for index, path in enumerate(images, start=1):
|
||
image = load_image(path)
|
||
height, width = image.shape[:2]
|
||
|
||
if args.tile_size > 0:
|
||
tiles = plan_tiles(width, height, args.tile_size, args.tile_overlap)
|
||
patches = [cut(image, tile, args.tile_size) for tile in tiles]
|
||
boxes = merge_tiles(
|
||
detector(patches, args.conf, args.iou), tiles, args.tile_size, width, height
|
||
)
|
||
else:
|
||
if (height, width) != (detector.size, detector.size):
|
||
mismatched += 1
|
||
continue
|
||
boxes = detector([image], args.conf, args.iou)[0]
|
||
|
||
with_boxes += bool(boxes)
|
||
for box in boxes:
|
||
per_class[detector.classes[box.class_id] if detector.classes else box.class_id] += 1
|
||
|
||
if out_dir:
|
||
lines = [
|
||
f"{b.class_id} {(b.x1 + b.x2) / 2:.6f} {(b.y1 + b.y2) / 2:.6f} "
|
||
f"{b.x2 - b.x1:.6f} {b.y2 - b.y1:.6f} {b.score:.4f}"
|
||
for b in boxes
|
||
]
|
||
(out_dir / f"{path.stem}.txt").write_text(
|
||
"".join(f"{line}\n" for line in lines), encoding="utf-8"
|
||
)
|
||
|
||
if labels_dir:
|
||
truth_file = labels_dir / f"{path.stem}.txt"
|
||
truth = parse_label(truth_file.read_text(encoding="utf-8")) if truth_file.exists() else []
|
||
match = match_boxes(boxes, truth, iou_threshold=0.5)
|
||
name = lambda i: detector.classes[i] if detector.classes else i # noqa: E731
|
||
for _, gt_index in match.pairs:
|
||
stats["tp"][name(truth[gt_index].class_id)] += 1
|
||
for i in match.unmatched_predictions:
|
||
stats["fp"][name(boxes[i].class_id)] += 1
|
||
for i in match.unmatched_ground_truth:
|
||
stats["fn"][name(truth[i].class_id)] += 1
|
||
|
||
if index % 2000 == 0:
|
||
logger.info("处理 [%d/%d]", index, len(images))
|
||
|
||
if mismatched:
|
||
logger.warning(
|
||
"%d 张图不是 %dx%d，已跳过——它们需要切片，加 --tile-size %d",
|
||
mismatched, detector.size, detector.size, detector.size,
|
||
)
|
||
logger.info(
|
||
"完成：%d 张，有框 %d，框数 %d", len(images) - mismatched, with_boxes, sum(per_class.values())
|
||
)
|
||
logger.info("逐类框数: %s", dict(sorted(per_class.items(), key=lambda kv: -kv[1])))
|
||
|
||
if labels_dir:
|
||
logger.info("%-12s %8s %8s %8s %8s", "类别", "TP", "FP", "FN", "召回")
|
||
for name in detector.classes:
|
||
tp, fp, fn = (stats[k][name] for k in ("tp", "fp", "fn"))
|
||
if tp or fp or fn:
|
||
logger.info(
|
||
"%-12s %8d %8d %8d %7.1f%%", name, tp, fp, fn, 100 * tp / max(tp + fn, 1)
|
||
)
|
||
if out_dir:
|
||
(out_dir / "summary.json").write_text(
|
||
json.dumps(
|
||
{k: dict(v) for k, v in stats.items()}, ensure_ascii=False, indent=2
|
||
),
|
||
encoding="utf-8",
|
||
)
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|