@echo off
chcp 65001 >nul
echo ============================================
echo YOLO INT8 TensorRT Engine Builder
echo ============================================
echo.

REM 激活虚拟环境
echo Activating yolo virtual environment...
call conda activate yolo
echo.

REM 检查ONNX模型
if not exist yolo11n.onnx (
    echo ERROR: yolo11n.onnx not found!
    echo Please export the ONNX model first.
    pause
    exit /b 1
)

REM 检查校准缓存
if not exist yolo11n_int8.cache (
    echo ERROR: yolo11n_int8.cache not found!
    echo Please run calibration_gen.py first to generate the calibration cache.
    pause
    exit /b 1
)

echo Building TensorRT Engine with INT8 calibration...
echo ONNX: yolo11n.onnx
echo Engine: yolo11n_int8_b1_8.engine
echo Cache: yolo11n_int8.cache
echo.

trtexec ^
   --onnx=yolo11n.onnx ^
   --saveEngine=yolo11n_int8_b1_8.engine ^
   --explicitBatch ^
   --int8 ^
   --fp16 ^
   --workspace=4096 ^
   --builderOptimizationLevel=5 ^
   --profilingVerbosity=detailed ^
   --calib=yolo11n_int8.cache ^
   --minShapes=input:1x3x480x640 ^
   --optShapes=input:4x3x640x640 ^
   --maxShapes=input:8x3x640x736 ^
   --useCudaGraph ^
   --useSpinWait ^
   --noTF32

if %errorlevel% equ 0 (
    echo.
    echo ============================================
    echo Engine built successfully!
    echo Output: yolo11n_int8_b1_8.engine
    echo ============================================
) else (
    echo.
    echo ============================================
    echo Engine build failed!
    echo ============================================
    pause
    exit /b 1
)

pause