@echo off
chcp 65001 >nul
echo ============================================
echo YOLO INT8 Quantization Complete Workflow
echo ============================================
echo.

REM 激活虚拟环境
echo [1/4] Activating yolo virtual environment...
call conda activate yolo
echo.

REM 步骤1: 检查并导出ONNX模型
echo [2/4] Checking ONNX model...
if not exist yolo11n.onnx (
    echo ONNX model not found. Exporting from yolo11n.pt...
    if not exist yolo11n.pt (
        echo ERROR: yolo11n.pt not found!
        echo Please place yolo11n.pt in the current directory.
        pause
        exit /b 1
    )
    python -c "from ultralytics import YOLO; model = YOLO('yolo11n.pt'); model.export(format='onnx', dynamic=True, simplify=True, opset=12)"
    if not exist yolo11n.onnx (
        echo ERROR: ONNX export failed!
        pause
        exit /b 1
    )
    echo ONNX model exported successfully.
) else (
    echo ONNX model already exists: yolo11n.onnx
)
echo.

REM 步骤2: 生成校准缓存
echo [3/4] Generating calibration cache...
if not exist yolo11n_int8.cache (
    echo Running calibration_gen.py...
    python calibration_gen.py
    if not exist yolo11n_int8.cache (
        echo ERROR: Calibration cache generation failed!
        pause
        exit /b 1
    )
    echo Calibration cache generated: yolo11n_int8.cache
) else (
    echo Calibration cache already exists: yolo11n_int8.cache
)
echo.

REM 步骤3: 构建TensorRT Engine
echo [4/4] Building TensorRT Engine...
echo.
echo Running trtexec with INT8 calibration...
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:1x3x640x640 ^
   --optShapes=input:4x3x640x640 ^
   --maxShapes=input:8x3x640x640 ^
   --useCudaGraph ^
   --useSpinWait ^
   --noTF32

if %errorlevel% equ 0 (
    echo.
    echo ============================================
    echo TensorRT Engine built successfully!
    echo Output: yolo11n_int8_b1_8.engine
    echo ============================================
    echo.
    echo Next steps:
    echo 1. Run: python quantize_yolo.py
    echo    to validate models and calculate mAP drop
    echo.
) else (
    echo.
    echo ============================================
    echo Engine build failed!
    ============================================
    pause
    exit /b 1
)

pause