116 lines
3.5 KiB
Python
116 lines
3.5 KiB
Python
# coding=utf-8
|
||
"""导出 YOLO / ConvNeXt 特征模型为 OpenVINO，供无 NVIDIA 的客户机加速。
|
||
|
||
用法（在项目根目录）::
|
||
|
||
pip install openvino onnx
|
||
python tools/export_openvino.py
|
||
python tools/export_openvino.py --yolo-only
|
||
python tools/export_openvino.py --feature-only
|
||
python tools/export_openvino.py --imgsz 416
|
||
|
||
导出后：
|
||
- best_openvino_model/ → 启动时自动被 YOLO 加载
|
||
- feature_openvino_model/ → 启动时自动被特征提取加载
|
||
|
||
环境变量（可选）:
|
||
DISH_YOLO_MODEL / DISH_FEATURE_OPENVINO / DISH_OPENVINO_DEVICE / DISH_YOLO_IMGSZ
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import shutil
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
_ROOT = Path(__file__).resolve().parents[1]
|
||
if str(_ROOT) not in sys.path:
|
||
sys.path.insert(0, str(_ROOT))
|
||
|
||
|
||
def export_yolo(imgsz: int = 640) -> Path:
|
||
from ultralytics import YOLO
|
||
|
||
pt = _ROOT / "best.pt"
|
||
if not pt.exists():
|
||
raise FileNotFoundError(f"找不到权重: {pt}")
|
||
model = YOLO(str(pt))
|
||
out = model.export(format="openvino", imgsz=imgsz, half=False)
|
||
out_path = Path(out)
|
||
# ultralytics 通常输出 best_openvino_model/
|
||
target = _ROOT / "best_openvino_model"
|
||
if out_path.resolve() != target.resolve():
|
||
if target.exists():
|
||
shutil.rmtree(target)
|
||
if out_path.is_dir():
|
||
shutil.move(str(out_path), str(target))
|
||
else:
|
||
target.mkdir(parents=True, exist_ok=True)
|
||
shutil.copy2(out_path, target / out_path.name)
|
||
print(f"[OK] YOLO OpenVINO → {target}")
|
||
return target
|
||
|
||
|
||
def export_feature(imgsz: int = 224) -> Path:
|
||
"""导出 ConvNeXt-Tiny 特征骨干为 OpenVINO（输入 NCHW 已归一化）。"""
|
||
import numpy as np
|
||
import torch
|
||
from openvino import convert_model, save_model
|
||
|
||
from core.common_classifier import FeatureExtractor
|
||
|
||
extractor = FeatureExtractor().eval()
|
||
dummy = torch.randn(1, 3, imgsz, imgsz)
|
||
|
||
onnx_path = _ROOT / "feature_extractor.onnx"
|
||
torch.onnx.export(
|
||
extractor,
|
||
dummy,
|
||
str(onnx_path),
|
||
input_names=["input"],
|
||
output_names=["features"],
|
||
dynamic_axes={"input": {0: "batch"}, "features": {0: "batch"}},
|
||
opset_version=17,
|
||
)
|
||
print(f"[OK] ONNX → {onnx_path}")
|
||
|
||
ov_model = convert_model(str(onnx_path))
|
||
out_dir = _ROOT / "feature_openvino_model"
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
xml_path = out_dir / "feature_extractor.xml"
|
||
save_model(ov_model, str(xml_path))
|
||
print(f"[OK] feature OpenVINO → {xml_path}")
|
||
|
||
# 快速校验
|
||
from openvino import Core
|
||
compiled = Core().compile_model(str(xml_path), "CPU")
|
||
test = np.random.randn(1, 3, imgsz, imgsz).astype(np.float32)
|
||
out = compiled([test])[compiled.outputs[0]]
|
||
print(f"[OK] smoke test output shape={np.asarray(out).shape}")
|
||
return out_dir
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="导出 OpenVINO 模型")
|
||
parser.add_argument("--yolo-only", action="store_true")
|
||
parser.add_argument("--feature-only", action="store_true")
|
||
parser.add_argument("--imgsz", type=int, default=640, help="YOLO 输入边长")
|
||
args = parser.parse_args()
|
||
|
||
do_yolo = not args.feature_only
|
||
do_feat = not args.yolo_only
|
||
if args.yolo_only:
|
||
do_feat = False
|
||
if args.feature_only:
|
||
do_yolo = False
|
||
|
||
if do_yolo:
|
||
export_yolo(imgsz=args.imgsz)
|
||
if do_feat:
|
||
export_feature(imgsz=224)
|
||
print("完成。重启识别程序即可自动加载 OpenVINO 模型。")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|