#!/usr/bin/env python """用导出的 OpenVINO 模型推理，不依赖 ultralytics 与 torch。 # 已经切好的 160 片：直接推理 uv run python scripts/openvino_predict.py \\ --model runs/detect/current_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/current_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())