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>
104 lines
3.5 KiB
Python
104 lines
3.5 KiB
Python
#!/usr/bin/env python
|
||
"""把训练好的权重按统一命名归档，并导出 OpenVINO 格式。
|
||
|
||
uv run python scripts/export_openvino.py --weights runs/detect/film_20260924_0817/weights/best.pt
|
||
|
||
产出：
|
||
|
||
```
|
||
runs/detect/models/
|
||
├── film_<边长>_<时间戳>.pt
|
||
└── film_<边长>_<时间戳>_openvino_model/ # .xml / .bin / metadata.yaml
|
||
```
|
||
|
||
**边长从权重里读**（ultralytics 把训练参数存在 checkpoint 里），不用手填，也就不会出现
|
||
文件名写 128、实际是 160 训练的情况。
|
||
|
||
**时间戳默认取训练目录名里的那一串**（`film_20260924_0817` → `20260924_0817`），
|
||
这样同一次训练无论何时归档，名字都一样、可重复；取不到才退回当前时刻。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import re
|
||
import shutil
|
||
import sys
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||
|
||
from utils import setup_logging # noqa: E402
|
||
|
||
logger = None
|
||
|
||
#: 训练目录名里的时间戳，如 film_20260924_0817 / film128_20260930_0415
|
||
STAMP_PATTERN = re.compile(r"(\d{8}_\d{4})")
|
||
|
||
|
||
def stamp_for(weights: Path) -> str:
|
||
"""从权重路径里找训练时间戳；找不到就用当前时刻。"""
|
||
for part in reversed(weights.resolve().parts):
|
||
found = STAMP_PATTERN.search(part)
|
||
if found:
|
||
return found.group(1)
|
||
return datetime.now().strftime("%Y%m%d_%H%M")
|
||
|
||
|
||
def main() -> int:
|
||
global logger
|
||
parser = argparse.ArgumentParser(
|
||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
|
||
)
|
||
parser.add_argument("--weights", required=True, help="训练产出的 .pt")
|
||
parser.add_argument(
|
||
"--out", default="runs/detect/models", help="归档目录，默认 runs/detect/models",
|
||
)
|
||
parser.add_argument("--stamp", default=None, help="覆盖时间戳，缺省从权重路径推断")
|
||
parser.add_argument("--size", type=int, default=0, help="覆盖输入边长，缺省从权重里读")
|
||
parser.add_argument("--prefix", default="film", help="文件名前缀")
|
||
parser.add_argument(
|
||
"--no-openvino", action="store_true", help="只归档 .pt，不导出 OpenVINO",
|
||
)
|
||
parser.add_argument("--overwrite", action="store_true")
|
||
args = parser.parse_args()
|
||
|
||
logger = setup_logging()
|
||
weights = Path(args.weights)
|
||
if not weights.exists():
|
||
logger.error("找不到权重 %s", weights)
|
||
return 1
|
||
|
||
from ultralytics import YOLO
|
||
|
||
model = YOLO(str(weights))
|
||
size = args.size or int(model.overrides.get("imgsz") or 0)
|
||
if not size:
|
||
logger.error("权重里读不到 imgsz，请用 --size 指定")
|
||
return 1
|
||
|
||
stamp = args.stamp or stamp_for(weights)
|
||
out_dir = Path(args.out)
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
target = out_dir / f"{args.prefix}_{size}_{stamp}.pt"
|
||
|
||
if target.exists() and not args.overwrite:
|
||
logger.error("%s 已存在；确认要覆盖请加 --overwrite", target)
|
||
return 1
|
||
shutil.copy2(weights, target)
|
||
logger.info("归档权重: %s（%.1f MB）", target, target.stat().st_size / 1e6)
|
||
|
||
if args.no_openvino:
|
||
return 0
|
||
|
||
# 从归档后的文件导出，这样 OpenVINO 目录名与 .pt 同名
|
||
exported = YOLO(str(target)).export(format="openvino", imgsz=size, batch=1, dynamic=False)
|
||
logger.info("OpenVINO: %s", exported)
|
||
logger.info("输入固定为 %dx%d；部署时切片的 overlap 需与训练时一致", size, size)
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|