diff --git a/scripts/openvino_predict.py b/scripts/openvino_predict.py new file mode 100644 index 0000000..bfbe320 --- /dev/null +++ b/scripts/openvino_predict.py @@ -0,0 +1,246 @@ +#!/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())