#!/usr/bin/env bash
# ATC 模型转换脚本
# 在高性能电脑（CANN 环境）上执行
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
cd "$SCRIPT_DIR"

MODEL=${MODEL:-yolov8n}
INPUT_SIZE=${INPUT_SIZE:-640}
OUTPUT_TYPE=${OUTPUT_TYPE:-FP16}

echo "=========================================="
echo "  ATC ONNX -> OM 模型转换"
echo "=========================================="
echo "模型: $MODEL.onnx"
echo "输入尺寸: ${INPUT_SIZE}x${INPUT_SIZE}"
echo "输出精度: $OUTPUT_TYPE"
echo ""

# 限制并行度，避免内存溢出
export TBE_IMPL_MODE=AI_CORE
export TE_PARALLEL_COMPILER=1

atc \
  --model=${MODEL}.onnx \
  --framework=5 \
  --output=${MODEL} \
  --soc_version=Ascend310B4 \
  --input_format=NCHW \
  --input_shape="images:1,3,${INPUT_SIZE},${INPUT_SIZE}" \
  --output_type=${OUTPUT_TYPE}

echo ""
if [ -f ${MODEL}.om ]; then
    echo "转换成功!"
    echo "输出: ${MODEL}.om"
    ls -lh ${MODEL}.om
else
    echo "转换失败，请检查日志"
    exit 1
fi