feat(deploy): 模型归档统一为 film_<边长>_<时间戳>

之前 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>
This commit is contained in:
Shaun Zhou
2026-09-30 15:11:45 +08:00
co-authored by Claude Opus 5
parent bbd5fa167e
commit c2bd5301b4
3 changed files with 124 additions and 7 deletions
+103
View File
@@ -0,0 +1,103 @@
#!/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())