#!/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"