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
co-authored by
Claude Opus 5
Claude Opus 5
parent
bbd5fa167e
commit
c2bd5301b4
@@ -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())
|
||||
@@ -3,12 +3,12 @@
|
||||
|
||||
# 已经切好的 160 片：直接推理
|
||||
uv run python scripts/openvino_predict.py \\
|
||||
--model runs/detect/current_openvino_model \\
|
||||
--model runs/detect/models/film_160_20260924_0817_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 \\
|
||||
--model runs/detect/models/film_160_20260924_0817_openvino_model \\
|
||||
--images runs/boxes_shap_label_studio/images --tile-size 160 --out /tmp/ov_pred
|
||||
|
||||
产出 YOLO 格式的标签（`class cx cy w h conf`，坐标归一化），`--labels` 指向真值目录时
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
# bash scripts/retrain_from_label_studio.sh
|
||||
#
|
||||
# 也由 ML backend 在点击 Start Training 时自动调用（输出见 $RETRAIN_LOG）。
|
||||
# 五步：导出标注 -> 装配数据集 -> 切成固定边长的片 -> 训练 -> 切换权重并重启后端。
|
||||
# 五步：导出标注 -> 装配数据集 -> 切成固定边长的片 -> 训练 -> 切换权重、归档、重启后端。
|
||||
#
|
||||
# 产线要求模型输入固定为 TILE_SIZE×TILE_SIZE 且禁止 resize，所以训练数据必须先切片；
|
||||
# 推理侧（可视化、批量推送、ML backend）也要用同一个尺寸，否则输入分布对不上。
|
||||
@@ -17,10 +17,6 @@ cd "$REPO"
|
||||
PROJECT="${LS_PROJECT_ID:-39}"
|
||||
IMAGES="${LS_IMAGES_DIR:-$REPO/runs/boxes_shap_label_studio/images}"
|
||||
STAMP=$(date +%Y%m%d_%H%M)
|
||||
EXPORT="${LS_EXPORT_DIR:-/tmp}/ls_export_${STAMP}.zip"
|
||||
DATASET="$REPO/data/film_yolo_${STAMP}"
|
||||
TILED="$REPO/data/film_yolo_${STAMP}_tiled"
|
||||
WORK_DIR="$REPO/runs/detect/film_${STAMP}"
|
||||
|
||||
# 切片参数。改尺寸时这里与 configs/detect/film.yaml 的 imgsz 会一起生效（见第 4 步）
|
||||
TILE_SIZE="${TILE_SIZE:-160}"
|
||||
@@ -28,6 +24,17 @@ TILE_OVERLAP="${TILE_OVERLAP:-32}"
|
||||
# >0.5 保证一个框最多出现在一片里，不会在相邻片各留一个残框
|
||||
MIN_VISIBILITY="${MIN_VISIBILITY:-0.6}"
|
||||
|
||||
# 边长写进目录名：同时留着 128 和 160 两套时，一眼能分清谁是谁
|
||||
SIZE_TAG=$([ "$TILE_SIZE" -gt 0 ] && echo "$TILE_SIZE" || echo "raw")
|
||||
EXPORT="${LS_EXPORT_DIR:-/tmp}/ls_export_${STAMP}.zip"
|
||||
DATASET="$REPO/data/film_yolo_${STAMP}"
|
||||
TILED="$REPO/data/film_yolo_${SIZE_TAG}_${STAMP}"
|
||||
WORK_DIR="$REPO/runs/detect/film_${SIZE_TAG}_${STAMP}"
|
||||
# 归档目录：每次训练留一份 film_<边长>_<时间戳>.pt 与同名的 OpenVINO 导出，
|
||||
# 供部署机取用。软链接会被下次训练改写，这里的带戳副本不会。
|
||||
MODELS_DIR="$REPO/runs/detect/models"
|
||||
EXPORT_OPENVINO="${EXPORT_OPENVINO:-1}"
|
||||
|
||||
PYTHON="${REPO_PYTHON:-$REPO/.venv/bin/python}"
|
||||
# 用跑着 Label Studio 的那个环境导出，确保读的是同一个库
|
||||
LS_BIN="${LS_BIN:-/home/deep/.v/label/bin/label-studio}"
|
||||
@@ -96,6 +103,12 @@ else
|
||||
echo " best.pt 保持不变（本次 mAP50 $NEW_MAP < 此前 $OLD_MAP）"
|
||||
fi
|
||||
|
||||
echo "--- 归档 -> $MODELS_DIR/film_${SIZE_TAG}_${STAMP}.pt"
|
||||
ARCHIVE_ARGS=(--weights "$WORK_DIR/weights/best.pt" --out "$MODELS_DIR" --stamp "$STAMP" --overwrite)
|
||||
[ "$EXPORT_OPENVINO" = "1" ] || ARCHIVE_ARGS+=(--no-openvino)
|
||||
# 归档失败不该让整轮重训算失败：权重和软链接此时已经就位
|
||||
"$PYTHON" scripts/export_openvino.py "${ARCHIVE_ARGS[@]}" || echo " 归档失败，跳过（权重仍在 $WORK_DIR/weights/best.pt）"
|
||||
|
||||
if [ "$RESTART_BACKEND" = "1" ]; then
|
||||
echo "--- 重启 ML backend"
|
||||
# 必须显式传 YOLO_WEIGHTS 并清掉继承来的变量：本脚本可能是旧后端的子进程，
|
||||
@@ -112,6 +125,7 @@ fi
|
||||
|
||||
echo "===== [$(date '+%F %T')] 重训完成：$WORK_DIR ====="
|
||||
echo "指标见 $WORK_DIR/metrics.json，曲线见 $WORK_DIR/results.png"
|
||||
echo "归档 -> $MODELS_DIR/film_${SIZE_TAG}_${STAMP}.pt（部署取这份，带戳不会被改写）"
|
||||
echo "current.pt -> $(readlink "$REPO/runs/detect/current.pt")"
|
||||
echo "best.pt -> $(readlink "$REPO/runs/detect/best.pt")"
|
||||
if [ "$TILE_SIZE" -gt 0 ]; then
|
||||
|
||||
Reference in New Issue
Block a user