From bb35a8fcffbecbd6660cdfc13a5fc9464dc3d05f Mon Sep 17 00:00:00 2001 From: Ankle Date: Tue, 26 May 2026 13:28:05 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E8=BD=AC=E6=8D=A2=E4=BB=A3?= =?UTF-8?q?=E7=A0=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- gen_wts_yoloV5.py => pt_to_cfg_wts.py | 24 +++++--- pt_to_cfg_wts.sh | 86 +++++++++++++++++++++++++++ 2 files changed, 103 insertions(+), 7 deletions(-) rename gen_wts_yoloV5.py => pt_to_cfg_wts.py (89%) create mode 100644 pt_to_cfg_wts.sh diff --git a/gen_wts_yoloV5.py b/pt_to_cfg_wts.py similarity index 89% rename from gen_wts_yoloV5.py rename to pt_to_cfg_wts.py index ebddf51..63fa03f 100644 --- a/gen_wts_yoloV5.py +++ b/pt_to_cfg_wts.py @@ -304,11 +304,21 @@ now = datetime.now() time = now.strftime("%Y%m") -if not os.path.exists("/data6/zhangw/yolov5/Detection_onnx/{}/{}/".format(time,str(model_name)) ): - print("/data6/zhangw/yolov5/Detection_onnx/{}/{}/".format(time,str(model_name))," is created") - os.makedirs("/data6/zhangw/yolov5/Detection_onnx/{}/{}/".format(time,str(model_name)) ) -wts_file = f"/data6/zhangw/yolov5/Detection_onnx/{time}/" + model_name + '.wts' if 'yolov55' in model_name else "/data6/zhangw/yolov5/Detection_onnx/{}/{}/".format(time,str(model_name)) + 'yolov5_' + model_name + '.wts' -cfg_file = f"/data6/zhangw/yolov5/Detection_onnx/{time}" + model_name + '.cfg' if 'yolov55' in model_name else "/data6/zhangw/yolov5/Detection_onnx/{}/{}/".format(time,str(model_name)) + 'yolov5_' + model_name + '.cfg' +# if not os.path.exists("/data6/zhangw/yolov5/Detection_onnx/{}/{}/".format(time,str(model_name)) ): +# print("/data6/zhangw/yolov5/Detection_onnx/{}/{}/".format(time,str(model_name))," is created") +# os.makedirs("/data6/zhangw/yolov5/Detection_onnx/{}/{}/".format(time,str(model_name)) ) +# wts_file = f"/data6/zhangw/yolov5/Detection_onnx/{time}/" + model_name + '.wts' if 'yolov55' in model_name else "/data6/zhangw/yolov5/Detection_onnx/{}/{}/".format(time,str(model_name)) + 'yolov5_' + model_name + '.wts' +# cfg_file = f"/data6/zhangw/yolov5/Detection_onnx/{time}" + model_name + '.cfg' if 'yolov55' in model_name else "/data6/zhangw/yolov5/Detection_onnx/{}/{}/".format(time,str(model_name)) + 'yolov5_' + model_name + '.cfg' +save_dir = os.environ.get("YOLOV5_CFG_WTS_SAVE_DIR") + +if save_dir is None: + raise SystemExit("ERROR: YOLOV5_CFG_WTS_SAVE_DIR is not set. Please run this script through the bash wrapper.") + +os.makedirs(save_dir, exist_ok=True) + +wts_file = os.path.join(save_dir, "yolov5_" + model_name + ".wts") +cfg_file = os.path.join(save_dir, "yolov5_" + model_name + ".cfg") + device = select_device('cpu') model = torch.load(pt_file, map_location=device)['model'].float() @@ -345,8 +355,8 @@ with open(wts_file, 'w') as fw, open(cfg_file, 'w') as fc: raise SystemExit('Model not supported') print("ok") os.system('echo "%d" | cat - %s > temp && mv temp %s' % (layers.wc, wts_file, wts_file)) -print("save" + f"/data6/zhangw/yolov5/Detection_onnx/{time}/" + model_name) - +# print("save" + f"/data6/zhangw/yolov5/Detection_onnx/{time}/" + model_name) +print("save:", save_dir) diff --git a/pt_to_cfg_wts.sh b/pt_to_cfg_wts.sh new file mode 100644 index 0000000..e73206d --- /dev/null +++ b/pt_to_cfg_wts.sh @@ -0,0 +1,86 @@ +#!/bin/bash +set -e + +# ========================== +# 固定配置 +# ========================== + +# cfg/wts 总输出根目录（和 yolov5_birdproject 平级） +RESULT_ROOT="/data2/zhangw_data/yolov5_result_cts_wts" + +# Python 转换脚本名 +PY_SCRIPT="pt_to_cfg_wts.py" + +# 输入尺寸 +IMG_SIZE=960 + +# ========================== +# 命令行参数 +# ========================== + +PT_PATH="$1" +RESULT_DATE="$2" +RESERVE_NAME="$3" + +# ========================== +# 参数检查 +# ========================== + +if [ -z "$PT_PATH" ] || [ -z "$RESULT_DATE" ] || [ -z "$RESERVE_NAME" ]; then + echo "用法：" + echo "bash pt_to_cfg_wts.sh 权重.pt 日期 保护区简称" + echo "" + echo "示例：" + echo "bash pt_to_cfg_wts.sh \\" + echo "/data2/zhangw_data/yolov5_birdproject/runs/train/exp/weights/best.pt \\" + echo "20260526 \\" + echo "HHK" + exit 1 +fi + +if [ ! -f "$PT_PATH" ]; then + echo "错误：找不到权重文件：$PT_PATH" + exit 1 +fi + +if [ ! -f "$PY_SCRIPT" ]; then + echo "错误：找不到Python转换脚本：$PY_SCRIPT" + echo "请确认当前目录是 yolov5_birdproject 根目录" + exit 1 +fi + +# ========================== +# 输出目录 +# ========================== + +RESULT_DIR="${RESULT_ROOT}/yolov5_result_cts_wts_${RESULT_DATE}_${RESERVE_NAME}" + +mkdir -p "$RESULT_DIR" + +export YOLOV5_CFG_WTS_SAVE_DIR="$RESULT_DIR" + +PT_BASENAME=$(basename "$PT_PATH") + +# ========================== +# 开始转换 +# ========================== + +echo "========== YOLOv5 PT 转 CFG/WTS ==========" +echo "权重文件：$PT_PATH" +echo "输入尺寸：${IMG_SIZE} x ${IMG_SIZE}" +echo "结果目录：$RESULT_DIR" +echo "日期：$RESULT_DATE" +echo "保护区简称：$RESERVE_NAME" +echo "=========================================" + +python "$PY_SCRIPT" \ + --weights "$PT_PATH" \ + --size "$IMG_SIZE" "$IMG_SIZE" \ + --name "$PT_BASENAME" + +echo "" +echo "========== 转换完成 ==========" +echo "输出目录：$RESULT_DIR" +echo "" + +ls -lh "$RESULT_DIR" \ No newline at end of file