修改转换代码

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)) ):
+# 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")
+#     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)) )
+#     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'
+# 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'
+# 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)
 
 
 

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"